SageMaker HyperPod クラスターの耐障害性

次のコード例を使用してスクリプトを作成し、 として保存しますtrain_auto_resume.sh。このスクリプトは、置き換えられたノードに対して以前に手動設定が行われていないと仮定して、トレーニング環境のセットアップをデプロイします。これにより、環境がノードに依存しないため、ノードが置き換えられると、ジョブを再開する前に同じ環境がノードにプロビジョニングされます。

#!/bin/bash

# Filename: train_auto_resume.sh
# Sample containerized script to launch a training job with a single srun which can be auto-resumed.

# Place your training environment setup here. 
# Example: Install conda, docker, activate virtual env, etc.

# Get the list of nodes for a given job
NODE_LIST=$(scontrol show jobid=$SLURM_JOBID | \ # Show details of the SLURM job
            awk -F= '/NodeList=/{print $2}' | \  # Extract NodeList field
            grep -v Exc)                         # Exclude nodes marked as excluded

# Determine the master node from the node list
MASTER_NODE=$(scontrol show hostname $NODE_LIST | \ # Convert node list to hostnames
              head -n 1)                            # Select the first hostname as master node

# Get the master node address
MASTER_ADDR=$(scontrol show node=$MASTER_NODE | \ # Show node information
              awk -F= '/NodeAddr=/{print $2}' | \ # Extract NodeAddr
              awk '{print $1}')                   # Print the first part of NodeAddr


# Torchrun command to launch the training job
torchrun_cmd="torchrun --nnodes=$SLURM_NNODES \
                       --nproc_per_node=1 \
                       --node_rank=$SLURM_NODE \
                       --master-addr=$MASTER_ADDR \
                       --master_port=1234 \
                       <your_training_script.py>"

# Execute the torchrun command in the 'pytorch' Conda environment, 
# streaming output live
/opt/conda/bin/conda run --live-stream -n pytorch $torchrun_cmd

ハードウェアに障害が発生した場合にsrunコマンドを自動的に再試行する必要があることを示す フラグを追加して--auto-resume=1、 SageMaker HyperPod 自動再開を有効にしてジョブを起動します。

auto-resume を有効にしたsrunコマンドの例を次に示します。

スバッチの使用

環境を設定するためのロジックのほとんどが既に にあるためtrain_auto_resume.sh、バッチスクリプトは次のコード例のようにシンプルでなければなりません。次のバッチスクリプトが として保存されていると仮定しますbatch.sh。

#!/bin/bash
#SBATCH --nodes 2
#SBATCH --exclusive
srun --auto-resume=1 train_auto_resume.sh

次のコマンドを使用して、前述のバッチスクリプトを実行します。

sbatch batch.sh

salloc の使用

まず排他的な割り当てを取得し、 --auto-resumeフラグとエントリポイントスクリプトを使用して srun コマンドを実行します。

salloc -N 2 --exclusive
srun --auto-resume=1 train_auto_resume.sh