wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-11 12:31:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

moved slurm flag resolution to init

Browse Source
This commit is contained in:
William Falcon
2019-07-24 18:42:22 -04:00
parent 18ce3e5a23
commit f4d8fe5d77
1 changed files with 13 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+13 -11
View File
@@ -161,17 +161,6 @@ class Trainer(TrainerIO):
self.nb_tng_batches = None
self.nb_test_batches = None
# manages slurm task
# whenever we have the correct number of tasks, we let slurm manage processes
# otherwise we launch the required number of processes
self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
self.nb_slurm_tasks = 0
try:
self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
except Exception as e:
# likely not on slurm, so set the slurm managed flag to false
self.is_slurm_managing_tasks = False
# gpus come in as a string.
# if gpus = -1 then use all available devices
@@ -208,6 +197,19 @@ class Trainer(TrainerIO):
'To silence this warning set distributed_backend=ddp'
warnings.warn(w)
# extract SLURM flag vars
# whenever we have the correct number of tasks, we let slurm manage processes
# otherwise we launch the required number of processes
if self.use_ddp:
self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
self.nb_slurm_tasks = 0
try:
self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
except Exception as e:
# likely not on slurm, so set the slurm managed flag to false
self.is_slurm_managing_tasks = False
# process info
self.proc_rank = 0