feat(pipeline): Add inference and lineage step types - #6224
Open
Rishabh0255 wants to merge 1 commit 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
requested a deployment
to
manual-approval
August 28, 2026 20:57 — with 
GitHub Actions
Waiting
GitHub Actions
Waiting
Rishabh0255
requested a deployment
to
manual-approval
August 28, 2026 20:57 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
Rishabh0255
requested a deployment
to
manual-approval
August 28, 2026 20:57 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
| 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
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.