Add VC get and list operations (Azure#28624)

* VC operations

Author: luigiw

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

@@ -40,7 +40,7 @@
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils._http_utils import HttpPipeline
 from azure.ai.ml._utils._registry_utils import RegistryDiscovery
-from azure.ai.ml._utils.utils import _is_https_url, _validate_missing_sub_or_rg_and_raise
+from azure.ai.ml._utils.utils import _is_https_url
 from azure.ai.ml.constants._common import AzureMLResourceType
 from azure.ai.ml.entities import (
     BatchDeployment,
@@ -74,6 +74,7 @@
     RegistryOperations,
     WorkspaceConnectionsOperations,
     WorkspaceOperations,
+    VirtualClusterOperations,
 )
 from azure.ai.ml.operations._code_operations import CodeOperations
 from azure.ai.ml.operations._local_deployment_helper import _LocalDeploymentHelper
@@ -158,8 +159,7 @@ def __init__(
                 target=ErrorTarget.GENERAL,
                 error_category=ErrorCategory.USER_ERROR,
             )
-        if not registry_name:
-            _validate_missing_sub_or_rg_and_raise(subscription_id, resource_group_name)
+
         self._credential = credential
 
         show_progress = kwargs.pop("show_progress", True)
@@ -429,6 +429,8 @@ def __init__(
         )
         self._operation_container.add(AzureMLResourceType.SCHEDULE, self._schedules)
 
+        self._virtual_clusters = VirtualClusterOperations(self._operation_scope, self._credential, **ops_kwargs)
+
     @classmethod
     def from_config(
         cls,

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

@@ -0,0 +1,67 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+from typing import List, Dict, Optional
+import azure.mgmt.resourcegraph as arg
+from azure.mgmt.resource import SubscriptionClient, ResourceManagementClient
+from azure.core.credentials import TokenCredential
+
+
+def get_resources_from_subscriptions(strQuery: str, credential: TokenCredential,
+    subscription_list: Optional[List[str]] = None):
+
+    # If a subscription list is passed in, use it. Otherwise, get all subscriptions
+    subsList = []
+    if subscription_list is not None:
+        subsList = subscription_list
+    else:
+        subsClient = SubscriptionClient(credential)
+        for sub in subsClient.subscriptions.list():
+            subsList.append(sub.as_dict().get('subscription_id'))
+
+    # Create Azure Resource Graph client and set options
+    argClient = arg.ResourceGraphClient(credential)
+    argQueryOptions = arg.models.QueryRequestOptions(result_format="objectArray")
+
+    # Create query
+    argQuery = arg.models.QueryRequest(subscriptions=subsList, query=strQuery, options=argQueryOptions)
+
+    # Allowing API version to be set is yet to be released by azure-mgmt-resourcegraph,
+    # hence the commented out code below. This is the API version Studio UX is using.
+    # return argClient.resources(argQuery, api_version="2021-03-01")
+
+    return argClient.resources(argQuery)
+
+
+def get_virtual_clusters_from_subscriptions(credential: TokenCredential,
+    subscription_list: Optional[List[str]] = None) -> List[Dict]:
+
+    # cspell:ignore tolower
+    strQuery = """resources
+    | where type == 'microsoft.machinelearningservices/virtualclusters'
+    | order by tolower(name) asc 
+    | project id, subscriptionId, resourceGroup, name, location, tags, type"""
+
+    return get_resources_from_subscriptions(strQuery, credential, subscription_list).data
+
+
+def get_generic_resource_by_id(arm_id: str, credential: TokenCredential,
+    subscription_id: str, api_version: Optional[str] = None) -> Dict:
+
+    resource_client = ResourceManagementClient(credential, subscription_id)
+    generic_resource = resource_client.resources.get_by_id(arm_id, api_version)
+
+    return generic_resource.as_dict()
+
+def get_virtual_cluster_by_id(name: str, resource_group: str,
+    subscription_id: str, credential: TokenCredential) -> Dict:
+
+    arm_id = (
+        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}"
+        f"/providers/Microsoft.MachineLearningServices/virtualClusters/{name}"
+    )
+
+    # This is the API version Studio UX is using.
+    return get_generic_resource_by_id(arm_id, credential, subscription_id, api_version="2021-03-01-preview")

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

@@ -908,32 +908,6 @@ def get_all_enum_values_iter(enum_type):
             yield getattr(enum_type, key)
 
 
-def _validate_missing_sub_or_rg_and_raise(subscription_id: Optional[str], resource_group: Optional[str]):
-    """Determine if subscription or resource group is missing and raise exception
-    as appropriate."""
-    from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
-
-    # These imports can't be placed in at top file level because it will cause a circular import in
-    # exceptions.py via _get_mfe_url_override
-
-    msg = "Both subscription id and resource group are required for this operation, missing {}"
-    sub_msg = None
-    if not subscription_id and not resource_group:
-        sub_msg = "subscription id and resource group"
-    elif not subscription_id and resource_group:
-        sub_msg = "subscription id"
-    elif subscription_id and not resource_group:
-        sub_msg = "resource group"
-
-    if sub_msg:
-        raise ValidationException(
-            message=msg.format(sub_msg),
-            no_personal_data_message=msg.format(sub_msg),
-            target=ErrorTarget.GENERAL,
-            error_category=ErrorCategory.USER_ERROR,
-        )
-
-
 def write_to_shared_file(file_path: Union[str, PathLike], content: str):
     """Open file with specific mode and return the file object.
 

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

@@ -23,6 +23,7 @@
 from ._schedule_operations import ScheduleOperations
 from ._workspace_connections_operations import WorkspaceConnectionsOperations
 from ._workspace_operations import WorkspaceOperations
+from ._virtual_cluster_operations import VirtualClusterOperations
 
 __all__ = [
     "ComputeOperations",
@@ -41,4 +42,5 @@
     "WorkspaceConnectionsOperations",
     "RegistryOperations",
     "ScheduleOperations",
+    "VirtualClusterOperations",
 ]

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

@@ -0,0 +1,75 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from typing import Dict, Iterable, Optional
+
+from azure.ai.ml._scope_dependent_operations import OperationScope
+
+# from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
+from azure.ai.ml._utils._logger_utils import OpsLogger
+from azure.ai.ml._utils.azure_resource_utils import get_virtual_clusters_from_subscriptions, get_virtual_cluster_by_id
+from azure.ai.ml.constants._common import Scope
+from azure.ai.ml.exceptions import UserErrorException
+from azure.core.tracing.decorator import distributed_trace
+from azure.core.credentials import TokenCredential
+
+ops_logger = OpsLogger(__name__)
+module_logger = ops_logger.module_logger
+
+
+class VirtualClusterOperations():
+    """VirtualClusterOperations.
+
+    You should not instantiate this class directly. Instead, you should
+    create an MLClient instance that instantiates it for you and
+    attaches it as an attribute.
+    """
+
+    def __init__(
+        self,
+        operation_scope: OperationScope,
+        credentials: TokenCredential,
+        **kwargs: Dict,
+    ):
+        self._resource_group_name = operation_scope.resource_group_name
+        self._subscription_id = operation_scope.subscription_id
+        self._credentials = credentials
+        self._init_kwargs = kwargs
+
+    @distributed_trace
+    # @monitor_with_activity(logger, "VirtualCluster.List", ActivityType.PUBLICAPI)
+    def list(self, *, scope: Optional[str] = None) -> Iterable[Dict]:
+        """List virtual clusters a user has access to.
+
+        :param scope: scope of the listing, "subscription" or None, defaults to None.
+            If None, list virtual clusters across all subscriptions a customer has access to.
+        :type scope: str, optional
+        :return: An iterator like instance of dictionaries.
+        :rtype: ~azure.core.paging.ItemPaged[Dict]
+        """
+
+        if scope is None:
+            subscription_list = None
+        elif scope.lower() == Scope.SUBSCRIPTION:
+            subscription_list = [self._subscription_id]
+        else:
+            message = f"Invalid scope: {scope}. Valid values are 'subscription' or None."
+            raise UserErrorException(message=message, no_personal_data_message=message)
+
+        return get_virtual_clusters_from_subscriptions(self._credentials, subscription_list=subscription_list)
+
+    @distributed_trace
+    # @monitor_with_activity(logger, "VirtualCluster.Get", ActivityType.PUBLICAPI)
+    def get(self, name: str) -> Dict:
+        """Get a virtual cluster resource.
+
+        :param name: Name of the virtual cluster.
+        :type name: str
+        :return: Virtual cluster object
+        :rtype: Dict
+        """
+
+
+        return get_virtual_cluster_by_id(name=name, resource_group=self._resource_group_name,
+            subscription_id=self._subscription_id, credential=self._credentials)

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

@@ -85,6 +85,8 @@
         "azure-common<2.0.0,>=1.1",
         "typing-extensions<5.0.0",
         # "opencensus-ext-azure<2.0.0", disabled until SDK logging re-activated
+        "azure-mgmt-resourcegraph<9.0.0,>=2.0.0",
+        "azure-mgmt-resource<23.0.0,>=3.0.0",
     ],
     extras_require={
         # user can run `pip install azure-ai-ml[designer]` to install mldesigner alone with this package

sdk/ml/azure-ai-ml/tests/conftest.py

@@ -834,6 +834,7 @@ def pytest_configure(config):
         ("data_experiences_test", "marks tests as data experience tests"),
         ("local_endpoint_local_assets", "marks tests as local_endpoint_local_assets"),
         ("local_endpoint_byoc", "marks tests as local_endpoint_byoc"),
+        ("virtual_cluster_test", "marks tests as virtual cluster tests"),
     ]:
         config.addinivalue_line("markers", f"{marker}: {description}")
 

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

@@ -421,35 +421,6 @@ def test_ml_client_with_invalid_cloud(self, mock_credential):
             assert ml_client._kwargs["cloud"] == "SomeInvalidCloudName"
         assert "Unknown cloud environment supplied" in str(e)
 
-
-    def test_ml_client_validation_rg_sub_missing_throws(
-        self, auth: ClientSecretCredential
-    ) -> None:
-        with pytest.raises(ValidationException) as exception:
-            MLClient(
-                credential=auth,
-            )
-        message = exception.value.args[0]
-        assert (
-            message
-            == "Both subscription id and resource group are required for this operation, missing subscription id and resource group"
-        )
-
-
-    def test_ml_client_with_no_rg_sub_for_ws_throws(
-        self, e2e_ws_scope: OperationScope, auth: ClientSecretCredential
-    ) -> None:
-        with pytest.raises(ValidationException) as exception:
-            MLClient(
-                credential=auth,
-                workspace_name=e2e_ws_scope.workspace_name,
-            )
-        message = exception.value.args[0]
-        assert (
-            message
-            == "Both subscription id and resource group are required for this operation, missing subscription id and resource group"
-        )
-
     def test_ml_client_with_both_workspace_registry_names_throws(self, e2e_ws_scope: OperationScope, auth: ClientSecretCredential) -> None:
         with pytest.raises(ValidationException) as exception:
             MLClient(