Start with base lifecycle scripts provided by HyperPod

provisioning_parameters.json – This is a Configuration form for provisioning Slurm nodes on HyperPod. The on_create.sh script finds this JSON file and defines environment variable for identifying the path to it. Through this JSON file, you can configure Slurm nodes and storage options such as Amazon FSx for Lustre for Slurm to communicate with. In provisioning_parameters.json, make sure that you assign the HyperPod cluster instance groups using the names you specified in create_cluster.json to the Slurm nodes appropriately based on how you plan to set them up.

The following diagram shows an example of how the two JSON configuration files create_cluster.json and provisioning_parameters.json should be written to assign HyperPod instance groups to Slurm nodes. In this example, we assume a case of setting up three Slurm nodes: controller (management) node, log-in node (which is optional), and compute (worker) node.

Figure: Direct comparison between create_cluster.json for HyperPod cluster creation and provisiong_params.json for Slurm configuration. The number of instance groups in create_cluster.json should match with the number of nodes you want to configure as Slurm nodes. In case of the example in the figure, three Slurm nodes will be configured on a HyperPod cluster of three instance groups. You should assign the HyperPod cluster instance groups to Slurm nodes by specifying the instance group names accordingly.

resource_config.json – During cluster creation, the lifecycle_script.py script is written to expect a resource_config.json file from HyperPod. This file contains information about the cluster, such as instance types and IP addresses.

When you run the CreateCluster API, HyperPod creates a resource configuration file at /opt/ml/config/resource_config.json based on the create_cluster.json file. The file path is saved to the environment variable named SAGEMAKER_RESOURCE_CONFIG_PATH.

{

    "ClusterConfig": {
        "ClusterArn": "arn:aws:sagemaker:us-west-2:111122223333:cluster/abcde01234yz",
        "ClusterName": "your-hyperpod-cluster"
    },
    "InstanceGroups": [
        {
            "Name": "controller-machine",
            "InstanceType": "ml.c5.xlarge",
            "Instances": [
                {
                    "InstanceName": "controller-machine-1",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                }
            ]
        },
        {
            "Name": "login-group",
            "InstanceType": "ml.m5.xlarge",
            "Instances": [
                {
                    "InstanceName": "login-group-1",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                }
            ]
        },
        {
            "Name": "compute-nodes",
            "InstanceType": "ml.trn1.32xlarge",
            "Instances": [
                {
                    "InstanceName": "compute-nodes-1",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                },
                {
                    "InstanceName": "compute-nodes-2",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                },
                {
                    "InstanceName": "compute-nodes-3",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                },
                {
                    "InstanceName": "compute-nodes-4",
                    "AgentIpAddress": "111.222.333.444",
                    "CustomerIpAddress": "111.222.333.444",
                    "InstanceId": "i-12345abcedfg67890"
                }
            ]
        }
    ]
}

lifecycle_script.py – This is the main Python script that collectively runs lifecycle scripts setting up Slurm on the HyperPod cluster while being provisioned. This script reads in provisioning_parameters.json and resource_config.json from the paths that are specified or identified in on_create.sh, passes the relevant information to each lifecycle script, and then runs the lifecycle scripts in order.

Lifecycle scripts are a set of scripts that you have a complete flexibility to customize to install software packages and set up necessary or custom configurations during cluster creation, such as setting up Slurm, creating users, installing Conda or Docker. The sample lifecycle_script.py script is prepared to run other base lifecycle scripts in the repository, such as launching Slurm deamons (start_slurm.sh), mounting Amazon FSx for Lustre (mount_fsx.sh), and setting up MariaDB accounting (setup_mariadb_accounting.sh) and RDS accounting (setup_rds_accounting.sh). You can also add more scripts, package them under the same directory, and add code lines to lifecycle_script.py to let HyperPod run the scripts. For more information about the base lifecycle scripts, see also 3.1 Lifecycle scripts in the Awsome Distributed Training GitHub repository.

In addition to the default setups, more scripts for installing the following software are available under the utils folder. The lifecycle_script.py file is already prepared to include code lines for running the installation scripts, so see the following items to search those lines and uncomment to activate them.