Added args/kwargs to idist barrier

ignite/distributed/comp_models/base.py

@@ -276,7 +276,7 @@ def _do_broadcast(self, tensor: torch.Tensor, src: int) -> torch.Tensor:
         pass
 
     @abstractmethod
-    def barrier(self) -> None:
+    def barrier(self, *args: Any, **kwargs: Any) -> None:
         pass
 
 
@@ -357,5 +357,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, *args: Any, **kwargs: Any) -> None:
         pass

ignite/distributed/comp_models/horovod.py

@@ -1,11 +1,13 @@
 import warnings
+from distutils.version import LooseVersion
 from typing import Any, Callable, Mapping, Optional, Tuple, cast
 
 import torch
 
 from ignite.distributed.comp_models.base import ComputationModel
 
 try:
+    import horovod
     import horovod.torch as hvd
 
     try:
@@ -23,6 +25,7 @@
 if has_hvd_support:
 
     HOROVOD = "horovod"
+    HOROVOD_VERSION = horovod.__version__
 
     class _HorovodDistModel(ComputationModel):
         """Private class for `Horovod <https://horovod.readthedocs.io/en/stable/>`_ distributed computation model."""
@@ -192,7 +195,15 @@ 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:
-            # 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")
+        def barrier(self, *args: Any, **kwargs: Any) -> None:
+            if LooseVersion(HOROVOD_VERSION) < LooseVersion("0.23.0"):
+                if len(args) or len(kwargs):
+                    warnings.warn(
+                        f"Arguments {list(args) + list(kwargs)} are not passed to horovod barrier method. "
+                        f"Please use horovod version>='0.23.0'"
+                    )
+                # 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")
+            else:
+                hvd.barrier(*args, **kwargs)

ignite/distributed/comp_models/native.py

@@ -431,8 +431,8 @@ 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, *args: Any, **kwargs: Any) -> None:
+            dist.barrier(*args, **kwargs)
 
     def _expand_hostlist(nodelist: str) -> List[str]:
         """Expand a compressed hostlist string and returns all hosts listed.

ignite/distributed/comp_models/xla.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Callable, Mapping, Optional, Tuple, cast
 
 import torch
@@ -159,5 +160,11 @@ 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, *args: Any, **kwargs: Any) -> None:
+            if not len(args) and "tag" not in kwargs:
+                warnings.warn(
+                    f"`tag` parameter is mandatory and is set by default to `barrier` for {self._backend} `rendezvous`"
+                    f" method."
+                )
+                kwargs["tag"] = "barrier"
+            xm.rendezvous(*args, **kwargs)

ignite/distributed/utils.py

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