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

您可以使用 gcloud ai custom-jobs local-run 命令构建基于训练代码的 Docker 容器映像,并将该映像作为容器在您的本地计算机上运行。此功能具有多种优势,包括:

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

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

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

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

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

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

准备工作

  1. 设置 Vertex AI 开发环境。

  2. 安装 Docker Engine。

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

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

使用 local-run 命令

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

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME

请替换以下内容:

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

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

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

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

    如果您省略 --local-package-path 标志(以及占位符),则 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-uri 标志(以及占位符),则 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 映像中。例如,如果您指定 --local-package-path=/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 ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --requirements=REQUIREMENTS

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

指定其他本地 Python 依赖项

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

例如:

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  --extra-packages=LOCAL_DEPENDENCIES

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

添加其他文件

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

gcloud ai custom-jobs local-run \
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=WORKING_DIRECTORY \
  --script=SCRIPT_PATH \
  --output-image-uri=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 ai custom-jobs local-run \\
  --executor-image-uri=BASE_IMAGE_URI \
  --local-package-path=/hello-world \
  --script=/trainer/task.py \
  --output-image-uri=OUTPUT_IMAGE_NAME \
  -- \
  --learning_rate=0.1 \
  --input_data=gs://BUCKET/small-dataset/

使用 GPU 加快模型训练速度

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

如果本地计算机运行 Linux 并且具有 GPU,您还可以将 local-run 命令配置为在本地运行容器时使用 GPU。此配置不是必需的,但如果您想测试训练代码在使用 GPU 时的表现,则此配置会非常有用。执行以下操作:

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

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

    gcloud ai custom-jobs local-run \
      --executor-image-uri=BASE_IMAGE_URI \
      --local-package-path=WORKING_DIRECTORY \
      --script=SCRIPT_PATH \
      --output-image-uri=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 ai custom-jobs local-run \
      --executor-image-uri=BASE_IMAGE_URI \
      --local-package-path=WORKING_DIRECTORY \
      --script=SCRIPT_PATH \
      --output-image-uri=OUTPUT_IMAGE_NAME \
      --service-account-key-file=KEY_PATH
    

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

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

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

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

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

后续步骤