fix: Pickle mixin & update args + test with Prod (#67)

* fix: Pickle mixin
* update args
* test with Prod
* Apply suggestions from code review

---------

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

Author: 3 people

src/litmodels/integrations/mixins.py

@@ -15,66 +15,67 @@ class ModelRegistryMixin(ABC):
     """Mixin for model registry integration."""
 
     def push_to_registry(
-        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
     ) -> None:
         """Push the model to the registry.
 
         Args:
-            model_name: The name of the model. If not use the class name.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model. If not use the class name.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
 
     @classmethod
-    def pull_from_registry(
-        cls, model_name: str, model_version: Optional[str] = None, temp_folder: Optional[str] = None
-    ) -> object:
+    def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folder: Optional[str] = None) -> object:
         """Pull the model from the registry.
 
         Args:
-            model_name: The name of the model.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
 
 
-class PickleRegistryMixin(ABC):
+class PickleRegistryMixin(ModelRegistryMixin):
     """Mixin for pickle registry integration."""
 
     def push_to_registry(
-        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
     ) -> None:
         """Push the model to the registry.
 
         Args:
-            model_name: The name of the model. If not use the class name.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model. If not use the class name.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
-        if model_name is None:
-            model_name = self.__class__.__name__
+        if name is None:
+            name = model_name = self.__class__.__name__
+        elif ":" in name:
+            raise ValueError(f"Invalid model name: '{name}'. It should not contain ':' associated with version.")
+        else:
+            model_name = name.split("/")[-1]
         if temp_folder is None:
-            temp_folder = tempfile.gettempdir()
+            temp_folder = tempfile.mkdtemp()
         pickle_path = Path(temp_folder) / f"{model_name}.pkl"
         with open(pickle_path, "wb") as fp:
             pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
-        model_registry = f"{model_name}:{model_version}" if model_version else model_name
-        upload_model(name=model_registry, model=pickle_path)
+        if version:
+            name = f"{name}:{version}"
+        upload_model(name=name, model=pickle_path)
 
     @classmethod
-    def pull_from_registry(
-        cls, model_name: str, model_version: Optional[str] = None, temp_folder: Optional[str] = None
-    ) -> object:
+    def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folder: Optional[str] = None) -> object:
         """Pull the model from the registry.
 
         Args:
-            model_name: The name of the model.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
         if temp_folder is None:
-            temp_folder = tempfile.gettempdir()
-        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+            temp_folder = tempfile.mkdtemp()
+        model_registry = f"{name}:{version}" if version else name
         files = download_model(name=model_registry, download_dir=temp_folder)
         pkl_files = [f for f in files if f.endswith(".pkl")]
         if not pkl_files:
@@ -89,7 +90,7 @@ def pull_from_registry(
         return obj
 
 
-class PyTorchRegistryMixin(ABC):
+class PyTorchRegistryMixin(ModelRegistryMixin):
     """Mixin for PyTorch model registry integration."""
 
     def __post_init__(self) -> None:
@@ -101,51 +102,51 @@ def __post_init__(self) -> None:
             raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
 
     def push_to_registry(
-        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
     ) -> None:
         """Push the model to the registry.
 
         Args:
-            model_name: The name of the model. If not use the class name.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model. If not use the class name.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
         import torch
 
         if not isinstance(self, torch.nn.Module):
             raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
 
-        if model_name is None:
-            model_name = self.__class__.__name__
+        if name is None:
+            name = self.__class__.__name__
         if temp_folder is None:
-            temp_folder = tempfile.gettempdir()
-        torch_path = Path(temp_folder) / f"{model_name}.pth"
+            temp_folder = tempfile.mkdtemp()
+        torch_path = Path(temp_folder) / f"{name}.pth"
         torch.save(self.state_dict(), torch_path)
         # todo: dump also object creation arguments so we can dump it and load with model for object instantiation
-        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+        model_registry = f"{name}:{version}" if version else name
         upload_model(name=model_registry, model=torch_path)
 
     @classmethod
     def pull_from_registry(
         cls,
-        model_name: str,
-        model_version: Optional[str] = None,
+        name: str,
+        version: Optional[str] = None,
         temp_folder: Optional[str] = None,
         torch_load_kwargs: Optional[dict] = None,
     ) -> "torch.nn.Module":
         """Pull the model from the registry.
 
         Args:
-            model_name: The name of the model.
-            model_version: The version of the model. If None, the latest version is used.
+            name: The name of the model.
+            version: The version of the model. If None, the latest version is used.
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
             torch_load_kwargs: Additional arguments to pass to `torch.load()`.
         """
         import torch
 
         if temp_folder is None:
-            temp_folder = tempfile.gettempdir()
-        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+            temp_folder = tempfile.mkdtemp()
+        model_registry = f"{name}:{version}" if version else name
         files = download_model(name=model_registry, download_dir=temp_folder)
         torch_files = [f for f in files if f.endswith(".pth")]
         if not torch_files:

tests/integrations/test_cloud.py

@@ -9,6 +9,7 @@
 from lightning_sdk.lightning_cloud.rest_client import GridRestClient
 from lightning_sdk.utils.resolve import _resolve_teamspace
 from litmodels import download_model, upload_model
+from litmodels.integrations.mixins import PickleRegistryMixin
 
 from tests.integrations import _SKIP_IF_LIGHTNING_BELLOW_2_5_1, _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1
 
@@ -213,3 +214,28 @@ def test_lightning_checkpoint_ddp(importing, tmp_path):
 
     # CLEANING
     _cleanup_model(teamspace, model_name, expected_num_versions=2)
+
+
+class DummyModel(PickleRegistryMixin):
+    def __init__(self, value):
+        self.value = value
+
+
+@pytest.mark.cloud()
+def test_pickle_mixin_push_and_pull():
+    # model name with random hash
+    teamspace, org_team, model_name = _prepare_variables("pickle_mixin")
+    model_registry = f"{org_team}/{model_name}"
+
+    # Create an instance of DummyModel and call push_to_registry.
+    dummy = DummyModel(42)
+    dummy.push_to_registry(model_registry)
+
+    # Call pull_from_registry and load the DummyModel instance.
+    loaded_dummy = DummyModel.pull_from_registry(model_registry)
+    # Verify that the unpickled instance has the expected value.
+    assert isinstance(loaded_dummy, DummyModel)
+    assert loaded_dummy.value == 42
+
+    # CLEANING
+    _cleanup_model(teamspace, model_name, expected_num_versions=1)

tests/integrations/test_mixins.py

@@ -18,17 +18,15 @@ def __eq__(self, other):
 def test_pickle_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
     # Create an instance of DummyModel and call push_to_registry.
     dummy = DummyModel(42)
-    dummy.push_to_registry(model_version="v1", temp_folder=str(tmp_path))
+    dummy.push_to_registry(version="v1", temp_folder=str(tmp_path))
     # The expected registry name is "dummy_model:v1" and the file should be placed in the temp folder.
     expected_path = tmp_path / "DummyModel.pkl"
     mock_upload_model.assert_called_once_with(name="DummyModel:v1", model=expected_path)
 
     # Set the mock to return the full path to the pickle file.
     mock_download_model.return_value = ["DummyModel.pkl"]
     # Call pull_from_registry and load the DummyModel instance.
-    loaded_dummy = DummyModel.pull_from_registry(
-        model_name="dummy_model", model_version="v1", temp_folder=str(tmp_path)
-    )
+    loaded_dummy = DummyModel.pull_from_registry(name="dummy_model", version="v1", temp_folder=str(tmp_path))
     # Verify that the unpickled instance has the expected value.
     assert loaded_dummy.value == 42
 
@@ -59,7 +57,7 @@ def test_pytorch_pull_updated(mock_download_model, mock_upload_model, tmp_path):
     torch.save(dummy.state_dict(), expected_path)
     # Prepare mocking for pull_from_registry.
     mock_download_model.return_value = [f"{dummy.__class__.__name__}.pth"]
-    loaded_dummy = DummyTorchModel.pull_from_registry(model_name="DummyTorchModel", temp_folder=str(tmp_path))
+    loaded_dummy = DummyTorchModel.pull_from_registry(name="DummyTorchModel", temp_folder=str(tmp_path))
     loaded_dummy.eval()
     output_after = loaded_dummy(input_tensor)