[ml] Cherry-pick type annotation fixes for ApiView (#34682)

* [ml] Fix datalake gen2 datastore type annotation (#34669)

* Remove Optional type annotation from get() methods

* Remove duplicate overloads

* Update type annotations

---------

Co-authored-by: Anna Tisch <[email protected]>

* [ML] Fix operation keyword arg (#34666)

* [ML] Fix operation keyword arg

* Update _model_operations.py docstrings

* Fix docstring

* Datastore docstrings

* Update image_classification_search_space.py type hints

---------

Co-authored-by: Anna Tisch <[email protected]>

Author: diondrapeck

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/azure_storage.py

@@ -17,7 +17,12 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
 from azure.ai.ml._schema._datastore import AzureBlobSchema, AzureDataLakeGen2Schema, AzureFileSchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
-from azure.ai.ml.entities._credentials import AccountKeyConfiguration, SasTokenConfiguration
+from azure.ai.ml.entities._credentials import (
+    AccountKeyConfiguration,
+    CertificateConfiguration,
+    SasTokenConfiguration,
+    ServicePrincipalConfiguration,
+)
 from azure.ai.ml.entities._datastore.datastore import Datastore
 from azure.ai.ml.entities._datastore.utils import from_rest_datastore_credentials
 from azure.ai.ml.entities._util import load_from_dict
@@ -145,7 +150,8 @@ class AzureBlobDatastore(Datastore):
     :param properties: The asset property dictionary.
     :type properties: dict[str, str]
     :param credentials: Credentials to use for Azure ML workspace to connect to the storage.
-    :type credentials: Union[AccountKeySection, SasSection]
+    :type credentials: Union[~azure.ai.ml.entities.AccountKeyConfiguration,
+        ~azure.ai.ml.entities.SasTokenConfiguration]
     :param kwargs: A dictionary of additional configuration parameters.
     :type kwargs: dict
     """
@@ -243,7 +249,11 @@ class AzureDataLakeGen2Datastore(Datastore):
     :param protocol: Protocol to use to connect with the Azure storage account
     :type protocol: str
     :param credentials: Credentials to use for Azure ML workspace to connect to the storage.
-    :type credentials: Union[ServicePrincipalConfiguration, CertificateConfiguration]
+    :type credentials: Union[
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration,
+        ~azure.ai.ml.entities.CertificateConfiguration
+
+        ]
     :param properties: The asset property dictionary.
     :type properties: dict[str, str]
     :param kwargs: A dictionary of additional configuration parameters.
@@ -261,7 +271,7 @@ def __init__(
         endpoint: str = _get_storage_endpoint_from_metadata(),
         protocol: str = HTTPS,
         properties: Optional[Dict] = None,
-        credentials: Optional[Union[AccountKeyConfiguration, SasTokenConfiguration]] = None,
+        credentials: Optional[Union[ServicePrincipalConfiguration, CertificateConfiguration]] = None,
         **kwargs: Any
     ):
         kwargs[TYPE] = DatastoreType.AZURE_DATA_LAKE_GEN2

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_datastore/datastore.py

@@ -13,7 +13,13 @@
 from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
 from azure.ai.ml._utils.utils import camel_to_snake, dump_yaml_to_file
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, CommonYamlFields
-from azure.ai.ml.entities._credentials import NoneCredentialConfiguration
+from azure.ai.ml.entities._credentials import (
+    AccountKeyConfiguration,
+    CertificateConfiguration,
+    NoneCredentialConfiguration,
+    SasTokenConfiguration,
+    ServicePrincipalConfiguration,
+)
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 from azure.ai.ml.entities._resource import Resource
 from azure.ai.ml.entities._util import find_type_in_override
@@ -28,7 +34,14 @@ class Datastore(Resource, RestTranslatableMixin, ABC):
     :param description: Description of the resource.
     :type description: str
     :param credentials: Credentials to use for Azure ML workspace to connect to the storage.
-    :type credentials: Union[ServicePrincipalSection, CertificateSection]
+    :type credentials: Optional[Union[
+        ~azure.ai.ml.entities.ServicePrincipalConfiguration,
+        ~azure.ai.ml.entities.CertificateConfiguration,
+        ~azure.ai.ml.entities.NoneCredentialConfiguration,
+        ~azure.ai.ml.entities.AccountKeyConfiguration,
+        ~azure.ai.ml.entities.SasTokenConfiguration
+
+        ]]
     :param tags: Tag dictionary. Tags can be added, removed, and updated.
     :type tags: dict[str, str]
     :param properties: The asset property dictionary.
@@ -39,7 +52,15 @@ class Datastore(Resource, RestTranslatableMixin, ABC):
 
     def __init__(
         self,
-        credentials: Any,
+        credentials: Optional[
+            Union[
+                ServicePrincipalConfiguration,
+                CertificateConfiguration,
+                NoneCredentialConfiguration,
+                AccountKeyConfiguration,
+                SasTokenConfiguration,
+            ]
+        ],
         name: Optional[str] = None,
         description: Optional[str] = None,
         tags: Optional[Dict] = None,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/image_classification_search_space.py

@@ -18,7 +18,7 @@ class ImageClassificationSearchSpace(RestTranslatableMixin):
     Multilabel tasks.
 
     :param ams_gradient: Enable AMSGrad when optimizer is 'adam' or 'adamw'.
-    :type ams_gradient: str or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
+    :type ams_gradient: bool or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
     :param beta1: Value of 'beta1' when optimizer is 'adam' or 'adamw'. Must be a float in the
         range [0, 1].
     :type beta1: float or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
@@ -85,7 +85,7 @@ class ImageClassificationSearchSpace(RestTranslatableMixin):
     :param training_batch_size: Training batch size. Must be a positive integer.
     :type training_batch_size: int or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
     :param validation_batch_size: Validation batch size. Must be a positive integer.
-    :type validation_batch_size: str or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
+    :type validation_batch_size: int or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
     :param warmup_cosine_lr_cycles: Value of cosine cycle when learning rate scheduler is
         'warmup_cosine'. Must be a float in the range [0, 1].
     :type warmup_cosine_lr_cycles: float or ~azure.ai.ml.entities._job.sweep.search_space.SweepDistribution
@@ -122,37 +122,37 @@ class ImageClassificationSearchSpace(RestTranslatableMixin):
     def __init__(
         self,
         *,
-        ams_gradient: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        beta1: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        beta2: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        distributed: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        early_stopping: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        early_stopping_delay: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        early_stopping_patience: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        enable_onnx_normalization: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        evaluation_frequency: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        gradient_accumulation_step: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        layers_to_freeze: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        learning_rate: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        learning_rate_scheduler: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        model_name: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        momentum: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        nesterov: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        number_of_epochs: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        number_of_workers: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        optimizer: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        random_seed: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        step_lr_gamma: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        step_lr_step_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        training_batch_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        validation_batch_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        warmup_cosine_lr_cycles: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        warmup_cosine_lr_warmup_epochs: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        weight_decay: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        training_crop_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        validation_crop_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        validation_resize_size: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
-        weighted_loss: Optional[Union[bool, int, float, str, SweepDistribution]] = None,
+        ams_gradient: Optional[Union[bool, SweepDistribution]] = None,
+        beta1: Optional[Union[float, SweepDistribution]] = None,
+        beta2: Optional[Union[float, SweepDistribution]] = None,
+        distributed: Optional[Union[bool, SweepDistribution]] = None,
+        early_stopping: Optional[Union[bool, SweepDistribution]] = None,
+        early_stopping_delay: Optional[Union[int, SweepDistribution]] = None,
+        early_stopping_patience: Optional[Union[int, SweepDistribution]] = None,
+        enable_onnx_normalization: Optional[Union[bool, SweepDistribution]] = None,
+        evaluation_frequency: Optional[Union[int, SweepDistribution]] = None,
+        gradient_accumulation_step: Optional[Union[int, SweepDistribution]] = None,
+        layers_to_freeze: Optional[Union[int, SweepDistribution]] = None,
+        learning_rate: Optional[Union[float, SweepDistribution]] = None,
+        learning_rate_scheduler: Optional[Union[str, SweepDistribution]] = None,
+        model_name: Optional[Union[str, SweepDistribution]] = None,
+        momentum: Optional[Union[float, SweepDistribution]] = None,
+        nesterov: Optional[Union[bool, SweepDistribution]] = None,
+        number_of_epochs: Optional[Union[int, SweepDistribution]] = None,
+        number_of_workers: Optional[Union[int, SweepDistribution]] = None,
+        optimizer: Optional[Union[str, SweepDistribution]] = None,
+        random_seed: Optional[Union[int, SweepDistribution]] = None,
+        step_lr_gamma: Optional[Union[float, SweepDistribution]] = None,
+        step_lr_step_size: Optional[Union[int, SweepDistribution]] = None,
+        training_batch_size: Optional[Union[int, SweepDistribution]] = None,
+        validation_batch_size: Optional[Union[int, SweepDistribution]] = None,
+        warmup_cosine_lr_cycles: Optional[Union[float, SweepDistribution]] = None,
+        warmup_cosine_lr_warmup_epochs: Optional[Union[int, SweepDistribution]] = None,
+        weight_decay: Optional[Union[float, SweepDistribution]] = None,
+        training_crop_size: Optional[Union[int, SweepDistribution]] = None,
+        validation_crop_size: Optional[Union[int, SweepDistribution]] = None,
+        validation_resize_size: Optional[Union[int, SweepDistribution]] = None,
+        weighted_loss: Optional[Union[int, SweepDistribution]] = None,
     ) -> None:
         self.ams_gradient = ams_gradient
         self.beta1 = beta1

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

@@ -285,10 +285,10 @@ def get(self, name: str, version: Optional[str] = None, label: Optional[str] = N
 
         :param name: Name of the model.
         :type name: str
-        :keyword version: Version of the model.
-        :paramtype version: str
-        :keyword label: Label of the model. (mutually exclusive with version)
-        :paramtype label: str
+        :param version: Version of the model.
+        :type version: str
+        :param label: Label of the model. (mutually exclusive with version)
+        :type label: str
         :raises ~azure.ai.ml.exceptions.ValidationException: Raised if Model cannot be successfully validated.
             Details will be provided in the error message.
         :return: Model asset object.
@@ -329,9 +329,9 @@ def download(self, name: str, version: str, download_path: Union[PathLike, str]
         :type name: str
         :param version: Version of the model.
         :type version: str
-        :keyword download_path: Local path as download destination, defaults to current working directory of the current
+        :param download_path: Local path as download destination, defaults to current working directory of the current
             user. Contents will be overwritten.
-        :paramtype download_path: Union[PathLike, str]
+        :type download_path: Union[PathLike, str]
         :raises ResourceNotFoundError: if can't find a model matching provided name.
         """
 
@@ -473,10 +473,10 @@ def list(
     ) -> Iterable[Model]:
         """List all model assets in workspace.
 
-        :keyword name: Name of the model.
-        :paramtype name: Optional[str]
-        :keyword stage: The Model stage
-        :paramtype stage: Optional[str]
+        :param name: Name of the model.
+        :type name: Optional[str]
+        :param stage: The Model stage
+        :type stage: Optional[str]
         :keyword list_view_type: View type for including/excluding (for example) archived models.
             Defaults to :attr:`ListViewType.ACTIVE_ONLY`.
         :paramtype list_view_type: ListViewType

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

@@ -133,15 +133,15 @@ def delete(self, name: str) -> None:
     @monitor_with_activity(ops_logger, "WorkspaceConnections.List", ActivityType.PUBLICAPI)
     def list(
         self,
-        *,
         connection_type: Optional[str] = None,
+        *,
         include_data_connections: bool = False,
         **kwargs: Any,
     ) -> Iterable[WorkspaceConnection]:
         """List all workspace connections for a workspace.
 
-        :keyword connection_type: Type of workspace connection to list.
-        :paramtype connection_type: Optional[str]
+        :param connection_type: Type of workspace connection to list.
+        :type connection_type: Optional[str]
         :keyword include_data_connections: If true, also return data connections. Defaults to False.
         :paramtype include_data_connections: bool
         :return: An iterator like instance of workspace connection objects