在本地容器化和运行训练代码

您可以使用 gcloud 命令行工具,根据您的训练代码构建 Docker 容器映像,并在本地计算机上将该映像作为容器运行。此功能具有多项优势:

  • 您只需最少地了解 Docker 即可构建容器映像。您无需编写自己的 Dockerfile。稍后,您可以将此映像推送到 Artifact Registry 并将其用于自定义容器训练

    对于高级使用场景,您可能仍需要编写自己的 Dockerfile

  • 您的容器映像可以运行 Python 训练应用或 Bash 脚本。

    您可以使用 Bash 脚本运行使用其他编程语言编写的训练代码(前提是您还指定了支持其他语言的基础容器映像)。

  • 在本地运行容器采用与在 Vertex AI 上运行容器类似的方式执行训练代码。

    在本地运行代码有助于在对 Vertex AI 执行自定义训练之前调试代码存在的问题。

本指南介绍如何使用 gcloud 工具通过 gcloud beta ai custom-jobs local-run 命令在本地容器化和运行代码。

准备工作

  1. 设置 Vertex AI 开发环境。

  2. 安装 Docker Engine。

  3. 如果您使用的是 Linux,请配置 Docker,以便您可以在不使用 sudo 的情况下运行它

    local-run 命令需要此配置才能使用 Docker。

使用 local-run 命令

运行以下命令,以根据您的训练代码构建容器映像并在本地运行容器:

gcloud beta ai custom-jobs local-run \
  --base-image=BASE_IMAGE_URI \
  --work-dir=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image=OUTPUT_IMAGE_NAME

请替换以下内容:

  • BASE_IMAGE_URI:要用作容器基础的 Docker 映像的 URI。选择包含训练代码所需的依赖项的基础映像。

    您可以将此 URI 用于预构建训练容器或其他任何对 DockerfileFROM 指令有效的值;例如,公开提供的 Docker 映像或 Artifact Registry 中有权访问的 Docker 映像。

  • WORKING_DIRECTORY:文件系统中最高级别的目录,其中包含需要用于训练的所有训练代码和本地依赖项。

    默认情况下,该命令仅会将 --script 标志指定的文件的父目录(请参阅以下列表项)复制到生成的 Docker 映像中。Docker 映像不一定包含 WORKING_DIRECTORY 中的所有文件;如需自定义要包含哪些文件,请参阅本文档中有关添加依赖项的部分。

    如果您省略 --work-dir 标志(以及占位符),则 local-run 命令会使用当前工作目录作为此值。

  • SCRIPT_PATH:相对于本地文件系统上的 WORKING_DIRECTORY 的路径,指向作为训练代码入口点的脚本。该脚本可以是 Python 脚本(以 .py 结尾)或 Bash 脚本。

    例如,如果要运行 /hello-world/trainer/task.pyWORKING_DIRECTORY/hello-world,请使用 trainer/task.py 作为此值。

    如果您指定 Python 脚本,则基础映像必须安装 Python;如果指定 bash 脚本,则基础映像必须安装 Bash。(所有预构建的训练容器以及许多其他公开提供的 Docker 映像都包含这两个依赖项。)

    使用 --python-module,而不是 --script

    如果您省略 --script 标志(和 SCRIPT_PATH),则必须使用 --python-module 标志在 WORKING_DIRECTORY 中指定要作为训练的入口点运行的 Python 模块的名称。例如,您可以指定 --python-module=trainer.task,而不是 --script=trainer/task.py

    在这种情况下,生成的 Docker 容器会将代码作为模块而不是作为脚本进行加载。如果您的入口点脚本在 WORKING_DIRECTORY 中导入其他 Python 模块,则建议您使用此选项。

  • OUTPUT_IMAGE_NAME:由此命令构建的生成的 Docker 映像的名称。您可以使用 docker build-t 标志接受的任何值。

    如果您打算稍后将映像推送到 Artifact Registry,则建议您使用符合 Artifact Registry 要求的映像名称。如果您打算稍后将映像推送到 Container Registry,则建议您使用符合 Container Registry 要求的映像名称。(或者,您也可以稍后使用其他名称标记映像)。

    如果您省略 --output-image 标志(以及占位符),则 local-run 命令会根据当前时间和入口点脚本的文件名,使用名称标记映像。

该命令会根据您的配置构建 Docker 容器映像。构建映像后,该命令会输出以下输出内容:

A training image is built.
Starting to run ...

然后,该命令会立即使用此容器映像在本地计算机上运行容器。容器退出时,该命令会输出以下输出内容:

A local run is finished successfully using custom image: OUTPUT_IMAGE_NAME

其他选项

以下部分介绍了可用于自定义 local-run 命令的行为的其他选项。

添加依赖项

您的训练代码可以依赖于基础映像上安装的任何依赖项(例如,预构建的训练容器映像包含许多适用于机器学习的 Python 库),以及在 local-run 命令创建的 Docker 映像中包含的任何文件。

使用 --script 标志或 --python-module 标志指定脚本时,该命令会将脚本的父目录(及其子目录)复制到 Docker 映像。例如,如果您指定 --work-dir=/hello-world--script=trainer/task.py,则该命令会将 /hello-world/trainer/ 复制到 Docker 映像中。

您还可以通过完成以下某个部分中介绍的额外步骤,添加其他 Python 依赖项或文件系统中的任意文件:

添加其他 Python 依赖项

您可以通过多种方式在 Docker 映像中添加其他 Python 依赖项:

使用 requirements.txt 文件

如果工作目录中有一个名为 requirements.txt 的文件,则 local-run 命令会将它视为 pip 需求文件,并使用该文件在 Docker 映像中安装 Python 依赖项。

使用 setup.py 文件

如果工作目录中有一个名为 setup.py 的文件,则 local-run 命令会将该文件视为 Python setup.py 文件,将该文件复制到 Docker 映像,并在该文件所在的 Docker 映像的目录中运行 pip install

例如,您可以将 install_requires 参数添加到 setup.py 中,以便在 Docker 映像中安装 Python 依赖项。

指定单个 PyPI 依赖项

您可以使用 --requirements 标志在 Docker 映像中安装特定 PyPI 依赖项。例如:

gcloud beta ai custom-jobs local-run \
  --base-image=BASE_IMAGE_URI \
  --work-dir=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image=OUTPUT_IMAGE_NAME \
  --requirements=REQUIREMENTS

REQUIREMENTS 替换为以英文逗号分隔的 Python 需求说明符列表。

指定其他本地 Python 依赖项

您可以使用 --extra-packages 标志安装特定的本地 Python 依赖项。这些 Python 依赖项必须位于工作目录下,并且每个依赖项都必须采用 pip install 支持的格式;例如,轮文件Python 源分发

例如:

gcloud beta ai custom-jobs local-run \
  --base-image=BASE_IMAGE_URI \
  --work-dir=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image=OUTPUT_IMAGE_NAME \
  --extra-packages=LOCAL_DEPENDENCIES

LOCAL_DEPENDENCIES 替换为以英文逗号分隔的本地文件路径列表(相对于工作目录表示)。

添加其他文件

如需将其他目录复制到 Docker 映像(而不是将其作为 Python 依赖项安装),您可以使用 --extra-dirs 标志。您只能指定工作目录下的目录。例如:

gcloud beta ai custom-jobs local-run \
  --base-image=BASE_IMAGE_URI \
  --work-dir=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image=OUTPUT_IMAGE_NAME \
  --extra-dirs=EXTRA_DIRECTORIES

EXTRA_DIRECTORIES 替换为以英文逗号分隔的本地目录列表(相对于工作目录表示)。

将参数传递给训练应用

如果训练应用的入口点脚本需要命令行参数,则您可以在运行 local-run 命令时指定这些参数。这些参数不会保存在 Docker 映像中;而是在映像作为容器运行时以参数形式传递。

如需将参数传递给入口点脚本,请将 -- 参数后接脚本的参数传递给 local-run 命令,放在该命令的所有其他标志之后。

例如,假设一个使用以下命令在本地运行的脚本:

python /hello-world/trainer/task.py \
  --learning_rate=0.1 \
  --input_data=gs://BUCKET/small-dataset/

使用 local-run 命令时,您可以使用以下标志在具有相同参数的容器中运行脚本:

gcloud beta ai custom-jobs local-run \\
  --base-image=BASE_IMAGE_URI \
  --work-dir=/hello-world \
  --script=/trainer/task.py \
  --output-image=OUTPUT_IMAGE_NAME \
  -- \
  --learning_rate=0.1 \
  --input_data=gs://BUCKET/small-dataset/

使用 GPU

如果您要将 local-run 命令创建的 Docker 映像最终部署到 Vertex AI,并使用 GPU 进行训练,请务必编写利用 GPU 的训练代码,并使用支持 GPU 的 Docker 映像作为 --base-image 标志的值。例如,您可以使用某个支持 GPU 的预构建训练容器映像。

如果本地计算机运行 Linux 并且具有 GPU,还可以配置 local-run 命令以在本地运行容器时使用 GPU。这是可选操作,但如果您想测试训练代码如何与 GPU 搭配使用,则此操作非常有用。执行以下操作:

  1. 在本地计算机上安装 NVIDIA 容器工具包 (nvidia-docker)(如果尚未安装)。

  2. 运行 local-run 命令时,请指定 --gpu 标志。例如:

    gcloud beta ai custom-jobs local-run \
      --base-image=BASE_IMAGE_URI \
      --work-dir=WORKING_DIRECTORY \
      --script=SCRIPT_PATH \
      --output-image=OUTPUT_IMAGE_NAME \
      --gpu
    

使用自定义服务帐号

默认情况下,当 local-run 命令在本地容器中运行训练代码时,它会通过应用默认凭据 (ADC) 将本地环境提供的 Google Cloud 凭据装载到该容器中,以便训练代码可以使用 ADC 通过相同的凭据进行身份验证。换句话说,当您运行 local-run 命令时,ADC 还会为您的代码提供在本地 shell 中 ADC 提供的凭据。

您可以使用 gcloud auth application-default login 命令将用户帐号用于 ADC,也可以在 shell 中设置环境变量以将服务帐号用于 ADC

如果您希望容器使用在本地 shell 中 ADC 提供的 Google Cloud 凭据之外的凭据运行,请执行以下操作:

  1. 创建或选择一个服务帐号,该帐号具有您希望训练代码有权访问的权限。

  2. 将此服务服务帐号的服务帐号密钥下载到本地计算机。

  3. 运行 local-run 命令时,请指定 --service-account-key-file 标志。例如:

    gcloud beta ai custom-jobs local-run \
      --base-image=BASE_IMAGE_URI \
      --work-dir=WORKING_DIRECTORY \
      --script=SCRIPT_PATH \
      --output-image=OUTPUT_IMAGE_NAME \
      --service-account-key-file=KEY_PATH
    

    KEY_PATH 替换为本地文件系统中的服务帐号密钥的路径。该路径必须是绝对路径,或相对于 shell 的当前工作目录的路径,而不是相对于 --work-dir 标志指定的目录的路径。

在生成的容器中,您的训练代码可以使用 ADC 通过指定的服务帐号凭据进行身份验证。

与在 Vertex AI 上执行训练进行比较

当您在 Vertex AI 上执行自定义训练时,Vertex AI 默认使用项目的 Vertex AI Custom Code Service Agent 来运行您的代码。您还可以关联其他服务帐号进行自定义训练。

使用 local-run 命令时,您无法以 Vertex AI Custom Code Service Agent 的身份进行身份验证,但可以创建具有类似权限的服务帐号并在本地使用。

后续步骤