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方法裏面加一次調用就能夠, 我也是省事偷懶慣了.