ckpt: use checkpoint name as version (#92)

Author: Borda

docs/source/conf.py

@@ -274,7 +274,7 @@ def _convert_markdown(path_in: str, path_out: str) -> None:
     "numpy": ("https://numpy.org/doc/stable/", None),
 }
 
-# -- Options for todo extension ----------------------------------------------
+# -- Options for to-do extension ----------------------------------------------
 
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = True

src/litmodels/integrations/checkpoints.py

@@ -1,4 +1,5 @@
 import inspect
+import os.path
 import queue
 import threading
 from abc import ABC
@@ -150,12 +151,16 @@ def __init__(self, model_name: Optional[str]) -> None:
 
     @rank_zero_only
     def _upload_model(self, filepath: str, metadata: Optional[dict] = None) -> None:
-        # todo: use filename as version but need to validate that such version does not exists yet
         if not self.model_registry:
             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."
             )
+        model_registry = self.model_registry
+        if os.path.isfile(filepath):
+            # parse the file name as version
+            version, _ = os.path.splitext(os.path.basename(filepath))
+            model_registry += f":{version}"
         if not metadata:
             metadata = {}
         # Add the integration name to the metadata
@@ -164,7 +169,7 @@ def _upload_model(self, filepath: str, metadata: Optional[dict] = None) -> None:
         ckpt_class = mro[abc_index - 1]
         metadata.update({"litModels_integration": ckpt_class.__name__})
         # Add to queue instead of uploading directly
-        get_model_manager().queue_upload(registry_name=self.model_registry, filepath=filepath, metadata=metadata)
+        get_model_manager().queue_upload(registry_name=model_registry, filepath=filepath, metadata=metadata)
 
     @rank_zero_only
     def _remove_model(self, trainer: "pl.Trainer", filepath: str) -> None:

tests/integrations/test_checkpoints.py

@@ -78,14 +78,16 @@ def test_lightning_checkpoint_callback(mock_auth, mock_upload_model, monkeypatch
     trainer.fit(BoringModel())
 
     assert mock_auth.call_count == 1
-    expected_call = mock.call(
-        name=f"{expected_org}/{expected_teamspace}/{expected_model}",
-        path=mock.ANY,
-        progress_bar=True,
-        cloud_account=None,
-        metadata={"litModels_integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__},
-    )
-    assert mock_upload_model.call_args_list == [expected_call] * 2
+    assert mock_upload_model.call_args_list == [
+        mock.call(
+            name=f"{expected_org}/{expected_teamspace}/{expected_model}:{v}",
+            path=mock.ANY,
+            progress_bar=True,
+            cloud_account=None,
+            metadata={"litModels_integration": LitModelCheckpoint.__name__, "litModels": litmodels.__version__},
+        )
+        for v in ("epoch=0-step=64", "epoch=1-step=128")
+    ]
 
     # Verify paths match the expected pattern
     for call_args in mock_upload_model.call_args_list: