Fix Distributed Initialization by Validating Legacy tcp:// and Updating Tests

ignite/distributed/comp_models/native.py

@@ -62,7 +62,8 @@ def create_from_backend(
             if dist.is_available() and dist.is_initialized():
                 raise RuntimeError("Can not create new distributed process group if default one is already initialized")
 
-            if init_method is None:
+            store = kwargs.get("store", None)
+            if init_method is None and store is None:
                 if world_size is not None or rank is not None:
                     raise ValueError("Arguments rank and world_size should be None if no init_method is provided")
             else:
@@ -93,7 +94,12 @@ def __init__(
             self._init_method: str | None = None
             if backend is not None:
                 self._create_from_backend(
-                    backend, timeout=timeout, init_method=init_method, world_size=world_size, rank=rank, **kwargs
+                    backend,
+                    timeout=timeout,
+                    init_method=init_method,
+                    world_size=world_size,
+                    rank=rank,
+                    **kwargs,
                 )
             else:
                 self._init_from_context()
@@ -107,6 +113,21 @@ def _create_from_backend(
             rank: int | None = None,
             **kwargs: Any,
         ) -> None:
+            store = kwargs.pop("store", None)
+            if init_method is not None and init_method.startswith("tcp://"):
+                raise ValueError(
+                    f"TCP initialization via init_method='{init_method}' will hang. "
+                    "To fix this, please configure MASTER_ADDR and MASTER_PORT in the environment and "
+                    "use 'env://' (or omit init_method). Alternatively, you can configure a TCPStore and "
+                    "pass it using the 'store' argument. For example:\n\n"
+                    "    import torch.distributed as dist\n"
+                    '    store = dist.TCPStore("<master_addr>", <master_port>, world_size, is_master)\n'
+                    "    # Then pass the store to idist.initialize or idist.Parallel:\n"
+                    "    # idist.initialize(backend=backend, store=store, ...)\n"
+                    "    # or\n"
+                    "    # with idist.Parallel(backend=backend, store=store, ...) as parallel:\n"
+                )
+
             if backend == dist.Backend.NCCL and not torch.cuda.is_available():
                 raise RuntimeError("Nccl backend is required but no cuda capable devices")
             self._backend = backend
@@ -116,15 +137,26 @@ def _create_from_backend(
             if timeout is not None:
                 init_pg_kwargs["timeout"] = timeout
 
-            if init_method is None:
+            if init_method is None and store is None:
                 init_method = "env://"
 
-            if "env" not in init_method:
-                init_pg_kwargs["world_size"] = int(os.environ["WORLD_SIZE"])
-                init_pg_kwargs["rank"] = int(os.environ["RANK"])
-            self._init_method = init_method
+            if init_method is not None:
+                if "env" not in init_method:
+                    init_pg_kwargs["world_size"] = int(os.environ["WORLD_SIZE"])
+                    init_pg_kwargs["rank"] = int(os.environ["RANK"])
+                self._init_method = init_method
+                dist.init_process_group(backend, init_method=init_method, **init_pg_kwargs)
+            else:
+                if rank is not None:
+                    init_pg_kwargs["rank"] = rank
+                else:
+                    init_pg_kwargs["rank"] = int(os.environ.get("RANK", 0))
+                if world_size is not None:
+                    init_pg_kwargs["world_size"] = world_size
+                else:
+                    init_pg_kwargs["world_size"] = int(os.environ.get("WORLD_SIZE", 1))
 
-            dist.init_process_group(backend, init_method=init_method, **init_pg_kwargs)
+                dist.init_process_group(backend, store=store, **init_pg_kwargs)
 
             if torch.cuda.is_available():
                 torch.cuda.set_device(self._local_rank)

ignite/distributed/launcher.py

@@ -132,11 +132,14 @@ def training(local_rank, config, **kwargs):
             with idist.Parallel(backend=backend, init_method='file:///d:/tmp/some_file', nproc_per_node=4) as parallel:
                 parallel.run(training, config, a=1, b=2)
 
-        Initializing the process using ``tcp://``
+        Initializing the process using ``store``
 
         .. code-block:: python
 
-            with idist.Parallel(backend=backend, init_method='tcp://10.1.1.20:23456', nproc_per_node=4) as parallel:
+            import torch.distributed as dist
+
+            store = dist.TCPStore("<master_addr>", <master_port>, world_size, is_master)
+            with idist.Parallel(backend=backend, store=store) as parallel:
                 parallel.run(training, config, a=1, b=2)
 
 
@@ -233,9 +236,22 @@ def __init__(
                 if value is not None:
                     raise ValueError(f"If backend is None, argument '{name}' should be also None, but given {value}")
 
+        if init_method is not None and init_method.startswith("tcp://"):
+            raise ValueError(
+                f"TCP initialization via init_method='{init_method}' will hang. "
+                "To fix this, please configure MASTER_ADDR and MASTER_PORT in the environment and "
+                "use 'env://' (or omit init_method). Alternatively, you can configure a TCPStore and "
+                "pass it using the 'store' argument. For example:\n\n"
+                "    import torch.distributed as dist\n"
+                '    store = dist.TCPStore("<master_addr>", <master_port>, world_size, is_master)\n'
+                "    # Then pass the store to idist.Parallel:\n"
+                "    # idist.Parallel(backend=backend, store=store, ...)"
+            )
+
         self.backend = backend
         self._spawn_params = None
         self.init_method = init_method
+        self.store = spawn_kwargs.pop("store", None)
 
         if self.backend is not None:
             if nproc_per_node is not None:
@@ -322,7 +338,7 @@ def training(local_rank, config, **kwargs):
 
     def __enter__(self) -> "Parallel":
         if self.backend is not None and self._spawn_params is None:
-            idist.initialize(self.backend, init_method=self.init_method)
+            idist.initialize(self.backend, init_method=self.init_method, store=self.store)
 
         # The logger can be setup from now since idist.initialize() has been called (if needed)
         self._logger = setup_logger(__name__ + "." + self.__class__.__name__)

ignite/distributed/utils.py

@@ -576,7 +576,7 @@ def initialize(backend: str, **kwargs: Any) -> None:
         kwargs: acceptable kwargs according to provided backend:
 
             - | "nccl" or "gloo" : ``timeout(=timedelta(minutes=30))``, ``init_method(=None)``,
-              | ``rank(=None)``, ``world_size(=None)``.
+              | ``rank(=None)``, ``world_size(=None)``, ``store(=None)``.
               | By default, ``init_method`` will be "env://". See more info about parameters: `torch_init`_.
 
             - | "horovod" : comm(=None), more info: `hvd_init`_.
@@ -604,6 +604,9 @@ def train_fn(local_rank, a, b, c):
             idist.initialize(backend)
             # or for torch native distributed on Windows:
             # idist.initialize("nccl", init_method="file://tmp/shared")
+            # or using store:
+            # store = dist.TCPStore("<master_addr>", <master_port>, world_size, is_master)
+            # idist.initialize("nccl", store=store)
             local_rank = idist.get_local_rank()
             train_fn(local_rank, a, b, c)
             idist.finalize()

tests/ignite/conftest.py

@@ -231,15 +231,19 @@ def _setup_free_port(local_rank):
 @pytest.fixture()
 def distributed_context_single_node_nccl(local_rank, world_size):
     free_port = _setup_free_port(local_rank)
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(free_port)
 
     dist_info = {
         "backend": "nccl",
         "world_size": world_size,
         "rank": local_rank,
-        "init_method": f"tcp://localhost:{free_port}",
+        "init_method": "env://",
     }
     yield _create_dist_context(dist_info, local_rank)
     _destroy_dist_context()
+    os.environ.pop("MASTER_ADDR", None)
+    os.environ.pop("MASTER_PORT", None)
 
 
 @pytest.fixture()
@@ -251,22 +255,32 @@ def distributed_context_single_node_gloo(local_rank, world_size):
         # can't use backslashes in f-strings
         backslash = "\\"
         init_method = f"file:///{temp_file.name.replace(backslash, '/')}"
+        dist_info = {
+            "backend": "gloo",
+            "world_size": world_size,
+            "rank": local_rank,
+            "init_method": init_method,
+            "timeout": timedelta(seconds=30),
+        }
+        yield _create_dist_context(dist_info, local_rank)
+        _destroy_dist_context()
+        temp_file.close()
     else:
         free_port = _setup_free_port(local_rank)
-        init_method = f"tcp://localhost:{free_port}"
-        temp_file = None
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(free_port)
 
-    dist_info = {
-        "backend": "gloo",
-        "world_size": world_size,
-        "rank": local_rank,
-        "init_method": init_method,
-        "timeout": timedelta(seconds=30),
-    }
-    yield _create_dist_context(dist_info, local_rank)
-    _destroy_dist_context()
-    if temp_file:
-        temp_file.close()
+        dist_info = {
+            "backend": "gloo",
+            "world_size": world_size,
+            "rank": local_rank,
+            "init_method": "env://",
+            "timeout": timedelta(seconds=30),
+        }
+        yield _create_dist_context(dist_info, local_rank)
+        _destroy_dist_context()
+        os.environ.pop("MASTER_ADDR", None)
+        os.environ.pop("MASTER_PORT", None)
 
 
 @pytest.fixture()
@@ -472,16 +486,21 @@ def distributed(request, local_rank, world_size):
             # can't use backslashes in f-strings
             backslash = "\\"
             init_method = f"file:///{temp_file.name.replace(backslash, '/')}"
+            dist_info = {
+                "world_size": world_size,
+                "rank": local_rank,
+                "init_method": init_method,
+            }
         else:
             temp_file = None
             free_port = _setup_free_port(local_rank)
-            init_method = f"tcp://localhost:{free_port}"
-
-        dist_info = {
-            "world_size": world_size,
-            "rank": local_rank,
-            "init_method": init_method,
-        }
+            os.environ["MASTER_ADDR"] = "localhost"
+            os.environ["MASTER_PORT"] = str(free_port)
+            dist_info = {
+                "world_size": world_size,
+                "rank": local_rank,
+                "init_method": "env://",
+            }
 
         if request.param == "nccl":
             dist_info["backend"] = "nccl"
@@ -494,6 +513,9 @@ def distributed(request, local_rank, world_size):
         _destroy_dist_context()
         if temp_file:
             temp_file.close()
+        else:
+            os.environ.pop("MASTER_ADDR", None)
+            os.environ.pop("MASTER_PORT", None)
 
     elif request.param == "horovod":
         request.node.stash[is_horovod_stash_key] = True

tests/ignite/distributed/comp_models/test_native.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 import pytest
 import torch
@@ -334,7 +335,8 @@ def _test__native_dist_model_create_from_context_set_local_rank(true_conf):
 def _test__native_dist_model_create_from_context_no_dist(true_backend, true_device):
     assert _NativeDistModel.create_from_context() is None
 
-    dist.init_process_group(true_backend, "tcp://0.0.0.0:2222", world_size=1, rank=0)
+    store = dist.TCPStore("0.0.0.0", 2222, world_size=1, is_master=True)
+    dist.init_process_group(true_backend, store=store, world_size=1, rank=0)
     dist.barrier()
 
     _test__native_dist_model_create_from_context_no_local_rank()
@@ -358,7 +360,9 @@ def _test__native_dist_model_create_from_context_no_dist(true_backend, true_devi
 def _test__native_dist_model_create_from_context_dist(local_rank, rank, world_size, true_backend, true_device):
     assert _NativeDistModel.create_from_context() is None
 
-    dist.init_process_group(true_backend, "tcp://0.0.0.0:2222", world_size=world_size, rank=rank)
+    is_master = rank == 0
+    store = dist.TCPStore("0.0.0.0", 2222, world_size=world_size, is_master=is_master)
+    dist.init_process_group(true_backend, store=store, world_size=world_size, rank=rank)
     dist.barrier()
     if torch.cuda.is_available():
         torch.cuda.set_device(local_rank)
@@ -397,7 +401,7 @@ def test__native_dist_model_create_no_dist_nccl(clean_env):
 
 
 @pytest.mark.distributed
-@pytest.mark.parametrize("init_method", [None, "tcp://0.0.0.0:22334", "FILE"])
+@pytest.mark.parametrize("init_method", [None, "FILE"])
 def test__native_dist_model_create_dist_gloo_1(init_method, get_fixed_dirname, local_rank, world_size):
     if init_method == "FILE":
         init_method = f"file://{get_fixed_dirname('native_dist_model_create_dist_gloo_1')}/shared"
@@ -418,7 +422,7 @@ def test__native_dist_model_create_dist_gloo_2(local_rank, world_size):
 
 @pytest.mark.distributed
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
-@pytest.mark.parametrize("init_method", [None, "tcp://0.0.0.0:22334", "FILE"])
+@pytest.mark.parametrize("init_method", [None, "FILE"])
 def test__native_dist_model_create_dist_nccl_1(init_method, get_fixed_dirname, local_rank, world_size):
     if init_method == "FILE":
         init_method = f"file://{get_fixed_dirname('native_dist_model_create_dist_nccl_1')}/shared"
@@ -444,7 +448,9 @@ def test__native_dist_model_create_dist_nccl_2(local_rank, world_size):
 def test__native_dist_model_warning_index_less_localrank(local_rank, world_size):
     assert _NativeDistModel.create_from_context() is None
 
-    dist.init_process_group("nccl", "tcp://0.0.0.0:2222", world_size=world_size, rank=local_rank)
+    is_master = local_rank == 0
+    store = dist.TCPStore("0.0.0.0", 2222, world_size=world_size, is_master=is_master)
+    dist.init_process_group("nccl", store=store, world_size=world_size, rank=local_rank)
     dist.barrier()
     # We deliberately incorrectly set cuda device to 0
     torch.cuda.set_device(0)
@@ -496,7 +502,7 @@ def _test__native_dist_model_spawn(backend, num_workers_per_machine, device, ini
 
 @pytest.mark.distributed
 @pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
-@pytest.mark.parametrize("init_method", [None, "CUSTOM_ADDR_PORT", "env://", "tcp://0.0.0.0:22334", "FILE"])
+@pytest.mark.parametrize("init_method", [None, "CUSTOM_ADDR_PORT", "env://", "FILE"])
 def test__native_dist_model_spawn_gloo(init_method, dirname):
     spawn_kwargs = {}
 
@@ -532,7 +538,7 @@ def test__native_dist_model_spawn_gloo(init_method, dirname):
 @pytest.mark.distributed
 @pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
-@pytest.mark.parametrize("init_method", [None, "CUSTOM_ADDR_PORT", "tcp://0.0.0.0:22334", "FILE"])
+@pytest.mark.parametrize("init_method", [None, "CUSTOM_ADDR_PORT", "FILE"])
 def test__native_dist_model_spawn_nccl(init_method, dirname):
     spawn_kwargs = {}
 
@@ -720,3 +726,30 @@ def test__setup_ddp_vars_from_slurm_env_bad_configs():
             "SLURM_JOB_ID": "12345",
         }
         _setup_ddp_vars_from_slurm_env(environ)
+
+
+def test__native_dist_model_tcp_init_method_error():
+    with pytest.raises(ValueError, match="will hang. To fix this, please configure MASTER_ADDR"):
+        _NativeDistModel.create_from_backend(backend="gloo", init_method="tcp://10.1.1.20:23456", rank=0, world_size=1)
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="Skip on Windows due to TCPStore libuv bugs")
+@pytest.mark.parametrize(
+    "backend,device",
+    [
+        ("gloo", torch.device("cpu")),
+        pytest.param(
+            "nccl",
+            torch.device("cuda:0"),
+            marks=pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU"),
+        ),
+    ],
+)
+def test__native_dist_model_store_init(clean_env, backend, device):
+    store = dist.TCPStore("0.0.0.0", 2222, world_size=1, is_master=True)
+
+    model = _NativeDistModel.create_from_backend(backend=backend, store=store, rank=0, world_size=1)
+    assert model.backend() == backend
+    assert model.get_world_size() == 1
+    assert model.get_rank() == 0
+    dist.destroy_process_group()