Repository: pytorch/ignite

Issue:
Provide a context manager for running distributed code on a given process first
## 🚀 Feature

```python
import ignite.distibuted as idist

with idist.rank_process_first(rank=1):
    data_setup_code(...)
```
where behind we would be doing something like:
```python
    if idist.get_local_rank() != rank:
        idist.barrier()

    data_setup_code(...)

    if idist.get_local_rank() == rank:
        idist.barrier()
```




Discussion:

Resolving patch:
diff --git a/examples/contrib/cifar10/main.py b/examples/contrib/cifar10/main.py
--- a/examples/contrib/cifar10/main.py
+++ b/examples/contrib/cifar10/main.py
@@ -211,16 +211,8 @@ def run(
 
 def get_dataflow(config):
     # - Get train/test datasets
-    if idist.get_local_rank() > 0:
-        # Ensure that only local rank 0 download the dataset
-        # Thus each node will download a copy of the dataset
-        idist.barrier()
-
-    train_dataset, test_dataset = utils.get_train_test_datasets(config["data_path"])
-
-    if idist.get_local_rank() == 0:
-        # Ensure that only local rank 0 download the dataset
-        idist.barrier()
+    with idist.one_rank_first(local=True):
+        train_dataset, test_dataset = utils.get_train_test_datasets(config["data_path"])
 
     # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
     train_loader = idist.auto_dataloader(
diff --git a/examples/contrib/cifar10_qat/main.py b/examples/contrib/cifar10_qat/main.py
--- a/examples/contrib/cifar10_qat/main.py
+++ b/examples/contrib/cifar10_qat/main.py
@@ -195,16 +195,8 @@ def run(
 
 def get_dataflow(config):
     # - Get train/test datasets
-    if idist.get_local_rank() > 0:
-        # Ensure that only local rank 0 download the dataset
-        # Thus each node will download a copy of the dataset
-        idist.barrier()
-
-    train_dataset, test_dataset = utils.get_train_test_datasets(config["data_path"])
-
-    if idist.get_local_rank() == 0:
-        # Ensure that only local rank 0 download the dataset
-        idist.barrier()
+    with idist.one_rank_first(local=True):
+        train_dataset, test_dataset = utils.get_train_test_datasets(config["data_path"])
 
     # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
     train_loader = idist.auto_dataloader(
diff --git a/ignite/distributed/utils.py b/ignite/distributed/utils.py
--- a/ignite/distributed/utils.py
+++ b/ignite/distributed/utils.py
@@ -1,4 +1,5 @@
 import socket
+from contextlib import contextmanager
 from functools import wraps
 from typing import Any, Callable, List, Mapping, Optional, Tuple, Union
 
@@ -41,6 +42,7 @@
     "registered_computation_models",
     "one_rank_only",
     "new_group",
+    "one_rank_first",
 ]
 
 _model = _SerialModel()
@@ -635,3 +637,44 @@ def wrapper(*args: Any, **kwargs: Any) -> Optional[Any]:
         return wrapper
 
     return _one_rank_only
+
+
+@contextmanager
+def one_rank_first(rank: int = 0, local: bool = False) -> Any:
+    """Context manager that ensures a specific rank runs first before others in a distributed
+    environment.
+
+    Args:
+        rank: rank of the process that should execute the code
+            block inside the context manager first. Default, 0.
+        local: flag to specify local rank or global rank.
+            If True ``rank`` argument will define a local rank to run first.
+            Default, False
+
+    Examples:
+        .. code-block:: python
+
+            def download_dataset():
+                ...
+
+            with idist.one_rank_first():
+                ds = download_dataset()
+
+            dp = ds[0]
+
+    .. versionadded:: 0.5.0
+    """
+
+    current_rank = get_local_rank() if local else get_rank()
+    size = get_nproc_per_node() if local else get_world_size()
+
+    if rank >= size or rank < 0:
+        raise ValueError(f"rank should be between 0 and {size - 1}, but given {rank}")
+
+    if current_rank != rank:
+        barrier()
+
+    yield
+
+    if current_rank == rank:
+        barrier()


Test patch:
diff --git a/tests/ignite/distributed/utils/test_native.py b/tests/ignite/distributed/utils/test_native.py
--- a/tests/ignite/distributed/utils/test_native.py
+++ b/tests/ignite/distributed/utils/test_native.py
@@ -370,3 +370,45 @@ def test_idist_one_rank_only_nccl(local_rank, distributed_context_single_node_nc
     device = idist.device()
     _test_distrib_one_rank_only(device=device)
     _test_distrib_one_rank_only_with_engine(device=device)
+
+
+@pytest.mark.distributed
+@pytest.mark.parametrize("rank", range(int(os.environ.get("WORLD_SIZE", 1))))
+@pytest.mark.parametrize("local", [True, False])
+def test_one_rank_first(distributed, get_rank_zero_dirname, rank, local):
+    def get_ds(file_path):
+        rank = idist.get_local_rank() if local else idist.get_rank()
+        if not file_path.exists():
+            with open(file_path, "w") as f:
+                f.write("readed")
+            return f"{rank} not readed"
+        else:
+            return f"{rank} readed"
+
+    folder = get_rank_zero_dirname()
+    file_path = folder / "res.txt"
+
+    with idist.one_rank_first(rank, local=local):
+        x = get_ds(file_path)
+
+    output = idist.all_gather(x)
+
+    if local:
+        expected = [
+            f"{x} not readed" if x == rank else f"{x} readed" for x in range(idist.get_nproc_per_node())
+        ] * idist.get_nnodes()
+    else:
+        expected = [f"{x} not readed" if x == rank else f"{x} readed" for x in range(idist.get_world_size())]
+
+    print("expected:", expected, idist.get_nnodes())
+    assert set(expected) == set(output)
+
+
+@pytest.mark.distributed
+def test_one_rank_first_asserts():
+    rank = 100
+    with pytest.raises(
+        ValueError, match=f"rank should be between 0 and {idist.get_world_size() - 1}, but given {rank}"
+    ):
+        with idist.one_rank_first(rank):
+            pass