Fix MLClient Registry base_url value for model operations (Azure#40375)

* Fix MLClient Registry base_url value for model operations

* add missing changes

* refactor the model client init code

* fix black

* refactor code again

* fix pylint

Author: chaoyu-msft

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -21,7 +21,6 @@
     _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,
 )
@@ -252,22 +251,17 @@ def __init__(
                 self._service_client_10_2021_dataplanepreview,
                 resource_group_name,
                 subscription_id,
+                self._service_client_model_dataplane,
             ) = get_registry_client(
                 self._credential,
                 registry_name if registry_name else registry_reference,
                 workspace_location,
                 **kwargs,
             )
+
             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),

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_registry_utils.py

@@ -8,6 +8,7 @@
 from typing_extensions import Literal
 
 from azure.ai.ml._azure_environments import _get_default_cloud_name, _get_registry_discovery_endpoint_from_metadata
+from azure.ai.ml._restclient.model_dataplane import AzureMachineLearningWorkspaces as ServiceClientModelDataPlane
 from azure.ai.ml._restclient.registry_discovery import AzureMachineLearningWorkspaces as ServiceClientRegistryDiscovery
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import AzureMachineLearningWorkspaces
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview.models import (
@@ -63,7 +64,13 @@ def get_registry_service_client(self) -> AzureMachineLearningWorkspaces:
             base_url=self._base_url,
             **self.kwargs,
         )
-        return service_client_10_2021_dataplanepreview
+        service_model_client_10_2021_dataplanepreview = ServiceClientModelDataPlane(
+            credential=self.credential,
+            subscription_id=self._subscription_id,
+            base_url=self._base_url,
+            **self.kwargs,
+        )
+        return service_client_10_2021_dataplanepreview, service_model_client_10_2021_dataplanepreview
 
     @property
     def subscription_id(self) -> str:
@@ -207,10 +214,17 @@ def get_registry_client(credential, registry_name, workspace_location: Optional[
     registry_discovery = RegistryDiscovery(
         credential, registry_name, service_client_registry_discovery_client, **kwargs
     )
-    service_client_10_2021_dataplanepreview = registry_discovery.get_registry_service_client()
+    service_client_10_2021_dataplanepreview, service_model_client_10_2021_dataplanepreview = (
+        registry_discovery.get_registry_service_client()
+    )
     subscription_id = registry_discovery.subscription_id
     resource_group_name = registry_discovery.resource_group
-    return service_client_10_2021_dataplanepreview, resource_group_name, subscription_id
+    return (
+        service_client_10_2021_dataplanepreview,
+        resource_group_name,
+        subscription_id,
+        service_model_client_10_2021_dataplanepreview,
+    )
 
 
 def _check_region_fqdn(workspace_region, response):

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py

@@ -853,7 +853,9 @@ def _set_registry_client(self, registry_name: str) -> Generator:
         data_versions_operation_ = self._operation
 
         try:
-            _client, _rg, _sub = get_registry_client(self._service_client._config.credential, registry_name)
+            _client, _rg, _sub, _model_client = get_registry_client(
+                self._service_client._config.credential, registry_name
+            )
             self._operation_scope.registry_name = registry_name
             self._operation_scope._resource_group_name = _rg
             self._operation_scope._subscription_id = _sub

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py

@@ -548,7 +548,9 @@ def _set_registry_client(self, registry_name: str) -> Generator:
         environment_versions_operation_ = self._version_operations
 
         try:
-            _client, _rg, _sub = get_registry_client(self._service_client._config.credential, registry_name)
+            _client, _rg, _sub, _model_client = get_registry_client(
+                self._service_client._config.credential, registry_name
+            )
             self._operation_scope.registry_name = registry_name
             self._operation_scope._resource_group_name = _rg
             self._operation_scope._subscription_id = _sub

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py

@@ -686,7 +686,9 @@ def _set_registry_client(self, registry_name: str) -> Generator:
         model_versions_operation_ = self._model_versions_operation
 
         try:
-            _client, _rg, _sub = get_registry_client(self._service_client._config.credential, registry_name)
+            _client, _rg, _sub, _model_client = get_registry_client(
+                self._service_client._config.credential, registry_name
+            )
             self._operation_scope.registry_name = registry_name
             self._operation_scope._resource_group_name = _rg
             self._operation_scope._subscription_id = _sub