From 81f70ac374ba7f6c152f7075114d11d306d84259 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Tue, 4 Aug 2026 19:16:05 +0000 Subject: [PATCH 1/2] feat: Add Zimmer deployment and lineage step types (v2 backport) Backports the 8 Zimmer pipeline step classes from v3 (PR #6140) to v2: - EndpointConfigStep, EndpointStep (SageMaker inference deployment) - InferenceComponentStep (multi-model endpoint support) - BedrockCustomModelStep, BedrockCustomModelDeploymentStep, BedrockModelImportStep, BedrockProvisionedModelThroughputStep - LineageStep (ML governance tracking) Also adds 8 new values to StepTypeEnum. 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 Create* API references for accepted fields. Retryability per Zimmer step contract: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config. sim: https://issues.amazon.com/issues/P424919850 --- X-AI-Prompt: Backport the Zimmer pipeline step types from the v3 PR (aws/sagemaker-python-sdk#6140) to the v2 SDK on master-v2 X-AI-Tool: kiro-cli --- src/sagemaker/workflow/bedrock_steps.py | 295 ++++++++++++++++ src/sagemaker/workflow/endpoint_step.py | 162 +++++++++ .../workflow/inference_component_step.py | 83 +++++ src/sagemaker/workflow/lineage_step.py | 94 +++++ src/sagemaker/workflow/steps.py | 8 + .../workflow/test_zimmer_lineage_step.py | 128 +++++++ .../sagemaker/workflow/test_zimmer_steps.py | 332 ++++++++++++++++++ 7 files changed, 1102 insertions(+) create mode 100644 src/sagemaker/workflow/bedrock_steps.py create mode 100644 src/sagemaker/workflow/endpoint_step.py create mode 100644 src/sagemaker/workflow/inference_component_step.py create mode 100644 src/sagemaker/workflow/lineage_step.py create mode 100644 tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py create mode 100644 tests/unit/sagemaker/workflow/test_zimmer_steps.py diff --git a/src/sagemaker/workflow/bedrock_steps.py b/src/sagemaker/workflow/bedrock_steps.py new file mode 100644 index 0000000000..b81b9df073 --- /dev/null +++ b/src/sagemaker/workflow/bedrock_steps.py @@ -0,0 +1,295 @@ +# 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.""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.workflow.entities import RequestType +from sagemaker.workflow.properties import Properties +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.workflow.steps import Step, StepTypeEnum + +# 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 + forwarded to the service. Typical fields: ``ModelName``, ``RoleArn``, + ``ModelSourceConfig``, ``ClientRequestToken``, ``ModelKmsKeyArn``. + """ + + 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``. ``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): + """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 forwarded to the service. + """ + + 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``. + 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): + """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 forwarded to the service. + """ + + 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): + """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 forwarded to the service. + """ + + 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): + """Fields from ``GetProvisionedModelThroughputResponse``.""" + return self._properties diff --git a/src/sagemaker/workflow/endpoint_step.py b/src/sagemaker/workflow/endpoint_step.py new file mode 100644 index 0000000000..49a44e6999 --- /dev/null +++ b/src/sagemaker/workflow/endpoint_step.py @@ -0,0 +1,162 @@ +# 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.""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.workflow.entities import RequestType +from sagemaker.workflow.properties import Properties +from sagemaker.workflow.retry import RetryPolicy +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.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 forwarded to the service — refer to the + `CreateEndpointConfig API reference + `_ + for accepted fields. Values may be pipeline variables (parameter + references, step property references). + """ + + 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. Values may be pipeline + variables. + 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. Refer + to the `CreateEndpoint API reference + `_ + for accepted fields. + """ + + 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. Values may + be pipeline variables. + 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/src/sagemaker/workflow/inference_component_step.py b/src/sagemaker/workflow/inference_component_step.py new file mode 100644 index 0000000000..1113f91d5a --- /dev/null +++ b/src/sagemaker/workflow/inference_component_step.py @@ -0,0 +1,83 @@ +# 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.""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.workflow.entities import RequestType +from sagemaker.workflow.properties import Properties +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.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. Refer to the + `CreateInferenceComponent API reference + `_ + for accepted fields. + """ + + 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. Values may be pipeline variables. Note that + ``ComputeResourceRequirements.NumberOfCpuCoresRequired`` + is a JSON float; pass ``2.0`` rather than ``2`` from + Python. + 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/src/sagemaker/workflow/lineage_step.py b/src/sagemaker/workflow/lineage_step.py new file mode 100644 index 0000000000..45468a1360 --- /dev/null +++ b/src/sagemaker/workflow/lineage_step.py @@ -0,0 +1,94 @@ +# 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.""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.workflow.entities import RequestType +from sagemaker.workflow.properties import Properties +from sagemaker.workflow.step_collections import StepCollection +from sagemaker.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. 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/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py index dbc37371db..3de5989d43 100644 --- a/src/sagemaker/workflow/steps.py +++ b/src/sagemaker/workflow/steps.py @@ -71,6 +71,14 @@ class StepTypeEnum(Enum): EMR = "EMR" 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/tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py b/tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py new file mode 100644 index 0000000000..e8dbc004aa --- /dev/null +++ b/tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py @@ -0,0 +1,128 @@ +# 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 (v2).""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker import get_execution_role, utils +from sagemaker.workflow.lineage_step import LineageStep +from sagemaker.workflow.pipeline import Pipeline + + +@pytest.fixture +def role(sagemaker_session): + return get_execution_role(sagemaker_session) + + +@pytest.fixture +def pipeline_name(): + return utils.unique_name_from_base("my-pipeline-lineage") + + +def test_lineage_step_execute_end_to_end(sagemaker_session, role, pipeline_name): + """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 + that the server-reported step metadata contains the created action + ARN. Cleans up the Action, the pipeline, and any lingering S3 + definition artifact. + + Requires the execution role to have ``sagemaker:CreateAction``. The + ``SageMakerRole`` fixture role has broad SageMaker access and + satisfies this requirement. + """ + stamp = uuid.uuid4().hex[:8] + action_name = f"zimmer-integ-{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 v2 integ test action", + } + ] + }, + ) + pipeline = Pipeline( + name=pipeline_name, + steps=[step], + sagemaker_session=sagemaker_session, + ) + + sm_client = sagemaker_session.sagemaker_client + try: + pipeline.create(role_arn=role) + execution = pipeline.start() + + # LineageStep is metadata-only and completes quickly. Poll up to + # 5 minutes for headroom under load. + timeout = 300 + start_time = time.time() + final_status = None + while time.time() - start_time < timeout: + description = execution.describe() + status = description["PipelineExecutionStatus"] + if status in ("Succeeded", "Failed", "Stopped"): + final_status = status + break + time.sleep(10) + + if final_status != "Succeeded": + steps = sm_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 = sm_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: + sm_client.delete_action(ActionName=action_name) + except Exception: + pass + # Delete the pipeline. + try: + sm_client.delete_pipeline(PipelineName=pipeline_name) + except Exception: + pass diff --git a/tests/unit/sagemaker/workflow/test_zimmer_steps.py b/tests/unit/sagemaker/workflow/test_zimmer_steps.py new file mode 100644 index 0000000000..e6e03fe448 --- /dev/null +++ b/tests/unit/sagemaker/workflow/test_zimmer_steps.py @@ -0,0 +1,332 @@ +# 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 (v2). + +Passthrough ``arguments: Dict[str, Any]`` API. The service validates +the argument schema server-side; this SDK does not duplicate that. +""" + +from __future__ import absolute_import + +import pytest + +from sagemaker.workflow.bedrock_steps import ( + BedrockCustomModelDeploymentStep, + BedrockCustomModelStep, + BedrockModelImportStep, + BedrockProvisionedModelThroughputStep, +) +from sagemaker.workflow.endpoint_step import EndpointConfigStep, EndpointStep +from sagemaker.workflow.inference_component_step import InferenceComponentStep +from sagemaker.workflow.lineage_step import LineageStep +from sagemaker.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(): + 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:..." + + +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(): + 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 + + +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_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 34b16d3dc636b54803edb457e04d6bcb8ce8da94 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Wed, 5 Aug 2026 23:21:08 +0000 Subject: [PATCH 2/2] fix(pipeline): Add client-side argument validation for Zimmer steps Mirrors the v3 change (commit fbe6fceb on feat/zimmer-pipeline-steps): 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). --- X-AI-Prompt: Mirror the Fortress-driven argument validation change to the v2 backport branch X-AI-Tool: kiro-cli --- .../workflow/_argument_validation.py | 134 ++++++++++++++++++ src/sagemaker/workflow/bedrock_steps.py | 41 ++++++ src/sagemaker/workflow/endpoint_step.py | 29 ++++ .../workflow/inference_component_step.py | 13 ++ src/sagemaker/workflow/lineage_step.py | 8 ++ .../sagemaker/workflow/test_zimmer_steps.py | 118 ++++++++++++++- 6 files changed, 341 insertions(+), 2 deletions(-) create mode 100644 src/sagemaker/workflow/_argument_validation.py diff --git a/src/sagemaker/workflow/_argument_validation.py b/src/sagemaker/workflow/_argument_validation.py new file mode 100644 index 0000000000..ec19b6c905 --- /dev/null +++ b/src/sagemaker/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/src/sagemaker/workflow/bedrock_steps.py b/src/sagemaker/workflow/bedrock_steps.py index b81b9df073..4431438333 100644 --- a/src/sagemaker/workflow/bedrock_steps.py +++ b/src/sagemaker/workflow/bedrock_steps.py @@ -16,11 +16,32 @@ from typing import Any, Dict, List, Optional, Union +from sagemaker.workflow._argument_validation import validate_step_arguments from sagemaker.workflow.entities import RequestType from sagemaker.workflow.properties import Properties from sagemaker.workflow.step_collections import StepCollection from sagemaker.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 the pipeline service'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, + ) + + # Property paths for each Bedrock step, sourced from each ``Get*Response`` shape. # Users reference these via ``step.properties.``. _BEDROCK_CUSTOM_MODEL_FIELDS = [ @@ -132,12 +153,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 @@ -181,12 +204,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 @@ -230,12 +259,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 @@ -279,6 +312,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 @@ -287,6 +323,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/src/sagemaker/workflow/endpoint_step.py b/src/sagemaker/workflow/endpoint_step.py index 49a44e6999..7af7694e82 100644 --- a/src/sagemaker/workflow/endpoint_step.py +++ b/src/sagemaker/workflow/endpoint_step.py @@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Union +from sagemaker.workflow._argument_validation import validate_step_arguments from sagemaker.workflow.entities import RequestType from sagemaker.workflow.properties import Properties from sagemaker.workflow.retry import RetryPolicy @@ -73,6 +74,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( @@ -82,6 +90,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 @@ -138,6 +153,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( @@ -147,6 +169,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/src/sagemaker/workflow/inference_component_step.py b/src/sagemaker/workflow/inference_component_step.py index 1113f91d5a..2c61eab258 100644 --- a/src/sagemaker/workflow/inference_component_step.py +++ b/src/sagemaker/workflow/inference_component_step.py @@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Union +from sagemaker.workflow._argument_validation import validate_step_arguments from sagemaker.workflow.entities import RequestType from sagemaker.workflow.properties import Properties from sagemaker.workflow.step_collections import StepCollection @@ -67,6 +68,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" @@ -75,6 +82,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/src/sagemaker/workflow/lineage_step.py b/src/sagemaker/workflow/lineage_step.py index 45468a1360..96cb5fec95 100644 --- a/src/sagemaker/workflow/lineage_step.py +++ b/src/sagemaker/workflow/lineage_step.py @@ -70,12 +70,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/tests/unit/sagemaker/workflow/test_zimmer_steps.py b/tests/unit/sagemaker/workflow/test_zimmer_steps.py index e6e03fe448..8a5b0ab0c4 100644 --- a/tests/unit/sagemaker/workflow/test_zimmer_steps.py +++ b/tests/unit/sagemaker/workflow/test_zimmer_steps.py @@ -12,8 +12,11 @@ # language governing permissions and limitations under the License. """Unit tests for Zimmer pipeline step types (v2). -Passthrough ``arguments: Dict[str, Any]`` API. The service validates -the argument schema server-side; this SDK does not duplicate that. +Passthrough ``arguments: Dict[str, Any]`` API. Top-level argument keys +are validated client-side against the public AWS API input shape +(botocore service model); fields known to be rejected by SageMaker +Pipelines fail fast at construction. Values are not validated -- they +may be pipeline variables. Full schema validation remains server-side. """ from __future__ import absolute_import @@ -330,3 +333,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": []})