update checkpointing (#20653)

Author: Borda

src/lightning/pytorch/trainer/__init__.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""""""
 
 from lightning.fabric.utilities.seed import seed_everything
 from lightning.pytorch.trainer.trainer import Trainer

src/lightning/pytorch/trainer/connectors/callback_connector.py

@@ -11,13 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import inspect
 import logging
 import os
 from collections.abc import Sequence
 from datetime import timedelta
 from typing import Optional, Union
 
+from lightning_utilities import module_available
+
 import lightning.pytorch as pl
 from lightning.fabric.utilities.registry import _load_external_callbacks
 from lightning.pytorch.callbacks import (
@@ -91,7 +93,24 @@ def _configure_checkpoint_callbacks(self, enable_checkpointing: bool) -> None:
                     " but found `ModelCheckpoint` in callbacks list."
                 )
         elif enable_checkpointing:
-            self.trainer.callbacks.append(ModelCheckpoint())
+            if module_available("litmodels") and self.trainer._model_registry:
+                trainer_source = inspect.getmodule(self.trainer)
+                if trainer_source is None or not isinstance(trainer_source.__package__, str):
+                    raise RuntimeError("Unable to determine the source of the trainer.")
+                # this need to imported based on the actual package lightning/pytorch_lightning
+                if "pytorch_lightning" in trainer_source.__package__:
+                    from litmodels.integrations.checkpoints import PytorchLightningModelCheckpoint as LitModelCheckpoint
+                else:
+                    from litmodels.integrations.checkpoints import LightningModelCheckpoint as LitModelCheckpoint
+
+                model_checkpoint = LitModelCheckpoint(model_name=self.trainer._model_registry)
+            else:
+                rank_zero_info(
+                    "You are using the plain ModelCheckpoint callback."
+                    " Consider using LitModelCheckpoint which with seamless uploading to Model registry."
+                )
+                model_checkpoint = ModelCheckpoint()
+            self.trainer.callbacks.append(model_checkpoint)
 
     def _configure_model_summary_callback(self, enable_model_summary: bool) -> None:
         if not enable_model_summary:

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -19,6 +19,7 @@
 import torch
 from fsspec.core import url_to_fs
 from fsspec.implementations.local import LocalFileSystem
+from lightning_utilities import module_available
 from torch import Tensor
 
 import lightning.pytorch as pl
@@ -33,6 +34,10 @@
 from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE
 from lightning.pytorch.utilities.migration import pl_legacy_patch
 from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
+from lightning.pytorch.utilities.model_registry import (
+    _is_registry,
+    find_model_local_ckpt_path,
+)
 from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
 
 log = logging.getLogger(__name__)
@@ -48,8 +53,7 @@ def __init__(self, trainer: "pl.Trainer") -> None:
 
     @property
     def _hpc_resume_path(self) -> Optional[str]:
-        dir_path_hpc = self.trainer.default_root_dir
-        dir_path_hpc = str(dir_path_hpc)
+        dir_path_hpc = str(self.trainer.default_root_dir)
         fs, path = url_to_fs(dir_path_hpc)
         if not _is_dir(fs, path):
             return None
@@ -194,10 +198,17 @@ def _parse_ckpt_path(
             if not self._hpc_resume_path:
                 raise ValueError(
                     f'`.{fn}(ckpt_path="hpc")` is set but no HPC checkpoint was found.'
-                    " Please pass an exact checkpoint path to `.{fn}(ckpt_path=...)`"
+                    f" Please pass an exact checkpoint path to `.{fn}(ckpt_path=...)`"
                 )
             ckpt_path = self._hpc_resume_path
 
+        elif _is_registry(ckpt_path) and module_available("litmodels"):
+            ckpt_path = find_model_local_ckpt_path(
+                ckpt_path,
+                default_model_registry=self.trainer._model_registry,
+                default_root_dir=self.trainer.default_root_dir,
+            )
+
         if not ckpt_path:
             raise ValueError(
                 f"`.{fn}()` found no path for the best weights: {ckpt_path!r}. Please"

src/lightning/pytorch/trainer/trainer.py

@@ -30,6 +30,7 @@
 from weakref import proxy
 
 import torch
+from lightning_utilities import module_available
 from torch.optim import Optimizer
 
 import lightning.pytorch as pl
@@ -70,6 +71,7 @@
 from lightning.pytorch.utilities.compile import _maybe_unwrap_optimized, _verify_strategy_supports_compile
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.model_registry import _is_registry, download_model_from_registry
 from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
 from lightning.pytorch.utilities.seed import isolate_rng
 from lightning.pytorch.utilities.types import (
@@ -128,6 +130,7 @@ def __init__(
         sync_batchnorm: bool = False,
         reload_dataloaders_every_n_epochs: int = 0,
         default_root_dir: Optional[_PATH] = None,
+        model_registry: Optional[str] = None,
     ) -> None:
         r"""Customize every aspect of training via flags.
 
@@ -290,6 +293,8 @@ def __init__(
                 Default: ``os.getcwd()``.
                 Can be remote file paths such as `s3://mybucket/path` or 'hdfs://path/'
 
+            model_registry: The name of the model being uploaded to Model hub.
+
         Raises:
             TypeError:
                 If ``gradient_clip_val`` is not an int or float.
@@ -304,6 +309,9 @@ def __init__(
         if default_root_dir is not None:
             default_root_dir = os.fspath(default_root_dir)
 
+        # remove version if accidentally passed
+        self._model_registry = model_registry.split(":")[0] if model_registry else None
+
         self.barebones = barebones
         if barebones:
             # opt-outs
@@ -519,7 +527,20 @@ def fit(
                 the :class:`~lightning.pytorch.core.hooks.DataHooks.train_dataloader` hook.
 
             ckpt_path: Path/URL of the checkpoint from which training is resumed. Could also be one of two special
-                keywords ``"last"`` and ``"hpc"``. If there is no checkpoint file at the path, an exception is raised.
+                keywords ``"last"``, ``"hpc"`` and ``"registry"``.
+                Otherwise, if there is no checkpoint file at the path, an exception is raised.
+
+                    - best: the best model checkpoint from the previous ``trainer.fit`` call will be loaded
+                    - last: the last model checkpoint from the previous ``trainer.fit`` call will be loaded
+                    - registry: the model will be downloaded from the Lightning Model Registry with following notations:
+
+                        - ``'registry'``: uses the latest/default version of default model set
+                          with ``Tainer(..., model_registry="my-model")``
+                        - ``'registry:model-name'``: uses the latest/default version of this model `model-name`
+                        - ``'registry:model-name:version:v2'``: uses the specific version 'v2' of the model `model-name`
+                        - ``'registry:version:v2'``: uses the default model set
+                          with ``Tainer(..., model_registry="my-model")`` and version 'v2'
+
 
         Raises:
             TypeError:
@@ -567,6 +588,8 @@ def _fit_impl(
         )
 
         assert self.state.fn is not None
+        if _is_registry(ckpt_path) and module_available("litmodels"):
+            download_model_from_registry(ckpt_path, self)
         ckpt_path = self._checkpoint_connector._select_ckpt_path(
             self.state.fn,
             ckpt_path,
@@ -596,8 +619,8 @@ def validate(
                 Alternatively, a :class:`~lightning.pytorch.core.datamodule.LightningDataModule` that defines
                 the :class:`~lightning.pytorch.core.hooks.DataHooks.val_dataloader` hook.
 
-            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"`` or path to the checkpoint you wish to validate.
-                If ``None`` and the model instance was passed, use the current weights.
+            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"``, ``"registry"`` or path to the checkpoint you wish
+                to validate. If ``None`` and the model instance was passed, use the current weights.
                 Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
                 if a checkpoint callback is configured.
 
@@ -675,6 +698,8 @@ def _validate_impl(
         self._data_connector.attach_data(model, val_dataloaders=dataloaders, datamodule=datamodule)
 
         assert self.state.fn is not None
+        if _is_registry(ckpt_path) and module_available("litmodels"):
+            download_model_from_registry(ckpt_path, self)
         ckpt_path = self._checkpoint_connector._select_ckpt_path(
             self.state.fn, ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
         )
@@ -705,8 +730,8 @@ def test(
                 Alternatively, a :class:`~lightning.pytorch.core.datamodule.LightningDataModule` that defines
                 the :class:`~lightning.pytorch.core.hooks.DataHooks.test_dataloader` hook.
 
-            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"`` or path to the checkpoint you wish to test.
-                If ``None`` and the model instance was passed, use the current weights.
+            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"``, ``"registry"`` or path to the checkpoint you wish
+                to test. If ``None`` and the model instance was passed, use the current weights.
                 Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
                 if a checkpoint callback is configured.
 
@@ -784,6 +809,8 @@ def _test_impl(
         self._data_connector.attach_data(model, test_dataloaders=dataloaders, datamodule=datamodule)
 
         assert self.state.fn is not None
+        if _is_registry(ckpt_path) and module_available("litmodels"):
+            download_model_from_registry(ckpt_path, self)
         ckpt_path = self._checkpoint_connector._select_ckpt_path(
             self.state.fn, ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
         )
@@ -820,8 +847,8 @@ def predict(
             return_predictions: Whether to return predictions.
                 ``True`` by default except when an accelerator that spawns processes is used (not supported).
 
-            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"`` or path to the checkpoint you wish to predict.
-                If ``None`` and the model instance was passed, use the current weights.
+            ckpt_path: Either ``"best"``, ``"last"``, ``"hpc"``, ``"registry"`` or path to the checkpoint you wish
+                to predict. If ``None`` and the model instance was passed, use the current weights.
                 Otherwise, the best model checkpoint from the previous ``trainer.fit`` call will be loaded
                 if a checkpoint callback is configured.
 
@@ -893,6 +920,8 @@ def _predict_impl(
         self._data_connector.attach_data(model, predict_dataloaders=dataloaders, datamodule=datamodule)
 
         assert self.state.fn is not None
+        if _is_registry(ckpt_path) and module_available("litmodels"):
+            download_model_from_registry(ckpt_path, self)
         ckpt_path = self._checkpoint_connector._select_ckpt_path(
             self.state.fn, ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
         )