[train][lightning] change remaining pytorch_lightning imports

[liulehui]

Description

1. With version>2.0, the import path of PyTorch Lightning has been updated from import pytorch_lightning to import lightning.pytorch. Although import pytorch_lightning still exists for backward compatibility, we have to modify our docker environment and lightning integration utilities to use the new import path.
2. previously in [Train] Update PyTorch Lightning import path #39841, some of the lightning util/import path are migrated. this PR is to follow up and clean the rest.

Related issues

38200

Additional information

release test: https://buildkite.com/ray-project/release/builds/81617#

[gemini-code-assist]

Code Review

This pull request correctly updates the PyTorch Lightning imports from pytorch_lightning to lightning.pytorch for compatibility with version 2.0 and later. The changes are well-executed across various files, including documentation, examples, and tests. I have a couple of suggestions to enhance maintainability and consistency.

[gemini-code-assist]

For consistency with other examples in this PR, consider using a try-except block to support both old and new PyTorch Lightning import paths. Most other examples (e.g., in doc/source/train/doc_code/checkpoints.py and python/ray/train/examples/) use a compatibility block, while this notebook and tune-vanilla-pytorch-lightning.ipynb perform a hard switch to lightning.pytorch. Using the try-except pattern would make the examples more robust for users with different environments.

# Support both Lightning 2.x (`lightning.pytorch`) and 1.x (`pytorch_lightning`).
try:
    import lightning.pytorch as pl
except ModuleNotFoundError:
    import pytorch_lightning as pl

[gemini-code-assist]

To improve maintainability and avoid code duplication, consider centralizing the PyTorch Lightning import logic. This try-except block for importing pl is repeated in several test files (e.g., python/ray/train/tests/test_train_usage.py).

You could enhance ray.tune.integration.pytorch_lightning to handle this, as it already contains compatibility logic for other Lightning components.

For example, in python/ray/tune/integration/pytorch_lightning.py:

try:
    import lightning.pytorch as pl
    from lightning.pytorch import Callback, LightningModule, Trainer
except ModuleNotFoundError:
    import pytorch_lightning as pl
    from pytorch_lightning import Callback, LightningModule, Trainer

Then, test files could simply import pl from this central location, making future updates easier:

from ray.tune.integration.pytorch_lightning import pl, TuneReportCheckpointCallback

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.