From f4905494264d749fa2b3a8dace925dbaa3f2add6 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Sat, 8 Aug 2026 00:03:23 +0000 Subject: [PATCH 1/3] feat(pipeline): Add inference and lineage step types Add 4 pipeline step classes: - EndpointConfigStep, EndpointStep (SageMaker inference deployment) - InferenceComponentStep (multi-model endpoint support) - LineageStep (ML governance tracking) Design: each step accepts an 'arguments: Dict[str, Any]' forwarded to the pipeline service. Top-level argument keys are validated client-side against the corresponding public AWS API input shape (botocore service model) at construction and at serialization; fields the service is known to reject fail fast with actionable errors (EndpointConfig: DataCaptureConfig, ExplainerConfig; Endpoint: DeploymentConfig). Values are not validated -- they may be pipeline variables resolved at compile time. Full schema validation remains server-side. If the installed botocore does not know an operation, shape validation is skipped and the service remains the authority. Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config. Includes 23 unit tests and a LineageStep end-to-end integration test. --- X-AI-Prompt: Add the inference and lineage pipeline step types to the Python SDK with client-side argument validation X-AI-Tool: kiro-cli --- .../src/sagemaker/mlops/workflow/__init__.py | 8 + .../mlops/workflow/_argument_validation.py | 124 +++++++ .../sagemaker/mlops/workflow/endpoint_step.py | 223 ++++++++++++ .../workflow/inference_component_step.py | 107 ++++++ .../sagemaker/mlops/workflow/lineage_step.py | 117 ++++++ .../src/sagemaker/mlops/workflow/steps.py | 4 + .../tests/integ/workflow/test_lineage_step.py | 145 ++++++++ .../workflow/test_inference_lineage_steps.py | 332 ++++++++++++++++++ 8 files changed, 1060 insertions(+) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_lineage_step.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 129abb1c76..a5ab5ba851 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -14,6 +14,7 @@ functions, conditions, properties) and can import from sagemaker.train and sagemaker.serve for orchestration purposes. """ + from __future__ import absolute_import __version__ = "0.1.0" @@ -46,8 +47,11 @@ from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep from sagemaker.mlops.workflow.condition_step import ConditionStep from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep from sagemaker.mlops.workflow.fail_step import FailStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput +from sagemaker.mlops.workflow.lineage_step import LineageStep from sagemaker.mlops.workflow.model_step import ModelStep from sagemaker.mlops.workflow.monitor_batch_transform_step import MonitorBatchTransformStep from sagemaker.mlops.workflow.notebook_job_step import NotebookJobStep @@ -98,9 +102,13 @@ "ConditionStep", "EMRStep", "EMRStepConfig", + "EndpointConfigStep", + "EndpointStep", "FailStep", + "InferenceComponentStep", "LambdaStep", "LambdaOutput", + "LineageStep", "ModelStep", "MonitorBatchTransformStep", "NotebookJobStep", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py new file mode 100644 index 0000000000..24072c4d4a --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py @@ -0,0 +1,124 @@ +# 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 release), 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) -> allowed top-level keys. +# ``None`` means botocore does not know the operation; skip shape checks. +_SHAPE_CACHE: Dict[Tuple[str, str], 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``). + + 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) + 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() + _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] = (), +) -> 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. + + 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) + if allowed is None: + return + unknown = sorted(set(arguments) - allowed) + if unknown: + raise ValueError( + f"{step_class_name}: unknown argument field(s) {unknown}. " + f"Allowed top-level fields (from {service_name}.{operation_name}): " + f"{sorted(allowed)}." + ) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py new file mode 100644 index 0000000000..b3588f359a --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -0,0 +1,223 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definitions for SageMaker Endpoint deployment in Pipelines. + +Design note: the pipeline service models each step's +``Arguments`` block as an opaque structure validated against +the underlying SageMaker request model (``CreateEndpointConfigInput`` +or ``CreateEndpointInput``) minus a small exclusion set. This SDK +validates the **top-level keys** of the ``arguments`` dict against the +public ``CreateEndpointConfig``/``CreateEndpoint`` API input shape at +construction time (values are not validated -- they may be pipeline +variables) and forwards the dict to the service, which remains the +authority on full schema validation. + +Excluded fields (the pipeline service rejects the pipeline if present): + +* ``EndpointConfig``: ``DataCaptureConfig``, ``ExplainerConfig`` +* ``Endpoint``: ``DeploymentConfig`` +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow._argument_validation import validate_step_arguments +from sagemaker.mlops.workflow.retry import RetryPolicy +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import ( + CacheConfig, + ConfigurableRetryStep, + Step, + StepTypeEnum, +) + + +class EndpointConfigStep(ConfigurableRetryStep): + """Creates a SageMaker EndpointConfig within a pipeline. + + Wraps the SageMaker ``CreateEndpointConfig`` API. The ``arguments`` + dict is passed through to the service; it accepts any field of + ``CreateEndpointConfigInput`` **except** ``DataCaptureConfig`` and + ``ExplainerConfig``, which are rejected by the pipeline service. + + Per the pipeline service's step contract, ``EndpointConfig`` is structurally + cacheable (``cache_config``) and retryable (``retry_policies``). + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: Optional[CacheConfig] = None, + retry_policies: Optional[List[RetryPolicy]] = None, + ): + """Construct an ``EndpointConfigStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateEndpointConfig`` call. Required fields: + ``EndpointConfigName``, ``ProductionVariants``. Optional + fields include ``KmsKeyId``, ``AsyncInferenceConfig``, + ``ShadowProductionVariants``, ``ExecutionRoleArn``, + ``VpcConfig``, ``EnableNetworkIsolation``, + ``MetricsConfig``. Values may be pipeline variables + (parameter references, step property references) — the + pipeline compiler resolves them at definition time. + Do not include ``DataCaptureConfig`` or ``ExplainerConfig`` + (the pipeline service rejects them). + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + cache_config (CacheConfig): Optional cache configuration. + retry_policies (List[RetryPolicy]): Optional retry policies. + """ + super().__init__( + name=name, + step_type=StepTypeEnum.ENDPOINT_CONFIG, + display_name=display_name, + description=description, + depends_on=depends_on, + retry_policies=retry_policies, + ) + if arguments is None: + raise ValueError("arguments is required for EndpointConfigStep.") + 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. The + ``arguments`` dict is passed through to the service; it accepts any + field of ``CreateEndpointInput`` **except** ``DeploymentConfig``, + which is rejected by the pipeline service. + + Per the pipeline service's step contract, ``Endpoint`` is structurally cacheable + but not retryable at the pipeline level. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + cache_config: Optional[CacheConfig] = None, + ): + """Construct an ``EndpointStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateEndpoint`` / ``UpdateEndpoint`` call. Required + fields: ``EndpointName``, ``EndpointConfigName``. Optional + fields: ``GraphConfigName``, ``DeletionCondition``. + Values may be pipeline variables. Do not include + ``DeploymentConfig`` (the pipeline service rejects it). + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + cache_config (CacheConfig): Optional cache configuration. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.ENDPOINT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for EndpointStep.") + 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/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py new file mode 100644 index 0000000000..935aab1078 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py @@ -0,0 +1,107 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definition for SageMaker InferenceComponent in Pipelines. + +Design note: the ``Arguments`` block is validated server-side against SageMaker's +``CreateInferenceComponentInput`` request model with no field +exclusions — any field the AWS API accepts, the pipeline service +accepts. +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow._argument_validation import validate_step_arguments +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class InferenceComponentStep(Step): + """Creates or updates a SageMaker Inference Component within a pipeline. + + Wraps the SageMaker ``CreateInferenceComponent``/``UpdateInferenceComponent`` + API — the pipeline chooses create-vs-update based on component existence. + Inference components enable multi-model endpoint deployments with + independent scaling per model. + + The ``arguments`` dict is passed through to the service; it accepts + any field of ``CreateInferenceComponentInput`` (no exclusions). + + Per the pipeline service's step contract, ``InferenceComponent`` is neither + cacheable nor retryable at the pipeline level. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct an ``InferenceComponentStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block for the + ``CreateInferenceComponent``/``UpdateInferenceComponent`` + call. Typical fields: ``InferenceComponentName``, + ``EndpointName``, ``VariantName``, ``Specification``, + ``Specifications`` (plural, for multi-spec deployments), + ``RuntimeConfig``. Values may be pipeline variables. + Note: ``ComputeResourceRequirements.NumberOfCpuCoresRequired`` + is a float — pass ``2.0`` not ``2``. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.INFERENCE_COMPONENT, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for InferenceComponentStep.") + 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/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py new file mode 100644 index 0000000000..20b053d9e1 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -0,0 +1,117 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Step definition for SageMaker Lineage tracking in Pipelines. + +Design note: the ``Arguments`` block is a structure of — four optional +lists: + +* ``Actions`` — list of ``CreateActionRequest`` shapes +* ``Artifacts`` — list of ``CreateArtifactRequest`` shapes +* ``Contexts`` — list of ``CreateContextRequest`` shapes +* ``Associations`` — list of ``LineageAssociation`` shapes + (``Source``/``Destination``/``AssociationType``) + +The SDK validates that the ``arguments`` dict contains only these four +top-level keys (at least one required) and forwards it to the service. +""" + +from __future__ import absolute_import + +from typing import Any, Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.properties import Properties + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class LineageStep(Step): + """Creates and associates lineage entities in SageMaker's lineage system. + + Wraps SageMaker's ``CreateAction``/``CreateArtifact``/``CreateContext`` + and lineage ``AddAssociation`` APIs. A single step may create + multiple entities of any of the four types (Actions, Artifacts, + Contexts, Associations). Property references use + ``Steps..ActionArns['']``, + ``Steps..ArtifactArns['']``, + ``Steps..ContextArns['']``, and + ``Steps..Associations``. + """ + + def __init__( + self, + name: str, + arguments: Dict[str, Any], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Construct a ``LineageStep``. + + Args: + name (str): The name of the step. + arguments (Dict[str, Any]): The ``Arguments`` block. Recognized + top-level keys: ``Actions``, ``Artifacts``, ``Contexts``, + ``Associations`` — each is a list of dicts conforming to + the corresponding SageMaker API shape (or the service's + ``LineageAssociation`` for ``Associations``). At least + one of the four keys must be present. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + + Raises: + ValueError: If ``arguments`` is None or contains none of the + recognized keys. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.LINEAGE, + depends_on=depends_on, + ) + if arguments is None: + raise ValueError("arguments is required for LineageStep.") + 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/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 76e90a5309..60b7420844 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -62,6 +62,10 @@ class StepTypeEnum(Enum): EMR_SERVERLESS = "EMRServerless" FAIL = "Fail" AUTOML = "AutoML" + ENDPOINT_CONFIG = "EndpointConfig" + ENDPOINT = "Endpoint" + INFERENCE_COMPONENT = "InferenceComponent" + LINEAGE = "Lineage" class Step(Entity): diff --git a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py new file mode 100644 index 0000000000..3d2d30efe2 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -0,0 +1,145 @@ +# 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 LineageStep. + +Creates a pipeline containing a single ``LineageStep`` that records a +SageMaker lineage Action, executes it end-to-end against the real +service, and asserts the execution reaches ``Succeeded``. Cleans up the +Action, the pipeline, and the S3 pipeline definition artifact. + +Requires the execution role to have ``sagemaker:CreateAction`` (and +related lineage permissions). ``SageMakerRole`` — the standard fixture +role used across the SDK's integ tests — has broad SageMaker access and +satisfies this requirement. + +This test represents the SDK-side end-to-end validation of the +inference and lineage step family. See ``endpoint_step.py`` and +``inference_component_step.py`` for the other step types; those are not +integ-tested here because they provision paid resources +(Endpoint/InferenceComponent). +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.pipeline import Pipeline + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role): + """Full end-to-end run of a LineageStep pipeline against the real service. + + Builds a pipeline with a single ``LineageStep`` that creates one + lineage ``Action``. Verifies the pipeline execution succeeds and the + server-reported step metadata contains the created action ARN. + """ + stamp = uuid.uuid4().hex[:8] + action_name = f"lineage-integ-{stamp}" + pipeline_name = f"integ-lineage-{stamp}" + + step = LineageStep( + name="RecordLineage", + arguments={ + "Actions": [ + { + "ActionName": action_name, + "ActionType": "ModelTraining", + "Status": "Completed", + "Source": { + "SourceUri": f"s3://lineage-integ-test/{stamp}/model.tar.gz", + "SourceType": "MODEL", + }, + "Description": "Lineage integ test action", + } + ] + }, + ) + pipeline = Pipeline( + name=pipeline_name, + steps=[step], + sagemaker_session=pipeline_session, + ) + + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + # LineageStep is metadata-only; execution completes quickly. Poll + # up to 5 minutes to give the service plenty of headroom under load. + timeout = 300 + start_time = time.time() + final_status = None + while time.time() - start_time < timeout: + execution_desc = execution.describe() + status = execution_desc["PipelineExecutionStatus"] + if status in ("Succeeded", "Failed", "Stopped"): + final_status = status + break + time.sleep(10) + + if final_status != "Succeeded": + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + failure_details = "\n".join( + f"{s['StepName']}: {s.get('FailureReason', 'no reason')}" + for s in steps + if s.get("StepStatus") == "Failed" + ) + pytest.fail(f"Pipeline execution status={final_status}. Details:\n{failure_details}") + + # Verify the step metadata reports the created action ARN. + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + lineage_step = next(s for s in steps if s["StepName"] == "RecordLineage") + assert lineage_step["StepStatus"] == "Succeeded" + metadata = lineage_step.get("Metadata", {}) + action_arns = metadata.get("Lineage", {}).get("ActionArns", {}) + assert ( + action_name in action_arns + ), f"expected {action_name} in ActionArns, got: {action_arns}" + assert action_arns[action_name].endswith(f":action/{action_name}") + + finally: + # Delete the lineage Action. + try: + sagemaker_session.sagemaker_client.delete_action(ActionName=action_name) + except Exception: + pass + # Delete the pipeline. + try: + sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=pipeline_name) + except Exception: + pass diff --git a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py new file mode 100644 index 0000000000..af32493e3c --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -0,0 +1,332 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for the inference and lineage pipeline step types. + +These steps use a passthrough ``arguments: Dict[str, Any]`` API, +mirroring ``LambdaStep``/``CallbackStep``. Top-level argument keys are validated +client-side against the corresponding public AWS API input shape +(botocore service model), and fields known to be rejected by SageMaker +Pipelines fail fast at construction. Values are not validated -- they +may be pipeline variables resolved at compile time. Full schema +validation remains server-side. +""" + +from __future__ import absolute_import + +import pytest + +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep +from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum + +# ---------- EndpointConfigStep ---------- + + +def test_endpoint_config_step_basic(): + step = EndpointConfigStep( + name="Cfg", + arguments={ + "EndpointConfigName": "MyCfg", + "ProductionVariants": [ + { + "VariantName": "AllTraffic", + "ModelName": "m", + "InstanceType": "ml.m5.large", + "InitialInstanceCount": 1, + } + ], + }, + ) + assert step.step_type == StepTypeEnum.ENDPOINT_CONFIG + assert step.arguments["EndpointConfigName"] == "MyCfg" + + +def test_endpoint_config_step_to_request_includes_cache_and_retry(): + step = EndpointConfigStep( + name="Cfg", + arguments={"EndpointConfigName": "MyCfg", "ProductionVariants": []}, + display_name="Create Config", + description="desc", + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), + ) + req = step.to_request() + assert req["Type"] == "EndpointConfig" + assert req["DisplayName"] == "Create Config" + assert req["Description"] == "desc" + assert req["CacheConfig"] == {"Enabled": True, "ExpireAfter": "P30D"} + + +def test_endpoint_config_step_accepts_full_api_surface(): + """User can pass any CreateEndpointConfigInput field (except the ones + the service excludes — that's a server-side rejection, not client-side).""" + step = EndpointConfigStep( + name="Cfg", + arguments={ + "EndpointConfigName": "MyCfg", + "ProductionVariants": [], + "KmsKeyId": "arn:aws:kms:...", + "ExecutionRoleArn": "arn:aws:iam:...", + "AsyncInferenceConfig": {"OutputConfig": {"S3OutputPath": "s3://x/"}}, + "VpcConfig": {"SecurityGroupIds": ["sg-0"], "Subnets": ["subnet-0"]}, + "EnableNetworkIsolation": False, + "ShadowProductionVariants": [], + }, + ) + args = step.arguments + assert args["KmsKeyId"] == "arn:aws:kms:..." + assert args["ExecutionRoleArn"] == "arn:aws:iam:..." + assert "OutputConfig" in args["AsyncInferenceConfig"] + + +def test_endpoint_config_step_requires_arguments(): + with pytest.raises(ValueError): + EndpointConfigStep(name="Cfg", arguments=None) + + +# ---------- EndpointStep ---------- + + +def test_endpoint_step_basic(): + step = EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + ) + assert step.step_type == StepTypeEnum.ENDPOINT + assert step.arguments == {"EndpointName": "ep", "EndpointConfigName": "cfg"} + + +def test_endpoint_step_cache_config(): + step = EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + cache_config=CacheConfig(enable_caching=True), + ) + req = step.to_request() + assert req["Type"] == "Endpoint" + assert req["CacheConfig"] == {"Enabled": True} + + +def test_endpoint_step_rejects_retry_policies_kwarg(): + """EndpointStep is not retryable — constructor must not accept retry_policies.""" + with pytest.raises(TypeError): + EndpointStep( + name="Deploy", + arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, + retry_policies=[], + ) + + +# ---------- InferenceComponentStep ---------- + + +def test_inference_component_step_basic(): + step = InferenceComponentStep( + name="IC", + arguments={ + "InferenceComponentName": "ic", + "EndpointName": "ep", + "VariantName": "v", + "Specification": { + "ModelName": "m", + "ComputeResourceRequirements": { + "MinMemoryRequiredInMb": 1024, + "NumberOfCpuCoresRequired": 2.0, + }, + }, + "RuntimeConfig": {"CopyCount": 1}, + }, + ) + assert step.step_type == StepTypeEnum.INFERENCE_COMPONENT + assert step.arguments["Specification"]["ModelName"] == "m" + + +def test_inference_component_step_rejects_retry_policies_kwarg(): + with pytest.raises(TypeError): + InferenceComponentStep( + name="IC", + arguments={}, + retry_policies=[], + ) + + +# ---------- LineageStep ---------- + + +def test_lineage_step_basic(): + step = LineageStep( + name="Rec", + arguments={ + "Actions": [ + { + "ActionName": "a1", + "ActionType": "ModelTraining", + "Status": "Completed", + } + ], + "Artifacts": [ + { + "ArtifactName": "art1", + "ArtifactType": "Model", + "Source": {"SourceUri": "s3://x/y"}, + } + ], + "Associations": [ + { + "Source": {"Name": "a1", "Type": "Action"}, + "Destination": {"Name": "art1", "Type": "Artifact"}, + "AssociationType": "Produced", + } + ], + }, + ) + assert step.step_type == StepTypeEnum.LINEAGE + assert len(step.arguments["Actions"]) == 1 + assert len(step.arguments["Associations"]) == 1 + + +def test_lineage_step_partial_arguments(): + step = LineageStep( + name="Rec", + arguments={"Actions": [{"ActionName": "a", "ActionType": "T", "Status": "Completed"}]}, + ) + assert "Actions" in step.arguments + assert "Artifacts" not in step.arguments + + +def test_lineage_step_requires_at_least_one_recognized_key(): + with pytest.raises(ValueError): + LineageStep(name="Rec", arguments={}) + with pytest.raises(ValueError): + LineageStep(name="Rec", arguments={"Bogus": []}) + + +def test_lineage_step_properties(): + step = LineageStep(name="Rec", arguments={"Actions": []}) + for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): + assert hasattr(step.properties, field) + + +# ---------- Cross-cutting ---------- + + +def test_all_steps_importable_from_init(): + from sagemaker.mlops.workflow import ( # noqa: F401 + EndpointConfigStep, + EndpointStep, + InferenceComponentStep, + LineageStep, + ) + + +def test_step_type_enum_values(): + assert StepTypeEnum.ENDPOINT_CONFIG.value == "EndpointConfig" + assert StepTypeEnum.ENDPOINT.value == "Endpoint" + assert StepTypeEnum.INFERENCE_COMPONENT.value == "InferenceComponent" + assert StepTypeEnum.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_empty_arguments_rejected(): + for cls, valid_key in ( + (EndpointConfigStep, "EndpointConfigName"), + (EndpointStep, "EndpointName"), + (InferenceComponentStep, "InferenceComponentName"), + ): + with pytest.raises(ValueError): + cls(name="x", arguments={}) + # sanity: a single valid key constructs fine + assert cls(name="x", arguments={valid_key: "v"}).arguments == {valid_key: "v"} + + +def test_pipeline_variable_values_pass_validation(): + """Only top-level keys are validated -- values may be pipeline + variables (Get expressions) at any position.""" + step = EndpointStep( + name="Deploy", + arguments={ + "EndpointName": {"Get": "Parameters.EndpointName"}, + "EndpointConfigName": {"Get": "Steps.Cfg.EndpointConfigName"}, + }, + ) + assert step.arguments["EndpointName"] == {"Get": "Parameters.EndpointName"} + + +def test_post_construction_mutation_caught_at_serialization(): + """Injecting an unsupported field after construction is caught when + the arguments property is read (i.e., at pipeline serialization).""" + step = EndpointConfigStep( + name="Cfg", + arguments={"EndpointConfigName": "cfg", "ProductionVariants": []}, + ) + step._arguments["DataCaptureConfig"] = {} + with pytest.raises(ValueError, match="DataCaptureConfig"): + _ = step.arguments + + +def test_lineage_step_rejects_unknown_keys_alongside_recognized(): + with pytest.raises(ValueError, match="Bogus"): + LineageStep(name="Rec", arguments={"Actions": [], "Bogus": []}) From 85001b1bb2e91c9c042538da64f042f966db9baf Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Mon, 31 Aug 2026 21:46:15 +0000 Subject: [PATCH 2/3] refactor(pipeline): Use step_args from PipelineSession for inference steps Address review feedback: EndpointConfigStep, EndpointStep, and InferenceComponentStep now take step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep, instead of a raw arguments dict. - Session.endpoint_from_production_variants, Session.create_endpoint, and Session.create_inference_component now route their service calls through _intercept_create_request. Under a plain Session the behavior is unchanged (the intercept is a pass-through); under a PipelineSession the request is captured and returned as step arguments, and no service call is made. - Each step validates the provenance of its step_args via validate_step_args_input (wrong producer or a raw dict is rejected). - The _argument_validation module is removed: requests are now built by the session methods rather than hand-authored, so client-side key validation is no longer needed. - Adds an integration test chaining EndpointConfigStep -> EndpointStep -> InferenceComponentStep in a single pipeline execution, with full resource cleanup. LineageStep is unchanged pending a design decision on multi-entity step arguments. --- X-AI-Prompt: Rework the inference step types to use step_args captured via PipelineSession per review feedback X-AI-Tool: kiro-cli --- .../sagemaker/core/helper/session_helper.py | 73 +++- .../mlops/workflow/_argument_validation.py | 124 ------ .../sagemaker/mlops/workflow/endpoint_step.py | 146 +++---- .../workflow/inference_component_step.py | 79 ++-- .../integ/workflow/test_deployment_steps.py | 172 +++++++++ .../workflow/test_inference_lineage_steps.py | 363 +++++++++--------- 6 files changed, 508 insertions(+), 449 deletions(-) delete mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index ecdd4b95eb..ab55cfa4e9 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -1091,15 +1091,23 @@ def endpoint_from_production_variants( config_options["ExecutionRoleArn"] = role logger.info("Creating endpoint-config with name %s", name) - self.sagemaker_client.create_endpoint_config(**config_options) - - return self.create_endpoint( - endpoint_name=name, - config_name=name, - tags=endpoint_tags, - wait=wait, - live_logging=live_logging, + + def submit(request): + self.sagemaker_client.create_endpoint_config(**request) + return self.create_endpoint( + endpoint_name=name, + config_name=name, + tags=endpoint_tags, + wait=wait, + live_logging=live_logging, + ) + + result = self._intercept_create_request( + config_options, submit, self.endpoint_from_production_variants.__name__ ) + if self._is_pipeline_context(): + return self.context + return result def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False): """Create an Amazon SageMaker ``Endpoint`` according to the configuration in the request. @@ -1129,16 +1137,28 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live tags = self._append_sagemaker_config_tags( tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT, TAGS) ) - try: - res = self.sagemaker_client.create_endpoint( - EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags - ) + create_endpoint_request = { + "EndpointName": endpoint_name, + "EndpointConfigName": config_name, + "Tags": tags, + } + + def submit(request): + res = self.sagemaker_client.create_endpoint(**request) if res: self.endpoint_arn = res["EndpointArn"] if wait: self.wait_for_endpoint(endpoint_name, live_logging=live_logging) return endpoint_name + + try: + result = self._intercept_create_request( + create_endpoint_request, submit, self.create_endpoint.__name__ + ) + if self._is_pipeline_context(): + return self.context + return result except Exception as e: troubleshooting = ( "https://docs.aws.amazon.com/sagemaker/latest/dg/" @@ -1257,10 +1277,18 @@ def create_inference_component( if tags and len(tags) != 0: request["Tags"] = tags - self.sagemaker_client.create_inference_component(**request) - if wait: - self.wait_for_inference_component(inference_component_name) - return inference_component_name + def submit(req): + self.sagemaker_client.create_inference_component(**req) + if wait: + self.wait_for_inference_component(inference_component_name) + return inference_component_name + + result = self._intercept_create_request( + request, submit, self.create_inference_component.__name__ + ) + if self._is_pipeline_context(): + return self.context + return result def wait_for_inference_component(self, inference_component_name, poll=20): """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete. @@ -1439,6 +1467,19 @@ def _intercept_create_request( """ return create(request) + def _is_pipeline_context(self) -> bool: + """Whether this session is a pipeline session capturing requests. + + Producer methods that support composing pipeline steps use this to + return the captured step arguments (``self.context``) instead of the + result of a service call. Always ``False`` for a plain ``Session``. + """ + # Lazy import to avoid a circular dependency: pipeline_context imports + # from this module at import time. + from sagemaker.core.workflow.pipeline_context import PipelineSession + + return isinstance(self, PipelineSession) + def _create_inference_recommendations_job_request( self, role: str, diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py deleted file mode 100644 index 24072c4d4a..0000000000 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_argument_validation.py +++ /dev/null @@ -1,124 +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. -"""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 release), 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) -> allowed top-level keys. -# ``None`` means botocore does not know the operation; skip shape checks. -_SHAPE_CACHE: Dict[Tuple[str, str], 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``). - - 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) - 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() - _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] = (), -) -> 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. - - 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) - if allowed is None: - return - unknown = sorted(set(arguments) - allowed) - if unknown: - raise ValueError( - f"{step_class_name}: unknown argument field(s) {unknown}. " - f"Allowed top-level fields (from {service_name}.{operation_name}): " - f"{sorted(allowed)}." - ) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py index b3588f359a..9d21af441b 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -12,30 +12,38 @@ # language governing permissions and limitations under the License. """Step definitions for SageMaker Endpoint deployment in Pipelines. -Design note: the pipeline service models each step's -``Arguments`` block as an opaque structure validated against -the underlying SageMaker request model (``CreateEndpointConfigInput`` -or ``CreateEndpointInput``) minus a small exclusion set. This SDK -validates the **top-level keys** of the ``arguments`` dict against the -public ``CreateEndpointConfig``/``CreateEndpoint`` API input shape at -construction time (values are not validated -- they may be pipeline -variables) and forwards the dict to the service, which remains the -authority on full schema validation. - -Excluded fields (the pipeline service rejects the pipeline if present): - -* ``EndpointConfig``: ``DataCaptureConfig``, ``ExplainerConfig`` -* ``Endpoint``: ``DeploymentConfig`` +These steps follow the ``step_args`` convention used by ``TrainingStep`` +and ``ModelStep``: call the corresponding session method under a +:class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and +pass the returned step arguments to the step. The request is captured at +call time and the service call is deferred to pipeline execution. + +Example:: + + pipeline_session = PipelineSession() + + config_step_args = pipeline_session.endpoint_from_production_variants( + name="my-endpoint-config", + production_variants=[...], + ) + config_step = EndpointConfigStep(name="CreateConfig", step_args=config_step_args) + + endpoint_step_args = pipeline_session.create_endpoint( + endpoint_name="my-endpoint", + config_name="my-endpoint-config", + ) + endpoint_step = EndpointStep(name="CreateEndpoint", step_args=endpoint_step_args) """ from __future__ import absolute_import -from typing import Any, Dict, List, Optional, Union +from typing import List, Optional, Union from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.pipeline_context import _JobStepArguments from sagemaker.core.workflow.properties import Properties +from sagemaker.core.workflow.utilities import validate_step_args_input -from sagemaker.mlops.workflow._argument_validation import validate_step_arguments from sagemaker.mlops.workflow.retry import RetryPolicy from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import ( @@ -49,19 +57,19 @@ class EndpointConfigStep(ConfigurableRetryStep): """Creates a SageMaker EndpointConfig within a pipeline. - Wraps the SageMaker ``CreateEndpointConfig`` API. The ``arguments`` - dict is passed through to the service; it accepts any field of - ``CreateEndpointConfigInput`` **except** ``DataCaptureConfig`` and - ``ExplainerConfig``, which are rejected by the pipeline service. + Wraps the SageMaker ``CreateEndpointConfig`` API. The ``step_args`` + must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.endpoint_from_production_variants` + on a ``PipelineSession``. - Per the pipeline service's step contract, ``EndpointConfig`` is structurally - cacheable (``cache_config``) and retryable (``retry_policies``). + ``EndpointConfig`` is structurally cacheable (``cache_config``) and + retryable (``retry_policies``). """ def __init__( self, name: str, - arguments: Dict[str, Any], + step_args: _JobStepArguments, display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, @@ -72,17 +80,9 @@ def __init__( Args: name (str): The name of the step. - arguments (Dict[str, Any]): The ``Arguments`` block for the - ``CreateEndpointConfig`` call. Required fields: - ``EndpointConfigName``, ``ProductionVariants``. Optional - fields include ``KmsKeyId``, ``AsyncInferenceConfig``, - ``ShadowProductionVariants``, ``ExecutionRoleArn``, - ``VpcConfig``, ``EnableNetworkIsolation``, - ``MetricsConfig``. Values may be pipeline variables - (parameter references, step property references) — the - pipeline compiler resolves them at definition time. - Do not include ``DataCaptureConfig`` or ``ExplainerConfig`` - (the pipeline service rejects them). + step_args (_JobStepArguments): The arguments for this step, + obtained from + ``pipeline_session.endpoint_from_production_variants()``. display_name (str): Optional display name. description (str): Optional description. depends_on (List[Union[str, Step, StepCollection]]): Optional @@ -98,16 +98,15 @@ def __init__( 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"), + validate_step_args_input( + step_args=step_args, + expected_caller={"endpoint_from_production_variants"}, + error_message=( + "The step_args of EndpointConfigStep must be obtained from " + "pipeline_session.endpoint_from_production_variants()." + ), ) - self._arguments = arguments + self.step_args = step_args self.cache_config = cache_config self._properties = Properties( step_name=name, step=self, shape_name="DescribeEndpointConfigOutput" @@ -115,15 +114,8 @@ def __init__( @property def arguments(self) -> RequestType: - """The ``Arguments`` block for the ``CreateEndpointConfig`` call.""" - validate_step_arguments( - "EndpointConfigStep", - self._arguments, - service_name="sagemaker", - operation_name="CreateEndpointConfig", - unsupported_fields=("DataCaptureConfig", "ExplainerConfig"), - ) - return self._arguments + """The arguments dictionary that is used to call ``create_endpoint_config``.""" + return self.step_args.args @property def properties(self): @@ -141,20 +133,20 @@ def to_request(self) -> RequestType: class EndpointStep(Step): """Creates or updates a SageMaker Endpoint within a pipeline. - Wraps the SageMaker ``CreateEndpoint``/``UpdateEndpoint`` API — the + Wraps the SageMaker ``CreateEndpoint``/``UpdateEndpoint`` API -- the pipeline chooses create-vs-update based on endpoint existence. The - ``arguments`` dict is passed through to the service; it accepts any - field of ``CreateEndpointInput`` **except** ``DeploymentConfig``, - which is rejected by the pipeline service. + ``step_args`` must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.create_endpoint` + on a ``PipelineSession``. - Per the pipeline service's step contract, ``Endpoint`` is structurally cacheable - but not retryable at the pipeline level. + ``Endpoint`` is structurally cacheable but not retryable at the + pipeline level. """ def __init__( self, name: str, - arguments: Dict[str, Any], + step_args: _JobStepArguments, display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, @@ -164,12 +156,8 @@ def __init__( Args: name (str): The name of the step. - arguments (Dict[str, Any]): The ``Arguments`` block for the - ``CreateEndpoint`` / ``UpdateEndpoint`` call. Required - fields: ``EndpointName``, ``EndpointConfigName``. Optional - fields: ``GraphConfigName``, ``DeletionCondition``. - Values may be pipeline variables. Do not include - ``DeploymentConfig`` (the pipeline service rejects it). + step_args (_JobStepArguments): The arguments for this step, + obtained from ``pipeline_session.create_endpoint()``. display_name (str): Optional display name. description (str): Optional description. depends_on (List[Union[str, Step, StepCollection]]): Optional @@ -183,16 +171,15 @@ def __init__( 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",), + validate_step_args_input( + step_args=step_args, + expected_caller={"create_endpoint"}, + error_message=( + "The step_args of EndpointStep must be obtained from " + "pipeline_session.create_endpoint()." + ), ) - self._arguments = arguments + self.step_args = step_args self.cache_config = cache_config self._properties = Properties( step_name=name, step=self, shape_name="DescribeEndpointOutput" @@ -200,15 +187,8 @@ def __init__( @property def arguments(self) -> RequestType: - """The ``Arguments`` block for the ``CreateEndpoint``/``UpdateEndpoint`` call.""" - validate_step_arguments( - "EndpointStep", - self._arguments, - service_name="sagemaker", - operation_name="CreateEndpoint", - unsupported_fields=("DeploymentConfig",), - ) - return self._arguments + """The arguments dictionary that is used to call ``create_endpoint``.""" + return self.step_args.args @property def properties(self): diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py index 935aab1078..39cbce2388 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py @@ -12,20 +12,33 @@ # language governing permissions and limitations under the License. """Step definition for SageMaker InferenceComponent in Pipelines. -Design note: the ``Arguments`` block is validated server-side against SageMaker's -``CreateInferenceComponentInput`` request model with no field -exclusions — any field the AWS API accepts, the pipeline service -accepts. +Follows the ``step_args`` convention: call +:meth:`~sagemaker.core.helper.session_helper.Session.create_inference_component` +under a :class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` +and pass the returned step arguments to the step. + +Example:: + + pipeline_session = PipelineSession() + + step_args = pipeline_session.create_inference_component( + inference_component_name="my-component", + endpoint_name="my-endpoint", + variant_name="AllTraffic", + specification={...}, + ) + step = InferenceComponentStep(name="CreateComponent", step_args=step_args) """ from __future__ import absolute_import -from typing import Any, Dict, List, Optional, Union +from typing import List, Optional, Union from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.pipeline_context import _JobStepArguments from sagemaker.core.workflow.properties import Properties +from sagemaker.core.workflow.utilities import validate_step_args_input -from sagemaker.mlops.workflow._argument_validation import validate_step_arguments from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum @@ -34,21 +47,22 @@ 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. + API -- the pipeline chooses create-vs-update based on component + existence. Inference components enable multi-model endpoint + deployments with independent scaling per model. - The ``arguments`` dict is passed through to the service; it accepts - any field of ``CreateInferenceComponentInput`` (no exclusions). + The ``step_args`` must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.create_inference_component` + on a ``PipelineSession``. - Per the pipeline service's step contract, ``InferenceComponent`` is neither - cacheable nor retryable at the pipeline level. + ``InferenceComponent`` is neither cacheable nor retryable at the + pipeline level. """ def __init__( self, name: str, - arguments: Dict[str, Any], + step_args: _JobStepArguments, display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, @@ -57,14 +71,9 @@ def __init__( Args: name (str): The name of the step. - arguments (Dict[str, Any]): The ``Arguments`` block for the - ``CreateInferenceComponent``/``UpdateInferenceComponent`` - call. Typical fields: ``InferenceComponentName``, - ``EndpointName``, ``VariantName``, ``Specification``, - ``Specifications`` (plural, for multi-spec deployments), - ``RuntimeConfig``. Values may be pipeline variables. - Note: ``ComputeResourceRequirements.NumberOfCpuCoresRequired`` - is a float — pass ``2.0`` not ``2``. + step_args (_JobStepArguments): The arguments for this step, + obtained from + ``pipeline_session.create_inference_component()``. display_name (str): Optional display name. description (str): Optional description. depends_on (List[Union[str, Step, StepCollection]]): Optional @@ -77,29 +86,23 @@ def __init__( 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", + validate_step_args_input( + step_args=step_args, + expected_caller={"create_inference_component"}, + error_message=( + "The step_args of InferenceComponentStep must be obtained from " + "pipeline_session.create_inference_component()." + ), ) - self._arguments = arguments + self.step_args = step_args 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 + """The arguments dictionary that is used to call ``create_inference_component``.""" + return self.step_args.args @property def properties(self): diff --git a/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py new file mode 100644 index 0000000000..142252b9b8 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py @@ -0,0 +1,172 @@ +# 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 inference deployment step types. + +Runs a single pipeline chaining ``EndpointConfigStep`` -> +``EndpointStep`` -> ``InferenceComponentStep`` end-to-end against the +real service: an inference-component-style endpoint config (no model +name on the variant, execution role on the config), an endpoint, and an +inference component carrying the container specification. + +This test provisions a real endpoint instance for its duration; all +resources are deleted in the ``finally`` block. +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker.core import image_uris +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep +from sagemaker.mlops.workflow.pipeline import Pipeline + +INSTANCE_TYPE = "ml.m5.xlarge" +EXECUTION_TIMEOUT_SECONDS = 45 * 60 +POLL_SECONDS = 30 + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def test_deployment_steps_execute_end_to_end(sagemaker_session, pipeline_session, role): + """Chained EndpointConfig -> Endpoint -> InferenceComponent pipeline run.""" + stamp = uuid.uuid4().hex[:8] + config_name = f"integ-deploy-cfg-{stamp}" + endpoint_name = f"integ-deploy-ep-{stamp}" + component_name = f"integ-deploy-ic-{stamp}" + pipeline_name = f"integ-deploy-{stamp}" + + image_uri = image_uris.retrieve( + framework="sklearn", + region=sagemaker_session.boto_region_name, + version="1.2-1", + instance_type=INSTANCE_TYPE, + ) + + config_step_args = pipeline_session.endpoint_from_production_variants( + name=config_name, + production_variants=[ + { + "VariantName": "AllTraffic", + "InstanceType": INSTANCE_TYPE, + "InitialInstanceCount": 1, + "ManagedInstanceScaling": { + "Status": "ENABLED", + "MinInstanceCount": 1, + "MaxInstanceCount": 1, + }, + "RoutingConfig": {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"}, + } + ], + role=role, + ) + config_step = EndpointConfigStep(name="CreateConfig", step_args=config_step_args) + + endpoint_step_args = pipeline_session.create_endpoint( + endpoint_name=endpoint_name, config_name=config_name + ) + endpoint_step = EndpointStep( + name="CreateEndpoint", step_args=endpoint_step_args, depends_on=[config_step] + ) + + component_step_args = pipeline_session.create_inference_component( + inference_component_name=component_name, + endpoint_name=endpoint_name, + variant_name="AllTraffic", + specification={ + "Container": {"Image": image_uri}, + "ComputeResourceRequirements": { + "NumberOfCpuCoresRequired": 1.0, + "MinMemoryRequiredInMb": 1024, + }, + }, + runtime_config={"CopyCount": 1}, + ) + component_step = InferenceComponentStep( + name="CreateComponent", step_args=component_step_args, depends_on=[endpoint_step] + ) + + pipeline = Pipeline( + name=pipeline_name, + steps=[config_step, endpoint_step, component_step], + sagemaker_session=pipeline_session, + ) + + sm_client = sagemaker_session.sagemaker_client + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + deadline = time.time() + EXECUTION_TIMEOUT_SECONDS + status = None + while time.time() < deadline: + status = execution.describe()["PipelineExecutionStatus"] + if status not in ("Executing", "Stopping"): + break + time.sleep(POLL_SECONDS) + + assert status == "Succeeded", f"Pipeline execution ended in status {status}" + + # Server-side resources exist with the expected names. + config_desc = sm_client.describe_endpoint_config(EndpointConfigName=config_name) + assert config_desc["EndpointConfigName"] == config_name + + endpoint_desc = sm_client.describe_endpoint(EndpointName=endpoint_name) + assert endpoint_desc["EndpointName"] == endpoint_name + + component_desc = sm_client.describe_inference_component( + InferenceComponentName=component_name + ) + assert component_desc["InferenceComponentName"] == component_name + assert component_desc["EndpointName"] == endpoint_name + finally: + for cleanup in ( + lambda: sm_client.delete_inference_component(InferenceComponentName=component_name), + lambda: _wait_component_deleted(sm_client, component_name), + lambda: sm_client.delete_endpoint(EndpointName=endpoint_name), + lambda: sm_client.delete_endpoint_config(EndpointConfigName=config_name), + lambda: pipeline.delete(), + ): + try: + cleanup() + except Exception: # noqa: BLE001 -- best-effort cleanup + pass + + +def _wait_component_deleted(sm_client, component_name, timeout_seconds=10 * 60): + """The endpoint cannot be deleted until its inference component is gone.""" + deadline = time.time() + timeout_seconds + while time.time() < deadline: + try: + sm_client.describe_inference_component(InferenceComponentName=component_name) + except Exception: + return + time.sleep(15) diff --git a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py index af32493e3c..e17e5ff966 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -12,154 +12,225 @@ # language governing permissions and limitations under the License. """Unit tests for the inference and lineage pipeline step types. -These steps use a passthrough ``arguments: Dict[str, Any]`` API, -mirroring ``LambdaStep``/``CallbackStep``. Top-level argument keys are validated -client-side against the corresponding public AWS API input shape -(botocore service model), and fields known to be rejected by SageMaker -Pipelines fail fast at construction. Values are not validated -- they -may be pipeline variables resolved at compile time. Full schema -validation remains server-side. +The inference steps (EndpointConfigStep, EndpointStep, +InferenceComponentStep) follow the ``step_args`` convention: the step +arguments are captured by calling the corresponding session method under +a ``PipelineSession``, which intercepts the request instead of calling +the service. """ from __future__ import absolute_import +from unittest.mock import Mock + import pytest +from sagemaker.core.workflow.pipeline_context import PipelineSession, _JobStepArguments from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.retry import ( + StepExceptionTypeEnum, + StepRetryPolicy, +) from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum -# ---------- EndpointConfigStep ---------- +ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" -def test_endpoint_config_step_basic(): - step = EndpointConfigStep( - name="Cfg", - arguments={ - "EndpointConfigName": "MyCfg", - "ProductionVariants": [ - { - "VariantName": "AllTraffic", - "ModelName": "m", - "InstanceType": "ml.m5.large", - "InitialInstanceCount": 1, - } - ], - }, +@pytest.fixture +def pipeline_session(): + """A PipelineSession with a mocked client -- no AWS calls are made.""" + return PipelineSession( + boto_session=Mock(region_name="us-west-2"), + sagemaker_client=Mock(), ) - 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"), +@pytest.fixture +def endpoint_config_step_args(pipeline_session): + return pipeline_session.endpoint_from_production_variants( + name="my-config", + production_variants=[ + { + "ModelName": "my-model", + "VariantName": "AllTraffic", + "InstanceType": "ml.m5.large", + "InitialInstanceCount": 1, + } + ], + kms_key="arn:aws:kms:us-west-2:123456789012:key/abc", ) + + +@pytest.fixture +def endpoint_step_args(pipeline_session): + return pipeline_session.create_endpoint(endpoint_name="my-endpoint", config_name="my-config") + + +@pytest.fixture +def inference_component_step_args(pipeline_session): + return pipeline_session.create_inference_component( + inference_component_name="my-component", + endpoint_name="my-endpoint", + variant_name="AllTraffic", + specification={"ModelName": "my-model"}, + runtime_config={"CopyCount": 2}, + ) + + +# ---------- step_args capture via PipelineSession ---------- + + +def test_capture_does_not_call_service(pipeline_session, endpoint_config_step_args): + assert isinstance(endpoint_config_step_args, _JobStepArguments) + assert not pipeline_session.sagemaker_client.create_endpoint_config.called + assert not pipeline_session.sagemaker_client.create_endpoint.called + + +def test_captured_request_content(endpoint_config_step_args): + args = endpoint_config_step_args.args + assert args["EndpointConfigName"] == "my-config" + assert args["KmsKeyId"] == "arn:aws:kms:us-west-2:123456789012:key/abc" + assert args["ProductionVariants"][0]["ModelName"] == "my-model" + + +# ---------- EndpointConfigStep ---------- + + +def test_endpoint_config_step_basic(endpoint_config_step_args): + step = EndpointConfigStep(name="Cfg", step_args=endpoint_config_step_args) + assert step.step_type == StepTypeEnum.ENDPOINT_CONFIG + assert step.arguments["EndpointConfigName"] == "my-config" req = step.to_request() assert req["Type"] == "EndpointConfig" - assert req["DisplayName"] == "Create Config" - assert req["Description"] == "desc" - assert req["CacheConfig"] == {"Enabled": True, "ExpireAfter": "P30D"} + assert req["Name"] == "Cfg" -def test_endpoint_config_step_accepts_full_api_surface(): - """User can pass any CreateEndpointConfigInput field (except the ones - the service excludes — that's a server-side rejection, not client-side).""" +def test_endpoint_config_step_to_request_includes_cache_and_retry( + endpoint_config_step_args, +): 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": [], - }, + step_args=endpoint_config_step_args, + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), + retry_policies=[ + StepRetryPolicy(exception_types=[StepExceptionTypeEnum.THROTTLING], max_attempts=3) + ], ) - args = step.arguments - assert args["KmsKeyId"] == "arn:aws:kms:..." - assert args["ExecutionRoleArn"] == "arn:aws:iam:..." - assert "OutputConfig" in args["AsyncInferenceConfig"] + req = step.to_request() + assert req["CacheConfig"]["Enabled"] is True + assert req["RetryPolicies"][0]["MaxAttempts"] == 3 -def test_endpoint_config_step_requires_arguments(): - with pytest.raises(ValueError): - EndpointConfigStep(name="Cfg", arguments=None) +def test_endpoint_config_step_rejects_wrong_producer(endpoint_step_args): + with pytest.raises(ValueError, match="endpoint_from_production_variants"): + EndpointConfigStep(name="Cfg", step_args=endpoint_step_args) + + +def test_endpoint_config_step_rejects_raw_dict(): + with pytest.raises(TypeError): + EndpointConfigStep(name="Cfg", step_args={"EndpointConfigName": "x"}) + + +def test_endpoint_config_step_properties(endpoint_config_step_args): + step = EndpointConfigStep(name="Cfg", step_args=endpoint_config_step_args) + assert step.properties.EndpointConfigName.expr == {"Get": "Steps.Cfg.EndpointConfigName"} # ---------- EndpointStep ---------- -def test_endpoint_step_basic(): - step = EndpointStep( - name="Deploy", - arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, - ) +def test_endpoint_step_basic(endpoint_step_args): + step = EndpointStep(name="Deploy", step_args=endpoint_step_args) assert step.step_type == StepTypeEnum.ENDPOINT - assert step.arguments == {"EndpointName": "ep", "EndpointConfigName": "cfg"} + assert step.arguments["EndpointName"] == "my-endpoint" + assert step.arguments["EndpointConfigName"] == "my-config" + assert step.to_request()["Type"] == "Endpoint" -def test_endpoint_step_cache_config(): +def test_endpoint_step_cache_config(endpoint_step_args): step = EndpointStep( name="Deploy", - arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, - cache_config=CacheConfig(enable_caching=True), + step_args=endpoint_step_args, + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), ) - req = step.to_request() - assert req["Type"] == "Endpoint" - assert req["CacheConfig"] == {"Enabled": True} + assert step.to_request()["CacheConfig"]["Enabled"] is True -def test_endpoint_step_rejects_retry_policies_kwarg(): - """EndpointStep is not retryable — constructor must not accept retry_policies.""" +def test_endpoint_step_rejects_retry_policies_kwarg(endpoint_step_args): + """EndpointStep is not retryable -- constructor must not accept retry_policies.""" with pytest.raises(TypeError): - EndpointStep( - name="Deploy", - arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, - retry_policies=[], - ) + EndpointStep(name="Deploy", step_args=endpoint_step_args, retry_policies=[]) + + +def test_endpoint_step_rejects_wrong_producer(endpoint_config_step_args): + with pytest.raises(ValueError, match="create_endpoint"): + EndpointStep(name="Deploy", step_args=endpoint_config_step_args) + + +def test_endpoint_step_properties(endpoint_step_args): + step = EndpointStep(name="Deploy", step_args=endpoint_step_args) + assert step.properties.EndpointName.expr == {"Get": "Steps.Deploy.EndpointName"} # ---------- 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}, - }, - ) +def test_inference_component_step_basic(inference_component_step_args): + step = InferenceComponentStep(name="IC", step_args=inference_component_step_args) assert step.step_type == StepTypeEnum.INFERENCE_COMPONENT - assert step.arguments["Specification"]["ModelName"] == "m" + args = step.arguments + assert args["InferenceComponentName"] == "my-component" + assert args["EndpointName"] == "my-endpoint" + assert args["VariantName"] == "AllTraffic" + assert args["Specification"] == {"ModelName": "my-model"} + assert args["RuntimeConfig"] == {"CopyCount": 2} + + +def test_inference_component_step_default_runtime_config(pipeline_session): + step_args = pipeline_session.create_inference_component( + inference_component_name="ic", + endpoint_name="ep", + variant_name="v", + specification={"ModelName": "m"}, + ) + step = InferenceComponentStep(name="IC", step_args=step_args) + assert step.arguments["RuntimeConfig"] == {"CopyCount": 1} -def test_inference_component_step_rejects_retry_policies_kwarg(): +def test_inference_component_step_rejects_retry_policies_kwarg( + inference_component_step_args, +): with pytest.raises(TypeError): InferenceComponentStep( - name="IC", - arguments={}, - retry_policies=[], + name="IC", step_args=inference_component_step_args, retry_policies=[] ) +def test_inference_component_step_rejects_wrong_producer(endpoint_step_args): + with pytest.raises(ValueError, match="create_inference_component"): + InferenceComponentStep(name="IC", step_args=endpoint_step_args) + + +def test_inference_component_step_properties(inference_component_step_args): + step = InferenceComponentStep(name="IC", step_args=inference_component_step_args) + assert step.properties.InferenceComponentName.expr == {"Get": "Steps.IC.InferenceComponentName"} + + +# ---------- plain Session behavior is unchanged ---------- + + +def test_plain_session_still_calls_service(): + from sagemaker.core.helper.session_helper import Session + + session = Session(boto_session=Mock(region_name="us-west-2"), sagemaker_client=Mock()) + session.sagemaker_client.create_endpoint.return_value = {"EndpointArn": "arn:x"} + name = session.create_endpoint(endpoint_name="ep", config_name="cfg", wait=False) + assert name == "ep" + assert session.sagemaker_client.create_endpoint.called + + # ---------- LineageStep ---------- @@ -211,6 +282,11 @@ def test_lineage_step_requires_at_least_one_recognized_key(): LineageStep(name="Rec", arguments={"Bogus": []}) +def test_lineage_step_rejects_unknown_keys_alongside_recognized(): + with pytest.raises(ValueError, match="Bogus"): + LineageStep(name="Rec", arguments={"Actions": [], "Bogus": []}) + + def test_lineage_step_properties(): step = LineageStep(name="Rec", arguments={"Actions": []}) for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): @@ -236,97 +312,8 @@ def test_step_type_enum_values(): assert StepTypeEnum.LINEAGE.value == "Lineage" -def test_depends_on_accepts_string_list(): - step = EndpointStep( - name="Deploy", - arguments={"EndpointName": "ep", "EndpointConfigName": "cfg"}, - depends_on=["Prev"], - ) +def test_depends_on_accepts_step_and_string(endpoint_config_step_args, endpoint_step_args): + cfg_step = EndpointConfigStep(name="Cfg", step_args=endpoint_config_step_args) + step = EndpointStep(name="Deploy", step_args=endpoint_step_args, depends_on=[cfg_step, "Other"]) 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_empty_arguments_rejected(): - for cls, valid_key in ( - (EndpointConfigStep, "EndpointConfigName"), - (EndpointStep, "EndpointName"), - (InferenceComponentStep, "InferenceComponentName"), - ): - 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": []}) + assert req["DependsOn"] == [cfg_step, "Other"] From 00feb86e6e9adb53f5b38b1009447f0c605eb5b1 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Tue, 1 Sep 2026 00:22:23 +0000 Subject: [PATCH 3/3] refactor(pipeline): Use step_args from PipelineSession for LineageStep LineageStep now creates one lineage entity per step, with step_args obtained by calling Action.create(), Artifact.create(), Context.create(), or Association.create() from sagemaker.core.lineage under a PipelineSession -- completing the step_args convention across all four new step types. - Record._invoke_api in apiutils captures the request as step arguments when the session is a PipelineSession and the call is one of the four lineage create methods; behavior under a plain Session is unchanged. - The step derives the service Arguments block from the captured call: Actions/Artifacts/Contexts wrap the create request; AddAssociation's SourceArn/DestinationArn are translated to entity references, so associations can reference entities created in other steps via step property references (Steps..ActionArns['']). - ActionArns/ArtifactArns/ContextArns properties support ['name'] item access for cross-step references. Adds 8 unit tests; updates the LineageStep integration test to the new API. --- X-AI-Prompt: Rework LineageStep to one-entity-per-step with step_args captured via PipelineSession, completing the review feedback X-AI-Tool: kiro-cli --- .../sagemaker/core/apiutils/_base_types.py | 19 +++ .../sagemaker/mlops/workflow/lineage_step.py | 145 ++++++++++++------ .../tests/integ/workflow/test_lineage_step.py | 26 ++-- .../workflow/test_inference_lineage_steps.py | 130 ++++++++++------ 4 files changed, 212 insertions(+), 108 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py index 3b762be826..5b5de4bd10 100644 --- a/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py +++ b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py @@ -219,10 +219,29 @@ def with_boto(self, boto_dict): ) return self + # Lineage entity creation methods whose requests are captured as pipeline + # step arguments when invoked under a ``PipelineSession`` (instead of + # calling the service). Used by ``sagemaker.mlops.workflow.LineageStep``. + _PIPELINE_CAPTURABLE_METHODS = frozenset( + {"create_action", "create_artifact", "create_context", "add_association"} + ) + def _invoke_api(self, boto_method, boto_method_members): """Invoke a SageMaker API.""" api_values = {k: v for k, v in vars(self).items() if k in boto_method_members} api_kwargs = self.to_boto(api_values) + + if boto_method in self._PIPELINE_CAPTURABLE_METHODS: + # Lazy import to avoid a circular dependency at module load time. + from sagemaker.core.workflow.pipeline_context import ( + PipelineSession, + _JobStepArguments, + ) + + if isinstance(self.sagemaker_session, PipelineSession): + self.sagemaker_session.context = _JobStepArguments(boto_method, api_kwargs) + return self.sagemaker_session.context + api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method) api_boto_response = api_method(**api_kwargs) return self.with_boto(api_boto_response) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py index 20b053d9e1..8ea5b55c3b 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -12,37 +12,87 @@ # language governing permissions and limitations under the License. """Step definition for SageMaker Lineage tracking in Pipelines. -Design note: the ``Arguments`` block is a structure of — four optional -lists: - -* ``Actions`` — list of ``CreateActionRequest`` shapes -* ``Artifacts`` — list of ``CreateArtifactRequest`` shapes -* ``Contexts`` — list of ``CreateContextRequest`` shapes -* ``Associations`` — list of ``LineageAssociation`` shapes - (``Source``/``Destination``/``AssociationType``) - -The SDK validates that the ``arguments`` dict contains only these four -top-level keys (at least one required) and forwards it to the service. +Follows the ``step_args`` convention: create a lineage entity with the +corresponding class from :mod:`sagemaker.core.lineage` under a +:class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and +pass the returned step arguments to the step. Each step creates one +lineage entity; associations between entities created in different +steps reference them by ARN via step property references. + +Example:: + + pipeline_session = PipelineSession() + + action_args = Action.create( + action_name="my-action", + source_uri="s3://bucket/model.tar.gz", + source_type="S3ETag", + action_type="ModelTraining", + status="Completed", + sagemaker_session=pipeline_session, + ) + action_step = LineageStep(name="RecordAction", step_args=action_args) + + association_args = Association.create( + source_arn=action_step.properties.ActionArns["my-action"], + destination_arn="arn:aws:sagemaker:...:artifact/abc", + association_type="Produced", + sagemaker_session=pipeline_session, + ) + association_step = LineageStep( + name="RecordAssociation", + step_args=association_args, + depends_on=[action_step], + ) """ from __future__ import absolute_import -from typing import Any, Dict, List, Optional, Union +from typing import List, Optional, Union from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.pipeline_context import _JobStepArguments from sagemaker.core.workflow.properties import Properties +from sagemaker.core.workflow.utilities import validate_step_args_input from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum +# Maps the captured lineage create call to the Arguments key the pipeline +# service expects. +_CALLER_TO_ARGUMENTS_KEY = { + "create_action": "Actions", + "create_artifact": "Artifacts", + "create_context": "Contexts", + "add_association": "Associations", +} + + +class _EntityArnMap(Properties): + """Map-style property access for service-native ARN maps. + + The lineage step's ``ActionArns``/``ArtifactArns``/``ContextArns`` + outputs are maps keyed by entity name. They are pipeline-service + outputs with no botocore shape, so this supports ``['name']`` access + without a shape lookup. + """ + + def __getitem__(self, item: str) -> Properties: + """Reference the ARN of the entity created under the given name.""" + return Properties(step_name=self.step_name, path=f"{self.path}['{item}']") + class LineageStep(Step): - """Creates and associates lineage entities in SageMaker's lineage system. + """Creates a lineage entity or association 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 + and lineage ``AddAssociation`` APIs. Each step creates one entity; + the ``step_args`` must be obtained by calling ``Action.create()``, + ``Artifact.create()``, ``Context.create()``, or + ``Association.create()`` from :mod:`sagemaker.core.lineage` with a + ``PipelineSession``. + + Property references expose the created entity: ``Steps..ActionArns['']``, ``Steps..ArtifactArns['']``, ``Steps..ContextArns['']``, and @@ -52,7 +102,7 @@ class LineageStep(Step): def __init__( self, name: str, - arguments: Dict[str, Any], + step_args: _JobStepArguments, display_name: Optional[str] = None, description: Optional[str] = None, depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, @@ -61,20 +111,14 @@ def __init__( Args: name (str): The name of the step. - arguments (Dict[str, Any]): The ``Arguments`` block. Recognized - top-level keys: ``Actions``, ``Artifacts``, ``Contexts``, - ``Associations`` — each is a list of dicts conforming to - the corresponding SageMaker API shape (or the service's - ``LineageAssociation`` for ``Associations``). At least - one of the four keys must be present. + step_args (_JobStepArguments): The arguments for this step, + obtained from ``Action.create()``, ``Artifact.create()``, + ``Context.create()``, or ``Association.create()`` called + with a ``PipelineSession``. 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, @@ -83,33 +127,38 @@ def __init__( 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 + validate_step_args_input( + step_args=step_args, + expected_caller=set(_CALLER_TO_ARGUMENTS_KEY), + error_message=( + "The step_args of LineageStep must be obtained from " + "Action.create(), Artifact.create(), Context.create(), or " + "Association.create() called with a PipelineSession." + ), + ) + self.step_args = step_args root = Properties(step_name=name, step=self) - for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): - root.__dict__[field] = Properties(step_name=name, path=field) + for field in ("ActionArns", "ArtifactArns", "ContextArns"): + root.__dict__[field] = _EntityArnMap(step_name=name, path=field) + root.__dict__["Associations"] = Properties(step_name=name, path="Associations") self._properties = root @property def arguments(self) -> RequestType: - """The ``Arguments`` block describing lineage entities and associations.""" - return self._arguments + """The ``Arguments`` block describing the lineage entity to create.""" + key = _CALLER_TO_ARGUMENTS_KEY[self.step_args.caller_name] + entity = self.step_args.args + if key == "Associations": + # The AddAssociation API uses SourceArn/DestinationArn; the + # pipeline service models associations as entity references. + entity = { + "Source": {"Arn": entity["SourceArn"]}, + "Destination": {"Arn": entity["DestinationArn"]}, + } + if "AssociationType" in self.step_args.args: + entity["AssociationType"] = self.step_args.args["AssociationType"] + return {key: [entity]} @property def properties(self): diff --git a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py index 3d2d30efe2..b96f86bb47 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -37,6 +37,7 @@ import pytest from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.lineage.action import Action from sagemaker.core.workflow.pipeline_context import PipelineSession from sagemaker.mlops.workflow.lineage_step import LineageStep from sagemaker.mlops.workflow.pipeline import Pipeline @@ -68,23 +69,16 @@ def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, ro action_name = f"lineage-integ-{stamp}" pipeline_name = f"integ-lineage-{stamp}" - step = LineageStep( - name="RecordLineage", - arguments={ - "Actions": [ - { - "ActionName": action_name, - "ActionType": "ModelTraining", - "Status": "Completed", - "Source": { - "SourceUri": f"s3://lineage-integ-test/{stamp}/model.tar.gz", - "SourceType": "MODEL", - }, - "Description": "Lineage integ test action", - } - ] - }, + step_args = Action.create( + action_name=action_name, + source_uri=f"s3://lineage-integ-test/{stamp}/model.tar.gz", + source_type="MODEL", + action_type="ModelTraining", + status="Completed", + description="Lineage integ test action", + sagemaker_session=pipeline_session, ) + step = LineageStep(name="RecordLineage", step_args=step_args) pipeline = Pipeline( name=pipeline_name, steps=[step], diff --git a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py index e17e5ff966..c4f94272eb 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -28,6 +28,10 @@ from sagemaker.core.workflow.pipeline_context import PipelineSession, _JobStepArguments from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep +from sagemaker.core.lineage.action import Action +from sagemaker.core.lineage.artifact import Artifact +from sagemaker.core.lineage.association import Association +from sagemaker.core.lineage.context import Context from sagemaker.mlops.workflow.lineage_step import LineageStep from sagemaker.mlops.workflow.retry import ( StepExceptionTypeEnum, @@ -234,63 +238,101 @@ def test_plain_session_still_calls_service(): # ---------- 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", - } - ], - }, +@pytest.fixture +def action_step_args(pipeline_session): + return Action.create( + action_name="act1", + source_uri="s3://bucket/model.tar.gz", + source_type="S3ETag", + action_type="ModelTraining", + status="Completed", + sagemaker_session=pipeline_session, ) + + +def test_lineage_step_action(pipeline_session, action_step_args): + step = LineageStep(name="RecA", step_args=action_step_args) assert step.step_type == StepTypeEnum.LINEAGE - assert len(step.arguments["Actions"]) == 1 - assert len(step.arguments["Associations"]) == 1 + args = step.arguments + assert list(args.keys()) == ["Actions"] + assert args["Actions"][0]["ActionName"] == "act1" + assert args["Actions"][0]["Source"]["SourceUri"] == "s3://bucket/model.tar.gz" + assert not pipeline_session.sagemaker_client.create_action.called + + +def test_lineage_step_artifact(pipeline_session): + step_args = Artifact.create( + artifact_name="art1", + source_uri="s3://bucket/data", + artifact_type="Model", + sagemaker_session=pipeline_session, + ) + step = LineageStep(name="RecB", step_args=step_args) + assert list(step.arguments.keys()) == ["Artifacts"] + assert step.arguments["Artifacts"][0]["ArtifactName"] == "art1" -def test_lineage_step_partial_arguments(): - step = LineageStep( - name="Rec", - arguments={"Actions": [{"ActionName": "a", "ActionType": "T", "Status": "Completed"}]}, +def test_lineage_step_context(pipeline_session): + step_args = Context.create( + context_name="ctx1", + source_uri="s3://bucket/ctx", + context_type="Endpoint", + sagemaker_session=pipeline_session, + ) + step = LineageStep(name="RecC", step_args=step_args) + assert list(step.arguments.keys()) == ["Contexts"] + assert step.arguments["Contexts"][0]["ContextName"] == "ctx1" + + +def test_lineage_step_association_translates_arns(pipeline_session, action_step_args): + action_step = LineageStep(name="RecA", step_args=action_step_args) + step_args = Association.create( + source_arn=action_step.properties.ActionArns["act1"], + destination_arn="arn:aws:sagemaker:us-west-2:123456789012:artifact/abc", + association_type="Produced", + sagemaker_session=pipeline_session, ) - assert "Actions" in step.arguments - assert "Artifacts" not in step.arguments + step = LineageStep(name="RecD", step_args=step_args, depends_on=[action_step]) + entity = step.arguments["Associations"][0] + # AddAssociation's SourceArn/DestinationArn become entity references. + assert entity["Source"]["Arn"].expr == {"Get": "Steps.RecA.ActionArns['act1']"} + assert entity["Destination"] == {"Arn": "arn:aws:sagemaker:us-west-2:123456789012:artifact/abc"} + assert entity["AssociationType"] == "Produced" + assert not pipeline_session.sagemaker_client.add_association.called -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_rejects_wrong_producer(endpoint_step_args): + with pytest.raises(ValueError, match="Action.create"): + LineageStep(name="Rec", step_args=endpoint_step_args) -def test_lineage_step_rejects_unknown_keys_alongside_recognized(): - with pytest.raises(ValueError, match="Bogus"): - LineageStep(name="Rec", arguments={"Actions": [], "Bogus": []}) +def test_lineage_step_rejects_raw_dict(): + with pytest.raises(TypeError): + LineageStep(name="Rec", step_args={"Actions": []}) -def test_lineage_step_properties(): - step = LineageStep(name="Rec", arguments={"Actions": []}) +def test_lineage_step_properties(action_step_args): + step = LineageStep(name="Rec", step_args=action_step_args) for field in ("ActionArns", "ArtifactArns", "ContextArns", "Associations"): assert hasattr(step.properties, field) + assert step.properties.ArtifactArns["x"].expr == {"Get": "Steps.Rec.ArtifactArns['x']"} + + +def test_lineage_create_on_plain_session_calls_service(): + from sagemaker.core.helper.session_helper import Session + + session = Session(boto_session=Mock(region_name="us-west-2"), sagemaker_client=Mock()) + session.sagemaker_client.create_action.return_value = {"ActionArn": "arn:x"} + result = Action.create( + action_name="a", + source_uri="s3://b", + source_type="S3ETag", + action_type="T", + status="Completed", + sagemaker_session=session, + ) + assert session.sagemaker_client.create_action.called + assert not isinstance(result, _JobStepArguments) # ---------- Cross-cutting ----------