torchrun_in_slurm.py

import os
import subprocess
import torch
import sys

local_rank = os.environ['SLURM_LOCALID']

nnodes = os.environ['SLURM_JOB_NUM_NODES']
node_id = os.environ['SLURM_NODEID']
node_list = os.environ['SLURM_JOB_NODELIST']
hostname = subprocess.getoutput(
    f'scontrol show hostname "{node_list}" | head -n1')
port = os.environ.get('PORT') or '27182'


def run_torchrun():
    cmd = [
        'torchrun',
        '--nnodes', nnodes,
        '--nproc_per_node', str(torch.cuda.device_count()),
        '--node_rank', node_id,
        '--master_addr', hostname,
        '--master_port', port,
        *sys.argv[1:]
    ]
    cmd_text = ' '.join(cmd)
    if local_rank == '0':
        curhost = subprocess.getoutput('hostname')
        print(curhost, cmd_text)
        subprocess.check_call(cmd)


if __name__ == '__main__':
    run_torchrun()