test_slurm_env_variables.py

import os
import torch

ENV_VARIABLES = [
    'SLURM_CPUS_ON_NODE',
    'SLURM_CPUS_PER_TASK',
    'SLURM_JOB_ID',
    'SLURM_JOB_NAME',
    'SLURM_JOB_NODELIST',
    'SLURM_JOB_NUM_NODES',
    'SLURM_LOCALID',
    'SLURM_NODEID',
    'SLURM_NTASKS',
    'SLURM_PROCID',
    'SLURM_SUBMIT_DIR',
    'SLURM_SUBMIT_HOST',
    'SLURM_TASKS_PER_NODE'
]

if __name__ == '__main__':
    env_dict = dict()
    for env in ENV_VARIABLES:
        env_dict[env] = os.environ.get(env, 'NOT FOUND')
    env_dict['torch.cuda.device_count'] = torch.cuda.device_count()
    print(env_dict)