Add support for workspace v2 endpoints

azure-quantum/README.md

@@ -29,15 +29,14 @@ To get started, visit the following Quickstart guides:
 
 ## General usage ##
 
-To connect to your Azure Quantum Workspace, go to the [Azure Portal](https://portal.azure.com), navigate to your Workspace and copy-paste the resource ID and location into the code snippet below.
+To connect to your Azure Quantum Workspace, go to the [Azure Portal](https://portal.azure.com), navigate to your Workspace and copy-paste the resource ID into the code snippet below.
 
 ```python
 from azure.quantum import Workspace
 
-# Enter your Workspace details (resource ID and location) below
+# Enter your Workspace resource ID below
 workspace = Workspace(
-    resource_id="",
-    location=""
+    resource_id=""
 )
 ```
 

azure-quantum/azure/quantum/_constants.py

@@ -55,6 +55,9 @@ class ConnectionConstants:
     DATA_PLANE_CREDENTIAL_SCOPE = "https://quantum.microsoft.com/.default"
     ARM_CREDENTIAL_SCOPE = "https://management.azure.com/.default"
 
+    DEFAULT_ARG_API_VERSION = "2021-03-01"
+    DEFAULT_WORKSPACE_API_VERSION = "2025-11-01-preview"
+
     MSA_TENANT_ID = "9188040d-6c67-4c5b-b112-36a304b66dad"
 
     AUTHORITY = AzureIdentityInternals.get_default_authority()
@@ -63,10 +66,14 @@ class ConnectionConstants:
     # pylint: disable=unnecessary-lambda-assignment
     GET_QUANTUM_PRODUCTION_ENDPOINT = \
         lambda location: f"https://{location}.quantum.azure.com/"
+    GET_QUANTUM_PRODUCTION_ENDPOINT_v2 = \
+        lambda location: f"https://{location}-v2.quantum.azure.com/"
     GET_QUANTUM_CANARY_ENDPOINT = \
         lambda location: f"https://{location or 'eastus2euap'}.quantum.azure.com/"
     GET_QUANTUM_DOGFOOD_ENDPOINT = \
         lambda location: f"https://{location}.quantum-test.azure.com/"
+    GET_QUANTUM_DOGFOOD_ENDPOINT_v2 = \
+        lambda location: f"https://{location}-v2.quantum-test.azure.com/"
 
     ARM_PRODUCTION_ENDPOINT = "https://management.azure.com/"
     ARM_DOGFOOD_ENDPOINT = "https://api-dogfood.resources.windows-int.net/"
@@ -93,3 +100,65 @@ class ConnectionConstants:
 GUID_REGEX_PATTERN = (
     r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}"
 )
+
+VALID_WORKSPACE_NAME_PATTERN = r"^[a-zA-Z0-9]+(-*[a-zA-Z0-9])*$"
+
+VALID_AZURE_REGIONS = {
+    "australiacentral",
+    "australiacentral2",
+    "australiaeast",
+    "australiasoutheast",
+    "austriaeast",
+    "belgiumcentral",
+    "brazilsouth",
+    "brazilsoutheast",
+    "canadacentral",
+    "canadaeast",
+    "centralindia",
+    "centralus",
+    "centraluseuap",
+    "chilecentral",
+    "eastasia",
+    "eastus",
+    "eastus2",
+    "eastus2euap",
+    "francecentral",
+    "francesouth",
+    "germanynorth",
+    "germanywestcentral",
+    "indonesiacentral",
+    "israelcentral",
+    "italynorth",
+    "japaneast",
+    "japanwest",
+    "koreacentral",
+    "koreasouth",
+    "malaysiawest",
+    "mexicocentral",
+    "newzealandnorth",
+    "northcentralus",
+    "northeurope",
+    "norwayeast",
+    "norwaywest",
+    "polandcentral",
+    "qatarcentral",
+    "southafricanorth",
+    "southafricawest",
+    "southcentralus",
+    "southindia",
+    "southeastasia",
+    "spaincentral",
+    "swedencentral",
+    "switzerlandnorth",
+    "switzerlandwest",
+    "uaecentral",
+    "uaenorth",
+    "uksouth",
+    "ukwest",
+    "westcentralus",
+    "westeurope",
+    "westindia",
+    "westus",
+    "westus2",
+    "westus3",
+}

azure-quantum/azure/quantum/_mgmt_client.py

@@ -0,0 +1,239 @@
+##
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+##
+"""
+Module providing the WorkspaceMgmtClient class for managing workspace operations.
+Created to do not add additional azure-mgmt-* dependencies that can conflict with existing ones.
+"""
+
+import logging
+from http import HTTPStatus
+from typing import Any, Optional, cast
+from azure.core import PipelineClient
+from azure.core.credentials import TokenProvider
+from azure.core.pipeline import policies
+from azure.core.rest import HttpRequest
+from azure.core.exceptions import HttpResponseError
+from azure.quantum._workspace_connection_params import WorkspaceConnectionParams
+from azure.quantum._constants import ConnectionConstants
+from azure.quantum._client._configuration import VERSION
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["WorkspaceMgmtClient"]
+
+
+class WorkspaceMgmtClient():
+    """
+    Client for Azure Quantum Workspace related ARM/ARG operations.
+    Uses PipelineClient under the hood which is standard for all Azure SDK clients,
+    see https://learn.microsoft.com/en-us/azure/developer/python/sdk/fundamentals/http-pipeline-retries.
+
+    :param credential:
+        The credential to use to connect to Azure services.
+
+    :param base_url:
+        The base URL for the ARM endpoint.
+
+    :param user_agent:
+        Add the specified value as a prefix to the HTTP User-Agent header.
+    """
+
+    # Constants
+    DEFAULT_RETRY_TOTAL = 3
+    CONTENT_TYPE_JSON = "application/json"
+    CONNECT_DOC_LINK = "https://learn.microsoft.com/en-us/azure/quantum/how-to-connect-workspace"
+    CONNECT_DOC_MESSAGE = f"To find details on how to connect to your workspace, please see {CONNECT_DOC_LINK}."
+
+    def __init__(self, credential: TokenProvider, base_url: str, user_agent: Optional[str] = None) -> None:
+        """
+        Initialize the WorkspaceMgmtClient.
+
+        :param credential:
+            The credential to use to connect to Azure services.
+
+        :param base_url:
+            The base URL for the ARM endpoint.
+        """
+        self._credential = credential
+        self._base_url = base_url
+        self._policies = [
+            policies.RequestIdPolicy(),
+            policies.HeadersPolicy({
+                "Content-Type": self.CONTENT_TYPE_JSON,
+                "Accept": self.CONTENT_TYPE_JSON,
+            }),
+            policies.UserAgentPolicy(user_agent=user_agent, sdk_moniker="quantum/{}".format(VERSION)),
+            policies.RetryPolicy(retry_total=self.DEFAULT_RETRY_TOTAL),
+            policies.BearerTokenCredentialPolicy(self._credential, ConnectionConstants.ARM_CREDENTIAL_SCOPE),
+        ]
+        self._client: PipelineClient = PipelineClient(base_url=cast(str, base_url), policies=self._policies)
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> 'WorkspaceMgmtClient':
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *exc_details: Any) -> None:
+        self._client.__exit__(*exc_details)
+
+    def load_workspace_from_arg(self, connection_params: WorkspaceConnectionParams) -> None:
+        """
+        Queries Azure Resource Graph to find a workspace by name and optionally location, resource group, subscription.
+        Provided workspace name, location, resource group, and subscription in connection params must be validated beforehand.
+
+        :param connection_params:
+            The workspace connection parameters to use and update.
+        """
+        if not connection_params.workspace_name:
+            raise ValueError("Workspace name must be specified to try to load workspace details from ARG.")
+
+        query = f"""
+            Resources
+            | where type =~ 'microsoft.quantum/workspaces'
+            | where name =~ '{connection_params.workspace_name}'
+        """
+
+        if connection_params.resource_group:
+            query += f"\n                | where resourceGroup =~ '{connection_params.resource_group}'"
+
+        if connection_params.location:
+            query += f"\n                | where location =~ '{connection_params.location}'"
+
+        query += """
+            | extend endpointUri = tostring(properties.endpointUri)
+            | project name, subscriptionId, resourceGroup, location, endpointUri
+        """
+
+        request_body = {
+            "query": query
+        }
+
+        if connection_params.subscription_id:
+            request_body["subscriptions"] = [connection_params.subscription_id]
+
+        # Create request to Azure Resource Graph API
+        request = HttpRequest(
+            method="POST",
+            url=self._client.format_url("/providers/Microsoft.ResourceGraph/resources"),
+            params={"api-version": ConnectionConstants.DEFAULT_ARG_API_VERSION},
+            json=request_body
+        )
+
+        try:
+            response = self._client.send_request(request)
+            response.raise_for_status()
+            result = response.json()
+        except Exception as e:
+            raise RuntimeError(
+                f"Could not load workspace details from Azure Resource Graph: {str(e)}.\n{self.CONNECT_DOC_MESSAGE}"
+            ) from e
+
+        data = result.get('data', [])
+
+        if not data:
+            raise ValueError(f"No matching workspace found with name '{connection_params.workspace_name}'. {self.CONNECT_DOC_MESSAGE}")
+
+        if len(data) > 1:
+            raise ValueError(
+                f"Multiple Azure Quantum workspaces found with name '{connection_params.workspace_name}'. "
+                f"Please specify additional connection parameters. {self.CONNECT_DOC_MESSAGE}"
+            )
+
+        workspace_data = data[0]
+
+        connection_params.subscription_id = workspace_data.get('subscriptionId')
+        connection_params.resource_group = workspace_data.get('resourceGroup')
+        connection_params.location = workspace_data.get('location')
+        connection_params.quantum_endpoint = workspace_data.get('endpointUri')
+
+        logger.debug(
+            "Found workspace '%s' in subscription '%s', resource group '%s', location '%s', endpoint '%s'",
+            connection_params.workspace_name,
+            connection_params.subscription_id,
+            connection_params.resource_group,
+            connection_params.location,
+            connection_params.quantum_endpoint
+        )
+
+        # If one of the required parameters is missing, probably workspace in failed provisioning state
+        if not connection_params.is_complete():
+            raise ValueError(
+                f"Failed to retrieve complete workspace details for workspace '{connection_params.workspace_name}'. "
+                "Please check that workspace is in valid state."
+            )
+
+    def load_workspace_from_arm(self, connection_params: WorkspaceConnectionParams) -> None:
+        """
+        Fetches the workspace resource from ARM and sets location and endpoint URI params.
+        Provided workspace name, resource group, and subscription in connection params must be validated beforehand.
+
+        :param connection_params:
+            The workspace connection parameters to use and update.
+        """
+        if not all([connection_params.subscription_id, connection_params.resource_group, connection_params.workspace_name]):
+            raise ValueError("Missing required connection parameters to load workspace details from ARM.")
+
+        api_version = connection_params.api_version or ConnectionConstants.DEFAULT_WORKSPACE_API_VERSION
+
+        url = (
+            f"/subscriptions/{connection_params.subscription_id}"
+            f"/resourceGroups/{connection_params.resource_group}"
+            f"/providers/Microsoft.Quantum/workspaces/{connection_params.workspace_name}"
+        )
+
+        request = HttpRequest(
+            method="GET",
+            url=self._client.format_url(url),
+            params={"api-version": api_version},
+        )
+
+        try:
+            response = self._client.send_request(request)
+            response.raise_for_status()
+            workspace_data = response.json()
+        except HttpResponseError as e:
+            if e.status_code == HTTPStatus.NOT_FOUND:
+                raise ValueError(
+                    f"Azure Quantum workspace '{connection_params.workspace_name}' "
+                    f"not found in resource group '{connection_params.resource_group}' "
+                    f"and subscription '{connection_params.subscription_id}'. "
+                    f"{self.CONNECT_DOC_MESSAGE}"
+                ) from e
+            # Re-raise for other HTTP errors
+            raise
+        except Exception as e:
+            raise RuntimeError(
+                f"Could not load workspace details from ARM: {str(e)}.\n{self.CONNECT_DOC_MESSAGE}"
+            ) from e
+
+        # Extract and apply location
+        location = workspace_data.get("location")
+        if location:
+            connection_params.location = location
+            logger.debug(
+                "Updated workspace location from ARM: %s",
+                location
+            )
+        else:
+            raise ValueError(
+                f"Failed to retrieve location for workspace '{connection_params.workspace_name}'. "
+                f"Please check that workspace is in valid state."
+            )
+
+        # Extract and apply endpoint URI from properties
+        properties = workspace_data.get("properties", {})
+        endpoint_uri = properties.get("endpointUri")
+        if endpoint_uri:
+            connection_params.quantum_endpoint = endpoint_uri
+            logger.debug(
+                "Updated workspace endpoint from ARM: %s", connection_params.quantum_endpoint
+            )
+        else:
+            raise ValueError(
+                f"Failed to retrieve endpoint uri for workspace '{connection_params.workspace_name}'. "
+                f"Please check that workspace is in valid state."
+            )