本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。

# 修改训练脚本以分配实例组
<a name="train-heterogeneous-cluster-modify-training-script"></a>

使用前几节中的异构集群配置，您已经为 SageMaker 训练作业准备好了训练环境和实例。要进一步将实例组分配给某些训练和数据处理任务，下一步是修改训练脚本。默认情况下，无论实例大小如何，训练作业都只为所有节点创建训练脚本副本，这可能会导致性能下降。

例如，如果您在异构集群中混合 CPU 实例和 GPU 实例，同时将深度神经网络训练脚本传递给 SageMaker AI 估算器的`entry_point`参数，则该`entry_point`脚本将复制到每个实例。这意味着，如果没有适当的任务分配，CPU 实例还会运行整个脚本并启动专为 GPU 实例上的分布式训练设计的训练作业。因此，您必须更改要卸载并在 CPU 实例上运行的特定处理函数。您可以使用 SageMaker AI 环境变量来检索异构集群的信息，并让特定的进程相应地运行。

当您的训练作业开始时，您的训练脚本会读取包括异构集群配置在内的 SageMaker 训练环境信息。该配置包含诸如当前实例组、每个组中的当前主机以及当前主机所在的组之类的信息。

在 SageMaker AI 训练作业的初始化阶段，您可以通过以下方式查询实例组信息。

**（推荐）使用 SageMaker 训练工具包读取实例组信息**

使用[SageMaker 训练工具库](https://github.com/aws/sagemaker-training-toolkit)提供的环境 Python 模块。工具包库已预先安装在 TensorFlow 和的[SageMaker 框架容器](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only)中 PyTorch，因此在使用预构建的容器时，您无需执行额外的安装步骤。这是检索 SageMaker AI 环境变量的推荐方法，只需对训练脚本进行较少的代码更改。

```
from sagemaker_training import environment

env = environment.Environment()
```

与常规 SageMaker 训练和异构集群相关的环境变量：
+ `env.is_hetero` – 返回一个布尔结果，指示是否配置了异构集群。
+ `env.current_host` – 返回当前主机。
+ `env.current_instance_type` – 返回当前主机的实例的类型。
+ `env.current_instance_group` – 返回当前实例组的名称。
+ `env.current_instance_group_hosts` – 返回当前实例组中的主机列表。
+ `env.instance_groups` – 返回用于训练的实例组名称的列表。
+ `env.instance_groups_dict` – 返回训练作业的整个异构集群配置。
+ `env.distribution_instance_groups`— 返回分配给 SageMaker AI 估算器`distribution`类参数的实例组列表。
+ `env.distribution_hosts`— 返回属于分配给 SageMaker AI 估算器类`distribution`参数的实例组的主机列表。

例如，考虑以下包含两个实例组的异构集群示例。

```
from sagemaker.instance_group import InstanceGroup

instance_group_1 = InstanceGroup(
    "instance_group_1", "ml.c5.18xlarge", 1)
instance_group_2 = InstanceGroup(
    "instance_group_2", "ml.p3dn.24xlarge", 2)
```

示例异构集群的 `env.instance_groups_dict` 的输出应类似于以下内容。

```
{
    "instance_group_1": {
        "hosts": [
            "algo-2"
        ],
        "instance_group_name": "instance_group_1",
        "instance_type": "ml.c5.18xlarge"
    },
    "instance_group_2": {
        "hosts": [
            "algo-3",
            "algo-1"
        ],
        "instance_group_name": "instance_group_2",
        "instance_type": "ml.p3dn.24xlarge"
    }
}
```

**（可选）从资源配置 JSON 文件中读取实例组信息**

如果您更喜欢以 JSON 格式检索环境变量，则可以直接使用资源配置 JSON 文件。默认情况下， SageMaker 训练实例中的 JSON 文件`/opt/ml/input/config/resourceconfig.json`位于。

```
file_path = '/opt/ml/input/config/resourceconfig.json'
config = read_file_as_json(file_path)
print(json.dumps(config, indent=4, sort_keys=True))
```