Repository: pytorch/ignite

Issue:
Add helper decorators to use event handler or method with specified rank 
## 🚀 Feature

Idea is to simplify the following code:
```python
import torch.distributed as dist

if dist.get_rank() == 0:
     
    @trainer.on(...)
    def some_handler(_):
        ...

    tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")        
```
into 

```python
from ignite.utils import one_rank_only

@trainer.on(...)
@one_rank_only   #  equivalent @one_rank_only(rank=0)
def some_handler(_):
    ...

# implicitly work only on rank 0
tb_logger = TensorboardLogger(log_dir="experiments/tb_logs") 
```



Discussion:

Resolving patch:
diff --git a/examples/contrib/mnist/mnist_with_neptune_logger.py b/examples/contrib/mnist/mnist_with_neptune_logger.py
--- a/examples/contrib/mnist/mnist_with_neptune_logger.py
+++ b/examples/contrib/mnist/mnist_with_neptune_logger.py
@@ -148,13 +148,18 @@ def compute_metrics(engine):
     npt_logger.attach(trainer, log_handler=GradsScalarHandler(model), event_name=Events.ITERATION_COMPLETED(every=100))
 
     def score_function(engine):
-        return engine.state.metrics['accuracy']
-
-    to_save = {'model': model}
-    handler = Checkpoint(to_save, NeptuneSaver(npt_logger), n_saved=2,
-                         filename_prefix='best', score_function=score_function,
-                         score_name="validation_accuracy",
-                         global_step_transform=global_step_from_engine(trainer))
+        return engine.state.metrics["accuracy"]
+
+    to_save = {"model": model}
+    handler = Checkpoint(
+        to_save,
+        NeptuneSaver(npt_logger),
+        n_saved=2,
+        filename_prefix="best",
+        score_function=score_function,
+        score_name="validation_accuracy",
+        global_step_transform=global_step_from_engine(trainer),
+    )
     validation_evaluator.add_event_handler(Events.COMPLETED, handler)
 
     # kick everything off
diff --git a/ignite/utils.py b/ignite/utils.py
--- a/ignite/utils.py
+++ b/ignite/utils.py
@@ -1,10 +1,12 @@
 import collections.abc as collections
 import logging
+from functools import wraps
 from typing import Union, Optional, Callable, Any, Type, Tuple
 
 import torch
+import torch.distributed as dist
 
-__all__ = ["convert_tensor", "apply_to_tensor", "apply_to_type", "to_onehot", "setup_logger"]
+__all__ = ["convert_tensor", "apply_to_tensor", "apply_to_type", "to_onehot", "setup_logger", "one_rank_only"]
 
 
 def convert_tensor(
@@ -125,3 +127,37 @@ def setup_logger(
         logger.addHandler(fh)
 
     return logger
+
+
+def one_rank_only(rank: int = 0, barrier: bool = False):
+    """Decorator to filter handlers wrt a rank number
+
+    Args:
+        rank (int): rank number of the handler (default: 0).
+        barrier (bool): synchronisation with a barrier (default: False).
+
+    .. code-block:: python
+        engine = ...
+
+        @engine.on(...)
+        @one_rank_only() # means @one_rank_only(rank=0)
+        def some_handler(_):
+            ...
+
+        @engine.on(...)
+        @one_rank_only(rank=1)
+        def some_handler(_):
+            ...
+    """
+    def _one_rank_only(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            ret = None
+            if dist.get_rank() == rank:
+                ret = func(*args, **kwargs)
+            if barrier:
+                dist.barrier()
+            return ret
+        return wrapper
+
+    return _one_rank_only


Test patch:
diff --git a/tests/ignite/contrib/metrics/test_precision_recall_curve.py b/tests/ignite/contrib/metrics/test_precision_recall_curve.py
--- a/tests/ignite/contrib/metrics/test_precision_recall_curve.py
+++ b/tests/ignite/contrib/metrics/test_precision_recall_curve.py
@@ -11,7 +11,7 @@ def test_precision_recall_curve():
     size = 100
     np_y_pred = np.random.rand(size, 1)
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     sk_precision, sk_recall, sk_thresholds = precision_recall_curve(np_y, np_y_pred)
 
     precision_recall_curve_metric = PrecisionRecallCurve()
@@ -32,7 +32,7 @@ def test_integration_precision_recall_curve_with_output_transform():
     size = 100
     np_y_pred = np.random.rand(size, 1)
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     np.random.shuffle(np_y)
 
     sk_precision, sk_recall, sk_thresholds = precision_recall_curve(np_y, np_y_pred)
@@ -41,8 +41,8 @@ def test_integration_precision_recall_curve_with_output_transform():
 
     def update_fn(engine, batch):
         idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx: idx + batch_size]
-        y_pred_batch = np_y_pred[idx: idx + batch_size]
+        y_true_batch = np_y[idx : idx + batch_size]
+        y_pred_batch = np_y_pred[idx : idx + batch_size]
         return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
 
     engine = Engine(update_fn)
@@ -65,7 +65,7 @@ def test_integration_precision_recall_curve_with_activated_output_transform():
     np_y_pred = np.random.rand(size, 1)
     np_y_pred_sigmoid = torch.sigmoid(torch.from_numpy(np_y_pred)).numpy()
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     np.random.shuffle(np_y)
 
     sk_precision, sk_recall, sk_thresholds = precision_recall_curve(np_y, np_y_pred_sigmoid)
@@ -74,8 +74,8 @@ def test_integration_precision_recall_curve_with_activated_output_transform():
 
     def update_fn(engine, batch):
         idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx: idx + batch_size]
-        y_pred_batch = np_y_pred[idx: idx + batch_size]
+        y_true_batch = np_y[idx : idx + batch_size]
+        y_pred_batch = np_y_pred[idx : idx + batch_size]
         return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
 
     engine = Engine(update_fn)
diff --git a/tests/ignite/contrib/metrics/test_roc_curve.py b/tests/ignite/contrib/metrics/test_roc_curve.py
--- a/tests/ignite/contrib/metrics/test_roc_curve.py
+++ b/tests/ignite/contrib/metrics/test_roc_curve.py
@@ -11,7 +11,7 @@ def test_roc_curve():
     size = 100
     np_y_pred = np.random.rand(size, 1)
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     sk_fpr, sk_tpr, sk_thresholds = roc_curve(np_y, np_y_pred)
 
     roc_curve_metric = RocCurve()
@@ -32,7 +32,7 @@ def test_integration_roc_curve_with_output_transform():
     size = 100
     np_y_pred = np.random.rand(size, 1)
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     np.random.shuffle(np_y)
 
     sk_fpr, sk_tpr, sk_thresholds = roc_curve(np_y, np_y_pred)
@@ -41,8 +41,8 @@ def test_integration_roc_curve_with_output_transform():
 
     def update_fn(engine, batch):
         idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx: idx + batch_size]
-        y_pred_batch = np_y_pred[idx: idx + batch_size]
+        y_true_batch = np_y[idx : idx + batch_size]
+        y_pred_batch = np_y_pred[idx : idx + batch_size]
         return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
 
     engine = Engine(update_fn)
@@ -65,7 +65,7 @@ def test_integration_roc_curve_with_activated_output_transform():
     np_y_pred = np.random.rand(size, 1)
     np_y_pred_sigmoid = torch.sigmoid(torch.from_numpy(np_y_pred)).numpy()
     np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2:] = 1
+    np_y[size // 2 :] = 1
     np.random.shuffle(np_y)
 
     sk_fpr, sk_tpr, sk_thresholds = roc_curve(np_y, np_y_pred_sigmoid)
@@ -74,8 +74,8 @@ def test_integration_roc_curve_with_activated_output_transform():
 
     def update_fn(engine, batch):
         idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx: idx + batch_size]
-        y_pred_batch = np_y_pred[idx: idx + batch_size]
+        y_true_batch = np_y[idx : idx + batch_size]
+        y_pred_batch = np_y_pred[idx : idx + batch_size]
         return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
 
     engine = Engine(update_fn)
diff --git a/tests/ignite/metrics/test_metric.py b/tests/ignite/metrics/test_metric.py
--- a/tests/ignite/metrics/test_metric.py
+++ b/tests/ignite/metrics/test_metric.py
@@ -342,9 +342,9 @@ def process_function(*args, **kwargs):
     m2.detach(engine)
     engine.run(range(10), 5)
 
-    assert 'm1' not in engine.state.metrics
-    assert 'm2_1' not in engine.state.metrics
-    assert 'm2_2' not in engine.state.metrics
+    assert "m1" not in engine.state.metrics
+    assert "m2_1" not in engine.state.metrics
+    assert "m2_2" not in engine.state.metrics
 
     assert not m1.is_attached(engine)
     assert not m2.is_attached(engine)
diff --git a/tests/ignite/test_utils.py b/tests/ignite/test_utils.py
--- a/tests/ignite/test_utils.py
+++ b/tests/ignite/test_utils.py
@@ -2,9 +2,11 @@
 import logging
 import pytest
 import torch
+import torch.distributed as dist
 
 from collections import namedtuple
-from ignite.utils import convert_tensor, to_onehot, setup_logger
+from ignite.utils import convert_tensor, to_onehot, setup_logger, one_rank_only
+from ignite.engine import Engine, Events
 
 
 def test_convert_tensor():
@@ -113,3 +115,86 @@ def _(_):
     for source in [err, data]:
         assert "trainer INFO: Engine run starting with max_epochs=5." in source[0]
         assert "evaluator INFO: Engine run starting with max_epochs=1." in source[2]
+
+
+def _test_distrib_one_rank_only():
+
+    def _test(barrier):
+        # last rank
+        rank = dist.get_world_size() - 1
+
+        value = torch.tensor(0)
+
+        @one_rank_only(rank=rank, barrier=barrier)
+        def initialize():
+            value.data = torch.tensor(100)
+
+        initialize()
+
+        value_list = [torch.tensor(0) for _ in range(dist.get_world_size())]
+
+        dist.all_gather(tensor=value, tensor_list=value_list)
+
+        for r in range(dist.get_world_size()):
+            if r == rank:
+                assert value_list[r].item() == 100
+            else:
+                assert value_list[r].item() == 0
+
+    _test(barrier=True)
+    _test(barrier=False)
+
+
+def _test_distrib_one_rank_only_with_engine():
+
+    def _test(barrier):
+        engine = Engine(lambda e, b: b)
+
+        batch_sum = torch.tensor(0)
+
+        @engine.on(Events.ITERATION_COMPLETED)
+        @one_rank_only(barrier=barrier)  # ie rank == 0
+        def _(_):
+            batch_sum.data += torch.tensor(engine.state.batch)
+
+        engine.run([1, 2, 3], max_epochs=2)
+
+        value_list = [torch.tensor(0) for _ in range(dist.get_world_size())]
+
+        dist.all_gather(tensor=batch_sum, tensor_list=value_list)
+
+        for r in range(dist.get_world_size()):
+            if r == 0:
+                assert value_list[r].item() == 12
+            else:
+                assert value_list[r].item() == 0
+
+    _test(barrier=True)
+    _test(barrier=False)
+
+
+@pytest.mark.distributed
+def test_distrib_cpu(distributed_context_single_node_gloo):
+    _test_distrib_one_rank_only()
+    _test_distrib_one_rank_only_with_engine()
+
+
+@pytest.mark.distributed
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
+def test_distrib_gpu(distributed_context_single_node_nccl):
+    _test_distrib_one_rank_only()
+    _test_distrib_one_rank_only_with_engine()
+
+
+@pytest.mark.multinode_distributed
+@pytest.mark.skipif("MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
+def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
+    _test_distrib_one_rank_only()
+    _test_distrib_one_rank_only_with_engine()
+
+
+@pytest.mark.multinode_distributed
+@pytest.mark.skipif("GPU_MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
+def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
+    _test_distrib_one_rank_only()
+    _test_distrib_one_rank_only_with_engine()