[Train] Update PyTorch Lightning import path

python/ray/train/lightning/__init__.py

@@ -1,11 +1,14 @@
 # isort: off
 try:
-    import pytorch_lightning  # noqa: F401
+    import lightning  # noqa: F401
 except ModuleNotFoundError:
-    raise ModuleNotFoundError(
-        "PyTorch Lightning isn't installed. To install PyTorch Lightning, "
-        "please run 'pip install pytorch-lightning'"
-    )
+    try:
+        import pytorch_lightning  # noqa: F401
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "PyTorch Lightning isn't installed. To install PyTorch Lightning, "
+            "please run 'pip install lightning'"
+        )
 # isort: on
 
 from ray.train.lightning.lightning_trainer import (

python/ray/train/lightning/_lightning_utils.py

@@ -20,19 +20,21 @@
 from typing import Any, Dict, Optional
 from torch.utils.data import IterableDataset, DataLoader
 
-import pytorch_lightning as pl
-from pytorch_lightning.callbacks import ModelCheckpoint, Callback
-from pytorch_lightning.plugins.environments import LightningEnvironment
-from pytorch_lightning.strategies import DDPStrategy, DeepSpeedStrategy
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
 
 _LIGHTNING_GREATER_EQUAL_2_0 = Version(pl.__version__) >= Version("2.0.0")
 _TORCH_GREATER_EQUAL_1_12 = Version(torch.__version__) >= Version("1.12.0")
 _TORCH_FSDP_AVAILABLE = _TORCH_GREATER_EQUAL_1_12 and torch.distributed.is_available()
 
 if _LIGHTNING_GREATER_EQUAL_2_0:
-    from pytorch_lightning.strategies import FSDPStrategy
+    from lightning.pytorch.strategies import FSDPStrategy
+    from lightning.pytorch.plugins.environments import LightningEnvironment
 else:
     from pytorch_lightning.strategies import DDPFullyShardedStrategy as FSDPStrategy
+    from pytorch_lightning.plugins.environments import LightningEnvironment
 
 if _TORCH_FSDP_AVAILABLE:
     from torch.distributed.fsdp import (
@@ -57,7 +59,7 @@ def get_worker_root_device():
 
 
 @PublicAPI(stability="beta")
-class RayDDPStrategy(DDPStrategy):
+class RayDDPStrategy(pl.strategies.DDPStrategy):
     """Subclass of DDPStrategy to ensure compatibility with Ray orchestration.
 
     For a full list of initialization arguments, please refer to:
@@ -124,7 +126,7 @@ def lightning_module_state_dict(self) -> Dict[str, Any]:
 
 
 @PublicAPI(stability="beta")
-class RayDeepSpeedStrategy(DeepSpeedStrategy):
+class RayDeepSpeedStrategy(pl.strategies.DeepSpeedStrategy):
     """Subclass of DeepSpeedStrategy to ensure compatibility with Ray orchestration.
 
     For a full list of initialization arguments, please refer to:
@@ -210,7 +212,7 @@ def prepare_trainer(trainer: pl.Trainer) -> pl.Trainer:
 
 
 @PublicAPI(stability="beta")
-class RayTrainReportCallback(Callback):
+class RayTrainReportCallback(pl.callbacks.Callback):
     """A simple callback that reports checkpoints to Ray on train epoch end."""
 
     def __init__(self) -> None:
@@ -288,7 +290,7 @@ def _val_dataloader() -> DataLoader:
             self.val_dataloader = _val_dataloader
 
 
-class RayModelCheckpoint(ModelCheckpoint):
+class RayModelCheckpoint(pl.callbacks.ModelCheckpoint):
     """
     AIR customized ModelCheckpoint callback.
 
@@ -306,7 +308,7 @@ def setup(
         super().setup(trainer, pl_module, stage)
         self.is_checkpoint_step = False
 
-        if isinstance(trainer.strategy, DeepSpeedStrategy):
+        if isinstance(trainer.strategy, pl.strategies.DeepSpeedStrategy):
             # For DeepSpeed, each node has a unique set of param and optimizer states,
             # so the local rank 0 workers report the checkpoint shards for all workers
             # on their node.

python/ray/train/lightning/lightning_predictor.py

@@ -5,7 +5,12 @@
 from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
 from ray.train.torch.torch_predictor import TorchPredictor
 from ray.util.annotations import PublicAPI
-import pytorch_lightning as pl
+
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
+
 
 logger = logging.getLogger(__name__)
 

python/ray/train/lightning/lightning_trainer.py

@@ -1,5 +1,10 @@
 import os
-import pytorch_lightning as pl
+
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
+
 
 from copy import copy
 from inspect import isclass

python/ray/train/tests/lightning_test_utils.py

@@ -1,7 +1,11 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import pytorch_lightning as pl
+
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
 
 from torch.utils.data import DataLoader
 from torchmetrics import Accuracy

python/ray/train/tests/test_lightning_checkpoint.py

@@ -1,4 +1,8 @@
-import pytorch_lightning as pl
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
+
 import torch
 import torch.nn as nn
 import tempfile

python/ray/train/tests/test_lightning_predictor.py

@@ -3,9 +3,14 @@
 import torch
 
 import numpy as np
-import pytorch_lightning as pl
 from torch.utils.data import DataLoader
 
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
+
+
 from ray.air.constants import MAX_REPR_LENGTH, MODEL_KEY
 from ray.train.tests.conftest import *  # noqa
 from ray.train.lightning import LightningCheckpoint, LightningPredictor

python/ray/train/tests/test_lightning_trainer_restore.py

@@ -2,8 +2,11 @@
 import numpy as np
 from pathlib import Path
 import pytest
-import pytorch_lightning as pl
-from pytorch_lightning.callbacks import ModelCheckpoint
+
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
 
 import ray
 from ray.train import RunConfig, CheckpointConfig
@@ -147,7 +150,7 @@ def test_air_trainer_restore(
 
     if resume_from_ckpt_path:
         ckpt_dir = f"{tmpdir}/ckpts"
-        callback = ModelCheckpoint(dirpath=ckpt_dir, save_last=True)
+        callback = pl.callbacks.ModelCheckpoint(dirpath=ckpt_dir, save_last=True)
         pl_trainer = pl.Trainer(
             max_epochs=init_epoch, accelerator="cpu", callbacks=[callback]
         )

python/ray/train/tests/test_torch_lightning_train.py

@@ -17,7 +17,11 @@
     LinearModule,
     DummyDataModule,
 )
-import pytorch_lightning as pl
+
+try:
+    import lightning.pytorch as pl
+except ModuleNotFoundError:
+    import pytorch_lightning as pl
 
 
 @pytest.fixture