enable empty checkpoint name (#72)

* enable default checkpoint name
* tests + mocks
* linting

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Borda

src/litmodels/integrations/checkpoints.py

@@ -1,31 +1,42 @@
-from typing import TYPE_CHECKING, Any
+from abc import ABC
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Optional
 
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_utilities.core.rank_zero import rank_zero_only
+from lightning_utilities.core.rank_zero import rank_zero_only, rank_zero_warn
 
 from litmodels import upload_model
 from litmodels.integrations.imports import _LIGHTNING_AVAILABLE, _PYTORCHLIGHTNING_AVAILABLE
 
 if _LIGHTNING_AVAILABLE:
     from lightning.pytorch.callbacks import ModelCheckpoint as _LightningModelCheckpoint
 
-    if TYPE_CHECKING:
-        from lightning.pytorch import Trainer
-
 
 if _PYTORCHLIGHTNING_AVAILABLE:
     from pytorch_lightning.callbacks import ModelCheckpoint as _PytorchLightningModelCheckpoint
 
-    if TYPE_CHECKING:
-        from pytorch_lightning import Trainer
+
+if TYPE_CHECKING:
+    if _LIGHTNING_AVAILABLE:
+        import lightning.pytorch as pl
+    if _PYTORCHLIGHTNING_AVAILABLE:
+        import pytorch_lightning as pl
 
 
 # Base class to be inherited
-class LitModelCheckpointMixin:
+class LitModelCheckpointMixin(ABC):
     """Mixin class for LitModel checkpoint functionality."""
 
-    def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
+    # mainly ofr mocking reasons
+    _datetime_stamp: str = datetime.now().strftime("%Y%m%d-%H%M")
+    model_name: Optional[str] = None
+
+    def __init__(self, model_name: Optional[str]) -> None:
         """Initialize with model name."""
+        if not model_name:
+            rank_zero_warn(
+                "The model is not defined so we will continue with LightningModule names and timestamp of now"
+            )
         self.model_name = model_name
 
         try:  # authenticate before anything else starts
@@ -38,8 +49,19 @@ def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
     def _upload_model(self, filepath: str) -> None:
         # todo: uploading on background so training does nt stops
         # todo: use filename as version but need to validate that such version does not exists yet
+        if not self.model_name:
+            raise RuntimeError(
+                "Model name is not specified neither updated by `setup` method via Trainer."
+                " Please set the model name before uploading or ensure that `setup` method is called."
+            )
         upload_model(name=self.model_name, model=filepath)
 
+    def _update_model_name(self, pl_model: "pl.LightningModule") -> None:
+        if self.model_name:
+            return
+        # setting the model name as Lightning module with some time hash
+        self.model_name = pl_model.__class__.__name__ + f"_{self._datetime_stamp}"
+
 
 # Create specific implementations
 if _LIGHTNING_AVAILABLE:
@@ -53,15 +75,20 @@ class LightningModelCheckpoint(LitModelCheckpointMixin, _LightningModelCheckpoin
             kwargs: Additional keyword arguments to pass to the parent class.
         """
 
-        def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
+        def __init__(self, *args: Any, model_name: Optional[str] = None, **kwargs: Any) -> None:
             """Initialize the checkpoint with model name and other parameters."""
             _LightningModelCheckpoint.__init__(self, *args, **kwargs)
             LitModelCheckpointMixin.__init__(self, model_name)
 
-        def _save_checkpoint(self, trainer: "Trainer", filepath: str) -> None:
+        def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
+            """Setup the checkpoint callback."""
+            super().setup(trainer, pl_module, stage)
+            self._update_model_name(pl_module)
+
+        def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+            """Extend the save checkpoint method to upload the model."""
             super()._save_checkpoint(trainer, filepath)
-            if trainer.is_global_zero:
-                # Only upload from the main process
+            if trainer.is_global_zero:  # Only upload from the main process
                 self._upload_model(filepath)
 
 
@@ -76,13 +103,18 @@ class PytorchLightningModelCheckpoint(LitModelCheckpointMixin, _PytorchLightning
             kwargs: Additional keyword arguments to pass to the parent class.
         """
 
-        def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
+        def __init__(self, *args: Any, model_name: Optional[str] = None, **kwargs: Any) -> None:
             """Initialize the checkpoint with model name and other parameters."""
             _PytorchLightningModelCheckpoint.__init__(self, *args, **kwargs)
             LitModelCheckpointMixin.__init__(self, model_name)
 
-        def _save_checkpoint(self, trainer: "Trainer", filepath: str) -> None:
+        def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
+            """Setup the checkpoint callback."""
+            super().setup(trainer, pl_module, stage)
+            self._update_model_name(pl_module)
+
+        def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+            """Extend the save checkpoint method to upload the model."""
             super()._save_checkpoint(trainer, filepath)
-            if trainer.is_global_zero:
-                # Only upload from the main process
+            if trainer.is_global_zero:  # Only upload from the main process
                 self._upload_model(filepath)

tests/integrations/__init__.py

@@ -16,3 +16,6 @@
 _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1 = pytest.mark.skipif(
     not _PYTORCHLIGHTNING_GREATER_EQUAL_2_5_1, reason="PyTorch Lightning without integration introduced in 2.5.1"
 )
+
+LIT_ORG = "lightning-ai"
+LIT_TEAMSPACE = "LitModels"

tests/integrations/test_checkpoints.py

@@ -14,9 +14,17 @@
         pytest.param("pytorch_lightning", marks=_SKIP_IF_PYTORCHLIGHTNING_MISSING),
     ],
 )
+@pytest.mark.parametrize("with_model_name", [True, False])
+@mock.patch("litmodels.integrations.checkpoints.LitModelCheckpointMixin._datetime_stamp", return_value="20250102-1213")
+@mock.patch(
+    "lightning_sdk.models._resolve_teamspace",
+    return_value=mock.MagicMock(owner=mock.MagicMock(name="my-org"), name="dream-team"),
+)
 @mock.patch("litmodels.io.cloud.sdk_upload_model")
 @mock.patch("litmodels.integrations.checkpoints.Auth")
-def test_lightning_checkpoint_callback(mock_auth, mock_upload_model, importing, tmp_path):
+def test_lightning_checkpoint_callback(
+    mock_auth, mock_upload_model, mock_resolve_teamspace, mock_datetime_stamp, importing, with_model_name, tmp_path
+):
     if importing == "lightning":
         from lightning import Trainer
         from lightning.pytorch.callbacks import ModelCheckpoint
@@ -31,20 +39,24 @@ def test_lightning_checkpoint_callback(mock_auth, mock_upload_model, importing,
     # Validate inheritance
     assert issubclass(LitModelCheckpoint, ModelCheckpoint)
 
-    mock_upload_model.return_value.name = "org-name/teamspace/model-name"
+    ckpt_args = {"model_name": "org-name/teamspace/model-name"} if with_model_name else {}
+    expected_model_registry = ckpt_args.get("model_name", f"BoringModel_{LitModelCheckpoint._datetime_stamp}")
+    mock_upload_model.return_value.name = expected_model_registry
 
     trainer = Trainer(
         max_epochs=2,
-        callbacks=LitModelCheckpoint(model_name="org-name/teamspace/model-name"),
+        callbacks=LitModelCheckpoint(**ckpt_args),
     )
     trainer.fit(BoringModel())
 
-    # expected_path = model_path % str(tmpdir) if "%" in model_path else model_path
-    assert mock_upload_model.call_count == 2
+    assert mock_auth.call_count == 1
     assert mock_upload_model.call_args_list == [
-        mock.call(name="org-name/teamspace/model-name", path=mock.ANY, progress_bar=True, cloud_account=None),
-        mock.call(name="org-name/teamspace/model-name", path=mock.ANY, progress_bar=True, cloud_account=None),
+        mock.call(name=expected_model_registry, path=mock.ANY, progress_bar=True, cloud_account=None),
+        mock.call(name=expected_model_registry, path=mock.ANY, progress_bar=True, cloud_account=None),
     ]
+    called_name_related_mocks = 0 if with_model_name else 1
+    mock_datetime_stamp.call_count == called_name_related_mocks
+    mock_resolve_teamspace.call_count == called_name_related_mocks
 
     # Verify paths match the expected pattern
     for call_args in mock_upload_model.call_args_list:

tests/integrations/test_cloud.py

@@ -12,10 +12,12 @@
 from litmodels import download_model, upload_model
 from litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
 
-from tests.integrations import _SKIP_IF_LIGHTNING_BELLOW_2_5_1, _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1
-
-LIT_ORG = "lightning-ai"
-LIT_TEAMSPACE = "LitModels"
+from tests.integrations import (
+    _SKIP_IF_LIGHTNING_BELLOW_2_5_1,
+    _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1,
+    LIT_ORG,
+    LIT_TEAMSPACE,
+)
 
 
 def _prepare_variables(test_name: str) -> tuple[Teamspace, str, str]: