Repository: optuna/optuna

Issue:
Support latest PyTorch lightning
### Expected behavior

Optuna works with the latest PyTorch-Lightning (1.6.0).

### Environment

- Optuna version: 3.0.0b0.dev
- Python version: 
- OS: Independent
- (Optional) Other libraries and their versions: PyTorch-Lightning 1.6.0


### Error messages, stack traces, or logs

```shell
tests/integration_tests/test_pytorch_lightning.py:158: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
optuna/study/study.py:399: in optimize
    _optimize(
optuna/study/_optimize.py:68: in _optimize
    _optimize_sequential(
optuna/study/_optimize.py:162: in _optimize_sequential
    trial = _run_trial(study, func, catch)
optuna/study/_optimize.py:262: in _run_trial
    raise func_err
optuna/study/_optimize.py:211: in _run_trial
    value_or_values = func(trial)
tests/integration_tests/test_pytorch_lightning.py:143: in objective
    trainer = pl.Trainer(
/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/utilities/argparse.py:339: in insert_env_defaults
    return fn(self, **kwargs)
/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:561: in __init__
    self._call_callback_hooks("on_init_start")
/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1617: in _call_callback_hooks
    fn(self, *args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <optuna.integration.pytorch_lightning.PyTorchLightningPruningCallback object at 0x7f892c1c5730>
trainer = <pytorch_lightning.trainer.trainer.Trainer object at 0x7f892c1c5040>

    def on_init_start(self, trainer: Trainer) -> None:
>       self.is_ddp_backend = trainer._accelerator_connector.distributed_backend is not None
E       AttributeError: 'AcceleratorConnector' object has no attribute 'distributed_backend'

optuna/integration/pytorch_lightning.py:60: AttributeError
```

https://github.com/optuna/optuna/runs/5745775785?check_suite_focus=true

```
tests/integration_tests/test_pytorch_lightning.py:6: in <module>
    import pytorch_lightning as pl
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/__init__.py:30: in <module>
    from pytorch_lightning.callbacks import Callback  # noqa: E402
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/callbacks/__init__.py:26: in <module>
    from pytorch_lightning.callbacks.pruning import ModelPruning
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/callbacks/pruning.py:31: in <module>
    from pytorch_lightning.core.lightning import LightningModule
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/core/__init__.py:16: in <module>
    from pytorch_lightning.core.lightning import LightningModule
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py:40: in <module>
    from pytorch_lightning.loggers import LightningLoggerBase
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/loggers/__init__.py:18: in <module>
    from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/loggers/tensorboard.py:26: in <module>
    from torch.utils.tensorboard import SummaryWriter
../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/torch/utils/tensorboard/__init__.py:4: in <module>
    LooseVersion = distutils.version.LooseVersion
E   AttributeError: module 'distutils' has no attribute 'version'
```

https://github.com/optuna/optuna/runs/5745734509?check_suite_focus=true


### Steps to reproduce

See our CI failures.

### Additional context (optional)

It may be the simplest way to support PyTorch v1.11.0 (https://github.com/PyTorchLightning/pytorch-lightning/issues/12324).

🔗 https://github.com/optuna/optuna/pull/3417


Discussion:
I ran pytests related to pytorch on my local machine (M1 Mac). The following integration tests failed.

- ~fastai: the latest fastai requires torch<1.11,>=1.7.0~: latest stable fastai supports pytorch 1.11.0
- allelnlp: `Fatal Python error: Aborte` 🤔  and see also https://github.com/optuna/optuna/issues/3366
- pytorch-lightning: as described in this issue, we need to update our callback


Note that skorch has not supported pytorch 1.11.0 yet officially, but the latest stable w/ pytorch 1.11.0 passed Optuna's integration test.
This issue has not seen any recent activity.
Let's summarise the current candidates:
- Reimplement this feature using `Strategy`: see also https://github.com/optuna/optuna/pull/3431#issuecomment-1087290722
- Update the implementation as in https://github.com/optuna/optuna/pull/3431. The approach is similar to CatboosCallabck.
- try to approach shared by ptl developer in https://github.com/PyTorchLightning/pytorch-lightning/discussions/12569

This seems to be solved in #3510. We will close this issue.
I'm afraid not. As in https://github.com/optuna/optuna/blob/master/setup.py#L109, Optuna doesn't support the latest stable version: `1.6.*`.
Sorry, I mentioned the wrong line. It should have been https://github.com/optuna/optuna/blob/fd575d04cebae8b471e72aa890e4a4cb052aa083/.github/workflows/tests-integration.yml#L79. 
This issue has not seen any recent activity.
@nzw0301 would be happy to have a look if I find a bit of time. Did you narrow down to some preferred way of solving the problem?
@hrzn Yes, https://github.com/optuna/optuna/pull/3431 was my attempt to resolve the PL so it might be helpful.
This issue has not seen any recent activity.
This issue was closed automatically because it had not seen any recent activity. If you want to discuss it, you can reopen it freely.
This issue is still open for contribution.

The goal is to remove [this line](https://github.com/optuna/optuna/blob/594858fd3f43fbb713f411a47ff264b08b8d2337/.github/workflows/tests-integration.yml#L63) in the CI setting, after fixing the test failures.
@higucheese is kindly working on this issue now.
The current latest version of pytorch lightning is [1.8.0.post1](https://github.com/Lightning-AI/lightning/releases/tag/1.8.0.post1).

Resolving patch:
diff --git a/optuna/integration/pytorch_lightning.py b/optuna/integration/pytorch_lightning.py
--- a/optuna/integration/pytorch_lightning.py
+++ b/optuna/integration/pytorch_lightning.py
@@ -56,13 +56,11 @@ def __init__(self, trial: optuna.trial.Trial, monitor: str) -> None:
         self.monitor = monitor
         self.is_ddp_backend = False
 
-    def on_init_start(self, trainer: Trainer) -> None:
-        self.is_ddp_backend = (
-            trainer._accelerator_connector.distributed_backend is not None  # type: ignore
-        )
+    def on_fit_start(self, trainer: Trainer, pl_module: "pl.LightningModule") -> None:
+        self.is_ddp_backend = trainer._accelerator_connector.is_distributed
         if self.is_ddp_backend:
-            if version.parse(pl.__version__) < version.parse("1.5.0"):  # type: ignore
-                raise ValueError("PyTorch Lightning>=1.5.0 is required in DDP.")
+            if version.parse(pl.__version__) < version.parse("1.6.0"):  # type: ignore
+                raise ValueError("PyTorch Lightning>=1.6.0 is required in DDP.")
             if not (
                 isinstance(self._trial.study._storage, _CachedStorage)
                 and isinstance(self._trial.study._storage._backend, RDBStorage)
@@ -95,7 +93,7 @@ def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> Non
         if trainer.is_global_zero:
             self._trial.report(current_score.item(), step=epoch)
             should_stop = self._trial.should_prune()
-        should_stop = trainer.training_type_plugin.broadcast(should_stop)
+        should_stop = trainer.strategy.broadcast(should_stop)
         if not should_stop:
             return
 


Test patch:
diff --git a/tests/integration_tests/test_pytorch_lightning.py b/tests/integration_tests/test_pytorch_lightning.py
--- a/tests/integration_tests/test_pytorch_lightning.py
+++ b/tests/integration_tests/test_pytorch_lightning.py
@@ -143,14 +143,16 @@ def test_pytorch_lightning_pruning_callback_monitor_is_invalid() -> None:
         callback.on_validation_end(trainer, model)
 
 
+@pytest.mark.skip(reason="Currently DDP is not supported")
 @pytest.mark.parametrize("storage_mode", ["sqlite", "cached_sqlite"])
 def test_pytorch_lightning_pruning_callback_ddp_monitor(
     storage_mode: str,
 ) -> None:
     def objective(trial: optuna.trial.Trial) -> float:
         trainer = pl.Trainer(
-            max_epochs=2,
-            accelerator="ddp_cpu",
+            strategy="ddp",
+            accelerator="cpu",
+            devices=2,
             num_processes=2,
             enable_checkpointing=False,
             callbacks=[PyTorchLightningPruningCallback(trial, monitor="accuracy")],
@@ -177,14 +179,16 @@ def objective(trial: optuna.trial.Trial) -> float:
         np.testing.assert_almost_equal(study.trials[0].intermediate_values[1], 0.45)
 
 
+@pytest.mark.skip(reason="Currently DDP is not supported")
 def test_pytorch_lightning_pruning_callback_ddp_unsupported_storage() -> None:
     storage_mode = "inmemory"
 
     def objective(trial: optuna.trial.Trial) -> float:
         trainer = pl.Trainer(
             max_epochs=1,
-            accelerator="ddp_cpu",
-            num_processes=2,
+            strategy="ddp",
+            accelerator="cpu",
+            devices=2,
             enable_checkpointing=False,
             callbacks=[PyTorchLightningPruningCallback(trial, monitor="accuracy")],
         )