Repository: pytorch/ignite

Issue:
Alternative to the command scontrol
## 🚀 Feature

Using `idist` on a slurm cluster relies on the command `scontrol` for selecting an hostname as `MASTER_ADDR`.

https://github.com/pytorch/ignite/blob/86cec1165441a945c93c938c23544ac52a17a170/ignite/distributed/comp_models/native.py#L258

As discussed in #2048, we should provide another way whether the command missing, for instance, using container launcher as `pyxis`.


Discussion:
@sdesrozis any ideas how quickly we could add a support for that ?
Sources
* https://sphinx-python.readthedocs.io/en/latest/source/hostlist.html
* https://github.com/LLNL/py-hostlist 

A relevant base could be [here](https://github.com/LLNL/py-hostlist/blob/master/hostlist/hostlist.py)


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
@@ -1,4 +1,5 @@
 import os
+import re
 import subprocess
 import warnings
 from distutils.version import LooseVersion
@@ -254,12 +255,12 @@ def _setup_env_in_slurm(self) -> None:
             slurm_port = slurm_port[-4:]
             os.environ["MASTER_PORT"] = str(int(slurm_port) + 15000)
             try:
-                # master address is the first hostname of nodes list
+                # use scontrol to expand hostname list
                 hostnames = subprocess.check_output(["scontrol", "show", "hostnames", os.environ["SLURM_JOB_NODELIST"]])
-            except FileNotFoundError as e:
-                # restore the environment before raising the exception
-                self._restore_env()
-                raise e
+            except FileNotFoundError:
+                # expand hostname list as scontrol
+                hostnames = " ".join(_expand_hostlist(os.environ["SLURM_JOB_NODELIST"])).encode("utf-8")
+            # master address is the first hostname of nodes list
             os.environ["MASTER_ADDR"] = hostnames.split()[0].decode("utf-8")
 
         def get_local_rank(self) -> int:
@@ -427,3 +428,69 @@ def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
 
         def barrier(self) -> None:
             dist.barrier()
+
+    def _expand_hostlist(nodelist: str) -> List[str]:
+        """Expand a compressed hostlist string and returns all hosts listed.
+
+        Source : https://github.com/LLNL/py-hostlist/blob/master/hostlist/hostlist.py
+
+        Args:
+            nodelist: Compressed hostlist string
+
+        .. versionadded:: 0.5.1
+        """
+        node_list = nodelist.split(", ")
+
+        result_hostlist = []
+        for node in node_list:
+            nodelist_match = r"(\w+-?)\[((,?[0-9]+-?,?-?){0,})\](.*)?"
+            if re.search(nodelist_match, node):
+                match = re.search(nodelist_match, node)
+
+                if match is None:
+                    raise ValueError(f"hostlist unvalid : {nodelist}")
+
+                # holds the ranges of nodes as a string
+                # now we can manipulate the string and cast it to a list of numbers
+                num = str(match.group(2)).replace("[", "").replace("]", "")
+
+                if len(num) == 0:
+                    raise ValueError(f"hostlist unvalid : {nodelist}")
+
+                num_list = num.split(",")
+
+                # find range of node numbers
+                ranges = [elem.split("-") for elem in num_list if "-" in elem]
+
+                # if the node numbers contain leading zeros, store them to be
+                if len(ranges):
+                    lead_zeros = max([len(s) - len(s.lstrip("0")) for s, _ in ranges])
+                else:
+                    lead_zeros = 0
+
+                # list of expanded ranges of node numbers
+                nodes_list = [list(range(int(s), int(e) + 1)) for s, e in ranges]
+
+                # add list of single node numbers
+                nodes_list += [[int(elem)] for elem in num_list if "-" not in elem]
+
+                # flat the list
+                final_list = [item for sublist in nodes_list for item in sublist]
+
+                # put final list in ascending order and append cluster name to each node number
+                final_list = list(set(sorted(final_list)))
+
+                # prepend leading zeros to numbers required
+                hostlist_tmp = [str(elem).zfill(lead_zeros + 1) for elem in final_list]
+
+                # append hostname to the node numbers
+                hostlist_no_suffix = [match.group(1) + elem for elem in hostlist_tmp]
+
+                # append suffix to hostlist if there is one
+                final_hostlist = [elem + match.group(4) for elem in hostlist_no_suffix]
+
+                result_hostlist += final_hostlist
+            else:
+                result_hostlist.append(node)
+
+        return result_hostlist


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
@@ -9,7 +9,45 @@
 if not has_native_dist_support:
     pytest.skip("Skip if no native dist support", allow_module_level=True)
 else:
-    from ignite.distributed.comp_models.native import _NativeDistModel
+    from ignite.distributed.comp_models.native import _expand_hostlist, _NativeDistModel
+
+
+# tests from https://github.com/LLNL/py-hostlist/blob/master/hostlist/unittest_hostlist.py
+@pytest.mark.parametrize(
+    "hostlist, expected",
+    [
+        ("localhost", "localhost"),
+        ("quartz[4-8]", "quartz4,quartz5,quartz6,quartz7,quartz8"),
+        (
+            "node[18-19,1-16,21-22]",
+            "node1,node2,node3,node4,node5,"
+            "node6,node7,node8,node9,node10,"
+            "node11,node12,node13,node14,node15,"
+            "node16,node18,node19,node21,node22",
+        ),
+        (
+            "node[4-8,12,16-20,22,24-26]",
+            "node4,node5,node6,node7,node8,"
+            "node12,node16,node17,node18,"
+            "node19,node20,node22,node24,"
+            "node25,node26",
+        ),
+        ("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-[009-11]vm1", "machine2-009vm1,machine2-010vm1,machine2-011vm1"),
+        ("node[1,2,3]", "node1,node2,node3"),
+    ],
+)
+def test_expand_hostlist(hostlist, expected):
+    assert _expand_hostlist(hostlist) == expected.split(",")
+
+
+def test_expand_hostlist_unvalid():
+    with pytest.raises(ValueError, match=r"hostlist unvalid"):
+        _expand_hostlist("unvalid[]")
 
 
 @pytest.mark.distributed
@@ -79,9 +117,6 @@ def test__native_dist_model_create_from_backend_bad_slurm_config():
     os.environ["SLURM_NTASKS"] = "1"
     os.environ["SLURM_JOB_NODELIST"] = "localhost"
 
-    with pytest.raises(FileNotFoundError, match=r"No such file or directory: 'scontrol'"):
-        _NativeDistModel.create_from_backend(backend="gloo", timeout=timedelta(seconds=10))
-
     os.environ["RANK"] = "1"
 
     with pytest.raises(RuntimeError, match=r"Defined env variables"):
@@ -177,6 +212,62 @@ def _test__native_dist_model_create_from_backend_dist(init_method, local_rank, r
     assert "RANK" not in os.environ
 
 
+def _test__native_dist_model_create_from_backend_slurm(local_rank, rank, world_size, backend, true_device):
+
+    import os
+    from datetime import timedelta
+
+    timeout = timedelta(seconds=20)
+
+    assert "MASTER_ADDR" not in os.environ
+    assert "MASTER_PORT" not in os.environ
+
+    del os.environ["WORLD_SIZE"]
+    del os.environ["LOCAL_RANK"]
+
+    os.environ["SLURM_JOB_ID"] = "15000"
+    os.environ["SLURM_PROCID"] = str(rank)
+    os.environ["SLURM_LOCALID"] = str(local_rank)
+    os.environ["SLURM_NTASKS"] = str(world_size)
+    os.environ["SLURM_JOB_NODELIST"] = "localhost"
+
+    model = _NativeDistModel.create_from_backend(backend=backend, timeout=timeout)
+
+    assert dist.is_available() and dist.is_initialized()
+    assert dist.get_backend() == backend
+
+    with pytest.raises(RuntimeError, match=r"Can not create new distributed process group if default one is"):
+        _NativeDistModel.create_from_backend(backend=backend, timeout=timeout)
+
+    _assert_model(
+        model,
+        {
+            "device": true_device,
+            "local_rank": local_rank,
+            "rank": rank,
+            "world_size": world_size,
+            "node_index": 0,
+            "nnodes": 1,
+            "nproc_per_node": world_size,
+        },
+    )
+
+    model.finalize()
+
+    del os.environ["SLURM_JOB_ID"]
+    del os.environ["SLURM_PROCID"]
+    del os.environ["SLURM_LOCALID"]
+    del os.environ["SLURM_NTASKS"]
+    del os.environ["SLURM_JOB_NODELIST"]
+
+    assert "MASTER_ADDR" not in os.environ
+    assert "MASTER_PORT" not in os.environ
+    assert "RANK" not in os.environ
+
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["LOCAL_RANK"] = str(local_rank)
+
+
 def _test__native_dist_model_create_from_context_no_local_rank():
 
     if "LOCAL_RANK" in os.environ:
@@ -302,12 +393,16 @@ def test__native_dist_model_create_dist_gloo_1(init_method, get_fixed_dirname, l
     device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
     _test__native_dist_model_create_from_backend_dist(init_method, local_rank, local_rank, world_size, "gloo", device)
 
+    if init_method is None:
+        _test__native_dist_model_create_from_backend_slurm(local_rank, local_rank, world_size, "gloo", device)
+
 
 @pytest.mark.distributed
 def test__native_dist_model_create_dist_gloo_2(local_rank, world_size):
 
     device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
     _test__native_dist_model_create_from_context_dist(local_rank, local_rank, world_size, "gloo", device)
+    _test__native_dist_model_create_from_backend_slurm(local_rank, local_rank, world_size, "gloo", device)
 
 
 @pytest.mark.distributed
@@ -321,6 +416,11 @@ def test__native_dist_model_create_dist_nccl_1(init_method, get_fixed_dirname, l
         init_method, local_rank, local_rank, world_size, "nccl", f"cuda:{local_rank}"
     )
 
+    if init_method is None:
+        _test__native_dist_model_create_from_backend_slurm(
+            local_rank, local_rank, world_size, "nccl", f"cuda:{local_rank}"
+        )
+
 
 @pytest.mark.distributed
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")