fix: Torch mixin + test with Prod (#68)

* test
* fixing

---------

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

Author: Borda

src/litmodels/integrations/mixins.py

@@ -3,7 +3,7 @@
 import warnings
 from abc import ABC
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Tuple
 
 from litmodels import download_model, upload_model
 
@@ -35,6 +35,18 @@ def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folde
             temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
         """
 
+    def _setup(self, name: Optional[str] = None, temp_folder: Optional[str] = None) -> Tuple[str, str, str]:
+        """Parse and validate the model name and temporary folder."""
+        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.mkdtemp()
+        return name, model_name, temp_folder
+
 
 class PickleRegistryMixin(ModelRegistryMixin):
     """Mixin for pickle registry integration."""
@@ -49,14 +61,7 @@ def push_to_registry(
             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 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.mkdtemp()
+        name, model_name, temp_folder = self._setup(name, temp_folder)
         pickle_path = Path(temp_folder) / f"{model_name}.pkl"
         with open(pickle_path, "wb") as fp:
             pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
@@ -93,14 +98,6 @@ def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folde
 class PyTorchRegistryMixin(ModelRegistryMixin):
     """Mixin for PyTorch model registry integration."""
 
-    def __post_init__(self) -> None:
-        """Post-initialization method to set up the model."""
-        import torch
-
-        # Ensure that the model is in evaluation mode
-        if not isinstance(self, torch.nn.Module):
-            raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
-
     def push_to_registry(
         self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
     ) -> None:
@@ -116,11 +113,8 @@ def push_to_registry(
         if not isinstance(self, torch.nn.Module):
             raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
 
-        if name is None:
-            name = self.__class__.__name__
-        if temp_folder is None:
-            temp_folder = tempfile.mkdtemp()
-        torch_path = Path(temp_folder) / f"{name}.pth"
+        name, model_name, temp_folder = self._setup(name, temp_folder)
+        torch_path = Path(temp_folder) / f"{model_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"{name}:{version}" if version else name

tests/integrations/test_cloud.py

@@ -5,11 +5,12 @@
 from typing import Optional
 
 import pytest
+import torch
 from lightning_sdk import Teamspace
 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 litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
 
 from tests.integrations import _SKIP_IF_LIGHTNING_BELLOW_2_5_1, _SKIP_IF_PYTORCHLIGHTNING_BELLOW_2_5_1
 
@@ -239,3 +240,39 @@ def test_pickle_mixin_push_and_pull():
 
     # CLEANING
     _cleanup_model(teamspace, model_name, expected_num_versions=1)
+
+
+class DummyTorchModel(torch.nn.Module, PyTorchRegistryMixin):
+    def __init__(self, input_size=784):
+        super().__init__()
+        self.fc = torch.nn.Linear(input_size, 10)
+
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+
+@pytest.mark.cloud()
+def test_pytorch_mixin_push_and_pull():
+    # model name with random hash
+    teamspace, org_team, model_name = _prepare_variables("torch_mixin")
+    model_registry = f"{org_team}/{model_name}"
+
+    # Create an instance, push the model and record its forward output.
+    dummy = DummyTorchModel(784)
+    dummy.eval()
+    input_tensor = torch.randn(1, 784)
+    output_before = dummy(input_tensor)
+
+    dummy.push_to_registry(model_registry)
+
+    loaded_dummy = DummyTorchModel.pull_from_registry(model_registry)
+    loaded_dummy.eval()
+    output_after = loaded_dummy(input_tensor)
+
+    assert isinstance(loaded_dummy, DummyTorchModel)
+    # Compare the outputs as a verification.
+    assert torch.allclose(output_before, output_after), "Loaded model output differs from original."
+
+    # CLEANING
+    _cleanup_model(teamspace, model_name, expected_num_versions=1)

tests/integrations/test_mixins.py

@@ -43,7 +43,7 @@ def forward(self, x):
 
 @mock.patch("litmodels.integrations.mixins.upload_model")
 @mock.patch("litmodels.integrations.mixins.download_model")
-def test_pytorch_pull_updated(mock_download_model, mock_upload_model, tmp_path):
+def test_pytorch_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
     # Create an instance, push the model and record its forward output.
     dummy = DummyTorchModel(784)
     dummy.eval()