ElasticDL：一个原生 Kubernetes 的深度学习框架

ElasticDL 是一个原生 Kubernetes 的深度学习框架，支持容错和弹性调度。

主要特性

弹性调度与容错

通过原生 Kubernetes 的设计，ElasticDL 实现了容错能力，并结合 Kubernetes 的基于优先级的抢占机制，为深度学习任务提供弹性调度功能。

支持 TensorFlow 和 PyTorch

• TensorFlow Estimator。
• TensorFlow Keras。
• PyTorch

极简接口

只需使用 Keras API 定义一个 模型，即可通过命令行实现分布式训练。

elasticdl train \
  --image_name=elasticdl:mnist \
  --model_zoo=model_zoo \
  --model_def=mnist.mnist_functional_api.custom_model \
  --training_data=/data/mnist/train \
  --job_name=test-mnist \
  --volume="host_path=/data,mount_path=/data"

快速入门

请参阅我们的分步教程，了解如何在本地笔记本电脑、自建集群或 Google Kubernetes Engine 等公有云上运行 ElasticDL。

TensorFlow Estimator 在 MiniKube 上

TensorFlow Keras 在 MiniKube 上

PyTorch 在 MiniKube 上

背景

TensorFlow/PyTorch 自带的分布式计算功能具有故障恢复能力。当部分进程失败时，分布式计算作业会失败；不过，我们可以重启作业，并从最近的检查点文件中恢复其状态。

ElasticDL 在分布式训练过程中支持容错。即使部分进程失败，作业仍将继续运行。因此，ElasticDL 无需保存检查点或从检查点恢复。

这种容错特性使 ElasticDL 能够与 Kubernetes 的基于优先级的抢占机制结合，实现弹性调度。当 Kubernetes 杀死某个作业的部分进程以释放资源给更高优先级的新作业时，当前作业不会失败，而是以更少的资源继续运行。

弹性调度可以显著提高集群的整体利用率。假设一个集群有 N 个 GPU，而某个作业正在使用其中一个。如果没有弹性调度，另一个申请 N 个 GPU 的新作业必须等到第一个作业完成才能开始。这段时间可能长达数小时、数天甚至数周，在此期间集群的利用率仅为 1/N。而有了弹性调度，新作业可以立即使用 N-1 个 GPU 开始运行，待第一个作业完成后，Kubernetes 可能会再为其分配 1 个 GPU。在这种情况下，集群的整体利用率将达到 100%。

ElasticDL 的弹性调度特性源于其原生 Kubernetes 的设计——它不依赖 Kubeflow 等 Kubernetes 扩展来运行 TensorFlow/PyTorch 程序。相反，ElasticDL 作业的主进程会调用 Kubernetes API 来启动工作节点和参数服务器，并监听进程/Pod 被杀死等事件，从而实现容错。

简而言之，如果你拥有一个 Kubernetes 集群，ElasticDL 可以增强 TensorFlow/PyTorch 的容错能力和弹性调度功能。我们提供了一个教程，展示如何在 Google Cloud 上搭建 Kubernetes 集群并在其中运行 ElasticDL 作业。我们尊重 TensorFlow 原生的分布式计算功能，该功能并不依赖于特定的计算平台（如 Kubernetes），允许 TensorFlow 在任何平台上运行。

开发指南

请参阅 本文档 获取开发指南。

ElasticDL 快速上手指南

ElasticDL 是一个云原生（Kubernetes-native）的深度学习框架，专为支持容错和弹性调度而设计。它允许深度学习任务在 Kubernetes 集群中根据资源优先级动态调整，无需依赖检查点（Checkpoint）即可从节点故障中恢复。

环境准备

在开始之前，请确保满足以下系统要求和前置依赖：

• 操作系统：Linux 或 macOS（Windows 用户建议使用 WSL2 或虚拟机）。
• Kubernetes 集群：需要一个可用的 K8s 集群。
  • 本地开发：推荐安装 Minikube 或 Kind。
  • 云端环境：支持 Google Kubernetes Engine (GKE)、阿里云 ACK 等主流云厂商集群。
• 命令行工具：
  • kubectl：用于与 Kubernetes 集群交互。
  • docker：用于构建和推送镜像。
• Python 环境：Python 3.6+。
• 深度学习框架：根据需求安装 TensorFlow 或 PyTorch（ElasticDL 客户端会处理大部分依赖，但本地调试模型定义时可能需要）。

提示：国内开发者若拉取 Docker 镜像或 Python 包较慢，建议配置阿里云或腾讯云的镜像加速器，并使用国内 PyPI 源（如 pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...）。

安装步骤

ElasticDL 主要通过 Python pip 安装客户端工具，并通过 Docker 镜像运行训练任务。

1. 安装 ElasticDL 客户端

   使用 pip 安装最新版本的客户端：

   pip install elasticdl-client
   

   国内加速安装：

   pip install elasticdl-client -i https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 验证安装

   安装完成后，检查命令是否可用：

   elasticdl --help
   

3. 准备 Docker 镜像

   ElasticDL 需要特定的 Docker 镜像来运行 Worker 和 Master。你可以从 Docker Hub 拉取官方镜像，或根据项目提供的 Dockerfile 自行构建。

   docker pull elasticdl:mnist
   

   注意：如果网络受限，建议自行克隆源码并在本地构建镜像。

基本使用

ElasticDL 的核心优势在于极简的接口。只需定义好模型（支持 Keras、TensorFlow Estimator 或 PyTorch），即可通过一行命令启动分布式训练。

以下是一个基于 Keras API 定义的 MNIST 模型进行分布式训练的示例：

elasticdl train \
  --image_name=elasticdl:mnist \
  --model_zoo=model_zoo \
  --model_def=mnist.mnist_functional_api.custom_model \
  --training_data=/data/mnist/train \
  --job_name=test-mnist \
  --volume="host_path=/data,mount_path=/data"

参数说明：

• --image_name：指定包含运行环境的 Docker 镜像名称。
• --model_zoo：包含模型定义文件的目录路径。
• --model_def：模型定义的具体入口（格式为 目录名.文件名.类名/函数名）。
• --training_data：训练数据在容器内的挂载路径。
• --job_name：当前训练任务在 Kubernetes 中的名称。
• --volume：将宿主机的数据目录挂载到容器中，格式为 host_path=<宿主机路径>,mount_path=<容器内路径>。

运行流程： 执行上述命令后，ElasticDL Master 会自动调用 Kubernetes API 启动 Worker 和 Parameter Server Pod。即使部分 Pod 因资源抢占被杀死，任务也会自动继续运行，无需人工干预或重新加载检查点。

更多针对 Minikube、本地集群或公有云的详细分步教程，请参考官方文档中的 Step-by-step Tutorial。