Jupyter生態二次開發系列(二)

這篇文章記錄一下基於 jupyterlab作自定義接口和插件的二次開發過程和關鍵點

目前咱們給甲方提供的機器學習平臺是基於k8s + jupyterlab實現的, 這樣的好處是數據科學家能夠在一個相對隔離的環境裏開發本身的數據應用, 可是缺點是每一個人之間沒法共享本身開發的腳本給其餘人. jupyter生態並不提供這樣的功能, hub這種多用戶系統也沒有. 因此咱們的思路是用第三方雲存儲來實現文件的共享. A用戶將本身的開發的腳本上傳到雲存儲並決定共享給B, 而後B用戶接到A用戶共享文件給他的通知, 從雲存儲上下載該文件. 由於我不作前端開發, 因此, 我只記錄我這邊python端實現的思路和方法.

1. A 決定共享一個文件給B

2. A 上傳文件到S3

3. python端記錄A的username和文件名和共享給B的用戶名信息並存入數據庫

4. B的lab中經過自定義接口獲取數據庫中共享給本身的信息

5. B經過S3下載該文件到本身的文件夾裏面

這裏面比較複雜的地方在於hack jupyterlab的源碼, 增長自定義的接口. 固然, 我沒有耐心去看jupyterlab的插件開發文檔, 直接讀源碼, 分析他的API入口, 而後本身寫方法實現是我等糙人的一向行事風格.也許這樣不夠規範, 可是足夠簡單粗暴直接. 

首先在 jupyterlab/handlers/建立一個custom_handler.py的文件.

# 引入notebook的APIHandler, 做用是繼承環境變量, 以及jupyter的各類配置項
from notebook.base.handlers import APIHandler

class ShareNotebookHandler(APIHandler):
    # APIHandler繼承自Tornado的web.RequestHandler, 因此, 繼承的類不能用__init__作入口, 須要用initialize方法作入口, 這是tornado規定的.
    def initialize(self, uploader):
        self.uploader = uploader
        # 這裏獲取用戶的 notebook 所在的文件夾, 並替換爲絕對路徑, 這是個 Lazy 的 Config
        # 因爲環境是k8s, notebook_dir 是有 lab 啓動參數設定的, 是個 LazyConfig, 因此能夠獲取到, 若是是普通環境, 須要用 server_dir 代替
        # Jupyter 的 LazyConfig 的本質是一個 traitlets 對象, 因此須要str轉換一下變量類型.
        self.working_dir = str(self.application.settings['config']['LabApp']['notebook_dir'].
                               replace('~', os.path.expanduser('~')))
        custom_config_dir = str(self.application.settings['config_dir'])
        # CustomConfig 和 ShareNotebookMetadata 是我本身寫的獲取custom配置的類, 在別的文件裏, 做用是獲取postgres, s3的配置信息, 能夠不用理會
        # s3 使用的是 minio 作的本地存儲集羣, 兼容 s3協議
        cc = CustomConfig(custom_config_dir)
        self.custom_config = cc.get_config()
        self.db = ShareNotebookMetadata()
        self.minio = MinioUtils(custom_config_dir, 'model-share')

    @gen.coroutine
    @web.authenticated
    # http put方法是被分享用戶從 S3 存儲下載文件用的
    def put(self): # Download shared notebooks from s3
        data = json.loads(self.request.body)
        filename = data['filename'] # shao.zs/xxxx.ipynb
        group = data['group']
        # 從環境變量中獲取當前進程的用戶, k8s須要傳遞用戶名環境變量到pod中
        cur_user = os.environ['USER'] # meng
        filename_split = filename.split('/') # ['shao.zs', 'xxxx.ipynb']
        file_user = filename_split[0] # shao.zs
        # 設定用戶下載所使用的文件夾, 有則寫入文件, 沒有則建立文件夾再寫入文件
        file_path = self.working_dir + '/shared/' + file_user # /home/meng/notebooks/shared/shao.zs
        if not os.path.exists(file_path):
            try:
                os.makedirs(file_path, 0o775)
                self.log.info('mkdir share dir %s' % file_path)
            except IOError as e:
                self.log.error('%s' % e)
        ret = []
        # 下載文件, 經過子文件夾名稱, 代表文件是共享自誰, 成功刪除文件記錄, 並返回200, 失敗返回500
        if self.minio.download_minio(filename, file_path): # shao.zs/xxxx.ipynb, /home/meng/notebooks/shared/shao.zs
            self.db.delete(cur_user, file_user, filename, group)
            minio = {'code': 200, 'messages': 'Downloaded to ' + file_path}
        else:
            minio = {'code': 500, 'messages': 'Check log for details'}
        ret.append(minio)
        self.set_header('Content-Type', 'application/json')
        self.write(json.dumps(ret, ensure_ascii=False))

    # 獲取當前環境變量中的用戶名, 並根據用戶名獲取被共享的文件列表, 寫入到http restful接口, group在這裏沒啥意義, 是咱們本身內部使用的標示, 能夠忽略
    def get(self):
        cur_user = os.environ['USER']
        group = self.get_argument('group', '')
        shared = self.db.select(cur_user, group, 0)
        print(cur_user)
        self.set_header('Content-Type', 'application/json')
        self.set_header("Access-Control-Allow-Origin", "*")
        self.set_header("Access-Control-Allow-Headers", "x-requested-with")
        self.set_header('Access-Control-Allow-Methods', 'POST, GET, PUT, DELETE')
        self.write(json.dumps(shared, ensure_ascii=False))


    @gen.coroutine
    @web.authenticated
    # A用戶上傳文件的 post 方法
    def post(self):
        data = json.loads(self.request.body)
        filenames = data['filenames']
        groups = data['groups']
        cur_user = os.environ['USER']
        ug_reflect = UserGroupReflection() # 前端傳遞的實際上是組名, 一個組名可能對應多個用戶名, 但一般組名與用戶名相同, 參考linux用戶和組的概念
        token = self.get_cookie('Token')
        if len(filenames) > 0:
            for filename in filenames:
                extend_name = filename.split('.')[-1]
                real_filename = filename.split('/')[-1]
                abs_filename = self.working_dir + '/' + filename
                new_filename = cur_user + '/' + real_filename # shao.zs/xxxx.ipynb
                if extend_name != filename:
                    new_filename = new_filename + '.' + extend_name
                else:
                    new_filename = new_filename
                ret = list()
                # 上傳一個或多個文件
                if self.minio.upload_minio(abs_filename, new_filename):
                    minio = {
                        'code': 200,
                        'module': 'minio',
                        'messages': 'ok',
                        'username': cur_user,
                        'filename': new_filename,
                        'groups': groups,
                        'url': '/lab/custom/sharenb?filename=' + new_filename
                    }
                    for group in groups:
                        reflect = ug_reflect.get_user_from_groupname(group, token)
                        for ug in reflect:
                            self.db.insert(ug['user'], cur_user, new_filename, ug['gname'])
                else:
                    minio = {
                        'code': 500,
                        'module': 'minio',
                        'messages': 'Check log for details',
                        'username': cur_user,
                        'filename': new_filename,
                        'groups': groups
                    }
                ret.append(minio)
                self.set_header('Content-Type', 'application/json')
                self.write(json.dumps(ret, ensure_ascii=False))

而後把文件上傳的結果寫回到restful接口供前端使用.

接下來, 修改jupyterlab/extension.py, 添加自定義接口的自定義路由到tornado裏面.

def load_jupyter_server_extension(nbapp):
    from .handlers.custom_handler import (
        ShareNotebookHandler
    )
    
    # 參考添加路由的位置
    build_url = ujoin(base_url, build_path)
    builder = Builder(core_mode, app_options=build_handler_options)
    build_handler = (build_url, BuildHandler, {'builder': builder})
    handlers = [build_handler]

    ###############
    # custom handler added here by xianglei
    ###############

    # ujoin爲jupyterlab內部方法, 做用是append web路由給tornado
    # 這個uploader目前沒搞明白是幹嗎使的, 不寫還不成, 給個空的就能夠
    # base_url是jupyter啓動時的一個配置項, 定義路由前綴, 默認是空
    custom_sharenb_url = (ujoin(base_url, '/lab/custom/sharenb'))
    custom_sharenb_handler = (custom_sharenb_url, ShareNotebookHandler, {'uploader': ''})
    handlers.append(custom_sharenb_handler)

    ###############

這樣改造完以後 前端能夠訪問 http://xxxx.com/lab/custom/sharenb 來進行文件共享的操做, get 是獲取文件列表, post是發佈共享文件, put是被分享人下載共享文件, 效果以下, 這個右鍵菜單是前端開發的, 跟我不要緊. 

而後被分享文件列表頁

而後是已下載的文件位置

代表 jinjianbing分享了 ty.pmml文件給當前用戶, 而且已經下載到了本地文件夾.

順便展現一下給尊貴的甲方科學家搭建的基於hadoop集羣的笛卡爾積開發環境

我仍然是一個老工程師