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

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

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

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

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

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

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

    在本地运行代码有助于在对 Agent Platform 执行无服务器训练之前调试代码存在的问题。

准备工作

  1. 设置 Agent Platform 开发环境。

  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 要求的映像名称。(或者,您也可以稍后使用其他名称标记映像)。

    如果您省略 --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 映像最终部署到 Agent Platform,并使用 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 通过指定的服务账号凭据进行身份验证。

与在 Agent Platform 上执行训练进行比较

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

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

后续步骤