支持弹性伸缩的Jupyter （elastic-jupyter-operator项目）

前言

Jupyter Notebooks 在 Kubernetes 上部署往往需要绑定一张 GPU，而大多数时候 GPU 并没有被使用，因此利用率低下。为了解决这一问题，开源项目 elastic-jupyter-operator 将占用 GPU 的 Kernel 组件单独部署，在长期空闲的情况下自动回收，释放占用的 GPU。这篇文章主要介绍了这一开源项目的使用方式以及工作原理。

如何提高资源利用率（base K8S）

最大的挑战来自于GPU资源利用率。在运行的过程中即使没有代码在运行，Notebook 也会长期占用着 GPU，造成 GPU 的空置等问题。在大规模部署 Jupyter 实例的场景下，一般会通过 Kubernetes 创建多个 Notebook 实例，分配给不同的算法工程师使用。而在这样的情况下，我们需要在对应的 Deployment 中事先申请 GPU，这样 GPU 会与对应的 Notebook 实例绑定，每个 Notebook 实例都会占用一张 GPU 显卡。

然而同一时间，并不是所有的算法工程师都在使用 GPU。在 Jupyter 中，编辑代码的过程是不需要使用计算资源的，只有在执行 Cell 中的代码片段时，才会使用 CPU 或 GPU 等硬件资源，执行并返回结果。由此可以预见，如果通过这样的部署方式会造成相当程度的资源浪费。

 造成这一问题的原因主要是原生的Jupyter Notebooks没有很好地适配 Kubernetes。在介绍问题原因之前，先简单概述一下 Jupyter Notebook 的技术架构。如下图所示，Jupyter Notebook 主要由三部分组成，分别是用户和浏览器端，Notebook Server 和 Kernel。

 其中用户和浏览器端是 Jupyter 的前端，主要负责展示代码和执行结果等。Notebook Server 是它的后端服务器，来自浏览器的代码执行请求会被 Notebook Server 处理，分派给 Kernel 执行。Kernel 是真正负责执行代码，返回结果。

在传统的使用方式中，用户会通过 jupyter notebook $CODE_PATH 等命令，在本地运行 Jupyter Notebook Server，随后访问浏览器中的 Jupyter 交互式开发界面。当代码需要执行时，Notebook Server 会创建一个独立的 Kernel 进程，这一进程会使用 GPU 等运行。在 Kernel 长期空闲，没有代码需要执行时，这一进程会被终止，GPU 也就不再会被占用。

而当部署在 Kuberenetes 之上后，问题就产生了。Notebook Server 和 Kernel 运行在同一个 Pod 的同一个容器下，尽管只有执行代码时才需要运行的 Kernel 组件是需要 GPU 的，而长期运行的 Notebook Server 是不需要的，但是受限于 Kubernetes 的资源管理机制，还是需要给其提前申请 GPU 资源。

在 Notebook Server 的整个生命周期中，这一块 GPU 始终与 Pod 绑定。在 Kernel 进程空闲时虽然会被回收，但是已经分配给 Pod 的 GPU 卡却不能再交还给 Kubernetes 进行调度了

解决方案

为了解决这一问题，就有了开源项目elatic-jupyter-operator。思路非常朴素：问题源于 Notebook Server 和 Kernel 在同一个 Pod 中，导致我们无法分别为这两个组件申请计算资源。那只要将他们分开部署，让 Notebook Server 在单独的 Pod 中，Kernel 也在单独的 Pod 中，相互之间通过 ZeroMQ 通信即可。

通过这样的方式，Kernel 会在空闲时被释放。在需要时会再次被临时性地申请 GPU，运行起来。为了实现这一目的，我们在 Kubernetes 中实现了 5 个 CRD，同时为 Jupyter 引入了一个新的 KernelLauncher 实现。通过它们，用户可以在 GPU 空闲时将 Kernel 回收释放，在需要执行代码时再动态地申请 GPU 资源，创建 Kernel Pod 进行代码执行。

简单的例子

下面我们将通过一个例子介绍使用方式。首先我们需要创建 JupyterNotebook CR（CustomResource），这一个 CR 会创建出对应的 Notebook Server：

apiVersion: kubeflow.tkestack.io/v1alpha1
kind: JupyterNotebook
metadata:
  name: jupyternotebook-elastic
spec:
  gateway:
    name: jupytergateway-elastic
    namespace: default
  auth:
    mode: disable

其中指定了gateway，这是另外一个JupyterGateway CR。为了能够让Jupyter支持远程的Kernel，需要这们一个网关进行请求的转发。我们同样需要创建这样一个CR：

apiVersion: kubeflow.tkestack.io/v1alpha1
kind: JupyterGateway
metadata:
  name: jupytergateway-elastic
spec:
  cullIdleTimeout: 3600
  image: ccr.ccs.tencentyun.com/kubeflow-oteam/enterprise-gateway:2.5.0

JupyterGateway CR 中的配置 cullIdleTimeout 指定了经过多久的空闲时间后，其管理的 Kernel Pod 会被系统回收释放。在例子中是 1 个小时。创建完这两个资源后，就可以体验到弹性伸缩的 Jupyter Notebook 了。如果在一个小时内一直没有使用的话，Kernel 会被回收。

安装及使用

1、下载 elastic-jupyter-operator

git clone https://github.com/skai-x/elastic-jupyter-operator

项目结构如下：

2、安装 elastic-jupyter-operator

进入项目根目录，根据 README.md 指引安装elastic-jupyter-operator

# 安装所需的Namespace、ServiceAccount等资源
kubectl apply -f ./hack/enterprise_gateway/prepare.yaml
# 安装elastic-jupyter-operator
kustomize edit set image controller=ghcr.io/skai-x/elastic-jupyter-operator:latest
kustomize build config/default | kubectl apply -f -

验证安装结果

kubectl get po -n elastic-jupyter-operator-system

3、创建 jupytergateway-elastic、jupyternotebook-elastic

kubectl apply -f ./examples/elastic/kubeflow.tkestack.io_v1alpha1_jupyternotebook.yaml
kubectl apply -f ./examples/elastic/kubeflow.tkestack.io_v1alpha1_jupytergateway.yaml

验证创建结果

4、为 jupyternotebook-elastic 创建svc

 为了让浏览器能访问Notebook Server，需要为jupyternotebook-elastic创建“NodePort”类型的svc

5、使用notebook

查看k8s集群，会发现当我们使用Notebook运行代码的时候，自动生成运行Kernel的pod（亲测，此过程很快，几乎不会影响用户使用体验）

 设计与实现

在介绍完使用方式后，我们简单介绍其设计与实现。

当用户在浏览器中选择执行代码时，首先请求会发送给在Kubernetes上运行的Notebook Server。由于集群上没有正在运行的 Kernel，代码执行任务无法分配下去，所以 Notebook Server 会向 Gateway 发送一个创建 Kernel 的请求。Gateway负责管理远端的Kernel的生命周期，它会在Kubernetes集群中创建对应的JupyterKernel CR。随后会与集群中已经创建好的Kernel通过ZeroMQ进行交互，然后将代码执行的请求发送给 Kernel 进行执行，随后将结果发送给 Notebook Server 再将其返回给前端进行渲染和展示。

而Gateway会根据在JupyterGateway CR中定义的有关资源回收的参数，定时检查目前管理的Kernel中有没有满足要求，需要被回收的实例。当Kernel空闲时间达到了定义的阈值时，Gateway会删除对应的JupyterKernel CR，将其回收，释放GPU等资源。

 参考文献：

云原生的弹性 AI 训练系列之三：借助弹性伸缩的 JUPYTER NOTEBOOK，大幅提高 GPU 利用率 | 高策