py: client: add update method for supported types

clients/python/README.md

@@ -37,6 +37,12 @@ model = registry.get_registered_model("my-model")
 version = registry.get_model_version("my-model", "2.0.0")
 
 experiment = registry.get_model_artifact("my-model", "2.0.0")
+
+# change is not reflected on pushed model version
+version.description = "Updated model version"
+
+# you can update it using
+registry.update(version)
 ```
 
 ### Importing from S3

clients/python/src/model_registry/_client.py

@@ -4,7 +4,7 @@
 
 import os
 from pathlib import Path
-from typing import Any, get_args
+from typing import Any, TypeVar, Union, get_args
 from warnings import warn
 
 from .core import ModelRegistryAPIClient
@@ -18,6 +18,9 @@
     SupportedTypes,
 )
 
+ModelTypes = Union[RegisteredModel, ModelVersion, ModelArtifact]
+TModel = TypeVar("TModel", bound=ModelTypes)
+
 
 class ModelRegistry:
     """Model registry client."""
@@ -191,6 +194,20 @@ def register_model(
 
         return rm
 
+    def update(self, model: TModel) -> TModel:
+        """Update a model."""
+        if not model.id:
+            msg = "Model must have an ID"
+            raise StoreError(msg)
+        if not isinstance(model, get_args(ModelTypes)):
+            msg = f"Model must be one of {get_args(ModelTypes)}"
+            raise StoreError(msg)
+        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))
+
     def register_hf_model(
         self,
         repo: str,

clients/python/tests/test_client.py

@@ -77,6 +77,35 @@ def test_register_existing_version(client: ModelRegistry):
         client.register_model(**params)
 
 
+@pytest.mark.e2e
+async def test_update_models(client: ModelRegistry):
+    name = "test_model"
+    version = "1.0.0"
+    rm = client.register_model(
+        name,
+        "s3",
+        model_format_name="test_format",
+        model_format_version="test_version",
+        version=version,
+    )
+    assert rm.id
+
+    mr_api = client._api
+    mv = await mr_api.get_model_version_by_params(rm.id, version)
+    assert mv
+    assert mv.id
+    ma = await mr_api.get_model_artifact_by_params(name, mv.id)
+    assert ma
+
+    new_description = "updated description"
+    rm.description = new_description
+    mv.description = new_description
+    ma.description = new_description
+    assert client.update(rm).description == new_description
+    assert client.update(mv).description == new_description
+    assert client.update(ma).description == new_description
+
+
 @pytest.mark.e2e
 async def test_get(client: ModelRegistry):
     name = "test_model"