enable Torch mixin with init args (#69)

* fixing PT mixin
* update + test
* 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

@@ -1,11 +1,15 @@
+import inspect
+import json
 import pickle
 import tempfile
 import warnings
 from abc import ABC
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
 
-from litmodels import download_model, upload_model
+from lightning_utilities.core.rank_zero import rank_zero_warn
+
+from litmodels.io.cloud import download_model_files, upload_model_files
 
 if TYPE_CHECKING:
     import torch
@@ -15,7 +19,7 @@ class ModelRegistryMixin(ABC):
     """Mixin for model registry integration."""
 
     def push_to_registry(
-        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
     ) -> None:
         """Push the model to the registry.
 
@@ -26,7 +30,9 @@ def push_to_registry(
         """
 
     @classmethod
-    def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folder: Optional[str] = None) -> object:
+    def pull_from_registry(
+        cls, name: str, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
+    ) -> object:
         """Pull the model from the registry.
 
         Args:
@@ -35,7 +41,9 @@ 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]:
+    def _setup(
+        self, name: Optional[str] = None, temp_folder: Union[str, Path, None] = None
+    ) -> Tuple[str, str, Union[str, Path]]:
         """Parse and validate the model name and temporary folder."""
         if name is None:
             name = model_name = self.__class__.__name__
@@ -52,7 +60,7 @@ class PickleRegistryMixin(ModelRegistryMixin):
     """Mixin for pickle registry integration."""
 
     def push_to_registry(
-        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
     ) -> None:
         """Push the model to the registry.
 
@@ -67,10 +75,12 @@ def push_to_registry(
             pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
         if version:
             name = f"{name}:{version}"
-        upload_model(name=name, model=pickle_path)
+        upload_model_files(name=name, path=pickle_path)
 
     @classmethod
-    def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folder: Optional[str] = None) -> object:
+    def pull_from_registry(
+        cls, name: str, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
+    ) -> object:
         """Pull the model from the registry.
 
         Args:
@@ -81,7 +91,7 @@ def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folde
         if temp_folder is None:
             temp_folder = tempfile.mkdtemp()
         model_registry = f"{name}:{version}" if version else name
-        files = download_model(name=model_registry, download_dir=temp_folder)
+        files = download_model_files(name=model_registry, download_dir=temp_folder)
         pkl_files = [f for f in files if f.endswith(".pkl")]
         if not pkl_files:
             raise RuntimeError(f"No pickle file found for model: {model_registry} with {files}")
@@ -98,8 +108,27 @@ def pull_from_registry(cls, name: str, version: Optional[str] = None, temp_folde
 class PyTorchRegistryMixin(ModelRegistryMixin):
     """Mixin for PyTorch model registry integration."""
 
+    def __new__(cls, *args: Any, **kwargs: Any) -> "torch.nn.Module":
+        """Create a new instance of the class without calling __init__."""
+        instance = super().__new__(cls)
+
+        # Get __init__ signature excluding 'self'
+        init_sig = inspect.signature(cls.__init__)
+        params = list(init_sig.parameters.values())[1:]  # Skip self
+
+        # Create temporary signature for binding
+        temp_sig = init_sig.replace(parameters=params)
+
+        # Bind and apply defaults
+        bound_args = temp_sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        # Store unified kwargs
+        instance.__init_kwargs = bound_args.arguments
+        return instance
+
     def push_to_registry(
-        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Optional[str] = None
+        self, name: Optional[str] = None, version: Optional[str] = None, temp_folder: Union[str, Path, None] = None
     ) -> None:
         """Push the model to the registry.
 
@@ -110,22 +139,43 @@ def push_to_registry(
         """
         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)}")
 
         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
+
+        if self.__init_kwargs:
+            try:
+                # Save the model arguments to a JSON file
+                init_kwargs_path = Path(temp_folder) / f"{model_name}__init_kwargs.json"
+                with open(init_kwargs_path, "w") as fp:
+                    json.dump(self.__init_kwargs, fp)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to save model arguments: {e}."
+                    " Ensure the model's arguments are JSON serializable or use `PickleRegistryMixin`."
+                ) from e
+        elif not hasattr(self, "__init_kwargs"):
+            rank_zero_warn(
+                "The child class is missing `__init_kwargs`."
+                " Ensure `PyTorchRegistryMixin` is first in the inheritance order"
+                " or call `PyTorchRegistryMixin.__init__` explicitly in the child class."
+            )
+
+        torch_state_dict_path = Path(temp_folder) / f"{model_name}.pth"
+        torch.save(self.state_dict(), torch_state_dict_path)
         model_registry = f"{name}:{version}" if version else name
-        upload_model(name=model_registry, model=torch_path)
+        # todo: consider creating another temp folder and copying these two files
+        # todo: updating SDK to support uploading just specific files
+        upload_model_files(name=model_registry, path=temp_folder)
 
     @classmethod
     def pull_from_registry(
         cls,
         name: str,
         version: Optional[str] = None,
-        temp_folder: Optional[str] = None,
+        temp_folder: Union[str, Path, None] = None,
         torch_load_kwargs: Optional[dict] = None,
     ) -> "torch.nn.Module":
         """Pull the model from the registry.
@@ -141,7 +191,8 @@ def pull_from_registry(
         if temp_folder is None:
             temp_folder = tempfile.mkdtemp()
         model_registry = f"{name}:{version}" if version else name
-        files = download_model(name=model_registry, download_dir=temp_folder)
+        files = download_model_files(name=model_registry, download_dir=temp_folder)
+
         torch_files = [f for f in files if f.endswith(".pth")]
         if not torch_files:
             raise RuntimeError(f"No torch file found for model: {model_registry} with {files}")
@@ -153,8 +204,18 @@ def pull_from_registry(
             warnings.simplefilter("ignore", category=FutureWarning)
             state_dict = torch.load(state_dict_path, **(torch_load_kwargs if torch_load_kwargs else {}))
 
+        init_files = [fp for fp in files if fp.endswith("__init_kwargs.json")]
+        if not init_files:
+            init_kwargs = {}
+        elif len(init_files) > 1:
+            raise RuntimeError(f"Multiple init files found for model: {model_registry} with {init_files}")
+        else:
+            init_kwargs_path = Path(temp_folder) / init_files[0]
+            with open(init_kwargs_path) as fp:
+                init_kwargs = json.load(fp)
+
         # Create a new model instance without calling __init__
-        instance = cls()  # todo: we need to add args used when created dumped model
+        instance = cls(**init_kwargs)
         if not isinstance(instance, torch.nn.Module):
             raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(instance)}")
         # Now load the state dict on the instance

src/litmodels/io/cloud.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (the "License");
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
+from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, Union
 
 from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
@@ -41,7 +42,7 @@ def _print_model_link(name: str, verbose: Union[bool, int]) -> None:
 
 def upload_model_files(
     name: str,
-    path: str,
+    path: Union[str, Path],
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     verbose: Union[bool, int] = 1,
@@ -71,7 +72,7 @@ def upload_model_files(
 
 def download_model_files(
     name: str,
-    download_dir: str = ".",
+    download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
 ) -> Union[str, List[str]]:
     """Download a checkpoint from the model store.

src/litmodels/io/gateway.py

@@ -10,7 +10,6 @@
 
 if module_available("torch"):
     import torch
-    from torch.nn import Module
 else:
     torch = None
 
@@ -20,7 +19,7 @@
 
 def upload_model(
     name: str,
-    model: Union[str, Path, "Module", Any],
+    model: Union[str, Path, "torch.nn.Module", Any],
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     staging_dir: Optional[str] = None,
@@ -42,19 +41,17 @@ def upload_model(
     """
     if not staging_dir:
         staging_dir = tempfile.mkdtemp()
+    if isinstance(model, (str, Path)):
+        path = model
     # if LightningModule and isinstance(model, LightningModule):
     #     path = os.path.join(staging_dir, f"{model.__class__.__name__}.ckpt")
     #     model.save_checkpoint(path)
-    if torch and isinstance(model, torch.jit.ScriptModule):
+    elif torch and isinstance(model, torch.jit.ScriptModule):
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.ts")
         model.save(path)
-    elif torch and isinstance(model, Module):
+    elif torch and isinstance(model, torch.nn.Module):
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.pth")
         torch.save(model.state_dict(), path)
-    elif isinstance(model, str):
-        path = model
-    elif isinstance(model, Path):
-        path = str(model)
     else:
         path = os.path.join(staging_dir, f"{model.__class__.__name__}.pkl")
         joblib.dump(model, path)
@@ -70,7 +67,7 @@ def upload_model(
 
 def download_model(
     name: str,
-    download_dir: str = ".",
+    download_dir: Union[str, Path] = ".",
     progress_bar: bool = True,
 ) -> Union[str, List[str]]:
     """Download a checkpoint from the model store.
@@ -109,7 +106,7 @@ def load_model(name: str, download_dir: str = ".") -> Any:
     download_paths = [p for p in download_paths if Path(p).suffix.lower() not in {".md", ".txt", ".rst"}]
     if len(download_paths) > 1:
         raise NotImplementedError("Downloaded model with multiple files is not supported yet.")
-    model_path = Path(os.path.join(download_dir, download_paths[0]))
+    model_path = Path(download_dir) / download_paths[0]
     if model_path.suffix.lower() == ".pkl":
         return joblib.load(model_path)
     if model_path.suffix.lower() == ".ts":

tests/integrations/test_cloud.py

@@ -242,10 +242,14 @@ def test_pickle_mixin_push_and_pull():
     _cleanup_model(teamspace, model_name, expected_num_versions=1)
 
 
-class DummyTorchModel(torch.nn.Module, PyTorchRegistryMixin):
-    def __init__(self, input_size=784):
+# This is a dummy model for PyTorch that uses the PyTorchRegistryMixin.
+# This mixin has to be first in the inheritance order.
+# Otherwise, `PyTorchRegistryMixin.__init__` need to be called explicitly.
+class DummyTorchModel(PyTorchRegistryMixin, torch.nn.Module):
+    def __init__(self, input_size: int, output_size: int = 10):
+        # PyTorchRegistryMixin.__init__ will capture these arguments
         super().__init__()
-        self.fc = torch.nn.Linear(input_size, 10)
+        self.fc = torch.nn.Linear(input_size, output_size)
 
     def forward(self, x):
         x = x.view(x.size(0), -1)

tests/integrations/test_mixins.py

@@ -1,5 +1,6 @@
 from unittest import mock
 
+import pytest
 import torch
 from litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
 from torch import nn
@@ -13,15 +14,15 @@ def __eq__(self, other):
         return isinstance(other, DummyModel) and self.value == other.value
 
 
-@mock.patch("litmodels.integrations.mixins.upload_model")
-@mock.patch("litmodels.integrations.mixins.download_model")
+@mock.patch("litmodels.integrations.mixins.upload_model_files")
+@mock.patch("litmodels.integrations.mixins.download_model_files")
 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(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)
+    mock_upload_model.assert_called_once_with(name="DummyModel:v1", path=expected_path)
 
     # Set the mock to return the full path to the pickle file.
     mock_download_model.return_value = ["DummyModel.pkl"]
@@ -31,36 +32,52 @@ def test_pickle_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
     assert loaded_dummy.value == 42
 
 
-class DummyTorchModel(nn.Module, PyTorchRegistryMixin):
-    def __init__(self, input_size=784):
+class DummyTorchModelFirst(PyTorchRegistryMixin, nn.Module):
+    def __init__(self, input_size: int, output_size: int = 10):
+        # PyTorchRegistryMixin.__init__ will capture these arguments
         super().__init__()
-        self.fc = nn.Linear(input_size, 10)
+        self.fc = nn.Linear(input_size, output_size)
 
     def forward(self, x):
         x = x.view(x.size(0), -1)
         return self.fc(x)
 
 
-@mock.patch("litmodels.integrations.mixins.upload_model")
-@mock.patch("litmodels.integrations.mixins.download_model")
-def test_pytorch_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
+class DummyTorchModelSecond(nn.Module, PyTorchRegistryMixin):
+    def __init__(self, input_size: int, output_size: int = 10):
+        PyTorchRegistryMixin.__init__(input_size, output_size)
+        super().__init__()
+        self.fc = nn.Linear(input_size, output_size)
+
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+
+@pytest.mark.parametrize("torch_class", [DummyTorchModelFirst, DummyTorchModelSecond])
+@mock.patch("litmodels.integrations.mixins.upload_model_files")
+@mock.patch("litmodels.integrations.mixins.download_model_files")
+def test_pytorch_push_and_pull(mock_download_model, mock_upload_model, torch_class, tmp_path):
     # Create an instance, push the model and record its forward output.
-    dummy = DummyTorchModel(784)
+    dummy = torch_class(784)
     dummy.eval()
     input_tensor = torch.randn(1, 784)
     output_before = dummy(input_tensor)
 
     dummy.push_to_registry(temp_folder=str(tmp_path))
-    expected_path = tmp_path / f"{dummy.__class__.__name__}.pth"
-    mock_upload_model.assert_called_once_with(name="DummyTorchModel", model=expected_path)
+    mock_upload_model.assert_called_once_with(name=torch_class.__name__, path=str(tmp_path))
 
-    torch.save(dummy.state_dict(), expected_path)
+    torch_file = f"{dummy.__class__.__name__}.pth"
+    torch.save(dummy.state_dict(), tmp_path / torch_file)
+    json_file = f"{dummy.__class__.__name__}__init_kwargs.json"
+    with open(tmp_path / json_file, "w") as fp:
+        fp.write('{"input_size": 784, "output_size": 10}')
     # Prepare mocking for pull_from_registry.
-    mock_download_model.return_value = [f"{dummy.__class__.__name__}.pth"]
-    loaded_dummy = DummyTorchModel.pull_from_registry(name="DummyTorchModel", temp_folder=str(tmp_path))
+    mock_download_model.return_value = [torch_file, json_file]
+    loaded_dummy = torch_class.pull_from_registry(name=torch_class.__name__, temp_folder=str(tmp_path))
     loaded_dummy.eval()
     output_after = loaded_dummy(input_tensor)
 
-    assert isinstance(loaded_dummy, DummyTorchModel)
+    assert isinstance(loaded_dummy, torch_class)
     # Compare the outputs as a verification.
     assert torch.allclose(output_before, output_after), "Loaded model output differs from original."