Keep online endpoint/deployment operations pinned to workspace scope for registry-backed clients

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

@@ -223,14 +223,14 @@ def __init__(
         workspace_location = None
         workspace_id = None
         registry_reference = kwargs.pop("registry_reference", None)
+        workspace_reference = kwargs.pop("workspace_reference", None)
         if registry_name or registry_reference:
             # get the workspace location here if workspace_reference is provided
             self._ws_operation_scope = OperationScope(
                 str(subscription_id),
                 str(resource_group_name),
                 workspace_name,
             )
-            workspace_reference = kwargs.pop("workspace_reference", None)
             if workspace_reference or registry_reference:
                 ws_ops = WorkspaceOperations(
                     OperationScope(str(subscription_id), str(resource_group_name), workspace_reference),
@@ -273,6 +273,17 @@ def __init__(
             workspace_id,
             workspace_location,
         )
+        self._online_operation_scope = (
+            OperationScope(
+                self._ws_operation_scope.subscription_id,
+                self._ws_operation_scope.resource_group_name,
+                workspace_name,
+                workspace_id=workspace_id,
+                workspace_location=workspace_location,
+            )
+            if workspace_reference or registry_reference
+            else self._operation_scope
+        )
 
         # Cannot send multiple base_url as azure-cli sets the base_url automatically.
         kwargs.pop("base_url", None)
@@ -319,6 +330,12 @@ def __init__(
             base_url=base_url,
             **kwargs,
         )
+        self._service_client_02_2022_preview_online = ServiceClient022022Preview(
+            subscription_id=self._online_operation_scope._subscription_id,
+            credential=self._credential,
+            base_url=base_url,
+            **kwargs,
+        )
 
         self._service_client_05_2022 = ServiceClient052022(
             credential=self._credential,
@@ -420,6 +437,12 @@ def __init__(
             base_url=base_url,
             **kwargs,
         )
+        self._service_client_04_2023_preview_online = ServiceClient042023Preview(
+            credential=self._credential,
+            subscription_id=self._online_operation_scope._subscription_id,
+            base_url=base_url,
+            **kwargs,
+        )
 
         self._service_client_06_2023_preview = ServiceClient062023Preview(
             credential=self._credential,
@@ -604,9 +627,9 @@ def __init__(
         self._local_endpoint_helper = _LocalEndpointHelper(requests_pipeline=self._requests_pipeline)
         self._local_deployment_helper = _LocalDeploymentHelper(self._operation_container)
         self._online_endpoints = OnlineEndpointOperations(
-            self._ws_operation_scope if registry_reference else self._operation_scope,
+            self._online_operation_scope,
             self._operation_config,
-            self._service_client_02_2022_preview,
+            self._service_client_02_2022_preview_online,
             self._operation_container,
             self._local_endpoint_helper,
             self._credential,
@@ -625,9 +648,9 @@ def __init__(
         self._operation_container.add(AzureMLResourceType.BATCH_ENDPOINT, self._batch_endpoints)
         self._operation_container.add(AzureMLResourceType.ONLINE_ENDPOINT, self._online_endpoints)
         self._online_deployments = OnlineDeploymentOperations(
-            self._ws_operation_scope if registry_reference else self._operation_scope,
+            self._online_operation_scope,
             self._operation_config,
-            self._service_client_04_2023_preview,
+            self._service_client_04_2023_preview_online,
             self._operation_container,
             self._local_deployment_helper,
             self._credential,

sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_ml_client.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from contextlib import ExitStack
 from unittest import mock
 from unittest.mock import Mock, patch
 
@@ -471,6 +472,98 @@ def test_ml_client_with_both_workspace_registry_names_throws(
         message = exception.value.args[0]
         assert message == "Both workspace_name and registry_name cannot be used together, for the ml_client."
 
+    def test_registry_backed_online_operations_keep_workspace_scope(self, mock_credential) -> None:
+        workspace_sub = "workspace-sub"
+        workspace_rg = "workspace-rg"
+        workspace_name = "workspace-name"
+        registry_sub = "registry-sub"
+        registry_rg = "registry-rg"
+        captured = {}
+        workspace_details = Mock(location="eastus", _workspace_id="workspace-id")
+
+        def make_fake_operation(name):
+            class FakeOperation:
+                def __init__(self, *args, **kwargs):
+                    captured[name] = {
+                        "args": args,
+                        "kwargs": kwargs,
+                        "scope": args[0] if args else kwargs.get("operation_scope"),
+                    }
+
+                def get(self, *args, **kwargs):
+                    return workspace_details
+
+            return FakeOperation
+
+        operation_names = [
+            "WorkspaceOperations",
+            "WorkspaceOutboundRuleOperations",
+            "RegistryOperations",
+            "WorkspaceConnectionsOperations",
+            "CapabilityHostsOperations",
+            "ComputeOperations",
+            "DatastoreOperations",
+            "ModelOperations",
+            "EvaluatorOperations",
+            "CodeOperations",
+            "EnvironmentOperations",
+            "OnlineEndpointOperations",
+            "BatchEndpointOperations",
+            "OnlineDeploymentOperations",
+            "BatchDeploymentOperations",
+            "DeploymentTemplateOperations",
+            "DataOperations",
+            "ComponentOperations",
+            "JobOperations",
+            "ScheduleOperations",
+            "IndexOperations",
+            "FeatureStoreOperations",
+            "FeatureSetOperations",
+            "FeatureStoreEntityOperations",
+            "AzureOpenAIDeploymentOperations",
+            "ServerlessEndpointOperations",
+            "MarketplaceSubscriptionOperations",
+        ]
+        with ExitStack() as stack:
+            for operation_name in operation_names:
+                stack.enter_context(
+                    patch(f"azure.ai.ml._ml_client.{operation_name}", make_fake_operation(operation_name))
+                )
+
+            stack.enter_context(patch("azure.ai.ml._ml_client.get_deployments_operation", return_value=Mock()))
+            stack.enter_context(
+                patch(
+                    "azure.ai.ml._ml_client.get_registry_client",
+                    return_value=(Mock(), registry_rg, registry_sub, Mock(), Mock()),
+                )
+            )
+
+            MLClient(
+                credential=mock_credential,
+                subscription_id=workspace_sub,
+                resource_group_name=workspace_rg,
+                registry_name="test-registry",
+                workspace_reference=workspace_name,
+            )
+
+        online_endpoint_scope = captured["OnlineEndpointOperations"]["scope"]
+        online_deployment_scope = captured["OnlineDeploymentOperations"]["scope"]
+        model_scope = captured["ModelOperations"]["scope"]
+        data_scope = captured["DataOperations"]["scope"]
+
+        assert online_endpoint_scope.subscription_id == workspace_sub
+        assert online_endpoint_scope.resource_group_name == workspace_rg
+        assert online_endpoint_scope.workspace_name == workspace_name
+        assert online_deployment_scope.subscription_id == workspace_sub
+        assert online_deployment_scope.resource_group_name == workspace_rg
+        assert online_deployment_scope.workspace_name == workspace_name
+        assert model_scope.subscription_id == registry_sub
+        assert model_scope.resource_group_name == registry_rg
+        assert data_scope.subscription_id == registry_sub
+        assert data_scope.resource_group_name == registry_rg
+        assert online_endpoint_scope is not model_scope
+        assert online_deployment_scope is not data_scope
+
     def test_ml_client_with_cli_config(self, mock_credential):
         # This cloud config should not work and it should NOT overwrite the hardcoded AzureCloud
         kwargs = {