[ML][Pipelines] Remove prefix in pipeline component jobs (Azure#28153)

* Remove prefix in pipeline component jobs

Signed-off-by: Brynn Yin <[email protected]>

* Fix bug

Signed-off-by: Brynn Yin <[email protected]>

Signed-off-by: Brynn Yin <[email protected]>

Author: brynn-code

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

@@ -6,6 +6,7 @@
 
 import json
 import logging
+import re
 import typing
 from collections import Counter
 from typing import Dict, Optional, Tuple, Union
@@ -16,7 +17,7 @@
 from azure.ai.ml._schema import PathAwareSchema
 from azure.ai.ml._schema.pipeline.pipeline_component import PipelineComponentSchema
 from azure.ai.ml._utils.utils import is_data_binding_expression, hash_dict
-from azure.ai.ml.constants._common import COMPONENT_TYPE
+from azure.ai.ml.constants._common import COMPONENT_TYPE, ARM_ID_PREFIX, ASSET_ARM_ID_REGEX_FORMAT
 from azure.ai.ml.constants._component import ComponentSource, NodeType
 from azure.ai.ml.constants._job.pipeline import ValidationErrorCode
 from azure.ai.ml.entities._builders import BaseNode, Command
@@ -350,6 +351,10 @@ def _resolve_sub_nodes(cls, rest_jobs):
         if rest_jobs is None:
             return sub_nodes
         for node_name, node in rest_jobs.items():
+            # TODO: Remove this ad-hoc fix after unified arm id format in object
+            component_id = node.get("componentId", "")
+            if isinstance(component_id, str) and re.match(ASSET_ARM_ID_REGEX_FORMAT, component_id):
+                node["componentId"] = component_id[len(ARM_ID_PREFIX):]
             if not LoopNode._is_loop_node_dict(node):
                 # skip resolve LoopNode first since it may reference other nodes
                 # use node factory instead of BaseNode._from_rest_object here as AutoMLJob is not a BaseNode

sdk/ml/azure-ai-ml/tests/component/e2etests/test_component.py

@@ -737,7 +737,6 @@ def test_simple_pipeline_component_create(self, client: MLClient, randstr: Calla
             dict(rest_pipeline_component._to_dict()),
             "name",
             "creation_context",
-            # jobs not returned now
             "jobs",
             "id",
         )
@@ -762,6 +761,10 @@ def test_simple_pipeline_component_create(self, client: MLClient, randstr: Calla
             "type": "pipeline",
         }
         assert component_dict == expected_dict
+        jobs_dict = rest_pipeline_component._to_dict()["jobs"]
+        # Assert full componentId extra azureml prefix has been removed and parsed to versioned arm id correctly.
+        assert "azureml:azureml_anonymous" in jobs_dict["component_a_job"]["component"]
+        assert jobs_dict["component_a_job"]["type"] == "command"
 
     def test_helloworld_nested_pipeline_component(self, client: MLClient, randstr: Callable[[str], str]) -> None:
         component_path = "./tests/test_configs/components/helloworld_nested_pipeline_component.yml"