[AIR] <Part 2> Support metric logging and checkpointing for LightningTrainer

python/ray/train/BUILD

@@ -400,6 +400,14 @@ py_test(
     deps = [":train_lib"]
 )
 
+py_test(
+    name = "test_lightning_checkpoint",
+    size = "medium",
+    srcs = ["tests/test_lightning_checkpoint.py"],
+    tags = ["team:ml", "exclusive", "ray_air", "gpu"],
+    deps = [":train_lib"]
+)
+
 py_test(
     name = "test_lightning_trainer",
     size = "large",

python/ray/train/lightning/__init__.py

@@ -8,10 +8,10 @@
     )
 # isort: on
 
+from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
 from ray.train.lightning.lightning_trainer import (
     LightningTrainer,
     LightningConfigBuilder,
 )
 
-
-__all__ = ["LightningTrainer", "LightningConfigBuilder"]
+__all__ = ["LightningTrainer", "LightningConfigBuilder", "LightningCheckpoint"]

python/ray/train/lightning/_lightning_utils.py

@@ -1,14 +1,20 @@
+import os
 import logging
 import torch
-from typing import Any, Dict, Optional
-
 import pytorch_lightning as pl
+
+from torch import Tensor
+from copy import deepcopy
+from typing import Any, Dict, Optional
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.utilities.types import STEP_OUTPUT
 from pytorch_lightning.strategies import DDPStrategy
 from pytorch_lightning.plugins.environments import LightningEnvironment
 
 import ray
 from ray.air import session
-
+from ray.air.constants import MODEL_KEY
+from ray.train.lightning.lightning_checkpoint import LightningCheckpoint
 from torch.utils.data import IterableDataset, DataLoader
 from ray.data.dataset import DatasetIterator
 
@@ -46,11 +52,13 @@ def node_rank(self) -> int:
         return session.get_node_rank()
 
     def set_world_size(self, size: int) -> None:
-        logger.warning("world_size setter is disabled in AIR LightningTrainer.")
+        if self.global_rank() == 0:
+            logger.warning("world_size setter is disabled in AIR LightningTrainer.")
         pass
 
     def set_global_rank(self, rank: int) -> None:
-        logger.warning("global_rank setter is disabled in AIR LightningTrainer.")
+        if self.global_rank() == 0:
+            logger.warning("global_rank setter is disabled in AIR LightningTrainer.")
         pass
 
     def teardown(self):
@@ -94,3 +102,115 @@ def _val_dataloader() -> DataLoader:
         # setting, we only override this method when `val_dataset` is not `None`.
         if val_dataset:
             self.val_dataloader = _val_dataloader
+
+
+class RayModelCheckpoint(ModelCheckpoint):
+    """
+    AIR customized ModelCheckpoint callback.
+
+    A subclass of ``pytorch_lightning.callbacks.ModelCheckpoint``.
+    This callback function reports the latest metrics to the AIR session and
+    creates an AIR checkpoint whenever a lightning checkpoint is saved.
+    """
+
+    def setup(self, *args, **kwargs) -> None:
+        super().setup(*args, **kwargs)
+        self.last_best_k_models = {}
+        self.last_best_model_path = None
+        self.is_checkpoint_step = False
+
+    def format_checkpoint_name(
+        self,
+        metrics: Dict[str, Tensor],
+        filename: Optional[str] = None,
+        ver: Optional[int] = None,
+    ) -> str:
+        """
+        Change checkpoint files path to align with AIR checkpoint format.
+
+        e.g. './epoch=2-loss=0.12.ckpt' -> './epoch=2-loss=0.12.ckpt/model'
+        """
+        filepath = super().format_checkpoint_name(metrics, filename, ver)
+        return f"{filepath}/{MODEL_KEY}"
+
+    def _session_report(self, trainer: "pl.Trainer", stage: str):
+        """Report latest metrics dict and checkpoint to AIR training session."""
+
+        # Align the frequency of session.report() and checkpointing.
+        if not self.is_checkpoint_step:
+            return
+        self.is_checkpoint_step = False
+
+        # Report latest logged metrics
+        kwargs = {}
+        metrics = {}
+        for k, v in self._monitor_candidates(trainer).items():
+            if k == "_stage":
+                logger.warning(
+                    "'_stage' is a reserved key in AIR report metrics. "
+                    "Original values are overwritten!"
+                )
+                continue
+            if isinstance(v, torch.Tensor):
+                metrics[k] = v.item()
+
+        metrics["_stage"] = stage
+        kwargs["metrics"] = metrics
+
+        filepath = None
+        if self.monitor:
+            # Capture metric-based top-k checkpoint
+            new_checkpoint = self.best_k_models.keys() - self.last_best_k_models.keys()
+            if new_checkpoint:
+                filepath = new_checkpoint.pop()
+        else:
+            # Capture frequency-based checkpoint
+            if self.last_best_model_path != self.best_model_path:
+                filepath = self.best_model_path
+
+        # Report latest saved checkpoint
+        # Note that AIR only takes the checkpoint of rank 0.
+        # Save a dummy checkpoint on the other workers to avoid blocking.
+        if filepath:
+            if trainer.global_rank == 0:
+                kwargs["checkpoint"] = LightningCheckpoint.from_directory(
+                    path=os.path.dirname(filepath)
+                )
+            else:
+                kwargs["checkpoint"] = LightningCheckpoint.from_dict(
+                    {"rank": session.get_world_rank()}
+                )
+
+        self.last_best_k_models = deepcopy(self.best_k_models)
+        self.last_best_model_path = self.best_model_path
+
+        session.report(**kwargs)
+
+    def _save_topk_checkpoint(
+        self, trainer: "pl.Trainer", monitor_candidates: Dict[str, Tensor]
+    ) -> None:
+        self.is_checkpoint_step = True
+        return super()._save_topk_checkpoint(trainer, monitor_candidates)
+
+    def on_train_batch_end(
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)
+        self._session_report(trainer=trainer, stage="train_batch_end")
+
+    def on_train_epoch_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        super().on_train_epoch_end(trainer, pl_module)
+        self._session_report(trainer=trainer, stage="train_epoch_end")
+
+    def on_validation_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> None:
+        super().on_validation_end(trainer, pl_module)
+        self._session_report(trainer=trainer, stage="validation_end")

python/ray/train/lightning/lightning_checkpoint.py

@@ -0,0 +1,97 @@
+import os
+import logging
+import pytorch_lightning as pl
+import tempfile
+import shutil
+
+from inspect import isclass
+from typing import Optional, Type
+
+from ray.air.constants import MODEL_KEY
+from ray.air._internal.checkpointing import save_preprocessor_to_dir
+from ray.data import Preprocessor
+from ray.train.torch import TorchCheckpoint
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+class LightningCheckpoint(TorchCheckpoint):
+    """A :class:`~ray.air.checkpoint.Checkpoint` with Lightning-specific functionality.
+
+    LightningCheckpoint only support file based checkpoint loading.
+    Create this by calling ``LightningCheckpoint.from_directory(ckpt_dir)``,
+    ``LightningCheckpoint.from_uri(uri)`` or ``LightningCheckpoint.from_path(path)``
+
+    LightningCheckpoint loads file named ``model`` under the specified directory.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._cache_dir = None
+
+    @classmethod
+    def from_path(
+        cls,
+        path: str,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+    ) -> "LightningCheckpoint":
+        """Create a ``ray.air.lightning.LightningCheckpoint`` from a checkpoint path.
+        Args:
+            path: The file path to the PyTorch Lightning checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+        Returns:
+            An :py:class:`LightningCheckpoint` containing the model.
+        Examples:
+            >>> from ray.train.lightning import LightningCheckpoint
+            >>>
+            >>> checkpoint = LightningCheckpoint.from_path("/path/to/checkpoint.ckpt")
+        """
+
+        assert os.path.exists(path), f"Lightning checkpoint {path} doesn't exists!"
+
+        cache_dir = tempfile.mkdtemp()
+        new_checkpoint_path = os.path.join(cache_dir, MODEL_KEY)
+        shutil.copy(path, new_checkpoint_path)
+        if preprocessor:
+            save_preprocessor_to_dir(preprocessor, cache_dir)
+        checkpoint = cls.from_directory(cache_dir)
+        checkpoint._cache_dir = cache_dir
+        return checkpoint
+
+    def get_model(
+        self, model_class: Type[pl.LightningModule], **load_from_checkpoint_kwargs
+    ) -> pl.LightningModule:
+        """Retrieve the model stored in this checkpoint.
+
+        Args:
+            model_class: A subclass of ``pytorch_lightning.LightningModule`` that
+            defines your model and training logic.
+            load_from_checkpoint_kwargs: Arguments to pass into
+            ``model_cls.load_from_checkpoint``
+
+        Returns:
+            pl.LightningModule: An instance of the loaded model.
+        """
+        if not isclass(model_class):
+            raise ValueError(
+                "'lightning_module' must be a class, not a class instance."
+            )
+
+        with self.as_directory() as checkpoint_dir:
+            ckpt_path = os.path.join(checkpoint_dir, MODEL_KEY)
+            if not os.path.exists(ckpt_path):
+                raise RuntimeError(
+                    f"File {ckpt_path} not found under the checkpoint directory."
+                )
+
+            model = model_class.load_from_checkpoint(
+                ckpt_path, **load_from_checkpoint_kwargs
+            )
+        return model
+
+    def __del__(self):
+        if self._cache_dir and os.path.exists(self._cache_dir):
+            shutil.rmtree(self._cache_dir)