这篇文章记录一下基于 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是被分享人下载共享文件, 效果如下, 这个右键菜单是前端开发的, 跟我没关系. 


image.png


然后被分享文件列表页

image.png

然后是已下载的文件位置


image.png


表明 jinjianbing分享了 ty.pmml文件给当前用户, 并且已经下载到了本地文件夹.


顺便展示一下给尊贵的甲方科学家搭建的基于hadoop集群的笛卡尔积开发环境

image.png



我仍然是一个老工程师