feat(pipeline): Add inference and lineage step types - #6224
Open
Rishabh0255 wants to merge 3 commits into
Open
Conversation
Add 4 pipeline step classes: - EndpointConfigStep, EndpointStep (SageMaker inference deployment) - InferenceComponentStep (multi-model endpoint support) - LineageStep (ML governance tracking) Design: each step accepts an 'arguments: Dict[str, Any]' forwarded to the pipeline service. Top-level argument keys are validated client-side against the corresponding public AWS API input shape (botocore service model) at construction and at serialization; fields the service is known to reject fail fast with actionable errors (EndpointConfig: DataCaptureConfig, ExplainerConfig; Endpoint: DeploymentConfig). Values are not validated -- they may be pipeline variables resolved at compile time. Full schema validation remains server-side. If the installed botocore does not know an operation, shape validation is skipped and the service remains the authority. Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config. Includes 23 unit tests and a LineageStep end-to-end integration test. --- X-AI-Prompt: Add the inference and lineage pipeline step types to the Python SDK with client-side argument validation X-AI-Tool: kiro-cli
Rishabh0255
had a problem deploying
to
manual-approval
August 28, 2026 20:57 — with 
GitHub Actions
Error
GitHub Actions
Error
Rishabh0255
had a problem deploying
to
manual-approval
August 28, 2026 20:57 — with
GitHub Actions
Error
GitHub Actions
Error
Rishabh0255
had a problem deploying
to
manual-approval
August 28, 2026 20:57 — with
GitHub Actions
Error
GitHub Actions
Error
| return _SHAPE_CACHE[cache_key] | ||
|
|
||
|
|
||
| def validate_step_arguments( |
Member
There was a problem hiding this comment.
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?
| def __init__( | ||
| self, | ||
| name: str, | ||
| arguments: Dict[str, Any], |
Member
There was a problem hiding this comment.
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
| return get_execution_role() | ||
|
|
||
|
|
||
| def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role): |
Member
There was a problem hiding this comment.
please add integ tests for other steps as well
…steps Address review feedback: EndpointConfigStep, EndpointStep, and InferenceComponentStep now take step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep, instead of a raw arguments dict. - Session.endpoint_from_production_variants, Session.create_endpoint, and Session.create_inference_component now route their service calls through _intercept_create_request. Under a plain Session the behavior is unchanged (the intercept is a pass-through); under a PipelineSession the request is captured and returned as step arguments, and no service call is made. - Each step validates the provenance of its step_args via validate_step_args_input (wrong producer or a raw dict is rejected). - The _argument_validation module is removed: requests are now built by the session methods rather than hand-authored, so client-side key validation is no longer needed. - Adds an integration test chaining EndpointConfigStep -> EndpointStep -> InferenceComponentStep in a single pipeline execution, with full resource cleanup. LineageStep is unchanged pending a design decision on multi-entity step arguments. --- X-AI-Prompt: Rework the inference step types to use step_args captured via PipelineSession per review feedback X-AI-Tool: kiro-cli
Rishabh0255
had a problem deploying
to
manual-approval
August 31, 2026 22:59 — with
GitHub Actions
Error
GitHub Actions
Error
Rishabh0255
had a problem deploying
to
manual-approval
August 31, 2026 22:59 — with
GitHub Actions
Error
GitHub Actions
Error
Rishabh0255
had a problem deploying
to
manual-approval
August 31, 2026 23:00 — with
GitHub Actions
Error
GitHub Actions
Error
LineageStep now creates one lineage entity per step, with step_args obtained by calling Action.create(), Artifact.create(), Context.create(), or Association.create() from sagemaker.core.lineage under a PipelineSession -- completing the step_args convention across all four new step types. - Record._invoke_api in apiutils captures the request as step arguments when the session is a PipelineSession and the call is one of the four lineage create methods; behavior under a plain Session is unchanged. - The step derives the service Arguments block from the captured call: Actions/Artifacts/Contexts wrap the create request; AddAssociation's SourceArn/DestinationArn are translated to entity references, so associations can reference entities created in other steps via step property references (Steps.<name>.ActionArns['<entity>']). - ActionArns/ArtifactArns/ContextArns properties support ['name'] item access for cross-step references. Adds 8 unit tests; updates the LineageStep integration test to the new API. --- X-AI-Prompt: Rework LineageStep to one-entity-per-step with step_args captured via PipelineSession, completing the review feedback X-AI-Tool: kiro-cli
Rishabh0255
requested a deployment
to
manual-approval
September 1, 2026 21:04 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
Rishabh0255
requested a deployment
to
manual-approval
September 1, 2026 21:04 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add 4 pipeline step classes:
Design: each step accepts an 'arguments: Dict[str, Any]' forwarded to the pipeline service. Top-level argument keys are validated client-side against the corresponding public AWS API input shape (botocore service model) at construction and at serialization; fields the service is known to reject fail fast with actionable errors (EndpointConfig: DataCaptureConfig, ExplainerConfig; Endpoint: DeploymentConfig). Values are not validated -- they may be pipeline variables resolved at compile time. Full schema validation remains server-side. If the installed botocore does not know an operation, shape validation is skipped and the service remains the authority.
Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config.
Includes 23 unit tests and a LineageStep end-to-end integration test. ---
X-AI-Prompt: Add the inference and lineage pipeline step types to the Python SDK with client-side argument validation
X-AI-Tool: kiro-cli
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.