refactor(pipeline): Take lineage entities per kind on LineageStep

Replaces LineageStep's single step_args argument with actions, artifacts
and contexts, as requested in review. The caller states which kind each
captured argument is instead of the step inferring it from the producer.

This makes a misplaced entity a construction error that names the right
producer -- passing an Artifact.create() result as actions now fails with
"the actions of LineageStep must be obtained from Action.create()" --
rather than something the step silently files under the wrong key.

The emitted Arguments are unchanged: Actions, Artifacts, Contexts and
Associations, with empty kinds omitted. Each kind accepts a single
captured value or a list.

Testing: 39 unit tests, and the lineage integration test run end to end
against the service, which creates an Action, an Artifact and a Context
in one step with two associations referencing them by name and type.
---
X-AI-Prompt: more comments
X-AI-Tool: kiro-cli

Author: Rishabh Devnani

sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py

@@ -17,7 +17,8 @@
 create each one with the corresponding class from
 :mod:`sagemaker.core.lineage` under a
 :class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and pass
-the captured arguments to the step.
+the captured arguments to the matching argument of the step -- ``actions``,
+``artifacts`` or ``contexts``.
 
 Associations are declared rather than captured. ``Association.create()``
 takes source and destination ARNs, but an entity created by the same step
@@ -31,14 +32,16 @@
 
     step = LineageStep(
         name="RecordLineage",
-        step_args=[
+        actions=[
             Action.create(
                 action_name="training-run",
                 source_uri="s3://bucket/run",
                 source_type="S3ETag",
                 action_type="ModelTraining",
                 sagemaker_session=pipeline_session,
             ),
+        ],
+        artifacts=[
             Artifact.create(
                 artifact_name="trained-model",
                 source_uri="s3://bucket/model.tar.gz",
@@ -72,16 +75,34 @@
 ENTITY_TYPE_ARTIFACT = "Artifact"
 ENTITY_TYPE_CONTEXT = "Context"
 
-# Maps the captured lineage create call to the ``Arguments`` key the pipeline
+# One entry per kind of entity a LineageStep can create, in the order the
+# entities appear in the request. Each maps the step argument to the lineage
+# create call that must have produced it, the ``Arguments`` key the pipeline
 # service expects, the entity type used when referencing the entity from an
 # association, and the request field holding the entity name.
-_ENTITY_SPECS = {
-    "create_action": ("Actions", ENTITY_TYPE_ACTION, "ActionName"),
-    "create_artifact": ("Artifacts", ENTITY_TYPE_ARTIFACT, "ArtifactName"),
-    "create_context": ("Contexts", ENTITY_TYPE_CONTEXT, "ContextName"),
+_ENTITY_KINDS = (
+    ("actions", "create_action", "Actions", ENTITY_TYPE_ACTION, "ActionName"),
+    ("artifacts", "create_artifact", "Artifacts", ENTITY_TYPE_ARTIFACT, "ArtifactName"),
+    ("contexts", "create_context", "Contexts", ENTITY_TYPE_CONTEXT, "ContextName"),
+)
+
+# The producer each kind must come from, for the "wrong argument" error message.
+_EXPECTED_PRODUCER = {
+    "actions": "Action.create()",
+    "artifacts": "Artifact.create()",
+    "contexts": "Context.create()",
 }
 
 
+def _as_list(value) -> list:
+    """Normalise a single captured argument or a sequence of them to a list."""
+    if value is None:
+        return []
+    if isinstance(value, (list, tuple)):
+        return list(value)
+    return [value]
+
+
 class LineageEntityReference:
     """Reference to a lineage entity, used as an association endpoint.
 
@@ -207,7 +228,9 @@ class LineageStep(Step):
     def __init__(
         self,
         name: str,
-        step_args: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None,
+        actions: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None,
+        artifacts: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None,
+        contexts: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None,
         associations: Optional[List[LineageAssociation]] = None,
         display_name: Optional[str] = None,
         description: Optional[str] = None,
@@ -217,11 +240,13 @@ def __init__(
 
         Args:
             name (str): The name of the step.
-            step_args (_JobStepArguments or list): The captured arguments for
-                each entity this step creates, obtained from
-                ``Action.create()``, ``Artifact.create()`` or
-                ``Context.create()`` called with a ``PipelineSession``. A
-                single value is accepted for a step that creates one entity.
+            actions (_JobStepArguments or list): Captured arguments from
+                ``Action.create()`` called with a ``PipelineSession``, for each
+                action this step creates.
+            artifacts (_JobStepArguments or list): Captured arguments from
+                ``Artifact.create()``, for each artifact this step creates.
+            contexts (_JobStepArguments or list): Captured arguments from
+                ``Context.create()``, for each context this step creates.
             associations (List[LineageAssociation]): Associations to add after
                 the entities are created.
             display_name (str): Optional display name.
@@ -237,30 +262,30 @@ def __init__(
             depends_on=depends_on,
         )
 
-        if step_args is None:
-            entities = []
-        elif isinstance(step_args, (list, tuple)):
-            entities = list(step_args)
-        else:
-            entities = [step_args]
-
         self.associations = list(associations) if associations else []
 
-        if not entities and not self.associations:
-            raise ValueError(
-                "A LineageStep requires at least one entity in step_args, or one association."
-            )
+        supplied = {"actions": actions, "artifacts": artifacts, "contexts": contexts}
+        self._entities: Dict[str, List[_JobStepArguments]] = {}
+        for param, caller, _, _, _ in _ENTITY_KINDS:
+            entities = _as_list(supplied[param])
+            for entity in entities:
+                validate_step_args_input(
+                    step_args=entity,
+                    expected_caller={caller},
+                    error_message=(
+                        f"The {param} of LineageStep must be obtained from "
+                        f"{_EXPECTED_PRODUCER[param]} called with a PipelineSession. "
+                        "Associations are passed to the associations argument instead, "
+                        "because Association.create() cannot reference an entity created "
+                        "by the same step."
+                    ),
+                )
+            self._entities[param] = entities
 
-        for entity in entities:
-            validate_step_args_input(
-                step_args=entity,
-                expected_caller=set(_ENTITY_SPECS),
-                error_message=(
-                    "The step_args of LineageStep must be obtained from Action.create(), "
-                    "Artifact.create() or Context.create() called with a PipelineSession. "
-                    "Associations are passed to the associations argument instead, because "
-                    "Association.create() cannot reference an entity created by the same step."
-                ),
+        if not any(self._entities.values()) and not self.associations:
+            raise ValueError(
+                "A LineageStep requires at least one entity in actions, artifacts or "
+                "contexts, or one association."
             )
 
         for association in self.associations:
@@ -270,7 +295,6 @@ def __init__(
                     f"{type(association).__name__}."
                 )
 
-        self.step_args = entities
         self._validate_sibling_references()
 
         root = Properties(step_name=name, step=self)
@@ -279,14 +303,29 @@ def __init__(
         root.__dict__["Associations"] = Properties(step_name=name, path="Associations")
         self._properties = root
 
+    @property
+    def actions(self) -> List[_JobStepArguments]:
+        """The captured arguments for the actions this step creates."""
+        return self._entities["actions"]
+
+    @property
+    def artifacts(self) -> List[_JobStepArguments]:
+        """The captured arguments for the artifacts this step creates."""
+        return self._entities["artifacts"]
+
+    @property
+    def contexts(self) -> List[_JobStepArguments]:
+        """The captured arguments for the contexts this step creates."""
+        return self._entities["contexts"]
+
     def _created_entities(self) -> Dict[str, str]:
         """Map the name of each entity created by this step to its type."""
         created = {}
-        for entity in self.step_args:
-            _, entity_type, name_field = _ENTITY_SPECS[entity.caller_name]
-            name = entity.args.get(name_field)
-            if isinstance(name, str):
-                created[name] = entity_type
+        for param, _, _, entity_type, name_field in _ENTITY_KINDS:
+            for entity in self._entities[param]:
+                name = entity.args.get(name_field)
+                if isinstance(name, str):
+                    created[name] = entity_type
         return created
 
     def _validate_sibling_references(self) -> None:
@@ -319,9 +358,10 @@ def _validate_sibling_references(self) -> None:
     def arguments(self) -> RequestType:
         """The ``Arguments`` block: the entities and associations for this step."""
         request: RequestType = {}
-        for entity in self.step_args:
-            key, _, _ = _ENTITY_SPECS[entity.caller_name]
-            request.setdefault(key, []).append(entity.args)
+        for param, _, key, _, _ in _ENTITY_KINDS:
+            entities = self._entities[param]
+            if entities:
+                request[key] = [entity.args for entity in entities]
         if self.associations:
             request["Associations"] = [a.to_request() for a in self.associations]
         return request

sagemaker-mlops/tests/integ/workflow/test_lineage_step.py

@@ -106,7 +106,9 @@ def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, ro
     # so a single step covers create-then-associate end to end.
     step = LineageStep(
         name="RecordLineage",
-        step_args=[action_args, artifact_args, context_args],
+        actions=[action_args],
+        artifacts=[artifact_args],
+        contexts=[context_args],
         associations=[
             LineageAssociation(
                 source=LineageEntityReference(name=action_name, type="Action"),

sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py

@@ -255,7 +255,7 @@ def action_step_args(pipeline_session):
 
 
 def test_lineage_step_action(pipeline_session, action_step_args):
-    step = LineageStep(name="RecA", step_args=action_step_args)
+    step = LineageStep(name="RecA", actions=action_step_args)
     assert step.step_type == StepTypeEnum.LINEAGE
     args = step.arguments
     assert list(args.keys()) == ["Actions"]
@@ -271,7 +271,7 @@ def test_lineage_step_artifact(pipeline_session):
         artifact_type="Model",
         sagemaker_session=pipeline_session,
     )
-    step = LineageStep(name="RecB", step_args=step_args)
+    step = LineageStep(name="RecB", artifacts=step_args)
     assert list(step.arguments.keys()) == ["Artifacts"]
     assert step.arguments["Artifacts"][0]["ArtifactName"] == "art1"
 
@@ -283,7 +283,7 @@ def test_lineage_step_context(pipeline_session):
         context_type="Endpoint",
         sagemaker_session=pipeline_session,
     )
-    step = LineageStep(name="RecC", step_args=step_args)
+    step = LineageStep(name="RecC", contexts=step_args)
     assert list(step.arguments.keys()) == ["Contexts"]
     assert step.arguments["Contexts"][0]["ContextName"] == "ctx1"
 
@@ -313,7 +313,7 @@ def test_lineage_step_batches_multiple_entities(pipeline_session):
         sagemaker_session=pipeline_session,
     )
 
-    step = LineageStep(name="Rec", step_args=actions + [artifact, context])
+    step = LineageStep(name="Rec", actions=actions, artifacts=artifact, contexts=context)
     args = step.arguments
 
     assert [a["ActionName"] for a in args["Actions"]] == ["act0", "act1"]
@@ -342,7 +342,8 @@ def test_lineage_step_association_references_sibling_by_name_and_type(pipeline_s
     )
     step = LineageStep(
         name="Rec",
-        step_args=[action, artifact],
+        actions=action,
+        artifacts=artifact,
         associations=[
             LineageAssociation(
                 source=LineageEntityReference(name="act1", type="Action"),
@@ -364,7 +365,7 @@ def test_lineage_step_association_accepts_literal_arn(pipeline_session, action_s
     existing = "arn:aws:sagemaker:us-west-2:123456789012:artifact/abc"
     step = LineageStep(
         name="Rec",
-        step_args=action_step_args,
+        actions=action_step_args,
         associations=[
             LineageAssociation(
                 source=LineageEntityReference(name="act1", type="Action"),
@@ -385,7 +386,7 @@ def test_lineage_step_rejects_unresolvable_sibling_reference(action_step_args):
     with pytest.raises(ValueError, match="does not create"):
         LineageStep(
             name="Rec",
-            step_args=action_step_args,
+            actions=action_step_args,
             associations=[
                 LineageAssociation(
                     source=LineageEntityReference(name="act1", type="Action"),
@@ -396,15 +397,15 @@ def test_lineage_step_rejects_unresolvable_sibling_reference(action_step_args):
 
 
 def test_lineage_step_rejects_captured_association(pipeline_session, action_step_args):
-    """Association.create() cannot express a sibling, so it is not valid step_args."""
+    """Association.create() cannot express a sibling, so it is not a valid entity."""
     captured = Association.create(
         source_arn="arn:aws:sagemaker:us-west-2:123456789012:action/a",
         destination_arn="arn:aws:sagemaker:us-west-2:123456789012:artifact/b",
         association_type="Produced",
         sagemaker_session=pipeline_session,
     )
     with pytest.raises(ValueError, match="associations argument"):
-        LineageStep(name="Rec", step_args=[action_step_args, captured])
+        LineageStep(name="Rec", actions=[action_step_args, captured])
 
 
 def test_lineage_step_requires_an_entity_or_association():
@@ -431,16 +432,35 @@ def test_lineage_association_rejects_non_reference_endpoint():
 
 def test_lineage_step_rejects_wrong_producer(endpoint_step_args):
     with pytest.raises(ValueError, match="Action.create"):
-        LineageStep(name="Rec", step_args=endpoint_step_args)
+        LineageStep(name="Rec", actions=endpoint_step_args)
+
+
+def test_lineage_step_rejects_entity_in_the_wrong_argument(pipeline_session):
+    """An artifact passed as an action is rejected, naming the right producer.
+
+    The per-kind arguments make the expected producer unambiguous, so a
+    misplaced entity fails at construction instead of building a request the
+    service would reject.
+    """
+    artifact = Artifact.create(
+        artifact_name="art1",
+        source_uri="s3://bucket/model.tar.gz",
+        artifact_type="Model",
+        sagemaker_session=pipeline_session,
+    )
+    with pytest.raises(ValueError, match="actions of LineageStep must be obtained from"):
+        LineageStep(name="Rec", actions=artifact)
+    with pytest.raises(ValueError, match="contexts of LineageStep must be obtained from"):
+        LineageStep(name="Rec", contexts=artifact)
 
 
 def test_lineage_step_rejects_raw_dict():
     with pytest.raises(TypeError):
-        LineageStep(name="Rec", step_args={"Actions": []})
+        LineageStep(name="Rec", actions={"ActionName": "a"})
 
 
 def test_lineage_step_properties(action_step_args):
-    step = LineageStep(name="Rec", step_args=action_step_args)
+    step = LineageStep(name="Rec", actions=action_step_args)
     for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"):
         assert hasattr(step.properties, field)
     assert step.properties.ArtifactArns["x"].expr == {"Get": "Steps.Rec.ArtifactArns['x']"}