SLURM 단위 (Job, Task) 이해

• JOB: sbatch 명령어로 실험 1개의 단위. jobId가 부여되는 하나의 단위임. QoS, Partition 자원같은게 Job단위로 일단 크게 구성됨.
• Task: Job내에서 SLURM이 띄우는 프로세스의 단위. 개수는 "#SABTACH--ntask=..."으로 결정. 직접 띄우는 OS프로세스임. 즉. "ps -u [사용자명]"으로 띄우는 테스크임. slurm이 관리하는 병렬 프로세스 수임. 즉, slurm이 띄운 프로세스가 자식프로세스로 여러개 병렬로 돌리는 것은 해당되지 않음.

 

아래의 경우는 테스크가 Job 하나에 Task가 한 개인 상태

#SBATCH --nodes=1
#SBATCH --ntasks=1              # SLURM task는 1개만
#SBATCH --gres=gpu:2            # GPU 각각 1대씩 DDP에 참여 (총2대)

python3 \
  torchrun ... \
  --nnodes=1 \
  --nproc_per_node=2 \
  --node_rank=0 \
  train.py

 

Job하나에 2개의 Task가 있는 경우도 있음. 예를 들어 DDP을 SLURM으로 쏘는 경우, 2개의 노드에서 2대의 GPU을 이용해서 SLURM을 쏜다고 가정하면 아래와 같음.

#SBATCH --nodes=2
#SBATCH --ntasks=2                 # 보통 노드당 1 task
#SBATCH --gpus-per-task=2

srun bash -c '
  torchrun \
    --nnodes=2 \
    --nproc_per_node=2 \
    --node_rank=$SLURM_NODEID \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    train.py
'

 

또 다른 예시: --nodes=1, --ntasks=4 인경우, "--ntasks"가 4이기 때문에, "python worker.py"가 4개가 동시에 돌아가게됨. 즉, 동일한 worker.py가 4번 병렬로 동시에 돌아감. 단, Task마다 환경변수가 다른데, 예를 들어, 4개는 SLURM_PROCID=0, SLURM_PROCID=1, SLURM_PROCID=2, SLURM_PROCID=3을 갖게되고, "SLURM_NTASKS"라는 환경변수는 4이게됨. 따라서, 아래의  파이썬에서 files라는 변수를 알아서 슬라이스해서 분배하게됨.

#SBATCH --ntasks=4

srun python worker.py

# worker.py
import os, glob
import time

rank = int(os.environ["SLURM_PROCID"])
world_size = int(os.environ["SLURM_NTASKS"])

files = sorted(glob.glob("/data/input/*.h5"))
chunk = files[rank::world_size]  # 슬라이스로 파일 분배

for f in chunk:
    print(f"[rank {rank}] processing {f}")
    time.sleep(1)

 

즉, 환경변수를 안쓴다면 완전 동일한 task가 여러번 실행되게 됨. 이런 코드라면 완전 동일하게 4개가 처음부터 100초 동안 자게됨.

import time
print("hello")
time.sleep(100)