[AzureML] Fix BatchDeployment typevar signature + other API review comments (Azure#30354)

* Fix typevar signature

* Fix pylint

* black

* resolve last few api comments

* fix circular import

* kwargs passthrough

---------

Co-authored-by: Neehar Duvvuri <[email protected]>

Author: annatisch

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/core/intellectual_property.py

@@ -9,7 +9,7 @@
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils.utils import camel_to_snake
-from azure.ai.ml.constants._common import IPProtectionLevel
+from azure.ai.ml.constants._assets import IPProtectionLevel
 from azure.ai.ml.entities._assets.intellectual_property import IntellectualProperty
 
 

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

@@ -5,6 +5,7 @@
 
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
+from ._assets import IPProtectionLevel
 from ._common import AssetTypes, InputOutputModes, ModelType, Scope, TimeZone
 from ._component import ParallelTaskType
 from ._deployment import BatchDeploymentOutputAction
@@ -56,4 +57,5 @@
     "MonitorModelType",
     "MonitorFeatureType",
     "MonitorDatasetContext",
+    "IPProtectionLevel",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_assets.py

@@ -0,0 +1,13 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from enum import Enum
+
+from azure.ai.ml._utils._experimental import experimental
+from azure.core import CaseInsensitiveEnumMeta
+
+
+@experimental
+class IPProtectionLevel(str, Enum, metaclass=CaseInsensitiveEnumMeta):
+    ALL = "all"
+    NONE = "none"

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -852,11 +852,6 @@ class InferenceServerType:
     CUSTOM = "custom"
 
 
-class IPProtectionLevel(str, Enum, metaclass=CaseInsensitiveEnumMeta):
-    ALL = "all"
-    NONE = "none"
-
-
 class AzureDevopsArtifactsType:
     ARTIFACT = "artifact"
 

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

@@ -67,10 +67,12 @@
     OnlineDeployment,
 )
 from ._deployment.data_collector import DataCollector
+from ._deployment.deployment_collection import DeploymentCollection
 from ._deployment.model_batch_deployment import ModelBatchDeployment
 from ._deployment.model_batch_deployment_settings import ModelBatchDeploymentSettings
 from ._deployment.pipeline_component_batch_deployment import PipelineComponentBatchDeployment
 from ._deployment.resource_requirements_settings import ResourceRequirementsSettings
+from ._deployment.request_logging import RequestLogging
 from ._deployment.scale_settings import DefaultScaleSettings, TargetUtilizationScaleSettings, OnlineScaleSettings
 from ._endpoint.batch_endpoint import BatchEndpoint
 from ._endpoint.endpoint import Endpoint
@@ -408,4 +410,6 @@
     "CustomMonitoringMetricThreshold",
     "DataCollector",
     "IntellectualProperty",
+    "DeploymentCollection",
+    "RequestLogging",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/intellectual_property.py

@@ -4,7 +4,7 @@
 
 from typing import Any
 from azure.ai.ml._utils._experimental import experimental
-from azure.ai.ml.constants._common import IPProtectionLevel
+from azure.ai.ml.constants._assets import IPProtectionLevel
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 from azure.ai.ml._restclient.v2023_04_01_preview.models import IntellectualProperty as RestIntellectualProperty
 
@@ -15,7 +15,7 @@ class IntellectualProperty(RestTranslatableMixin):
     :param publisher: Publisher name
     :type publisher: str
     :param protection_level: Asset Protection Level. Allowed values: ALL, NONE
-    :type protection_level: IPProtectionLevel
+    :type protection_level: ~azure.ai.ml.constants.IPProtectionLevel
     """
 
     def __init__(self, *, publisher: str = None, protection_level: IPProtectionLevel = IPProtectionLevel.ALL):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_deployment/deployment_collection.py

@@ -7,8 +7,10 @@
 from azure.ai.ml._schema._deployment.online.deployment_collection_schema import DeploymentCollectionSchema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml._restclient.v2023_04_01_preview.models import Collection as RestCollection
+from azure.ai.ml._utils._experimental import experimental
 
 
+@experimental
 class DeploymentCollection:
     """Collection entity
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_deployment/request_logging.py

@@ -7,8 +7,10 @@
 from azure.ai.ml._schema._deployment.online.request_logging_schema import RequestLoggingSchema
 from azure.ai.ml._restclient.v2023_04_01_preview.models import RequestLogging as RestRequestLogging
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
+from azure.ai.ml._utils._experimental import experimental
 
 
+@experimental
 class RequestLogging:
     """Request Logging deployment entity.
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_registry/registry.py

@@ -58,8 +58,8 @@ def __init__(
         :type public_network_access: str
         :param discovery_url: Backend service base url for the registry.
         :type discovery_url: str
-        :param intellectual_property_publisher: Intellectual property publisher.
-        :type intellectual_property_publisher: str
+        :param intellectual_property: **Experimental** Intellectual property publisher.
+        :type intellectual_property: ~azure.ai.ml.entities.IntellectualProperty
         :param managed_resource_group: Managed resource group created for the registry.
         :type managed_resource_group: str
         :param mlflow_registry_uri: Ml flow tracking uri for the registry.

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

@@ -39,7 +39,9 @@
 
 ops_logger = OpsLogger(__name__)
 logger, module_logger = ops_logger.package_logger, ops_logger.module_logger
-DeploymentType = TypeVar("DeploymentType", bound=Union[BatchDeployment, PipelineComponentBatchDeployment])
+DeploymentType = TypeVar(
+    "DeploymentType", bound=Union[BatchDeployment, PipelineComponentBatchDeployment, ModelBatchDeployment]
+)
 
 
 class BatchDeploymentOperations(_ScopeDependentOperations):
@@ -76,7 +78,7 @@ def __init__(
     @monitor_with_activity(logger, "BatchDeployment.BeginCreateOrUpdate", ActivityType.PUBLICAPI)
     def begin_create_or_update(
         self,
-        deployment: Union[BatchDeployment, PipelineComponentBatchDeployment, ModelBatchDeployment],
+        deployment: DeploymentType,
         *,
         skip_script_validation: bool = False,
         **kwargs,