在 AWS 和 Amazon SageMaker 上使用 Kubeflow 构建灵活且可扩展的分布式训练架构

在这篇文章中，我们演示了如何 AWS 上的 Kubeflow （一种特定于 AWS 的 Kubeflow 发行版）与 AWS深度学习容器 和 亚马逊弹性文件系统 (Amazon EFS) 简化了协作，并为在两个平台上大规模训练深度学习模型提供了灵活性 Amazon Elastic Kubernetes服务 （亚马逊 EKS）和 亚马逊SageMaker 利用混合架构方法。

机器学习 (ML) 开发依赖于复杂且不断发展的开源框架和工具包，以及复杂且不断发展的硬件生态系统。 这在将 ML 开发扩展到集群时提出了挑战。 容器提供了一种解决方案，因为它们不仅可以完全封装训练代码，还可以将整个依赖堆栈封装到硬件库。 这确保了一致且可移植的 ML 环境，并有助于在训练集群的每个单独节点上重复训练环境。

Kubernetes 是一种广泛采用的系统，用于自动化基础设施部署、资源扩展和这些容器化应用程序的管理。 但是，Kubernetes 在构建时并未考虑到 ML，因此由于严重依赖 YAML 规范文件，因此对数据科学家来说可能会违反直觉。 没有 Jupyter 体验，也没有很多特定于 ML 的功能，例如工作流管理和管道，以及 ML 专家期望的其他功能，例如超参数调整、模型托管等。 可以构建这样的功能，但 Kubernetes 并非旨在将其作为其主要目标。

开源社区注意到并在 Kubernetes 之上开发了一个名为 Kubeflow 的层。 Kubeflow 旨在使端到端 ML 工作流在 Kubernetes 上的部署变得简单、可移植和可扩展。 您可以使用 Kubeflow 将用于 ML 的同类最佳开源系统部署到不同的基础设施。

Kubeflow 和 Kubernetes 为数据科学家团队提供了灵活性和控制力。 然而，确保在减少运营开销的情况下大规模运行的训练集群的高利用率仍然具有挑战性。

这篇博文展示了具有本地限制或现有 Kubernetes 投资的客户如何通过在 AWS 上使用 Amazon EKS 和 Kubeflow 基于自我管理的方法实施用于分布式训练的 ML 管道，并使用完全托管的 SageMaker 来应对这一挑战。成本优化、全面管理和生产规模的培训基础设施。 这包括逐步实施混合分布式训练架构，允许您在运行时在两种方法之间进行选择，从而在严格的部署需求下提供最大的控制和灵活性。 您将了解如何在深度学习训练脚本中继续使用开源库，并使其兼容以与平台无关的方式在 Kubernetes 和 SageMaker 上运行。

AWS 上的 Kubeflow 和 SageMaker 有何帮助？

使用 TensorFlow、PyTorch、MXNet 等深度学习框架构建的神经网络模型通过使用更大的训练数据集提供了更高的准确度，尤其是在计算机视觉和自然语言处理用例中。 然而，对于大型训练数据集，训练深度学习模型需要更长的时间，这最终会减慢上市时间。 如果我们可以扩展集群并将模型训练时间从几周缩短到几天或几小时，它可能会对生产力和业务速度产生巨大影响。

Amazon EKS 帮助预置托管的 Kubernetes 控制平面。 您可以使用 Amazon EKS 创建具有 CPU 和 GPU 实例的大型训练集群，并使用 Kubeflow 工具包提供对 ML 友好的开源工具，并使用 Kubeflow Pipelines 操作可移植和可扩展的 ML 工作流，以提高您团队的生产力和缩短上市时间。

但是，这种方法可能存在一些挑战：

• 确保跨数据科学团队最大限度地利用集群。 例如，您应该按需配置 GPU 实例并确保其在要求苛刻的生产规模任务（如深度学习训练）中的高利用率，并将 CPU 实例用于要求不高的任务（如数据预处理）
• 确保部署在 Kubernetes 集群工作节点中的重量级 Kubeflow 基础架构组件（包括数据库、存储和身份验证）的高可用性。 例如，Kubeflow 控制平面生成的工件（例如 MySQL 实例、Pod 日志或 MinIO 存储）会随着时间的推移而增长，并且需要具有持续监控功能的可调整大小的存储卷。
• 在开发人员、训练集群和项目之间共享训练数据集、代码和计算环境具有挑战性。 例如，如果您正在开发自己的一组库，并且这些库具有很强的相互依赖性，那么在同一团队的数据科学家之间共享和运行同一段代码就变得非常困难。 此外，每次训练运行都需要您下载训练数据集并使用新的代码更改构建训练图像。

AWS 上的 Kubeflow 有助于应对这些挑战，并提供企业级的半托管 Kubeflow 产品。 借助 AWS 上的 Kubeflow，您可以将一些 Kubeflow 控制平面服务（如数据库、存储、监控和用户管理）替换为 AWS 托管服务，例如 亚马逊关系数据库服务 （亚马逊 RDS）， 亚马逊简单存储服务 （Amazon S3）， 亚马逊弹性文件系统 （Amazon EFS）， 亚马逊FSx, 亚马逊CloudWatch及 亚马逊Cognito.

替换这些 Kubeflow 组件可将 Kubeflow 控制平面的关键部分与 Kubernetes 分离，从而提供安全、可扩展、弹性和成本优化的设计。 这种方法还从 EKS 数据平面释放存储和计算资源，分布式模型训练或用户笔记本服务器等应用程序可能需要这些资源。 AWS 上的 Kubeflow 还提供 Jupyter 笔记本与深度学习容器 (DLC) 图像的原生集成，这些图像预先打包并预配置了 AWS 优化的深度学习框架，例如 PyTorch 和 TensorFlow，让您无需处理即可立即开始编写训练代码具有依赖性解决方案和框架优化。 此外，Amazon EFS 与训练集群和开发环境的集成允许您共享代码和处理后的训练数据集，从而避免在每次代码更改后构建容器映像和加载大量数据集。 这些与 AWS 上 Kubeflow 的集成可帮助您加快模型构建和训练时间，并通过更轻松的数据和代码共享实现更好的协作。

AWS 上的 Kubeflow 有助于构建高度可用且强大的 ML 平台。 该平台提供了构建和训练深度学习模型的灵活性，并提供了对许多开源工具包的访问、对日志的洞察以及用于实验的交互式调试。 然而，在数百个 GPU 上训练深度学习模型的同时实现基础设施资源的最大利用仍然涉及大量的运营开销。 这可以通过使用 SageMaker 来解决，这是一项完全托管的服务，专为处理性能和成本优化的培训集群而设计和优化，这些集群仅在请求时进行配置，根据需要进行扩展，并在作业完成时自动关闭，从而提供近 100 ％ 资源利用率。 您可以使用托管 SageMaker 组件将 SageMaker 与 Kubeflow Pipelines 集成。 这允许您将 ML 工作流作为 Kubeflow 管道的一部分进行操作，您可以在其中使用 Kubernetes 进行本地培训，并使用 SageMaker 在混合架构中进行产品规模的培训。

解决方案概述

以下架构描述了我们如何使用 Kubeflow Pipelines 构建和部署可移植且可扩展的端到端 ML 工作流，以基于运行时参数使用 Kubeflow 训练或 SageMaker 在 Kubernetes 上有条件地运行分布式训练。

Kubeflow 训练是一组 Kubernetes Operator，它们为 Kubeflow 添加了对使用不同框架（如 TensorFlow、PyTorch 等）的 ML 模型的分布式训练的支持。 pytorch-operator 是 Kubernetes 的 Kubeflow 实现 自定义资源 (PyTorchJob) 在 Kubernetes 上运行分布式 PyTorch 训练作业。

我们使用 PyTorchJob Launcher 组件作为 Kubeflow 管道的一部分，在我们需要灵活性并访问所有底层资源以进行交互式调试和分析的实验阶段运行 PyTorch 分布式训练。

我们还使用 Kubeflow Pipelines 的 SageMaker 组件在生产规模上运行我们的模型训练。 这使我们能够利用强大的 SageMaker 功能，例如完全托管的服务、具有最大 GPU 利用率的分布式训练作业，以及通过 亚马逊弹性计算云 （Amazon EC2）竞价型实例。

作为工作流创建过程的一部分，您需要完成以下步骤（如上图所示）来创建此管道：

1. 使用 Kubeflow 清单文件创建 Kubeflow 仪表板并从 Kubeflow 中央仪表板访问 Jupyter 笔记本。
2. 使用 Kubeflow 管道 SDK 使用 Python 代码创建和编译 Kubeflow 管道。 管道编译将 Python 函数转换为工作流资源，这是一种与 Argo 兼容的 YAML 格式。
3. 使用 Kubeflow Pipelines SDK 客户端调用管道服务端点来运行管道。
4. 管道评估条件运行时变量，并在 SageMaker 或 Kubernetes 作为目标运行环境之间做出决定。
5. 使用 Kubeflow PyTorch Launcher 组件在原生 Kubernetes 环境上运行分布式训练，或使用 SageMaker 组件在 SageMaker 托管平台上提交训练。

下图显示了架构中涉及的 Kubeflow Pipelines 组件，使我们能够灵活地在 Kubernetes 或 SageMaker 分布式环境之间进行选择。

用例工作流程

我们使用以下分步方法在 AWS 上使用 Amazon EKS 和 SageMaker 安装和运行分布式训练用例。

先决条件

对于本演练，您应该具有以下先决条件：

• An AWS账户.
• 一台带有 Docker 和 AWS命令行界面 (AWS CLI) 已安装。
• 或者，您可以使用 AWS 云9，一个基于云的集成开发环境 (IDE)，可以通过 Web 浏览器完成所有工作。 有关设置说明，请参阅 设置 Cloud9 IDE. 在您的 Cloud9 环境中，选择加号并打开新终端。
• 创建一个角色 具有名称 sagemakerrole. 添加托管策略 AmazonSageMakerFullAccess 和 AmazonS3FullAccess 授予 SageMaker 访问 S3 存储桶的权限。 此角色由作为 Kubeflow Pipelines 步骤的一部分提交的 SageMaker 作业使用。
• 确保您的帐户具有 SageMaker 培训资源类型限制 ml.p3.2xlarge 增加到 2 使用 服务配额控制台

1. 在 AWS 上安装 Amazon EKS 和 Kubeflow

您可以使用多种不同的方法来构建 Kubernetes 集群并部署 Kubeflow。 在这篇文章中，我们专注于一种我们认为可以简化流程的方法。 首先，我们创建一个 EKS 集群，然后在其上部署 Kubeflow on AWS v1.5。 对于这些任务中的每一项，我们都使用了一个相应的开源项目，该项目遵循以下原则 做框架. 我们不是为每个任务安装一组先决条件，而是构建具有所有必要工具并在容器内执行任务的 Docker 容器。

我们在这篇文章中使用了 Do Framework，它使用 Amazon EFS 作为附加组件来自动执行 Kubeflow 部署。 有关用于生产部署的官方 Kubeflow on AWS 部署选项，请参阅 部署.

配置当前工作目录和 AWS CLI

我们配置一个工作目录，以便我们可以将其作为以下步骤的起点：

export working_dir=$PWD

我们还配置了 AWS CLI 配置文件。 为此，您需要一个访问密钥 ID 和秘密访问密钥 AWS身份和访问管理 （我是） 用户 具有管理权限（附加现有托管策略）和编程访问权限的帐户。 请参阅以下代码：

aws configure --profile=kubeflow
AWS Access Key ID [None]: 
AWS Secret Access Key [None]: 
Default region name [None]: us-west-2
Default output format [None]: json

# (In Cloud9, select “Cancel” and “Permanently disable” when the AWS managed temporary credentials dialog pops up)

export AWS_PROFILE=kubeflow

1.1 创建EKS集群

如果您已经有可用的 EKS 集群，则可以跳到下一部分。 对于这篇文章，我们使用 aws-do-eks 项目 创建我们的集群。

1. 首先将项目克隆到您的工作目录中

   cd ${working_dir}
   git clone https://github.com/aws-samples/aws-do-eks
   cd aws-do-eks/

2. 然后构建并运行 aws-do-eks 容器：

   ./build.sh
   ./run.sh

   build.sh 脚本创建一个 Docker 容器映像，其中包含用于配置和操作 EKS 集群的所有必要工具和脚本。 这 run.sh 脚本使用创建的 Docker 映像启动一个容器并保持运行，因此我们可以将其用作我们的 EKS 管理环境。 查看您的状态 aws-do-eks 容器，你可以运行 ./status.sh. 如果容器处于退出状态，您可以使用 ./start.sh 启动容器或重新启动容器的脚本，您可以运行 ./stop.sh 其次是 ./run.sh.

3. 在运行中打开一个shell aws-do-eks 容器：
4. 要查看我们的 KubeFlow 部署的 EKS 集群配置，请运行以下命令：

   vi ./eks-kubeflow.yaml

   默认情况下，此配置会创建一个名为 eks-kubeflow ，在 us-west-2 具有六个 m5.xlarge 节点的区域。 此外，默认情况下未启用 EBS 卷加密。 您可以通过添加启用它 "volumeEncrypted: true" 到节点组，它将使用默认密钥进行加密。 如果需要，修改其他配置设置。

5. 要创建集群，请运行以下命令：

   export AWS_PROFILE=kubeflow
   eksctl create cluster -f ./eks-kubeflow.yaml

   集群配置过程最多可能需要 30 分钟。

6. 要验证集群是否已成功创建，请运行以下命令：

   kubectl get nodes

   对于已成功创建的集群，上述命令的输出类似于以下代码：

   root@cdf4ecbebf62:/eks# kubectl get nodes
   NAME                                           STATUS   ROLES    AGE   VERSION
   ip-192-168-0-166.us-west-2.compute.internal    Ready       23m   v1.21.14-eks-ba74326
   ip-192-168-13-28.us-west-2.compute.internal    Ready       23m   v1.21.14-eks-ba74326
   ip-192-168-45-240.us-west-2.compute.internal   Ready       23m   v1.21.14-eks-ba74326
   ip-192-168-63-84.us-west-2.compute.internal    Ready       23m   v1.21.14-eks-ba74326
   ip-192-168-75-56.us-west-2.compute.internal    Ready       23m   v1.21.14-eks-ba74326
   ip-192-168-85-226.us-west-2.compute.internal   Ready       23m   v1.21.14-eks-ba74326

为 SageMaker 训练作业创建 EFS 卷

在此用例中，您通过使用已存储在 Amazon EFS 中的数据训练深度学习模型来加快 SageMaker 训练作业。 这种选择的好处是直接从 Amazon EFS 中的数据启动您的训练作业，无需移动数据，从而缩短训练开始时间。

我们创建一个 EFS 卷并部署 EFS 容器存储接口 (CSI) 驱动程序。 这是通过位于 /eks/deployment/csi/efs 字幕可视电话用于 aws-do-eks 容器。

此脚本假设您的账户中有一个 EKS 集群。 放 CLUSTER_NAME= 如果您有多个 EKS 集群。

cd /eks/deployment/csi/efs
./deploy.sh

此脚本预置 EFS 卷并为集群 VPC 的子网创建挂载目标。 然后部署 EFS CSI 驱动程序并创建 efs-sc 存储类和 efs-pv EKS 集群中的持久卷。

成功完成脚本后，您应该会看到如下输出：

Generating efs-sc.yaml ...

Applying efs-sc.yaml ...
storageclass.storage.k8s.io/efs-sc created
NAME            PROVISIONER             RECLAIMPOLICY   VOLUMEBINDINGMODE      ALLOWVOLUMEEXPANSION   AGE
efs-sc          efs.csi.aws.com         Delete          Immediate              false                  1s
gp2 (default)   kubernetes.io/aws-ebs   Delete          WaitForFirstConsumer   false                  36m

Generating efs-pv.yaml ...
Applying efs-pv.yaml ...
persistentvolume/efs-pv created
NAME     CAPACITY   ACCESS MODES   RECLAIM POLICY   STATUS      CLAIM   STORAGECLASS   REASON   AGE
efs-pv   5Gi        RWX            Retain           Available           efs-sc                  10s

Done ...

创建 Amazon S3 VPC 终端节点

您使用 SageMaker 训练作业和 EFS 文件系统可以访问的私有 VPC。 要让 SageMaker 训练集群从您的私有 VPC 访问 S3 存储桶，您需要创建一个 VPC 终端节点：

cd /eks/vpc 
export CLUSTER_NAME= 
export REGION= 
./vpc-endpoint-create.sh

您现在可以退出 aws-do-eks 容器外壳并继续下一部分：

exit

root@cdf4ecbebf62:/eks/deployment/csi/efs# exit
exit
TeamRole:~/environment/aws-do-eks (main) $

1.2 在 Amazon EKS 上的 AWS 上部署 Kubeflow

要在 Amazon EKS 上部署 Kubeflow，我们使用 aws-do-kubeflow 项目.

1. 使用以下命令克隆存储库：

   cd ${working_dir}
   git clone https://github.com/aws-samples/aws-do-kubeflow
   cd aws-do-kubeflow

2. 然后配置项目：

   ./config.sh

   此脚本在文本编辑器中打开项目配置文件。 这对 AWS_REGION 设置为您的集群所在的区域，以及 AWS_CLUSTER_NAME 以匹配您之前创建的集群的名称。 默认情况下，您的配置已经正确设置，因此如果您不需要进行任何更改，只需关闭编辑器即可。

   ./build.sh
   ./run.sh
   ./exec.sh

   build.sh 脚本创建一个 Docker 容器镜像，其中包含在现有 Kubernetes 集群上部署和管理 Kubeflow 所需的所有工具。 这 run.sh 脚本使用 Docker 镜像启动一个容器，exec.sh 脚本在容器中打开一个命令 shell，我们可以将其用作我们的 Kubeflow 管理环境。 您可以使用 ./status.sh 脚本看看是否 aws-do-kubeflow 容器已启动并运行，并且 ./stop.sh 和 ./run.sh 脚本以根据需要重新启动它。

3. 在你打开一个shell之后 aws-do-eks 容器，您可以验证配置的集群上下文是否符合预期：

   root@ip-172-31-43-155:/kubeflow# kubectx
   kubeflow@eks-kubeflow.us-west-2.eksctl.io

4. 要在 EKS 集群上部署 Kubeflow，请运行 deploy.sh 脚本：

   ./kubeflow-deploy.sh

   当 kubeflow 命名空间中的所有 Pod 进入 Running 状态时，部署成功。 典型的输出类似于以下代码：

   Waiting for all Kubeflow pods to start Running ...
   
   Waiting for all Kubeflow pods to start Running ...
   
   Restarting central dashboard ...
   pod "centraldashboard-79f489b55-vr6lp" deleted
   /kubeflow/deploy/distro/aws/kubeflow-manifests /kubeflow/deploy/distro/aws
   /kubeflow/deploy/distro/aws
   
   Kubeflow deployment succeeded
   Granting cluster access to kubeflow profile user ...
   Argument not provided, assuming default user namespace kubeflow-user-example-com ...
   clusterrolebinding.rbac.authorization.k8s.io/kubeflow-user-example-com-cluster-admin-binding created
   Setting up access to Kubeflow Pipelines ...
   Argument not provided, assuming default user namespace kubeflow-user-example-com ...
   
   Creating pod-default for namespace kubeflow-user-example-com ...
   poddefault.kubeflow.org/access-ml-pipeline created

5. 要监控 KubeFlow pod 的状态，可以在单独的窗口中使用以下命令：

   watch kubectl -n kubeflow get pods

6. 媒体 按Ctrl + C 当所有 pod 都在运行时，通过运行以下命令将 Kubeflow 仪表板暴露在集群外部：

   ./kubeflow-expose.sh

您应该看到类似于以下代码的输出：

root@ip-172-31-43-155:/kubeflow# ./kubeflow-expose.sh
root@ip-172-31-43-155:/kubeflow# Forwarding from 127.0.0.1:8080 -> 8080
Forwarding from [::1]:8080 -> 8080

此命令将 Istio 入口网关服务从集群端口转发到本地端口 8080。要访问 Kubeflow 仪表板，请访问 http://localhost:8080 并使用默认用户凭据登录（user@example.com/12341234）。 如果您正在运行 aws-do-kubeflow AWS Cloud9 中的容器，然后您可以选择 预览，然后选择 预览正在运行的应用程序. 如果您在 Docker Desktop 上运行，则可能需要运行 ./kubeflow-expose.sh 脚本之外的 aws-do-kubeflow 容器。

2. 设置 Kubeflow on AWS 环境

为了在 AWS 环境中设置您的 Kubeflow，我们创建了一个 EFS 卷和一个 Jupyter 笔记本。

2.1 创建 EFS 卷

要创建 EFS 卷，请完成以下步骤：

• 在 Kubeflow 仪表板上，选择 卷 在导航窗格中。
• 选择 新卷.
• 针对 名字，输入 efs-sc-claim.
• 针对 体积大小，输入 10.
• 针对 存储类，选择 EFS-SC.
• 针对 存取模式，选择 一次读写.
• 创建.

2.2 创建一个 Jupyter 笔记本

要创建新笔记本，请完成以下步骤：

• 在 Kubeflow 仪表板上，选择 笔记本电脑 在导航窗格中。
• 新笔记本.
• 针对 名字，输入 aws-hybrid-nb.
• 针对 Jupyter 案卷图像, 选择图像 c9e4w0g3/notebook-servers/jupyter-pytorch:1.11.0-cpu-py38-ubuntu20.04-e3-v1.1 （最新可用的 jupyter-pytorch DLC 图像）。
• 针对 中央处理器，输入 1.
• 针对 内存，输入 5.
• 针对 图形处理器, 离开为 不包含.
• 不要对 箱内容积 部分。
• 在 数据量 部分中，选择 附加现有卷 并展开现有卷部分
• 针对 名字，选择 efs-sc-claim.
• 针对 挂载路径，输入 /home/jovyan/efs-sc-claim.
  这会将 EFS 卷安装到您的 Jupyter notebook pod，您可以看到该文件夹 efs-sc-claim 在您的 Jupyter 实验室界面中。 您将训练数据集和训练代码保存到此文件夹，以便训练集群可以访问它，而无需重新构建容器映像进行测试。
• 选择 允许访问 Kubeflow Pipelines 在配置部分。
• 实行.
  验证您的笔记本是否已成功创建（可能需要几分钟）。
• 点击 笔记本电脑 页面，选择 分享链接 登录 JupyterLab 环境。
• 点击 混帐 菜单中选择 克隆存储库.
• 针对 克隆一个仓库，输入 https://github.com/aws-samples/aws-do-kubeflow.
  

3. 运行分布式训练

设置 Jupyter 笔记本后，您可以使用文件夹中的以下高级步骤运行整个演示 aws-do-kubeflow/workshop 在克隆的存储库中：

• PyTorch 分布式数据并行 (DDP) 训练脚本： 请参阅 PyTorch DDP 训练脚本 cifar10-distributed-gpu-final.py，其中包括示例卷积神经网络和逻辑，用于在多节点 CPU 和 GPU 集群上进行分布训练。 （详见 3.1）
• 安装库： 运行笔记本 0_initialize_dependencies.ipynb 初始化所有依赖项。 （详见 3.2）
• 在 Kubernetes 上运行分布式 PyTorch 作业训练： 运行笔记本 1_submit_pytorchdist_k8s.ipynb 使用 Kubernetes 自定义资源 PyTorchJob YAML 文件，使用 Python 代码创建和提交一个主容器和两个工作容器的分布式训练。 （详见 3.3）
• 创建一个混合 Kubeflow 管道： 运行笔记本 2_create_pipeline_k8s_sagemaker.ipynb 使用运行时变量创建在 SageMaker 或 Amazon EKS 上运行分布式训练的混合 Kubeflow 管道 training_runtime. （详见 3.4）

确保你运行了笔记本 1_submit_pytorchdist_k8s.ipynb 在你开始笔记本之前 2_create_pipeline_k8s_sagemaker.ipynb.

在随后的部分中，我们将详细讨论这些步骤中的每一个。

3.1 PyTorch 分布式数据并行（DDP）训练脚本

作为分布式训练的一部分，我们训练了一个由在 CIFAR10 数据集上运行的简单卷积神经网络创建的分类模型。 训练脚本 cifar10-distributed-gpu-final.py 仅包含开源库，并且兼容在 GPU 设备或 CPU 实例上的 Kubernetes 和 SageMaker 训练集群上运行。 在运行笔记本示例之前，让我们先看看训练脚本的几个重要方面。

我们使用 torch.distributed 模块，其中包含 PyTorch 支持和通信原语，用于集群中跨节点的多进程并行：

...
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision
from torchvision import datasets, transforms
...

我们使用卷积层、最大池化层和线性层的组合创建了一个简单的图像分类模型，在模型训练的前向传递中应用了 relu 激活函数：

# Define models
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x

我们使用结合了数据集和 DistributedSampler （使用分布式方式加载数据子集 torch.nn.parallel.DistributedDataParallel) 并在数据上提供单进程或多进程迭代器：

# Define data loader for training dataset
def _get_train_data_loader(batch_size, training_dir, is_distributed):
logger.info("Get train data loader")

train_set = torchvision.datasets.CIFAR10(root=training_dir,
train=True,
download=False,
transform=_get_transforms())

train_sampler = (
torch.utils.data.distributed.DistributedSampler(train_set) if is_distributed else None
)

return torch.utils.data.DataLoader(
train_set,
batch_size=batch_size,
shuffle=train_sampler is None,
sampler=train_sampler)
...

如果训练集群有 GPU，则脚本在 CUDA 设备上运行训练，并且 device 变量保存默认的 CUDA 设备：

device = "cuda" if torch.cuda.is_available() else "cpu"
...

在使用 PyTorch 运行分布式训练之前 DistributedDataParallel 要在多个节点上运行分布式处理，您需要通过调用初始化分布式环境 init_process_group. 这是在训练集群的每台机器上初始化的。

dist.init_process_group(backend=args.backend, rank=host_rank, world_size=world_size)
...

我们实例化分类器模型并将模型复制到目标设备。 如果启用分布式训练以在多个节点上运行，则 DistributedDataParallel class 用作模型对象的包装对象，它允许跨多台机器进行同步分布式训练。 输入数据在批次维度上进行拆分，模型的副本放置在每台机器和每台设备上。

model = Net().to(device)

if is_distributed:
model = torch.nn.parallel.DistributedDataParallel(model)

...

3.2 安装库

您将安装所有必要的库来运行 PyTorch 分布式训练示例。 这包括 Kubeflow Pipelines SDK、Training Operator Python SDK、适用于 Kubernetes 的 Python 客户端和 Amazon SageMaker Python SDK。

#Please run the below commands to install necessary libraries

!pip install kfp==1.8.4

!pip install kubeflow-training

!pip install kubernetes

!pip install sagemaker

3.3 在 Kubernetes 上运行分布式 PyTorch 作业训练

笔记本 1_submit_pytorchdist_k8s.ipynb 使用 Kubeflow 训练和 Kubernetes 客户端 Python SDK 创建 Kubernetes 自定义资源 PyTorchJob YAML 文件。 以下是此笔记本中的一些重要片段。

我们使用主容器和工作容器创建 PyTorchJob YAML，如以下代码所示：

# Define PyTorchJob custom resource manifest
pytorchjob = V1PyTorchJob(
api_version="kubeflow.org/v1",
kind="PyTorchJob",
metadata=V1ObjectMeta(name=pytorch_distributed_jobname,namespace=user_namespace),
spec=V1PyTorchJobSpec(
run_policy=V1RunPolicy(clean_pod_policy="None"),
pytorch_replica_specs={"Master": master,
"Worker": worker}
)
)

这是使用提交到 Kubernetes 控制平面的 PyTorchJobClient:

# Creates and Submits PyTorchJob custom resource file to Kubernetes
pytorchjob_client = PyTorchJobClient()

pytorch_job_manifest=pytorchjob_client.create(pytorchjob):

查看 Kubernetes 训练日志

您可以使用 Python 代码从同一个 Jupyter notebook 或从 Kubernetes 客户端 shell 查看训练日志。

3.4 创建混合 Kubeflow 管道

笔记本 2_create_pipeline_k8s_sagemaker.ipynb 基于条件运行时变量创建混合 Kubeflow 管道 training_runtime，如以下代码所示。 笔记本使用 Kubeflow管道SDK 它提供了一组 Python 包来指定和运行 ML 工作流管道。 作为此 SDK 的一部分，我们使用以下软件包：

• 领域特定语言 (DSL) 包装饰器 dsl.pipeline，它装饰 Python 函数以返回管道
• dsl.Condition 包，它代表一组操作，只有在满足某个条件时才运行，例如检查 training_runtime 价值为 sagemaker or kubernetes

请参见以下代码：

# Define your training runtime value with either 'sagemaker' or 'kubernetes'
training_runtime='sagemaker'

# Create Hybrid Pipeline using Kubeflow PyTorch Training Operators and Amazon SageMaker Service
@dsl.pipeline(name="PyTorch Training pipeline", description="Sample training job test")
def pytorch_cnn_pipeline():

# Pipeline Step 1: to evaluate the condition. You can enter any logic here. For demonstration we are checking if GPU is needed for training
condition_result = check_condition_op(training_runtime)

# Pipeline Step 2: to run training on Kuberentes using PyTorch Training Operators. This will be executed if gpus are not needed
with dsl.Condition(condition_result.output == 'kubernetes', name="PyTorch_Comp"):
train_task = pytorch_job_op(
name=training_job_name,
namespace=user_namespace,
master_spec=json.dumps(master_spec_loaded), # Please refer file at pipeline_yaml_specifications/pipeline_master_spec.yml
worker_spec=json.dumps(worker_spec_loaded), # Please refer file at pipeline_yaml_specifications/pipeline_worker_spec.yml
delete_after_done=False
).after(condition_result)

# Pipeline Step 3: to run training on SageMaker using SageMaker Components for Pipeline. This will be executed if gpus are needed
with dsl.Condition(condition_result.output == 'sagemaker', name="SageMaker_Comp"):
training = sagemaker_train_op(
region=region,
image=train_image,
job_name=training_job_name,
training_input_mode=training_input_mode,
hyperparameters='{ 
"backend": "'+str(pytorch_backend)+'", 
"batch-size": "64", 
"epochs": "3", 
"lr": "'+str(learning_rate)+'", 
"model-type": "custom", 
"sagemaker_container_log_level": "20", 
"sagemaker_program": "cifar10-distributed-gpu-final.py", 
"sagemaker_region": "us-west-2", 
"sagemaker_submit_directory": "'+source_s3+'" 
}',
channels=channels,
instance_type=instance_type,
instance_count=instance_count,
volume_size=volume_size,
max_run_time=max_run_time,
model_artifact_path=f's3://{bucket_name}/jobs',
network_isolation=network_isolation,
traffic_encryption=traffic_encryption,
role=role,
vpc_subnets=subnet_id,
vpc_security_group_ids=security_group_id
).after(condition_result)

我们使用两个 ml.p3.2xlarge 实例配置 SageMaker 分布式训练。

定义管道后，您可以使用 Kubeflow Pipelines SDK 将管道编译为 Argo YAML 规范 kfp.compiler 包裹。 您可以使用 Kubeflow Pipeline SDK 客户端运行此管道，该客户端调用 Pipelines 服务端点并直接从笔记本传递适当的身份验证标头。 请参阅以下代码：

# DSL Compiler that compiles pipeline functions into workflow yaml.
kfp.compiler.Compiler().compile(pytorch_cnn_pipeline, "pytorch_cnn_pipeline.yaml")

# Connect to Kubeflow Pipelines using the Kubeflow Pipelines SDK client
client = kfp.Client()

experiment = client.create_experiment(name="kubeflow")

# Run a specified pipeline
my_run = client.run_pipeline(experiment.id, "pytorch_cnn_pipeline", "pytorch_cnn_pipeline.yaml")

# Please click “Run details” link generated below this cell to view your pipeline. You can click every pipeline step to see logs.

如果你得到一个 sagemaker import 错误，运行 !pip install sagemaker 并重新启动内核（在 核心 菜单中选择 重启内核).

选择 运行详情 最后一个单元格下的链接以查看 Kubeflow 管道。

重复管道创建步骤 training_runtime='kubernetes' 测试在 Kubernetes 环境中运行的管道。 这 training_runtime 变量也可以在生产场景中的 CI/CD 管道中传递。

查看 SageMaker 组件的 Kubeflow 管道运行日志

以下屏幕截图显示了 SageMaker 组件的管道详细信息。

选择训练作业步骤并在 日志 选项卡，选择 CloudWatch 日志链接以访问 SageMaker 日志。

以下屏幕截图显示了两个 ml.p3.2xlarge 实例中每个实例的 CloudWatch 日志。

选择任何组以查看日志。

查看 Kubeflow PyTorchJob Launcher 组件的 Kubeflow 管道运行日志

以下屏幕截图显示了我们的 Kubeflow 组件的管道详细信息。

使用以下命令运行 Kubectl 在连接到 Kubernetes 集群的 Kubernetes 客户端 shell 上查看日志（替换您的命名空间和 pod 名称）：

kubectl get pods -n kubeflow-user-example-com
kubectl logs  -n kubeflow-user-example-com -f

4.1 清理

要清理我们在帐户中创建的所有资源，我们需要以相反的顺序将它们删除。

1. 通过运行删除 Kubeflow 安装 ./kubeflow-remove.sh ，在 aws-do-kubeflow 容器。 第一组命令是可选的，可以在您还没有命令 shell 的情况下使用 aws-do-kubeflow 容器打开。

   cd aws-do-kubeflow
   ./status.sh
   ./start.sh
   ./exec.sh
   
   ./kubeflow-remove.sh

2. 来自 aws-do-eks 容器文件夹，删除 EFS 卷。 第一组命令是可选的，可以在您还没有命令 shell 的情况下使用 aws-do-eks 容器打开。

   cd aws-do-eks
   ./status.sh
   ./start.sh
   ./exec.sh
   
   cd /eks/deployment/csi/efs
   ./delete.sh
   ./efs-delete.sh

   为了释放与我们为集群创建的 VPC 关联的网络接口，必须删除 Amazon EFS。 请注意，删除 EFS 卷会破坏其上存储的所有数据。

3. 来自 aws-do-eks 容器，运行 eks-delete.sh 删除集群和与之关联的任何其他资源的脚本，包括 VPC：

   cd /eks
   ./eks-delete.sh

总结

在这篇文章中，我们讨论了分布式模型训练和 ML 工作流程的一些典型挑战。 我们概述了 AWS 上的 Kubeflow 发行版，并分享了两个开源项目 (AWS-do-eks 和 AWS-do-kubeflow) 简化了基础架构的配置和 Kubeflow 在其上的部署。 最后，我们描述并演示了一种混合架构，它使工作负载能够在自我管理的 Kubernetes 和完全管理的 SageMaker 基础架构上运行之间无缝转换。 我们鼓励您将这种混合架构用于您自己的用例。

你可以按照 AWS 实验室存储库 跟踪所有 AWS 对 Kubeflow 的贡献。 您也可以在 Kubeflow #AWS Slack 频道; 您的反馈将帮助我们优先考虑为 Kubeflow 项目做出贡献的下一个功能。

特别感谢 Sree Arasanagatta（软件开发经理 AWS ML）和 Suraj Kota（软件开发工程师）对本文发布的支持。

关于作者

坎瓦尔吉特·库尔米 是 Amazon Web Services 的 AI/ML 专家解决方案架构师。 他与 AWS 产品、工程和客户合作，提供指导和技术援助，帮助他们在使用 AWS 时提高混合 ML 解决方案的价值。 Kanwaljit 专门帮助客户使用容器化和机器学习应用程序。

高塔姆·库玛（Gautam Kumar） 是 AWS AI 深度学习的软件工程师。 他开发了 AWS Deep Learning Containers 和 AWS Deep Learning AMI。 他热衷于为 AI 构建工具和系统。 在业余时间，他喜欢骑自行车和看书。

亚历克斯·扬库尔斯基 是一位全栈软件和基础架构架构师，喜欢做深入的动手工作。 他目前是 AWS 自我管理机器学习的首席解决方案架构师。 在他的职位上，他专注于帮助客户在容器驱动的 AWS 服务上实现 ML 和 AI 工作负载的容器化和编排。 他也是开源的作者 做框架 还有一位喜欢应用容器技术来加快创新步伐，同时解决世界上最大挑战的 Docker 船长。 在过去的 10 年里，Alex 一直致力于应对气候变化、使 AI 和 ML 民主化、使旅行更安全、医疗保健更好、能源更智能。