자동 노드 복구 및 자동 재개

다음 코드 예제를 사용하여 스크립트를 생성하고 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 명령의 예입니다.

sbatch 사용

환경을 설정하기 위한 대부분의 로직은 이미 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