Adding cloud configuration via CLI keyword args (Azure#28780)

* Adding cloud configuration via CLI keyword args

* Fixing a couple comments

* updated sdk changes to accommodate the cli differences and changed function to throw exception instead of returning false

* fixed test and lint error

* updated tests based on PR suggestions. Fixed lint errors. Removed changelog updates

* Some last minute fixes

* Small test fix

* Update CHANGELOG.md

* retrigger checks

* Updating some formatting for the Black test

* Updating some formatting for the Black test again

---------

Co-authored-by: Ronald Shaw <[email protected]>
Co-authored-by: Ankit Singhal <[email protected]>

Author: 3 people

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -3,7 +3,6 @@
 ## 1.5.0 (Unreleased)
 
 ### Features Added
-
 - Added support for `tags` on Compute Resources.
 
 ### Bugs Fixed

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

@@ -33,6 +33,10 @@ class EndpointURLS:  # pylint: disable=too-few-public-methods,no-init
     REGISTRY_DISCOVERY_ENDPOINT = "registry_discovery_endpoint"
 
 
+class CloudArgumentKeys:
+    CLOUD_METADATA = "cloud_metadata"
+
+
 _environments = {
     AzureEnvironments.ENV_DEFAULT: {
         EndpointURLS.AZURE_PORTAL_ENDPOINT: "https://portal.azure.com/",
@@ -60,8 +64,6 @@ class EndpointURLS:  # pylint: disable=too-few-public-methods,no-init
     },
 }
 
-_requests_pipeline = None
-
 
 def _get_cloud(cloud: str):
     if cloud in _environments:
@@ -120,7 +122,6 @@ def _get_base_url_from_metadata(cloud_name: Optional[str] = None, is_local_mfe:
     base_url = None
     if is_local_mfe:
         base_url = _get_mfe_url_override()
-
     if base_url is None:
         cloud_details = _get_cloud_details(cloud_name)
         base_url = cloud_details.get(EndpointURLS.RESOURCE_MANAGER_ENDPOINT).strip("/")
@@ -224,7 +225,8 @@ def _get_registry_discovery_url(cloud, cloud_suffix=""):
         return _environments[cloud_name].registry_url
 
     registry_discovery_region = os.environ.get(
-        ArmConstants.REGISTRY_DISCOVERY_REGION_ENV_NAME, ArmConstants.REGISTRY_DISCOVERY_DEFAULT_REGION
+        ArmConstants.REGISTRY_DISCOVERY_REGION_ENV_NAME,
+        ArmConstants.REGISTRY_DISCOVERY_DEFAULT_REGION,
     )
     registry_discovery_region_default = "https://{}{}.api.azureml.{}/".format(
         cloud_name.lower(), registry_discovery_region, cloud_suffix
@@ -244,7 +246,10 @@ def _get_clouds_by_metadata_url(metadata_url):
         with client.send_request(HttpRequest("GET", metadata_url)) as meta_response:
             arm_cloud_dict = meta_response.json()
             cli_cloud_dict = _convert_arm_to_cli(arm_cloud_dict)
-            module_logger.debug("Finish : Loading cloud metadata from the url specified by %s", metadata_url)
+            module_logger.debug(
+                "Finish : Loading cloud metadata from the url specified by %s",
+                metadata_url,
+            )
             return cli_cloud_dict
     except Exception as ex:  # pylint: disable=broad-except
         module_logger.warning(
@@ -283,3 +288,20 @@ def _convert_arm_to_cli(arm_cloud_metadata):
             module_logger.warning("Property on cloud not found in arm cloud metadata: %s", ex)
             continue
     return cli_cloud_metadata_dict
+
+
+def _add_cloud_to_environments(kwargs):
+    cloud_name = kwargs["cloud"]
+    if cloud_name in _environments:
+        raise AttributeError(f"Cannot overwrite existing cloud: {cloud_name}")
+    cloud_metadata = kwargs[CloudArgumentKeys.CLOUD_METADATA]
+    if cloud_metadata is None:
+        raise LookupError(f"{CloudArgumentKeys.CLOUD_METADATA} not present in kwargs, no environment to add!")
+    _environments[kwargs["cloud"]] = {
+        EndpointURLS.AZURE_PORTAL_ENDPOINT: cloud_metadata[EndpointURLS.AZURE_PORTAL_ENDPOINT],
+        EndpointURLS.RESOURCE_MANAGER_ENDPOINT: cloud_metadata[EndpointURLS.RESOURCE_MANAGER_ENDPOINT],
+        EndpointURLS.ACTIVE_DIRECTORY_ENDPOINT: cloud_metadata[EndpointURLS.ACTIVE_DIRECTORY_ENDPOINT],
+        EndpointURLS.AML_RESOURCE_ID: cloud_metadata[EndpointURLS.AML_RESOURCE_ID],
+        EndpointURLS.STORAGE_ENDPOINT: cloud_metadata[EndpointURLS.STORAGE_ENDPOINT],
+        EndpointURLS.REGISTRY_DISCOVERY_ENDPOINT: cloud_metadata[EndpointURLS.REGISTRY_DISCOVERY_ENDPOINT],
+    }

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

@@ -16,25 +16,47 @@
 from azure.core.polling import LROPoller
 
 from azure.ai.ml._azure_environments import (
+    CloudArgumentKeys,
     _get_base_url_from_metadata,
     _get_cloud_information_from_metadata,
     _get_default_cloud_name,
     _get_registry_discovery_endpoint_from_metadata,
     _set_cloud,
+    _add_cloud_to_environments,
 )
 from azure.ai.ml._file_utils.file_utils import traverse_up_path_and_find_file
-from azure.ai.ml._restclient.registry_discovery import AzureMachineLearningWorkspaces as ServiceClientRegistryDiscovery
+from azure.ai.ml._restclient.registry_discovery import (
+    AzureMachineLearningWorkspaces as ServiceClientRegistryDiscovery,
+)
 from azure.ai.ml._restclient.v2020_09_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient092020DataplanePreview,
 )
-from azure.ai.ml._restclient.v2022_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012022Preview
-from azure.ai.ml._restclient.v2022_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022022Preview
-from azure.ai.ml._restclient.v2022_05_01 import AzureMachineLearningWorkspaces as ServiceClient052022
-from azure.ai.ml._restclient.v2022_10_01 import AzureMachineLearningWorkspaces as ServiceClient102022
-from azure.ai.ml._restclient.v2022_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102022Preview
-from azure.ai.ml._restclient.v2022_12_01_preview import AzureMachineLearningWorkspaces as ServiceClient122022Preview
-from azure.ai.ml._restclient.v2023_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022023Preview
-from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationsContainer, OperationScope
+from azure.ai.ml._restclient.v2022_01_01_preview import (
+    AzureMachineLearningWorkspaces as ServiceClient012022Preview,
+)
+from azure.ai.ml._restclient.v2022_02_01_preview import (
+    AzureMachineLearningWorkspaces as ServiceClient022022Preview,
+)
+from azure.ai.ml._restclient.v2022_05_01 import (
+    AzureMachineLearningWorkspaces as ServiceClient052022,
+)
+from azure.ai.ml._restclient.v2022_10_01 import (
+    AzureMachineLearningWorkspaces as ServiceClient102022,
+)
+from azure.ai.ml._restclient.v2022_10_01_preview import (
+    AzureMachineLearningWorkspaces as ServiceClient102022Preview,
+)
+from azure.ai.ml._restclient.v2022_12_01_preview import (
+    AzureMachineLearningWorkspaces as ServiceClient122022Preview,
+)
+from azure.ai.ml._restclient.v2023_02_01_preview import (
+    AzureMachineLearningWorkspaces as ServiceClient022023Preview,
+)
+from azure.ai.ml._scope_dependent_operations import (
+    OperationConfig,
+    OperationsContainer,
+    OperationScope,
+)
 
 # from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
 from azure.ai.ml._user_agent import USER_AGENT
@@ -166,7 +188,21 @@ def __init__(
         show_progress = kwargs.pop("show_progress", True)
         self._operation_config = OperationConfig(show_progress=show_progress)
 
-        cloud_name = kwargs.get("cloud", _get_default_cloud_name())
+        if "cloud" in kwargs:
+            cloud_name = kwargs["cloud"]
+            if CloudArgumentKeys.CLOUD_METADATA in kwargs:
+                try:
+                    _add_cloud_to_environments(kwargs)
+                except AttributeError as e:
+                    module_logger.debug("Cloud already exists: %s", e)
+                except LookupError as e:
+                    module_logger.debug("Missing keyword: %s", e)
+            else:
+                module_logger.debug("%s key not found in kwargs", CloudArgumentKeys.CLOUD_METADATA)
+        else:
+            module_logger.debug("cloud key not found in kwargs")
+            cloud_name = _get_default_cloud_name()
+
         self._cloud = cloud_name
         _set_cloud(cloud_name)
         if "cloud" not in kwargs:
@@ -187,7 +223,10 @@ def __init__(
                 credential=self._credential, base_url=base_url, **kwargs_registry
             )
             registry_discovery = RegistryDiscovery(
-                self._credential, registry_name, self._service_client_registry_discovery_client, **kwargs_registry
+                self._credential,
+                registry_name,
+                self._service_client_registry_discovery_client,
+                **kwargs_registry,
             )
             self._service_client_10_2021_dataplanepreview = registry_discovery.get_registry_service_client()
             subscription_id = registry_discovery.subscription_id
@@ -802,7 +841,15 @@ def create_or_update(
     # R = valid inputs/outputs for begin_create_or_update
     # Each entry here requires a registered _begin_create_or_update function below
     R = TypeVar(
-        "R", Workspace, Registry, Compute, OnlineDeployment, OnlineEndpoint, BatchDeployment, BatchEndpoint, JobSchedule
+        "R",
+        Workspace,
+        Registry,
+        Compute,
+        OnlineDeployment,
+        OnlineEndpoint,
+        BatchDeployment,
+        BatchEndpoint,
+        JobSchedule,
     )
 
     def begin_create_or_update(

sdk/ml/azure-ai-ml/samples/ml_samples_cloud_configurations.py

@@ -0,0 +1,83 @@
+# coding: utf-8
+
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+"""
+FILE: ml_samples_cloud_configurations.py
+DESCRIPTION:
+    These samples demonstrate different ways to configure clouds for the MLClient.
+USAGE:
+    python ml_samples_cloud_configurations.py
+
+"""
+
+import os
+from azure.ai.ml import MLClient
+from azure.ai.ml.constants._common import ArmConstants
+from azure.ai.ml.constants._common import AZUREML_CLOUD_ENV_NAME
+from azure.identity import AzureAuthorityHosts, DefaultAzureCredential
+
+
+class CloudConfigurationOptions(object):
+    def ml_cloud_config_from_environment_arm(self):
+        subscription_id = "AZURE_SUBSCRIPTION_ID"
+        resource_group = "RESOURCE_GROUP_NAME"
+        credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_PUBLIC_CLOUD)
+
+        # This environment variable should be set when running python.  If it is set, and the configured cloud is not
+        # one of the default ones in _azure_environments.py, it will follow the url to find the cloud config.  If it
+        # does not find it anywhere, it will throw an "Unknown cloud environment" error.
+        os.environ[
+            ArmConstants.METADATA_URL_ENV_NAME
+        ] = "https://management.azure.com/metadata/endpoints?api-version=2019-05-01"
+        kwargs = {"cloud": "AzureCloud"}
+
+        ml_client = MLClient(
+            credential=credential,
+            subscription_id=subscription_id,
+            resource_group_name=resource_group,
+            workspace_name="test-ws1",
+            **kwargs,
+        )
+        # The client will use the cloud that we passed in
+        print("Client is using cloud:", ml_client._cloud)
+        # In specifying this cloud, we've also set the environment variable to match it
+        print("Cloud name environment variable:", os.environ[AZUREML_CLOUD_ENV_NAME])
+
+    def ml_cloud_config_from_keyword_args(self):
+        subscription_id = "AZURE_SUBSCRIPTION_ID"
+        resource_group = "RESOURCE_GROUP_NAME"
+        credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_PUBLIC_CLOUD)
+
+        # All of these configurations are needed in the kwargs, otherwise the SDK will not recognize the configuration
+        # and will throw an "Unknown cloud environment" error.
+        kwargs = {
+            "cloud": "TestCloud",
+            "azure_portal": "https://test.portal.azure.com/",
+            "resource_manager": "https://test.management.azure.com/",
+            "active_directory": "https://test.login.microsoftonline.com/",
+            "aml_resource_id": "https://test.ml.azure.com/",
+            "storage_endpoint": "test.core.windows.net",
+            "registry_discovery_endpoint": "https://test.eastus.api.azureml.ms/",
+        }
+        ml_client = MLClient(
+            credential=credential,
+            subscription_id=subscription_id,
+            resource_group_name=resource_group,
+            workspace_name="test-ws1",
+            **kwargs,
+        )
+        # The client will use the cloud that we passed in
+        print("Client is using cloud:", ml_client._cloud)
+        # In configuring this cloud, we've also set the environment variable to match it
+        print("Cloud name environment variable:", os.environ[AZUREML_CLOUD_ENV_NAME])
+
+
+if __name__ == "__main__":
+    sample = CloudConfigurationOptions()
+    sample.ml_cloud_config_from_environment_arm()
+    sample.ml_cloud_config_from_keyword_args()