Repository: optuna/optuna

Issue:
TorchDistributedTrial properties can make torch.distributed hang
### Expected behavior

Reading TorchDistributedTrial properties should have no side effects, in particular it should not call torch.distributed methods. 

```
# ...
if rank == 0:
    print(hp_search_trial.number)

torch.distributed.barrier()
# execution continues
```

### Environment

- Optuna version: 2.10.0
- Python version: 3.9.12
- OS: Linux-4.18.0-240.el8.x86_64-x86_64-with-glibc2.28
- (Optional) Other libraries and their versions: Pytorch 1.11


### Error messages, stack traces, or logs

Processes hang after a trial property is read on one node only.
In the example above, reading the trial number on only one node causes a call to a torch.distributed method on one node only. The next distributed functions will hang because it expects the other nodes to call the same method first.


### Steps to reproduce

```py
import torch
import optuna
from torch import multiprocessing


def train(rank, world_size, trial):
    print(f"process rank {rank} started.")

    torch.distributed.init_process_group(
        backend="gloo",
        init_method="tcp://127.0.0.1:29400",
        world_size=world_size,
        rank=rank)

    print(f"process rank {rank} initialized.")

    trial = optuna.integration.TorchDistributedTrial(trial if rank == 0 else None)

    torch.distributed.barrier()
    print(f"process rank {rank} ready.")

    if rank == 0:
        print(trial.number)  # <<<------ rank 0 makes distributed op alone

    torch.distributed.barrier()  # <<<------ will hang here because rank 1 is late by one distributed op
    print(f"process rank {rank} exiting.")


def run_trial(trial):
    multiprocessing.spawn(train, (2, trial), nprocs=2)
    return 0


def main():
    study = optuna.create_study()
    study.optimize(run_trial, n_trials=1)


if __name__ == "__main__":
    main()
```

### Additional context (optional)

torch.distributed essentially works by keeping a history of called functions. For instance on two nodes, it may look like this:

| index | rank0 | rank1 |
|--|--|--|
| 0 | dist.all_reduce | dist.all_reduce |
| 1 | dist.all_reduce | dist.all_reduce |
| 2 | dist.broadcast |  |

In this example rank1 is a bit late so rank0 will hang before executing the next distributed operation.

Back to optuna, a common use-case in a torch.distributed setup is to log or save a checkpoint on rank 0:

```
if torch.distributed.get_rank() == 0:
    print(hp_search_trial.number)

torch.distributed.barrier()
# execution continues
```

Before this code, all nodes participating in the distributed environment have issued the same distributed functions calls.
The current implementation of `hp_search_trial.number` causes a call to torch.distributed.broadcast **but only on rank 0**.
Upon reaching the next distributed operation (`torch.distributed.barrier()` in this example), pytorch hangs on rank0 until rank1 executes a broadcast to mirror rank0 operations but this never happens.

Optuna documentation already states that trial methods use the distributed api and thus must be called on all nodes, but properties are not mentioned.
I think it is inconvenient and non-pythonic that reading properties has a side effect.


Discussion:
I'm sorry I could not understand what is issue. Could you elaborate a little bit more about the bug?  We really appreciate you if you share minimal reproducible codes with us.
Sure! I'm working on it and a PR too.
Amazing. Thanks!

Resolving patch:
diff --git a/optuna/integration/pytorch_distributed.py b/optuna/integration/pytorch_distributed.py
--- a/optuna/integration/pytorch_distributed.py
+++ b/optuna/integration/pytorch_distributed.py
@@ -1,10 +1,14 @@
 from datetime import datetime
+import functools
 import pickle
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import Optional
 from typing import Sequence
+from typing import TypeVar
+
+from typing_extensions import ParamSpec
 
 import optuna
 from optuna._deprecated import deprecated
@@ -19,11 +23,56 @@
     import torch.distributed as dist
 
 
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
 _suggest_deprecated_msg = (
     "Use :func:`~optuna.integration.TorchDistributedTrial.suggest_float` instead."
 )
 
 
+def broadcast_properties(f: Callable[_P, _T]) -> Callable[_P, _T]:
+    """Method decorator to fetch updated trial properties from rank 0 after ``f`` is run.
+
+    This decorator ensures trial properties (params, distributions, etc.) on all distributed
+    processes are up-to-date with the wrapped trial stored on rank 0.
+    It should be applied to all :class:`~optuna.integration.TorchDistributedTrial`
+    methods that update property values.
+    """
+
+    @functools.wraps(f)
+    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+        # TODO(nlgranger): Remove type ignore after mypy includes
+        # https://github.com/python/mypy/pull/12668
+        self: TorchDistributedTrial = args[0]  # type: ignore
+
+        def fetch_properties() -> Sequence:
+            assert self._delegate is not None
+            return (
+                self._delegate.number,
+                self._delegate.params,
+                self._delegate.distributions,
+                self._delegate.user_attrs,
+                self._delegate.system_attrs,
+                self._delegate.datetime_start,
+            )
+
+        try:
+            return f(*args, **kwargs)
+        finally:
+            (
+                self._number,
+                self._params,
+                self._distributions,
+                self._user_attrs,
+                self._system_attrs,
+                self._datetime_start,
+            ) = self._call_and_communicate_obj(fetch_properties)
+
+    return wrapped
+
+
 @experimental("2.6.0")
 class TorchDistributedTrial(optuna.trial.BaseTrial):
     """A wrapper of :class:`~optuna.trial.Trial` to incorporate Optuna with PyTorch distributed.
@@ -74,6 +123,14 @@ def __init__(
         self._delegate = trial
         self._device = device
 
+        self._number = self._broadcast(getattr(self._delegate, "number", None))
+        self._params = self._broadcast(getattr(self._delegate, "params", None))
+        self._distributions = self._broadcast(getattr(self._delegate, "distributions", None))
+        self._user_attrs = self._broadcast(getattr(self._delegate, "user_attrs", None))
+        self._system_attrs = self._broadcast(getattr(self._delegate, "system_attrs", None))
+        self._datetime_start = self._broadcast(getattr(self._delegate, "datetime_start", None))
+
+    @broadcast_properties
     def suggest_float(
         self,
         name: str,
@@ -105,6 +162,7 @@ def suggest_discrete_uniform(self, name: str, low: float, high: float, q: float)
 
         return self.suggest_float(name, low, high, step=q)
 
+    @broadcast_properties
     def suggest_int(self, name: str, low: int, high: int, step: int = 1, log: bool = False) -> int:
         def func() -> float:
 
@@ -113,6 +171,7 @@ def func() -> float:
 
         return self._call_and_communicate(func, torch.int)
 
+    @broadcast_properties
     def suggest_categorical(self, name: str, choices: Sequence["CategoricalChoiceType"]) -> Any:
         def func() -> CategoricalChoiceType:
 
@@ -121,6 +180,7 @@ def func() -> CategoricalChoiceType:
 
         return self._call_and_communicate_obj(func)
 
+    @broadcast_properties
     def report(self, value: float, step: int) -> None:
         err = None
         if dist.get_rank() == 0:
@@ -136,6 +196,7 @@ def report(self, value: float, step: int) -> None:
         if err is not None:
             raise err
 
+    @broadcast_properties
     def should_prune(self) -> bool:
         def func() -> bool:
 
@@ -147,6 +208,7 @@ def func() -> bool:
         # due to the RuntimeError.
         return self._call_and_communicate(func, torch.uint8)
 
+    @broadcast_properties
     def set_user_attr(self, key: str, value: Any) -> None:
         err = None
         if dist.get_rank() == 0:
@@ -162,6 +224,7 @@ def set_user_attr(self, key: str, value: Any) -> None:
         if err is not None:
             raise err
 
+    @broadcast_properties
     def set_system_attr(self, key: str, value: Any) -> None:
         err = None
 
@@ -180,57 +243,27 @@ def set_system_attr(self, key: str, value: Any) -> None:
 
     @property
     def number(self) -> int:
-        def func() -> int:
-
-            assert self._delegate is not None
-            return self._delegate.number
-
-        return self._call_and_communicate(func, torch.int)
+        return self._number
 
     @property
     def params(self) -> Dict[str, Any]:
-        def func() -> Dict[str, Any]:
-
-            assert self._delegate is not None
-            return self._delegate.params
-
-        return self._call_and_communicate_obj(func)
+        return self._params
 
     @property
     def distributions(self) -> Dict[str, BaseDistribution]:
-        def func() -> Dict[str, BaseDistribution]:
-
-            assert self._delegate is not None
-            return self._delegate.distributions
-
-        return self._call_and_communicate_obj(func)
+        return self._distributions
 
     @property
     def user_attrs(self) -> Dict[str, Any]:
-        def func() -> Dict[str, Any]:
-
-            assert self._delegate is not None
-            return self._delegate.user_attrs
-
-        return self._call_and_communicate_obj(func)
+        return self._user_attrs
 
     @property
     def system_attrs(self) -> Dict[str, Any]:
-        def func() -> Dict[str, Any]:
-
-            assert self._delegate is not None
-            return self._delegate.system_attrs
-
-        return self._call_and_communicate_obj(func)
+        return self._system_attrs
 
     @property
     def datetime_start(self) -> Optional[datetime]:
-        def func() -> Optional[datetime]:
-
-            assert self._delegate is not None
-            return self._delegate.datetime_start
-
-        return self._call_and_communicate_obj(func)
+        return self._datetime_start
 
     def _call_and_communicate(self, func: Callable, dtype: "torch.dtype") -> Any:
         buffer = torch.empty(1, dtype=dtype)


Test patch:
diff --git a/tests/integration_tests/test_pytorch_distributed.py b/tests/integration_tests/test_pytorch_distributed.py
--- a/tests/integration_tests/test_pytorch_distributed.py
+++ b/tests/integration_tests/test_pytorch_distributed.py
@@ -26,7 +26,7 @@ def init_process_group() -> None:
     os.environ["MASTER_ADDR"] = "127.0.0.1"
     os.environ["MASTER_PORT"] = "20000"
 
-    dist.init_process_group("gloo")
+    dist.init_process_group("gloo", timeout=datetime.timedelta(seconds=15))
 
 
 def test_torch_distributed_trial_experimental_warning() -> None:
@@ -334,3 +334,37 @@ def test_distributions(storage_mode: str) -> None:
         assert distributions["i"] == optuna.distributions.IntDistribution(0, 1)
         assert distributions["il"] == optuna.distributions.IntDistribution(1, 128, log=True)
         assert distributions["c"] == optuna.distributions.CategoricalDistribution(("a", "b", "c"))
+
+
+@pytest.mark.filterwarnings("ignore::optuna.exceptions.ExperimentalWarning")
+@pytest.mark.parametrize("storage_mode", STORAGE_MODES)
+def test_updates_properties(storage_mode: str) -> None:
+    """Check for any distributed deadlock following a property read."""
+    with StorageSupplier(storage_mode) as storage:
+        if dist.get_rank() == 0:
+            study = optuna.create_study(storage=storage)
+            trial = TorchDistributedTrial(study.ask())
+        else:
+            trial = TorchDistributedTrial(None)
+
+        trial.suggest_float("f", 0, 1)
+        trial.suggest_int("i", 0, 1)
+        trial.suggest_categorical("c", ("a", "b", "c"))
+
+        property_names = [
+            p
+            for p in dir(TorchDistributedTrial)
+            if isinstance(getattr(TorchDistributedTrial, p), property)
+        ]
+
+        # Rank 0 can read properties without deadlock.
+        if dist.get_rank() == 0:
+            [getattr(trial, p) for p in property_names]
+
+        dist.barrier()
+
+        # Same with rank 1.
+        if dist.get_rank() == 1:
+            [getattr(trial, p) for p in property_names]
+
+        dist.barrier()