slurm2torch

#!/usr/bin/env bash

set -euo pipefail

signal_worker () {
    kill -"$1" "$trpid"
    wait "$trpid"
    exit $?  # use exit code of worker instead of default 128 + signal_number
}
trap 'signal_worker USR1' USR1  # to be used with e.g. --signal=SIGUSR1@60

# set env
export RANK="$SLURM_PROCID"
# export LOCAL_RANK="$SLURM_LOCALID"
export LOCAL_RANK=0  # slurm on this cluster already constrains CUDA_VISIBLE_DEVICES to a single GPU
export WORLD_SIZE="$SLURM_NTASKS"
export MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)"
export MASTER_PORT="$((10000 + 10#$(echo -n "$SLURM_JOBID" | tail -c 4)))"

# exec worker
"$@" &
trpid=$!
wait "$trpid"