Migrate 6 API versions to shared arm_ml_service client

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_datastore/one_lake.py

@@ -8,7 +8,7 @@
 
 from marshmallow import Schema, fields, post_load
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType, OneLakeArtifactType
+from azure.ai.ml._restclient.arm_ml_service.models import DatastoreType, OneLakeArtifactType
 from azure.ai.ml._schema.core.fields import NestedField, PathAwareSchema, StringTransformedEnum, UnionField
 from azure.ai.ml._utils.utils import camel_to_snake
 

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

@@ -27,28 +27,28 @@
 from azure.ai.ml._restclient.arm_ml_service.models import (
     WorkspaceConnectionUsernamePassword as RestWorkspaceConnectionUsernamePassword,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AccountKeyDatastoreCredentials as RestAccountKeyDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AccountKeyDatastoreSecrets as RestAccountKeyDatastoreSecrets,
 )
 from azure.ai.ml._restclient.v2023_04_01_preview.models import AmlToken as RestAmlToken
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     CertificateDatastoreCredentials as RestCertificateDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import CertificateDatastoreSecrets, CredentialsType
+from azure.ai.ml._restclient.arm_ml_service.models import CertificateDatastoreSecrets, CredentialsType
 from azure.ai.ml._restclient.v2023_04_01_preview.models import IdentityConfiguration as RestJobIdentityConfiguration
 from azure.ai.ml._restclient.v2023_04_01_preview.models import IdentityConfigurationType
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ManagedIdentity as RestJobManagedIdentity
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ManagedServiceIdentity as RestRegistryManagedIdentity
-from azure.ai.ml._restclient.v2023_04_01_preview.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
-from azure.ai.ml._restclient.v2023_04_01_preview.models import SasDatastoreCredentials as RestSasDatastoreCredentials
-from azure.ai.ml._restclient.v2023_04_01_preview.models import SasDatastoreSecrets as RestSasDatastoreSecrets
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
+from azure.ai.ml._restclient.arm_ml_service.models import SasDatastoreCredentials as RestSasDatastoreCredentials
+from azure.ai.ml._restclient.arm_ml_service.models import SasDatastoreSecrets as RestSasDatastoreSecrets
+from azure.ai.ml._restclient.arm_ml_service.models import (
     ServicePrincipalDatastoreCredentials as RestServicePrincipalDatastoreCredentials,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     ServicePrincipalDatastoreSecrets as RestServicePrincipalDatastoreSecrets,
 )
 from azure.ai.ml._restclient.v2023_04_01_preview.models import UserIdentity as RestUserIdentity

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

@@ -7,11 +7,11 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AzureDataLakeGen1Datastore as RestAzureDatalakeGen1Datastore,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
+from azure.ai.ml._restclient.arm_ml_service.models import Datastore as DatastoreData
+from azure.ai.ml._restclient.arm_ml_service.models import DatastoreType
 from azure.ai.ml._schema._datastore.adls_gen1 import AzureDataLakeGen1Schema
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE
 from azure.ai.ml.entities._credentials import CertificateConfiguration, ServicePrincipalConfiguration

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

@@ -8,13 +8,13 @@
 from typing import Any, Dict, Optional, Union
 
 from azure.ai.ml._azure_environments import _get_storage_endpoint_from_metadata
-from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureBlobDatastore as RestAzureBlobDatastore
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import AzureBlobDatastore as RestAzureBlobDatastore
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AzureDataLakeGen2Datastore as RestAzureDataLakeGen2Datastore,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureFileDatastore as RestAzureFileDatastore
-from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
+from azure.ai.ml._restclient.arm_ml_service.models import AzureFileDatastore as RestAzureFileDatastore
+from azure.ai.ml._restclient.arm_ml_service.models import Datastore as DatastoreData
+from azure.ai.ml._restclient.arm_ml_service.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 (

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

@@ -9,8 +9,8 @@
 from pathlib import Path
 from typing import IO, Any, AnyStr, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
+from azure.ai.ml._restclient.arm_ml_service.models import Datastore as DatastoreData
+from azure.ai.ml._restclient.arm_ml_service.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 (

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

@@ -8,11 +8,11 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2023_04_01_preview.models import Datastore as DatastoreData
-from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType
-from azure.ai.ml._restclient.v2023_04_01_preview.models import LakeHouseArtifact as RestLakeHouseArtifact
-from azure.ai.ml._restclient.v2023_04_01_preview.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
-from azure.ai.ml._restclient.v2023_04_01_preview.models import OneLakeDatastore as RestOneLakeDatastore
+from azure.ai.ml._restclient.arm_ml_service.models import Datastore as DatastoreData
+from azure.ai.ml._restclient.arm_ml_service.models import DatastoreType
+from azure.ai.ml._restclient.arm_ml_service.models import LakeHouseArtifact as RestLakeHouseArtifact
+from azure.ai.ml._restclient.arm_ml_service.models import NoneDatastoreCredentials as RestNoneDatastoreCredentials
+from azure.ai.ml._restclient.arm_ml_service.models import OneLakeDatastore as RestOneLakeDatastore
 from azure.ai.ml._schema._datastore.one_lake import OneLakeSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, TYPE

sdk/ml/azure-ai-ml/tests/datastore/unittests/test_datastore_schema.py

@@ -3,16 +3,17 @@
 
 import azure.ai.ml._schema._datastore as DatastoreSchemaDir
 from azure.ai.ml import load_datastore
+from azure.ai.ml._restclient.arm_ml_service import models as models_arm
 from azure.ai.ml._restclient.v2023_04_01_preview import models as models_preview
-from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureBlobDatastore as RestAzureBlobDatastore
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import AzureBlobDatastore as RestAzureBlobDatastore
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AzureDataLakeGen1Datastore as RestAzureDataLakeGen1Datastore,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import (
     AzureDataLakeGen2Datastore as RestAzureDataLakeGen2Datastore,
 )
-from azure.ai.ml._restclient.v2023_04_01_preview.models import AzureFileDatastore as RestAzureFileDatastore
-from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+from azure.ai.ml._restclient.arm_ml_service.models import AzureFileDatastore as RestAzureFileDatastore
+from azure.ai.ml._restclient.arm_ml_service.models import (
     NoneDatastoreCredentials,
     ServicePrincipalDatastoreCredentials,
 )
@@ -271,7 +272,7 @@ def test_credential_less_one_lake_schema(self):
         datastore_resource.name = internal_ds.name
         ds_properties = datastore_resource.properties
         assert ds_properties
-        assert isinstance(ds_properties, models_preview.OneLakeDatastore)
+        assert isinstance(ds_properties, models_arm.OneLakeDatastore)
         assert isinstance(ds_properties.credentials, NoneDatastoreCredentials)
         assert ds_properties.one_lake_workspace_name == cfg["one_lake_workspace_name"]
         assert ds_properties.endpoint == cfg["endpoint"]
@@ -315,7 +316,7 @@ def validate_one_lake_schema_with_service_principal(self, test_yaml_file_path, a
         datastore_resource.name = internal_ds.name
         ds_properties = datastore_resource.properties
         assert ds_properties
-        assert isinstance(ds_properties, models_preview.OneLakeDatastore)
+        assert isinstance(ds_properties, models_arm.OneLakeDatastore)
         assert isinstance(ds_properties.credentials, ServicePrincipalDatastoreCredentials)
         assert ds_properties.one_lake_workspace_name == cfg["one_lake_workspace_name"]
         assert ds_properties.endpoint == cfg["endpoint"]

sdk/ml/azure-ai-ml/tests/datastore/unittests/test_datastore_serialization_regression.py

@@ -0,0 +1,95 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Regression test for the datastore create_or_update mixed-tree serialization bug.
+
+Context (production regression, azure-ai-ml 1.34.0/1.35.0):
+``DatastoreOperations.create_or_update`` was migrated to the arm_ml_service client (its operation
+serializes the request body with ``json.dumps(body, cls=SdkJSONEncoder, ...)``), but the datastore
+ENTITIES still returned an old ``v2023_04_01_preview`` msrest ``Datastore`` model. ``SdkJSONEncoder``
+can only serialize arm hybrid models, so a msrest datastore body raised
+``TypeError: Object of type Datastore is not JSON serializable`` on every create/update — a hard
+failure for customers, invisible to the mocked unit tests and the (skipped/live-only) e2e tests.
+
+This test freezes the fix: every datastore entity's ``_to_rest_object()`` must serialize cleanly
+through the arm operation's encoder. It is the offline Class-A (mixed-tree) guard the original
+migration lacked.
+"""
+import json
+
+import pytest
+
+from azure.ai.ml._restclient.arm_ml_service._utils.model_base import SdkJSONEncoder
+from azure.ai.ml.entities._credentials import (
+    AccountKeyConfiguration,
+    ServicePrincipalConfiguration,
+)
+from azure.ai.ml.entities._datastore.adls_gen1 import AzureDataLakeGen1Datastore
+from azure.ai.ml.entities._datastore.azure_storage import (
+    AzureBlobDatastore,
+    AzureDataLakeGen2Datastore,
+    AzureFileDatastore,
+)
+from azure.ai.ml.entities._datastore.one_lake import OneLakeDatastore
+
+
+def _serializes_via_arm_operation(rest_obj) -> dict:
+    """Serialize exactly the way ``DatastoreOperations.create_or_update`` does on the arm client."""
+    return json.loads(json.dumps(rest_obj, cls=SdkJSONEncoder, exclude_readonly=True))
+
+
+@pytest.mark.unittest
+@pytest.mark.data_experiences_test
+class TestDatastoreSerializationRegression:
+    def test_blob_datastore_body_serializes(self) -> None:
+        ds = AzureBlobDatastore(
+            name="ds",
+            account_name="acct",
+            container_name="container",
+            credentials=AccountKeyConfiguration(account_key="fake"),
+        )
+        wire = _serializes_via_arm_operation(ds._to_rest_object())
+        assert wire["properties"]["datastoreType"] == "AzureBlob"
+        assert wire["properties"]["accountName"] == "acct"
+
+    def test_file_datastore_body_serializes(self) -> None:
+        ds = AzureFileDatastore(
+            name="ds",
+            account_name="acct",
+            file_share_name="share",
+            credentials=AccountKeyConfiguration(account_key="fake"),
+        )
+        wire = _serializes_via_arm_operation(ds._to_rest_object())
+        assert wire["properties"]["datastoreType"] == "AzureFile"
+
+    def test_adls_gen2_datastore_body_serializes(self) -> None:
+        ds = AzureDataLakeGen2Datastore(
+            name="ds",
+            account_name="acct",
+            filesystem="fs",
+            credentials=ServicePrincipalConfiguration(tenant_id="t", client_id="c", client_secret="s"),
+        )
+        wire = _serializes_via_arm_operation(ds._to_rest_object())
+        assert wire["properties"]["datastoreType"] == "AzureDataLakeGen2"
+
+    def test_adls_gen1_datastore_body_serializes(self) -> None:
+        ds = AzureDataLakeGen1Datastore(
+            name="ds",
+            store_name="store",
+            credentials=ServicePrincipalConfiguration(tenant_id="t", client_id="c", client_secret="s"),
+        )
+        wire = _serializes_via_arm_operation(ds._to_rest_object())
+        assert wire["properties"]["datastoreType"] == "AzureDataLakeGen1"
+
+    def test_one_lake_datastore_body_serializes(self) -> None:
+        from azure.ai.ml.entities._datastore.one_lake import LakeHouseArtifact
+
+        ds = OneLakeDatastore(
+            name="ds",
+            one_lake_workspace_name="ws",
+            endpoint="onelake.dfs.fabric.microsoft.com",
+            artifact=LakeHouseArtifact(name="artifact"),
+            credentials=ServicePrincipalConfiguration(tenant_id="t", client_id="c", client_secret="s"),
+        )
+        wire = _serializes_via_arm_operation(ds._to_rest_object())
+        assert wire["properties"]["datastoreType"] == "OneLake"