创建用于训练的自定义容器映像

使用自定义容器映像可为 Vertex AI 中的训练提供最大的灵活性。如需了解使用自定义容器映像与将 Python 训练应用与预构建容器搭配使用的区别，请参阅训练代码要求。

本指南逐步演示以下步骤：

  1. 创建自定义容器：
    1. 编写一个 Dockerfile，用于设置容器以使用 Vertex AI，并包含训练应用所需的依赖项。
    2. 在本地构建和运行 Docker 容器。
  2. 将容器映像推送到 Artifact Registry。

准备工作

要配置 Artifact Registry API 代码库并在开发环境中设置 Docker，请按照 Artifact Registry 的 Docker 快速入门中的说明操作。具体而言，请务必完成快速入门中的以下步骤：

  • 准备工作
  • 选择 shell
  • 创建 Docker 代码库
  • 配置身份验证

创建自定义容器映像

我们建议通过两个可能的工作流创建自定义容器映像：

  • 编写训练代码。然后，使用 gcloud CLI 的 local-run 命令，根据您的训练代码构建和测试自定义容器映像，而无需自行编写 Dockerfile。

    如果您不熟悉 Docker，则此工作流可能更直接。如果您按照此工作流操作，则可以跳过本部分的其余内容。

  • 编写训练代码。然后，编写 Dockerfile 并基于其构建容器映像。最后，在本地测试容器。

    此工作流提供了更大的灵活性，因为您可以根据需要自定义容器映像。

本部分的其余内容演示了后一种工作流的示例。

训练代码

您可以使用任何编程语言的依赖项编写训练代码。请确保您的代码符合训练代码要求。如果您打算使用超参数调节、GPU 或分布式训练，请务必阅读该文档的相应部分；这些部分介绍了将这些功能与自定义容器搭配使用的具体注意事项。

创建 Dockerfile

创建 Dockerfile 以指定构建容器映像所需的所有说明。

本部分介绍如何创建用于自定义训练的 Dockerfile 通用示例。如需详细了解如何创建容器映像，请参阅 Docker 文档的快速入门。

要与 Vertex AI 搭配使用，您的 Dockerfile 需要包含涵盖以下任务的命令：

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

您的 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/training-app.py dest/training-app.py 将训练应用的代码复制到映像中。根据您的训练应用的结构，要复制的内容可能包括多个文件。

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

  • model.py
  • task.py
  • data_utils.py
ENTRYPOINT ["exec", "file"] 设置入口点以调用要运行的训练代码。 开始自定义训练时，您可以通过指定 ContainerSpec 中 command 字段来替换此入口点。您还可以指定 ContainerSpec 中的 args 字段，为入口点提供其他参数（并替换容器映像的 CMD 指令（如有）） ["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"]

（可选）调整 TPU 虚拟机的 Dockerfile

如果要使用 TPU 虚拟机在 Vertex AI 上进行训练，则必须调整 Dockerfile 以安装特别构建的 tensorflow 和 libtpu 库版本。详细了解如何调整容器以用于 TPU 虚拟机。

构建容器映像

使用环境变量创建正确的映像 URI，然后构建 Docker 映像：

export PROJECT_ID=$(gcloud config list project --format "value(core.project)")
export REPO_NAME=REPOSITORY_NAME
export IMAGE_NAME=IMAGE_NAME
export IMAGE_TAG=IMAGE_TAG
export IMAGE_URI=us-central1-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/${IMAGE_NAME}:${IMAGE_TAG}

docker build -f Dockerfile -t ${IMAGE_URI} ./

在这些命令中，替换以下内容：

  • REPOSITORY_NAME：您在准备工作部分创建的 Artifact Registry 代码库的名称。
  • IMAGE_NAME：您为容器映像选择的名称。
  • IMAGE_TAG：您为此容器映像版本选择的标记。

详细了解 Artifact Registry 对容器映像命名的要求。

在本地运行容器（可选）

将容器映像作为容器在本地运行，以验证映像。与在 Vertex AI 上的运行计划相比，您可以在较小的数据集上运行训练代码，或者运行较少的迭代次数。例如，如果容器映像中的入口点脚本接受 --epochs 标志以控制运行周期数，您可以运行以下命令：

docker run ${IMAGE_URI} --epochs 1

将容器推送到 Artifact Registry

如果本地运行一切正常，您可以将容器推送到 Artifact Registry。

如果您还没有在开发环境中运行 gcloud auth configure-docker us-central1-docker.pkg.dev，请先运行该命令。然后运行以下命令：

docker push ${IMAGE_URI}

Artifact Registry 权限

如果您使用的 Artifact Registry 映像与您使用的 Vertex AI 都在同一 Google Cloud 项目中，则无需进一步配置权限。您可以立即创建使用您的容器映像的自定义训练作业。

但是，如果您已将容器映像推送到 Artifact Registry 并且其所在 Google Cloud 项目与计划使用 Vertex AI 的项目不同，则必须为 Vertex AI 项目的 Vertex AI Service Agent 授予从其他项目拉取映像的权限。详细了解 Vertex AI Service Agent 以及如何授予其权限。

Artifact Registry

如需了解如何为 Vertex AI Service Agent 授予 Artifact Registry 代码库的访问权限，请参阅有关授予特定于代码库的权限的 Artifact Registry 文档。

后续步骤