Repository: optuna/optuna

Issue:
ValueError: Rank 0 node expects an optuna.trial.Trial instance as the trial argument.
### Expected behavior

`FrozenTrial` can be used instead of `Trial` in PyTorch distributed integration.

### Environment

- Optuna version:3.1.0
- Python version:3.10.5
- OS:Linux-5.15.0-56-generic-x86_64-with-glibc2.35

### Error messages, stack traces, or logs

```shell
Traceback (most recent call last):
  File "/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/train_optuna.py", line 491, in <module>
    main()
  File "/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/train_optuna.py", line 471, in main
    train(trial)
  File "/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/train_optuna.py", line 124, in __call__
    trial = optuna.integration.TorchDistributedTrial(
  File "/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/_experimental.py", line 115, in wrapped_init
    _original_init(self, *args, **kwargs)
  File "/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/integration/pytorch_distributed.py", line 141, in __init__
    raise ValueError(
ValueError: Rank 0 node expects an optuna.trial.Trial instance as the trial argument.
```


### Steps to reproduce


```python
import os
import optuna
import torch
import torch.distributed

def objective(trial):
    rank = 0
    trial = optuna.integration.TorchDistributedTrial(
        trial if rank == 0 else None
    )
    x = trial.suggest_float("x", low=-100, high=100)
    return x * x


os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "1234"

torch.distributed.init_process_group(init_method="env://", world_size=1, rank=0, backend="gloo")


study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=1)

best_trial = study.best_trial

objective(best_trial)

```

===>
```
[I 2023-05-19 23:27:49,066] A new study created in memory with name: no-name-3d3ee40c-9011-4e59-aff2-b28606662609
/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/tmp.py:8: ExperimentalWarning: TorchDistributedTrial is experimental (supported from v2.6.0). The interface can change in the future.
  trial = optuna.integration.TorchDistributedTrial(
/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/integration/pytorch_distributed.py:157: FutureWarning: system_attrs has been deprecated in v3.1.0. This feature will be removed in v6.0.0. See https://github.com/optuna/optuna/releases/tag/v3.1.0.
  self._system_attrs = self._broadcast(getattr(self._delegate, "system_attrs", None))
/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/integration/pytorch_distributed.py:64: FutureWarning: system_attrs has been deprecated in v3.1.0. This feature will be removed in v6.0.0. See https://github.com/optuna/optuna/releases/tag/v3.1.0.
  self._delegate.system_attrs,
[I 2023-05-19 23:27:49,087] Trial 0 finished with value: 3632.3258131131442 and parameters: {'x': -60.26878014780663}. Best is trial 0 with value: 3632.3258131131442.
Traceback (most recent call last):
  File "/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/tmp.py", line 26, in <module>
    objective(best_trial)
  File "/mnt/nfs-mnj-archive-03/group/nict-compression/work/kfukuda/pfn-weather/LSSM/LSSM/tmp.py", line 8, in objective
    trial = optuna.integration.TorchDistributedTrial(
  File "/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/_experimental.py", line 115, in wrapped_init
    _original_init(self, *args, **kwargs)
  File "/home/team-nict-comp/.local/lib/python3.10/site-packages/optuna/integration/pytorch_distributed.py", line 141, in __init__
    raise ValueError(
ValueError: Rank 0 node expects an optuna.trial.Trial instance as the trial argument.
```


### Additional context (optional)

_No response_


Discussion:

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
@@ -112,7 +112,7 @@ class TorchDistributedTrial(optuna.trial.BaseTrial):
 
     def __init__(
         self,
-        trial: Optional[optuna.trial.Trial],
+        trial: Optional[optuna.trial.BaseTrial],
         group: Optional["ProcessGroup"] = None,
     ) -> None:
         _imports.check()
@@ -135,7 +135,7 @@ def __init__(
             self._group = _g_pg
 
         if dist.get_rank(self._group) == 0:
-            if not isinstance(trial, optuna.trial.Trial):
+            if not isinstance(trial, optuna.trial.BaseTrial):
                 raise ValueError(
                     "Rank 0 node expects an optuna.trial.Trial instance as the trial argument."
                 )
@@ -278,7 +278,9 @@ def set_system_attr(self, key: str, value: Any) -> None:
         if dist.get_rank(self._group) == 0:
             try:
                 assert self._delegate is not None
-                self._delegate.storage.set_trial_system_attr(self._delegate._trial_id, key, value)
+                self._delegate.storage.set_trial_system_attr(  # type: ignore[attr-defined]
+                    self._delegate._trial_id, key, value  # type: ignore[attr-defined]
+                )
             except Exception as e:
                 err = e
             err = self._broadcast(err)


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
@@ -351,3 +351,20 @@ def test_updates_properties(storage_mode: str) -> None:
             [getattr(trial, p) for p in property_names]
 
         dist.barrier()  # type: ignore[no-untyped-call]
+
+
+@pytest.mark.filterwarnings("ignore::optuna.exceptions.ExperimentalWarning")
+@pytest.mark.parametrize("storage_mode", STORAGE_MODES)
+def test_pass_frozen_trial_to_torch_distributed(storage_mode: str) -> None:
+    # Regression test of #4697
+    def objective(trial: optuna.trial.BaseTrial) -> float:
+        trial = optuna.integration.TorchDistributedTrial(trial if dist.get_rank() == 0 else None)
+        x = trial.suggest_float("x", low=-100, high=100)
+        return x * x
+
+    with StorageSupplier(storage_mode) as storage:
+        study = optuna.create_study(direction="minimize", storage=storage)
+        study.optimize(objective, n_trials=1)
+        best_trial = study.best_trial
+
+        objective(best_trial)