enable standalone artifacts (#376)

* core: enable standalone artifacts Fixes: #231 Signed-off-by: Isabella do Amaral <[email protected]> * core: improve tests Signed-off-by: Alessio Pragliola <[email protected]> Co-authored-by: Alessio Pragliola <[email protected]> Signed-off-by: Isabella do Amaral <[email protected]> * py: update core bindings to match core API Signed-off-by: Isabella do Amaral <[email protected]> * oapi: override readOnly id in response models This is a workaround to allow upserting models using a `POST` request to `/model_versions/{id}/artifacts`. Signed-off-by: Isabella do Amaral <[email protected]> * docs: update go lib docs Signed-off-by: Isabella do Amaral <[email protected]> * api: fix upsert MV artifact response status Signed-off-by: Isabella do Amaral <[email protected]> * docs: apply suggestion Signed-off-by: Isabella do Amaral <[email protected]> --------- Signed-off-by: Isabella do Amaral <[email protected]> Signed-off-by: Alessio Pragliola <[email protected]> Co-authored-by: Alessio Pragliola <[email protected]>
kubeflow · Sep 18, 2024 · e5422f7 · e5422f7
1 parent 8170dc4
commit e5422f7
Show file tree

Hide file tree

Showing 43 changed files with 1,156 additions and 856 deletions.
diff --git a/api/openapi/model-registry.yaml b/api/openapi/model-registry.yaml
@@ -391,8 +391,8 @@ paths:
           $ref: "#/components/responses/NotFound"
         "500":
           $ref: "#/components/responses/InternalServerError"
-      operationId: createModelVersionArtifact
-      summary: Create an Artifact in a ModelVersion
+      operationId: upsertModelVersionArtifact
+      summary: Upsert an Artifact in a ModelVersion
       description: Creates a new instance of an Artifact if needed and associates it with `ModelVersion`.
     parameters:
       - name: modelversionId
@@ -1224,9 +1224,8 @@ components:
           properties:
             id:
               format: int64
-              description: Output only. The unique server generated id of the resource.
+              description: The unique server generated id of the resource.
               type: string
-              readOnly: true
             createTimeSinceEpoch:
               format: int64
               description: Output only. Create time of the resource in millisecond since epoch.

diff --git a/clients/python/Makefile b/clients/python/Makefile
@@ -28,8 +28,7 @@ test:
 
 .PHONY: lint
 lint:
-	poetry run ruff check
-	poetry run black src/mr_openapi --check
+	poetry run ruff check src/model_registry
 
 .PHONY: tidy
 tidy:

diff --git a/clients/python/src/model_registry/_client.py b/clients/python/src/model_registry/_client.py
@@ -121,7 +121,7 @@ async def _register_model_artifact(
         self, mv: ModelVersion, name: str, uri: str, /, **kwargs
     ) -> ModelArtifact:
         assert mv.id is not None, "Model version must have an ID"
-        return await self._api.upsert_model_artifact(
+        return await self._api.upsert_model_version_artifact(
             ModelArtifact(name=name, uri=uri, **kwargs), mv.id
         )
 
@@ -206,8 +206,8 @@ def update(self, model: TModel) -> TModel:
         if isinstance(model, RegisteredModel):
             return self.async_runner(self._api.upsert_registered_model(model))
         if isinstance(model, ModelVersion):
-            return self.async_runner(self._api.upsert_model_version(model, model.id))
-        return self.async_runner(self._api.upsert_model_artifact(model, model.id))
+            return self.async_runner(self._api.upsert_model_version(model, None))
+        return self.async_runner(self._api.upsert_model_artifact(model))
 
     def register_hf_model(
         self,

diff --git a/clients/python/src/model_registry/core.py b/clients/python/src/model_registry/core.py
@@ -183,15 +183,15 @@ async def get_registered_models(
         return [RegisteredModel.from_basemodel(rm) for rm in rm_list.items or []]
 
     async def upsert_model_version(
-        self, model_version: ModelVersion, registered_model_id: str
+        self, model_version: ModelVersion, registered_model_id: str | None = None
     ) -> ModelVersion:
         """Upsert a model version.
 
         Updates or creates a model version on the server.
 
         Args:
             model_version: Model version to upsert.
-            registered_model_id: ID of the registered model this version will be associated to.
+            registered_model_id: ID of the registered model this version will be associated to. Can be None when updating an existing model version.
 
         Returns:
             New model version.
@@ -201,10 +201,13 @@ async def upsert_model_version(
                 mv = await client.update_model_version(
                     model_version.id, model_version.update()
                 )
-            else:
+            elif registered_model_id:
                 mv = await client.create_model_version(
                     model_version.create(registered_model_id=registered_model_id)
                 )
+            else:
+                msg = f"Registered model ID required for creating a new model version: {model_version}"
+                raise ValueError(msg)
 
         return ModelVersion.from_basemodel(mv)
 
@@ -295,7 +298,7 @@ async def get_model_version_by_params(
         return ModelVersion.from_basemodel(mv)
 
     async def upsert_model_artifact(
-        self, model_artifact: ModelArtifact, model_version_id: str
+        self, model_artifact: ModelArtifact
     ) -> ModelArtifact:
         """Upsert a model artifact.
 
@@ -308,19 +311,16 @@ async def upsert_model_artifact(
         Returns:
             New model artifact.
         """
-        if not model_artifact.id:
-            return await self.create_model_version_artifact(
-                model_artifact, model_version_id
-            )
-
         async with self.get_client() as client:
-            return ModelArtifact.from_basemodel(
-                await client.update_model_artifact(
+            if not model_artifact.id:
+                ma = await client.create_model_artifact(model_artifact.create())
+            else:
+                ma = await client.update_model_artifact(
                     model_artifact.id, model_artifact.update()
                 )
-            )
+        return ModelArtifact.from_basemodel(ma)
 
-    async def create_model_version_artifact(
+    async def upsert_model_version_artifact(
         self, artifact: ArtifactT, model_version_id: str
     ) -> ArtifactT:
         """Creates a model version artifact.
@@ -338,7 +338,7 @@ async def create_model_version_artifact(
             return cast(
                 ArtifactT,
                 Artifact.validate_artifact(
-                    await client.create_model_version_artifact(
+                    await client.upsert_model_version_artifact(
                         model_version_id, artifact.wrap()
                     )
                 ),

diff --git a/clients/python/src/mr_openapi/README.md b/clients/python/src/mr_openapi/README.md
@@ -79,7 +79,6 @@ Class | Method | HTTP request | Description
 *ModelRegistryServiceApi* | [**create_inference_service_serve**](mr_openapi/docs/ModelRegistryServiceApi.md#create_inference_service_serve) | **POST** /api/model_registry/v1alpha3/inference_services/{inferenceserviceId}/serves | Create a ServeModel action in a InferenceService
 *ModelRegistryServiceApi* | [**create_model_artifact**](mr_openapi/docs/ModelRegistryServiceApi.md#create_model_artifact) | **POST** /api/model_registry/v1alpha3/model_artifacts | Create a ModelArtifact
 *ModelRegistryServiceApi* | [**create_model_version**](mr_openapi/docs/ModelRegistryServiceApi.md#create_model_version) | **POST** /api/model_registry/v1alpha3/model_versions | Create a ModelVersion
-*ModelRegistryServiceApi* | [**create_model_version_artifact**](mr_openapi/docs/ModelRegistryServiceApi.md#create_model_version_artifact) | **POST** /api/model_registry/v1alpha3/model_versions/{modelversionId}/artifacts | Create an Artifact in a ModelVersion
 *ModelRegistryServiceApi* | [**create_registered_model**](mr_openapi/docs/ModelRegistryServiceApi.md#create_registered_model) | **POST** /api/model_registry/v1alpha3/registered_models | Create a RegisteredModel
 *ModelRegistryServiceApi* | [**create_registered_model_version**](mr_openapi/docs/ModelRegistryServiceApi.md#create_registered_model_version) | **POST** /api/model_registry/v1alpha3/registered_models/{registeredmodelId}/versions | Create a ModelVersion in RegisteredModel
 *ModelRegistryServiceApi* | [**create_serving_environment**](mr_openapi/docs/ModelRegistryServiceApi.md#create_serving_environment) | **POST** /api/model_registry/v1alpha3/serving_environments | Create a ServingEnvironment
@@ -109,6 +108,7 @@ Class | Method | HTTP request | Description
 *ModelRegistryServiceApi* | [**update_model_version**](mr_openapi/docs/ModelRegistryServiceApi.md#update_model_version) | **PATCH** /api/model_registry/v1alpha3/model_versions/{modelversionId} | Update a ModelVersion
 *ModelRegistryServiceApi* | [**update_registered_model**](mr_openapi/docs/ModelRegistryServiceApi.md#update_registered_model) | **PATCH** /api/model_registry/v1alpha3/registered_models/{registeredmodelId} | Update a RegisteredModel
 *ModelRegistryServiceApi* | [**update_serving_environment**](mr_openapi/docs/ModelRegistryServiceApi.md#update_serving_environment) | **PATCH** /api/model_registry/v1alpha3/serving_environments/{servingenvironmentId} | Update a ServingEnvironment
+*ModelRegistryServiceApi* | [**upsert_model_version_artifact**](mr_openapi/docs/ModelRegistryServiceApi.md#upsert_model_version_artifact) | **POST** /api/model_registry/v1alpha3/model_versions/{modelversionId}/artifacts | Upsert an Artifact in a ModelVersion
 
 
 ## Documentation For Models