Full support for shared model registry

fix

fix

fix

fix

fix

fix

fix

fix

fix

move endpoint id to backend

fix

fix

fix

refactor

fix

fix

remove self

fix

fix

fix

fix

fix input example name

fix

fix

put newline back

fix formatting

test

add read program

refactor

fixes

Author: robzor92

python/hsml/core/model_api.py

@@ -47,7 +47,7 @@ def put(self, model_instance, query_params):
             )
         )
 
-    def get(self, name, version):
+    def get(self, name, version, endpoint_id, shared_registry_project=None):
         """Get the metadata of a model with a certain name and version.
 
         :param name: name of the model
@@ -64,11 +64,23 @@ def get(self, name, version):
             "models",
             name + "_" + str(version),
         ]
-        query_params = {"expand": "trainingdatasets"}
+        query_params = {"expand": "trainingdatasets", "endpointId": str(endpoint_id)}
+
         model_json = _client._send_request("GET", path_params, query_params)
-        return model.Model.from_response_json(model_json)
+        model_meta = model.Model.from_response_json(model_json)
+
+        model_meta.shared_registry_project = shared_registry_project
+
+        return model_meta
 
-    def get_models(self, name, metric=None, direction=None):
+    def get_models(
+        self,
+        name,
+        endpoint_id,
+        shared_registry_project=None,
+        metric=None,
+        direction=None,
+    ):
         """Get the metadata of models based on the name or optionally the best model given a metric and direction.
 
         :param name: name of the model
@@ -83,8 +95,11 @@ def get_models(self, name, metric=None, direction=None):
 
         _client = client.get_instance()
         path_params = ["project", _client._project_id, "models"]
+        query_params = {
+            "expand": "trainingdatasets",
+            "filter_by": ["name_eq:" + name, "endpoint_id:" + str(endpoint_id)],
+        }
 
-        query_params = {"expand": "trainingdatasets", "filter_by": "name_eq:" + name}
         if metric is not None and direction is not None:
             if direction.lower() == "max":
                 direction = "desc"
@@ -95,7 +110,12 @@ def get_models(self, name, metric=None, direction=None):
             query_params["limit"] = "1"
 
         model_json = _client._send_request("GET", path_params, query_params)
-        return model.Model.from_response_json(model_json)
+        models_meta = model.Model.from_response_json(model_json)
+
+        for model_meta in models_meta:
+            model_meta.shared_registry_project = shared_registry_project
+
+        return models_meta
 
     def delete(self, model_instance):
         """Delete the model and metadata.
@@ -104,5 +124,6 @@ def delete(self, model_instance):
         :type model_instance: Model
         """
         _client = client.get_instance()
+        query_params = {"endpointId": str(model_instance.endpoint_id)}
         path_params = ["project", _client._project_id, "models", model_instance.id]
-        _client._send_request("DELETE", path_params)
+        _client._send_request("DELETE", path_params, query_params)

python/hsml/core/model_registry_api.py

@@ -50,6 +50,14 @@ def get(self, project=None):
                 )
             )
 
+        endpoint_id = _client._project_id
+        if project is not None:
+            project_info = _client._get_project_info(project)
+            endpoint_id = str(project_info["projectId"])
+
         return ModelRegistry(
-            _client._project_name, _client._project_id, shared_registry_project=project
+            _client._project_name,
+            _client._project_id,
+            endpoint_id,
+            shared_registry_project=project,
         )

python/hsml/engine/local_engine.py

@@ -25,4 +25,4 @@ def mkdir(self, model_instance):
         self._dataset_api.mkdir(model_instance.path)
 
     def delete(self, model_instance):
-        self._dataset_api.delete(model_instance.path)
+        self._dataset_api.rm(model_instance.path)

python/hsml/engine/model_engine.py

@@ -43,17 +43,26 @@ def __init__(self):
         else:
             self._engine = hopsworks_engine.HopsworksEngine()
 
-    def _poll_model_available(self, model_instance, await_registration):
+    def _poll_model_available(self, _client, model_instance, await_registration):
         if await_registration > 0:
+            endpoint_id = _client._project_id
+            if model_instance.shared_registry_project is not None:
+                project_info = _client._get_project_info(
+                    model_instance.shared_registry_project
+                )
+                endpoint_id = str(project_info["projectId"])
             sleep_seconds = 5
             for i in range(int(await_registration / sleep_seconds)):
                 try:
                     time.sleep(sleep_seconds)
-                    model = self._model_api.get(
-                        name=model_instance.name, version=model_instance.version
+                    model_meta = self._model_api.get(
+                        model_instance.name,
+                        model_instance.version,
+                        endpoint_id,
+                        model_instance.shared_registry_project,
                     )
-                    if model is not None:
-                        return model
+                    if model_meta is not None:
+                        return model_meta
                 except RestAPIError as e:
                     if e.response.status_code != 404:
                         raise e
@@ -88,7 +97,7 @@ def _upload_additional_resources(self, model_instance, dataset_model_version_pat
             model_instance.signature = dataset_model_version_path + "/signature.json"
         return model_instance
 
-    def _copy_hopsfs_model(self, model_path, dataset_model_version_path, client):
+    def _copy_hopsfs_model(self, model_path, dataset_model_version_path):
         # Strip hdfs prefix
         if model_path.startswith("hdfs:/"):
             projects_index = model_path.find("/Projects", 0)
@@ -168,18 +177,9 @@ def _set_model_version(
             )
         return model_instance
 
-    def _build_registry_path(self, model_instance, artifact_path):
-        models_path = None
-        if model_instance.shared_registry_project is not None:
-            models_path = "{}::{}".format(
-                model_instance.shared_registry_project,
-                constants.MODEL_SERVING.MODELS_DATASET,
-            )
-        else:
-            models_path = constants.MODEL_SERVING.MODELS_DATASET
-        return artifact_path.replace(
-            constants.MODEL_SERVING.MODELS_DATASET, models_path
-        )
+    def _build_artifact_path(self, model_instance, artifact):
+        artifact_path = model_instance.path + "/" + artifact
+        return artifact_path
 
     def save(self, model_instance, model_path, await_registration=480):
 
@@ -287,9 +287,7 @@ def save(self, model_instance, model_path, await_registration=480):
                     elif self._dataset_api.path_exists(
                         model_path
                     ):  # check hdfs relative and absolute
-                        self._copy_hopsfs_model(
-                            model_path, dataset_model_version_path, _client
-                        )
+                        self._copy_hopsfs_model(model_path, dataset_model_version_path)
                     else:
                         raise IOError(
                             "Could not find path {} in the local filesystem or in HopsFS".format(
@@ -299,11 +297,9 @@ def save(self, model_instance, model_path, await_registration=480):
                 if step["id"] == 3:
                     self._model_api.put(model_instance, model_query_params)
                 if step["id"] == 4:
-                    # We do not necessarily have access to the Models REST API for the shared model registry, so we do not know if it is registered or not
-                    if not is_shared_registry:
-                        model_instance = self._poll_model_available(
-                            model_instance, await_registration
-                        )
+                    model_instance = self._poll_model_available(
+                        _client, model_instance, await_registration
+                    )
                 if step["id"] == 5:
                     pass
             except BaseException as be:
@@ -326,12 +322,7 @@ def download(self, model_instance):
         zip_path = model_version_path + ".zip"
         os.makedirs(model_name_path)
 
-        dataset_model_name_path = (
-            constants.MODEL_SERVING.MODELS_DATASET + "/" + model_instance._name
-        )
-        dataset_model_version_path = (
-            dataset_model_name_path + "/" + str(model_instance._version)
-        )
+        dataset_model_version_path = model_instance.path
 
         temp_download_dir = "/Resources" + "/" + str(uuid.uuid4())
         try:
@@ -340,7 +331,7 @@ def download(self, model_instance):
                 dataset_model_version_path,
                 destination_path=temp_download_dir,
                 block=True,
-                timeout=480,
+                timeout=600,
             )
             self._dataset_api.download(
                 temp_download_dir + "/" + str(model_instance._version) + ".zip",
@@ -358,44 +349,30 @@ def download(self, model_instance):
 
         return model_version_path
 
-    def read_input_example(self, model_instance):
-        try:
-            tmp_dir = tempfile.TemporaryDirectory(dir=os.getcwd())
-            self._dataset_api.download(
-                self._build_registry_path(
-                    model_instance, model_instance._input_example
-                ),
-                tmp_dir.name + "/inputs.json",
-            )
-            with open(tmp_dir.name + "/inputs.json", "rb") as f:
-                return json.loads(f.read())
-        finally:
-            if tmp_dir is not None and os.path.exists(tmp_dir.name):
-                tmp_dir.cleanup()
-
-    def read_environment(self, model_instance):
+    def read_file(self, model_instance, artifact):
         try:
+            artifact = os.path.basename(artifact)
             tmp_dir = tempfile.TemporaryDirectory(dir=os.getcwd())
+            local_artifact_path = os.path.join(tmp_dir.name, artifact)
             self._dataset_api.download(
-                self._build_registry_path(
-                    model_instance, model_instance._environment[0]
-                ),
-                tmp_dir.name + "/environment.yml",
+                self._build_artifact_path(model_instance, artifact),
+                local_artifact_path,
             )
-            with open(tmp_dir.name + "/environment.yml", "r") as f:
+            with open(local_artifact_path, "r") as f:
                 return f.read()
         finally:
             if tmp_dir is not None and os.path.exists(tmp_dir.name):
                 tmp_dir.cleanup()
 
-    def read_signature(self, model_instance):
+    def read_json(self, model_instance, artifact):
         try:
             tmp_dir = tempfile.TemporaryDirectory(dir=os.getcwd())
+            local_artifact_path = os.path.join(tmp_dir.name, artifact)
             self._dataset_api.download(
-                self._build_registry_path(model_instance, model_instance._signature),
-                tmp_dir.name + "/signature.json",
+                self._build_artifact_path(model_instance, artifact),
+                os.path.join(tmp_dir.name, artifact),
             )
-            with open(tmp_dir.name + "/signature.json", "rb") as f:
+            with open(local_artifact_path, "rb") as f:
                 return json.loads(f.read())
         finally:
             if tmp_dir is not None and os.path.exists(tmp_dir.name):

python/hsml/model.py

@@ -45,13 +45,10 @@ def __init__(
         training_dataset=None,
         input_example=None,
         framework=None,
+        endpoint_id=None,
     ):
 
-        if id is None:
-            self._id = name + "_" + str(version)
-        else:
-            self._id = id
-
+        self._id = id
         self._name = name
         self._version = version
 
@@ -73,7 +70,11 @@ def __init__(
         self._signature = signature
         self._training_dataset = training_dataset
 
-        self._shared_registry_project = None
+        # This is needed for update_from_response_json function to not overwrite name of the shared registry this model originates from
+        if not hasattr(self, "_shared_registry_project"):
+            self._shared_registry_project = None
+
+        self._endpoint_id = endpoint_id
 
         self._model_api = model_api.ModelApi()
         self._dataset_api = dataset_api.DatasetApi()
@@ -186,8 +187,8 @@ def created(self, created):
     @property
     def environment(self):
         """Input example of the model."""
-        if self._environment is not None and isinstance(self._environment, list):
-            self._environment = self._model_engine.read_environment(self)
+        if self._environment is not None:
+            return self._model_engine.read_file(self, "environment.yml")
         return self._environment
 
     @environment.setter
@@ -215,6 +216,8 @@ def training_metrics(self, training_metrics):
     @property
     def program(self):
         """Executable used to export the model."""
+        if self._program is not None:
+            return self._model_engine.read_file(self, self._program)
         return self._program
 
     @program.setter
@@ -234,7 +237,9 @@ def user(self, user_full_name):
     def input_example(self):
         """input_example of the model."""
         if self._input_example is not None and isinstance(self._input_example, str):
-            self._input_example = self._model_engine.read_input_example(self)
+            self._input_example = self._model_engine.read_json(
+                self, "input_example.json"
+            )
         return self._input_example
 
     @input_example.setter
@@ -254,7 +259,7 @@ def framework(self, framework):
     def signature(self):
         """signature of the model."""
         if self._signature is not None and isinstance(self._signature, str):
-            self._signature = self._model_engine.read_signature(self)
+            self._signature = self._model_engine.read_json(self, "signature.json")
         return self._signature
 
     @signature.setter
@@ -289,6 +294,15 @@ def project_name(self):
     def project_name(self, project_name):
         self._project_name = project_name
 
+    @property
+    def endpoint_id(self):
+        """endpoint_id of the model."""
+        return self._endpoint_id
+
+    @endpoint_id.setter
+    def endpoint_id(self, endpoint_id):
+        self._endpoint_id = endpoint_id
+
     @property
     def experiment_project_name(self):
         """experiment_project_name of the model."""