Repository: pytorch/ignite

Issue:
SLURM and conflicting environment variables
## 🚀 Feature

Discussed here https://github.com/pytorch/ignite/issues/2048#issuecomment-921153418

Currently, if slurm environment variables are detected, an exception may be thrown if conflicting environment variables are defined (LOCAL_RANK, RANK, WORLD_SIZE). Typically, these variables are defined in `idist` using the slurm variables. This mutual exclusion is safe and avoid misuse of PyTorch launcher with slurm. 

https://github.com/pytorch/ignite/blob/922907217928fee7766046227506e54b28ae4283/ignite/distributed/comp_models/native.py#L219

However, some tools offered by nvidia (for docker/singularity, see https://github.com/NVIDIA/enroot/blob/v3.1.0/conf/hooks/extra/50-slurm-pytorch.sh) seem to do exactly the same things in the environment.

The idea is to verify the consistency of the environment (slurm wrt others), and raise an exception only if a conflict is detected.

https://github.com/pytorch/ignite/blob/922907217928fee7766046227506e54b28ae4283/ignite/distributed/comp_models/native.py#L250



Discussion:

Resolving patch:
diff --git a/ignite/distributed/comp_models/native.py b/ignite/distributed/comp_models/native.py
--- a/ignite/distributed/comp_models/native.py
+++ b/ignite/distributed/comp_models/native.py
@@ -431,21 +431,30 @@ def _expand_hostlist(nodelist: str) -> List[str]:
 
         Source : https://github.com/LLNL/py-hostlist/blob/master/hostlist/hostlist.py
 
+        .. note::
+            The host names can be composed by any character except the special ones `[`, `]`, `,`. Only one
+            sequence `[...]` is supported per hostname.
+
         Args:
             nodelist: Compressed hostlist string
 
         .. versionadded:: 0.4.6
         """
-        node_list = nodelist.split(", ")
-
         result_hostlist = []
-        for node in node_list:
-            nodelist_match = r"(.+)\[((,?[0-9]+-?,?-?){0,})\](.*)?"
 
-            match = re.search(nodelist_match, node)
+        nodelist_match = r"([^,\[\]]+\[[^\[\]]*\][^,\[\]]*|[^,\[\]]*),?"
+
+        nodelist = nodelist.replace(" ", "")
+
+        for node in re.findall(nodelist_match, nodelist):
+
+            node_match = r"(.+)\[((,?[0-9]+-?,?-?){0,})\](.*)?"
+
+            match = re.search(node_match, node)
 
             if match is None:
-                result_hostlist.append(node)
+                if node:
+                    result_hostlist.append(node)
             else:
                 # holds the ranges of nodes as a string
                 # now we can manipulate the string and cast it to a list of numbers
@@ -554,14 +563,21 @@ def _setup_ddp_vars_from_slurm_env(environ: Dict[str, str]) -> Dict[str, Union[s
             )
 
         if ddp_vars["MASTER_ADDR"] is None:
+            nodelist = environ["SLURM_JOB_NODELIST"]
             try:
                 # use scontrol to expand hostname list
-                hostnames = subprocess.check_output(["scontrol", "show", "hostnames", environ["SLURM_JOB_NODELIST"]])
+                hostnames = subprocess.check_output(["scontrol", "show", "hostnames", nodelist])
+                method = "scontrol"
             except FileNotFoundError:
                 # expand hostname list as scontrol
-                hostnames = " ".join(_expand_hostlist(environ["SLURM_JOB_NODELIST"])).encode("utf-8")
+                hostnames = " ".join(_expand_hostlist(nodelist)).encode("utf-8")
+                method = "ignite"
+            # at least one hostname should be defined
+            hostname_list = hostnames.split()
+            if len(hostname_list) < 1:
+                raise RuntimeError(f"No hostname detected in SLURM_JOB_NODELIST by {method} (nodelist={nodelist})")
             # master address is the first hostname of nodes list
-            ddp_vars["MASTER_ADDR"] = str(hostnames.split()[0].decode("utf-8"))
+            ddp_vars["MASTER_ADDR"] = str(hostname_list[0].decode("utf-8"))
 
         if ddp_vars["MASTER_PORT"] is None:
             # port should be the same over all process


Test patch:
diff --git a/tests/ignite/distributed/comp_models/test_native.py b/tests/ignite/distributed/comp_models/test_native.py
--- a/tests/ignite/distributed/comp_models/test_native.py
+++ b/tests/ignite/distributed/comp_models/test_native.py
@@ -17,9 +17,11 @@
     "hostlist, expected",
     [
         ("localhost", "localhost"),
+        ("compute!:b24_[1-2].r", "compute!:b24_1.r,compute!:b24_2.r"),
         ("quartz[4-8]", "quartz4,quartz5,quartz6,quartz7,quartz8"),
         ("c1001a-[11,17]", "c1001a-11,c1001a-17"),
         ("c1001a-s[11,17]", "c1001a-s11,c1001a-s17"),
+        ("c1009a-s17,c1010a-s11", "c1009a-s17,c1010a-s11"),
         (
             "gpu-compute-on-demand-dy-g4dnxlarge-[1-4]",
             "gpu-compute-on-demand-dy-g4dnxlarge-1,"
@@ -44,10 +46,17 @@
         ("machine2-[02-4]vm1", "machine2-02vm1,machine2-03vm1,machine2-04vm1"),
         (
             "machine2-[02-3]vm1, machine4-[0003-5].vml2",
-            "machine2-02vm1,machine2-03vm1," "machine4-0003.vml2," "machine4-0004.vml2," "machine4-0005.vml2",
+            "machine2-02vm1,machine2-03vm1,machine4-0003.vml2,machine4-0004.vml2,machine4-0005.vml2",
         ),
         ("machine2-[009-11]vm1", "machine2-009vm1,machine2-010vm1,machine2-011vm1"),
         ("node[1,2,3]", "node1,node2,node3"),
+        (
+            "compute-b24-[1-3,5-9], compute-b25-[1,4,8],compute-b25-[2-9,13]",
+            "compute-b24-1,compute-b24-2,compute-b24-3,compute-b24-5,compute-b24-6,"
+            "compute-b24-7,compute-b24-8,compute-b24-9,compute-b25-1,compute-b25-4,"
+            "compute-b25-8,compute-b25-2,compute-b25-3,compute-b25-4,compute-b25-5,"
+            "compute-b25-6,compute-b25-7,compute-b25-8,compute-b25-9,compute-b25-13",
+        ),
     ],
 )
 def test_expand_hostlist(hostlist, expected):
@@ -647,3 +656,14 @@ def test__setup_ddp_vars_from_slurm_env_bad_configs():
             "WORLD_SIZE": "2",
         }
         _setup_ddp_vars_from_slurm_env(environ)
+
+    with pytest.raises(RuntimeError, match=r"No hostname detected in SLURM_JOB_NODELIST by ignite"):
+        environ = {
+            "SLURM_PROCID": "1",
+            "SLURM_LOCALID": "1",
+            "SLURM_NTASKS": "4",
+            "SLURM_JOB_NUM_NODES": "1",
+            "SLURM_JOB_NODELIST": "[]",
+            "SLURM_JOB_ID": "12345",
+        }
+        _setup_ddp_vars_from_slurm_env(environ)