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 new file mode 100644 index 0000000000..4431438333 --- /dev/null +++ b/src/sagemaker/workflow/bedrock_steps.py @@ -0,0 +1,336 @@ +# 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/endpoint_step.py b/src/sagemaker/workflow/endpoint_step.py new file mode 100644 index 0000000000..7af7694e82 --- /dev/null +++ b/src/sagemaker/workflow/endpoint_step.py @@ -0,0 +1,191 @@ +# 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._argument_validation import validate_step_arguments +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.") + 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( + step_name=name, step=self, shape_name="DescribeEndpointConfigOutput" + ) + + @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 + 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.") + 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( + step_name=name, step=self, shape_name="DescribeEndpointOutput" + ) + + @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 + 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..2c61eab258 --- /dev/null +++ b/src/sagemaker/workflow/inference_component_step.py @@ -0,0 +1,96 @@ +# 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._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 + + +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.") + 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" + ) + + @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 + 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..96cb5fec95 --- /dev/null +++ b/src/sagemaker/workflow/lineage_step.py @@ -0,0 +1,102 @@ +# 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.") + 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) + 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..8a5b0ab0c4 --- /dev/null +++ b/tests/unit/sagemaker/workflow/test_zimmer_steps.py @@ -0,0 +1,446 @@ +# 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. 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 + +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"] + + +# ---------- 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": []})