Fix root node resolution in slurm environment

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Nicki Skafte <[email protected]>
Co-authored-by: chaton <[email protected]>

Author: 4 people

pytorch_lightning/cluster_environments/slurm_environment.py

@@ -26,7 +26,7 @@ def __init__(self):
     def master_address(self):
         # figure out the root node addr
         try:
-            root_node = os.environ["SLURM_NODELIST"].split(" ")[0]
+            root_node = os.environ["SLURM_NODELIST"].split(" ")[0].split(",")[0]
         except Exception:
             root_node = "127.0.0.1"
 

pytorch_lightning/trainer/connectors/slurm_connector.py

@@ -137,7 +137,7 @@ def connect_ddp(self, global_rank: int, world_size: int) -> None:
 
         # figure out the root node addr
         try:
-            root_node = os.environ["SLURM_NODELIST"].split(" ")[0]
+            root_node = os.environ["SLURM_NODELIST"].split(" ")[0].split(",")[0]
         except Exception:
             root_node = "127.0.0.1"