release v0.2.6

pytorch_lightning/models/trainer.py

@@ -1,12 +1,12 @@
 """
 The trainer handles all the logic for running a val loop, training loop, distributing, etc...
 """
-from time import sleep
 import subprocess
 import traceback
 import warnings
 import os
 import pdb
+import re
 
 import torch
 from torch.utils.data.distributed import DistributedSampler
@@ -146,6 +146,7 @@ class Trainer(TrainerIO):
         self.print_nan_grads = print_nan_grads
         self.data_parallel_device_ids = None
         self.world_size = 1
+        self.node_rank = 0
         self.use_ddp = False
         self.use_dp = False
 
@@ -377,9 +378,14 @@ class Trainer(TrainerIO):
 
             # whenever we have the correct number of tasks, we let slurm manage processes
             # otherwise we launch the required number of processes
-            nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
-            nb_requested_gpus = len(self.data_parallel_device_ids)
-            is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
+            try:
+                nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+                is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
+            except Exception as e:
+                # likely not on slurm, so set the slurm managed flag to false
+                is_slurm_managing_tasks = False
+
             if is_slurm_managing_tasks:
                 task = int(os.environ['SLURM_LOCALID'])
                 self.ddp_train(task, model)
@@ -388,6 +394,7 @@ class Trainer(TrainerIO):
                 You requested {nb_requested_gpus} GPUs but launched {nb_slurm_tasks} slurm tasks. 
                 We will launch {nb_requested_gpus} processes for you. 
                 We recommend you let slurm manage the processes by setting: --ntasks-per-node={nb_requested_gpus}
+                If you're not using SLURM, ignore this message!
                 """
                 warnings.warn(msg)
                 mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
@@ -443,9 +450,10 @@ class Trainer(TrainerIO):
         # node rank using relative slurm id
         # otherwise default to node rank 0
         try:
-            node_rank = int(os.environ['SLURM_NODEID'])
-        except KeyError as e:
-            node_rank = 0
+            node_id = os.environ['SLURM_NODEID']
+            self.node_rank = int(node_id)
+        except Exception as e:
+            self.node_rank = 0
 
         # recover original exp before went into process
         # init in write mode only on proc 0
@@ -453,10 +461,10 @@ class Trainer(TrainerIO):
         self.experiment = self.experiment.get_non_ddp_exp()
 
         # show progbar only on prog_rank 0
-        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
+        self.prog_bar = self.prog_bar and self.node_rank == 0 and gpu_nb == 0
 
         # determine which process we are and world size
-        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
         self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
         # set up server using proc 0's ip address
@@ -501,15 +509,29 @@ class Trainer(TrainerIO):
             port = 12910
             os.environ['MASTER_PORT'] = f'{port}'
 
+        root_node = self.__resolve_root_node_address()
+        os.environ['MASTER_ADDR'] = root_node
+
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+
+    def __resolve_root_node_address(self):
         try:
             root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+
+            if '[' in root_node:
+                name = root_node.split('[')[0]
+                number = root_node.split(',')[0]
+                if '-' in number:
+                    number = number.split('-')[0]
+
+                number = re.sub('[^0-9]', '', number)
+                root_node = name + number
+
         except Exception as e:
             root_node = '127.0.0.2'
 
-        os.environ['MASTER_ADDR'] = root_node
-
-        sleep(self.proc_rank*0.5)
-        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+        return root_node
 
     def __run_pretrain_routine(self, model):
         """

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.2.5.1',
+    version='0.2.6',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",