Add kwargs to idist.barrier

ignite/distributed/comp_models/base.py

@@ -1,6 +1,8 @@
+import warnings
 from abc import ABCMeta, abstractmethod
+from inspect import signature
 from numbers import Number
-from typing import Any, Callable, List, Optional, Union, cast
+from typing import Any, Callable, Dict, List, Optional, Union, cast
 
 import torch
 
@@ -275,8 +277,18 @@ def _do_all_gather(self, tensor: torch.Tensor) -> torch.Tensor:
     def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
         pass
 
+    def _check_signature(self, fn: Callable, **kwargs: Any) -> Dict[str, Any]:
+        try:
+            fn_signature = signature(fn)
+            fn_signature.bind(**kwargs)
+        except TypeError:
+            extra_params = kwargs.keys() - set(fn_signature.parameters)
+            warnings.warn(f"Extra params : {extra_params} will not be used by {self._backend}.)")
+            kwargs = {key: kwargs[key] for key in kwargs.keys() if key not in extra_params}
+        return kwargs
+
     @abstractmethod
-    def barrier(self) -> None:
+    def barrier(self, **kwargs: Any) -> None:
         pass
 
 
@@ -358,5 +370,5 @@ def _do_all_gather(self, tensor: torch.Tensor) -> torch.Tensor:
     def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
         return tensor
 
-    def barrier(self) -> None:
+    def barrier(self, **kwargs: Any) -> None:
         pass

ignite/distributed/comp_models/horovod.py

@@ -194,7 +194,7 @@ def _do_all_gather(self, tensor: torch.Tensor) -> torch.Tensor:
         def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
             return hvd.broadcast(tensor, root_rank=src)
 
-        def barrier(self) -> None:
+        def barrier(self, **kwargs: Any) -> None:
+            kwargs = self._check_signature(fn=hvd.allreduce, **kwargs)
             # https://github.com/horovod/horovod/issues/159#issuecomment-424834603
-            # hvd.allreduce(torch.tensor(0, device=self.device()), name="barrier")
-            hvd.allreduce(torch.tensor(0, device="cpu"), name="barrier")
+            hvd.allreduce(**kwargs)

ignite/distributed/comp_models/native.py

@@ -432,8 +432,9 @@ def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
             dist.broadcast(tensor, src=src)
             return tensor
 
-        def barrier(self) -> None:
-            dist.barrier()
+        def barrier(self, **kwargs: Any) -> None:
+            kwargs = self._check_signature(fn=dist.barrier, **kwargs)
+            dist.barrier(**kwargs)
 
     def _expand_hostlist(nodelist: str) -> List[str]:
         """Expand a compressed hostlist string and returns all hosts listed.

ignite/distributed/comp_models/xla.py

@@ -160,5 +160,6 @@ def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
             xm.all_reduce("sum", [tensor,])
             return tensor
 
-        def barrier(self) -> None:
-            xm.rendezvous("barrier")
+        def barrier(self, **kwargs: Any) -> None:
+            kwargs = self._check_signature(fn=xm.rendezvous, **kwargs)
+            xm.rendezvous(**kwargs)

ignite/distributed/utils.py

@@ -421,13 +421,28 @@ def broadcast(
     return _model.broadcast(tensor, src=src, safe_mode=safe_mode)
 
 
-def barrier() -> None:
+def barrier(**kwargs: Any) -> None:
     """Helper method to synchronize all processes.
+
+    Args:
+        kwargs: acceptable kwargs according to provided backend:
+
+            - | "nccl" or "gloo" : ``group`` (default, GroupMember.WORLD), ``async_op`` (default, False),
+              | ``device_ids`` (default, None).
+
+            - | "horovod" : ``average`` (default, None), ``compression`` (default, Compression.none),
+              | ``op`` (default, None), ``prescale_factor`` (default, 1.0), ``postscale_factor`` (default, 1.0),
+              | ``process_set`` (default, global_process_set).
+
+            - | "xla-tpu" : ``payload`` (default, b""), ``replicas`` (default, []).
+
+    .. versionchanged:: 0.5.1
+        Method now accepts ``kwargs`` for all supported backends.
     """
     if _need_to_sync and isinstance(_model, _SerialModel):
         sync(temporary=True)
 
-    _model.barrier()
+    _model.barrier(**kwargs)
 
 
 def set_local_rank(index: int) -> None:

tests/ignite/conftest.py

@@ -259,14 +259,14 @@ def distributed_context_multi_node_nccl(multi_node_conf):
     _destroy_mnodes_dist_context()
 
 
-def _xla_template_worker_task(index, fn, args):
+def _xla_template_worker_task(index, fn, args, kwargs=None):
     import torch_xla.core.xla_model as xm
 
     xm.rendezvous("init")
-    fn(index, *args)
+    fn(index, *args, **kwargs)
 
 
-def _xla_execute(fn, args, nprocs):
+def _xla_execute(fn, args, nprocs, kwargs_dict=None):
 
     import torch_xla.distributed.xla_multiprocessing as xmp
 
@@ -275,7 +275,7 @@ def _xla_execute(fn, args, nprocs):
         spawn_kwargs["start_method"] = "fork"
 
     try:
-        xmp.spawn(_xla_template_worker_task, args=(fn, args), nprocs=nprocs, **spawn_kwargs)
+        xmp.spawn(_xla_template_worker_task, args=(fn, args, kwargs_dict), nprocs=nprocs, **spawn_kwargs)
     except SystemExit as ex_:
         assert ex_.code == 0, "Didn't successfully exit in XLA test"
 
@@ -294,19 +294,19 @@ def mock_gpu_is_not_available():
         yield mock_cuda
 
 
-def _hvd_task_with_init(func, args):
+def _hvd_task_with_init(func, args, kwargs):
     import horovod.torch as hvd
 
     hvd.init()
     lrank = hvd.local_rank()
     if torch.cuda.is_available():
         torch.cuda.set_device(lrank)
 
-    func(*args)
+    func(*args, **kwargs)
     hvd.shutdown()
 
 
-def _gloo_hvd_execute(func, args, np=1, do_init=False):
+def _gloo_hvd_execute(func, args, np=1, do_init=False, kwargs_dict=None):
     try:
         # old API
         from horovod.run.runner import run
@@ -317,7 +317,7 @@ def _gloo_hvd_execute(func, args, np=1, do_init=False):
     kwargs = dict(use_gloo=True, np=np)
 
     if do_init:
-        return run(_hvd_task_with_init, args=(func, args), **kwargs)
+        return run(_hvd_task_with_init, args=(func, args, kwargs_dict), **kwargs)
 
     return run(func, args=args, **kwargs)
 

tests/ignite/distributed/utils/__init__.py

@@ -215,14 +215,15 @@ def _test(data_src, data_others, safe_mode):
             idist.broadcast(None, src=0)
 
 
-def _test_distrib_barrier(device):
+def _test_distrib_barrier(device, **kwargs):
 
     t = torch.tensor([idist.get_rank()], device=device, dtype=torch.float)
     true_res = sum([i for i in range(idist.get_world_size())])
 
     if idist.get_rank() == 0:
         t += 10.0
-    idist.barrier()
+
+    idist.barrier(**kwargs)
 
     tt = idist.all_reduce(t)
     assert tt.item() == true_res + 10.0

tests/ignite/distributed/utils/test_horovod.py

@@ -184,6 +184,29 @@ def test_idist_barrier_hvd(gloo_hvd_executor):
     gloo_hvd_executor(_test_distrib_barrier, (device,), np=np, do_init=True)
 
 
+@pytest.mark.distributed
+@pytest.mark.skipif(not has_hvd_support, reason="Skip if no Horovod dist support")
+@pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
+def test_idist_barrier_kwargs_hvd(gloo_hvd_executor):
+    from horovod.torch.compression import Compression
+    from horovod.torch.mpi_ops import global_process_set
+
+    device = "cpu" if not torch.cuda.is_available() else "cuda"
+    np = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
+
+    kwargs_dict = dict(
+        tensor=torch.tensor(0, device="cpu"),
+        average=None,
+        name=None,
+        compression=Compression.none,
+        op=None,
+        prescale_factor=1.0,
+        postscale_factor=1.0,
+        process_set=global_process_set,
+    )
+    gloo_hvd_executor(_test_distrib_barrier, (device,), np=np, do_init=True, kwargs_dict=kwargs_dict)
+
+
 def _test_idist_methods_overhead(ok_factor, sync_model):
     import time
 

tests/ignite/distributed/utils/test_native.py

@@ -272,6 +272,41 @@ def test_idist_barrier_gloo(distributed_context_single_node_gloo):
     _test_distrib_barrier(device)
 
 
+@pytest.mark.distributed
+@pytest.mark.skipif(not has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
+def test_idist_barrier_kwargs_nccl(distributed_context_single_node_nccl):
+
+    device = idist.device()
+    from torch.distributed import GroupMember
+
+    kwargs_dict = {"group": GroupMember.WORLD, "async_op": False, "device_ids": None}
+    _test_distrib_barrier(device, **kwargs_dict)
+
+    kwargs_dict.update({"tag": "barrier", "payload": b"", "replicas": []})
+    with pytest.warns(
+        UserWarning, match=r"Extra params : \{((, )?('payload'|'replicas'|'tag')(, )?)+\} will not be used by nccl."
+    ):
+        _test_distrib_barrier(device, **kwargs_dict)
+
+
+@pytest.mark.distributed
+@pytest.mark.skipif(not has_native_dist_support, reason="Skip if no native dist support")
+def test_idist_barrier_kwargs_gloo(distributed_context_single_node_gloo):
+
+    device = idist.device()
+    from torch.distributed import GroupMember
+
+    kwargs_dict = {"group": GroupMember.WORLD, "async_op": False, "device_ids": None}
+    _test_distrib_barrier(device, **kwargs_dict)
+
+    kwargs_dict.update({"tag": "barrier", "payload": b"", "replicas": []})
+    with pytest.warns(
+        UserWarning, match=r"Extra params : \{((, )?('payload'|'replicas'|'tag')(, )?)+\} will not be used by gloo."
+    ):
+        _test_distrib_barrier(device, **kwargs_dict)
+
+
 def _test_idist_methods_overhead(ok_factor):
     import time
 

tests/ignite/distributed/utils/test_xla.py

@@ -184,9 +184,28 @@ def test_idist_barrier_xla():
     _test_distrib_barrier(device)
 
 
-def _test_idist_barrier_xla_in_child_proc(index):
+def _test_idist_barrier_xla_in_child_proc(index, kwargs_dict=None):
     device = idist.device()
-    _test_distrib_barrier(device)
+    _test_distrib_barrier(device, **kwargs_dict)
+
+
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS is in env vars")
+@pytest.mark.skipif(not has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_idist_barrier_kwargs_xla():
+
+    device = idist.device()
+    kwargs_dict = {"tag": "barrier", "payload": b"", "replicas": []}
+    _test_distrib_barrier(device, **kwargs_dict)
+
+    from torch.distributed import GroupMember
+
+    kwargs_dict.update({"group": GroupMember.WORLD, "async_op": False, "device_ids": None})
+    with pytest.warns(
+        UserWarning,
+        match=r"Extra params : \{((, )?('async_op'|'group'|'device_ids')(, )?)+\} will not be used by xla-tpu.",
+    ):
+        _test_distrib_barrier(device, **kwargs_dict)
 
 
 @pytest.mark.tpu
@@ -197,6 +216,15 @@ def test_idist_barrier_xla_in_child_proc(xmp_executor):
     xmp_executor(_test_idist_barrier_xla_in_child_proc, args=(), nprocs=n)
 
 
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" not in os.environ, reason="Skip if no NUM_TPU_WORKERS in env vars")
+@pytest.mark.skipif(not has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_idist_barrier_kwargs_xla_in_child_proc(xmp_executor):
+    n = int(os.environ["NUM_TPU_WORKERS"])
+    kwargs_dict = {"tag": "barrier", "payload": b"", "replicas": []}
+    xmp_executor(_test_idist_barrier_xla_in_child_proc, args=(), nprocs=n, kwargs_dict=kwargs_dict)
+
+
 @pytest.mark.tpu
 @pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS is in env vars")
 @pytest.mark.skipif(not has_xla_support, reason="Skip if no PyTorch XLA package")