Repository: optuna/optuna
Pull request: Remove `torch.distributed` calls from `TorchDistributedTrial` properties
Fixes #3484

## Motivation

Removes all calls to torch.distributed from TorchDistributedTrial properties as they have undesired side-effects: namely they break the synchronization of distributed processes if one forgets to call them simultaneously on all distributed processes.

## Description of the changes
<!-- Describe the changes in this PR. -->
Prior to this PR, reading properties on TorchDistributedTrial triggered a call to torch.distributed.broadcast in order to retrieve the values, this was necessary because they are solely stored on the process with rank 0.
This PR instead broadcasts the values to all processes each time they are updated by calls to suggest_uniform, etc. This is not inconvenient because these functions are already called on all distributed processes anyways and Optuna documentation is clear about this.

Problem statement:
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.

Hints from issue 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!
I have updated the issue with an MRE and submitted the PR.
I'm afraid that this is not a bug, but an expected behavior of the current `TorchDistriubtedTrial`.

[The official example](https://github.com/optuna/optuna-examples/blob/e3ea68dece4bf2f59a5d1731139e13b416aa9cb2/pytorch/pytorch_distributed_simple.py#L99-L100) distinguishes `single_trial: Optional[Trial]` and `trial: TorchDistributedTrial`. Note that `single_trial` is an instance of `Trial` if rank==0, otherwise, it is `None`.

The rank-0 node has a Trial instance, so the code should refer to it there.
```
    if rank == 0:
        assert single_trial is not None
        print(single_trial.number)  # <<<------ Use the  Trial instance instead of `TorchDistributedTrial`
```

Currently, users need to choose single trials and distributed trials carefully, and it may be annoying. So, I appreciate your PR that will improve usability.
Let me change the label to `enhancement`.
The doc does not mention attributes actually being properties with calls to torch.distributed.broadcast, that got me confused.

Your comment makes me think we could simply remove the proxy to the single trial properties when rank != 0 and return the readily available single trial properties when rank == 0. It would be a lot simpler than the proposed PR and would likely not impact many users. Are there any other optuna functions that would need access to the trial properties on rank != 0 ?
> simply remove the proxy to the single trial properties when rank != 0 and return the readily available single trial properties when rank == 0.

Thank you for your proposal. I couldn't imagine the implementation, so could you share some code examples?
More specifically, I think a `TorchDistributedTrial` instance in rank-0 node needs to broadcast properties in order for the other rank nodes to get the value of properties.
This issue has not seen any recent activity.

Commit message: Ignore false positive type check error

Co-authored-by: Toshihiko Yanase <toshihiko.yanase@gmail.com>

Patch to review:
diff --git a/optuna/integration/pytorch_distributed.py b/optuna/integration/pytorch_distributed.py
index 2e62b9a1e0c..7ea8d95bddb 100644
--- 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 updates_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 state() -> 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(state)
+
+    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))
+
+    @updates_properties
     def suggest_float(
         self,
         name: str,
@@ -91,20 +148,24 @@ def func() -> float:
         return self._call_and_communicate(func, torch.float)
 
     @deprecated("3.0.0", "6.0.0", text=_suggest_deprecated_msg)
+    @updates_properties
     def suggest_uniform(self, name: str, low: float, high: float) -> float:
 
         return self.suggest_float(name, low, high)
 
     @deprecated("3.0.0", "6.0.0", text=_suggest_deprecated_msg)
+    @updates_properties
     def suggest_loguniform(self, name: str, low: float, high: float) -> float:
 
         return self.suggest_float(name, low, high, log=True)
 
     @deprecated("3.0.0", "6.0.0", text=_suggest_deprecated_msg)
+    @updates_properties
     def suggest_discrete_uniform(self, name: str, low: float, high: float, q: float) -> float:
 
         return self.suggest_float(name, low, high, step=q)
 
+    @updates_properties
     def suggest_int(self, name: str, low: int, high: int, step: int = 1, log: bool = False) -> int:
         def func() -> float:
 
@@ -113,6 +174,7 @@ def func() -> float:
 
         return self._call_and_communicate(func, torch.int)
 
+    @updates_properties
     def suggest_categorical(self, name: str, choices: Sequence["CategoricalChoiceType"]) -> Any:
         def func() -> CategoricalChoiceType:
 
@@ -121,6 +183,7 @@ def func() -> CategoricalChoiceType:
 
         return self._call_and_communicate_obj(func)
 
+    @updates_properties
     def report(self, value: float, step: int) -> None:
         err = None
         if dist.get_rank() == 0:
@@ -136,6 +199,7 @@ def report(self, value: float, step: int) -> None:
         if err is not None:
             raise err
 
+    @updates_properties
     def should_prune(self) -> bool:
         def func() -> bool:
 
@@ -147,6 +211,7 @@ def func() -> bool:
         # due to the RuntimeError.
         return self._call_and_communicate(func, torch.uint8)
 
+    @updates_properties
     def set_user_attr(self, key: str, value: Any) -> None:
         err = None
         if dist.get_rank() == 0:
@@ -162,6 +227,7 @@ def set_user_attr(self, key: str, value: Any) -> None:
         if err is not None:
             raise err
 
+    @updates_properties
     def set_system_attr(self, key: str, value: Any) -> None:
         err = None
 
@@ -180,57 +246,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)
diff --git a/tests/integration_tests/test_pytorch_distributed.py b/tests/integration_tests/test_pytorch_distributed.py
index 9257ffc03ca..0324c5dd4ba 100644
--- 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()


Code review comments:
diff_hunk: @@ -91,20 +148,24 @@ def func() -> float:
         return self._call_and_communicate(func, torch.float)
 
     @deprecated("3.0.0", "6.0.0", text=_suggest_deprecated_msg)
+    @updates_properties
path: optuna/integration/pytorch_distributed.py
text: @user1:
Since this method calls `suggest_float` that is decorated by `updates_properties`, can we remove `updates_properties`?
```suggestion
```

I suppose the other deprecated `suggest_*` methods too.

@author:
Yes it is possible but I though if all methods are decorated one can change the implementation without having to remember about adding it.

@user1:
Thanks. Hmm, but in my understanding, when we call one of these deprecated suggest methods, `updates_properties` is executed twice, which is supposed to be unnecessary communications among processes. 

@user2  What do you think about this?


@user2:
Good catch! `suggest_uniform` simply calls `suggest_float` and `suggest_float` actually update the properties. So, @user1  suggestion sounds reasonable to me.

> one can change the implementation without having to remember about adding it.

That makes sense, but we will not update the deprecated methods so frequently.

@author:
I have removed the redundant decorators in the latest commit.

@user1:
Thanks!