Adding in some functionality to use different clouds

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

@@ -10,6 +10,7 @@
 
 from azure.ai.ml._utils.utils import _get_mfe_url_override
 from azure.ai.ml.constants._common import AZUREML_CLOUD_ENV_NAME
+from azure.ai.ml.constants._common import ArmConstants
 
 module_logger = logging.getLogger(__name__)
 
@@ -56,6 +57,17 @@ class EndpointURLS:  # pylint: disable=too-few-public-methods,no-init
     },
 }
 
+def _get_cloud(cloud: str):
+    if cloud in _environments:
+        return _environments[cloud]
+    arm_url = os.environ.get(ArmConstants.METADATA_URL_ENV_NAME,ArmConstants.DEFAULT_URL)
+    arm_clouds = _get_clouds_by_metadata_url(arm_url)
+    try:
+        new_cloud = arm_clouds[cloud]
+        _environments.update(new_cloud)
+        return new_cloud
+    except KeyError:
+        raise Exception('Unknown cloud environment "{0}".'.format(cloud))
 
 def _get_default_cloud_name():
     """Return AzureCloud as the default cloud."""
@@ -74,17 +86,14 @@ def _get_cloud_details(cloud: str = AzureEnvironments.ENV_DEFAULT):
             AzureEnvironments.ENV_DEFAULT,
         )
         cloud = _get_default_cloud_name()
-    try:
-        azure_environment = _environments[cloud]
-        module_logger.debug("Using the cloud configuration: '%s'.", azure_environment)
-    except KeyError:
-        raise Exception('Unknown cloud environment "{0}".'.format(cloud))
-    return azure_environment
+    return _get_cloud(cloud)
 
 
 def _set_cloud(cloud: str = AzureEnvironments.ENV_DEFAULT):
     if cloud is not None:
-        if cloud not in _environments:
+        try:
+            _get_cloud(cloud)
+        except Exception:
             raise Exception('Unknown cloud environment supplied: "{0}".'.format(cloud))
     else:
         cloud = _get_default_cloud_name()
@@ -189,3 +198,46 @@ def _resource_to_scopes(resource):
     """
     scope = resource + "/.default"
     return [scope]
+
+def _get_clouds_by_metadata_url(metadata_url, timeout=ArmConstants.DEFAULT_TIMEOUT):
+    """Get all the clouds by the specified metadata url
+
+        :return: list of the clouds
+    """
+    import requests
+
+    try:
+        module_logger.debug('Start : Loading cloud metadata from the url specified by %s', metadata_url)
+        with requests.get(metadata_url, timeout=timeout) 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)
+            return cli_cloud_dict
+    except Exception as ex:  # pylint: disable=broad-except
+        module_logger.warning("Error: Azure ML was unable to load cloud metadata from the url specified by %s. %s. "
+                        "This may be due to a misconfiguration of networking controls. Azure Machine Learning Python "
+                        "SDK requires outbound access to Azure Resource Manager. Please contact your networking team "
+                        "to configure outbound access to Azure Resource Manager on both Network Security Group and "
+                        "Firewall. For more details on required configurations, see "
+                        "https://docs.microsoft.com/azure/machine-learning/how-to-access-azureml-behind-firewall.",
+                        metadata_url, ex)
+        return {}
+
+def _convert_arm_to_cli(arm_cloud_metadata):
+    cli_cloud_metadata_dict = {}
+    if isinstance(arm_cloud_metadata, dict):
+        arm_cloud_metadata = [arm_cloud_metadata]
+    for cloud in arm_cloud_metadata:
+        try:
+            cli_cloud_metadata_dict[cloud['name']] = {
+                EndpointURLS.AZURE_PORTAL_ENDPOINT: cloud["portal"],
+                EndpointURLS.RESOURCE_MANAGER_ENDPOINT: cloud["resourceManager"],
+                EndpointURLS.ACTIVE_DIRECTORY_ENDPOINT: cloud["authentication"]["loginEndpoint"],
+                EndpointURLS.AML_RESOURCE_ID: cloud["resourceManager"],
+                EndpointURLS.STORAGE_ENDPOINT: cloud["suffixes"]["storage"]
+
+            }
+        except KeyError as ex:
+            module_logger.warning("Property on cloud ot found in arm cloud metadata: {}".format(ex))
+            continue
+    return cli_cloud_metadata_dict

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -273,6 +273,10 @@ class ArmConstants(object):
     AZURE_MGMT_KEYVAULT_API_VERSION = "2019-09-01"
     AZURE_MGMT_CONTAINER_REG_API_VERSION = "2019-05-01"
 
+    DEFAULT_TIMEOUT = 30
+    DEFAULT_URL = "https://management.azure.com/metadata/endpoints?api-version=2019-05-01"
+    METADATA_URL_ENV_NAME = "ARM_CLOUD_METADATA_URL"
+
 
 class HttpResponseStatusCode(object):
     NOT_FOUND = 404

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

@@ -1,7 +1,8 @@
 import os
-
 import mock
 import pytest
+from mock import MagicMock, patch
+import requests
 
 from azure.ai.ml._azure_environments import (
     AzureEnvironments,
@@ -13,12 +14,16 @@
     _get_storage_endpoint_from_metadata,
     _set_cloud,
 )
-from azure.ai.ml.constants._common import AZUREML_CLOUD_ENV_NAME
+from azure.ai.ml.constants._common import  ArmConstants, AZUREML_CLOUD_ENV_NAME
+
+
+
 
 
 @pytest.mark.unittest
 @pytest.mark.core_sdk_test
 class TestCloudEnvironments:
+
     @mock.patch.dict(os.environ, {AZUREML_CLOUD_ENV_NAME: AzureEnvironments.ENV_DEFAULT}, clear=True)
     def test_set_valid_cloud_details_china(self):
         cloud_environment = AzureEnvironments.ENV_CHINA
@@ -88,4 +93,54 @@ def test_get_registry_endpoint_from_us_gov(self):
         cloud_environment = AzureEnvironments.ENV_US_GOVERNMENT
         _set_cloud(cloud_environment)
         base_url = _get_registry_discovery_endpoint_from_metadata(cloud_environment)
-        assert "https://usgovarizona.api.ml.azure.us/" in base_url
+        assert "https://usgovarizona.api.ml.azure.us/" in base_url
+
+
+
+
+    def mock_arm_response(*args, **kwargs):
+        class MockResponse:
+            def __init__(self, json_data):
+                self.json_data = json_data
+                self.status_code = 201
+            def __enter__(self):
+                return self;
+
+            def __exit__(self, exc_type, exc_value, traceback):
+                return
+
+            def json(self):
+                return self.json_data
+
+        json_data = [
+            {
+                "name": "TEST_ENV", 
+                "portal": "testportal", 
+                "resourceManager": "testresourcemanager",
+                "authentication": {
+                    "loginEndpoint": "testdirectoryendpoint"
+                },
+                "resourceManager": "testresourcemanager",
+                "suffixes": {
+                    "storage": "teststorageendpoint"
+                }
+            },
+            {
+                "name": "MISCONFIGURED"
+            }
+        ]
+        response = MockResponse(json_data)
+        return response
+
+    @mock.patch.dict(os.environ, {}, clear=True)
+    @mock.patch('requests.get', side_effect=mock_arm_response)
+    def test_get_cloud_from_arm(self, mock_get):
+        _set_cloud('TEST_ENV')
+        cloud_details = _get_cloud_information_from_metadata("TEST_ENV")
+        assert cloud_details.get("cloud") == "TEST_ENV"
+
+    @mock.patch.dict(os.environ, {}, clear=True)
+    @mock.patch('requests.get', side_effect=mock_arm_response)
+    def test_arm_misconfigured(self, mock_get):
+        with pytest.raises(Exception) as e_info:
+            _set_cloud("MISCONFIGURED")