aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py')
-rw-r--r--.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py67
1 files changed, 67 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py
new file mode 100644
index 00000000..2f2be676
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/_schema/job/identity.py
@@ -0,0 +1,67 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument
+
+import logging
+
+from marshmallow import fields, post_load
+
+from azure.ai.ml._restclient.v2023_04_01_preview.models import (
+ ConnectionAuthType,
+ IdentityConfigurationType,
+)
+from azure.ai.ml._schema.core.fields import StringTransformedEnum
+from azure.ai.ml._utils.utils import camel_to_snake
+from azure.ai.ml.entities._credentials import (
+ AmlTokenConfiguration,
+ ManagedIdentityConfiguration,
+ UserIdentityConfiguration,
+)
+
+from ..core.schema import PatchedSchemaMeta
+
+module_logger = logging.getLogger(__name__)
+
+
+class ManagedIdentitySchema(metaclass=PatchedSchemaMeta):
+ type = StringTransformedEnum(
+ required=True,
+ allowed_values=[IdentityConfigurationType.MANAGED, ConnectionAuthType.MANAGED_IDENTITY],
+ casing_transform=camel_to_snake,
+ )
+ client_id = fields.Str()
+ object_id = fields.Str()
+ msi_resource_id = fields.Str()
+
+ @post_load
+ def make(self, data, **kwargs):
+ data.pop("type")
+ return ManagedIdentityConfiguration(**data)
+
+
+class AMLTokenIdentitySchema(metaclass=PatchedSchemaMeta):
+ type = StringTransformedEnum(
+ required=True,
+ allowed_values=IdentityConfigurationType.AML_TOKEN,
+ casing_transform=camel_to_snake,
+ )
+
+ @post_load
+ def make(self, data, **kwargs):
+ data.pop("type")
+ return AmlTokenConfiguration(**data)
+
+
+class UserIdentitySchema(metaclass=PatchedSchemaMeta):
+ type = StringTransformedEnum(
+ required=True,
+ allowed_values=IdentityConfigurationType.USER_IDENTITY,
+ casing_transform=camel_to_snake,
+ )
+
+ @post_load
+ def make(self, data, **kwargs):
+ data.pop("type")
+ return UserIdentityConfiguration(**data)