feat(pipeline): Add inference and lineage step types by Rishabh0255 · Pull Request #6224 · aws/sagemaker-python-sdk

Rishabh0255 · 2026-08-28T20:46:22Z

Add 4 pipeline step classes:

EndpointConfigStep, EndpointStep (SageMaker inference deployment)
InferenceComponentStep (multi-model endpoint support)
LineageStep (ML governance tracking)

Design: each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep. Session.endpoint_from_production_variants, Session.create_endpoint and Session.create_inference_component route through _intercept_create_request -- a pass-through under a plain Session, capturing the request without a service call under a PipelineSession. LineageStep records one entity per step, with step_args from Action/Artifact/Context/Association.create(); an association can reference an entity created by another step via Steps..ActionArns['']. Each step validates step_args provenance with validate_step_args_input, rejecting a wrong producer or a raw dict.

Chaining note: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. The Endpoint step does not suffix.

Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config.

Includes 29 unit tests and two integration tests covering all four step types end to end against the service (2 passed in 473.64s).

Supersedes #6140 -- same feature, reworked to step_args per review feedback.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

rohangujarathi · 2026-08-28T22:47:51Z

+    return _SHAPE_CACHE[cache_key]
+
+
+def validate_step_arguments(


why do we need all of this additional validation here? Other steps do not have this explicit validation. How are these steps different from other steps?

rohangujarathi · 2026-08-28T22:53:23Z

+    def __init__(
+        self,
+        name: str,
+        arguments: Dict[str, Any],


This is not the right implementation for any of these steps. It will be very difficult to construct these arguments manually. We need to use the existing pysdk constructs and pass them as arguments. Please see how Training/Model steps are implemented and follow that pattern here. You must use step_args from PipelineSession instead of raw arguments: dict

SDK primitives exists for all four steps, and it eliminates the entire _argument_validation.py machinery

rohangujarathi · 2026-08-28T22:55:17Z

+    return get_execution_role()
+
+
+def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role):


please add integ tests for other steps as well

Adds four pipeline step classes -- EndpointConfigStep, EndpointStep, InferenceComponentStep and LineageStep -- along with their StepTypeEnum values and package exports. Each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep: - Session.endpoint_from_production_variants, Session.create_endpoint and Session.create_inference_component route their service calls through _intercept_create_request. Under a plain Session this is a pass-through; under a PipelineSession the request is captured and no service call is made. - LineageStep records one lineage entity per step, with step_args from Action.create(), Artifact.create(), Context.create() or Association.create(). Record._invoke_api captures those four calls under a PipelineSession. An association can reference an entity created by another step via Steps.<name>.ActionArns['<entity>']. - Each step validates the provenance of its step_args with validate_step_args_input, rejecting a wrong producer or a raw dict. Note for chaining: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. Tested with 29 unit tests plus two integration tests that exercise all four step types end to end against the service. --- X-AI-Prompt: Add the new inference and lineage pipeline step types to the public SDK, using step_args captured via PipelineSession per review feedback, and verify them with integration tests X-AI-Tool: kiro-cli

rohangujarathi

Inference steps look right and validation follows the existing pattern. Two blockers:
(1) capture is done three different ways — please route all four steps through PipelineSession._intercept_create_request (the seam Model.create/fit/run already use) instead of adding _is_pipeline_context() in session_helper.py and the _invoke_api hook.
(2) LineageStep bug needs to be fixed

rohangujarathi · 2026-09-10T00:29:39Z

+        """The ``Arguments`` block describing the lineage entity to create."""
+        key = _CALLER_TO_ARGUMENTS_KEY[self.step_args.caller_name]
+        entity = self.step_args.args
+        if key == "Associations":


Cross-step Source.Arn = {"Get": "Steps.X.ActionArns[...]"} won't resolve. Service only accepts a literal/param ARN or an intra-step Name+Type lookup (populated only from entities created in the same step). The service models a batch of entities+associations per step; this does one entity per step and chains across steps, which the runtime doesn't support. Reference siblings by Name+Type within one step, and let one step carry multiple entities.

rohangujarathi · 2026-09-10T00:32:08Z

        """
        return create(request)

+    def _is_pipeline_context(self) -> bool:


create_model doesn't do this. It routes through _intercept_create_request and returns the name; the pipeline branch lives in PipelineSession. Can you keep base Session pipeline-agnostic and drop this helper?

rohangujarathi · 2026-09-10T00:33:13Z

+    return get_execution_role()
+
+
+def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role):


Only covers a single Action. Please test with all the entities

rohangujarathi · 2026-09-10T00:44:06Z

+        )
+        validate_step_args_input(
+            step_args=step_args,
+            expected_caller={"endpoint_from_production_variants"},


can you check if this is the only helper class that creates endpointConfig? also this one seems to be creating both endpoint config as well are endpoint. are we handling this?

Addresses both review blockers. Single capture seam: PipelineSession._intercept_create_request now returns the captured context, so the three endpoint producers just return its result. The _is_pipeline_context() helper is removed from the base Session, which is pipeline-agnostic again, and Record._invoke_api routes the four lineage create calls through the same seam instead of doing its own PipelineSession check. LineageStep: a step now carries a batch of entities rather than one, which is what the pipeline service models. Associations reference siblings by name and type, resolved against the entities the same step creates; a literal ARN references an entity created elsewhere. The previous cross-step ARN reference could not resolve at runtime. Association.create() takes ARNs only and so cannot name a sibling, so associations are declared with LineageAssociation and LineageEntityReference; entities keep the step_args convention. A name and type reference that the step does not create is now rejected at construction rather than failing the execution. The lineage integration test covers an Action, an Artifact and a Context in one step, with associations between them. --- X-AI-Prompt: Rework the PR for both review blockers -- route all four steps through one capture seam, and fix the LineageStep design X-AI-Tool: kiro-cli

Rishabh0255 had a problem deploying to manual-approval August 28, 2026 20:57 — with GitHub Actions Error

rohangujarathi reviewed Aug 28, 2026

View reviewed changes

Rishabh0255 had a problem deploying to manual-approval August 31, 2026 22:59 — with GitHub Actions Error

Rishabh0255 had a problem deploying to manual-approval August 31, 2026 23:00 — with GitHub Actions Error

Rishabh0255 had a problem deploying to manual-approval September 1, 2026 21:04 — with GitHub Actions Error

Rishabh0255 deployed to manual-approval September 1, 2026 21:04 — with GitHub Actions Active

Rishabh0255 force-pushed the feat/inference-lineage-steps branch from 00feb86 to 00df5b4 Compare September 4, 2026 23:20

Rishabh0255 had a problem deploying to manual-approval September 4, 2026 23:20 — with GitHub Actions Error

rohangujarathi reviewed Sep 10, 2026

View reviewed changes

Rishabh0255 requested a deployment to manual-approval September 11, 2026 20:21 — with GitHub Actions Waiting

Rishabh0255 requested a deployment to manual-approval September 11, 2026 20:22 — with GitHub Actions Waiting

Rishabh0255 requested a deployment to manual-approval September 11, 2026 20:23 — with GitHub Actions Waiting

Navigation Menu

Uh oh!

feat(pipeline): Add inference and lineage step types - #6224

feat(pipeline): Add inference and lineage step types#6224
Rishabh0255 wants to merge 2 commits into
aws:masterfrom
Rishabh0255:feat/inference-lineage-steps

Rishabh0255 commented Aug 28, 2026 •

edited

Loading

Uh oh!

rohangujarathi Aug 28, 2026

Uh oh!

rohangujarathi Aug 28, 2026

Uh oh!

rohangujarathi Aug 28, 2026

Uh oh!

rohangujarathi left a comment

Uh oh!

rohangujarathi Sep 10, 2026

Uh oh!

rohangujarathi Sep 10, 2026

Uh oh!

rohangujarathi Sep 10, 2026

Uh oh!

rohangujarathi Sep 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return get_execution_role()


		def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role):

Uh oh!

Conversation

Rishabh0255 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rohangujarathi Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi left a comment

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

rohangujarathi Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Rishabh0255 commented Aug 28, 2026 •

edited

Loading