[ML][Pipelines] feat: handle corner serialization issues for internal components (Azure#28390)

* feat: handle corner case serialization for internal components

* feat: support attr datatransfer in internal components

* fix: bandit

* refactor: extract safe load with base resolver to a util

* ci: fix ci

Author: elliotzh

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

@@ -1,8 +1,8 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-# pylint: disable=unused-import
-from ._util import enable_internal_components_in_pipeline
+
+from ._setup import enable_internal_components_in_pipeline
 from .entities import (
     Ae365exepool,
     AISuperComputerConfiguration,

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

@@ -1,14 +1,17 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
+import os.path
 
+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
+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 ..._utils._arm_id_utils import parse_name_label
 from .environment import InternalEnvironmentSchema
 from .input_output import (
@@ -110,29 +113,24 @@ def _serialize(self, obj, *, many: bool = False):
                 ret[attr_name] = self.get_attribute(obj, attr_name, None)
         return ret
 
-    @pre_load()
-    def convert_input_value_to_str(self, data: dict, **kwargs) -> dict:
-        """
-        Convert the non-str value in input to str.
-
-        When load the v1.5 component yaml, true/false will be converted to bool type and yyyy-mm-dd will be
-        converted to date type. In order to be consistent with before, it needs to be converted to str type.
-        """
-        def convert_to_str(value):
-            if isinstance(value, bool):
-                return str(value).lower()
-            return str(value)
-
-        if "inputs" in data and isinstance(data["inputs"], dict):
-            for input_port in data["inputs"].values():
-                input_type = input_port["type"]
-                # input type can be a list for internal component
-                if isinstance(input_type, str) and input_type.lower() in ["string", "enum"]:
-                    if not isinstance(input_port.get("default", ""), str):
-                        input_port["default"] = convert_to_str(input_port["default"])
-                    if "enum" in input_port and any([not isinstance(item, str) for item in input_port["enum"]]):
-                        input_port["enum"] = [convert_to_str(item) for item in input_port["enum"]]
-        return data
+    # override param_override to ensure that param override happens after reloading the yaml
+    @pre_load
+    def add_param_overrides(self, data, **kwargs):
+        source_path = self.context.pop(SOURCE_PATH_CONTEXT_KEY, None)
+        if isinstance(data, dict) and source_path and os.path.isfile(source_path):
+            # do override here
+            with open(source_path, "r") as f:
+                origin_data = yaml_safe_load_with_base_resolver(f)
+                dot_keys = ["version"]
+                for input_key in data.get("inputs", {}).keys():
+                    # Keep value in float input as string to avoid precision issue.
+                    for attr_name in ["default", "enum", "min", "max"]:
+                        dot_keys.append(f"inputs.{input_key}.{attr_name}")
+
+                for dot_key in dot_keys:
+                    if pydash.has(data, dot_key) and pydash.has(origin_data, dot_key):
+                        pydash.set_(data, dot_key, pydash.get(origin_data, dot_key))
+        return super().add_param_overrides(data, **kwargs)
 
     @post_dump(pass_original=True)
     def simplify_input_output_port(self, data, original, **kwargs):  # pylint:disable=unused-argument, no-self-use
@@ -143,13 +141,6 @@ def simplify_input_output_port(self, data, original, **kwargs):  # pylint:disabl
 
         # hack, to match current serialization match expectation
         for port_name, port_definition in data["inputs"].items():
-            input_type = port_definition.get("type", None)
-            is_float_type = isinstance(input_type, str) and input_type.lower() == "float"
-            for key in ["default", "min", "max"]:
-                if key in port_definition:
-                    value = getattr(original.inputs[port_name], key)
-                    # Keep value in float input as string to avoid precision issue.
-                    data["inputs"][port_name][key] = str(value) if is_float_type else value
             if "mode" in port_definition:
                 del port_definition["mode"]
 

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_util.py renamed to sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_setup.py

@@ -0,0 +1,9 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from ._yaml_utils import yaml_safe_load_with_base_resolver
+
+__all__ = [
+    "yaml_safe_load_with_base_resolver"
+]

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_utils/__init__.py

@@ -0,0 +1,47 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+import strictyaml
+
+
+class _SafeLoaderWithBaseLoader(strictyaml.ruamel.SafeLoader):
+    """This is a SafeLoader with base resolver instead of version default resolver.
+
+    Differences between BaseResolver and VersionedResolver:
+    1) BaseResolver won't try to resolve node value. For example, "yes" and "no" will be resolved to "true"(bool)
+    and "false"(bool) by VersionedResolver, but won't be resolved by BaseResolver.
+    2) VersionedResolver will delay loading the pattern matching rules to pass yaml versions on loading.
+
+    Given SafeLoader inherits from VersionedResolver, we can't directly remove VersionedResolver
+    from the inheritance list. Instead, we overwrite add_version_implicit_resolver method to make
+    _SafeLoaderWithBaseLoader._version_implicit_resolver empty. Then the resolver will act like a BaseResolver.
+    """
+    def fetch_comment(self, comment):
+        pass
+
+    def add_version_implicit_resolver(self, version, tag, regexp, first):
+        """Overwrite the method to make the resolver act like a base resolver instead of version default resolver.
+
+        :param version: version of yaml, like (1, 1)(yaml 1.1) and (1, 2)(yaml 1.2)
+        :type version: VersionType
+        :param tag: a tag indicating the type of the resolved node, e.g., tag:yaml.org,2002:bool.
+        :param regexp: the regular expression to match the node to be resolved
+        :param first: a list of first characters to match
+        """
+        self._version_implicit_resolver.setdefault(version, {})
+
+
+def yaml_safe_load_with_base_resolver(stream):
+    """Load yaml string with base resolver instead of version default resolver.
+    For example:
+    1) "yes" and "no" will be loaded as "yes"(string) and "no"(string) instead of "true"(bool) and "false"(bool);
+    2) "0.10" will be loaded as "0.10"(string) instead of "0.1"(float).
+    3) "2019-01-01" will be loaded as "2019-01-01"(string) instead of "2019-01-01T00:00:00Z"(datetime).
+    4) "1" will be loaded as "1"(string) instead of "1"(int).
+    5) "1.0" will be loaded as "1.0"(string) instead of "1.0"(float).
+    6) "~" will be loaded as "~"(string) instead of "None"(NoneType).
+    Please refer to strictyaml.ruamel.resolver.implicit_resolvers for more details.
+    """
+    return strictyaml.ruamel.load(stream, Loader=_SafeLoaderWithBaseLoader)

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_utils/_yaml_utils.py

@@ -96,6 +96,7 @@ def __init__(
         starlite: Optional[Dict] = None,
         ae365exepool: Optional[Dict] = None,
         launcher: Optional[Dict] = None,
+        datatransfer: Optional[Dict] = None,
         **kwargs,
     ):
         type, self._type_label = parse_name_label(type)
@@ -135,6 +136,7 @@ def __init__(
         self.starlite = starlite
         self.ae365exepool = ae365exepool
         self.launcher = launcher
+        self.datatransfer = datatransfer
 
     @classmethod
     def _build_io(cls, io_dict: Union[Dict, Input, Output], is_input: bool):

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

@@ -26,6 +26,7 @@
     REGISTRY_URI_FORMAT,
     CommonYamlFields,
     AzureMLResourceType,
+    SOURCE_PATH_CONTEXT_KEY,
 )
 from azure.ai.ml.constants._component import ComponentSource, NodeType, IOConstants
 from azure.ai.ml.entities._assets import Code
@@ -328,6 +329,7 @@ def _load(
                 create_schema_func(
                     {
                         BASE_PATH_CONTEXT_KEY: base_path,
+                        SOURCE_PATH_CONTEXT_KEY: yaml_path,
                         PARAMS_OVERRIDE_KEY: params_override,
                     }
                 ).load(data, unknown=INCLUDE, **kwargs)

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

@@ -486,4 +486,4 @@ def _map_from_rest_type(cls, _type):
     def _from_rest_object(cls, obj: Dict) -> "Input":
         obj["type"] = cls._map_from_rest_type(obj["type"])
 
-        return Input(**obj)
+        return cls(**obj)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/input.py

@@ -854,7 +854,7 @@ def disable_internal_components():
     and enable_private_preview_features, as the execution order of fixtures is not guaranteed.
     """
     from azure.ai.ml._internal._schema.component import NodeType
-    from azure.ai.ml._internal._util import _set_registered
+    from azure.ai.ml._internal._setup import _set_registered
     from azure.ai.ml.entities._component.component_factory import component_factory
     from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
 

sdk/ml/azure-ai-ml/tests/conftest.py

@@ -84,6 +84,8 @@ def test_component_load(
         randstr: Callable[[str], str],
         yaml_path: str,
     ) -> None:
+        if "ae365" not in yaml_path:
+            return
         omit_fields = ["id", "creation_context", "code", "name"]
         component_name = randstr("component_name")
 
@@ -101,6 +103,12 @@ def test_component_load(
             expected_dict = json.load(f)
             expected_dict["_source"] = "REMOTE.WORKSPACE.COMPONENT"
 
+            # default value for datatransfer
+            if expected_dict["type"] == "DataTransferComponent" and "datatransfer" not in expected_dict:
+                expected_dict["datatransfer"] = {
+                    'allow_overwrite': 'True'
+                }
+
             # TODO: check if loaded environment is expected to be an ordered dict
             assert pydash.omit(loaded_dict, *omit_fields) == pydash.omit(expected_dict, *omit_fields)