[NOT READY FOR MERGE] feat: support internal spark component (Azure#29271)

* feat: support internal spark component

* feat: support conda file

* refactor: only load internal Spark component as internal

* refactor: inherit serialization logic from 3p spark component

* feat: register internal spark component as public spark component (environment will be lost for now

Author: elliotzh

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/command.py

@@ -3,12 +3,12 @@
 # ---------------------------------------------------------
 from marshmallow import fields
 
-from azure.ai.ml._internal._schema.node import InternalBaseNodeSchema, NodeType
-from azure.ai.ml._schema import AnonymousEnvironmentSchema, ArmVersionedStr, NestedField, RegistryStr, UnionField
-from azure.ai.ml._schema.core.fields import DumpableEnumField
-from azure.ai.ml._schema.job import ParameterizedCommandSchema, ParameterizedParallelSchema
-from azure.ai.ml._schema.job.job_limits import CommandJobLimitsSchema
-from azure.ai.ml.constants._common import AzureMLResourceType
+from ..._schema import AnonymousEnvironmentSchema, ArmVersionedStr, NestedField, RegistryStr, UnionField
+from ..._schema.core.fields import DumpableEnumField
+from ..._schema.job import ParameterizedCommandSchema, ParameterizedParallelSchema
+from ..._schema.job.job_limits import CommandJobLimitsSchema
+from ...constants._common import AzureMLResourceType
+from .._schema.node import InternalBaseNodeSchema, NodeType
 
 
 class CommandSchema(InternalBaseNodeSchema, ParameterizedCommandSchema):

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/component.py

@@ -6,14 +6,15 @@
 import pydash
 from marshmallow import EXCLUDE, INCLUDE, fields, post_dump, pre_load
 
-from azure.ai.ml._schema import NestedField, StringTransformedEnum, UnionField
-from azure.ai.ml._schema.component.component import ComponentSchema
-from azure.ai.ml._schema.core.fields import ArmVersionedStr, CodeField
-from azure.ai.ml.constants._common import LABELLED_RESOURCE_NAME, AzureMLResourceType, SOURCE_PATH_CONTEXT_KEY
-
-from .._utils import yaml_safe_load_with_base_resolver
+from ..._schema import AnonymousEnvironmentSchema, NestedField, StringTransformedEnum, UnionField
+from ..._schema.component.component import ComponentSchema
+from ..._schema.core.fields import ArmVersionedStr, CodeField, RegistryStr
+from ..._schema.job.parameterized_spark import SparkConfSchema, SparkEntryClassSchema, SparkEntryFileSchema
 from ..._utils._arm_id_utils import parse_name_label
 from ..._utils.utils import get_valid_dot_keys_with_wildcard
+from ...constants._common import LABELLED_RESOURCE_NAME, SOURCE_PATH_CONTEXT_KEY, AzureMLResourceType
+from ...constants._component import NodeType as PublicNodeType
+from .._utils import yaml_safe_load_with_base_resolver
 from .environment import InternalEnvironmentSchema
 from .input_output import (
     InternalEnumParameterSchema,
@@ -37,6 +38,9 @@ class NodeType:
     HEMERA = "HemeraComponent"
     AE365EXEPOOL = "AE365ExePoolComponent"
     IPP = "IntellectualPropertyProtectedComponent"
+    # internal spake component got a type value conflict with spark component
+    # this enum is used to identify its create_function in factories
+    SPARK = "DummySpark"
 
     @classmethod
     def all_values(cls):
@@ -161,3 +165,49 @@ def add_back_type_label(self, data, original, **kwargs):  # pylint:disable=unuse
         if type_label:
             data["type"] = LABELLED_RESOURCE_NAME.format(data["type"], type_label)
         return data
+
+
+class InternalSparkComponentSchema(InternalComponentSchema):
+    # type field is required for registration
+    type = StringTransformedEnum(
+        allowed_values=PublicNodeType.SPARK,
+        casing_transform=lambda x: parse_name_label(x)[0].lower(),
+        pass_original=True,
+    )
+
+    environment = UnionField(
+        [
+            # unlike other internal component, internal spark component do not use internal environment schema
+            NestedField(AnonymousEnvironmentSchema),
+            RegistryStr(azureml_type=AzureMLResourceType.ENVIRONMENT),
+            ArmVersionedStr(azureml_type=AzureMLResourceType.ENVIRONMENT, allow_default_version=True),
+            NestedField(InternalEnvironmentSchema),
+        ],
+        allow_none=True,
+    )
+
+    jars = UnionField(
+        [
+            fields.List(fields.Str()),
+            fields.Str(),
+        ],
+    )
+    py_files = UnionField(
+        [
+            fields.List(fields.Str()),
+            fields.Str(),
+        ],
+        data_key="pyFiles",
+        attribute="py_files",
+    )
+
+    entry = UnionField(
+        [NestedField(SparkEntryFileSchema), NestedField(SparkEntryClassSchema)],
+        required=True,
+        metadata={"description": "Entry."},
+    )
+
+    files = fields.List(fields.Str(required=True))
+    archives = fields.List(fields.Str(required=True))
+    conf = NestedField(SparkConfSchema, unknown=INCLUDE)
+    args = fields.Str(metadata={"description": "Command Line arguments."})

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/environment.py

@@ -4,8 +4,8 @@
 
 from marshmallow import fields
 
-from azure.ai.ml._schema import PathAwareSchema
-from azure.ai.ml._schema.core.fields import DumpableEnumField, VersionField
+from ..._schema import PathAwareSchema
+from ..._schema.core.fields import DumpableEnumField, VersionField
 
 
 class InternalEnvironmentSchema(PathAwareSchema):

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/input_output.py

@@ -4,9 +4,9 @@
 
 from marshmallow import fields, post_dump, post_load
 
-from azure.ai.ml._schema import PatchedSchemaMeta, StringTransformedEnum, UnionField
-from azure.ai.ml._schema.component.input_output import InputPortSchema, ParameterSchema
-from azure.ai.ml._schema.core.fields import DumpableEnumField, PrimitiveValueField
+from ..._schema import PatchedSchemaMeta, StringTransformedEnum, UnionField
+from ..._schema.component.input_output import InputPortSchema, ParameterSchema
+from ..._schema.core.fields import DumpableEnumField, PrimitiveValueField
 
 SUPPORTED_INTERNAL_PARAM_TYPES = [
     "integer",

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/node.py

@@ -4,11 +4,10 @@
 
 from marshmallow import INCLUDE, fields, post_load, pre_dump
 
-from azure.ai.ml._schema import ArmVersionedStr, NestedField, RegistryStr, UnionField
-from azure.ai.ml._schema.pipeline.component_job import BaseNodeSchema, _resolve_inputs_outputs
-from azure.ai.ml.constants._common import AzureMLResourceType
-
+from ..._schema import ArmVersionedStr, NestedField, RegistryStr, UnionField
 from ..._schema.core.fields import DumpableEnumField
+from ..._schema.pipeline.component_job import BaseNodeSchema, _resolve_inputs_outputs
+from ...constants._common import AzureMLResourceType
 from .component import InternalComponentSchema, NodeType
 
 
@@ -33,13 +32,13 @@ class Meta:
 
     @post_load
     def make(self, data, **kwargs):  # pylint: disable=unused-argument, no-self-use
-        from azure.ai.ml.entities._builders import parse_inputs_outputs
+        from ...entities._builders import parse_inputs_outputs
 
         # parse inputs/outputs
         data = parse_inputs_outputs(data)
 
         # dict to node object
-        from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
+        from ...entities._job.pipeline._load_component import pipeline_node_factory
 
         return pipeline_node_factory.load_from_dict(data=data)
 

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_setup.py

@@ -3,13 +3,15 @@
 # ---------------------------------------------------------
 
 # pylint: disable=protected-access
-
 from marshmallow import INCLUDE
 
-from azure.ai.ml._internal._schema.command import CommandSchema, DistributedSchema, ParallelSchema
-from azure.ai.ml._internal._schema.component import NodeType
-from azure.ai.ml._internal._schema.node import HDInsightSchema, InternalBaseNodeSchema, ScopeSchema
-from azure.ai.ml._internal.entities import (
+from .._schema import NestedField
+from ..entities._component.component_factory import component_factory
+from ..entities._job.pipeline._load_component import pipeline_node_factory
+from ._schema.command import CommandSchema, DistributedSchema, ParallelSchema
+from ._schema.component import NodeType
+from ._schema.node import HDInsightSchema, InternalBaseNodeSchema, ScopeSchema
+from .entities import (
     Command,
     DataTransfer,
     Distributed,
@@ -21,9 +23,7 @@
     Scope,
     Starlite,
 )
-from azure.ai.ml._schema import NestedField
-from azure.ai.ml.entities._component.component_factory import component_factory
-from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
+from .entities.spark import InternalSparkComponent
 
 _registered = False
 
@@ -41,6 +41,11 @@ def _enable_internal_components():
             create_instance_func=lambda: InternalComponent.__new__(InternalComponent),
             create_schema_func=create_schema_func,
         )
+    component_factory.register_type(
+        _type=NodeType.SPARK,
+        create_instance_func=lambda: InternalSparkComponent.__new__(InternalSparkComponent),
+        create_schema_func=InternalSparkComponent._create_schema_for_validation,
+    )
 
 
 def _register_node(_type, node_cls, schema_cls):

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_additional_includes.py

@@ -11,10 +11,9 @@
 
 import yaml
 
-from azure.ai.ml._utils._asset_utils import IgnoreFile, traverse_directory
-from azure.ai.ml.entities._util import _general_copy
-from azure.ai.ml.entities._validation import MutableValidationResult, _ValidationResultBuilder
-
+from ..._utils._asset_utils import IgnoreFile, traverse_directory
+from ...entities._util import _general_copy
+from ...entities._validation import MutableValidationResult, _ValidationResultBuilder
 from ._artifact_cache import ArtifactCache
 from .code import InternalComponentIgnoreFile
 

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_input_outputs.py

@@ -3,12 +3,12 @@
 # ---------------------------------------------------------
 from typing import Dict, Optional, Union
 
-from azure.ai.ml import Input, Output
-from azure.ai.ml._internal._schema.input_output import SUPPORTED_INTERNAL_PARAM_TYPES
-from azure.ai.ml._utils.utils import get_all_enum_values_iter
-from azure.ai.ml.constants import AssetTypes
-from azure.ai.ml.constants._common import InputTypes
-from azure.ai.ml.constants._component import ComponentParameterTypes, IOConstants
+from ... import Input, Output
+from ..._utils.utils import get_all_enum_values_iter
+from ...constants import AssetTypes
+from ...constants._common import InputTypes
+from ...constants._component import ComponentParameterTypes, IOConstants
+from .._schema.input_output import SUPPORTED_INTERNAL_PARAM_TYPES
 
 _INPUT_TYPE_ENUM = "enum"
 _INPUT_TYPE_ENUM_CAP = "Enum"

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/command.py

@@ -6,16 +6,16 @@
 
 from marshmallow import INCLUDE, Schema
 
-from azure.ai.ml import MpiDistribution, PyTorchDistribution, TensorFlowDistribution
-from azure.ai.ml._internal._schema.component import NodeType
-from azure.ai.ml._internal.entities.component import InternalComponent
-from azure.ai.ml._internal.entities.node import InternalBaseNode
-from azure.ai.ml._restclient.v2023_02_01_preview.models import CommandJobLimits as RestCommandJobLimits
-from azure.ai.ml._restclient.v2023_02_01_preview.models import JobResourceConfiguration as RestJobResourceConfiguration
-from azure.ai.ml._schema import PathAwareSchema
-from azure.ai.ml._schema.core.fields import DistributionField
-from azure.ai.ml.entities import CommandJobLimits, JobResourceConfiguration
-from azure.ai.ml.entities._util import get_rest_dict_for_node_attrs
+from ... import MpiDistribution, PyTorchDistribution, TensorFlowDistribution
+from ..._restclient.v2023_02_01_preview.models import CommandJobLimits as RestCommandJobLimits
+from ..._restclient.v2023_02_01_preview.models import JobResourceConfiguration as RestJobResourceConfiguration
+from ..._schema import PathAwareSchema
+from ..._schema.core.fields import DistributionField
+from ...entities import CommandJobLimits, JobResourceConfiguration
+from ...entities._util import get_rest_dict_for_node_attrs
+from .._schema.component import NodeType
+from ..entities.component import InternalComponent
+from ..entities.node import InternalBaseNode
 
 
 class Command(InternalBaseNode):

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py

@@ -11,20 +11,19 @@
 
 from marshmallow import Schema
 
-from azure.ai.ml._restclient.v2022_05_01.models import ComponentVersionData, ComponentVersionDetails
-from azure.ai.ml._schema import PathAwareSchema
-from azure.ai.ml.entities import Component
-from azure.ai.ml.entities._system_data import SystemData
-from azure.ai.ml.entities._util import convert_ordered_dict_to_dict
-from azure.ai.ml.entities._validation import MutableValidationResult
-
 from ... import Input, Output
+from ..._restclient.v2022_05_01.models import ComponentVersionData, ComponentVersionDetails
+from ..._schema import PathAwareSchema
 from ..._utils._arm_id_utils import parse_name_label
 from ..._utils._asset_utils import IgnoreFile
+from ...entities import Component
 from ...entities._assets import Code
 from ...entities._job.distribution import DistributionConfiguration
+from ...entities._system_data import SystemData
+from ...entities._util import convert_ordered_dict_to_dict
+from ...entities._validation import MutableValidationResult
 from .._schema.component import InternalComponentSchema
-from ._additional_includes import _AdditionalIncludes, ADDITIONAL_INCLUDES_SUFFIX
+from ._additional_includes import ADDITIONAL_INCLUDES_SUFFIX, _AdditionalIncludes
 from ._input_outputs import InternalInput, InternalOutput
 from ._merkle_tree import create_merkletree
 from .code import InternalCode, InternalComponentIgnoreFile