Do not use os.getcwd() for creating model archives (+ onnx issue)

CHANGELOG.md

@@ -15,6 +15,10 @@ With thanks to [Tom in the MLOps Community](https://mlops-community.slack.com/ar
 
 * Model store cleans up temporary files that it generates, even if downloads or uploads are cancelled mid flight [#258](https://github.com/operatorai/modelstore/pull/258).
 
+And thanks to [Michael in the MLOps Community](https://mlops-community.slack.com/archives/C0227QJCDS8/p1695397710224629) for the feedback:
+
+* Model store no longer creates its `artifacts.tar.gz` in the `os.getcwd()`, which was preventing modelstore from being used for parallel uploads [#266](https://github.com/operatorai/modelstore/pull/266).
+
 ## modelstore 0.0.79 ([June 2023](https://github.com/operatorai/modelstore/pull/243))
 
 **🆕  New functionality**

modelstore/model_store.py

@@ -262,14 +262,13 @@ def upload(self, domain: str, model_id: Optional[str] = None, **kwargs) -> dict:
             raise ValueError(f"model_id='{model_id}' contains invalid characters")
 
         # Figure out which library the kwargs match with
-        # We do this _before_ checking whether the model exists to raise
+        # We do this _before_ checking whether the model exists to
         # catch if the kwargs aren't quite right before potentially modifying
         # model state, below
         # pylint: disable=no-member
         managers = matching_managers(self._libraries, **kwargs)
-        if len(managers) == 1:
-            manager = managers[0]
-        else:
+        manager = managers[0]
+        if len(managers) > 1:
             # There are cases where we can match on more than one
             # manager (e.g., a model and an explainer)
             manager = MultipleModelsManager(managers, self.storage)

modelstore/models/model_manager.py

@@ -12,10 +12,9 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 import os
-import shutil
 import tarfile
-import tempfile
 import warnings
+import tempfile
 from abc import ABC, ABCMeta, abstractmethod
 from typing import Any, Optional
 
@@ -150,36 +149,28 @@ def _collect_extras(self, **kwargs) -> set:
         extra_paths = extras if isinstance(extras, list) else [extras]
         return set(f for f in extra_paths if os.path.isfile(f))
 
-    def _create_archive(self, **kwargs) -> str:
+    def _create_archive(self, tmp_dir: str, **kwargs) -> str:
         """
         Creates the `artifacts.tar.gz` archive which contains
         all of the files of the model
         """
         self._validate_kwargs(**kwargs)
-        archive_name = "artifacts.tar.gz"
-        archive_path = os.path.join(os.getcwd(), archive_name)
-        if os.path.exists(archive_path):
-            raise FileExistsError(f"modelstore cannot create an: {archive_name} file.")
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            result = os.path.join(tmp_dir, archive_name)
-            with tarfile.open(result, "w:gz") as tar:
-                # Add all of the model files to the top-level
-                # of the archive
-                for file_path in self._collect_files(tmp_dir, **kwargs):
-                    file_name = os.path.split(file_path)[1]
-                    tar.add(name=file_path, arcname=file_name)
-
-                # Add any extra files to a sub-directory of
-                # the archive
-                for file_path in self._collect_extras(**kwargs):
-                    file_name = os.path.split(file_path)[1]
-                    tar.add(
-                        name=file_path,
-                        arcname=os.path.join("extras", file_name),
-                    )
-
-            # Move the archive to the current working directory
-            shutil.move(result, archive_path)
+        archive_path = os.path.join(tmp_dir, "artifacts.tar.gz")
+        with tarfile.open(archive_path, "w:gz") as tar:
+            # Add all of the model files to the top-level
+            # of the archive
+            for file_path in self._collect_files(tmp_dir, **kwargs):
+                file_name = os.path.split(file_path)[1]
+                tar.add(name=file_path, arcname=file_name)
+
+            # Add any extra files to a sub-directory of
+            # the archive
+            for file_path in self._collect_extras(**kwargs):
+                file_name = os.path.split(file_path)[1]
+                tar.add(
+                    name=file_path,
+                    arcname=os.path.join("extras", file_name),
+                )
         return archive_path
 
     def upload(
@@ -207,13 +198,19 @@ def upload(
             data=self.model_data(**kwargs),
         )
 
-        try:
+        with tempfile.TemporaryDirectory() as tmp_dir:
             # Create the model archive and return
             # meta-data about its location
-            archive_path = self._create_archive(**kwargs)
+            archive_path = self._create_archive(tmp_dir, **kwargs)
 
             # Upload the model archive and any additional extras
-            storage_meta_data = self.storage.upload(domain, model_id, archive_path)
+            storage_meta_data = self.storage.upload(
+                domain,
+                model_id,
+                archive_path,
+            )
+
+            # Save the combined meta-data to storage
             meta_data = metadata.Summary.generate(
                 code_meta_data=metadata.Code.generate(
                     deps_list=self.get_dependencies()
@@ -222,13 +219,7 @@ def upload(
                 storage_meta_data=storage_meta_data,
                 extra_metadata=kwargs.get("extra_metadata"),
             )
-
-            # Save the combined meta-data to storage
             self.storage.set_meta_data(domain, model_id, meta_data)
-        finally:
-            # Clean up on exceptions (including KeyboardInterrupt)
-            if os.path.exists(archive_path):
-                os.remove(archive_path)
         return meta_data
 
 

requirements-dev1.txt

@@ -7,7 +7,7 @@ azure-storage-blob>=12.11.0
 boto3>=1.21.41
 google-cloud-storage>=2.3.0
 minio>=7.1.12
-pydoop<=2.0.0; sys_platform == 'darwin'
+#pydoop<=2.0.0; sys_platform == 'darwin'
 pystan>=2.19.1.1 # required to be installed before prophet
 
 # Machine Learning
@@ -19,7 +19,7 @@ Keras-Preprocessing>=1.1.2
 lightgbm>=3.3.2
 mxnet>=1.8.0.post0
 onnx==1.12.0 # onnx 1.13.0 depends on protobuf<4 and >=3.20.2, which breaks everything else
-onnxruntime>=1.11.0
+onnxruntime>=1.11.0,<1.16.0 # https://github.com/microsoft/onnxruntime/issues/17631
 prophet>=1.0.1
 pyspark>=3.3.1
 pytorch-lightning>=1.6.1

tests/models/test_model_manager.py

@@ -13,8 +13,9 @@
 #    limitations under the License.
 import os
 import tarfile
+import tempfile
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import pytest
 from mock import patch
@@ -38,8 +39,8 @@ def __init__(self, tmpdir):
     def upload(
         self,
         domain: str,
-        local_path: str,
-        extras: Optional[Union[str, list]] = None,
+        model_id: str,
+        local_path: str
     ):
         self.called = True
 
@@ -181,7 +182,7 @@ def test_upload(mock_manager):
 
 @patch("modelstore.models.model_manager.get_python_version")
 def test_load_from_different_python(mock_python_version, mock_manager):
-    mock_python_version.return_value = f"python:2.7.13"
+    mock_python_version.return_value = "python:2.7.13"
     meta_data = metadata.Summary.generate(
         code_meta_data=metadata.Code.generate(deps_list=[]),
         model_meta_data=None,
@@ -192,25 +193,24 @@ def test_load_from_different_python(mock_python_version, mock_manager):
 
 
 def test_create_archive(mock_manager, mock_file):
-    target = os.path.join(os.getcwd(), "artifacts.tar.gz")
-    if os.path.exists(target):
-        os.remove(target)
-
-    mock_manager._create_archive(
-        model="model",
-        config="config",
-        extra_files=mock_file,
-    )
-    exp = sorted(
-        [
-            "model-info.json",
-            "model.joblib",
-            "config.json",
-            os.path.join("extras", "extra-file.csv"),
-        ]
-    )
-    with tarfile.open(target) as tar:
-        files = sorted([f.name for f in tar.getmembers()])
-    assert len(files) == len(exp)
-    assert files == exp
-    os.remove(target)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        mock_manager._create_archive(
+            tmp_dir,
+            model="model",
+            config="config",
+            extra_files=mock_file,
+        )
+        exp = sorted(
+            [
+                "model-info.json",
+                "model.joblib",
+                "config.json",
+                os.path.join("extras", "extra-file.csv"),
+            ]
+        )
+
+        target = os.path.join(tmp_dir, "artifacts.tar.gz")
+        with tarfile.open(target) as tar:
+            files = sorted([f.name for f in tar.getmembers()])
+        assert len(files) == len(exp)
+        assert files == exp

tests/models/test_pyspark.py

@@ -12,7 +12,6 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 import os
-import platform
 
 import numpy as np
 import pytest
@@ -72,6 +71,7 @@ def test_model_data(spark_manager, spark_model):
 
 
 def test_required_kwargs(spark_manager):
+    # pylint: disable=protected-access
     assert spark_manager._required_kwargs() == ["model"]
 
 
@@ -82,6 +82,7 @@ def test_matches_with(spark_manager, spark_model):
 
 
 def test_get_functions(spark_manager, spark_model):
+    # pylint: disable=protected-access
     assert len(spark_manager._get_functions(model=spark_model)) == 1
 
 
@@ -98,12 +99,12 @@ def test_save_model(spark_model, tmp_path):
     ]
     assert exp == res
     exists_fn = os.path.exists
-    if platform.system() == "Darwin":
-        # Running hadoop locally, so need to check
-        # for the files in hdfs
-        import pydoop.hdfs as hdfs
+    # if platform.system() == "Darwin":
+    #     # Running hadoop locally, so need to check
+    #     # for the files in hdfs
+    #     import pydoop.hdfs as hdfs
 
-        exists_fn = hdfs.path.exists
+    #     exists_fn = hdfs.path.exists
     assert all(exists_fn(x) for x in exp)
 
 

tests/models/test_xgboost.py

@@ -18,7 +18,6 @@
 import xgboost as xgb
 
 from modelstore.metadata import metadata
-from modelstore.metadata.dataset.dataset import Features, Labels
 from modelstore.models import xgboost
 
 # pylint: disable=unused-import
@@ -113,10 +112,6 @@ def test_save_booster_config(xgb_booster, tmp_path):
 
 
 def test_load_model(tmp_path, xgb_manager, xgb_model, classification_data):
-    # Some fields in xgboost get_params change when loading
-    # or are nans; we cannot compare them in this test
-    ignore_params = ["missing", "tree_method"]
-
     # Get the model predictions
     X_train, _ = classification_data
     y_pred = xgb_model.predict(X_train)
@@ -125,8 +120,7 @@ def test_load_model(tmp_path, xgb_manager, xgb_model, classification_data):
     model_path = os.path.join(tmp_path, xgboost.MODEL_FILE)
     xgb_model.save_model(model_path)
     xgb_model_params = xgb_model.get_params()
-    for param in ignore_params:
-        xgb_model_params.pop(param)
+
 
     #  Load the model
     loaded_model = xgb_manager.load(
@@ -158,9 +152,19 @@ def test_load_model(tmp_path, xgb_manager, xgb_model, classification_data):
 
     # They should also have the same params
     loaded_model_params = loaded_model.get_params()
+    assert_same_xgboost_params(xgb_model_params, loaded_model_params)
+
+
+def assert_same_xgboost_params(a_params: dict, b_params: dict):
+    # Some fields in xgboost get_params change when loading
+    # or are nans; we cannot compare them in this test
+    ignore_params = ["missing", "tree_method", "use_label_encoder", "n_jobs"]
     for param in ignore_params:
-        loaded_model_params.pop(param)
-    assert xgb_model_params == loaded_model_params
+        if param in a_params:
+            a_params.pop(param)
+        if param in b_params:
+            b_params.pop(param)
+    assert a_params == b_params
 
 
 def test_load_booster(tmp_path, xgb_manager, xgb_booster, classification_data):

tests/storage/test_hdfs.py

@@ -12,7 +12,6 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 import os
-import platform
 
 import pytest
 
@@ -32,16 +31,12 @@
 # pylint: disable=missing-function-docstring
 
 
-def is_not_mac() -> bool:
-    return platform.system() != "Darwin"
-
-
 @pytest.fixture
 def storage(tmp_path):
     return HdfsStorage(root_prefix=str(tmp_path), create_directory=True)
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 def test_create_from_environment_variables(monkeypatch):
     # Does not fail when environment variables exist
     for key in HdfsStorage.BUILD_FROM_ENVIRONMENT.get("required", []):
@@ -52,12 +47,12 @@ def test_create_from_environment_variables(monkeypatch):
         pytest.fail("Failed to initialise storage from env variables")
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 def test_validate(storage):
     assert storage.validate()
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 def test_push_and_pull(storage, tmp_path):
     # pylint: disable=import-outside-toplevel
     import pydoop.hdfs as hdfs
@@ -73,7 +68,7 @@ def test_push_and_pull(storage, tmp_path):
     hdfs.rm(files[0])
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 @pytest.mark.parametrize(
     "file_exists,should_call_delete",
     [
@@ -96,7 +91,7 @@ def test_remove(storage, file_exists, should_call_delete):
     assert not os.path.exists(prefix)
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 def test_read_json_objects_ignores_non_json(storage):
     # pylint: disable=import-outside-toplevel
     import pydoop.hdfs as hdfs
@@ -112,7 +107,7 @@ def test_read_json_objects_ignores_non_json(storage):
     _ = [hdfs.rm(f) for f in hdfs.ls(prefix)]
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 def test_storage_location(storage):
     prefix = remote_path()
     # Asserts that the location meta data is correctly formatted
@@ -124,7 +119,7 @@ def test_storage_location(storage):
     assert storage._storage_location(prefix) == expected
 
 
-@pytest.mark.skipif(is_not_mac(), reason="no hadoop in ci")
+@pytest.mark.skip(reason="no hadoop installed")
 @pytest.mark.parametrize(
     "meta_data,should_raise,result",
     [