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方法里面加一次调用就可以, 我也是省事偷懒惯了.

猜你喜欢

转载自blog.51cto.com/slaytanic/2646820