We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optuna/integration/pytorch_lightning.py`
Content:
```
1 import warnings
2 
3 from packaging import version
4 
5 import optuna
6 from optuna.storages._cached_storage import _CachedStorage
7 from optuna.storages._rdb.storage import RDBStorage
8 
9 
10 # Define key names of `Trial.system_attrs`.
11 _PRUNED_KEY = "ddp_pl:pruned"
12 _EPOCH_KEY = "ddp_pl:epoch"
13 
14 
15 with optuna._imports.try_import() as _imports:
16     import pytorch_lightning as pl
17     from pytorch_lightning import LightningModule
18     from pytorch_lightning import Trainer
19     from pytorch_lightning.callbacks import Callback
20 
21 if not _imports.is_successful():
22     Callback = object  # type: ignore  # NOQA
23     LightningModule = object  # type: ignore  # NOQA
24     Trainer = object  # type: ignore  # NOQA
25 
26 
27 class PyTorchLightningPruningCallback(Callback):
28     """PyTorch Lightning callback to prune unpromising trials.
29 
30     See `the example <https://github.com/optuna/optuna-examples/blob/
31     main/pytorch/pytorch_lightning_simple.py>`__
32     if you want to add a pruning callback which observes accuracy.
33 
34     Args:
35         trial:
36             A :class:`~optuna.trial.Trial` corresponding to the current evaluation of the
37             objective function.
38         monitor:
39             An evaluation metric for pruning, e.g., ``val_loss`` or
40             ``val_acc``. The metrics are obtained from the returned dictionaries from e.g.
41             ``pytorch_lightning.LightningModule.training_step`` or
42             ``pytorch_lightning.LightningModule.validation_epoch_end`` and the names thus depend on
43             how this dictionary is formatted.
44 
45     .. note::
46         For the distributed data parallel training, the version of PyTorchLightning needs to be
47         higher than or equal to v1.5.0. In addition, :class:`~optuna.study.Study` should be
48         instantiated with RDB storage.
49     """
50 
51     def __init__(self, trial: optuna.trial.Trial, monitor: str) -> None:
52         _imports.check()
53         super().__init__()
54 
55         self._trial = trial
56         self.monitor = monitor
57         self.is_ddp_backend = False
58 
59     def on_init_start(self, trainer: Trainer) -> None:
60         self.is_ddp_backend = (
61             trainer._accelerator_connector.distributed_backend is not None  # type: ignore
62         )
63         if self.is_ddp_backend:
64             if version.parse(pl.__version__) < version.parse("1.5.0"):  # type: ignore
65                 raise ValueError("PyTorch Lightning>=1.5.0 is required in DDP.")
66             if not (
67                 isinstance(self._trial.study._storage, _CachedStorage)
68                 and isinstance(self._trial.study._storage._backend, RDBStorage)
69             ):
70                 raise ValueError(
71                     "optuna.integration.PyTorchLightningPruningCallback"
72                     " supports only optuna.storages.RDBStorage in DDP."
73                 )
74 
75     def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
76         # When the trainer calls `on_validation_end` for sanity check,
77         # do not call `trial.report` to avoid calling `trial.report` multiple times
78         # at epoch 0. The related page is
79         # https://github.com/PyTorchLightning/pytorch-lightning/issues/1391.
80         if trainer.sanity_checking:
81             return
82 
83         epoch = pl_module.current_epoch
84 
85         current_score = trainer.callback_metrics.get(self.monitor)
86         if current_score is None:
87             message = (
88                 "The metric '{}' is not in the evaluation logs for pruning. "
89                 "Please make sure you set the correct metric name.".format(self.monitor)
90             )
91             warnings.warn(message)
92             return
93 
94         should_stop = False
95         if trainer.is_global_zero:
96             self._trial.report(current_score.item(), step=epoch)
97             should_stop = self._trial.should_prune()
98         should_stop = trainer.training_type_plugin.broadcast(should_stop)
99         if not should_stop:
100             return
101 
102         if not self.is_ddp_backend:
103             message = "Trial was pruned at epoch {}.".format(epoch)
104             raise optuna.TrialPruned(message)
105         else:
106             # Stop every DDP process if global rank 0 process decides to stop.
107             trainer.should_stop = True
108             if trainer.is_global_zero:
109                 self._trial.storage.set_trial_system_attr(self._trial._trial_id, _PRUNED_KEY, True)
110                 self._trial.storage.set_trial_system_attr(self._trial._trial_id, _EPOCH_KEY, epoch)
111 
112     def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
113         if not self.is_ddp_backend:
114             return
115 
116         # Because on_validation_end is executed in spawned processes,
117         # _trial.report is necessary to update the memory in main process, not to update the RDB.
118         _trial_id = self._trial._trial_id
119         _study = self._trial.study
120         _trial = _study._storage._backend.get_trial(_trial_id)  # type: ignore
121         _trial_system_attrs = _study._storage.get_trial_system_attrs(_trial_id)
122         is_pruned = _trial_system_attrs.get(_PRUNED_KEY)
123         epoch = _trial_system_attrs.get(_EPOCH_KEY)
124         intermediate_values = _trial.intermediate_values
125         for step, value in intermediate_values.items():
126             self._trial.report(value, step=step)
127 
128         if is_pruned:
129             message = "Trial was pruned at epoch {}.".format(epoch)
130             raise optuna.TrialPruned(message)
131
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
         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

{"golden_diff": "diff --git a/optuna/integration/pytorch_lightning.py b/optuna/integration/pytorch_lightning.py\n--- a/optuna/integration/pytorch_lightning.py\n+++ b/optuna/integration/pytorch_lightning.py\n@@ -56,13 +56,11 @@\n         self.monitor = monitor\n         self.is_ddp_backend = False\n \n-    def on_init_start(self, trainer: Trainer) -> None:\n-        self.is_ddp_backend = (\n-            trainer._accelerator_connector.distributed_backend is not None  # type: ignore\n-        )\n+    def on_fit_start(self, trainer: Trainer, pl_module: \"pl.LightningModule\") -> None:\n+        self.is_ddp_backend = trainer._accelerator_connector.is_distributed\n         if self.is_ddp_backend:\n-            if version.parse(pl.__version__) < version.parse(\"1.5.0\"):  # type: ignore\n-                raise ValueError(\"PyTorch Lightning>=1.5.0 is required in DDP.\")\n+            if version.parse(pl.__version__) < version.parse(\"1.6.0\"):  # type: ignore\n+                raise ValueError(\"PyTorch Lightning>=1.6.0 is required in DDP.\")\n             if not (\n                 isinstance(self._trial.study._storage, _CachedStorage)\n                 and isinstance(self._trial.study._storage._backend, RDBStorage)\n@@ -95,7 +93,7 @@\n         if trainer.is_global_zero:\n             self._trial.report(current_score.item(), step=epoch)\n             should_stop = self._trial.should_prune()\n-        should_stop = trainer.training_type_plugin.broadcast(should_stop)\n+        should_stop = trainer.strategy.broadcast(should_stop)\n         if not should_stop:\n             return\n", "issue": "Support latest PyTorch lightning\n### Expected behavior\r\n\r\nOptuna works with the latest PyTorch-Lightning (1.6.0).\r\n\r\n### Environment\r\n\r\n- Optuna version: 3.0.0b0.dev\r\n- Python version: \r\n- OS: Independent\r\n- (Optional) Other libraries and their versions: PyTorch-Lightning 1.6.0\r\n\r\n\r\n### Error messages, stack traces, or logs\r\n\r\n```shell\r\ntests/integration_tests/test_pytorch_lightning.py:158: \r\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \r\noptuna/study/study.py:399: in optimize\r\n    _optimize(\r\noptuna/study/_optimize.py:68: in _optimize\r\n    _optimize_sequential(\r\noptuna/study/_optimize.py:162: in _optimize_sequential\r\n    trial = _run_trial(study, func, catch)\r\noptuna/study/_optimize.py:262: in _run_trial\r\n    raise func_err\r\noptuna/study/_optimize.py:211: in _run_trial\r\n    value_or_values = func(trial)\r\ntests/integration_tests/test_pytorch_lightning.py:143: in objective\r\n    trainer = pl.Trainer(\r\n/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/utilities/argparse.py:339: in insert_env_defaults\r\n    return fn(self, **kwargs)\r\n/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:561: in __init__\r\n    self._call_callback_hooks(\"on_init_start\")\r\n/opt/hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1617: in _call_callback_hooks\r\n    fn(self, *args, **kwargs)\r\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \r\n\r\nself = <optuna.integration.pytorch_lightning.PyTorchLightningPruningCallback object at 0x7f892c1c5730>\r\ntrainer = <pytorch_lightning.trainer.trainer.Trainer object at 0x7f892c1c5040>\r\n\r\n    def on_init_start(self, trainer: Trainer) -> None:\r\n>       self.is_ddp_backend = trainer._accelerator_connector.distributed_backend is not None\r\nE       AttributeError: 'AcceleratorConnector' object has no attribute 'distributed_backend'\r\n\r\noptuna/integration/pytorch_lightning.py:60: AttributeError\r\n```\r\n\r\nhttps://github.com/optuna/optuna/runs/5745775785?check_suite_focus=true\r\n\r\n```\r\ntests/integration_tests/test_pytorch_lightning.py:6: in <module>\r\n    import pytorch_lightning as pl\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/__init__.py:30: in <module>\r\n    from pytorch_lightning.callbacks import Callback  # noqa: E402\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/callbacks/__init__.py:26: in <module>\r\n    from pytorch_lightning.callbacks.pruning import ModelPruning\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/callbacks/pruning.py:31: in <module>\r\n    from pytorch_lightning.core.lightning import LightningModule\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/core/__init__.py:16: in <module>\r\n    from pytorch_lightning.core.lightning import LightningModule\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py:40: in <module>\r\n    from pytorch_lightning.loggers import LightningLoggerBase\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/loggers/__init__.py:18: in <module>\r\n    from pytorch_lightning.loggers.tensorboard import TensorBoardLogger\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/pytorch_lightning/loggers/tensorboard.py:26: in <module>\r\n    from torch.utils.tensorboard import SummaryWriter\r\n../../../hostedtoolcache/Python/3.8.12/x64/lib/python3.8/site-packages/torch/utils/tensorboard/__init__.py:4: in <module>\r\n    LooseVersion = distutils.version.LooseVersion\r\nE   AttributeError: module 'distutils' has no attribute 'version'\r\n```\r\n\r\nhttps://github.com/optuna/optuna/runs/5745734509?check_suite_focus=true\r\n\r\n\r\n### Steps to reproduce\r\n\r\nSee our CI failures.\r\n\r\n### Additional context (optional)\r\n\r\nIt may be the simplest way to support PyTorch v1.11.0 (https://github.com/PyTorchLightning/pytorch-lightning/issues/12324).\r\n\r\n\ud83d\udd17 https://github.com/optuna/optuna/pull/3417\n", "before_files": [{"content": "import warnings\n\nfrom packaging import version\n\nimport optuna\nfrom optuna.storages._cached_storage import _CachedStorage\nfrom optuna.storages._rdb.storage import RDBStorage\n\n\n# Define key names of `Trial.system_attrs`.\n_PRUNED_KEY = \"ddp_pl:pruned\"\n_EPOCH_KEY = \"ddp_pl:epoch\"\n\n\nwith optuna._imports.try_import() as _imports:\n    import pytorch_lightning as pl\n    from pytorch_lightning import LightningModule\n    from pytorch_lightning import Trainer\n    from pytorch_lightning.callbacks import Callback\n\nif not _imports.is_successful():\n    Callback = object  # type: ignore  # NOQA\n    LightningModule = object  # type: ignore  # NOQA\n    Trainer = object  # type: ignore  # NOQA\n\n\nclass PyTorchLightningPruningCallback(Callback):\n    \"\"\"PyTorch Lightning callback to prune unpromising trials.\n\n    See `the example <https://github.com/optuna/optuna-examples/blob/\n    main/pytorch/pytorch_lightning_simple.py>`__\n    if you want to add a pruning callback which observes accuracy.\n\n    Args:\n        trial:\n            A :class:`~optuna.trial.Trial` corresponding to the current evaluation of the\n            objective function.\n        monitor:\n            An evaluation metric for pruning, e.g., ``val_loss`` or\n            ``val_acc``. The metrics are obtained from the returned dictionaries from e.g.\n            ``pytorch_lightning.LightningModule.training_step`` or\n            ``pytorch_lightning.LightningModule.validation_epoch_end`` and the names thus depend on\n            how this dictionary is formatted.\n\n    .. note::\n        For the distributed data parallel training, the version of PyTorchLightning needs to be\n        higher than or equal to v1.5.0. In addition, :class:`~optuna.study.Study` should be\n        instantiated with RDB storage.\n    \"\"\"\n\n    def __init__(self, trial: optuna.trial.Trial, monitor: str) -> None:\n        _imports.check()\n        super().__init__()\n\n        self._trial = trial\n        self.monitor = monitor\n        self.is_ddp_backend = False\n\n    def on_init_start(self, trainer: Trainer) -> None:\n        self.is_ddp_backend = (\n            trainer._accelerator_connector.distributed_backend is not None  # type: ignore\n        )\n        if self.is_ddp_backend:\n            if version.parse(pl.__version__) < version.parse(\"1.5.0\"):  # type: ignore\n                raise ValueError(\"PyTorch Lightning>=1.5.0 is required in DDP.\")\n            if not (\n                isinstance(self._trial.study._storage, _CachedStorage)\n                and isinstance(self._trial.study._storage._backend, RDBStorage)\n            ):\n                raise ValueError(\n                    \"optuna.integration.PyTorchLightningPruningCallback\"\n                    \" supports only optuna.storages.RDBStorage in DDP.\"\n                )\n\n    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:\n        # When the trainer calls `on_validation_end` for sanity check,\n        # do not call `trial.report` to avoid calling `trial.report` multiple times\n        # at epoch 0. The related page is\n        # https://github.com/PyTorchLightning/pytorch-lightning/issues/1391.\n        if trainer.sanity_checking:\n            return\n\n        epoch = pl_module.current_epoch\n\n        current_score = trainer.callback_metrics.get(self.monitor)\n        if current_score is None:\n            message = (\n                \"The metric '{}' is not in the evaluation logs for pruning. \"\n                \"Please make sure you set the correct metric name.\".format(self.monitor)\n            )\n            warnings.warn(message)\n            return\n\n        should_stop = False\n        if trainer.is_global_zero:\n            self._trial.report(current_score.item(), step=epoch)\n            should_stop = self._trial.should_prune()\n        should_stop = trainer.training_type_plugin.broadcast(should_stop)\n        if not should_stop:\n            return\n\n        if not self.is_ddp_backend:\n            message = \"Trial was pruned at epoch {}.\".format(epoch)\n            raise optuna.TrialPruned(message)\n        else:\n            # Stop every DDP process if global rank 0 process decides to stop.\n            trainer.should_stop = True\n            if trainer.is_global_zero:\n                self._trial.storage.set_trial_system_attr(self._trial._trial_id, _PRUNED_KEY, True)\n                self._trial.storage.set_trial_system_attr(self._trial._trial_id, _EPOCH_KEY, epoch)\n\n    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:\n        if not self.is_ddp_backend:\n            return\n\n        # Because on_validation_end is executed in spawned processes,\n        # _trial.report is necessary to update the memory in main process, not to update the RDB.\n        _trial_id = self._trial._trial_id\n        _study = self._trial.study\n        _trial = _study._storage._backend.get_trial(_trial_id)  # type: ignore\n        _trial_system_attrs = _study._storage.get_trial_system_attrs(_trial_id)\n        is_pruned = _trial_system_attrs.get(_PRUNED_KEY)\n        epoch = _trial_system_attrs.get(_EPOCH_KEY)\n        intermediate_values = _trial.intermediate_values\n        for step, value in intermediate_values.items():\n            self._trial.report(value, step=step)\n\n        if is_pruned:\n            message = \"Trial was pruned at epoch {}.\".format(epoch)\n            raise optuna.TrialPruned(message)\n", "path": "optuna/integration/pytorch_lightning.py"}], "after_files": [{"content": "import warnings\n\nfrom packaging import version\n\nimport optuna\nfrom optuna.storages._cached_storage import _CachedStorage\nfrom optuna.storages._rdb.storage import RDBStorage\n\n\n# Define key names of `Trial.system_attrs`.\n_PRUNED_KEY = \"ddp_pl:pruned\"\n_EPOCH_KEY = \"ddp_pl:epoch\"\n\n\nwith optuna._imports.try_import() as _imports:\n    import pytorch_lightning as pl\n    from pytorch_lightning import LightningModule\n    from pytorch_lightning import Trainer\n    from pytorch_lightning.callbacks import Callback\n\nif not _imports.is_successful():\n    Callback = object  # type: ignore  # NOQA\n    LightningModule = object  # type: ignore  # NOQA\n    Trainer = object  # type: ignore  # NOQA\n\n\nclass PyTorchLightningPruningCallback(Callback):\n    \"\"\"PyTorch Lightning callback to prune unpromising trials.\n\n    See `the example <https://github.com/optuna/optuna-examples/blob/\n    main/pytorch/pytorch_lightning_simple.py>`__\n    if you want to add a pruning callback which observes accuracy.\n\n    Args:\n        trial:\n            A :class:`~optuna.trial.Trial` corresponding to the current evaluation of the\n            objective function.\n        monitor:\n            An evaluation metric for pruning, e.g., ``val_loss`` or\n            ``val_acc``. The metrics are obtained from the returned dictionaries from e.g.\n            ``pytorch_lightning.LightningModule.training_step`` or\n            ``pytorch_lightning.LightningModule.validation_epoch_end`` and the names thus depend on\n            how this dictionary is formatted.\n\n    .. note::\n        For the distributed data parallel training, the version of PyTorchLightning needs to be\n        higher than or equal to v1.5.0. In addition, :class:`~optuna.study.Study` should be\n        instantiated with RDB storage.\n    \"\"\"\n\n    def __init__(self, trial: optuna.trial.Trial, monitor: str) -> None:\n        _imports.check()\n        super().__init__()\n\n        self._trial = trial\n        self.monitor = monitor\n        self.is_ddp_backend = False\n\n    def on_fit_start(self, trainer: Trainer, pl_module: \"pl.LightningModule\") -> None:\n        self.is_ddp_backend = trainer._accelerator_connector.is_distributed\n        if self.is_ddp_backend:\n            if version.parse(pl.__version__) < version.parse(\"1.6.0\"):  # type: ignore\n                raise ValueError(\"PyTorch Lightning>=1.6.0 is required in DDP.\")\n            if not (\n                isinstance(self._trial.study._storage, _CachedStorage)\n                and isinstance(self._trial.study._storage._backend, RDBStorage)\n            ):\n                raise ValueError(\n                    \"optuna.integration.PyTorchLightningPruningCallback\"\n                    \" supports only optuna.storages.RDBStorage in DDP.\"\n                )\n\n    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:\n        # When the trainer calls `on_validation_end` for sanity check,\n        # do not call `trial.report` to avoid calling `trial.report` multiple times\n        # at epoch 0. The related page is\n        # https://github.com/PyTorchLightning/pytorch-lightning/issues/1391.\n        if trainer.sanity_checking:\n            return\n\n        epoch = pl_module.current_epoch\n\n        current_score = trainer.callback_metrics.get(self.monitor)\n        if current_score is None:\n            message = (\n                \"The metric '{}' is not in the evaluation logs for pruning. \"\n                \"Please make sure you set the correct metric name.\".format(self.monitor)\n            )\n            warnings.warn(message)\n            return\n\n        should_stop = False\n        if trainer.is_global_zero:\n            self._trial.report(current_score.item(), step=epoch)\n            should_stop = self._trial.should_prune()\n        should_stop = trainer.strategy.broadcast(should_stop)\n        if not should_stop:\n            return\n\n        if not self.is_ddp_backend:\n            message = \"Trial was pruned at epoch {}.\".format(epoch)\n            raise optuna.TrialPruned(message)\n        else:\n            # Stop every DDP process if global rank 0 process decides to stop.\n            trainer.should_stop = True\n            if trainer.is_global_zero:\n                self._trial.storage.set_trial_system_attr(self._trial._trial_id, _PRUNED_KEY, True)\n                self._trial.storage.set_trial_system_attr(self._trial._trial_id, _EPOCH_KEY, epoch)\n\n    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:\n        if not self.is_ddp_backend:\n            return\n\n        # Because on_validation_end is executed in spawned processes,\n        # _trial.report is necessary to update the memory in main process, not to update the RDB.\n        _trial_id = self._trial._trial_id\n        _study = self._trial.study\n        _trial = _study._storage._backend.get_trial(_trial_id)  # type: ignore\n        _trial_system_attrs = _study._storage.get_trial_system_attrs(_trial_id)\n        is_pruned = _trial_system_attrs.get(_PRUNED_KEY)\n        epoch = _trial_system_attrs.get(_EPOCH_KEY)\n        intermediate_values = _trial.intermediate_values\n        for step, value in intermediate_values.items():\n            self._trial.report(value, step=step)\n\n        if is_pruned:\n            message = \"Trial was pruned at epoch {}.\".format(epoch)\n            raise optuna.TrialPruned(message)\n", "path": "optuna/integration/pytorch_lightning.py"}]}