+ test

Author: Borda

src/litmodels/io/gateway.py

@@ -21,6 +21,7 @@ def upload_model(
     progress_bar: bool = True,
     cluster_id: Optional[str] = None,
     staging_dir: Optional[str] = None,
+    verbose: Union[bool, int] = 0,
 ) -> UploadedModelInfo:
     """Upload a checkpoint to the model store.
 
@@ -33,6 +34,7 @@ def upload_model(
             automatically.
         staging_dir: A directory where the model can be saved temporarily. If not provided, a temporary directory will
             be created and used.
+        verbose: Whether to print some additional information about the uploaded model.
 
     """
     if not staging_dir:
@@ -54,6 +56,7 @@ def upload_model(
         name=name,
         progress_bar=progress_bar,
         cluster_id=cluster_id,
+        verbose=verbose,
     )
 
 

tests/test_io_cloud.py

@@ -23,7 +23,7 @@ def test_wrong_model_name(name):
         (Module(), f"%s{os.path.sep}Module.pth", True),
     ],
 )
-def test_upload_model(mocker, tmpdir, model, model_path):
+def test_upload_model(mocker, tmpdir, model, model_path, verbose):
     # mocking the _get_teamspace to return another mock
     ts_mock = mock.MagicMock()
     mocker.patch("litmodels.io.cloud._get_teamspace", return_value=ts_mock)
@@ -34,6 +34,7 @@ def test_upload_model(mocker, tmpdir, model, model_path):
         name="org-name/teamspace/model-name",
         cluster_id="cluster_id",
         staging_dir=tmpdir,
+        verbose=verbose,
     )
     expected_path = model_path % str(tmpdir) if "%" in model_path else model_path
     ts_mock.upload_model.assert_called_once_with(