caicloud/ciao：在 Jupyter Notebook 中使用 Kubeflow 进行分布式模型训练

文末可获得 KubeCon China 2019 PPT 大礼包

作者 ｜ gaocegege（才云研发工程师）

编辑 ｜ 小君君（才云）

在这篇文章中，才云科技算法工程师 gaocegege 会介绍在 Jupyter Notebook 中支持 Kubeflow 的一些实现方案，以及一个为 Kubeflow 实现 Jupyter kernel 的开源组件 ciao。通过它，用户只需几行 magic commands 就可以发起一次分布式模型训练任务。

今天，我为大家介绍一款由才云科技开源的 Jupyter kernel 软件 ciao。这个软件的名字源于意大利语中的“你好”，我们把它起这个名字有两个原因：一是，我们公司的开源项目都喜欢用 C 打头，在翻阅字典后，我们发现 ciao 是一个比较合适的名字；二是，因为这个项目它是为算法工程师服务的，同时，“你好”这个单词可以让大家感觉很亲切。

背景与现状

ciao 是一个 Jupyter 的 kernel，它可以做一些什么事情呢？它能够允许我们在 Jupyter Notebook 中使用 Kubeflow 中的一些功能。关于 ciao，我们先介绍一下背景。

其实在 ML 领域，我们可以看到一个新的职位叫做算法工程师，这也是目前计算机领域比较火的一个职位。在工作职能方面，算法工程师与传统开发工程师相比会有一些工作流程上的区别。

目前算法工程师的开发方式往往与传统软件开发不同。比如，他们需要一个带 GPU 的虚拟机，或者一个本地的有 GPU 的物理机。他们可能会用一个 Jupyter Notebook，将资源利用起来，然后在周边做一个交互式的开发。再比如说，他们在 Jupyter Notebook 中，准备好数据集、ML 训练代码，然后再去其中发起一次训练任务。然而，这是一个较为偏向本地的开发过程。除此之外，它还是一个实验驱动的开发过程，这个是什么意思呢？

在传统软件开发过程中，工程师们往往都是实现一些 Feature 的，或者说是实现一个 System。而算法工程师实际是要构建一个机器学习或者深度学习的模型，这个模型就会涉及到一些不停的迭代式开发过程。比如，我现在用数据集去构建一个模型，这个模型训练出来后，它会给我一个结果。根据这个结果反馈，我们会在做下一个模型时调整一些东西。比如，在 CNN 中，我们可以调整 kernel 的 size 等等。在调整之后，我们会再发起一次新的训练。

在这样一个实验驱动的开发过程中，算法工程师们会比较习惯用一些交互式的开发环境，像 Jupyter Notebook。他们也有可能会直接用 IPython kernel 去做开发，或者是 R 语言 R Studio 之类的类似环境。这样的类似环境会给算法工程师造成一定的局限（ciao 就是为了解决这个局限而存在的）。

那么问题来了，我们应该如何为算法科学家提供一个更完善的开发环境呢 ? 我们如何在保证开发体验的同时，将本地的开发环境比较无痛地移到云端上呢？所谓的本地开发，目前存在一些什么问题呢？

第一个问题是硬件资源没有办法共享。 一方面，在 Jupyter Notebook 中，进程是跑在本地上的。比如，你需要用 GPU 去运行 TensorFlow 训练，但是这个 GPU 只能用你本地已有的 GPU。好比，在本地有两块 GPU，你就只能用这两块 GPU，而没有办法去共享。

第二个问题是模型分布式训练在 Jupyter Notebook 中很难落地。 比如，在成熟度不高的用户环境下，模型分布式训练是难以落地的。

在工业界，机器学习往往会有不同的分布式模型，像模型并行、还有数据并行，它们都是用不同的方式来支持。比如说我们在进行模型训练时，可以用不同机器上的不同 GPU 去做。在工业界，其实用的更多的是数据并行里面的 Parameter Server-Worke 分布式模型。比如，在一次分布式训练中，我们会有一些 Parameter Server（PS）、 worker 节点。Parameter Server 会负责收集模型训练时的一些参数、梯度等。worker 是一个利用 GPU 去做计算的节点，这两个节点在一起工作时，就需要进行一次分布式训练。

虽然，这个原理看起来很简单，Parameter Server 就像一个 K/V 数据库，它负责存储一些参数，然后 worker 负责去运行真正的计算过程。但是在实际的操作过程中，我们会发现分布式的训练其实并不是那么好运行的。

下图是分布式训练的一个示例，我们一共能看到 4 条命令，分别是 2 个 worker 和 2 个 PS。

大家可以看到在这上下各两条命令，这前两个命令是运行了两个 Parameter Server，后面的两个命令是起了两个 worker，worker 是负责计算的。 一般来说，我们都会让 worker 去挂一个 GPU 来做计算。 在这张图中，我们可以发现除了 PS 和 worker 两个参数以外还有其他参数，比如： ps_hosts 和 worker_hosts。 这两个参数是在进行  TensorFlow 分布式训练的时候，让 TensorFlow 的 worker 和 PS 知道其他的 worker 和 PS 在哪里运行。 这样在它们之间才能进行通信。

这是一个在物理机上进行的训练 。从这个示例中，我们可以看到原本工程师只是在 Jupyter Notebook 中写一些训练代码就可以了，但是如果让他去进行分布式训练，他就需要做很多手工操作，而且这个操作需要 ssh，还需要在不同的机器上执行。这是相当复杂的一个过程。

在这个问题上，我们的工程师就想利用一些组件进行优化。在 Jupyter Notebook 中，我们用到了云端资源，而且它提供了一个 0 侵入性的改动，也就是说你本地开发是什么样，你的云端开发还是怎么样的。而且它会带来一些额外的优势，比如说我们分布式训练不再需要去写一些额外的脚本来运行。

Kubeflow 如何帮助我们规避这些问题

我们是如何实现这件事情的？首先我们需要了解 Kubeflow 是如何去做的。Kubeflow 顾名思义，它其实就是 ML toolkit for Kubernetes。它主要就是将 TensorFlow、PyTorch、 Caffe、MXNet 这些框架训练跑在 Kubernetes 的平台上。Kubeflow 实现了一些不同的 CustomResourceDefinition（CRD）和一些独立的、运行在 Kubernetes 上的系统。举个例子，对于 TensorFlow 分布式训练的支持 ，就是通过一个名为 TFJob 的 CRD 实现的。当你定义了一个 TFJob 的 YAML 配置时，它就可以帮助你把配置中定义的 Parameter Server 和 worker 都运行起来。

在上图我们也看到了手工的操作方式。我们需要写 Python 训练脚本，然后在不同的机器上用不同的命令运行不同的任务。现在这种方式相当于是我们定义的一个 YAML 配置，里面会指定我们用哪个镜像，镜像里会有代码和你的使用的框架依赖，包括数据集。这边只需要指定说你需要几个 PS、几个 worker，它们的重启策略是什么，然后它就可以帮你去进行一个分布式的训练任务了。我们觉得这是一个架构，你只要定义一个 TFJob，它会帮你把 PS，还有 worker 都以 Pod 的方式去运行起来。

就像 我们刚提到的 Parameter Server 这种分布式的方式，在生产领域上，我们也会经常用到 AllReduce，它也会有一些不同的分布式模型，能够让你的训练更节约资源。

TFJob 也是支持不同的分布式模型，同时也可以隐藏底部细节。在上面的实例中，你可以选择用 CPU 或 GPU，而不需要通过一些脚本 Hack 的方式去做软件。在分布式训练中，worker 是分为两类的，一类是 Chief Worker，另一类是普通的 worker。当普通的 worker 挂了，在某些情况下是可以重启的。我们的这个 TFJob 可以自动帮你实现这个层级的策略等，但如果只做到这个程度是不够的。

为什么要 ciao

在这种方式下，算法科学家还需要把这个算法打包成镜像，写清数据的存储位置 。同时，算法科学家可能在不熟悉 Docker 的情况下产生一些问题。他还需要写一个 YAML 来定义 TFJob 的资源，然后把资源给到我们的 Kubernetes 进行处理。在这样的操作下，每次训练都需要构建镜像，在这里没有一个所见即所得的结果，就意味着你每次发起声明确定的时候，就需要去用比如说 kubectl logs 来得到训练日志。这些痛点也是目前 Kubeflow 无法解决的。

不得不说，我们现在在 Kubernetes 上运行分布式的 TF 训练任务已经非常方便了，但这种方式还是没有办法实现一种交互式的体验，所以能不能把分布式的任务，通过 Kubernetes 的资源调度能力把它运行在 Jupyter Notebook 中呢？ 这也是 ciao 想做的事情。 通过 ciao，Kubeflow 能够适应变更频繁的开发环境，以及让用户不再需要在 Kubernetes 中获取训练日志。

上图是我们整个 ciao 的架构图。我们想在 Jupyter Notebook 里也能够发起训练任务，并在处理之后把 log 返回给 Jupyter Notebook。这样的方式意味着你在 Jupyter Notebook 里面可以进行分布式的、交互式的开发，可以直接拿到你想要的结果，并且直接发起一次新的分布式训练任务。

其实，在 Kubeflow 中也有一个解决方案 Fairing。它做的事情相当于是，在 Python 中包装了一个 Client。 就比如在这个例子中，我们可以设置一个 Builder，然后就利用这一 Builder 依赖的基础镜像。 把 Python Code 直接打包成镜像，然后去进行训练。

这种方式会有一些好处也有坏处。好处是在于高通用性，就是说你只要写 Python 就可以了，它理论上能支持不同的开发环境。比如说，你用 VS Code 做算法开发，你也可以在 VS Code 中用这个库去做这件事，你也可以在 Jupyter 中做，这有好处也有坏处。因为它不是一个定值，单个 Jupyter 应用场景会引入一些额外的负担，同时，它的学习成本也比较高。

我们可以看到它自定义了一个配置，你需要学习它的一些 API 等。 同时，你还需要在你原本的训练代码中，再写一些额外的代码。 就像图中 Simple Model 是一些常规意义上模型训练的原有代码，而下面这个 main 函数是我们需要额外注入的一些新代码。有了这些新的代码之后，我们才能在 Python 代码里面直接发起分布式的训练任务，但是这对于我们来说并不方便。

Demo Show

我们想做的其实是一个适配 Jupyter Notebook 的方案，一个 Kubernetes Native 的实现。因此我们实现了一个 Jupyter kernel，也就是 ciao。 它利用 Magic Command 的方式，避免了对用户代码的侵入性改动。 Magic Command 是由 ciao 提供的独立的命令，以 % 开头。我们定义了 %framework、%ps 和 %worker 三个 Magic Command。

我们可以看到这样一个 Demo，它其实是完全不需要写任何多余的代码的。在 … 下面开始都是你原生的训练代码。而上面我们可以看到有三个 magic command 的，特别是比如说 framework=tensorflow、 ps=1、worker=1，这个就是 magic command。我们就是通过这样一种方式来实现不需要修改用户训练代码的想法。

大家也可以通过这个 Demo 来更深入地了解。在上面我们可以看到有三个 magic command（第一行、第二行、第三行），然后下面就是我们真正的分布式代码。在后边你可以看到右上角的一个选项。我们使用的是 Kubeflow 的一个 kernel，你可以看到后面在点运行之后，就会发起一次训练的任务。 在这次训练任务里面一共需要做两件事 。

第一件事： 因为我们是跑在 Kubernetes 中的，我们需要把下面用户定义的代码打包成一个镜像来运行。 ciao 会帮你做到这些事情，也就是你只要把代码写在 Jupyter 中，当运行起来时，我们就会帮你把代码打包成一个镜像，然后根据这个镜像构建出一个 TFJob，再交给 Kubernetes 做训练。

第二件事：我们需要根据这个镜像构建出一个 TFJob，再交给 Kubernetes 做训练。下面就是这次训练的一些日志，它会有一个前缀是 ps0 或 worker0，这个意义在于，同时从不同的任务上拿到一些日志。

Demo 的演示就到这里，接下来我们来介绍 ciao 是如何实现的。

Ciao 的实现

首先，我需要介绍一下 Jupyer 的实现。如下图所示，Jupyer 是一个 browser-server 架构，browser 是笑脸，Notebook server 充当一个 server 的作用。在后面它接了一些 kernel，就像我们所知道的有 IPython kernel、R 语言的 kernel。它们其实都是在处理用户提交的代码。如果你想实现我们刚刚讲的 ciao 的功能，有三种不同的可选实现方式。

第一种方式就是在 IPython kernel 中自定义库。 这种方法你在其他地方也可以用，可能会带来一定的侵入性，但是它有更好的通用性。 

另外一种方式就是修改 IPython kernel 的内容。 IPython kernel 本身就有一些扩展的方式。比如，你可以定义一些新的 magic command。因为我们的实现也是需要定义一些新的 command，所以之前我们团队也调研过，能不能直接进行一些修改来支持这个功能，但是我们发现这是没办法做到的。因为 IPython kernel 是跑在你本身进程的下面，也就意味着它没有办法跟 Kubernetes 去做一些通信。

而我们的实现是基于第三种方式，Custom kernel 架构。 我们实现的 Custom kernel 做的事情就是把请求给打包镜像，然后发给 Kubernetes，处理完的结果日志再返回给 Notebook，Notebook 再呈现给浏览器前的用户。

这样的架构会引入一个问题，现在支持 Kubeflow，但是不支持其他的框架。比如，原本支持所有的 Python 库，现在就只支持分布式训练，那你就不能支持其他的库，于是我们就引入了另外一个 sos kernel 来解决这个问题。

sos kernel 是社区里已经存在的一个开源方案，它的一个好处是：每一个 Notebook 里面都可以用不同的 kernel 去运行。比如，下图第二个 Cell 里面，我们是用 R 语言来实现的，就是在 Cell 的右上角直接选择需要用的语言。在这里是用三个不同的 kernel 运行不同的代码，这样一个好处就是有不同的 Cell 可以用不同的 kernel 来解释。当你需要 Kubeflow 的时候就需要选择 ciao，它可以帮你去运行分布式训练任务。当你不需要分布式训练时，你可以用 IPython kernel，直接去进行你下一步的单机训练。这对用户体验来说是一个比较好的提高。

上图是一个 sos 架构的实现，它像一个超级内核的概念，它的下面会挂着很多子 kernel，Kubeflow 是它其中的一个 kernel。我们引入它的好处就是在一个 Jupyter 上支持多个 kernel，包括支持一些更加高级的特性，比如 R 语言中的一些你想共享的数据结构；比如你想在 Python 中用到的数组。这时 sos kernel 也会提供一些简单封装，保证你可以做到这件事。

所以 ciao 的整个架构可以概括为：首先，它与 sos kernel 通信，sos kernel 负责分发来自前端的运行请求，如 Kubeflow 请求，它会发送到 ciao。

ciao会这样做：它会利用 S2I provider 将我们的代码转成镜像，之后根据镜像去构建出我们的 TFJob，然后用 Kubernetes Client 发给我们的 APIServer，通知它我现在需要一个新的 TFJob。 后面就是由 Kubeflow 社区里面的 tf-operator，还有 pytorch-operator 来处理不同的请求。 它一共有三个阶段：

• 第一阶段：解释那些 magic command（基于前缀匹配来做的）。

• 第二个阶段：将源码去打包成一个镜像。在 ciao 中，我们提供不同的后端实现支持：

• Knative 是用 kaniko 来做的；

• S2I 能够把源码直接打包成一个镜像；

• configmap，因为打包成镜像这个过程本身是很花时间的，意味着我们把源码得重新又得成一个镜像。configmap 可以省略打包这样一个过程。

• 第三个阶段：构建对应的分布式训练 CRD，发送给 Kubernetes 进行真正的训练，并且把日志输出到 Jupyter Notebook 中。

上图是日本的一个工程师做的一个比较不同源码到镜像的构建 工具 的性能图，蓝色表示第一次构建；黄色表示相同 Dockerfile 的第二次构建。在这其中，最好的工具应该是 Buildkit（属于 Docker 社区）。因为我们主要是做首次构建比较多，利用不到 BuildKit 和 img 后续构建加速的优势，所以我们现在推荐的方式还是用 Kubernetes Configmap。利用 Kubernetes Configmap，我们可以完全绕过镜像构建环节，只需要准备好基础镜像即可。

在第三个环节中，我们需要把构建好的镜像发给 Kubernetes，这个其实就用 Kubernetes Client 就可以了。此处不再赘述。

ciao 的部署

在 Jupyter 的部署中，我会遇到的问题是，比如我们会有很多用户，不同的用户的不同请求都会分发到同一个 kernel 上，这个时候 kernel 的压力会非常大，所以我们就需要一个方式让 kernel 能够横向扩展，这就需要用到了 Jupyter Enterprise Gateway。Jupyter 相当于是一个前端， Jupyter Enterprise Gateway 相当于一个网关。

所有都来自 Notebook 的请求，先走到 Jupyter Enterprise Gateway 里面，然后它会做一个分发， 不同的请求应该分发给不同的 kernel。这样实现还有一个好处是以前 Notebook 和 kernel 都是跑在同一个节点或者说同一个进程中。 现在 kernel 和 Notebook 可以跑在不同进程中，所以通过这种方式我们就可以让 ciao 变得无状态且可以横向扩展。

总结

今天我讲的内容主要是：如何能够在 Jupyter 中进行分布式训练，同时让用户能够尽可能的不修改它的代码，有一个比较好的用户体验。在之后，如果大家在这方面有什么问题，可以在本文下方留言。

同时，本文的 PPT 以及演讲视频，请通过以下链接浏览：

• 观看地址：

  https://dwz.cn/xVEiMxCi

• PPT 下载地址：

  http://www.k8smeetup.com/downloadAll?type=pdf

社区福利

KubeCon China 2019 PPT 回顾资料已经为大家整理好啦！在公众号后台回复“ 201902 ”，即可获取 KubeCon China 2019   资料大礼包。