[Unified TorchTrainer] Add PyTorch Lightning Trainer Utilities

python/ray/train/BUILD

@@ -454,6 +454,14 @@ py_test(
     deps = [":train_lib"]
 )
 
+py_test(
+    name = "test_torch_lightning_train",
+    size = "large",
+    srcs = ["tests/test_torch_lightning_train.py"],
+    tags = ["team:ml", "exclusive", "ray_air", "gpu", "ptl_v2"],
+    deps = [":train_lib"]
+)
+
 py_test(
     name = "test_minimal",
     size = "small",

python/ray/train/lightning/__init__.py

@@ -8,16 +8,30 @@
     )
 # isort: on
 
-from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
-from ray.train.lightning.lightning_predictor import LightningPredictor
 from ray.train.lightning.lightning_trainer import (
-    LightningTrainer,
     LightningConfigBuilder,
+    LightningTrainer,
+)
+from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
+from ray.train.lightning.lightning_predictor import LightningPredictor
+from ray.train.lightning.lightning_utils import (
+    prepare_trainer,
+    RayDDPStrategy,
+    RayFSDPStrategy,
+    RayDeepSpeedStrategy,
+    RayLightningEnvironment,
+    RayTrainReportCallback,
 )
 
 __all__ = [
     "LightningTrainer",
     "LightningConfigBuilder",
     "LightningCheckpoint",
     "LightningPredictor",
+    "prepare_trainer",
+    "RayDDPStrategy",
+    "RayFSDPStrategy",
+    "RayDeepSpeedStrategy",
+    "RayLightningEnvironment",
+    "RayTrainReportCallback",
 ]

python/ray/train/lightning/lightning_trainer.py

@@ -16,11 +16,11 @@
 from ray.train.torch import TorchTrainer
 from ray.train.torch.config import TorchConfig
 from ray.util import PublicAPI
-from ray.train.lightning._lightning_utils import (
+from ray.train.lightning.lightning_utils import (
     RayDDPStrategy,
     RayFSDPStrategy,
     RayDeepSpeedStrategy,
-    RayEnvironment,
+    RayLightningEnvironment,
     RayDataModule,
     RayModelCheckpoint,
     get_worker_root_device,
@@ -586,7 +586,7 @@ def _lightning_train_loop_per_worker(config):
         for plugin in trainer_config.get("plugins", [])
         if not isinstance(plugin, ClusterEnvironment)
     ]
-    trainer_config["plugins"].append(RayEnvironment())
+    trainer_config["plugins"].append(RayLightningEnvironment())
 
     # Setup ddp strategy for ray orchestration
     if "strategy" in trainer_config:

python/ray/train/lightning/_lightning_utils.py → python/ray/train/lightning/lightning_utils.py

@@ -4,17 +4,20 @@
 from ray.air.constants import MODEL_KEY
 from ray.data.dataset import DataIterator
 from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
+from ray.util import PublicAPI
 
 import logging
 import shutil
 import torch
 import tempfile
+from tempfile import TemporaryDirectory
+from ray.train import Checkpoint
 from packaging.version import Version
 from typing import Any, Dict, Optional
 from torch.utils.data import IterableDataset, DataLoader
 
 import pytorch_lightning as pl
-from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks import ModelCheckpoint, Callback
 from pytorch_lightning.plugins.environments import LightningEnvironment
 from pytorch_lightning.strategies import DDPStrategy, DeepSpeedStrategy
 
@@ -49,6 +52,7 @@ def get_worker_root_device():
         return devices
 
 
+@PublicAPI(stability="alpha")
 class RayDDPStrategy(DDPStrategy):
     """Subclass of DDPStrategy to ensure compatibility with Ray orchestration."""
 
@@ -64,6 +68,7 @@ def distributed_sampler_kwargs(self) -> Dict[str, Any]:
         )
 
 
+@PublicAPI(stability="alpha")
 class RayFSDPStrategy(FSDPStrategy):
     """Subclass of FSDPStrategy to ensure compatibility with Ray orchestration."""
 
@@ -98,19 +103,10 @@ def lightning_module_state_dict(self) -> Dict[str, Any]:
             return super().lightning_module_state_dict()
 
 
+@PublicAPI(stability="alpha")
 class RayDeepSpeedStrategy(DeepSpeedStrategy):
     """Subclass of DeepSpeedStrategy to ensure compatibility with Ray orchestration."""
 
-    def setup_distributed(self):
-        # We have to set the device ids for each node
-        # e.g. CUDA_VISIBLE_DEVICES = 2,3
-        # worker 0: LOCAL_RANK=0, parallel devices = [cuda:0, cuda:1]
-        # worker 1: LOCAL_RANK=1, parallel devices = [cuda:0, cuda:1]
-        self.parallel_devices = [
-            torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())
-        ]
-        super().setup_distributed()
-
     @property
     def root_device(self) -> torch.device:
         return get_worker_root_device()
@@ -123,7 +119,8 @@ def distributed_sampler_kwargs(self) -> Dict[str, Any]:
         )
 
 
-class RayEnvironment(LightningEnvironment):
+@PublicAPI(stability="alpha")
+class RayLightningEnvironment(LightningEnvironment):
     """Setup Lightning DDP training environment for Ray cluster."""
 
     def world_size(self) -> int:
@@ -150,6 +147,58 @@ def teardown(self):
         pass
 
 
+@PublicAPI(stability="alpha")
+def prepare_trainer(trainer: pl.Trainer) -> pl.Trainer:
+    """Prepare the PyTorch Lightning Trainer for distributed execution."""
+
+    # Check strategy class
+    valid_strategy_class = [RayDDPStrategy, RayFSDPStrategy, RayDeepSpeedStrategy]
+
+    if not any(isinstance(trainer.strategy, cls) for cls in valid_strategy_class):
+        raise RuntimeError(
+            f"Invalid strategy class: {type(trainer.strategy)}. To use "
+            "PyTorch Lightning with Ray, the strategy object should be one of "
+            f"{[cls.__name__ for cls in valid_strategy_class]} class "
+            "or its subclass."
+        )
+
+    # Check cluster environment
+    cluster_environment = getattr(trainer.strategy, "cluster_environment", None)
+    if cluster_environment and not isinstance(
+        cluster_environment, RayLightningEnvironment
+    ):
+        raise RuntimeError(
+            "Invalid cluster environment plugin. The expected class is"
+            "`ray.train.lightning.RayLightningEnvironment` "
+            f"but got {type(cluster_environment)}!"
+        )
+
+    return trainer
+
+
+@PublicAPI(stability="alpha")
+class RayTrainReportCallback(Callback):
+    """A simple callback that reports checkpoints to Ray on train epoch end."""
+
+    def on_train_epoch_end(self, trainer, pl_module) -> None:
+        with TemporaryDirectory() as tmpdir:
+            # Fetch metrics
+            metrics = trainer.callback_metrics
+            metrics = {k: v.item() for k, v in metrics.items()}
+
+            # (Optional) Add customized metrics
+            metrics["epoch"] = trainer.current_epoch
+            metrics["steps"] = trainer.global_step
+
+            # Save checkpoint to local
+            ckpt_path = os.path.join(tmpdir, f"ckpt_epoch_{trainer.current_epoch}")
+            trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+            # Report to train session
+            checkpoint = Checkpoint.from_directory(tmpdir)
+            ray.train.report(metrics=metrics, checkpoint=checkpoint)
+
+
 class RayIterableDataset(IterableDataset):
     def __init__(self, dataset: "DataIterator", config: Dict[str, Any]) -> None:
         super().__init__()

python/ray/train/tests/lightning_test_utils.py

@@ -2,8 +2,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import pytorch_lightning as pl
+
 from torch.utils.data import DataLoader
 from torchmetrics import Accuracy
+
 from ray import train