feat(pipeline): Add inference and lineage step types - #6224
Open
Rishabh0255 wants to merge 1 commit into
Open
Conversation
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
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
Rishabh0255
had a problem deploying
to
manual-approval
September 1, 2026 21:04 — with
GitHub Actions
Error
GitHub Actions
Error
Rishabh0255
had a problem deploying
to
manual-approval
September 1, 2026 21:04 — with
GitHub Actions
Error
GitHub Actions
Error
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
Rishabh0255
force-pushed
the
feat/inference-lineage-steps
branch
from
September 4, 2026 23:20
00feb86 to
00df5b4
Compare
Rishabh0255
requested a deployment
to
manual-approval
September 4, 2026 23:20 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
Rishabh0255
requested a deployment
to
manual-approval
September 4, 2026 23:20 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
Rishabh0255
requested a deployment
to
manual-approval
September 4, 2026 23:20 — 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.