KubeDL HostNetwork：加速分布式训练通信效率

作者：陈裘凯（ 求索）

前言

KubeDL 是阿里开源的基于 Kubernetes 的 AI 工作负载管理框架，取自"Kubernetes-Deep-Learning"的缩写，希望能够依托阿里巴巴的场景，将大规模机器学习作业调度与管理的经验反哺社区。目前 KubeDL 已经进入 CNCF Sandbox 项目孵化，我们会不断探索云原生 AI 场景中的最佳实践，助力算法科学家们简单高效地实现创新落地。

KubeDL 为分布式训练作业带来了 HostNetwork 网络模式，支持计算节点之间通过宿主机网络相互通信以提升网络性能，同时适应 RDMA/SCC 等新型高性能数据中心架构的网络环境，此外，KubeDL 针对 HostNetwork 模式带来的 FailOver 后新端口互相感知等问题也带来了新的解决思路。

Github 地址：
https://github.com/kubedl-io/kubedl

网站：
https://kubedl.io/model/intro/

Overlay 不是银弹

Kubernetes 原生的容器网络模型定义了一系列不依赖 NAT 的"Pod-Pod"间通信规约，基于 VxLAN 组建的 Overlay 网络很好地实现了这一模型（如经典的 Flannel）并解决了诸多大规模容器编排系统中的网络管理的痛点：

• Pod 的无感迁移：Overlay 网络是基于物理网络构建的虚拟二层网络，Pod IP 并不与任何节点绑定，当节点宕机或发生其他硬件异常时，对应的服务 Pod 可以通过相同的 IP 在其他节点上重新启动，只要底层的物理网络连通不中断就不影响服务的可用性。在大规模的分布式机器学习训练中。KubeDL 也是基于“Pod 可能漂移，但 Service 是固定的”这一前提实现的计算节点故障转移（FailOver）；
• 网络节点的规模：经典的 NAT 地址解析通常通过 ARP 广播协议来自动学习邻接节点 IP 与 MAC 地址的映射，但当节点规模庞大时，一次广播很容易造成 ARP 风暴并引起网络拥塞，而基于隧道穿越的 Overlay 网络只需知道少数的 VTEP 节点的 MAC 地址即能实现数据包的转发，极大的降低了网络的压力；
• 租户网络隔离：Kubernetes 强大的网络插件扩展性配合 VxLAN 的协议设计，很容易实现虚拟网络的再划分从而实现租户之间的网络隔离；

这些都是虚拟容器网络带来的好处，但虚拟化的代价是网络性能的损耗：Pod 与主机网络通过一对 Veth 虚拟网桥设备连接来实现网络 namespace 的互相隔离，每一次"Pod-Pod"间通信的数据包都需要经过”封包-路由-以太网-路由-拆包“等流程才能到达对端的 Pod，拖慢网络性能的同时还会增加宿主机内核网络栈的处理压力从而提升负载。

随着多模态模型训练、大规模稠密参数模型训练等分布式训练模式的兴起，以及数据集规模、特征参数的爆炸，网络通信已然成为分布式训练效率的一个“水桶短板”。最直接的优化网络性能的方法即使用主机网络（HostNetwork）通信，免去容器网络虚拟化的开销。同时随着 RDMA（RoCE），Nvidia GPU Direct 等技术的成熟，这些新型的高性能网络技术逐渐被应用于大规模的商业生产环境来大幅提升模型训练的效率，通过旁路内核网络栈的开销和零拷贝直读数据等技术充分利用网络带宽，Efficiency Is Money！这些原生的高性能网络通信库原语（如 RDMA_CM）也同样依赖主机网络实现，无法直接基于 Pod 虚拟网络通信。

KubeDL 在支持分布式训练基于标准容器网络通信的基础上扩展了主机网络的通信模型，同时解决了端口冲突和 FailOver 后新端口互相感知等分布式训练中的常见问题，实现高性能网络的轻松使能。

使能 Host 高性能网络

标准容器网络拓扑

在标准的容器网络通信模型中，Master/Worker/PS 等不同 Workload 角色之间通过 Headless Service 实现服务发现，Pod 之间通过恒定的域名相互通信，由 CoreDNS 实现域名到 Pod IP 的解析，由于 Pod 是可以漂移的但 Service 及其附属的域名是恒定的，即使部分 Pod 运行时异常了也能很好地实现 FailOver，在异常 Pod 重新拉起之后与其他 Pod 重连接。

apiVersion: training.kubedl.io/v1alpha1
kind: "TFJob"
metadata:
  name: "mnist"
  namespace: kubedl
spec:
  cleanPodPolicy: None
  tfReplicaSpecs:
    PS:
      replicas: 2
      restartPolicy: Never
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubedl/tf-mnist-with-summaries:1.0
              command:
                - "python"
                - "/var/tf_mnist/mnist_with_summaries.py"
                - "--log_dir=/train/logs"
                - "--learning_rate=0.01"
                - "--batch_size=150"
              volumeMounts:
                - mountPath: "/train"
                  name: "training"
              resources:
                limits:
                  cpu: 2048m
                  memory: 2Gi
                requests:
                  cpu: 1024m
                  memory: 1Gi
          volumes:
            - name: "training"
              hostPath:
                path: /tmp/data
                type: DirectoryOrCreate
    Worker:
      replicas: 3
      restartPolicy: ExitCode
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubedl/tf-mnist-with-summaries:1.0
              command:
                - "python"
                - "/var/tf_mnist/mnist_with_summaries.py"
                - "--log_dir=/train/logs"
                - "--learning_rate=0.01"
                - "--batch_size=150"
              volumeMounts:
                - mountPath: "/train"
                  name: "training"
              resources:
                limits:
                  cpu: 2048m
                  memory: 2Gi
                requests:
                  cpu: 1024m
                  memory: 1Gi
          volumes:
            - name: "training"
              hostPath:
                path: /tmp/data
                type: DirectoryOrCreate

以一个经典 PS-Worker 架构的 Tensorflow 分布式训练作业为例，Worker 负责计算参数的梯度，由 PS 负责聚合、更新并广播参数，因此每个 PS 都可能和所有 Worker 建立连接并通信，反之亦是。

在 Tensorflow 框架的实现中，这样一个作业间拓扑结构由一个 TF Cluster Spec 结构来描述，每个 Role（PS or Worker）实例都包含一个 Index 标识自身索引号，可以通过Role+Index 获取自身或其他Role实例的服务地址，即可建立连接开始通信。在标准容器网络模式中，用户提交以下 TFJob，KubeDL 会生成 TF Cluster Spec 并以环境变量的形式传入并被框架接收，同时为每个 Role 实例都准备好 Headless Service，它的 Endpoint 域名地址即对应 TF Cluster Spec 中的服务地址，每个 Pod 都拥有一份独立的 Linux Network Namespace，Pod 的端口地址空间也相互隔离，因此调度到相同的 Node 上也可以使用相同的容器端口。

至此不同 Role 的实例间就能通过 Tensorflow 原生的方式开始分布式训练及通信。

标准容器网络的好处显而易见，简单直观的网络设置，FailOver 友好的网络容错，都使得这一方案能够满足大多数场景下的需求。但对高性能网络有诉求的场景下又该如何运转呢？KubeDL 给出了主机网络的解决方案。

Host 容器网络拓扑

沿用以上的例子，启用主机网络的方式很简单，只要给 TFJob 追加一个 annotation 即可，其余的作业配置都无需特殊改造，如下所示：

apiVersion: training.kubedl.io/v1alpha1
kind: "TFJob"
metadata:
  name: "mnist"
  namespace: kubedl
  annotations:
    kubedl.io/network-mode: host
spec:
  cleanPodPolicy: None
  tfReplicaSpecs:
    PS:
    ...
    Worker:
    ...

当 KubeDL 发现该作业声明了使用主机网络后，会通过以下步骤完成网络的连接设置：

• 创建 Pod 时不再使用固定端口，而是在一定端口范围内随机出一个主机端口，并设置对应暴露的容器端口号，通过上下文的方式传递到后续的控制流中；
• 对 Pod 启用 HostNetwork 并设置 DNS 解析策略为 Host 优先；
• 不再创建 Headless Service，取而代之的是一个正常的流量转发 Service，暴露端口为原先的恒定值，目标端口为 Pod 的真实值；
• 生成的 TF Cluster Spec 中，自身对应的 Role+Index 可见 Local 地址端口为真实的主机端口，其他 Role 实例的地址端口都是恒定的，无论对方的 Pod 如何漂移都能通过 Service 正确转发；
• 当发生 FailOver 时，KubeDL 会为重建后的 Pod 重新选择端口，新启动的 Pod 会通过 TF_CONFIG 得到新的 Local 地址端口，同时 KubeDL 保证对应 Service 的目标端口得到正确更新，其他与之相连的 Role 也能在 Service 目标端口更新后继续通信；

这样一个根据训练作业拓扑结构搭建的主机网络就准备换好了，与之前的不同之处在于，所有的 Pod 都与主机共用了一个 Network Namespace，因此也共享了主机的端口号，而 Pod 之间的通信也从原先通过解析域名为 Pod IP 并建立连接，变成了通过 Service 实现流量的转发，另一方面 TF Cluster Spec 发生了变化但没有改变原生 Tensorflow 的模式，当前 Pod 直接获得 Local Port 监听，而其他的 Pod 地址看起来都是恒定的 Service 对应的域名和暴露的端口永远恒定，只有目标端口可能随着 FailOver 不断改变，这一切都通过 KubeDL 处理变得无感。

我们以 Tensorflow 作为主机网络的例子，因为它的 Cluster Spec 复杂性更具代表性，但 KubeDL 的内置工作负载（如 PyTorch，XGBoost 等）我们也都针对其框架的行为实现了对应主机网络模式的网络拓扑设置。

总结

KubeDL 通过扩展现有的分布式训练作业标准容器网络通信模式，实现了基于原生主机网络的通信模式，在常见训练场景下获得网络性能增益的同时，也完美适应了 RDMA/SCC 等高性能网络架构的环境，助力分布式训练作业运行效率的大幅提升，这一通信模式已经在阿里巴巴内部的生产集群中广泛使用，比如达摩院在云栖大会最新发布的 AliceMind 超大模型就是通过 KubeDL 主机网络+RDMA 在高性能计算集群中训练的产物。我们期待更多开发者参与 KubeDL 社区的建设，一起优化深度学习工作负载的调度及运行时效率！

戳此处，立即了解 KubeDL 项目！