Repository: pytorch/ignite

Issue:
Can't get `Parallel.run`  with NCCL in SLURM environment to work
## 🐛 Bug description
Dispatching a distributed multi-node/multi-gpu script via SLURM `sbatch` raises a `RuntimeError` 

## To reproduce:
#### Slurm invocation: 
```bash
 OMP_NUM_THREADS=1 sbatch -N1 -n2 -p gpu --gres=gpu:v100-32gb:02 --wrap "python -u test_dist.py run --nnodes=1 --nproc_per_node=2"
```
### `test_dist.py`
Python script
```python
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

import fire

import torch
import ignite
import ignite.distributed as idist

def run_diagnostic(local_rank):
    prefix = f"{local_rank}) "
    print(f"{prefix}Rank={idist.get_rank()}")
    print(f"{prefix}torch version: {torch.version.__version__}")
    print(f"{prefix}torch git version: {torch.version.git_version}")
    
    if torch.cuda.is_available():
        print(f"{prefix}torch version cuda: {torch.version.cuda}")
        print(f"{prefix}number of cuda devices: {torch.cuda.device_count()}")

        for i in range(torch.cuda.device_count()):
            print(f"{prefix}\t- device {i}: {torch.cuda.get_device_properties(i)}")
    else:
        print("{prefix}no cuda available")


    if "SLURM_JOBID" in os.environ:
        for k in ["SLURM_PROCID", "SLURM_LOCALID", "SLURM_NTASKS", "SLURM_JOB_NODELIST"]:
            print(f"{k}: {os.environ[k]}")
        
        if local_rank == 0:
            hostnames = subprocess.check_output(["scontrol", "show", "hostnames", os.environ["SLURM_JOB_NODELIST"]])
            print(f"hostnames: {hostnames}")


def run(**spawn_kwargs):
    with idist.Parallel(backend='nccl', **spawn_kwargs) as parallel:
        parallel.run(run_diagnostic)

if __name__ == '__main__':
    fire.Fire({'run': run})

```
### Error message (and logged output):
```
2021-04-25 19:06:35,565 ignite.distributed.launcher.Parallel INFO: Initialized distributed launcher with backend: 'nccl'
2021-04-25 19:06:35,566 ignite.distributed.launcher.Parallel INFO: - Parameters to spawn processes: 
	nproc_per_node: 2
	nnodes: 1
	node_rank: 0
2021-04-25 19:06:35,566 ignite.distributed.launcher.Parallel INFO: Spawn function '<function run_diagnostic at 0x1555554741e0>' in 2 processes
Traceback (most recent call last):
  File "test_dist.py", line 45, in <module>
    fire.Fire({'run': run})
  File "path_to_python_env/lib/python3.7/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "path_to_python_env/lib/python3.7/site-packages/fire/core.py", line 471, in _Fire
    target=component.__name__)
  File "path_to_python_env/lib/python3.7/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "test_dist.py", line 42, in run
    parallel.run(run_diagnostic)
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/launcher.py", line 309, in run
    idist.spawn(self.backend, func, args=args, kwargs_dict=kwargs, **self._spawn_params)
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/utils.py", line 324, in spawn
    fn, args=args, kwargs_dict=kwargs_dict, nproc_per_node=nproc_per_node, backend=backend, **kwargs
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/comp_models/native.py", line 380, in spawn
    **spawn_kwargs,
  File "path_to_python_env/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
    while not context.join():
  File "path_to_python_env/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 150, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "path_to_python_env/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/comp_models/native.py", line 323, in _dist_worker_task_fn
    backend, init_method=init_method, world_size=arg_world_size, rank=arg_rank, **kw
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/comp_models/native.py", line 72, in create_from_backend
    backend=backend, init_method=init_method, world_size=world_size, rank=rank, **kwargs
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/comp_models/native.py", line 93, in __init__
    backend, timeout=timeout, init_method=init_method, world_size=world_size, rank=rank, **kwargs
  File "path_to_python_env/lib/python3.7/site-packages/ignite/distributed/comp_models/native.py", line 124, in _create_from_backend
    dist.init_process_group(backend, init_method=init_method, **init_pg_kwargs)
  File "path_to_python_env/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 500, in init_process_group
    store, rank, world_size = next(rendezvous_iterator)
  File "path_to_python_env/lib/python3.7/site-packages/torch/distributed/rendezvous.py", line 190, in _env_rendezvous_handler
    store = TCPStore(master_addr, master_port, world_size, start_daemon, timeout)
RuntimeError: Address already in use

```
## Expected behavior
I would like the code to run `run_diagnostic` for each `local_rank` in the allocation. Right now it does not reach that point, it seems to complete parallel configuration on the master process but fails in the auxiliary processes.

## Environment

- PyTorch Version (e.g., 1.4): `1.8.1+cu102`
- Ignite Version: was using `0.4.4`, also tried on `0.5.0.dev20210423` 
- OS: Linux (CentOS)
- How you installed Ignite (`conda`, `pip`, source): `pip`
- Python version:  `3.7.3` 
- Any other relevant information:
1. I've checked that this isn't a problem of zombie processes lingering on the compute nodes from previous failed runs.
2. there is an `dist.barrier()` after `init_process_group()` in `ignite/distributed/comp_models/native.py` (`Parallel._create_from_backend`). Clearly the `init_process_group` completes for at least one process, but in syncing via `barrier()`, the other auxiliary processes fail.
3. Not sure if this is handled elsewhere (pretty sure not) or even relevant, but in `ignite.distributed.comp_models.native.Parallel.setup_env_vars`, upon discovering that the environment is a SLURM environment, the system makes a call to `self._setup_env_in_slurm` and returns without setting `self._local_rank`, `self._master_addr`, and `self._master_port` (as it would otherwise).
4. Thanks! 

Discussion:

Patch:
diff --git a/ignite/distributed/comp_models/native.py b/ignite/distributed/comp_models/native.py
index a6fd88b6..c723f418 100644
--- a/ignite/distributed/comp_models/native.py
+++ b/ignite/distributed/comp_models/native.py
@@ -85,6 +85,7 @@ if has_native_dist_support:
             """
             super(_NativeDistModel, self).__init__()
             self._env_backup = None  # type: Optional[Dict[str, str]]
+            self._local_rank = None  # type: Optional[int]
             self._master_port = None  # type: Optional[int]
             self._master_addr = None  # type: Optional[str]
             self._init_method = None  # type: Optional[str]
@@ -186,7 +187,7 @@ if has_native_dist_support:
 
         def _identify_local_rank(self) -> None:
 
-            if "SLURM_JOBID" in os.environ:
+            if "SLURM_JOB_ID" in os.environ:
                 os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"]
 
             if "LOCAL_RANK" in os.environ:
@@ -208,32 +209,40 @@ if has_native_dist_support:
 
             self._env_backup = os.environ.copy()
 
-            if "SLURM_JOBID" in os.environ:
+            # check whether all necessary env vars are set or not
+            env_vars = ["RANK", "LOCAL_RANK", "WORLD_SIZE"]
+            all_env_vars_defined = [k in os.environ for k in env_vars]
+
+            if "SLURM_JOB_ID" in os.environ:
+                if any(all_env_vars_defined):
+                    raise RuntimeError(
+                        f"Defined env variables '{env_vars}' should not be specified with SLURM. Typically, this "
+                        "happens when `torch.distributed.launch` or `torch.multiprocessing.spawn` are used. Please be "
+                        "sure to use the `srun` command instead."
+                    )
                 if rank is not None or world_size is not None:
                     raise ValueError("Arguments rank and world_size should not be specified with SLURM")
                 self._setup_env_in_slurm()
-                return
-
-            # check if all necessary env vars are set
-            # if partially defined raise an error
-            necessary_env_vars = ["RANK", "LOCAL_RANK", "WORLD_SIZE"]
-            all_env_vars_defined = [k in os.environ for k in necessary_env_vars]
-            if any(all_env_vars_defined) and not all(all_env_vars_defined):
-                raise RuntimeError(
-                    f"PyTorch distributed configuration should define env variables '{necessary_env_vars}'"
+            else:
+                # check if all necessary env vars are set
+                # if partially defined raise an error
+                if any(all_env_vars_defined) and not all(all_env_vars_defined):
+                    raise RuntimeError(f"PyTorch distributed configuration should define env variables '{env_vars}'")
+
+                os.environ["RANK"] = os.environ.get("RANK", f"{rank if rank is not None else 0}")
+                os.environ["WORLD_SIZE"] = os.environ.get(
+                    "WORLD_SIZE", f"{world_size if world_size is not None else 1}"
                 )
+                os.environ["LOCAL_RANK"] = os.environ.get("LOCAL_RANK", "0")
+                os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "15000")
+                os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "127.0.0.1")
 
-            os.environ["RANK"] = os.environ.get("RANK", f"{rank if rank is not None else 0}")
-            os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", f"{world_size if world_size is not None else 1}")
-            os.environ["LOCAL_RANK"] = os.environ.get("LOCAL_RANK", "0")
-            os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "15000")
-            os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "127.0.0.1")
             self._local_rank = int(os.environ["LOCAL_RANK"])
             self._master_addr = os.environ["MASTER_ADDR"]
             self._master_port = int(os.environ["MASTER_PORT"])
 
         def _setup_env_in_slurm(self) -> None:
-            for k in ["SLURM_PROCID", "SLURM_LOCALID", "SLURM_NTASKS", "SLURM_JOB_NODELIST"]:
+            for k in ["SLURM_JOB_ID", "SLURM_PROCID", "SLURM_LOCALID", "SLURM_NTASKS", "SLURM_JOB_NODELIST"]:
                 if k not in os.environ:
                     raise RuntimeError(f"SLURM distributed configuration is missing '{k}' in env variables")
 
@@ -244,12 +253,17 @@ if has_native_dist_support:
             slurm_port = os.environ["SLURM_JOB_ID"]
             slurm_port = slurm_port[-4:]
             os.environ["MASTER_PORT"] = str(int(slurm_port) + 15000)
-            # master address is the first hostname of nodes list
-            hostnames = subprocess.check_output(["scontrol", "show", "hostnames", os.environ["SLURM_JOB_NODELIST"]])
+            try:
+                # master address is the first hostname of nodes 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
             os.environ["MASTER_ADDR"] = hostnames.split()[0].decode("utf-8")
 
         def get_local_rank(self) -> int:
-            return self._local_rank
+            return cast(int, self._local_rank)
 
         def get_rank(self) -> int:
             return dist.get_rank()
@@ -282,6 +296,10 @@ if has_native_dist_support:
 
         def finalize(self) -> None:
             dist.destroy_process_group()
+            # restore backed-up env
+            self._restore_env()
+
+        def _restore_env(self) -> None:
             # restore backed-up env
             if self._env_backup is not None:
                 os.environ.clear()


Test patch:
diff --git a/tests/ignite/distributed/comp_models/test_native.py b/tests/ignite/distributed/comp_models/test_native.py
index 00b9df40..61499f35 100644
--- a/tests/ignite/distributed/comp_models/test_native.py
+++ b/tests/ignite/distributed/comp_models/test_native.py
@@ -58,6 +58,43 @@ def test__native_dist_model_create_from_backend_bad_config():
     del os.environ["RANK"]
 
 
+@pytest.mark.distributed
+@pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
+def test__native_dist_model_create_from_backend_bad_slurm_config():
+    import os
+    from datetime import timedelta
+
+    os.environ["SLURM_JOB_ID"] = "1"
+
+    with pytest.raises(RuntimeError, match=r"SLURM distributed configuration is missing"):
+        _NativeDistModel.create_from_backend(backend="gloo", timeout=timedelta(seconds=10))
+
+    with pytest.raises(ValueError, match=r"Arguments rank and world_size should not be specified with SLURM"):
+        _NativeDistModel.create_from_backend(
+            backend="gloo", timeout=timedelta(seconds=10), rank=1, init_method="", world_size=1
+        )
+
+    os.environ["SLURM_PROCID"] = "0"
+    os.environ["SLURM_LOCALID"] = "0"
+    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"):
+        _NativeDistModel.create_from_backend(backend="gloo", timeout=timedelta(seconds=10))
+
+    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"]
+    del os.environ["RANK"]
+
+
 def _assert_model(model, true_conf):
 
     assert model.device() == torch.device(true_conf["device"])