使用 Google Kubernetes Engine 在容器中进行训练

本页面介绍了如何在 Deep Learning Containers 实例中运行训练作业，以及如何在 Google Kubernetes Engine 集群上运行该容器映像。

准备工作

在开始之前，请确保您已完成以下步骤。

  1. 完成本地 Deep Learning Containers 容器使用入门的“准备工作”部分中的设置步骤。

  2. 确保您的 Google Cloud 项目已启用结算功能。

    了解如何启用结算功能

  3. 启用 Google Kubernetes Engine、Compute Engine 和 Artifact Registry API。

    启用 API

打开命令行工具

您可以在本地使用 Cloud Shell 或命令行工具按照本指南进行操作。Cloud Shell 预安装了本教程中使用的 gcloud、docker 和 kubectl 命令行工具。如果使用 Cloud Shell，则无需在工作站上安装这些命令行工具。

Cloud Shell

如需使用 Cloud Shell，请完成以下步骤。

  1. 前往 Google Cloud 控制台。

  2. 点击控制台窗口顶部的激活 Cloud Shell 按钮。

    Google Cloud Platform console

    一个 Cloud Shell 会话随即会在 Console 底部的新框内打开，并显示命令行提示符。

    Cloud Shell 会话

本地命令行

如需使用本地命令行，请完成以下步骤。

  1. 使用 gcloud CLI 安装 Kubernetes 命令行工具。 kubectl 用于与 Kubernetes 通信，Kubernetes 是 Deep Learning Containers 集群的集群编排系统：

    gcloud components install kubectl

    在完成使用入门步骤后，您已安装 Google Cloud CLI 和 Docker。

创建 GKE 集群

运行以下命令，在 GKE 中创建一个名为 pytorch-training-cluster 的双节点集群：

gcloud container clusters create pytorch-training-cluster \
    --num-nodes=2 \
    --zone=us-west1-b \
    --accelerator="type=nvidia-tesla-p100,count=1" \
    --machine-type="n1-highmem-2" \
    --scopes="gke-default,storage-rw"

如需详细了解这些设置，请参阅关于创建运行容器的集群的文档。

集群可能需要几分钟才能完成创建。

或者，您也可以不创建集群，而是使用 Google Cloud 项目中的现有集群。如果您这样做，则可能需要运行以下命令，以确保 kubectl 命令行工具具有访问集群的正确凭据：

gcloud container clusters get-credentials YOUR_EXISTING_CLUSTER

接下来，安装 NVIDIA GPU 设备驱动程序。

创建 Dockerfile

您可以通过多种方式构建容器映像。以下步骤展示如何构建一个容器映像来运行名为 trainer.py 的 Python 脚本。

如需查看可用容器映像的列表，请使用以下命令：

gcloud container images list \
    --repository="us-docker.pkg.dev/deeplearning-platform-release/gcr.io"

您可以查看选择容器，了解如何选择需要的容器。

以下示例展示如何将名为 trainer.py 的 Python 脚本放入特定的 PyTorch 深度学习容器类型。

要创建 dockerfile，请将以下命令写入名为 Dockerfile 的文件。此步骤假定您在名为 model-training-code 的目录中有训练机器学习模型的代码，并且该目录中的主 Python 模块名为 trainer.py。在这种情况下，容器将在作业完成后被移除，因此您的训练脚本应配置为输出到 Cloud Storage（请参阅输出到 Cloud Storage 的脚本的示例）或输出到永久性存储空间。

FROM us-docker.pkg.dev/deeplearning-platform-release/gcr.io/pytorch-gpu
COPY model-training-code /train
CMD ["python", "/train/trainer.py"]

构建和上传容器映像

如需构建容器映像并将其上传到 Artifact Registry，请使用以下命令：

export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
export IMAGE_REPO_NAME=pytorch_custom_container
export IMAGE_TAG=$(date +%Y%m%d_%H%M%S)
export IMAGE_URI=us-docker.pkg.dev/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG

docker build -f Dockerfile -t $IMAGE_URI ./

docker push $IMAGE_URI

部署应用

创建一个包含以下内容的名为 pod.yaml 的文件，并将 IMAGE_URI 替换为映像的 URI。

apiVersion: v1
kind: Pod
metadata:
  name: gke-training-pod
spec:
  containers:
  - name: my-custom-container
    image: IMAGE_URI
    resources:
      limits:
        nvidia.com/gpu: 1

使用 kubectl 命令行工具运行以下命令并部署应用：

kubectl apply -f ./pod.yaml

如需跟踪 pod 的状态，请运行以下命令：

kubectl describe pod gke-training-pod