add ValidationException for incorrect model id, default auth_model always to key, and add docstrings for serverless, aoai, and marketplace (#35471)

needuv · web-flow · commit f8e4e43818be · 2024-05-06T17:26:51.000Z
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_autogen_entities/models/_patch.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_autogen_entities/models/_patch.py
@@ -136,6 +136,7 @@ def _to_rest_object(self) -> RestServerlessEndpoint:
             properties=RestServerlessEndpointProperties(
                 model_settings=RestModelSettings(model_id=self.model_id),
             ),
+            auth_mode="key",  # only key is supported for now
             tags=self.tags,
             sku=RestSku(name="Consumption"),
             location=self.location,
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_azure_openai_deployment_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_azure_openai_deployment_operations.py
@@ -17,6 +17,13 @@
 
 
 class AzureOpenAIDeploymentOperations(_ScopeDependentOperations):
+    """AzureOpenAIDeploymentOperations.
+
+    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,
@@ -29,6 +36,13 @@ def __init__(
         self._connections_operations = connections_operations
 
     def list(self, connection_name: str, **kwargs) -> Iterable[AzureOpenAIDeployment]:
+        """List Azure OpenAI deployments of the workspace.
+
+        :param connection_name: Name of the connection from which to list deployments
+        :type connection_name: str
+        :return: A list of Azure OpenAI deployments
+        :rtype: ~typing.Iterable[~azure.ai.ml.entities.AzureOpenAIDeployment]
+        """
         connection = self._connections_operations.get(connection_name)
 
         def _from_rest_add_connection_name(obj):
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_marketplace_subscription_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_marketplace_subscription_operations.py
@@ -24,6 +24,13 @@
 
 
 class MarketplaceSubscriptionOperations(_ScopeDependentOperations):
+    """MarketplaceSubscriptionOperations.
+
+    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,
@@ -38,6 +45,13 @@ def __init__(
     def begin_create_or_update(
         self, marketplace_subscription: MarketplaceSubscription, **kwargs
     ) -> LROPoller[MarketplaceSubscription]:
+        """Create or update a Marketplace Subscription.
+
+        :param marketplace_subscription: The marketplace subscription entity.
+        :type marketplace_subscription: ~azure.ai.ml.entities.MarketplaceSubscription
+        :return: A poller to track the operation status
+        :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.MarketplaceSubscription]
+        """
         return self._service_client.begin_create_or_update(
             self._resource_group_name,
             self._workspace_name,
@@ -50,6 +64,13 @@ def begin_create_or_update(
     @experimental
     @monitor_with_activity(ops_logger, "MarketplaceSubscription.Get", ActivityType.PUBLICAPI)
     def get(self, name: str, **kwargs) -> MarketplaceSubscription:
+        """Get a Marketplace Subscription resource.
+
+        :param name: Name of the marketplace subscription.
+        :type name: str
+        :return: Marketplace subscription object retrieved from the service.
+        :rtype: ~azure.ai.ml.entities.MarketplaceSubscription
+        """
         return self._service_client.get(
             self._resource_group_name,
             self._workspace_name,
@@ -61,6 +82,11 @@ def get(self, name: str, **kwargs) -> MarketplaceSubscription:
     @experimental
     @monitor_with_activity(ops_logger, "MarketplaceSubscription.List", ActivityType.PUBLICAPI)
     def list(self, **kwargs) -> Iterable[MarketplaceSubscription]:
+        """List marketplace subscriptions of the workspace.
+
+        :return: A list of marketplace subscriptions
+        :rtype: ~typing.Iterable[~azure.ai.ml.entities.MarketplaceSubscription]
+        """
         return self._service_client.list(
             self._resource_group_name,
             self._workspace_name,
@@ -71,6 +97,13 @@ def list(self, **kwargs) -> Iterable[MarketplaceSubscription]:
     @experimental
     @monitor_with_activity(ops_logger, "MarketplaceSubscription.BeginDelete", ActivityType.PUBLICAPI)
     def begin_delete(self, name: str, **kwargs) -> LROPoller[None]:
+        """Delete a Marketplace Subscription.
+
+        :param name: Name of the marketplace subscription.
+        :type name: str
+        :return: A poller to track the operation status.
+        :rtype: ~azure.core.polling.LROPoller[None]
+        """
         return self._service_client.begin_delete(
             self._resource_group_name,
             self._workspace_name,
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_serverless_endpoint_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_serverless_endpoint_operations.py
@@ -5,9 +5,10 @@
 # pylint: disable=protected-access
 
 from typing import Iterable
+import re
 
 from azure.ai.ml.exceptions import ValidationException, ErrorTarget, ErrorCategory, ValidationErrorType
-from azure.ai.ml.constants._common import AzureMLResourceType
+from azure.ai.ml.constants._common import AzureMLResourceType, REGISTRY_VERSION_PATTERN
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient202401Preview
 from azure.ai.ml._restclient.v2024_01_01_preview.models import RegenerateEndpointKeysRequest, KeyType
@@ -30,6 +31,13 @@
 
 
 class ServerlessEndpointOperations(_ScopeDependentOperations):
+    """ServerlessEndpointOperations.
+
+    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,
@@ -50,8 +58,29 @@ def _get_workspace_location(self) -> str:
     @experimental
     @monitor_with_activity(ops_logger, "ServerlessEndpoint.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
     def begin_create_or_update(self, endpoint: ServerlessEndpoint, **kwargs) -> LROPoller[ServerlessEndpoint]:
+        """Create or update a serverless endpoint.
+
+        :param endpoint: The serverless endpoint entity.
+        :type endpoint: ~azure.ai.ml.entities.ServerlessEndpoint
+        :raises ~azure.ai.ml.exceptions.ValidationException: Raised if ServerlessEndpoint cannot be
+            successfully validated. Details will be provided in the error message.
+        :return: A poller to track the operation status
+        :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.ServerlessEndpoint]
+        """
         if not endpoint.location:
             endpoint.location = self._get_workspace_location()
+        if re.match(REGISTRY_VERSION_PATTERN, endpoint.model_id):
+            msg = (
+                "The given model_id {} points to a specific model version, which is not supported. "
+                "Please provide a model_id without the version information."
+            )
+            raise ValidationException(
+                message=msg.format(endpoint.model_id),
+                no_personal_data_message="Invalid model_id given for serverless endpoint",
+                target=ErrorTarget.SERVERLESS_ENDPOINT,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
         return self._service_client.begin_create_or_update(
             self._resource_group_name,
             self._workspace_name,
@@ -64,6 +93,13 @@ def begin_create_or_update(self, endpoint: ServerlessEndpoint, **kwargs) -> LROP
     @experimental
     @monitor_with_activity(ops_logger, "ServerlessEndpoint.Get", ActivityType.PUBLICAPI)
     def get(self, name: str, **kwargs) -> ServerlessEndpoint:
+        """Get a Serverless Endpoint resource.
+
+        :param name: Name of the serverless endpoint.
+        :type name: str
+        :return: Serverless endpoint object retrieved from the service.
+        :rtype: ~azure.ai.ml.entities.ServerlessEndpoint
+        """
         return self._service_client.get(
             self._resource_group_name,
             self._workspace_name,
@@ -75,6 +111,11 @@ def get(self, name: str, **kwargs) -> ServerlessEndpoint:
     @experimental
     @monitor_with_activity(ops_logger, "ServerlessEndpoint.list", ActivityType.PUBLICAPI)
     def list(self, **kwargs) -> Iterable[ServerlessEndpoint]:
+        """List serverless endpoints of the workspace.
+
+        :return: A list of serverless endpoints
+        :rtype: ~typing.Iterable[~azure.ai.ml.entities.ServerlessEndpoint]
+        """
         return self._service_client.list(
             self._resource_group_name,
             self._workspace_name,
@@ -85,6 +126,13 @@ def list(self, **kwargs) -> Iterable[ServerlessEndpoint]:
     @experimental
     @monitor_with_activity(ops_logger, "ServerlessEndpoint.BeginDelete", ActivityType.PUBLICAPI)
     def begin_delete(self, name: str, **kwargs) -> LROPoller[None]:
+        """Delete a Serverless Endpoint.
+
+        :param name: Name of the serverless endpoint.
+        :type name: str
+        :return: A poller to track the operation status.
+        :rtype: ~azure.core.polling.LROPoller[None]
+        """
         return self._service_client.begin_delete(
             self._resource_group_name,
             self._workspace_name,
@@ -95,6 +143,13 @@ def begin_delete(self, name: str, **kwargs) -> LROPoller[None]:
     @experimental
     @monitor_with_activity(ops_logger, "ServerlessEndpoint.GetKeys", ActivityType.PUBLICAPI)
     def get_keys(self, name: str, **kwargs) -> EndpointAuthKeys:
+        """Get serveless endpoint auth keys.
+
+        :param name: The serverless endpoint name
+        :type name: str
+        :return: Returns the keys of the serverless endpoint
+        :rtype: ~azure.ai.ml.entities.EndpointAuthKeys
+        """
         return self._service_client.list_keys(
             self._resource_group_name,
             self._workspace_name,
@@ -111,7 +166,18 @@ def begin_regenerate_keys(
         *,
         key_type: str = EndpointKeyType.PRIMARY_KEY_TYPE,
         **kwargs,
-    ) -> LROPoller[None]:
+    ) -> LROPoller[EndpointAuthKeys]:
+        """Regenerate keys for a serverless endpoint.
+
+        :param name: The endpoint name.
+        :type name: str
+        :keyword key_type: One of "primary", "secondary". Defaults to "primary".
+        :paramtype key_type: str
+        :raises ~azure.ai.ml.exceptions.ValidationException: Raised if key_type is not "primary"
+            or "secondary"
+        :return: A poller to track the operation status.
+        :rtype: ~azure.core.polling.LROPoller[EndpointAuthKeys]
+        """
         keys = self.get_keys(
             name=name,
         )
