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 index e227029d10e5..8bd5824aaccc 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py @@ -223,6 +223,7 @@ 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( @@ -230,7 +231,6 @@ def __init__( 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, diff --git a/sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_ml_client.py b/sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_ml_client.py index aa7f3a55942f..16f7e2b33433 100644 --- a/sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_ml_client.py +++ b/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 = {