Creating users on a Slurm cluster on SageMaker HyperPod

Connect to the controller node and download the script. Make sure you run it with sudo permissions.

$ curl -O https://raw.githubusercontent.com/awslabs/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/utils/create_users.sh
chmod +x create_users.sh

Run the script using one of its three modes, which it selects by priority from the arguments you pass and the files in the current directory: command line arguments first, then a user file, then interactive prompts. Every mode runs the same sequence: create the users on the current node, set up SSH keypairs on the shared file system, create the users on all remaining nodes (compute nodes in parallel), register them with Slurm accounting on the controller, append them to shared_users.txt, and offer to upload that file to Amazon S3.

The three modes are described in the following examples, each adding two users (user1 and user2).

Command line mode

Pass the user names as arguments. The script auto-assigns UIDs and prompts whether the new users should be sudoers. This is the quickest way to add users.

$ sudo ./create_users.sh user1 user2

The script validates the user names, then prompts for sudoer access:

========================================
 Step 1: User Configuration
========================================
  Users from CLI args: user1 user2
    ✓ user1 — will be created (auto-assign UID)
    ✓ user2 — will be created (auto-assign UID)
  Make these user(s) sudoer(s)? (y/N): n

It then creates the users on the current node, sets up SSH keypairs, creates them on the compute nodes, registers them with Slurm accounting on the controller, and writes shared_users.txt, prompting you at the end to upload that file to Amazon S3.

File mode

Create a user file in the same directory as the script, then run the script with no arguments. The script reads the user definitions, validates them (creating only users that don't already exist and skipping UID conflicts), and creates the users across all nodes. The script accepts two file formats.

For shared_users.txt, use a CSV file with one user per line in the format username,uid,home_directory:

user1,2001,/fsx/user1
user2,2002,/fsx/user2

Alternatively, for shared_users.yaml, use a YAML file with either a simple list of users or users organized into groups. Parsing this format requires PyYAML on the node. The following is the simple list format:

users:
  - username: user1
    uid: 2001
  - username: user2
    uid: 2002

If both files are present, shared_users.txt takes precedence. Run the script with no arguments:

$ sudo ./create_users.sh

The script validates the file entries, confirms creation, prompts for sudoer access, and offers to add users that are not in the file. It then creates user1 and user2 on all nodes with the UIDs from the file (2001 and 2002), sets up SSH keypairs, registers them with Slurm accounting, and updates shared_users.txt.

Interactive mode

Run the script with no arguments and no user file present. The script prompts you for the user names and, optionally, the UIDs to assign:

$ sudo ./create_users.sh

========================================
 Step 1: User Configuration
========================================
  No shared_users.txt or shared_users.yaml found.
  Entering interactive mode...

  Enter username(s), comma-separated (e.g. 'sean' or 'sean,alice,bob'): user1,user2
  Specify UIDs? (Enter for auto-assign, or comma-separated UIDs): 2001,2002
  Make these user(s) sudoer(s)? (y/N): n

After you answer the prompts, the script creates user1 and user2 on all nodes with the UIDs you entered, sets up SSH keypairs, registers them with Slurm accounting, and updates shared_users.txt.

Test a user by switching to it and running a command across the cluster.

$ sudo su - user1 && ssh $(srun hostname)

Copy the appropriate sample file in the add-users directory and edit it with your users:

$ cp shared_users_sample.yaml shared_users.yaml

Upload the add-users directory to your lifecycle scripts Amazon S3 bucket (the bucket path must start with s3://sagemaker-):

$ aws s3 cp add-users/ s3://DOC-EXAMPLE-BUCKET/add-users/ --recursive

Specify the extension in the LifeCycleConfig block of your CreateCluster request, using add_users.sh as the OnInitComplete script:

"LifeCycleConfig": {
    "OnInitComplete": "add_users.sh",
    "SourceS3Uri": "s3://DOC-EXAMPLE-BUCKET/add-users/"
}

Update your user file with the new users (keep the existing users in the file) and upload it to Amazon S3:

$ aws s3 cp add-users/shared_users.yaml s3://DOC-EXAMPLE-BUCKET/add-users/shared_users.yaml

Connect to the controller node and pull the scripts to the shared file system:

$ sudo mkdir -p /fsx/cluster-scripts/add-users
sudo aws s3 cp s3://DOC-EXAMPLE-BUCKET/add-users/ /fsx/cluster-scripts/add-users/ --recursive
sudo chmod +x /fsx/cluster-scripts/add-users/*.sh

Run the extension on the controller:

$ sudo bash /fsx/cluster-scripts/add-users/add_users.sh

Run the extension on the compute nodes with srun:

$ sudo srun --partition=partition-name bash /fsx/cluster-scripts/add-users/add_users.sh

Download the base-config lifecycle script set from the GitHub website. With this option, you upload the entire base-config folder as your lifecycle script set, not just add_users.sh.

In the base-config folder, create or edit a text file named shared_users.txt in the following format. The first column is the user name, the second column is the unique user ID, and the third column is the user directory in the Amazon FSx shared space.

username1,uid1,/fsx/username1
username2,uid2,/fsx/username2
...

add_users.sh reads shared_users.txt from the same folder, so the file must be inside base-config alongside the script.

Upload the entire base-config folder to the Amazon S3 bucket for your HyperPod lifecycle scripts. While cluster creation, cluster update, or cluster software update is in progress, the lifecycle script runner runs add_users.sh, which reads shared_users.txt and sets up the users and their home directories.