Data Collector Data field string fix (Azure#34739)

* data field for data collector should pass in a data id for the rest call * update import order
azure-sdk · Mar 13, 2024 · 04ad332 · 04ad332
1 parent f12a5b9
commit 04ad332
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_deployment/online/deployment_collection_schema.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_deployment/online/deployment_collection_schema.py
@@ -18,8 +18,8 @@ class DeploymentCollectionSchema(metaclass=PatchedSchemaMeta):
     enabled = StringTransformedEnum(required=True, allowed_values=[Boolean.TRUE, Boolean.FALSE])
     data = UnionField(
         [
-            fields.Str(),
             NestedField(DataAssetSchema),
+            fields.Str(),
         ]
     )
     client_id = fields.Str()

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_deployment/deployment_collection.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_deployment/deployment_collection.py
@@ -2,12 +2,13 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import Collection as RestCollection
 from azure.ai.ml._schema._deployment.online.deployment_collection_schema import DeploymentCollectionSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
+from .data_asset import DataAsset
 
 
 @experimental
@@ -27,7 +28,7 @@ def __init__(
         self,
         *,
         enabled: Optional[str] = None,
-        data: Optional[str] = None,
+        data: Optional[Union[str, DataAsset]] = None,
         client_id: Optional[str] = None,
         **kwargs: Any
     ):