From f431bf2bf4773b779471ed129a90ba9e60bed0c5 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Fri, 29 May 2026 22:00:42 +0000 Subject: [PATCH 1/3] feat(pipeline): Add Zimmer deployment and lineage step types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../src/sagemaker/mlops/workflow/__init__.py | 18 + .../sagemaker/mlops/workflow/bedrock_steps.py | 339 ++++++++++++++++ .../sagemaker/mlops/workflow/endpoint_step.py | 193 ++++++++++ .../workflow/inference_component_step.py | 95 +++++ .../sagemaker/mlops/workflow/lineage_step.py | 110 ++++++ .../src/sagemaker/mlops/workflow/steps.py | 9 + .../workflow/test_zimmer_lineage_step.py | 146 +++++++ .../tests/unit/workflow/test_zimmer_steps.py | 362 ++++++++++++++++++ 8 files changed, 1272 insertions(+) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_zimmer_lineage_step.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 129abb1c76..db190990d0 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -14,6 +14,7 @@ functions, conditions, properties) and can import from sagemaker.train and sagemaker.serve for orchestration purposes. """ + from __future__ import absolute_import __version__ = "0.1.0" @@ -42,12 +43,21 @@ # Step implementations from sagemaker.mlops.workflow.automl_step import AutoMLStep +from sagemaker.mlops.workflow.bedrock_steps import ( + BedrockCustomModelStep, + BedrockCustomModelDeploymentStep, + BedrockModelImportStep, + BedrockProvisionedModelThroughputStep, +) from sagemaker.mlops.workflow.callback_step import CallbackStep, CallbackOutput from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep from sagemaker.mlops.workflow.condition_step import ConditionStep from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep from sagemaker.mlops.workflow.fail_step import FailStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput +from sagemaker.mlops.workflow.lineage_step import LineageStep from sagemaker.mlops.workflow.model_step import ModelStep from sagemaker.mlops.workflow.monitor_batch_transform_step import MonitorBatchTransformStep from sagemaker.mlops.workflow.notebook_job_step import NotebookJobStep @@ -92,15 +102,23 @@ "TuningStep", # Step implementations "AutoMLStep", + "BedrockCustomModelStep", + "BedrockCustomModelDeploymentStep", + "BedrockModelImportStep", + "BedrockProvisionedModelThroughputStep", "CallbackStep", "CallbackOutput", "ClarifyCheckStep", "ConditionStep", "EMRStep", "EMRStepConfig", + "EndpointConfigStep", + "EndpointStep", "FailStep", + "InferenceComponentStep", "LambdaStep", "LambdaOutput", + "LineageStep", "ModelStep", "MonitorBatchTransformStep", "NotebookJobStep", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py new file mode 100644 index 0000000000..1842064315 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py @@ -0,0 +1,339 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definitions for Amazon Bedrock deployment steps in Pipelines. + +Design note: mirrors Tioga's Bedrock step definitions in +``IronmanTiogaPipelineDefinitionRepository``. Each step's ``Arguments`` +block is an opaque ``StructureArgument`` validated against the AWS +Bedrock SDK request class (``CreateCustomModelRequest``, +``CreateCustomModelDeploymentRequest``, ``CreateModelImportJobRequest``, +``CreateProvisionedModelThroughputRequest``) with no field exclusions. +Any field the AWS Bedrock API accepts, Tioga accepts. + +Two fields carry server-side validation that this SDK does not duplicate: +``BedrockCustomModelStep.arguments["ModelName"]`` and +``BedrockCustomModelDeploymentStep.arguments["ModelDeploymentName"]`` +must be pipeline parameter references (not hardcoded strings). Tioga +rejects hardcoded values with a clear error at pipeline creation time. + +Property references use PascalCase field names (e.g. +``step.properties.ModelArn``) because Tioga's property-path resolver +uses PascalCase, whereas the Bedrock JSON API uses camelCase member +names. Field lists are hand-populated below from each ``Get*Response``. +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + +# PascalCase property paths for each Bedrock step, sourced from each +# ``Get*Response`` shape. Users reference these via +# ``step.properties.``. +_BEDROCK_CUSTOM_MODEL_FIELDS = [ + "ModelArn", + "ModelName", + "JobArn", + "JobName", + "BaseModelArn", + "CustomizationType", + "ModelKmsKeyArn", + "HyperParameters", + "TrainingDataConfig", + "ValidationDataConfig", + "OutputDataConfig", + "TrainingMetrics", + "ValidationMetrics", + "CreationTime", + "CustomizationConfig", + "ModelStatus", + "FailureMessage", +] + +_BEDROCK_CUSTOM_MODEL_DEPLOYMENT_FIELDS = [ + "ModelDeploymentArn", + "ModelDeploymentName", + "ModelArn", + "CreatedAt", + "Status", + "FailureMessage", + "Description", + "Tags", +] + +_BEDROCK_MODEL_IMPORT_FIELDS = [ + "JobArn", + "JobName", + "ImportedModelName", + "ImportedModelArn", + "RoleArn", + "ModelDataSource", + "Status", + "FailureMessage", + "CreationTime", + "LastModifiedTime", + "EndTime", + "VpcConfig", + "ImportedModelKmsKeyArn", +] + +_BEDROCK_PROVISIONED_MODEL_THROUGHPUT_FIELDS = [ + "ModelUnits", + "DesiredModelUnits", + "ProvisionedModelName", + "ProvisionedModelArn", + "ModelArn", + "DesiredModelArn", + "FoundationModelArn", + "Status", + "CreationTime", + "LastModifiedTime", + "FailureMessage", + "CommitmentDuration", + "CommitmentExpirationTime", +] + + +def _bedrock_properties(step_name: str, step, fields: List[str]) -> Properties: + """Build a bare ``Properties`` root with the given top-level fields.""" + root = Properties(step_name=step_name, step=step) + for field in fields: + root.__dict__[field] = Properties(step_name=step_name, path=field) + return root + + +class BedrockCustomModelStep(Step): + """Creates a custom model in Amazon Bedrock. + + Wraps Bedrock's ``CreateCustomModel`` API. The ``arguments`` dict is + passed through to the service; it accepts any field of + ``CreateCustomModelRequest`` (no exclusions). Typical fields: + ``ModelName`` (required, must be pipeline parameter reference), + ``RoleArn`` (required), ``ModelSourceConfig`` (required), + ``ClientRequestToken``, ``ModelKmsKeyArn``. + + Server-side validation enforces ``ModelName`` as a pipeline parameter + reference — hardcoded strings are rejected at pipeline creation time. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``BedrockCustomModelStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for + ``CreateCustomModel``. ``ModelName`` must be a pipeline + parameter reference. ``ClientRequestToken`` is optional + — the pipeline service auto-generates one if omitted. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.BEDROCK_CUSTOM_MODEL, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for BedrockCustomModelStep.") + self._arguments = arguments + self._properties = _bedrock_properties(name, self, _BEDROCK_CUSTOM_MODEL_FIELDS) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateCustomModel`` call.""" + return self._arguments + + @property + def properties(self): + """PascalCase fields from ``GetCustomModelResponse``.""" + return self._properties + + +class BedrockCustomModelDeploymentStep(Step): + """Deploys a Bedrock custom model for inference. + + Wraps Bedrock's ``CreateCustomModelDeployment`` API. The ``arguments`` + dict is passed through; it accepts any field of + ``CreateCustomModelDeploymentRequest`` (no exclusions). Typical + fields: ``ModelDeploymentName`` (required, must be pipeline parameter + reference), ``ModelArn`` (required), ``Description``, + ``ClientRequestToken``, ``Tags``. + + Server-side validation enforces ``ModelDeploymentName`` as a pipeline + parameter reference. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``BedrockCustomModelDeploymentStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for + ``CreateCustomModelDeployment``. ``ModelDeploymentName`` + must be a pipeline parameter reference. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.BEDROCK_CUSTOM_MODEL_DEPLOYMENT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for BedrockCustomModelDeploymentStep.") + self._arguments = arguments + self._properties = _bedrock_properties(name, self, _BEDROCK_CUSTOM_MODEL_DEPLOYMENT_FIELDS) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateCustomModelDeployment`` call.""" + return self._arguments + + @property + def properties(self): + """PascalCase fields from ``GetCustomModelDeploymentResponse``.""" + return self._properties + + +class BedrockModelImportStep(Step): + """Imports a SageMaker-trained model into Bedrock. + + Wraps Bedrock's ``CreateModelImportJob`` API. The ``arguments`` dict + is passed through; it accepts any field of + ``CreateModelImportJobRequest`` (no exclusions). Typical fields: + ``ImportedModelName``, ``JobName``, ``RoleArn``, ``ModelDataSource``, + ``ClientRequestToken``, ``VpcConfig``, ``ImportedModelKmsKeyId``. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``BedrockModelImportStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for + ``CreateModelImportJob``. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.BEDROCK_MODEL_IMPORT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for BedrockModelImportStep.") + self._arguments = arguments + self._properties = _bedrock_properties(name, self, _BEDROCK_MODEL_IMPORT_FIELDS) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateModelImportJob`` call.""" + return self._arguments + + @property + def properties(self): + """PascalCase fields from ``GetModelImportJobResponse``.""" + return self._properties + + +class BedrockProvisionedModelThroughputStep(Step): + """Creates dedicated provisioned throughput for a Bedrock model. + + Wraps Bedrock's ``CreateProvisionedModelThroughput`` API. The + ``arguments`` dict is passed through; it accepts any field of + ``CreateProvisionedModelThroughputRequest`` (no exclusions). Typical + fields: ``ProvisionedModelName``, ``ModelId``, ``ModelUnits``, + ``CommitmentDuration`` (``OneMonth``/``SixMonths``/``NoCommitment``), + ``ClientRequestToken``, ``Tags``. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``BedrockProvisionedModelThroughputStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for + ``CreateProvisionedModelThroughput``. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.BEDROCK_PROVISIONED_MODEL_THROUGHPUT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for BedrockProvisionedModelThroughputStep.") + self._arguments = arguments + self._properties = _bedrock_properties( + name, self, _BEDROCK_PROVISIONED_MODEL_THROUGHPUT_FIELDS + ) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateProvisionedModelThroughput`` call.""" + return self._arguments + + @property + def properties(self): + """PascalCase fields from ``GetProvisionedModelThroughputResponse``.""" + return self._properties diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py new file mode 100644 index 0000000000..60e53007ce --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -0,0 +1,193 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definitions for SageMaker Endpoint deployment in Pipelines. + +Design note: these classes mirror Tioga's own step definitions +(``EndpointConfigStep``, ``EndpointStep`` in +``IronmanTiogaPipelineDefinitionRepository``). Tioga models each step's +``Arguments`` block as an opaque ``StructureArgument`` validated against +the underlying SageMaker Coral request model (``CreateEndpointConfigInput`` +or ``CreateEndpointInput``) minus a small exclusion set. This SDK +mirrors that: the caller passes an ``arguments`` dict, we forward it +verbatim. See the deserializer classes for the authoritative field lists. + +Excluded fields (Tioga will reject the pipeline if present): + +* ``EndpointConfig``: ``DataCaptureConfig``, ``ExplainerConfig`` +* ``Endpoint``: ``DeploymentConfig`` +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.retry import RetryPolicy +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import ( + CacheConfig, + ConfigurableRetryStep, + Step, + StepTypeEnum, +) + + +class EndpointConfigStep(ConfigurableRetryStep): + """Creates a SageMaker EndpointConfig within a pipeline. + + Wraps the SageMaker ``CreateEndpointConfig`` API. The ``arguments`` + dict is passed through to the service; it accepts any field of + ``CreateEndpointConfigInput`` **except** ``DataCaptureConfig`` and + ``ExplainerConfig``, which are rejected by Tioga. + + Per the Zimmer step contract, ``EndpointConfig`` is structurally + cacheable (``cache_config``) and retryable (``retry_policies``). + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: Optional[CacheConfig] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + ): + """Construct an ``EndpointConfigStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateEndpointConfig`` call. Required fields: + ``EndpointConfigName``, ``ProductionVariants``. Optional + fields include ``KmsKeyId``, ``AsyncInferenceConfig``, + ``ShadowProductionVariants``, ``ExecutionRoleArn``, + ``VpcConfig``, ``EnableNetworkIsolation``, + ``MetricsConfig``. Values may be pipeline variables + (parameter references, step property references) — the + pipeline compiler resolves them at definition time. + Do not include ``DataCaptureConfig`` or ``ExplainerConfig`` + (Tioga rejects them). + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + cache_config (CacheConfig): Optional cache configuration. + retry_policies (List[RetryPolicy]): Optional retry policies. + """ + super().__init__( + name=name, + step_type=StepTypeEnum.ENDPOINT_CONFIG, + display_name=display_name, + description=description, + depends_on=depends_on, + retry_policies=retry_policies, + ) + if arguments is None: + raise ValueError("arguments is required for EndpointConfigStep.") + self._arguments = arguments + self.cache_config = cache_config + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeEndpointConfigOutput" + ) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateEndpointConfig`` call.""" + return self._arguments + + @property + def properties(self): + """A ``Properties`` object shaped like ``DescribeEndpointConfigOutput``.""" + return self._properties + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + request_dict = super().to_request() + if self.cache_config: + request_dict.update(self.cache_config.config) + return request_dict + + +class EndpointStep(Step): + """Creates or updates a SageMaker Endpoint within a pipeline. + + Wraps the SageMaker ``CreateEndpoint``/``UpdateEndpoint`` API — the + pipeline chooses create-vs-update based on endpoint existence. The + ``arguments`` dict is passed through to the service; it accepts any + field of ``CreateEndpointInput`` **except** ``DeploymentConfig``, + which is rejected by Tioga. + + Per the Zimmer step contract, ``Endpoint`` is structurally cacheable + but not retryable at the pipeline level. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: Optional[CacheConfig] = None, + ): + """Construct an ``EndpointStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateEndpoint`` / ``UpdateEndpoint`` call. Required + fields: ``EndpointName``, ``EndpointConfigName``. Optional + fields: ``GraphConfigName``, ``DeletionCondition``. + Values may be pipeline variables. Do not include + ``DeploymentConfig`` (Tioga rejects it). + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + cache_config (CacheConfig): Optional cache configuration. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.ENDPOINT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for EndpointStep.") + self._arguments = arguments + self.cache_config = cache_config + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeEndpointOutput" + ) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the ``CreateEndpoint``/``UpdateEndpoint`` call.""" + return self._arguments + + @property + def properties(self): + """A ``Properties`` object shaped like ``DescribeEndpointOutput``.""" + return self._properties + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + request_dict = super().to_request() + if self.cache_config: + request_dict.update(self.cache_config.config) + return request_dict diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py new file mode 100644 index 0000000000..a23e6dde09 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py @@ -0,0 +1,95 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definition for SageMaker InferenceComponent in Pipelines. + +Design note: mirrors Tioga's ``InferenceComponentStep`` in +``IronmanTiogaPipelineDefinitionRepository``. The ``Arguments`` block is +an opaque ``StructureArgument`` validated against SageMaker's +``CreateInferenceComponentInput`` Coral model with no field exclusions — +any field the AWS API accepts, Tioga accepts. +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class InferenceComponentStep(Step): + """Creates or updates a SageMaker Inference Component within a pipeline. + + Wraps the SageMaker ``CreateInferenceComponent``/``UpdateInferenceComponent`` + API — the pipeline chooses create-vs-update based on component existence. + Inference components enable multi-model endpoint deployments with + independent scaling per model. + + The ``arguments`` dict is passed through to the service; it accepts + any field of ``CreateInferenceComponentInput`` (no exclusions). + + Per the Zimmer step contract, ``InferenceComponent`` is neither + cacheable nor retryable at the pipeline level. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct an ``InferenceComponentStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateInferenceComponent``/``UpdateInferenceComponent`` + call. Typical fields: ``InferenceComponentName``, + ``EndpointName``, ``VariantName``, ``Specification``, + ``Specifications`` (plural, for multi-spec deployments), + ``RuntimeConfig``. Values may be pipeline variables. + Note: ``ComputeResourceRequirements.NumberOfCpuCoresRequired`` + is a float — pass ``2.0`` not ``2``. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.INFERENCE_COMPONENT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for InferenceComponentStep.") + self._arguments = arguments + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeInferenceComponentOutput" + ) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block for the Create/Update InferenceComponent call.""" + return self._arguments + + @property + def properties(self): + """A ``Properties`` object shaped like ``DescribeInferenceComponentOutput``.""" + return self._properties diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py new file mode 100644 index 0000000000..614b28576a --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -0,0 +1,110 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definition for SageMaker Lineage tracking in Pipelines. + +Design note: mirrors Tioga's ``LineageStep`` in +``IronmanTiogaPipelineDefinitionRepository``. The ``Arguments`` block +conforms to Tioga's ``LineageStepArgument`` structure — four optional +lists: + +* ``Actions`` — list of ``CreateActionRequest`` shapes +* ``Artifacts`` — list of ``CreateArtifactRequest`` shapes +* ``Contexts`` — list of ``CreateContextRequest`` shapes +* ``Associations`` — list of ``LineageAssociation`` shapes + (``Source``/``Destination``/``AssociationType``) + +The SDK passes the ``arguments`` dict through verbatim. +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class LineageStep(Step): + """Creates and associates lineage entities in SageMaker's lineage system. + + Wraps SageMaker's ``CreateAction``/``CreateArtifact``/``CreateContext`` + and lineage ``AddAssociation`` APIs. A single step may create + multiple entities of any of the four types (Actions, Artifacts, + Contexts, Associations). Property references use + ``Steps..ActionArns['']``, + ``Steps..ArtifactArns['']``, + ``Steps..ContextArns['']``, and + ``Steps..Associations``. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``LineageStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block. Recognized + top-level keys: ``Actions``, ``Artifacts``, ``Contexts``, + ``Associations`` — each is a list of dicts conforming to + the corresponding SageMaker API shape (or Tioga's + ``LineageAssociation`` for ``Associations``). At least + one of the four keys must be present. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + + Raises: + ValueError: If ``arguments`` is None or contains none of the + recognized keys. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.LINEAGE, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for LineageStep.") + recognized = {"Actions", "Artifacts", "Contexts", "Associations"} + if not recognized & set(arguments.keys()): + raise ValueError( + "LineageStep.arguments must contain at least one of: " + + ", ".join(sorted(recognized)) + ) + self._arguments = arguments + + root = Properties(step_name=name, step=self) + for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): + root.__dict__[field] = Properties(step_name=name, path=field) + self._properties = root + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block describing lineage entities and associations.""" + return self._arguments + + @property + def properties(self): + """Exposes ``ActionArns``, ``ArtifactArns``, ``ContextArns``, ``Associations``.""" + return self._properties diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 76e90a5309..e96420a3a3 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """The `Step` definitions for SageMaker Pipelines Workflows.""" + from __future__ import absolute_import import abc @@ -62,6 +63,14 @@ class StepTypeEnum(Enum): EMR_SERVERLESS = "EMRServerless" FAIL = "Fail" AUTOML = "AutoML" + ENDPOINT_CONFIG = "EndpointConfig" + ENDPOINT = "Endpoint" + INFERENCE_COMPONENT = "InferenceComponent" + BEDROCK_CUSTOM_MODEL = "BedrockCustomModel" + BEDROCK_CUSTOM_MODEL_DEPLOYMENT = "BedrockCustomModelDeployment" + BEDROCK_MODEL_IMPORT = "BedrockModelImport" + BEDROCK_PROVISIONED_MODEL_THROUGHPUT = "BedrockProvisionedModelThroughput" + LINEAGE = "Lineage" class Step(Entity): diff --git a/sagemaker-mlops/tests/integ/workflow/test_zimmer_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_zimmer_lineage_step.py new file mode 100644 index 0000000000..4aa98902e2 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_zimmer_lineage_step.py @@ -0,0 +1,146 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integration test for the Zimmer LineageStep. + +Creates a pipeline containing a single ``LineageStep`` that records a +SageMaker lineage Action, executes it end-to-end against the real +service, and asserts the execution reaches ``Succeeded``. Cleans up the +Action, the pipeline, and the S3 pipeline definition artifact. + +Requires the execution role to have ``sagemaker:CreateAction`` (and +related lineage permissions). ``SageMakerRole`` — the standard fixture +role used across the SDK's integ tests — has broad SageMaker access and +satisfies this requirement. + +This test represents the SDK-side end-to-end validation of the entire +Zimmer step family. See ``bedrock_steps.py``, ``endpoint_step.py``, and +``inference_component_step.py`` for the other step types; those are not +integ-tested here because they either provision paid resources +(Endpoint/InferenceComponent/Bedrock ProvisionedThroughput) or need +external test data (real trained model artifacts for BedrockModelImport). +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.pipeline import Pipeline + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role): + """Full end-to-end run of a LineageStep pipeline against the real service. + + Builds a pipeline with a single ``LineageStep`` that creates one + lineage ``Action``. Verifies the pipeline execution succeeds and the + server-reported step metadata contains the created action ARN. + """ + stamp = uuid.uuid4().hex[:8] + action_name = f"zimmer-integ-{stamp}" + pipeline_name = f"zimmer-integ-lineage-{stamp}" + + step = LineageStep( + name="RecordLineage", + arguments={ + "Actions": [ + { + "ActionName": action_name, + "ActionType": "ModelTraining", + "Status": "Completed", + "Source": { + "SourceUri": f"s3://zimmer-integ-test/{stamp}/model.tar.gz", + "SourceType": "MODEL", + }, + "Description": "Zimmer integ test action", + } + ] + }, + ) + pipeline = Pipeline( + name=pipeline_name, + steps=[step], + sagemaker_session=pipeline_session, + ) + + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + # LineageStep is metadata-only; execution completes quickly. Poll + # up to 5 minutes to give the service plenty of headroom under load. + timeout = 300 + start_time = time.time() + final_status = None + while time.time() - start_time < timeout: + execution_desc = execution.describe() + status = execution_desc["PipelineExecutionStatus"] + if status in ("Succeeded", "Failed", "Stopped"): + final_status = status + break + time.sleep(10) + + if final_status != "Succeeded": + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + failure_details = "\n".join( + f"{s['StepName']}: {s.get('FailureReason', 'no reason')}" + for s in steps + if s.get("StepStatus") == "Failed" + ) + pytest.fail(f"Pipeline execution status={final_status}. Details:\n{failure_details}") + + # Verify the step metadata reports the created action ARN. + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + lineage_step = next(s for s in steps if s["StepName"] == "RecordLineage") + assert lineage_step["StepStatus"] == "Succeeded" + metadata = lineage_step.get("Metadata", {}) + action_arns = metadata.get("Lineage", {}).get("ActionArns", {}) + assert ( + action_name in action_arns + ), f"expected {action_name} in ActionArns, got: {action_arns}" + assert action_arns[action_name].endswith(f":action/{action_name}") + + finally: + # Delete the lineage Action. + try: + sagemaker_session.sagemaker_client.delete_action(ActionName=action_name) + except Exception: + pass + # Delete the pipeline. + try: + sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=pipeline_name) + except Exception: + pass diff --git a/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py b/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py new file mode 100644 index 0000000000..e0023fe033 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py @@ -0,0 +1,362 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for Zimmer pipeline step types. + +These steps use a passthrough ``arguments: Dict[str, Any]`` API, +mirroring Tioga's own step model. The service validates the argument +schema server-side against the underlying AWS Create*Request; this SDK +does not duplicate that validation. +""" + +from __future__ import absolute_import + +import pytest + +from sagemaker.mlops.workflow.bedrock_steps import ( + BedrockCustomModelDeploymentStep, + BedrockCustomModelStep, + BedrockModelImportStep, + BedrockProvisionedModelThroughputStep, +) +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep +from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum + +# ---------- EndpointConfigStep ---------- + + +def test_endpoint_config_step_basic(): + step = EndpointConfigStep( + name="Cfg", + arguments={ + "EndpointConfigName": "MyCfg", + "ProductionVariants": [ + { + "VariantName": "AllTraffic", + "ModelName": "m", + "InstanceType": "ml.m5.large", + "InitialInstanceCount": 1, + } + ], + }, + ) + assert step.step_type == StepTypeEnum.ENDPOINT_CONFIG + assert step.arguments["EndpointConfigName"] == "MyCfg" + + +def test_endpoint_config_step_to_request_includes_cache_and_retry(): + step = EndpointConfigStep( + name="Cfg", + arguments={"EndpointConfigName": "MyCfg", "ProductionVariants": []}, + display_name="Create Config", + description="desc", + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), + ) + req = step.to_request() + assert req["Type"] == "EndpointConfig" + assert req["DisplayName"] == "Create Config" + assert req["Description"] == "desc" + assert req["CacheConfig"] == {"Enabled": True, "ExpireAfter": "P30D"} + + +def test_endpoint_config_step_accepts_full_api_surface(): + """User can pass any CreateEndpointConfigInput field (except the ones + Tioga excludes — that's a server-side rejection, not client-side).""" + step = EndpointConfigStep( + name="Cfg", + arguments={ + "EndpointConfigName": "MyCfg", + "ProductionVariants": [], + "KmsKeyId": "arn:aws:kms:...", + "ExecutionRoleArn": "arn:aws:iam:...", + "AsyncInferenceConfig": {"OutputConfig": {"S3OutputPath": "s3://x/"}}, + "VpcConfig": {"SecurityGroupIds": ["sg-0"], "Subnets": ["subnet-0"]}, + "EnableNetworkIsolation": False, + "ShadowProductionVariants": [], + }, + ) + args = step.arguments + assert args["KmsKeyId"] == "arn:aws:kms:..." + assert args["ExecutionRoleArn"] == "arn:aws:iam:..." + assert "OutputConfig" in args["AsyncInferenceConfig"] + + +def test_endpoint_config_step_requires_arguments(): + with pytest.raises(ValueError): + EndpointConfigStep(name="Cfg", arguments=None) + + +# ---------- EndpointStep ---------- + + +def test_endpoint_step_basic(): + step = EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + ) + assert step.step_type == StepTypeEnum.ENDPOINT + assert step.arguments == {"EndpointName": "ep", "EndpointConfigName": "cfg"} + + +def test_endpoint_step_cache_config(): + step = EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + cache_config=CacheConfig(enable_caching=True), + ) + req = step.to_request() + assert req["Type"] == "Endpoint" + assert req["CacheConfig"] == {"Enabled": True} + + +def test_endpoint_step_rejects_retry_policies_kwarg(): + """EndpointStep is not retryable — constructor must not accept retry_policies.""" + with pytest.raises(TypeError): + EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + retry_policies=[], + ) + + +# ---------- InferenceComponentStep ---------- + + +def test_inference_component_step_basic(): + step = InferenceComponentStep( + name="IC", + arguments={ + "InferenceComponentName": "ic", + "EndpointName": "ep", + "VariantName": "v", + "Specification": { + "ModelName": "m", + "ComputeResourceRequirements": { + "MinMemoryRequiredInMb": 1024, + "NumberOfCpuCoresRequired": 2.0, + }, + }, + "RuntimeConfig": {"CopyCount": 1}, + }, + ) + assert step.step_type == StepTypeEnum.INFERENCE_COMPONENT + assert step.arguments["Specification"]["ModelName"] == "m" + + +def test_inference_component_step_rejects_retry_policies_kwarg(): + with pytest.raises(TypeError): + InferenceComponentStep( + name="IC", + arguments={}, + retry_policies=[], + ) + + +# ---------- Bedrock steps ---------- + + +def test_bedrock_custom_model_step_basic(): + step = BedrockCustomModelStep( + name="RegisterModel", + arguments={ + "ModelName": {"Get": "Parameters.ModelName"}, + "RoleArn": "arn:aws:iam:...", + "ModelSourceConfig": {"S3DataSource": {"S3Uri": "s3://x/y"}}, + }, + ) + assert step.step_type == StepTypeEnum.BEDROCK_CUSTOM_MODEL + assert step.arguments["ModelName"] == {"Get": "Parameters.ModelName"} + + +def test_bedrock_custom_model_deployment_step_basic(): + step = BedrockCustomModelDeploymentStep( + name="Deploy", + arguments={ + "ModelDeploymentName": {"Get": "Parameters.DepName"}, + "ModelArn": "arn:aws:bedrock:...", + }, + ) + assert step.step_type == StepTypeEnum.BEDROCK_CUSTOM_MODEL_DEPLOYMENT + + +def test_bedrock_model_import_step_basic(): + step = BedrockModelImportStep( + name="Import", + arguments={ + "ImportedModelName": "imp", + "JobName": "job", + "RoleArn": "arn:...", + "ModelDataSource": {"S3DataSource": {"S3Uri": "s3://x/y"}}, + }, + ) + assert step.step_type == StepTypeEnum.BEDROCK_MODEL_IMPORT + + +def test_bedrock_provisioned_model_throughput_step_basic(): + step = BedrockProvisionedModelThroughputStep( + name="Prov", + arguments={ + "ProvisionedModelName": "prov", + "ModelId": "m", + "ModelUnits": 1, + "CommitmentDuration": "OneMonth", + }, + ) + assert step.step_type == StepTypeEnum.BEDROCK_PROVISIONED_MODEL_THROUGHPUT + assert step.arguments["CommitmentDuration"] == "OneMonth" + + +def test_bedrock_steps_reject_none_arguments(): + for cls in ( + BedrockCustomModelStep, + BedrockCustomModelDeploymentStep, + BedrockModelImportStep, + BedrockProvisionedModelThroughputStep, + ): + with pytest.raises(ValueError): + cls(name="x", arguments=None) + + +# ---------- Bedrock Properties ---------- + + +def test_bedrock_custom_model_step_properties_typed(): + step = BedrockCustomModelStep( + name="R", + arguments={ + "ModelName": {"Get": "Parameters.ModelName"}, + "RoleArn": "r", + "ModelSourceConfig": {}, + }, + ) + assert step.properties.ModelArn.expr == {"Get": "Steps.R.ModelArn"} + assert step.properties.JobArn.expr == {"Get": "Steps.R.JobArn"} + + +def test_bedrock_model_import_step_properties_typed(): + step = BedrockModelImportStep( + name="I", + arguments={ + "ImportedModelName": "n", + "JobName": "j", + "RoleArn": "r", + "ModelDataSource": {}, + }, + ) + assert step.properties.ImportedModelArn.expr == {"Get": "Steps.I.ImportedModelArn"} + + +def test_bedrock_provisioned_model_throughput_step_properties_typed(): + step = BedrockProvisionedModelThroughputStep( + name="P", + arguments={"ProvisionedModelName": "p", "ModelId": "m", "ModelUnits": 1}, + ) + assert step.properties.ProvisionedModelArn.expr == {"Get": "Steps.P.ProvisionedModelArn"} + + +# ---------- LineageStep ---------- + + +def test_lineage_step_basic(): + step = LineageStep( + name="Rec", + arguments={ + "Actions": [ + { + "ActionName": "a1", + "ActionType": "ModelTraining", + "Status": "Completed", + } + ], + "Artifacts": [ + { + "ArtifactName": "art1", + "ArtifactType": "Model", + "Source": {"SourceUri": "s3://x/y"}, + } + ], + "Associations": [ + { + "Source": {"Name": "a1", "Type": "Action"}, + "Destination": {"Name": "art1", "Type": "Artifact"}, + "AssociationType": "Produced", + } + ], + }, + ) + assert step.step_type == StepTypeEnum.LINEAGE + assert len(step.arguments["Actions"]) == 1 + assert len(step.arguments["Associations"]) == 1 + + +def test_lineage_step_partial_arguments(): + step = LineageStep( + name="Rec", + arguments={"Actions": [{"ActionName": "a", "ActionType": "T", "Status": "Completed"}]}, + ) + assert "Actions" in step.arguments + assert "Artifacts" not in step.arguments + + +def test_lineage_step_requires_at_least_one_recognized_key(): + with pytest.raises(ValueError): + LineageStep(name="Rec", arguments={}) + with pytest.raises(ValueError): + LineageStep(name="Rec", arguments={"Bogus": []}) + + +def test_lineage_step_properties(): + step = LineageStep(name="Rec", arguments={"Actions": []}) + for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): + assert hasattr(step.properties, field) + + +# ---------- Cross-cutting ---------- + + +def test_all_steps_importable_from_init(): + from sagemaker.mlops.workflow import ( # noqa: F401 + BedrockCustomModelDeploymentStep, + BedrockCustomModelStep, + BedrockModelImportStep, + BedrockProvisionedModelThroughputStep, + EndpointConfigStep, + EndpointStep, + InferenceComponentStep, + LineageStep, + ) + + +def test_step_type_enum_values(): + assert StepTypeEnum.ENDPOINT_CONFIG.value == "EndpointConfig" + assert StepTypeEnum.ENDPOINT.value == "Endpoint" + assert StepTypeEnum.INFERENCE_COMPONENT.value == "InferenceComponent" + assert StepTypeEnum.BEDROCK_CUSTOM_MODEL.value == "BedrockCustomModel" + assert StepTypeEnum.BEDROCK_CUSTOM_MODEL_DEPLOYMENT.value == "BedrockCustomModelDeployment" + assert StepTypeEnum.BEDROCK_MODEL_IMPORT.value == "BedrockModelImport" + assert ( + StepTypeEnum.BEDROCK_PROVISIONED_MODEL_THROUGHPUT.value + == "BedrockProvisionedModelThroughput" + ) + assert StepTypeEnum.LINEAGE.value == "Lineage" + + +def test_depends_on_accepts_string_list(): + step = EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + depends_on=["Prev"], + ) + req = step.to_request() + assert req["DependsOn"] == ["Prev"] From fbe6fceba29c369c300db7f4cdad9f3b0d342c79 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Wed, 5 Aug 2026 23:16:43 +0000 Subject: [PATCH 2/3] fix(pipeline): Add client-side argument validation for Zimmer steps Address Fortress security scan findings on arguments()/to_request() passthrough: validate top-level argument keys against the public AWS API input shape (botocore service model) at step construction and at serialization, and fail fast on fields SageMaker Pipelines rejects (DataCaptureConfig, ExplainerConfig on EndpointConfig; DeploymentConfig on Endpoint). Bedrock member names are PascalCase-converted before comparison. LineageStep now rejects unrecognized top-level keys. Values are intentionally not validated: they may be pipeline variables resolved at compile time. If the installed botocore does not know an operation, shape validation is skipped and the service remains the authority. Adds 8 validation unit tests (32 total, all passing). --- X-AI-Prompt: Apply explicit field exclusions/validation on Zimmer step request payloads per Fortress scan feedback X-AI-Tool: kiro-cli --- .../mlops/workflow/_argument_validation.py | 134 ++++++++++++++++++ .../sagemaker/mlops/workflow/bedrock_steps.py | 40 ++++++ .../sagemaker/mlops/workflow/endpoint_step.py | 36 ++++- .../workflow/inference_component_step.py | 13 ++ .../sagemaker/mlops/workflow/lineage_step.py | 11 +- .../tests/unit/workflow/test_zimmer_steps.py | 120 +++++++++++++++- 6 files changed, 348 insertions(+), 6 deletions(-) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py new file mode 100644 index 0000000000..ec19b6c905 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py @@ -0,0 +1,134 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Client-side validation for pipeline step ``arguments`` blocks. + +Validates the **top-level keys** of a step's ``arguments`` dict against +the corresponding public AWS API input shape from botocore, and rejects +fields that SageMaker Pipelines is known not to support. This fails fast +at step construction with a clear error, instead of a server-side parse +failure at ``CreatePipeline`` time. + +Values are intentionally not validated: they may be pipeline variables +(parameter references, step property references, ``Join``/``JsonGet`` +expressions) that only resolve at pipeline compile or execution time. + +If the installed botocore release does not know the target operation +(for example, a very old botocore without newer Bedrock APIs), shape +validation is skipped and the service remains the authority. +""" + +from __future__ import absolute_import + +import logging +from typing import Any, Dict, FrozenSet, Optional, Sequence, Tuple + +import botocore.session +from botocore.exceptions import UnknownServiceError +from botocore.model import OperationNotFoundError + +logger = logging.getLogger(__name__) + +# Cache of (service, operation, pascal_case) -> allowed top-level keys. +# ``None`` means botocore does not know the operation; skip shape checks. +_SHAPE_CACHE: Dict[Tuple[str, str, bool], Optional[FrozenSet[str]]] = {} + + +def _allowed_top_level_keys( + service_name: str, operation_name: str, pascal_case: bool +) -> Optional[FrozenSet[str]]: + """Return the allowed top-level keys for an operation input shape. + + Args: + service_name (str): botocore service name (e.g. ``sagemaker``). + operation_name (str): operation name (e.g. ``CreateEndpointConfig``). + pascal_case (bool): If True, convert member names to PascalCase + (used for Bedrock, whose JSON API members are camelCase but + whose pipeline ``Arguments`` fields are PascalCase). + + Returns: + The allowed key set, or ``None`` if the installed botocore does + not know the operation (validation should then be skipped). + """ + cache_key = (service_name, operation_name, pascal_case) + if cache_key not in _SHAPE_CACHE: + try: + session = botocore.session.get_session() + service_model = session.get_service_model(service_name) + operation_model = service_model.operation_model(operation_name) + members = operation_model.input_shape.members.keys() + if pascal_case: + members = [m[0].upper() + m[1:] for m in members] + _SHAPE_CACHE[cache_key] = frozenset(members) + except (UnknownServiceError, OperationNotFoundError): + logger.warning( + "Installed botocore does not know %s.%s; skipping " + "client-side argument shape validation for this step.", + service_name, + operation_name, + ) + _SHAPE_CACHE[cache_key] = None + return _SHAPE_CACHE[cache_key] + + +def validate_step_arguments( + step_class_name: str, + arguments: Dict[str, Any], + service_name: str, + operation_name: str, + unsupported_fields: Sequence[str] = (), + pascal_case: bool = False, +) -> None: + """Validate the top-level keys of a step ``arguments`` dict. + + Args: + step_class_name (str): Step class name, used in error messages. + arguments (Dict[str, Any]): The user-provided ``arguments`` dict. + service_name (str): botocore service name of the wrapped API. + operation_name (str): Operation whose input shape defines the + allowed top-level fields. + unsupported_fields (Sequence[str]): Fields that exist in the + public API shape but are rejected by SageMaker Pipelines. + pascal_case (bool): Convert botocore member names to PascalCase + before comparison (Bedrock APIs). + + Raises: + ValueError: If ``arguments`` is not a non-empty dict with string + keys, contains an unsupported field, or contains a key that + is not part of the operation's input shape. + """ + if arguments is None: + raise ValueError(f"arguments is required for {step_class_name}.") + if not isinstance(arguments, dict) or not arguments: + raise ValueError(f"{step_class_name}: arguments must be a non-empty dict.") + non_string_keys = [key for key in arguments if not isinstance(key, str)] + if non_string_keys: + raise ValueError( + f"{step_class_name}: argument keys must be strings; got {non_string_keys!r}." + ) + rejected = sorted(field for field in unsupported_fields if field in arguments) + if rejected: + raise ValueError( + f"{step_class_name}: field(s) {rejected} are not supported by " + "SageMaker Pipelines and would be rejected at pipeline creation " + "time. Remove them from arguments." + ) + allowed = _allowed_top_level_keys(service_name, operation_name, pascal_case) + if allowed is None: + return + unknown = sorted(set(arguments) - allowed) + if unknown: + raise ValueError( + f"{step_class_name}: unknown argument field(s) {unknown}. " + f"Allowed top-level fields (from {service_name}.{operation_name}): " + f"{sorted(allowed)}." + ) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py index 1842064315..b68f629685 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/bedrock_steps.py @@ -39,9 +39,29 @@ from sagemaker.core.helper.pipeline_variable import RequestType from sagemaker.core.workflow.properties import Properties +from sagemaker.mlops.workflow._argument_validation import validate_step_arguments from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + +def _validate_bedrock_arguments( + step_class_name: str, arguments: Dict[str, Any], operation_name: str +) -> None: + """Validate a Bedrock step's arguments against the botocore input shape. + + Bedrock's JSON API members are camelCase, but pipeline ``Arguments`` + fields are PascalCase (matching Tioga's property-path resolver), so + shape member names are PascalCase-converted before comparison. + """ + validate_step_arguments( + step_class_name, + arguments, + service_name="bedrock", + operation_name=operation_name, + pascal_case=True, + ) + + # PascalCase property paths for each Bedrock step, sourced from each # ``Get*Response`` shape. Users reference these via # ``step.properties.``. @@ -161,12 +181,14 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for BedrockCustomModelStep.") + _validate_bedrock_arguments("BedrockCustomModelStep", arguments, "CreateCustomModel") self._arguments = arguments self._properties = _bedrock_properties(name, self, _BEDROCK_CUSTOM_MODEL_FIELDS) @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateCustomModel`` call.""" + _validate_bedrock_arguments("BedrockCustomModelStep", self._arguments, "CreateCustomModel") return self._arguments @property @@ -218,12 +240,18 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for BedrockCustomModelDeploymentStep.") + _validate_bedrock_arguments( + "BedrockCustomModelDeploymentStep", arguments, "CreateCustomModelDeployment" + ) self._arguments = arguments self._properties = _bedrock_properties(name, self, _BEDROCK_CUSTOM_MODEL_DEPLOYMENT_FIELDS) @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateCustomModelDeployment`` call.""" + _validate_bedrock_arguments( + "BedrockCustomModelDeploymentStep", self._arguments, "CreateCustomModelDeployment" + ) return self._arguments @property @@ -270,12 +298,16 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for BedrockModelImportStep.") + _validate_bedrock_arguments("BedrockModelImportStep", arguments, "CreateModelImportJob") self._arguments = arguments self._properties = _bedrock_properties(name, self, _BEDROCK_MODEL_IMPORT_FIELDS) @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateModelImportJob`` call.""" + _validate_bedrock_arguments( + "BedrockModelImportStep", self._arguments, "CreateModelImportJob" + ) return self._arguments @property @@ -323,6 +355,9 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for BedrockProvisionedModelThroughputStep.") + _validate_bedrock_arguments( + "BedrockProvisionedModelThroughputStep", arguments, "CreateProvisionedModelThroughput" + ) self._arguments = arguments self._properties = _bedrock_properties( name, self, _BEDROCK_PROVISIONED_MODEL_THROUGHPUT_FIELDS @@ -331,6 +366,11 @@ def __init__( @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateProvisionedModelThroughput`` call.""" + _validate_bedrock_arguments( + "BedrockProvisionedModelThroughputStep", + self._arguments, + "CreateProvisionedModelThroughput", + ) return self._arguments @property diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py index 60e53007ce..7f52a5acdb 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -18,8 +18,11 @@ ``Arguments`` block as an opaque ``StructureArgument`` validated against the underlying SageMaker Coral request model (``CreateEndpointConfigInput`` or ``CreateEndpointInput``) minus a small exclusion set. This SDK -mirrors that: the caller passes an ``arguments`` dict, we forward it -verbatim. See the deserializer classes for the authoritative field lists. +validates the **top-level keys** of the ``arguments`` dict against the +public ``CreateEndpointConfig``/``CreateEndpoint`` API input shape at +construction time (values are not validated -- they may be pipeline +variables) and forwards the dict to the service, which remains the +authority on full schema validation. Excluded fields (Tioga will reject the pipeline if present): @@ -34,6 +37,7 @@ from sagemaker.core.helper.pipeline_variable import RequestType from sagemaker.core.workflow.properties import Properties +from sagemaker.mlops.workflow._argument_validation import validate_step_arguments from sagemaker.mlops.workflow.retry import RetryPolicy from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import ( @@ -98,6 +102,13 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for EndpointConfigStep.") + validate_step_arguments( + "EndpointConfigStep", + arguments, + service_name="sagemaker", + operation_name="CreateEndpointConfig", + unsupported_fields=("DataCaptureConfig", "ExplainerConfig"), + ) self._arguments = arguments self.cache_config = cache_config self._properties = Properties( @@ -107,6 +118,13 @@ def __init__( @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateEndpointConfig`` call.""" + validate_step_arguments( + "EndpointConfigStep", + self._arguments, + service_name="sagemaker", + operation_name="CreateEndpointConfig", + unsupported_fields=("DataCaptureConfig", "ExplainerConfig"), + ) return self._arguments @property @@ -169,6 +187,13 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for EndpointStep.") + validate_step_arguments( + "EndpointStep", + arguments, + service_name="sagemaker", + operation_name="CreateEndpoint", + unsupported_fields=("DeploymentConfig",), + ) self._arguments = arguments self.cache_config = cache_config self._properties = Properties( @@ -178,6 +203,13 @@ def __init__( @property def arguments(self) -> RequestType: """The ``Arguments`` block for the ``CreateEndpoint``/``UpdateEndpoint`` call.""" + validate_step_arguments( + "EndpointStep", + self._arguments, + service_name="sagemaker", + operation_name="CreateEndpoint", + unsupported_fields=("DeploymentConfig",), + ) return self._arguments @property diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py index a23e6dde09..c2da087ae6 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py @@ -26,6 +26,7 @@ from sagemaker.core.helper.pipeline_variable import RequestType from sagemaker.core.workflow.properties import Properties +from sagemaker.mlops.workflow._argument_validation import validate_step_arguments from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum @@ -79,6 +80,12 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for InferenceComponentStep.") + validate_step_arguments( + "InferenceComponentStep", + arguments, + service_name="sagemaker", + operation_name="CreateInferenceComponent", + ) self._arguments = arguments self._properties = Properties( step_name=name, step=self, shape_name="DescribeInferenceComponentOutput" @@ -87,6 +94,12 @@ def __init__( @property def arguments(self) -> RequestType: """The ``Arguments`` block for the Create/Update InferenceComponent call.""" + validate_step_arguments( + "InferenceComponentStep", + self._arguments, + service_name="sagemaker", + operation_name="CreateInferenceComponent", + ) return self._arguments @property diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py index 614b28576a..073fed8479 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -23,7 +23,8 @@ * ``Associations`` — list of ``LineageAssociation`` shapes (``Source``/``Destination``/``AssociationType``) -The SDK passes the ``arguments`` dict through verbatim. +The SDK validates that the ``arguments`` dict contains only these four +top-level keys (at least one required) and forwards it to the service. """ from __future__ import absolute_import @@ -86,12 +87,20 @@ def __init__( ) if arguments is None: raise ValueError("arguments is required for LineageStep.") + if not isinstance(arguments, dict) or not arguments: + raise ValueError("LineageStep: arguments must be a non-empty dict.") recognized = {"Actions", "Artifacts", "Contexts", "Associations"} if not recognized & set(arguments.keys()): raise ValueError( "LineageStep.arguments must contain at least one of: " + ", ".join(sorted(recognized)) ) + unknown = sorted(set(arguments) - recognized) + if unknown: + raise ValueError( + f"LineageStep: unknown argument field(s) {unknown}. " + "Allowed top-level fields: " + ", ".join(sorted(recognized)) + "." + ) self._arguments = arguments root = Properties(step_name=name, step=self) diff --git a/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py b/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py index e0023fe033..b7a1b2bee6 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_zimmer_steps.py @@ -13,9 +13,12 @@ """Unit tests for Zimmer pipeline step types. These steps use a passthrough ``arguments: Dict[str, Any]`` API, -mirroring Tioga's own step model. The service validates the argument -schema server-side against the underlying AWS Create*Request; this SDK -does not duplicate that validation. +mirroring Tioga's own step model. Top-level argument keys are validated +client-side against the corresponding public AWS API input shape +(botocore service model), and fields known to be rejected by SageMaker +Pipelines fail fast at construction. Values are not validated -- they +may be pipeline variables resolved at compile time. Full schema +validation remains server-side. """ from __future__ import absolute_import @@ -360,3 +363,114 @@ def test_depends_on_accepts_string_list(): ) req = step.to_request() assert req["DependsOn"] == ["Prev"] + + +# ---------- Client-side argument validation ---------- + + +def test_endpoint_config_step_rejects_unsupported_fields(): + """DataCaptureConfig and ExplainerConfig exist in the public API but + are rejected by SageMaker Pipelines -- fail fast with a clear error.""" + for field in ("DataCaptureConfig", "ExplainerConfig"): + with pytest.raises(ValueError, match=field): + EndpointConfigStep( + name="Cfg", + arguments={ + "EndpointConfigName": "cfg", + "ProductionVariants": [], + field: {}, + }, + ) + + +def test_endpoint_step_rejects_unsupported_deployment_config(): + with pytest.raises(ValueError, match="DeploymentConfig"): + EndpointStep( + name="Deploy", + arguments={ + "EndpointName": "ep", + "EndpointConfigName": "cfg", + "DeploymentConfig": {}, + }, + ) + + +def test_unknown_argument_key_rejected(): + """Keys outside the operation's input shape fail fast at construction.""" + with pytest.raises(ValueError, match="Bogus"): + EndpointConfigStep( + name="Cfg", + arguments={"EndpointConfigName": "cfg", "Bogus": 1}, + ) + with pytest.raises(ValueError, match="Bogus"): + InferenceComponentStep( + name="IC", + arguments={"InferenceComponentName": "ic", "Bogus": 1}, + ) + + +def test_bedrock_steps_validate_pascal_case_keys(): + """Valid PascalCase keys (converted from Bedrock's camelCase API + members) are accepted; unknown keys are rejected.""" + step = BedrockCustomModelStep( + name="CM", + arguments={ + "ModelName": {"Get": "Parameters.ModelName"}, + "RoleArn": "arn:aws:iam:...", + "ModelSourceConfig": {}, + }, + ) + assert "ModelName" in step.arguments + with pytest.raises(ValueError, match="Bogus"): + BedrockCustomModelStep( + name="CM", + arguments={"ModelName": {"Get": "Parameters.ModelName"}, "Bogus": 1}, + ) + with pytest.raises(ValueError, match="Bogus"): + BedrockProvisionedModelThroughputStep( + name="PT", + arguments={"ProvisionedModelName": "pm", "Bogus": 1}, + ) + + +def test_empty_arguments_rejected(): + for cls, valid_key in ( + (EndpointConfigStep, "EndpointConfigName"), + (EndpointStep, "EndpointName"), + (InferenceComponentStep, "InferenceComponentName"), + (BedrockModelImportStep, "JobName"), + ): + with pytest.raises(ValueError): + cls(name="x", arguments={}) + # sanity: a single valid key constructs fine + assert cls(name="x", arguments={valid_key: "v"}).arguments == {valid_key: "v"} + + +def test_pipeline_variable_values_pass_validation(): + """Only top-level keys are validated -- values may be pipeline + variables (Get expressions) at any position.""" + step = EndpointStep( + name="Deploy", + arguments={ + "EndpointName": {"Get": "Parameters.EndpointName"}, + "EndpointConfigName": {"Get": "Steps.Cfg.EndpointConfigName"}, + }, + ) + assert step.arguments["EndpointName"] == {"Get": "Parameters.EndpointName"} + + +def test_post_construction_mutation_caught_at_serialization(): + """Injecting an unsupported field after construction is caught when + the arguments property is read (i.e., at pipeline serialization).""" + step = EndpointConfigStep( + name="Cfg", + arguments={"EndpointConfigName": "cfg", "ProductionVariants": []}, + ) + step._arguments["DataCaptureConfig"] = {} + with pytest.raises(ValueError, match="DataCaptureConfig"): + _ = step.arguments + + +def test_lineage_step_rejects_unknown_keys_alongside_recognized(): + with pytest.raises(ValueError, match="Bogus"): + LineageStep(name="Rec", arguments={"Actions": [], "Bogus": []}) From 395582be7c366b6eabcd164679a2e1e4640e032a Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Fri, 7 Aug 2026 20:34:28 +0000 Subject: [PATCH 3/3] chore: empty commit to re-trigger CI Re-trigger the PR check suite (integ tests run issue). --- X-AI-Prompt: Make dummy commits on v2 and v3 branches to fix an integ tests run issue X-AI-Tool: kiro-cli