simple lightning checkpoint callback (#39)

* checkpoint callback
* linter
* test

Author: Borda

.github/workflows/ci-testing.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-latest", "macOS-latest", "windows-latest"]
+        os: ["ubuntu-22.04", "macOS-13", "windows-2022"]
         python-version: ["3.9", "3.12"]
         requires: ["latest"]
         include:
@@ -44,7 +44,7 @@ jobs:
       - name: Install package & dependencies
         run: |
           pip --version
-          pip install -e '.[test]' -U -q --find-links $TORCH_URL
+          pip install -e '.[test,extra]' -U -q --find-links $TORCH_URL
           pip list
 
       - name: Tests
@@ -60,7 +60,7 @@ jobs:
         uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-          file: ./coverage.xml
+          files: ./coverage.xml
           flags: unittests
           env_vars: OS,PYTHON
           name: codecov-umbrella

_requirements/extra.txt

@@ -0,0 +1 @@
+lightning >= 2.0.0

examples/train-model-with-lightning-callback.py

@@ -4,22 +4,13 @@
 
 import torch.utils.data as data
 import torchvision as tv
-from lightning import Callback, Trainer
-from litmodels import upload_model
+from lightning import Trainer
+from litmodels.integrations.lightning_checkpoint import LitModelCheckpoint
 from sample_model import LitAutoEncoder
 
 # Define the model name - this should be unique to your model
 # The format is <organization>/<teamspace>/<model-name>
-MY_MODEL_NAME = "jirka/kaggle/lit-auto-encoder-callback"
-
-
-class UploadModelCallback(Callback):
-    def on_train_epoch_end(self, trainer, pl_module):
-        # Get the best model path from the checkpoint callback
-        best_model_path = trainer.checkpoint_callback.best_model_path
-        if best_model_path:
-            print(f"Uploading model: {best_model_path}")
-            upload_model(model=best_model_path, name=MY_MODEL_NAME)
+MY_MODEL_NAME = "lightning-ai/jirka/lit-auto-encoder-callback"
 
 
 if __name__ == "__main__":
@@ -30,7 +21,7 @@ def on_train_epoch_end(self, trainer, pl_module):
 
     trainer = Trainer(
         max_epochs=2,
-        callbacks=[UploadModelCallback()],
+        callbacks=LitModelCheckpoint(model_name=MY_MODEL_NAME),
     )
     trainer.fit(
         autoencoder,

src/litmodels/integrations/__init__.py

@@ -0,0 +1 @@
+"""Integrations with training frameworks like PyTorch Lightning, TensorFlow, and others."""

src/litmodels/integrations/imports.py

@@ -0,0 +1,4 @@
+from lightning_utilities import module_available
+
+_LIGHTNING_AVAILABLE = module_available("lightning")
+_PYTORCHLIGHTNING_AVAILABLE = module_available("pytorch_lightning")

src/litmodels/integrations/lightning_checkpoint.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+from litmodels import upload_model
+from litmodels.integrations.imports import _LIGHTNING_AVAILABLE, _PYTORCHLIGHTNING_AVAILABLE
+
+if _LIGHTNING_AVAILABLE:
+    from lightning.pytorch import Trainer
+    from lightning.pytorch.callbacks import ModelCheckpoint
+elif _PYTORCHLIGHTNING_AVAILABLE:
+    from pytorch_lightning.callbacks import ModelCheckpoint, Trainer
+else:
+    raise ModuleNotFoundError("No module named 'lightning' or 'pytorch_lightning'")
+
+
+class LitModelCheckpoint(ModelCheckpoint):
+    """Lightning ModelCheckpoint with LitModel support.
+
+    Args:
+        model_name: Name of the model to upload. Must be in the format 'organization/teamspace/modelname'
+            where entity is either your username or the name of an organization you are part of.
+        args: Additional arguments to pass to the parent class.
+        kwargs: Additional keyword arguments to pass to the parent class.
+
+    """
+
+    def __init__(self, model_name: str, *args: Any, **kwargs: Any) -> None:
+        """Initialize the LitModelCheckpoint."""
+        super().__init__(*args, **kwargs)
+        self.model_name = model_name
+
+    def _save_checkpoint(self, trainer: Trainer, filepath: str) -> None:
+        super()._save_checkpoint(trainer, filepath)
+        # 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
+        upload_model(name=self.model_name, model=filepath)

tests/integrations/test_lightning.py

@@ -0,0 +1,35 @@
+import re
+from unittest import mock
+
+from litmodels.integrations.imports import _LIGHTNING_AVAILABLE, _PYTORCHLIGHTNING_AVAILABLE
+from litmodels.integrations.lightning_checkpoint import LitModelCheckpoint
+
+if _LIGHTNING_AVAILABLE:
+    from lightning import Trainer
+    from lightning.pytorch.demos.boring_classes import BoringModel
+elif _PYTORCHLIGHTNING_AVAILABLE:
+    from pytorch_lightning import Trainer
+    from pytorch_lightning.demos.boring_classes import BoringModel
+
+
+@mock.patch("litmodels.io.cloud.sdk_upload_model")
+def test_lightning_checkpoint_callback(mock_upload_model, tmp_path):
+    mock_upload_model.return_value.name = "org-name/teamspace/model-name"
+
+    trainer = Trainer(
+        max_epochs=2,
+        callbacks=LitModelCheckpoint(model_name="org-name/teamspace/model-name"),
+    )
+    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_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),
+    ]
+
+    # Verify paths match the expected pattern
+    for call_args in mock_upload_model.call_args_list:
+        path = call_args[1]["path"]
+        assert re.match(r".*[/\\]lightning_logs[/\\]version_\d+[/\\]checkpoints[/\\]epoch=\d+-step=\d+\.ckpt$", path)

tests/test_io_cloud.py

@@ -27,18 +27,18 @@ def test_wrong_model_name(name):
     ],
 )
 @mock.patch("litmodels.io.cloud.sdk_upload_model")
-def test_upload_model(mock_upload_model, tmpdir, model, model_path, verbose):
+def test_upload_model(mock_upload_model, tmp_path, model, model_path, verbose):
     mock_upload_model.return_value.name = "org-name/teamspace/model-name"
 
     # The lit-logger function is just a wrapper around the SDK function
     upload_model(
         model=model,
         name="org-name/teamspace/model-name",
         cloud_account="cluster_id",
-        staging_dir=tmpdir,
+        staging_dir=str(tmp_path),
         verbose=verbose,
     )
-    expected_path = model_path % str(tmpdir) if "%" in model_path else model_path
+    expected_path = model_path % str(tmp_path) if "%" in model_path else model_path
     mock_upload_model.assert_called_once_with(
         path=expected_path,
         name="org-name/teamspace/model-name",