Heterogeneous Cluster - a hello world training job

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

This basic example on how to run a Heterogeneous Clusters training job consisting of two instance groups. Each instance group includes a different instance type. Each instance prints its environment information including its instance group and exits.

You can retrieve environment information in either of the following ways:

• Option 1: Read instance group information using the convenient sagemaker_training.environment.Environment class.
• Option 2: Read instance group information from /opt/ml/input/config/resourceconfig.json.

Note: This notebook does not demonstrate offloading of data preprocessing job to data group and deep neural network training to dnn_group. We will cover those examples in TensorFlow's tf.data.service based Amazon SageMaker Heterogeneous Clusters for training and PyTorch and gRPC distributed dataloader based Amazon SageMaker Heterogeneous Clusters for training notebooks.

A. Setting up SageMaker Studio notebook

Before you start

Ensure you have selected Python 3 (TensorFlow 2.6 Python 3.8 CPU Optimized) image for your SageMaker Studio Notebook instance, and running on ml.t3.medium instance type.

Step 1 - Upgrade SageMaker SDK and dependent packages

Heterogeneous Clusters for Amazon SageMaker model training was announced on 07/08/2022. This feature release requires you to have updated SageMaker SDK and boto3 client libraries.

Step 2 - Restart the notebook kernel

Step 3 - Validate SageMaker Python SDK and TensorFlow versions

Ensure the output of the cell below reflects:

• SageMaker Python SDK version 2.98.0 or above,
• boto3 1.24 or above
• botocore 1.27 or above
• TensorFlow 2.6 or above

B. Run a heterogeneous cluster training job

Step 1: Set up training environment

Import the required libraries that enable you to use Heterogeneous clusters for training. In this step, you are also inheriting this notebook's IAM role and SageMaker session.

Step 2: Define instance groups

Here we define instance groups. Each instance group includes a different instance type.

Step 3: Review the "hello world" training code

Step 4: Configure the Estimator

In order to use SageMaker to fit our algorithm, we'll create an Estimator that defines how to use the container to train. This includes the configuration we need to invoke SageMaker training.

Step 5: Submit the training job

Here you are submitting the heterogeneous cluster training job.

Step 6: Review the logs for environment information

Wait for the training job to finish, and review its logs in the AWS Console (click on View logs from the Training Jobs node in Amazon SageMaker Console) You'll find two logs: Algo1, Algo2. Examine the printouts on each node on how to retrieve instance group environment information. An example is shown here:

Option-1: Read instance group information from the sagemaker_training.environment.Environment class
env.is_hetero: True
env.current_host: algo-1
env.current_instance_type: ml.c5.xlarge
env.current_instance_group: data_group
env.current_instance_group_hosts: ['algo-1']
env.instance_groups: ['data_group', 'dnn_group']

Option-2: Read instance group information from {file_path}.            You'll need to parse the json yourself. This doesn't require an additional library.
/opt/ml/input/config/resourceconfig.json dump = {
    "current_group_name": "data_group",
    "current_host": "algo-1",
    "current_instance_type": "ml.c5.xlarge",
    "hosts": [
        "algo-1",
        "algo-2"
    ],
    "instance_groups": [
        {
            "hosts": [
                "algo-1"
            ],
            "instance_group_name": "data_group",
            "instance_type": "ml.c5.xlarge"
        },
        {
            "hosts": [
                "algo-2"
            ],
            "instance_group_name": "dnn_group",
            "instance_type": "ml.m4.xlarge"
        }
    ],
    "network_interface_name": "eth0"
}
env.is_hetero: True
current_host=algo-1
current_instance_type=ml.c5.xlarge
env.current_instance_group: data_group
env.current_instance_group_hosts: TODO
env.instance_groups: TODO
env.instance_groups_dict: [{'instance_group_name': 'data_group', 'instance_type': 'ml.c5.xlarge', 'hosts': ['algo-1']}, {'instance_group_name': 'dnn_group', 'instance_type': 'ml.m4.xlarge', 'hosts': ['algo-2']}]
env.distribution_hosts: TODO
env.distribution_instance_groups: TODO

C. Next steps

In this notebook, we demonstrated how to retrieve the environment information, and differentiate which instance group an instance belongs to. Based on this, you can build logic to offload data processing tasks in your training job to a dedicated instance group. To understand how that can be done with a real-world example, we suggest going through the following notebook examples:

• TensorFlow's tf.data.service based Amazon SageMaker Heterogeneous Clusters for training
• PyTorch and gRPC distributed dataloader based Amazon SageMaker Heterogeneous Clusters for training

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.