From acbbad1d0f008787994edd9380cdc637f03972ed Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Fri, 28 Aug 2026 00:19:10 +0000 Subject: [PATCH] refactor: Remove Bedrock step types and clean up step test modules Removes the four Bedrock pipeline step classes (BedrockCustomModelStep, BedrockCustomModelDeploymentStep, BedrockModelImportStep, BedrockProvisionedModelThroughputStep) and their StepTypeEnum values pending re-evaluation of SDK support for these step types. They can be reintroduced later if needed. The remaining step types (EndpointConfigStep, EndpointStep, InferenceComponentStep, LineageStep) are unchanged. Also renames the step test modules to descriptive names (test_inference_lineage_steps.py, test_lineage_step.py) and removes the now-unused camelCase-to-PascalCase member-name conversion from the argument validation helper (it existed only for the Bedrock APIs). --- X-AI-Prompt: Follow-up on the merged v2 backport: remove Bedrock step types and internal-codename file names X-AI-Tool: kiro-cli --- .../workflow/_argument_validation.py | 22 +- src/sagemaker/workflow/bedrock_steps.py | 336 ------------------ src/sagemaker/workflow/steps.py | 5 +- ...r_lineage_step.py => test_lineage_step.py} | 8 +- ...eps.py => test_inference_lineage_steps.py} | 142 +------- 5 files changed, 12 insertions(+), 501 deletions(-) delete mode 100644 src/sagemaker/workflow/bedrock_steps.py rename tests/integ/sagemaker/workflow/{test_zimmer_lineage_step.py => test_lineage_step.py} (94%) rename tests/unit/sagemaker/workflow/{test_zimmer_steps.py => test_inference_lineage_steps.py} (69%) diff --git a/src/sagemaker/workflow/_argument_validation.py b/src/sagemaker/workflow/_argument_validation.py index ec19b6c905..24072c4d4a 100644 --- a/src/sagemaker/workflow/_argument_validation.py +++ b/src/sagemaker/workflow/_argument_validation.py @@ -23,7 +23,7 @@ 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 +(for example, a very old botocore release), shape validation is skipped and the service remains the authority. """ @@ -38,36 +38,29 @@ logger = logging.getLogger(__name__) -# Cache of (service, operation, pascal_case) -> allowed top-level keys. +# Cache of (service, operation) -> 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]]] = {} +_SHAPE_CACHE: Dict[Tuple[str, str], Optional[FrozenSet[str]]] = {} -def _allowed_top_level_keys( - service_name: str, operation_name: str, pascal_case: bool -) -> Optional[FrozenSet[str]]: +def _allowed_top_level_keys(service_name: str, operation_name: str) -> 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) + cache_key = (service_name, operation_name) 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( @@ -86,7 +79,6 @@ def validate_step_arguments( 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. @@ -98,8 +90,6 @@ def validate_step_arguments( 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 @@ -122,7 +112,7 @@ def validate_step_arguments( "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) + allowed = _allowed_top_level_keys(service_name, operation_name) if allowed is None: return unknown = sorted(set(arguments) - allowed) diff --git a/src/sagemaker/workflow/bedrock_steps.py b/src/sagemaker/workflow/bedrock_steps.py deleted file mode 100644 index 4431438333..0000000000 --- a/src/sagemaker/workflow/bedrock_steps.py +++ /dev/null @@ -1,336 +0,0 @@ -# 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._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 = [ - "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.") - _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 - 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.") - _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 - 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.") - _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 - 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.") - _validate_bedrock_arguments( - "BedrockProvisionedModelThroughputStep", arguments, "CreateProvisionedModelThroughput" - ) - 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.""" - _validate_bedrock_arguments( - "BedrockProvisionedModelThroughputStep", - self._arguments, - "CreateProvisionedModelThroughput", - ) - return self._arguments - - @property - def properties(self): - """Fields from ``GetProvisionedModelThroughputResponse``.""" - return self._properties diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py index 3de5989d43..143af09ac0 100644 --- a/src/sagemaker/workflow/steps.py +++ b/src/sagemaker/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 @@ -74,10 +75,6 @@ class StepTypeEnum(Enum): 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" diff --git a/tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py b/tests/integ/sagemaker/workflow/test_lineage_step.py similarity index 94% rename from tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py rename to tests/integ/sagemaker/workflow/test_lineage_step.py index e8dbc004aa..f0ca4880d3 100644 --- a/tests/integ/sagemaker/workflow/test_zimmer_lineage_step.py +++ b/tests/integ/sagemaker/workflow/test_lineage_step.py @@ -10,7 +10,7 @@ # 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).""" +"""Integration test for the LineageStep (v2).""" from __future__ import absolute_import @@ -48,7 +48,7 @@ def test_lineage_step_execute_end_to_end(sagemaker_session, role, pipeline_name) satisfies this requirement. """ stamp = uuid.uuid4().hex[:8] - action_name = f"zimmer-integ-{stamp}" + action_name = f"lineage-integ-{stamp}" step = LineageStep( name="RecordLineage", @@ -59,10 +59,10 @@ def test_lineage_step_execute_end_to_end(sagemaker_session, role, pipeline_name) "ActionType": "ModelTraining", "Status": "Completed", "Source": { - "SourceUri": f"s3://zimmer-integ-test/{stamp}/model.tar.gz", + "SourceUri": f"s3://lineage-integ-test/{stamp}/model.tar.gz", "SourceType": "MODEL", }, - "Description": "Zimmer v2 integ test action", + "Description": "Lineage v2 integ test action", } ] }, diff --git a/tests/unit/sagemaker/workflow/test_zimmer_steps.py b/tests/unit/sagemaker/workflow/test_inference_lineage_steps.py similarity index 69% rename from tests/unit/sagemaker/workflow/test_zimmer_steps.py rename to tests/unit/sagemaker/workflow/test_inference_lineage_steps.py index 8a5b0ab0c4..710b2b429c 100644 --- a/tests/unit/sagemaker/workflow/test_zimmer_steps.py +++ b/tests/unit/sagemaker/workflow/test_inference_lineage_steps.py @@ -10,7 +10,7 @@ # 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). +"""Unit tests for the inference and lineage pipeline step types (v2). Passthrough ``arguments: Dict[str, Any]`` API. Top-level argument keys are validated client-side against the public AWS API input shape @@ -23,12 +23,6 @@ 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 @@ -156,108 +150,6 @@ def test_inference_component_step_rejects_retry_policies_kwarg(): 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 ---------- @@ -315,13 +207,6 @@ 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" @@ -379,36 +264,11 @@ def test_unknown_argument_key_rejected(): ) -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={})