Coverage for melissa/utility/idr_torch.py: 0%

1#!/usr/bin/env python 

2# coding: utf-8 

3 

4import os 

5from melissa.scheduler.slurm_parser import break2str 

6 

7# get SLURM variables 

8rank = int(os.environ['SLURM_PROCID']) 

9local_rank = int(os.environ['SLURM_LOCALID']) 

10size = int(os.environ['SLURM_NTASKS']) 

11cpus_per_task = int(os.environ['SLURM_CPUS_PER_TASK']) 

12 

13# get node list from slurm 

14hostnames = break2str(os.environ['SLURM_JOB_NODELIST']) 

15 

16# get IDs of reserved GPU 

17gpu_ids = os.environ['SLURM_STEP_GPUS'].split(",") 

18 

19# define MASTER_ADD & MASTER_PORT 

20os.environ['MASTER_ADDR'] = hostnames[0] 

21# to avoid port conflict on the same node 

22os.environ['MASTER_PORT'] = str(12345 + int(min(gpu_ids)))