[ML] feat: Add support for Index asset (#35251)

* refactor: Remove duplicatation in workspace api discovery

* chore: Generate restclient from AzureAI.Assets typespec

    Generated using eng/scripts/TypeSpec-Generate-Sdk.ps1 script
    from Azure/azure-rest-api-specs. Patched it to append the
    '@azure-tools/typespec-python.generate-packaging-files=false' option
    since we aren't generating a standalone SDK.

    Patched the tspconfig.yaml for AzureAI.Assets to:
    * Force the SDK to generate at sdk/ml/azure-ai-ml/azure/ai/ml/_restclient
    * Call it a different name from azurekai-resources-autogen

* fix: Fix the base genericasset's api endpoint

    The endpoint in the typespec was out of date.

* fix: Don't url encode the {endpoint} parameter.

    The {endpoint} parameter in the typespec is used to specify the
    base url for the generic asset api (depends on the workspace
    location).

    The python autogenerated sdk assumes that the parameter is part
    of the URL path, and tries to url encode the {endpoint} parameter.

    This is problematic since this encodes the forward slashes (//)
    joining the scheme and network location to `%2F%2F`.

* feat: Add Index Asset

* feat: Expose Index Asset publicly

* feat: Add IndexOperations

* feat: Add ml_client.indexes

* feat: Add latest label support to IndexOperations.get

* chore: Add __init__.py to ai assets rest client

* fix: Resolve mypy failures

* fix: Address FIXME about non-compliant service response

    The service response was updated to align with what's defined in
    typespec.

* feat: Support local upload of index files

* refactor: Replace `storage_uri` with `path`

    The rest api's concept is analogous with the `path` attribute
    for Assets in the SDK.

    Using `path` for the SDK Index asset keeps the semantics consistent.

* docs: Update types in Index docstring

* docs: Update CHANGELOG.md

* feat: Auto-increment index version

* feat: Mark index asset as experimental

* feat: Enable dynamic dispatch create_or_update

* docs: Update IndexOperations docstring

* docs: Remove datastore from docstring for consistency

* docs: Add a clarifying comment to Index._to_rest_object

* fix: Prevent Project._from_rest_object from dropping discovery_url

   The initial implementation of Project._from_rest_object seems to
   drop many of the fields supported by its parent class. One of which
   is `discovery_url` which this PR depends on to find the correct
   endpoint to call for the GenericAsset API.

   The implementation of Project._from_rest_object appears to do
   a subset of the assignments of Workspace._from_rest_object.

   This commit removes the override for Project._from_rest_object
   allowing the Project class to fall back on Workspace._from_rest_object

* test: Add unit tests for Index Asset

Author: kdestin

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

@@ -9,6 +9,7 @@
 - Project and Hub operations supported by workspace operations.
 - workspace list operation supports type filtering.
 - Add support for Microsoft Entra token (`aad_token`) auth in `invoke` and `get-credentials` operations.
+- Add experimental support for working with indexes: `ml_client.indexes`
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/_artifacts/_artifact_utilities.py

@@ -50,7 +50,13 @@
 from azure.storage.filedatalake import FileSasPermissions, generate_file_sas
 
 if TYPE_CHECKING:
-    from azure.ai.ml.operations import DataOperations, EnvironmentOperations, FeatureSetOperations, ModelOperations
+    from azure.ai.ml.operations import (
+        DataOperations,
+        EnvironmentOperations,
+        FeatureSetOperations,
+        IndexOperations,
+        ModelOperations,
+    )
     from azure.ai.ml.operations._code_operations import CodeOperations
 
 module_logger = logging.getLogger(__name__)
@@ -454,7 +460,9 @@ def _update_gen2_metadata(name, version, indicator_file, storage_client) -> None
 
 def _check_and_upload_path(
     artifact: T,
-    asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations", "FeatureSetOperations"],
+    asset_operations: Union[
+        "DataOperations", "ModelOperations", "CodeOperations", "FeatureSetOperations", "IndexOperations"
+    ],
     artifact_type: str,
     datastore_name: Optional[str] = None,
     sas_uri: Optional[str] = None,
@@ -466,7 +474,7 @@ def _check_and_upload_path(
     :param artifact: artifact to check and upload param
     :type artifact: T
     :param asset_operations: The asset operations to use for uploading
-    :type asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations"]
+    :type asset_operations: Union["DataOperations", "ModelOperations", "CodeOperations", "IndexOperations"]
     :param artifact_type: The artifact type
     :type artifact_type: str
     :param datastore_name: the name of the datastore to upload to

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

@@ -55,6 +55,7 @@
     Compute,
     Datastore,
     Environment,
+    Index,
     Job,
     Model,
     ModelBatchDeployment,
@@ -75,6 +76,7 @@
     DataOperations,
     DatastoreOperations,
     EnvironmentOperations,
+    IndexOperations,
     JobOperations,
     ModelOperations,
     OnlineDeploymentOperations,
@@ -617,6 +619,18 @@ def __init__(
         )
         self._operation_container.add(AzureMLResourceType.SCHEDULE, self._schedules)
 
+        self._indexes = IndexOperations(
+            operation_scope=self._operation_scope,
+            operation_config=self._operation_config,
+            credential=self._credential,
+            all_operations=self._operation_container,
+            datastore_operations=self._datastores,
+            _service_client_kwargs=kwargs,
+            requests_pipeline=self._requests_pipeline,
+            **ops_kwargs,
+        )
+        self._operation_container.add(AzureMLResourceType.INDEX, self._indexes)
+
         try:
             from azure.ai.ml.operations._virtual_cluster_operations import VirtualClusterOperations
 
@@ -970,6 +984,15 @@ def schedules(self) -> ScheduleOperations:
         """
         return self._schedules
 
+    @property
+    def indexes(self) -> IndexOperations:
+        """A collection of index related operations.
+
+        :return: Index operations.
+        :rtype: ~azure.ai.ml.operations.IndexOperations
+        """
+        return self._indexes
+
     @property
     def subscription_id(self) -> str:
         """Get the subscription ID of an MLClient object.
@@ -1178,6 +1201,12 @@ def _(entity: Datastore, operations):
     return operations[AzureMLResourceType.DATASTORE].create_or_update(entity)
 
 
+@_create_or_update.register(Index)
+def _(entity: Index, operations, *args, **kwargs):
+    module_logger.debug("Creating or updating indexes")
+    return operations[AzureMLResourceType.INDEX].begin_create_or_update(entity, **kwargs)
+
+
 @singledispatch
 def _begin_create_or_update(entity, operations, **kwargs):
     raise TypeError("Please refer to begin_create_or_update docstring for valid input types.")

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/azure_ai_assets_v2024_04_01/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+# Code generated by Microsoft (R) Python Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is regenerated.
+# --------------------------------------------------------------------------
+
+from ._client import MachineLearningServicesClient
+from ._version import VERSION
+
+__version__ = VERSION
+
+try:
+    from ._patch import __all__ as _patch_all
+    from ._patch import *  # pylint: disable=unused-wildcard-import
+except ImportError:
+    _patch_all = []
+from ._patch import patch_sdk as _patch_sdk
+
+__all__ = [
+    "MachineLearningServicesClient",
+]
+__all__.extend([p for p in _patch_all if p not in __all__])
+
+_patch_sdk()

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/azure_ai_assets_v2024_04_01/azureaiassetsv20240401/__init__.py

@@ -0,0 +1,127 @@
+# 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.
+# Code generated by Microsoft (R) Python Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is regenerated.
+# --------------------------------------------------------------------------
+
+from copy import deepcopy
+from typing import Any, TYPE_CHECKING
+
+from azure.core import PipelineClient
+from azure.core.pipeline import policies
+from azure.core.rest import HttpRequest, HttpResponse
+
+from ._configuration import MachineLearningServicesClientConfiguration
+from ._serialization import Deserializer, Serializer
+from .operations import IndexesOperations
+
+if TYPE_CHECKING:
+    # pylint: disable=unused-import,ungrouped-imports
+    from azure.core.credentials import TokenCredential
+
+
+class MachineLearningServicesClient:  # pylint: disable=client-accepts-api-version-keyword
+    """MachineLearningServicesClient.
+
+    :ivar indexes: IndexesOperations operations
+    :vartype indexes: azureaiassetsv20240401.operations.IndexesOperations
+    :param endpoint: Supported Azure-AI asset endpoints. Required.
+    :type endpoint: str
+    :param subscription_id: The ID of the target subscription. Required.
+    :type subscription_id: str
+    :param resource_group_name: The name of the Resource Group. Required.
+    :type resource_group_name: str
+    :param workspace_name: The name of the AzureML workspace or AI project. Required.
+    :type workspace_name: str
+    :param credential: Credential used to authenticate requests to the service. Required.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :keyword api_version: The API version to use for this operation. Default value is
+     "2024-04-01-preview". Note that overriding this default value may result in unsupported
+     behavior.
+    :paramtype api_version: str
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        subscription_id: str,
+        resource_group_name: str,
+        workspace_name: str,
+        credential: "TokenCredential",
+        **kwargs: Any
+    ) -> None:
+        _endpoint = "{endpoint}/genericasset/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.MachineLearningServices/workspaces/{workspaceName}"  # pylint: disable=line-too-long
+        self._config = MachineLearningServicesClientConfiguration(
+            endpoint=endpoint,
+            subscription_id=subscription_id,
+            resource_group_name=resource_group_name,
+            workspace_name=workspace_name,
+            credential=credential,
+            **kwargs
+        )
+        _policies = kwargs.pop("policies", None)
+        if _policies is None:
+            _policies = [
+                policies.RequestIdPolicy(**kwargs),
+                self._config.headers_policy,
+                self._config.user_agent_policy,
+                self._config.proxy_policy,
+                policies.ContentDecodePolicy(**kwargs),
+                self._config.redirect_policy,
+                self._config.retry_policy,
+                self._config.authentication_policy,
+                self._config.custom_hook_policy,
+                self._config.logging_policy,
+                policies.DistributedTracingPolicy(**kwargs),
+                policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
+                self._config.http_logging_policy,
+            ]
+        self._client: PipelineClient = PipelineClient(base_url=_endpoint, policies=_policies, **kwargs)
+
+        self._serialize = Serializer()
+        self._deserialize = Deserializer()
+        self._serialize.client_side_validation = False
+        self.indexes = IndexesOperations(self._client, self._config, self._serialize, self._deserialize)
+
+    def send_request(self, request: HttpRequest, *, stream: bool = False, **kwargs: Any) -> HttpResponse:
+        """Runs the network request through the client's chained policies.
+
+        >>> from azure.core.rest import HttpRequest
+        >>> request = HttpRequest("GET", "https://www.example.org/")
+        <HttpRequest [GET], url: 'https://www.example.org/'>
+        >>> response = client.send_request(request)
+        <HttpResponse: 200 OK>
+
+        For more information on this code flow, see https://aka.ms/azsdk/dpcodegen/python/send_request
+
+        :param request: The network request you want to make. Required.
+        :type request: ~azure.core.rest.HttpRequest
+        :keyword bool stream: Whether the response payload will be streamed. Defaults to False.
+        :return: The response of your network call. Does not do error handling on your response.
+        :rtype: ~azure.core.rest.HttpResponse
+        """
+
+        request_copy = deepcopy(request)
+        path_format_arguments = {
+            "endpoint": self._config.endpoint,
+            "subscriptionId": self._serialize.url("self._config.subscription_id", self._config.subscription_id, "str"),
+            "resourceGroupName": self._serialize.url(
+                "self._config.resource_group_name", self._config.resource_group_name, "str"
+            ),
+            "workspaceName": self._serialize.url("self._config.workspace_name", self._config.workspace_name, "str"),
+        }
+
+        request_copy.url = self._client.format_url(request_copy.url, **path_format_arguments)
+        return self._client.send_request(request_copy, stream=stream, **kwargs)  # type: ignore
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> "MachineLearningServicesClient":
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *exc_details: Any) -> None:
+        self._client.__exit__(*exc_details)

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/azure_ai_assets_v2024_04_01/azureaiassetsv20240401/_client.py

@@ -0,0 +1,88 @@
+# 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.
+# Code generated by Microsoft (R) Python Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is regenerated.
+# --------------------------------------------------------------------------
+
+from typing import Any, TYPE_CHECKING
+
+from azure.core.pipeline import policies
+
+from ._version import VERSION
+
+if TYPE_CHECKING:
+    # pylint: disable=unused-import,ungrouped-imports
+    from azure.core.credentials import TokenCredential
+
+
+class MachineLearningServicesClientConfiguration:  # pylint: disable=too-many-instance-attributes,name-too-long
+    """Configuration for MachineLearningServicesClient.
+
+    Note that all parameters used to create this instance are saved as instance
+    attributes.
+
+    :param endpoint: Supported Azure-AI asset endpoints. Required.
+    :type endpoint: str
+    :param subscription_id: The ID of the target subscription. Required.
+    :type subscription_id: str
+    :param resource_group_name: The name of the Resource Group. Required.
+    :type resource_group_name: str
+    :param workspace_name: The name of the AzureML workspace or AI project. Required.
+    :type workspace_name: str
+    :param credential: Credential used to authenticate requests to the service. Required.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :keyword api_version: The API version to use for this operation. Default value is
+     "2024-04-01-preview". Note that overriding this default value may result in unsupported
+     behavior.
+    :paramtype api_version: str
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        subscription_id: str,
+        resource_group_name: str,
+        workspace_name: str,
+        credential: "TokenCredential",
+        **kwargs: Any
+    ) -> None:
+        api_version: str = kwargs.pop("api_version", "2024-04-01-preview")
+
+        if endpoint is None:
+            raise ValueError("Parameter 'endpoint' must not be None.")
+        if subscription_id is None:
+            raise ValueError("Parameter 'subscription_id' must not be None.")
+        if resource_group_name is None:
+            raise ValueError("Parameter 'resource_group_name' must not be None.")
+        if workspace_name is None:
+            raise ValueError("Parameter 'workspace_name' must not be None.")
+        if credential is None:
+            raise ValueError("Parameter 'credential' must not be None.")
+
+        self.endpoint = endpoint
+        self.subscription_id = subscription_id
+        self.resource_group_name = resource_group_name
+        self.workspace_name = workspace_name
+        self.credential = credential
+        self.api_version = api_version
+        self.credential_scopes = kwargs.pop("credential_scopes", ["https://ml.azure.com/.default"])
+        kwargs.setdefault("sdk_moniker", "azure_ai_assets_v2024_04_01/{}".format(VERSION))
+        self.polling_interval = kwargs.get("polling_interval", 30)
+        self._configure(**kwargs)
+
+    def _configure(self, **kwargs: Any) -> None:
+        self.user_agent_policy = kwargs.get("user_agent_policy") or policies.UserAgentPolicy(**kwargs)
+        self.headers_policy = kwargs.get("headers_policy") or policies.HeadersPolicy(**kwargs)
+        self.proxy_policy = kwargs.get("proxy_policy") or policies.ProxyPolicy(**kwargs)
+        self.logging_policy = kwargs.get("logging_policy") or policies.NetworkTraceLoggingPolicy(**kwargs)
+        self.http_logging_policy = kwargs.get("http_logging_policy") or policies.HttpLoggingPolicy(**kwargs)
+        self.custom_hook_policy = kwargs.get("custom_hook_policy") or policies.CustomHookPolicy(**kwargs)
+        self.redirect_policy = kwargs.get("redirect_policy") or policies.RedirectPolicy(**kwargs)
+        self.retry_policy = kwargs.get("retry_policy") or policies.RetryPolicy(**kwargs)
+        self.authentication_policy = kwargs.get("authentication_policy")
+        if self.credential and not self.authentication_policy:
+            self.authentication_policy = policies.BearerTokenCredentialPolicy(
+                self.credential, *self.credential_scopes, **kwargs
+            )