Jupyter生态二次开发系列(六)

本次记录一下限制用户spark8s进程数量的方法, 咱们的jupyterlab是跑在pod里面的, sparkui是经过自定义jupyterlab url的方式来映射出来, 而lab url只有一个, 因此每次只能容许用户开启一个pyspark8s的notebook, 但使用过程当中发现, 用户能够开好几个spark8s kernel的notebook, 其实使用是没什么问题的, 只是后面几个开的spark是无法代理ui的. 并且很是占k8s集群资源, 原生的spark8s并不像spark跑在yarn上面, 一旦做业结束, executor就收回资源了, 在k8s环境里, spark做业结束, executor的pod并不会被k8s回收资源. 因此开着不关是很耗资源的, 而后再容许多开就更耗资源了, 因此须要限制死, 一个用户只能开一个 spark8s 的 notebook.

以前其实作过限制, 可是只能是用户在第一次建立spark8s notebook的时候不容许建立, 但下面的使用方式是无法限制的

1. 用户建立启动了第一个 spark8s notebook(后台会检测是否有spark进程, 没有则启动)

2. 用户shutdown 了刚刚建立的spark8s notebook的kernel

3. 用户又建立并启动了一个 spark8s notebook(后台会检测是否有spark进程, 因为用户shutdown掉了第一个, 因此没有进程, 正常启动)

4. 用户又启动了第一个 spark8s notebook的kernel(本应检查是否已存在spark进程, 但没有, 因此就变成了多开spark8s notebook)

发现这个状况以后, 我分析了一下代码和缘由, 发现是jupyter的流程问题致使的, jupyter在第一次建立notebook的时候会建立一个临时kernel文件并注册到server, 这个临时kernel文件包含了该notebook所指向的真实kernel.json文件位置, 而后启动时所使用的端口, kernel名称等信息. 这并非真正的kernel.json, 只能算是一个代理文件. 生成这个临时文件并注册以后, jupyter内部就再也不调用建立kernel时所使用的方法, 所以我以前在建立kernel时所写的spark进程检测程序就没有被调用, 所以就没有限制住用户使用上述流程启动多个spark8s notebook. 了解了缘由以后就须要去找notebook是在哪里启动kernel的. 

通过一番坚苦卓绝的查找, 最后仍是找到了 jupyter_client, 在 client里面有一个 start_kernel的方法, 这个方法是不管第一次启动仍是后面再次启动都会必须调用的, 因此在这段代码里加入以前检测 spark8s的代码, 就ok了, 可是仍然仍是有一个问题, jupyter_client和jupyter_server都不能直接返回报错字符串回给notebook和lab, 我只能经过 raise Error的方法返回, 这个在前端页面显示时是不太友好的, 可是目前还没找到好的办法解决.

附赠 kernelspec 文件增长的spark8s检测方法

def __find_spark8s_process(self):
    import psutil
    pids = psutil.pids()
    spark_instance = 0
    for pid in pids:
        p = psutil.Process(pid=pid)
        try:
            pcmd = p.cmdline()
        except:
            pcmd = ''
        if ' '.join(pcmd).find('spark.kubernetes.container.image.pullPolicy') >= 0:
            self.log.info('Finding spark8s kernel: [' + ' '.join(pcmd) + ']')
            spark_instance += 1
    if spark_instance >= 1:
        self.log.error('Spark k8s context already exists, do not start the kernel')
        return True
    else:
        self.log.info('Everything ok')
        return False

我经过 spark8s特有的关键字查找spark进程, 找到则True, 没有则False, 为何不用更短一点的诸如 pyspark-shell这样的关键字查找呢? 由于咱们的lab里面除了 spark8s以外, 还有 spark local的kernel, 我不能由于已经有spark local的进程存在就不让启动spark8s了.

而后在manager里面修改代码以下

# add by xianglei
# 因为kernelspec里面的KernelManager只在client第一次启动注册时使用，因此在这里只调用一下，注册后nb更换kernel或重启是直接调用这里,
# 不会再次调用 kernelspec的KernelManager方法
from .kernelspec import KernelSpecManager
ksm = KernelSpecManager()
try:
    self.log.error(ksm.get_kernel_spec(self.kernel_name).to_dict())
    kernel_cmd, kw = self.pre_start_kernel(**kw)
    # launch the kernel subprocess
    self.log.debug("Starting kernel: %s", kernel_cmd)
    self.kernel = self._launch_kernel(kernel_cmd, **kw)
    self.post_start_kernel(**kw)
except Exception as e:
    raise e

进程检测就是在KernelSpecManager().get_kernel_spec()方法里面调用的, 这个get_kernel_spec方法会在第一次建立notebook时被调用, 但后面再启动则不会, 但后面再启动notebook会调用 start_kernel方法, 所以在start_kernel方法里面加一次调用就能够, 我也是省事偷懒惯了.