Skip to content

feat(pipeline): Add Zimmer deployment and lineage step types - #6140

Open
Rishabh0255 wants to merge 1 commit into
aws:masterfrom
Rishabh0255:feat/zimmer-pipeline-steps
Open

feat(pipeline): Add Zimmer deployment and lineage step types#6140
Rishabh0255 wants to merge 1 commit into
aws:masterfrom
Rishabh0255:feat/zimmer-pipeline-steps

Conversation

@Rishabh0255

@Rishabh0255 Rishabh0255 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds 8 new pipeline step classes for Project Zimmer:

  • SageMaker inference deployment: EndpointConfigStep, EndpointStep, InferenceComponentStep
  • Amazon Bedrock: BedrockCustomModelStep, BedrockCustomModelDeploymentStep, BedrockModelImportStep, BedrockProvisionedModelThroughputStep
  • Lineage tracking: LineageStep

Also adds 8 new values to StepTypeEnum and exports the classes from sagemaker.mlops.workflow.

Design

Each step accepts an opaque arguments: Dict[str, Any] that is forwarded to the pipeline service verbatim, mirroring the existing SDK convention used by LambdaStep and CallbackStep. Refer to the corresponding AWS API reference (CreateEndpointConfig, CreateEndpoint, CreateInferenceComponent, Bedrock CreateCustomModel, etc.) for accepted fields.

Retryability follows the Zimmer step contract: only EndpointConfigStep accepts retry_policies. EndpointConfigStep and EndpointStep accept cache_config; nothing else does.

Testing

Unit tests (tests/unit/workflow/test_zimmer_steps.py, 24 tests): construction, to_request() shape, enum values, retry-policy rejection on non-retryable steps, Bedrock properties, LineageStep argument validation.

Integration test (tests/integ/workflow/test_zimmer_lineage_step.py): full end-to-end LineageStep pipeline execution, asserts Succeeded and verifies the created action ARN in step metadata.

Manual validation performed:

  • Byte-for-byte JSON parity with the 3 example pipelines from the Zimmer design doc
  • Gamma CreatePipeline + DescribePipeline round-trip for all 8 step types
  • Cross-step property references (existing SDK step → Zimmer step) resolve correctly
  • End-to-end LineageStep pipeline executed to Succeeded on gamma

Compatibility

No changes to existing step types, no changes to Pipeline, no changes to StepsCompiler beyond the enum additions. _LocalPipelineExecution correctly rejects the new step types with a clear error message consistent with other unsupported-in-local step types.

Add 8 pipeline step classes for Project Zimmer:
- EndpointConfigStep, EndpointStep (SageMaker inference deployment)
- InferenceComponentStep (multi-model endpoint support)
- BedrockCustomModelStep, BedrockCustomModelDeploymentStep,
  BedrockModelImportStep, BedrockProvisionedModelThroughputStep
- LineageStep (ML governance tracking)

Design: each step accepts an opaque 'arguments: Dict[str, Any]' that is
passed through to the Tioga service verbatim. This mirrors Tioga's own
step definitions in IronmanTiogaPipelineDefinitionRepository — the
service validates the Arguments block against the underlying AWS
Create*Request model (SageMaker Coral or Bedrock SDK) minus a small
exclusion set. The SDK does not duplicate that validation.

Excluded fields the service rejects (documented in each step's
docstring): EndpointConfig excludes DataCaptureConfig and ExplainerConfig;
Endpoint excludes DeploymentConfig; all others: no exclusions.

Retryability per Tioga step contract: only EndpointConfigStep is
retryable. Cacheability: EndpointConfigStep and EndpointStep are
structurally cacheable via cache_config.

Bedrock property references use PascalCase field names (matching
Tioga's step-property resolver) even though the underlying Bedrock JSON
API uses camelCase.

sim: https://issues.amazon.com/issues/P424919850
---
X-AI-Prompt: Implement all 8 Zimmer pipeline steps in the Python SDK
matching the service's actual accepted Arguments schema (source: Tioga
IronmanTiogaPipelineDefinitionRepository)
X-AI-Tool: kiro-cli
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant