在 AI Platform Training 上使用容器

本指南介绍了在 AI Platform Training 上如何构建自己的自定义容器以运行作业。

使用容器涉及的步骤

以下步骤演示了使用自定义容器进行训练的基本过程:

  1. 设置 Google Cloud 项目和您的本地环境。
  2. 创建自定义容器:
    1. 编写一个 Dockerfile,用于设置容器以使用 AI Platform Training,并包含训练应用所需的依赖项。
    2. 在本地构建和测试 Docker 容器。
  3. 将容器推送到 Container Registry。
  4. 提交在自定义容器上运行的训练作业。

使用超参数调节或 GPU 需要一些调整,但基本过程是相同的。

准备工作

使用 Cloud Shell 或任何安装了 gcloud CLI 的环境。

完成以下步骤以设置 GCP 账号,启用所需的 API,以及安装和激活 Cloud SDK。

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud 新手,请创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  2. 在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目

    转到“项目选择器”

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

  4. 启用 AI Platform Training & Prediction, Compute Engine and Container Registry API。

    启用 API

  5. 安装 Google Cloud CLI。
  6. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  7. 在 Google Cloud Console 中的项目选择器页面上,选择或创建一个 Google Cloud 项目

    转到“项目选择器”

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

  9. 启用 AI Platform Training & Prediction, Compute Engine and Container Registry API。

    启用 API

  10. 安装 Google Cloud CLI。
  11. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  12. 安装 Docker

    如果您使用的是基于 Linux 的操作系统(例如 Ubuntu 或 Debian),请将您的用户名添加到 docker 群组,这样您就可以在不使用 sudo 的情况下运行 Docker:

    sudo usermod -a -G docker ${USER}

    在将您自己添加到 docker 群组之后,您可能需要重启系统。

  13. 打开 Docker。如需确保 Docker 正在运行,请运行以下 Docker 命令以返回当前时间和日期:
    docker run busybox date
  14. 使用 gcloud 作为 Docker 的凭据帮助程序:
    gcloud auth configure-docker
  15. 可选:如果要在本地使用 GPU 运行容器,请安装 nvidia-docker

创建自定义容器

创建自定义容器涉及编写 Dockerfile 以设置您将用于训练作业的 Docker 映像。之后,您可以在本地构建和测试映像。

适用于 AI Platform Training 的 Dockerfile 基础知识

创建自定义容器时,使用 Dockerfile 指定构建映像所需的所有命令。

本节将介绍 Dockerfile 的通用示例。您可以在每个自定义容器教程和相关示例中找到具体示例。

要与 AI Platform Training 搭配使用,您的 Dockerfile 需要包含涵盖以下任务的命令:

  • 选择基础映像
  • 安装其他依赖项
  • 将训练代码复制到此映像
  • 配置 AI Platform Training 的入口点以调用您的训练代码

您的 Dockerfile 可能包含其他逻辑,具体取决于您的需求。详细了解编写 Dockerfile;如需详细了解每个具体命令,请参阅 Dockerfile 参考

Dockerfile 命令 说明 示例
FROM image:tag 指定基础映像及其标记。

带标记的示例基础映像:

  • pytorch/pytorch:latest
  • tensorflow/tensorflow:nightly
  • python:2.7.15-jessie
  • nvidia/cuda:9.0-cudnn7-runtime
WORKDIR /path/to/directory 在运行后续指令的映像上指定目录。 /root
RUN pip install pkg1 pkg2 pkg3 使用 pip 安装其他软件包。

注意:如果您的基础映像没有 pip,则必须在安装其他软件包之前使用相应命令来安装它。

示例软件包:

  • google-cloud-storage
  • cloudml-hypertune
  • pandas
COPY src/foo.py dest/foo.py 将训练应用的代码复制到映像中。根据您的训练应用的结构,要复制的内容可能包括多个文件。

训练应用中文件的示例名称:

  • model.py
  • task.py
  • data_utils.py
ENTRYPOINT ["exec", "file"] 设置入口点以调用要运行的训练代码。 ["python", "task.py"]

Dockerfile 中的逻辑可能会根据您的需要而有所不同,但通常包含的内容如下:

# Specifies base image and tag
FROM image:tag
WORKDIR /root

# Installs additional packages
RUN pip install pkg1 pkg2 pkg3

# Downloads training data
RUN curl https://example-url/path-to-data/data-filename --output /root/data-filename

# Copies the trainer code to the docker image.
COPY your-path-to/model.py /root/model.py
COPY your-path-to/task.py /root/task.py

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "task.py"]

在本地构建和测试 Docker 容器

  1. 使用环境变量创建正确的映像 URI,然后构建 Docker 映像。-t 标志使用您选择的 IMAGE_REPO_NAMEIMAGE_TAG 命名和标记映像。您可以为映像选择不同的名称和标记。

    export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
    export IMAGE_REPO_NAME=example_custom_container_image
    export IMAGE_TAG=example_image_tag
    export IMAGE_URI=gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG
    
    docker build -f Dockerfile -t $IMAGE_URI ./
    
  2. 通过在本地运行映像进行验证。请注意,--epochs 标志将传递给训练程序脚本。

    docker run $IMAGE_URI --epochs 1
    

将容器推送到 Container Registry

如果本地运行有效,您可以将容器推送到项目的 Container Registry。

将容器推送到 Container Registry。在此之前,请先运行 gcloud auth configure-docker(如果您尚未这样做)。

docker push $IMAGE_URI

管理 Container Registry 权限

如果您使用的 Container Registry 映像位于您用于在 AI Platform Training 上运行训练的同一项目中,则无需为本教程配置权限,您可以继续执行下一步。

Container Registry 的访问权限控制基于后台的 Cloud Storage 存储桶进行,因此配置 Container Registry 权限与配置 Cloud Storage 权限的过程非常相似。

如果要从其他项目中的 Container Registry 拉取映像,则需要允许 AI Platform Training 服务账号访问该项目中的映像。pull

如果要将 Docker 映像推送到与用于提交 AI Platform Training 训练作业的项目不同的项目,则您应该在具有 Container Registry 代码库的项目中向 AI Platform Training 服务账号授予映像拉取访问权限。push服务账号的格式为 service-$CMLE_PROJ_NUM@cloud-ml.google.com.iam.gserviceaccount.com,您可以在 IAM 控制台中找到该账号。

以下命令会将您的 AI Platform Training 服务账号添加到单独的 Container Registry 项目中:

export GCR_PROJ_ID=[YOUR-PROJECT-ID-FOR-GCR]
export CMLE_PROJ_NUM=[YOUR-PROJECT-NUMBER-FOR-CMLE-JOB-SUBMISSION]
export \
SVC_ACCT=service-$CMLE_PROJ_NUM@cloud-ml.google.com.iam.gserviceaccount.com

gcloud projects add-iam-policy-binding $GCR_PROJ_ID \
    --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent

详细了解如何为 Container Registry 配置访问权限控制

提交训练作业

使用 gcloud CLI 将训练作业提交到 AI Platform Training。使用 --master-image-uri 标志将 URI 传递至 Docker 映像:

export BUCKET_NAME=custom_containers
export MODEL_DIR=example_model_$(date +%Y%m%d_%H%M%S)

gcloud ai-platform jobs submit training $JOB_NAME \
  --region $REGION \
  --master-image-uri $IMAGE_URI \
  -- \
  --model-dir=gs://$BUCKET_NAME/$MODEL_DIR \
  --epochs=10

使用自定义容器进行超参数调节

要使用自定义容器进行超参数调节,您需要进行以下调整:

请参阅使用超参数调节的自定义容器训练示例

将 GPU 与自定义容器搭配使用

如要使用 GPU 进行训练,您的自定义容器需要满足一些特殊要求。您必须构建与用于 CPU 训练的映像不同的 Docker 映像。

  • 在容器中预安装 CUDA 工具包和 cuDNN。推荐使用 nvidia/cuda 映像作为基础映像来执行此操作,因为它预安装了 CUDA 工具包和 cuDNN,并且有助于您正确设置相关的环境变量。

    如果您的训练配置使用 NVIDIA A100 GPU,则您的容器必须使用 CUDA 11 或更高版本。

  • 安装其他依赖项,比如 wgetcurlpip 以及训练应用所需的任何其他依赖项。

请参阅 GPU 训练的 Dockerfile 示例

将 TPU 与自定义容器搭配使用

如果您使用 TensorFlow 执行分布式训练,则可以在工作器虚拟机上使用 TPU。为此,您必须将训练作业配置为使用 TPU,并在提交训练作业时指定 tpuTfVersion 字段

使用自定义容器进行分布式训练

使用自定义容器运行分布式训练作业时,只能指定一个要用作主实例、工作器和参数服务器的映像。您还可以选择为主实例、工作器和参数服务器构建并指定不同的映像。在这种情况下,所有三个映像中的依赖项可能相同,并且您可以在每个映像中运行不同的代码逻辑。

您可以在代码中使用环境变量 TF_CONFIGCLUSTER_SPEC。这些环境变量描述了集群的整体结构,AI Platform Training 会在训练集群的每个节点中为您填充这些环境变量。详细了解 CLUSTER_SPEC

您可以在提交作业时在 TrainingInput 对象中指定映像,也可以通过 gcloud ai-platform submit training 中相应的标志来指定映像。

在此示例中,假设您已经定义了三个单独的 Dockerfile,每种类型的机器(主实例、工作器和参数服务器)都有一个 Dockerfile。之后,您命名、构建、测试映像并将映像推送到 Container Registry。最后,您提交一份训练作业,指定您的不同映像以及集群的机器配置。

在此之前,请先运行 gcloud auth configure-docker(如果您尚未这样做)。

export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
export BUCKET_NAME=custom_containers
export MASTER_IMAGE_REPO_NAME=master_image_name
export MASTER_IMAGE_TAG=master_tag
export MASTER_IMAGE_URI=gcr.io/$PROJECT_ID/$MASTER_IMAGE_REPO_NAME:$MASTER_IMAGE_TAG
export WORKER_IMAGE_REPO_NAME=worker_image_name
export WORKER_IMAGE_TAG=worker_tag
export WORKER_IMAGE_URI=gcr.io/$PROJECT_ID/$WORKER_IMAGE_REPO_NAME:$WORKER_IMAGE_TAG
export PS_IMAGE_REPO_NAME=ps_image_name
export PS_IMAGE_TAG=ps_tag
export PS_IMAGE_URI=gcr.io/$PROJECT_ID/$PS_IMAGE_REPO_NAME:$PS_IMAGE_TAG
export MODEL_DIR=distributed_example_$(date +%Y%m%d_%H%M%S)
export REGION=us-central1
export JOB_NAME=distributed_container_job_$(date +%Y%m%d_%H%M%S)

docker build -f Dockerfile-master -t $MASTER_IMAGE_URI ./
docker build -f Dockerfile-worker -t $WORKER_IMAGE_URI ./
docker build -f Dockerfile-ps -t $PS_IMAGE_URI ./

docker run $MASTER_IMAGE_URI --epochs 1
docker run $WORKER_IMAGE_URI --epochs 1
docker run $PS_IMAGE_URI --epochs 1

docker push $MASTER_IMAGE_URI
docker push $WORKER_IMAGE_URI
docker push $PS_IMAGE_URI

gcloud ai-platform jobs submit training $JOB_NAME \
  --region $REGION \
  --master-machine-type complex_model_m \
  --master-image-uri $MASTER_IMAGE_URI \
  --worker-machine-type complex_model_m \
  --worker-image-uri $WORKER_IMAGE_URI \
  --worker-count 9 \
  --parameter-server-machine-type large_model \
  --parameter-server-image-uri $PS_IMAGE_URI \
  --parameter-server-count 3 \
  -- \
  --model-dir=gs://$BUCKET_NAME/$MODEL_DIR \
  --epochs=10

自定义容器中的默认凭据

使用自定义容器运行训练作业时,应用默认以 Cloud ML Service Agent 身份运行。您可以在 Google Cloud 控制台中的 IAM 页面上找到项目的 Cloud ML Service Agent 的服务账号 ID。此 ID 的格式如下:

service-PROJECT_NUMBER@cloud-ml.google.com.iam.gserviceaccount.com

PROJECT_NUMBER 替换为您的 Google Cloud 项目的项目编号

如果您使用 Tensorflow tfds、Google Cloud 客户端库或其他使用应用默认凭据策略的工具,则 AI Platform Training 会自动使用 Cloud ML Servce Agent 凭据来设置身份验证和授权。

但是,如果您希望自定义容器作业以其他方式访问 Google Cloud,则可能需要执行其他配置。例如,如果您使用 gsutil 从 Cloud Storage 复制数据并使用 boto 库从配置文件加载凭据,则添加命令到您的 Dockerfile,以确保 gsutil 使用默认的 Cloud ML Servce Agent 凭据:

RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg

后续步骤