Add method for registry model create/update with system metadata (Azure#40064)

chaoyu-msft · web-flow · commit 5b61a9534435 · 2025-04-02T09:09:26.000-04:00
* Add method for registry model create/update with system metadata

* add e2e test

* fix test

* fix test recording

* fix test recording

* fix test recording

* fix tests

* fix mypy

* run black

* fix pylint

* move request method

* refactor code

* fix tests

* fix compiling

* update recording

* fix black
diff --git a/sdk/ml/azure-ai-ml/assets.json b/sdk/ml/azure-ai-ml/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_305b890d5b"
+  "Tag": "python/ml/azure-ai-ml_9b9dd0f330"
 }
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py
@@ -21,6 +21,7 @@
     _set_cloud,
 )
 from azure.ai.ml._file_utils.file_utils import traverse_up_path_and_find_file
+from azure.ai.ml._restclient.model_dataplane import AzureMachineLearningWorkspaces as ServiceClientModelDataPlane
 from azure.ai.ml._restclient.v2020_09_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient092020DataplanePreview,
 )
@@ -260,6 +261,13 @@ def __init__(
             if not workspace_name:
                 workspace_name = workspace_reference
 
+            self._service_client_model_dataplane = ServiceClientModelDataPlane(
+                credential=self._credential,
+                subscription_id=subscription_id,
+                base_url=self._service_client_10_2021_dataplanepreview._client._base_url,
+                **kwargs,
+            )
+
         self._operation_scope = OperationScope(
             str(subscription_id),
             str(resource_group_name),
@@ -577,6 +585,7 @@ def __init__(
                 else self._service_client_08_2023_preview
             ),
             self._datastores,
+            (self._service_client_model_dataplane if registry_name or registry_reference else None),
             self._operation_container,
             requests_pipeline=self._requests_pipeline,
             control_plane_client=self._service_client_08_2023_preview,
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/model_dataplane/operations/_models_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/model_dataplane/operations/_models_operations.py
@@ -12,9 +12,11 @@
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError, ResourceExistsError, ResourceNotFoundError, map_error
 from azure.core.pipeline import PipelineResponse
 from azure.core.pipeline.transport import HttpResponse
+from azure.core.polling import LROPoller, NoPolling
 from azure.core.rest import HttpRequest
 from azure.core.tracing.decorator import distributed_trace
 from azure.mgmt.core.exceptions import ARMErrorFormat
+from azure.mgmt.core.polling.arm_polling import ARMPolling
 from msrest import Serializer
 
 from .. import models as _models
@@ -457,6 +459,50 @@ def build_deployment_settings_request(
         **kwargs
     )
 
+
+def build_create_or_update_request_initial(
+    name,  # type: str
+    version,  # type: str
+    subscription_id,  # type: str
+    resource_group_name,  # type: str
+    registry_name,  # type: str
+    **kwargs  # type: Any
+):
+    # type: (...) -> HttpRequest
+    api_version = kwargs.pop('api_version', "2021-10-01-dataplanepreview")  # type: str
+    content_type = kwargs.pop('content_type', None)  # type: Optional[str]
+
+    accept = "application/json"
+    # Construct URL
+    url = kwargs.pop("template_url", '/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.MachineLearningServices/registries/{registryName}/models/{name}/versions/{version}')
+    path_format_arguments = {
+        "name": _SERIALIZER.url("name", name, 'str', pattern=r'^(?![\-_.])[a-zA-Z0-9\-_.]{1,255}(?<!\.)$'),
+        "version": _SERIALIZER.url("version", version, 'str'),
+        "subscriptionId": _SERIALIZER.url("subscription_id", subscription_id, 'str', min_length=1),
+        "resourceGroupName": _SERIALIZER.url("resource_group_name", resource_group_name, 'str', max_length=90, min_length=1),
+        "registryName": _SERIALIZER.url("registry_name", registry_name, 'str'),
+    }
+
+    url = _format_url_section(url, **path_format_arguments)
+
+    # Construct parameters
+    query_parameters = kwargs.pop("params", {})  # type: Dict[str, Any]
+    query_parameters['api-version'] = _SERIALIZER.query("api_version", api_version, 'str')
+
+    # Construct headers
+    header_parameters = kwargs.pop("headers", {})  # type: Dict[str, Any]
+    if content_type is not None:
+        header_parameters['Content-Type'] = _SERIALIZER.header("content_type", content_type, 'str')
+    header_parameters['Accept'] = _SERIALIZER.header("accept", accept, 'str')
+
+    return HttpRequest(
+        method="PUT",
+        url=url,
+        params=query_parameters,
+        headers=header_parameters,
+        **kwargs
+    )
+
 # fmt: on
 class ModelsOperations(object):
     """ModelsOperations operations.
@@ -825,6 +871,107 @@ def create_unregistered_output_model(
     create_unregistered_output_model.metadata = {'url': '/modelregistry/v1.0/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.MachineLearningServices/workspaces/{workspaceName}/models/createUnregisteredOutput'}  # type: ignore
 
 
+    # this method is only used for model create/update with system metadata
+    @distributed_trace
+    def begin_create_or_update_model_with_system_metadata(
+        self,
+        subscription_id, # type: str
+        name,  # type: str
+        version,  # type: str
+        resource_group_name,  # type: str
+        registry_name,  # type: str
+        body,  # type: "ModelVersion"
+        **kwargs,  # type: Any
+    ):
+        # type: (...) -> LROPoller[None]
+        """Create or update version.
+
+        Create or update version.
+
+        :param name: Container name.
+        :type name: str
+        :param version: Version identifier.
+        :type version: str
+        :param resource_group_name: The name of the resource group. The name is case insensitive.
+        :type resource_group_name: str
+        :param registry_name: Name of Azure Machine Learning registry.
+        :type registry_name: str
+        :param body: Version entity to create or update.
+        :type body: ~azure.mgmt.machinelearningservices.models.ModelVersionData
+        :keyword api_version: Api Version. The default value is "2021-10-01-dataplanepreview". Note
+         that overriding this default value may result in unsupported behavior.
+        :paramtype api_version: str
+        :keyword callable cls: A custom type or function that will be passed the direct response
+        :keyword str continuation_token: A continuation token to restart a poller from a saved state.
+        :keyword polling: By default, your polling method will be ARMPolling. Pass in False for this
+         operation to not poll, or pass in your own initialized polling object for a personal polling
+         strategy.
+        :paramtype polling: bool or ~azure.core.polling.PollingMethod
+        :keyword int polling_interval: Default waiting time between two polls for LRO operations if no
+         Retry-After header is present.
+        :return: An instance of LROPoller that returns either None or the result of cls(response)
+        :rtype: ~azure.core.polling.LROPoller[None]
+        :raises: ~azure.core.exceptions.HttpResponseError
+        """
+
+        error_map = {401: ClientAuthenticationError, 404: ResourceNotFoundError, 409: ResourceExistsError}
+        error_map.update(kwargs.pop("error_map", {}))
+
+        _json = self._serialize.body(body, "ModelVersionData")
+        _json["properties"]["system_metadata"] = body.properties.system_metadata
+
+        request = build_create_or_update_request_initial(
+            name=name,
+            version=version,
+            subscription_id=subscription_id,
+            resource_group_name=resource_group_name,
+            registry_name=registry_name,
+            json=_json,
+            template_url=self.begin_create_or_update_model_with_system_metadata.metadata["url"],
+        )
+        request = _convert_request(request)
+        request.url = self._client.format_url(request.url)
+
+        pipeline_response = self._client._pipeline.run(request, stream=False, **kwargs)
+        response = pipeline_response.http_response
+
+        if response.status_code not in [202]:
+            map_error(status_code=response.status_code, response=response, error_map=error_map)
+            raise HttpResponseError(response=response, error_format=ARMErrorFormat)
+
+        response_headers = {}
+        response_headers["x-ms-async-operation-timeout"] = self._deserialize(
+            "duration", response.headers.get("x-ms-async-operation-timeout")
+        )
+        response_headers["Location"] = self._deserialize(
+            "str", response.headers.get("Location")
+        )
+        response_headers["Retry-After"] = self._deserialize(
+            "int", response.headers.get("Retry-After")
+        )
+
+        cls = kwargs.pop("cls", None)
+
+        def get_long_running_output(pipeline_response):
+            if cls:
+                return cls(pipeline_response, None, {})
+            return None
+
+        polling = kwargs.pop("polling", True)
+        if polling is True:
+            lro_delay = kwargs.pop("polling_interval", self._config.polling_interval)
+            polling_method = ARMPolling(lro_delay, **kwargs)
+        elif polling is False:
+            polling_method = NoPolling()
+        else:
+            polling_method = polling
+
+        return LROPoller(
+            self._client, pipeline_response, get_long_running_output, polling_method
+        )
+
+    begin_create_or_update_model_with_system_metadata.metadata = {"url": "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.MachineLearningServices/registries/{registryName}/models/{name}/versions/{version}"}  # type: ignore
+
     @distributed_trace
     def batch_get_resolved_uris(
         self,
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py
@@ -22,6 +22,8 @@
     CHANGED_ASSET_PATH_MSG_NO_PERSONAL_DATA,
 )
 from azure.ai.ml._exception_helper import log_and_raise_error
+
+from azure.ai.ml._restclient.model_dataplane import AzureMachineLearningWorkspaces as ServiceClientModelDataPlane
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient102021Dataplane,
 )
@@ -105,13 +107,16 @@ def __init__(
         operation_config: OperationConfig,
         service_client: Union[ServiceClient082023Preview, ServiceClient102021Dataplane],
         datastore_operations: DatastoreOperations,
+        service_client_model_dataplane: ServiceClientModelDataPlane = None,
         all_operations: Optional[OperationsContainer] = None,
         **kwargs,
     ):
         super(ModelOperations, self).__init__(operation_scope, operation_config)
         ops_logger.update_filter()
         self._model_versions_operation = service_client.model_versions
         self._model_container_operation = service_client.model_containers
+        if service_client_model_dataplane is not None:
+            self._model_dataplane_operation = service_client_model_dataplane.models
         self._service_client = service_client
         self._datastore_operation = datastore_operations
         self._all_operations = all_operations
@@ -251,14 +256,14 @@ def create_or_update(  # type: ignore
             model_version_resource = model._to_rest_object()
             auto_increment_version = model._auto_increment_version
             try:
+                cont_token = self._scope_kwargs.pop("continuation_token", None)  # type: Optional[str]
                 result = (
-                    self._model_versions_operation.begin_create_or_update(
-                        name=name,
-                        version=version,
-                        body=model_version_resource,
-                        registry_name=self._registry_name,
-                        **self._scope_kwargs,
-                    ).result()
+                    self._begin_create_or_update_registry_model(
+                        name,
+                        version,
+                        model_version_resource,
+                        cont_token,
+                    )
                     if self._registry_name
                     else self._model_versions_operation.create_or_update(
                         name=name,
@@ -295,6 +300,31 @@ def create_or_update(  # type: ignore
             else:
                 raise ex
 
+    def _begin_create_or_update_registry_model(self, name, version, model_version_resource, cont_token):
+        # if continuation token is None and system_metadata attribute found
+        # we need to send the system_metadata values in the request
+        return (
+            self._model_dataplane_operation.begin_create_or_update_model_with_system_metadata(
+                subscription_id=self._operation_scope._subscription_id,
+                name=str(name),
+                version=str(version),
+                body=model_version_resource,
+                registry_name=self._registry_name,
+                **self._scope_kwargs,
+            ).result()
+            if self._registry_name
+            and cont_token is None
+            and hasattr(model_version_resource.properties, "system_metadata")
+            and self._model_dataplane_operation is not None
+            else self._model_versions_operation.begin_create_or_update(
+                name=name,
+                version=version,
+                body=model_version_resource,
+                registry_name=self._registry_name,
+                **self._scope_kwargs,
+            ).result()
+        )
+
     def _get(self, name: str, version: Optional[str] = None) -> ModelVersion:  # name:latest
         if version:
             return (
diff --git a/sdk/ml/azure-ai-ml/tests/model/e2etests/test_model.py b/sdk/ml/azure-ai-ml/tests/model/e2etests/test_model.py
@@ -13,6 +13,7 @@
 from azure.ai.ml._restclient.v2022_05_01.models import ListViewType
 from azure.ai.ml.constants._common import LONG_URI_REGEX_FORMAT
 from azure.ai.ml.entities._assets import Model
+from azure.core.exceptions import HttpResponseError
 from azure.core.paging import ItemPaged
 
 
@@ -97,6 +98,19 @@ def test_crud_model_with_stage(self, client: MLClient, randstr: Callable[[], str
         model_stage_list = [m.stage for m in model_list if m is not None]
         assert model.stage in model_stage_list
 
+    def test_crud_model_with_system_metadata(self, registry_client: MLClient, randstr: Callable[[], str]) -> None:
+        path = Path("./tests/test_configs/model/model_with_system_metadata.yml")
+        model_name = randstr("model_with_system_metadata")
+
+        model = load_model(path)
+        model.name = model_name
+
+        model = registry_client.models.create_or_update(model)
+        assert model.name == model_name
+        assert model.version == "1"
+        assert model.description == "this is my test model with system metadata"
+        assert model.type == "mlflow_model"
+
     def test_list_no_name(self, client: MLClient) -> None:
         models = client.models.list()
         assert isinstance(models, Iterator)
diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/model/model_with_system_metadata.yml b/sdk/ml/azure-ai-ml/tests/test_configs/model/model_with_system_metadata.yml
@@ -1,9 +1,8 @@
 name: my-production-model
 type: mlflow_model
 path: ./lightgbm_mlflow_model
-version: 3
-description: "this is my test model with stage"
-stage: "Production"
+version: 1
+description: "this is my test model with system metadata"
 tags:
   foo: bar
   abc: 123

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"AssetsRepo": "Azure/azure-sdk-assets",`
`3`	`3`	`"AssetsRepoPrefixPath": "python",`
`4`	`4`	`"TagPrefix": "python/ml/azure-ai-ml",`
`5`		`- "Tag": "python/ml/azure-ai-ml_305b890d5b"`
	`5`	`+ "Tag": "python/ml/azure-ai-ml_9b9dd0f330"`
`6`	`6`	`}`