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={})