fix: pop _source and yaml_str (Azure#29335)

Author: elliotzh

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

@@ -323,20 +323,25 @@ def _load(
             for_load=True,
         )
         new_instance = create_instance_func()
-        init_kwargs = new_instance._load_with_schema(  # pylint: disable=protected-access
-            data,
-            context={
-                BASE_PATH_CONTEXT_KEY: base_path,
-                SOURCE_PATH_CONTEXT_KEY: yaml_path,
-                PARAMS_OVERRIDE_KEY: params_override,
-            },
-            unknown=INCLUDE,
-            raise_original_exception=True,
-            **kwargs,
+        # specific keys must be popped before loading with schema using kwargs
+        init_kwargs = {
+            "yaml_str": kwargs.pop("yaml_str", None),
+            "_source": kwargs.pop("_source", ComponentSource.YAML_COMPONENT),
+        }
+        init_kwargs.update(
+            new_instance._load_with_schema(  # pylint: disable=protected-access
+                data,
+                context={
+                    BASE_PATH_CONTEXT_KEY: base_path,
+                    SOURCE_PATH_CONTEXT_KEY: yaml_path,
+                    PARAMS_OVERRIDE_KEY: params_override,
+                },
+                unknown=INCLUDE,
+                raise_original_exception=True,
+                **kwargs,
+            )
         )
         new_instance.__init__(
-            yaml_str=kwargs.pop("yaml_str", None),
-            _source=kwargs.pop("_source", ComponentSource.YAML_COMPONENT),
             **init_kwargs,
         )
         # Set base path separately to avoid doing this in post load, as return types of post load are not unified,

sdk/ml/azure-ai-ml/tests/component/unittests/test_component_schema.py

@@ -7,11 +7,11 @@
 import pydash
 import pytest
 import yaml
-
 from azure.ai.ml import MLClient, load_component
 from azure.ai.ml._restclient.v2022_05_01.models import ComponentVersionData
 from azure.ai.ml._utils._arm_id_utils import PROVIDER_RESOURCE_ID_WITH_VERSION
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, AssetTypes, LegacyAssetTypes
+from azure.ai.ml.constants._component import ComponentSource
 from azure.ai.ml.entities import CommandComponent, Component, PipelineComponent
 from azure.ai.ml.entities._assets import Code
 from azure.ai.ml.entities._component.component import COMPONENT_CODE_PLACEHOLDER, COMPONENT_PLACEHOLDER
@@ -316,6 +316,7 @@ def test_component_factory(self):
             context={
                 "source_path": test_path,
             },
+            _source=ComponentSource.YAML_COMPONENT,
         )
         assert recreated_component._to_dict() == component_entity._to_dict()