本页面介绍了自定义容器如何支持 AI Platform Training 上的分布式训练结构。
借助自定义容器，您可以使用任何支持分布的机器学习框架进行分布式训练。虽然此处使用的术语是以 TensorFlow 的分布式模型为基础，但您可以使用其他任何具有类似分布式结构的机器学习框架。例如，MXNet 中的分布式训练使用调度器、工作器和服务器。这与 AI Platform Training 自定义容器的分布式训练结构相对应；后者使用主实例、工作器和参数服务器。
训练集群的结构
如果您使用 AI Platform Training 运行分布式训练作业，则需要在训练集群中指定多个机器（节点）。训练服务会为您指定的机器类型分配资源。给定节点上正在运行的作业称为副本。根据分布式 TensorFlow 模型，训练集群中的每个副本将在分布式训练中获得单个角色或任务：
主实例工作器：您只能将一个副本指定为主实例工作器（也称为主工作器）。这项任务可管理其他任务，并报告作业的整体状态。
工作器：您可以将一个或多个副本指定为工作器。这些副本执行您在作业配置中为它们指定的工作。
参数服务器：可以将一个或多个副本指定为参数服务器。这些副本会存储模型参数，并协调工作器之间的共享模型状态。
评估程序：您可以将一个或多个副本指定为评估器。这些副本可用于评估您的模型。如果您使用的是 TensorFlow，请注意 TensorFlow 通常要求您只能使用一个评估器。
API 映射
您可以为训练集群中的机器分配四种不同角色，这些角色对应于您可以在 TrainingInput（表示训练作业的输入参数）中指定的四个字段：
masterConfig.imageUri表示将要在主实例工作器上运行的容器映像 URI。workerConfig.imageUri、parameterServerConfig.imageUri和evaluatorConfig.imageUri分别表示将要在工作器、参数服务器和评估器上运行的容器映像 URI。如果没有为这些字段设置值，AI Platform Training 会使用masterConfig.imageUri的值。
您还可以在 gcloud ai-platform jobs submit training 中为每个字段设置相应标志的值：
- 对于主实例工作器配置，请使用
--master-image-uri。 - 对于工作器配置，请使用
--worker-image-uri。 - 对于参数服务器配置，请使用
--parameter-server-image-uri。 - 目前没有用于为评估器指定容器映像 URI 的标志。您可以在 config.yaml 配置文件中指定
evaluatorConfig.imageUri。
请参阅如何使用自定义容器提交分布式训练作业的示例。
了解 CLUSTER_SPEC
AI Platform Training 会针对每个副本填充一个 CLUSTER_SPEC 环境变量，以便描述整个集群的设置方式。与 TensorFlow 的 TF_CONFIG 类似，CLUSTER_SPEC 描述集群中的每个副本，包括其索引和角色（主实例工作器、工作器、参数服务器或评估器）。
当您使用 TensorFlow 运行分布式训练时，系统会解析 TF_CONFIG 以构建 tf.train.ClusterSpec。同样，当您使用其他机器学习框架运行分布式训练时，必须解析 CLUSTER_SPEC 以填充框架所需的任何环境变量或设置。
CLUSTER_SPEC 的格式
CLUSTER_SPEC 环境变量是一个 JSON 字符串，格式如下：
| 键 | 说明 | |
|---|---|---|
"cluster"
|
自定义容器的集群说明。与“TF_CONFIG”一样，此对象也会采用 TensorFlow 集群规范中的格式，且可以传递给 tf.train.ClusterSpec 的构造函数。 | |
"task"
|
描述代码运行所在特定节点的任务。您可以使用此信息为分布式作业中的特定工作器编写代码。此条目是一个含有以下键的字典： | |
"type"
|
此节点执行的任务类型。可能的值包括 master、worker、ps 和 evaluator。 |
|
"index"
|
任务的索引（从零开始）。大多数分布式训练作业都有一个主实例任务、一个或多个参数服务器，以及一个或多个工作器。 | |
"trial"
|
当前运行的超参数调节试验的标识符。为作业配置超参数调节时，可以设置多个要训练的试验。此值为您提供了一种区分代码中正在运行的试验的方法。标识符是一个包含试验编号的字符串值，从 1 开始。 | |
"job"
|
启动作业时使用的作业参数。在大多数情况下，您可以忽略此条目，因为其中包含的数据与通过命令行参数传递给应用的数据完全一致。 | |
与 TF_CONFIG 比较
请注意，AI Platform Training 还会对所有训练作业的每个副本设置 TF_CONFIG 环境变量。AI Platform Training 仅会对自定义容器训练作业的副本设置 CLUSTER_SPEC 环境变量。这两个环境变量共用一些值，但格式不同。
使用自定义容器训练时，默认情况下主实例副本在 TF_CONFIG 环境变量中标有 master 任务名称。您可以将主实例副本配置为标有 chief 任务名称，方法是在创建训练作业时将 trainingInput.useChiefInTfConfig 字段设置为 true，或者在作业中使用一个或多个评估器副本。这在您的自定义容器使用 TensorFlow 2 的情况下尤为有用。
除了此配置选项之外，在您使用自定义容器时与在您使用 AI Platform Training 运行时版本时，使用 TensorFlow 进行分布式训练的方法是相同的。查看有关如何在 AI Platform Training 上使用 TF_CONFIG 进行分布式训练的详细信息和示例。