update uploading multiple files for mixins (#71)

* update
* mock
* typing

Author: Borda

README.md

@@ -192,6 +192,7 @@ class MyModel(PickleRegistryMixin):
         self.param1 = param1
         self.param2 = param2
         # Your model initialization code
+        ...
 
 
 # Create and push a model instance

requirements.txt

@@ -1,5 +1,5 @@
 # NOTE: once we add more dependencies, consider update dependabot to check for updates
 
-lightning-sdk >=0.2.5
+lightning-sdk >=0.2.7
 lightning-utilities
 joblib

src/litmodels/integrations/mixins.py

@@ -168,7 +168,7 @@ def push_to_registry(
         model_registry = f"{name}:{version}" if version else name
         # 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)
+        upload_model_files(name=model_registry, path=[torch_state_dict_path, init_kwargs_path])
 
     @classmethod
     def pull_from_registry(

src/litmodels/io/cloud.py

@@ -42,7 +42,7 @@ def _print_model_link(name: str, verbose: Union[bool, int]) -> None:
 
 def upload_model_files(
     name: str,
-    path: Union[str, Path],
+    path: Union[str, Path, List[Union[str, Path]]],
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     verbose: Union[bool, int] = 1,

tests/integrations/test_mixins.py

@@ -64,14 +64,18 @@ def test_pytorch_push_and_pull(mock_download_model, mock_upload_model, torch_cla
     input_tensor = torch.randn(1, 784)
     output_before = dummy(input_tensor)
 
-    dummy.push_to_registry(temp_folder=str(tmp_path))
-    mock_upload_model.assert_called_once_with(name=torch_class.__name__, path=str(tmp_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:
+    json_path = tmp_path / json_file
+    with open(json_path, "w") as fp:
         fp.write('{"input_size": 784, "output_size": 10}')
+
+    dummy.push_to_registry(temp_folder=str(tmp_path))
+    mock_upload_model.assert_called_once_with(
+        name=torch_class.__name__, path=[tmp_path / f"{torch_class.__name__}.pth", json_path]
+    )
+
     # Prepare mocking for pull_from_registry.
     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))