From 71972f0c22900fccb59ed3b2d94e5b078cfbea73 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Fri, 11 Sep 2026 20:48:32 +0000 Subject: [PATCH 1/3] feat(pipeline): Add inference and lineage step types Adds four pipeline step classes -- EndpointConfigStep, EndpointStep, InferenceComponentStep and LineageStep -- along with their StepTypeEnum values and package exports. Each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep. Capture goes through one seam: PipelineSession._intercept_create_request returns the captured context, so a producer returns its result directly, and the base Session holds no pipeline-specific branch. Producers: - Session.create_endpoint_config, added here, creates only the endpoint configuration, matching this step's one-resource-per-step shape. endpoint_from_production_variants also creates the endpoint and so is not a producer; the request building is shared between the two. - Session.create_endpoint and Session.create_inference_component. - Action.create(), Artifact.create(), Context.create() for lineage, whose requests Record._invoke_api routes through the same seam. LineageStep records a batch of entities in one step, which is what the pipeline service models: it creates every entity in the step, then adds the associations between them. An association references a sibling by name and type, resolved against the entities that same step creates, or an existing entity by ARN. Because Association.create() takes ARNs only and a sibling has no ARN until the step runs, associations are declared with LineageAssociation and LineageEntityReference. A name and type reference that the step does not create is rejected at construction. Chaining note: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. The Endpoint step does not suffix. Testing: 38 unit tests, plus two integration tests run end to end against the service -- a chained EndpointConfig -> Endpoint -> InferenceComponent pipeline, and a LineageStep recording an Action, an Artifact and a Context with associations between them. --- X-AI-Prompt: Add the inference and lineage pipeline step types to the public SDK, reworked per review feedback to use one capture seam and to batch lineage entities per step X-AI-Tool: kiro-cli --- .../sagemaker/core/apiutils/_base_types.py | 23 +- .../sagemaker/core/helper/session_helper.py | 163 +++++- .../core/workflow/pipeline_context.py | 8 + .../src/sagemaker/mlops/workflow/__init__.py | 14 + .../sagemaker/mlops/workflow/endpoint_step.py | 205 +++++++ .../workflow/inference_component_step.py | 110 ++++ .../sagemaker/mlops/workflow/lineage_step.py | 332 +++++++++++ .../src/sagemaker/mlops/workflow/steps.py | 4 + .../integ/workflow/test_deployment_steps.py | 247 ++++++++ .../tests/integ/workflow/test_lineage_step.py | 217 +++++++ .../workflow/test_inference_lineage_steps.py | 539 ++++++++++++++++++ 11 files changed, 1834 insertions(+), 28 deletions(-) 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_deployment_steps.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-core/src/sagemaker/core/apiutils/_base_types.py b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py index 3b762be826..6e45fc3ffe 100644 --- a/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py +++ b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py @@ -219,10 +219,27 @@ def with_boto(self, boto_dict): ) return self + # Lineage entity creation methods that participate in pipeline step + # composition. Requests for these are routed through the session's + # ``_intercept_create_request`` seam, so a ``PipelineSession`` captures them + # as step arguments while a plain ``Session`` calls 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) - api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method) - api_boto_response = api_method(**api_kwargs) - return self.with_boto(api_boto_response) + + def submit(request): + api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method) + return self.with_boto(api_method(**request)) + + if boto_method in self._PIPELINE_CAPTURABLE_METHODS: + return self.sagemaker_session._intercept_create_request( + api_kwargs, submit, boto_method + ) + + return submit(api_kwargs) diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index ecdd4b95eb..be9ca8beed 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -1006,6 +1006,52 @@ def endpoint_from_production_variants( str: The name of the created ``Endpoint``. """ + config_options = self._build_endpoint_config_request( + name=name, + production_variants=production_variants, + tags=tags, + kms_key=kms_key, + data_capture_config_dict=data_capture_config_dict, + async_inference_config_dict=async_inference_config_dict, + explainer_config_dict=explainer_config_dict, + vpc_config=vpc_config, + enable_network_isolation=enable_network_isolation, + role=role, + ) + endpoint_tags = _append_project_tags(format_tags(tags)) + + 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 _build_endpoint_config_request( + self, + name, + production_variants, + tags=None, + kms_key=None, + data_capture_config_dict=None, + async_inference_config_dict=None, + explainer_config_dict=None, + vpc_config=None, + enable_network_isolation=None, + role=None, + ): + """Build the ``CreateEndpointConfig`` request, applying config defaults. + + Shared by ``create_endpoint_config`` and + ``endpoint_from_production_variants`` so the two cannot drift. + + Returns: + dict: The ``CreateEndpointConfig`` request. + """ supports_kms = any( [ instance_supports_kms(production_variant["InstanceType"]) @@ -1060,7 +1106,6 @@ def endpoint_from_production_variants( role = self.expand_role(role) config_options["ExecutionRoleArn"] = role endpoint_config_tags = _append_project_tags(format_tags(tags)) - endpoint_tags = _append_project_tags(format_tags(tags)) endpoint_config_tags = self._append_sagemaker_config_tags( endpoint_config_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS) @@ -1090,15 +1135,71 @@ def endpoint_from_production_variants( if role is not None: config_options["ExecutionRoleArn"] = role - logger.info("Creating endpoint-config with name %s", name) - self.sagemaker_client.create_endpoint_config(**config_options) + return config_options - return self.create_endpoint( - endpoint_name=name, - config_name=name, - tags=endpoint_tags, - wait=wait, - live_logging=live_logging, + def create_endpoint_config( + self, + name, + production_variants, + tags=None, + kms_key=None, + data_capture_config_dict=None, + async_inference_config_dict=None, + explainer_config_dict=None, + vpc_config=None, + enable_network_isolation=None, + role=None, + ): + """Create an Amazon SageMaker endpoint configuration and nothing else. + + Unlike ``endpoint_from_production_variants``, which creates an endpoint + configuration *and* the endpoint that uses it, this creates only the + configuration. That makes it the producer for a pipeline + ``EndpointConfigStep``, which composes one resource per step. + + Args: + name (str): Name of the endpoint configuration to create. + production_variants (list[dict[str, str]]): The list of production + variants to host at this endpoint. + tags (Optional[Tags]): Tags to apply to the endpoint configuration. + kms_key (str): KMS key ARN for encrypting the volume attached to the + instances hosting the endpoint. + data_capture_config_dict (dict): Configuration for capturing + inference request and response data. + async_inference_config_dict (dict): Configuration for asynchronous + inference. + explainer_config_dict (dict): Configuration for online explainers. + vpc_config (dict): The VpcConfig for the endpoint configuration. + enable_network_isolation (bool): Whether to isolate the model + containers deployed to the endpoint. + role (str): An execution role ARN, required when the production + variants carry no model name, as with inference components. + + Returns: + str: Name of the endpoint configuration created. Under a + ``PipelineSession`` the captured step arguments are returned + instead and no service call is made. + """ + config_options = self._build_endpoint_config_request( + name=name, + production_variants=production_variants, + tags=tags, + kms_key=kms_key, + data_capture_config_dict=data_capture_config_dict, + async_inference_config_dict=async_inference_config_dict, + explainer_config_dict=explainer_config_dict, + vpc_config=vpc_config, + enable_network_isolation=enable_network_isolation, + role=role, + ) + + def submit(request): + logger.info("Creating endpoint-config with name %s", name) + self.sagemaker_client.create_endpoint_config(**request) + return name + + return self._intercept_create_request( + config_options, submit, self.create_endpoint_config.__name__ ) def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False): @@ -1122,23 +1223,31 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live botocore.exceptions.ClientError: If Sagemaker throws an exception while creating endpoint. """ - logger.info("Creating endpoint with name %s", endpoint_name) - tags = format_tags(tags) or [] tags = _append_project_tags(tags) 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): + logger.info("Creating endpoint with name %s", endpoint_name) + 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: + return self._intercept_create_request( + create_endpoint_request, submit, self.create_endpoint.__name__ + ) except Exception as e: troubleshooting = ( "https://docs.aws.amazon.com/sagemaker/latest/dg/" @@ -1232,12 +1341,6 @@ def create_inference_component( Returns: str: Name of the Amazon SageMaker ``InferenceComponent`` if created. """ - LOGGER.info( - "Creating inference component with name %s for endpoint %s", - inference_component_name, - endpoint_name, - ) - if runtime_config is None: runtime_config = {"CopyCount": 1} @@ -1257,10 +1360,20 @@ 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): + LOGGER.info( + "Creating inference component with name %s for endpoint %s", + inference_component_name, + endpoint_name, + ) + self.sagemaker_client.create_inference_component(**req) + if wait: + self.wait_for_inference_component(inference_component_name) + return inference_component_name + + return self._intercept_create_request( + request, submit, self.create_inference_component.__name__ + ) def wait_for_inference_component(self, inference_component_name, poll=20): """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete. diff --git a/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py b/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py index a6f3ffe171..a2152df02e 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py +++ b/sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py @@ -220,6 +220,13 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None request (dict): the create job request create (functor): a functor calls the sagemaker client create method func_name (str): the name of the function needed intercepting + + Returns: + The captured pipeline context, so that a producer method can simply + return the result of this call: under a plain ``Session`` the base + implementation returns the service result, and here it returns the + step arguments instead. This keeps the pipeline branch inside + ``PipelineSession`` rather than in the base ``Session``. """ if func_name == "create_model": self.context.create_model_request = request @@ -229,6 +236,7 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None self.context.caller_name = func_name else: self.context = _JobStepArguments(func_name, request) + return self.context def init_model_step_arguments(self, model): """Create a `_ModelStepArguments` (if not exist) as pipeline context diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 129abb1c76..31e8a5e474 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,15 @@ 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 ( + LineageAssociation, + LineageEntityReference, + 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 +106,15 @@ "ConditionStep", "EMRStep", "EMRStepConfig", + "EndpointConfigStep", + "EndpointStep", "FailStep", + "InferenceComponentStep", "LambdaStep", "LambdaOutput", + "LineageAssociation", + "LineageEntityReference", + "LineageStep", "ModelStep", "MonitorBatchTransformStep", "NotebookJobStep", 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..a17d999229 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -0,0 +1,205 @@ +# 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. + +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.create_endpoint_config( + 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=config_step.properties.EndpointConfigName, + ) + endpoint_step = EndpointStep(name="CreateEndpoint", step_args=endpoint_step_args) +""" + +from __future__ import absolute_import + +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.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 ``step_args`` + must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.create_endpoint_config` + on a ``PipelineSession``. That method creates only the endpoint + configuration, matching this step's one-resource-per-step shape; + ``endpoint_from_production_variants`` also creates the endpoint and so + is not a valid producer here. + + ``EndpointConfig`` is structurally cacheable (``cache_config``) and + retryable (``retry_policies``). + """ + + def __init__( + self, + name: str, + step_args: _JobStepArguments, + 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. + step_args (_JobStepArguments): The arguments for this step, + obtained from ``pipeline_session.create_endpoint_config()``. + 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, + ) + validate_step_args_input( + step_args=step_args, + expected_caller={"create_endpoint_config"}, + error_message=( + "The step_args of EndpointConfigStep must be obtained from " + "pipeline_session.create_endpoint_config()." + ), + ) + self.step_args = step_args + self.cache_config = cache_config + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeEndpointConfigOutput" + ) + + @property + def arguments(self) -> RequestType: + """The arguments dictionary that is used to call ``create_endpoint_config``.""" + return self.step_args.args + + @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 + ``step_args`` must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.create_endpoint` + on a ``PipelineSession``. + + ``Endpoint`` is structurally cacheable but not retryable at the + pipeline level. + """ + + def __init__( + self, + name: str, + step_args: _JobStepArguments, + 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. + 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 + 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, + ) + 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.step_args = step_args + self.cache_config = cache_config + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeEndpointOutput" + ) + + @property + def arguments(self) -> RequestType: + """The arguments dictionary that is used to call ``create_endpoint``.""" + return self.step_args.args + + @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..39cbce2388 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/inference_component_step.py @@ -0,0 +1,110 @@ +# 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. + +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 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 + + +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 ``step_args`` must be obtained by calling + :meth:`~sagemaker.core.helper.session_helper.Session.create_inference_component` + on a ``PipelineSession``. + + ``InferenceComponent`` is neither cacheable nor retryable at the + pipeline level. + """ + + def __init__( + self, + name: str, + step_args: _JobStepArguments, + 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. + 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 + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.INFERENCE_COMPONENT, + depends_on=depends_on, + ) + 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.step_args = step_args + self._properties = Properties( + step_name=name, step=self, shape_name="DescribeInferenceComponentOutput" + ) + + @property + def arguments(self) -> RequestType: + """The arguments dictionary that is used to call ``create_inference_component``.""" + return self.step_args.args + + @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..4a9ba2f615 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.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. +"""Step definition for SageMaker Lineage tracking in Pipelines. + +A single ``LineageStep`` records a batch of lineage entities and the +associations between them. Entities follow the ``step_args`` convention: +create each one with the corresponding class from +:mod:`sagemaker.core.lineage` under a +:class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and pass +the captured arguments to the step. + +Associations are declared rather than captured. ``Association.create()`` +takes source and destination ARNs, but an entity created by the same step +has no ARN until the step runs, so a sibling is referenced by name and +type instead. An association may also reference a pre-existing entity by +literal ARN. + +Example:: + + pipeline_session = PipelineSession() + + step = LineageStep( + name="RecordLineage", + step_args=[ + Action.create( + action_name="training-run", + source_uri="s3://bucket/run", + source_type="S3ETag", + action_type="ModelTraining", + sagemaker_session=pipeline_session, + ), + Artifact.create( + artifact_name="trained-model", + source_uri="s3://bucket/model.tar.gz", + artifact_type="Model", + sagemaker_session=pipeline_session, + ), + ], + associations=[ + LineageAssociation( + source=LineageEntityReference(name="training-run", type="Action"), + destination=LineageEntityReference(name="trained-model", type="Artifact"), + association_type="Produced", + ), + ], + ) +""" + +from __future__ import absolute_import + +from typing import Dict, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import PipelineVariable, 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 + +ENTITY_TYPE_ACTION = "Action" +ENTITY_TYPE_ARTIFACT = "Artifact" +ENTITY_TYPE_CONTEXT = "Context" + +# Maps the captured lineage create call to the ``Arguments`` key the pipeline +# service expects, the entity type used when referencing the entity from an +# association, and the request field holding the entity name. +_ENTITY_SPECS = { + "create_action": ("Actions", ENTITY_TYPE_ACTION, "ActionName"), + "create_artifact": ("Artifacts", ENTITY_TYPE_ARTIFACT, "ArtifactName"), + "create_context": ("Contexts", ENTITY_TYPE_CONTEXT, "ContextName"), +} + + +class LineageEntityReference: + """Reference to a lineage entity, used as an association endpoint. + + Reference an entity created by the same ``LineageStep`` with ``name`` + and ``type``, or an entity that already exists with ``arn``. + + A name and type pair is resolved by the pipeline service against the + entities created by that same step. It cannot refer to an entity + created by a different step; use ``arn`` for anything created + elsewhere. + """ + + def __init__( + self, + name: Optional[Union[str, PipelineVariable]] = None, + type: Optional[str] = None, # pylint: disable=redefined-builtin + arn: Optional[Union[str, PipelineVariable]] = None, + ): + """Construct a ``LineageEntityReference``. + + Args: + name (str or PipelineVariable): Name of an entity created by the + same step. Must be given together with ``type``. + type (str): One of ``"Action"``, ``"Artifact"`` or ``"Context"``. + arn (str or PipelineVariable): ARN of an existing entity. Mutually + exclusive with ``name``/``type``. + """ + if arn is not None: + if name is not None or type is not None: + raise ValueError( + "A LineageEntityReference takes either arn, or name and type -- not both." + ) + else: + if name is None or type is None: + raise ValueError( + "A LineageEntityReference requires either arn, or both name and type." + ) + if type not in (ENTITY_TYPE_ACTION, ENTITY_TYPE_ARTIFACT, ENTITY_TYPE_CONTEXT): + raise ValueError( + f"Unsupported lineage entity type '{type}'. Expected one of " + f"{ENTITY_TYPE_ACTION}, {ENTITY_TYPE_ARTIFACT}, {ENTITY_TYPE_CONTEXT}." + ) + self.name = name + self.type = type + self.arn = arn + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + if self.arn is not None: + return {"Arn": self.arn} + return {"Name": self.name, "Type": self.type} + + +class LineageAssociation: + """An association between two lineage entities.""" + + def __init__( + self, + source: LineageEntityReference, + destination: LineageEntityReference, + association_type: Optional[str] = None, + ): + """Construct a ``LineageAssociation``. + + Args: + source (LineageEntityReference): The source entity. + destination (LineageEntityReference): The destination entity. + association_type (str): The association type, for example + ``ContributedTo``, ``AssociatedWith``, ``DerivedFrom`` or + ``Produced``. + """ + for role, ref in (("source", source), ("destination", destination)): + if not isinstance(ref, LineageEntityReference): + raise TypeError( + f"The {role} of a LineageAssociation must be a " + f"LineageEntityReference, got {type(ref).__name__}." + ) + self.source = source + self.destination = destination + self.association_type = association_type + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + request = { + "Source": self.source.to_request(), + "Destination": self.destination.to_request(), + } + if self.association_type is not None: + request["AssociationType"] = self.association_type + return request + + +class _EntityArnMap(Properties): + """Map-style property access for the step's ARN outputs. + + ``ActionArns``, ``ArtifactArns`` and ``ContextArns`` 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): + """Records lineage entities and their associations in one pipeline step. + + Wraps SageMaker's ``CreateAction``, ``CreateArtifact``, ``CreateContext`` + and ``AddAssociation`` APIs. The pipeline service creates every entity in + the step, then adds the associations between them, so associations can + reference siblings by name and type. + + The step exposes the ARNs of what it created as + ``Steps..ActionArns['']``, + ``Steps..ArtifactArns['']`` and + ``Steps..ContextArns['']``, for downstream steps to + consume. Note these cannot be used as an association endpoint in another + ``LineageStep``: the service resolves name and type only against the + entities created by the same step. + """ + + def __init__( + self, + name: str, + step_args: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None, + associations: Optional[List[LineageAssociation]] = None, + 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. + step_args (_JobStepArguments or list): The captured arguments for + each entity this step creates, obtained from + ``Action.create()``, ``Artifact.create()`` or + ``Context.create()`` called with a ``PipelineSession``. A + single value is accepted for a step that creates one entity. + associations (List[LineageAssociation]): Associations to add after + the entities are created. + 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.LINEAGE, + depends_on=depends_on, + ) + + if step_args is None: + entities = [] + elif isinstance(step_args, (list, tuple)): + entities = list(step_args) + else: + entities = [step_args] + + self.associations = list(associations) if associations else [] + + if not entities and not self.associations: + raise ValueError( + "A LineageStep requires at least one entity in step_args, or one association." + ) + + for entity in entities: + validate_step_args_input( + step_args=entity, + expected_caller=set(_ENTITY_SPECS), + error_message=( + "The step_args of LineageStep must be obtained from Action.create(), " + "Artifact.create() or Context.create() called with a PipelineSession. " + "Associations are passed to the associations argument instead, because " + "Association.create() cannot reference an entity created by the same step." + ), + ) + + for association in self.associations: + if not isinstance(association, LineageAssociation): + raise TypeError( + "Each entry in associations must be a LineageAssociation, got " + f"{type(association).__name__}." + ) + + self.step_args = entities + self._validate_sibling_references() + + root = Properties(step_name=name, step=self) + 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 + + def _created_entities(self) -> Dict[str, str]: + """Map the name of each entity created by this step to its type.""" + created = {} + for entity in self.step_args: + _, entity_type, name_field = _ENTITY_SPECS[entity.caller_name] + name = entity.args.get(name_field) + if isinstance(name, str): + created[name] = entity_type + return created + + def _validate_sibling_references(self) -> None: + """Reject a name and type reference that no entity in this step creates. + + The pipeline service resolves a name and type pair only against the + entities created by the same step, and fails the execution otherwise. + Catching it here turns a runtime failure into a construction error. + """ + created = self._created_entities() + for association in self.associations: + endpoints = ( + ("source", association.source), + ("destination", association.destination), + ) + for role, ref in endpoints: + if ref.arn is not None or not isinstance(ref.name, str): + # An ARN needs no lookup, and a pipeline variable cannot be + # compared against the names known at construction time. + continue + if created.get(ref.name) != ref.type: + raise ValueError( + f"The {role} of an association references {ref.type} '{ref.name}', " + f"which this step does not create. A name and type reference must " + f"name an entity created by the same LineageStep; use arn to " + f"reference an entity created elsewhere." + ) + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block: the entities and associations for this step.""" + request: RequestType = {} + for entity in self.step_args: + key, _, _ = _ENTITY_SPECS[entity.caller_name] + request.setdefault(key, []).append(entity.args) + if self.associations: + request["Associations"] = [a.to_request() for a in self.associations] + return request + + @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_deployment_steps.py b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py new file mode 100644 index 0000000000..e0807ddaf8 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py @@ -0,0 +1,247 @@ +# 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 os +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}" + model_name = f"integ-deploy-model-{stamp}" + + # An inference component must reference a model that actually serves + # ``/ping``, so use the XGBoost serving image with the checked-in churn + # model artifact (the same artifact the transform-job integ test uses). + # The model is created outside the pipeline: the step under test is + # InferenceComponentStep, not model creation. + region = sagemaker_session.boto_region_name + image_uri = image_uris.retrieve("xgboost", region, "0.90-1") + model_data_url = sagemaker_session.upload_data( + path=os.path.join( + os.path.dirname(os.path.dirname(__file__)), + "data", + "model", + "transform_job", + "xgb-churn-prediction-model.tar.gz", + ), + key_prefix=f"integ-deploy/{stamp}", + ) + sagemaker_session.create_model( + name=model_name, + role=role, + container_defs={"Image": image_uri, "ModelDataUrl": model_data_url}, + ) + + config_step_args = pipeline_session.create_endpoint_config( + 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) + + # The service appends an execution-unique suffix to names created by these + # steps, so downstream steps must reference the *created* resource via step + # properties rather than the requested name. + endpoint_step_args = pipeline_session.create_endpoint( + endpoint_name=endpoint_name, config_name=config_step.properties.EndpointConfigName + ) + 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_step.properties.EndpointName, + variant_name="AllTraffic", + specification={ + "ModelName": model_name, + "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) + + if status != "Succeeded": + steps = sm_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + details = "\n".join( + f"{s['StepName']}: {s.get('StepStatus')} {s.get('FailureReason', '')}" + for s in steps + ) + pytest.fail(f"Pipeline execution ended in status {status}. Steps:\n{details}") + + # Resolve the actual server-side names (the steps suffix them). + actual_config = _resolve( + sm_client.list_endpoint_configs(NameContains=stamp)["EndpointConfigs"], + "EndpointConfigName", + ) + actual_endpoint = _resolve( + sm_client.list_endpoints(NameContains=stamp)["Endpoints"], "EndpointName" + ) + actual_component = _resolve( + sm_client.list_inference_components(NameContains=stamp)["InferenceComponents"], + "InferenceComponentName", + ) + + assert ( + sm_client.describe_endpoint_config(EndpointConfigName=actual_config)[ + "EndpointConfigName" + ] + == actual_config + ) + assert sm_client.describe_endpoint(EndpointName=actual_endpoint)["EndpointStatus"] == ( + "InService" + ) + component_desc = sm_client.describe_inference_component( + InferenceComponentName=actual_component + ) + assert component_desc["EndpointName"] == actual_endpoint + finally: + _cleanup(sagemaker_session, sm_client, stamp, pipeline, model_name) + + +def _resolve(items, key): + """Return the single matching resource name from a list_* response.""" + assert len(items) == 1, f"expected exactly one {key} for this run, got {items}" + return items[0][key] + + +def _cleanup(sagemaker_session, sm_client, stamp, pipeline, model_name): + """Delete every resource this run created, in dependency order.""" + for component in sm_client.list_inference_components(NameContains=stamp)["InferenceComponents"]: + name = component["InferenceComponentName"] + try: + sm_client.delete_inference_component(InferenceComponentName=name) + _wait_component_deleted(sm_client, name) + except Exception: # noqa: BLE001 -- best-effort cleanup + pass + for endpoint in sm_client.list_endpoints(NameContains=stamp)["Endpoints"]: + try: + sm_client.delete_endpoint(EndpointName=endpoint["EndpointName"]) + except Exception: # noqa: BLE001 + pass + for config in sm_client.list_endpoint_configs(NameContains=stamp)["EndpointConfigs"]: + try: + sm_client.delete_endpoint_config(EndpointConfigName=config["EndpointConfigName"]) + except Exception: # noqa: BLE001 + pass + try: + pipeline.delete() + except Exception: # noqa: BLE001 + pass + try: + sm_client.delete_model(ModelName=model_name) + except Exception: # noqa: BLE001 + pass + try: + sagemaker_session.boto_session.client("s3").delete_object( + Bucket=sagemaker_session.default_bucket(), + Key=f"integ-deploy/{stamp}/xgb-churn-prediction-model.tar.gz", + ) + except Exception: # noqa: BLE001 + 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/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py new file mode 100644 index 0000000000..d3adbf53fb --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -0,0 +1,217 @@ +# 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 Action, Artifact and Context together, plus associations between +them referenced by name and type, executes it end-to-end against the real +service, and asserts the execution reaches ``Succeeded``. Cleans up every +created entity, the associations, and the pipeline. + +Requires the execution role to have ``sagemaker:CreateAction``, +``CreateArtifact``, ``CreateContext`` and ``AddAssociation`` (and the +matching delete permissions for cleanup). ``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 +LineageStep. The inference deployment steps are covered separately by +``test_deployment_steps.py``. +""" + +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.lineage.action import Action +from sagemaker.core.lineage.artifact import Artifact +from sagemaker.core.lineage.context import Context +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.lineage_step import ( + LineageAssociation, + LineageEntityReference, + 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 an Action, + an Artifact and a Context, plus associations between them referenced by + name and type. Verifies the execution succeeds and the server-reported + step metadata contains an ARN for every created entity. + """ + stamp = uuid.uuid4().hex[:8] + action_name = f"lineage-integ-action-{stamp}" + artifact_name = f"lineage-integ-artifact-{stamp}" + context_name = f"lineage-integ-context-{stamp}" + pipeline_name = f"integ-lineage-{stamp}" + + action_args = Action.create( + action_name=action_name, + source_uri=f"s3://lineage-integ-test/{stamp}/run", + source_type="MODEL", + action_type="ModelTraining", + status="Completed", + description="Lineage integ test action", + sagemaker_session=pipeline_session, + ) + artifact_args = Artifact.create( + artifact_name=artifact_name, + source_uri=f"s3://lineage-integ-test/{stamp}/model.tar.gz", + artifact_type="Model", + sagemaker_session=pipeline_session, + ) + context_args = Context.create( + context_name=context_name, + source_uri=f"s3://lineage-integ-test/{stamp}/experiment", + context_type="Experiment", + description="Lineage integ test context", + sagemaker_session=pipeline_session, + ) + + # Associations reference entities created by this same step by name and + # type. The service resolves them against the entities it just created, + # so a single step covers create-then-associate end to end. + step = LineageStep( + name="RecordLineage", + step_args=[action_args, artifact_args, context_args], + associations=[ + LineageAssociation( + source=LineageEntityReference(name=action_name, type="Action"), + destination=LineageEntityReference(name=artifact_name, type="Artifact"), + association_type="Produced", + ), + LineageAssociation( + source=LineageEntityReference(name=context_name, type="Context"), + destination=LineageEntityReference(name=action_name, type="Action"), + association_type="AssociatedWith", + ), + ], + ) + pipeline = Pipeline( + name=pipeline_name, + steps=[step], + sagemaker_session=pipeline_session, + ) + + # Bound before the try so cleanup never raises NameError if the execution + # fails before the metadata is read. + action_arns: dict = {} + artifact_arns: dict = {} + context_arns: dict = {} + + 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" + lineage_metadata = lineage_step.get("Metadata", {}).get("Lineage", {}) + + action_arns = lineage_metadata.get("ActionArns", {}) + artifact_arns = lineage_metadata.get("ArtifactArns", {}) + context_arns = lineage_metadata.get("ContextArns", {}) + 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}") + assert ( + artifact_name in artifact_arns + ), f"expected {artifact_name} in ArtifactArns, got: {artifact_arns}" + assert ":artifact/" in artifact_arns[artifact_name] + assert ( + context_name in context_arns + ), f"expected {context_name} in ContextArns, got: {context_arns}" + assert context_arns[context_name].endswith(f":context/{context_name}") + + # Both associations were added, proving the name-and-type references + # resolved against the entities this same step created. + associations = lineage_metadata.get("Associations", []) + assert len(associations) == 2, f"expected 2 associations, got: {associations}" + + finally: + client = sagemaker_session.sagemaker_client + action_arn = action_arns.get(action_name) + artifact_arn = artifact_arns.get(artifact_name) + context_arn = context_arns.get(context_name) + + deletes = [] + if action_arn and artifact_arn: + deletes.append( + lambda: client.delete_association(SourceArn=action_arn, DestinationArn=artifact_arn) + ) + if context_arn and action_arn: + deletes.append( + lambda: client.delete_association(SourceArn=context_arn, DestinationArn=action_arn) + ) + deletes.append(lambda: client.delete_action(ActionName=action_name)) + if artifact_arn: + deletes.append(lambda: client.delete_artifact(ArtifactArn=artifact_arn)) + deletes.append(lambda: client.delete_context(ContextName=context_name)) + deletes.append(lambda: client.delete_pipeline(PipelineName=pipeline_name)) + + for delete in deletes: + try: + delete() + except Exception: # noqa: BLE001 -- best-effort cleanup + 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..638912aeb1 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -0,0 +1,539 @@ +# 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. + +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.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 ( + LineageAssociation, + LineageEntityReference, + LineageStep, +) +from sagemaker.mlops.workflow.retry import ( + StepExceptionTypeEnum, + StepRetryPolicy, +) +from sagemaker.mlops.workflow.steps import CacheConfig, StepTypeEnum + +ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" + + +@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(), + ) + + +@pytest.fixture +def endpoint_config_step_args(pipeline_session): + return pipeline_session.create_endpoint_config( + 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["Name"] == "Cfg" + + +def test_endpoint_config_step_to_request_includes_cache_and_retry( + endpoint_config_step_args, +): + step = EndpointConfigStep( + name="Cfg", + 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) + ], + ) + req = step.to_request() + assert req["CacheConfig"]["Enabled"] is True + assert req["RetryPolicies"][0]["MaxAttempts"] == 3 + + +def test_endpoint_config_step_rejects_wrong_producer(endpoint_step_args): + with pytest.raises(ValueError, match="create_endpoint_config"): + 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(endpoint_step_args): + step = EndpointStep(name="Deploy", step_args=endpoint_step_args) + assert step.step_type == StepTypeEnum.ENDPOINT + assert step.arguments["EndpointName"] == "my-endpoint" + assert step.arguments["EndpointConfigName"] == "my-config" + assert step.to_request()["Type"] == "Endpoint" + + +def test_endpoint_step_cache_config(endpoint_step_args): + step = EndpointStep( + name="Deploy", + step_args=endpoint_step_args, + cache_config=CacheConfig(enable_caching=True, expire_after="P30D"), + ) + assert step.to_request()["CacheConfig"]["Enabled"] is True + + +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", 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(inference_component_step_args): + step = InferenceComponentStep(name="IC", step_args=inference_component_step_args) + assert step.step_type == StepTypeEnum.INFERENCE_COMPONENT + 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( + inference_component_step_args, +): + with pytest.raises(TypeError): + InferenceComponentStep( + 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 ---------- + + +@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 + 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_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_batches_multiple_entities(pipeline_session): + """One step carries Actions, Artifacts and Contexts together.""" + actions = [ + Action.create( + action_name=f"act{i}", + source_uri="s3://bucket/run", + source_type="S3ETag", + action_type="ModelTraining", + sagemaker_session=pipeline_session, + ) + for i in range(2) + ] + artifact = Artifact.create( + artifact_name="art1", + source_uri="s3://bucket/model.tar.gz", + artifact_type="Model", + sagemaker_session=pipeline_session, + ) + context = Context.create( + context_name="ctx1", + source_uri="s3://bucket/ctx", + context_type="Endpoint", + sagemaker_session=pipeline_session, + ) + + step = LineageStep(name="Rec", step_args=actions + [artifact, context]) + args = step.arguments + + assert [a["ActionName"] for a in args["Actions"]] == ["act0", "act1"] + assert [a["ArtifactName"] for a in args["Artifacts"]] == ["art1"] + assert [c["ContextName"] for c in args["Contexts"]] == ["ctx1"] + assert "Associations" not in args + assert not pipeline_session.sagemaker_client.create_action.called + assert not pipeline_session.sagemaker_client.create_artifact.called + assert not pipeline_session.sagemaker_client.create_context.called + + +def test_lineage_step_association_references_sibling_by_name_and_type(pipeline_session): + """Associations reference same-step entities by Name+Type, not by ARN.""" + action = Action.create( + action_name="act1", + source_uri="s3://bucket/run", + source_type="S3ETag", + action_type="ModelTraining", + sagemaker_session=pipeline_session, + ) + artifact = Artifact.create( + artifact_name="art1", + source_uri="s3://bucket/model.tar.gz", + artifact_type="Model", + sagemaker_session=pipeline_session, + ) + step = LineageStep( + name="Rec", + step_args=[action, artifact], + associations=[ + LineageAssociation( + source=LineageEntityReference(name="act1", type="Action"), + destination=LineageEntityReference(name="art1", type="Artifact"), + association_type="Produced", + ) + ], + ) + + association = step.arguments["Associations"][0] + assert association["Source"] == {"Name": "act1", "Type": "Action"} + assert association["Destination"] == {"Name": "art1", "Type": "Artifact"} + assert association["AssociationType"] == "Produced" + assert not pipeline_session.sagemaker_client.add_association.called + + +def test_lineage_step_association_accepts_literal_arn(pipeline_session, action_step_args): + """A pre-existing entity is referenced by ARN.""" + existing = "arn:aws:sagemaker:us-west-2:123456789012:artifact/abc" + step = LineageStep( + name="Rec", + step_args=action_step_args, + associations=[ + LineageAssociation( + source=LineageEntityReference(name="act1", type="Action"), + destination=LineageEntityReference(arn=existing), + association_type="Produced", + ) + ], + ) + assert step.arguments["Associations"][0]["Destination"] == {"Arn": existing} + + +def test_lineage_step_rejects_unresolvable_sibling_reference(action_step_args): + """A Name+Type reference this step does not create fails at construction. + + The service resolves Name+Type only against entities created by the same + step, and would otherwise fail the execution at runtime. + """ + with pytest.raises(ValueError, match="does not create"): + LineageStep( + name="Rec", + step_args=action_step_args, + associations=[ + LineageAssociation( + source=LineageEntityReference(name="act1", type="Action"), + destination=LineageEntityReference(name="ghost", type="Artifact"), + ) + ], + ) + + +def test_lineage_step_rejects_captured_association(pipeline_session, action_step_args): + """Association.create() cannot express a sibling, so it is not valid step_args.""" + captured = Association.create( + source_arn="arn:aws:sagemaker:us-west-2:123456789012:action/a", + destination_arn="arn:aws:sagemaker:us-west-2:123456789012:artifact/b", + association_type="Produced", + sagemaker_session=pipeline_session, + ) + with pytest.raises(ValueError, match="associations argument"): + LineageStep(name="Rec", step_args=[action_step_args, captured]) + + +def test_lineage_step_requires_an_entity_or_association(): + with pytest.raises(ValueError, match="at least one entity"): + LineageStep(name="Rec") + + +def test_lineage_entity_reference_validation(): + with pytest.raises(ValueError, match="not both"): + LineageEntityReference(name="a", type="Action", arn="arn:x") + with pytest.raises(ValueError, match="both name and type"): + LineageEntityReference(name="a") + with pytest.raises(ValueError, match="Unsupported lineage entity type"): + LineageEntityReference(name="a", type="Endpoint") + + +def test_lineage_association_rejects_non_reference_endpoint(): + with pytest.raises(TypeError, match="LineageEntityReference"): + LineageAssociation( + source="arn:aws:sagemaker:us-west-2:123456789012:action/a", + destination=LineageEntityReference(arn="arn:x"), + ) + + +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_raw_dict(): + with pytest.raises(TypeError): + LineageStep(name="Rec", step_args={"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_all_four_producers_route_through_the_intercept_seam(pipeline_session): + """Every step's capture goes through PipelineSession._intercept_create_request. + + Guards against re-introducing a second capture mechanism: the seam is the + only path, so patching it is enough to observe all four producers. + """ + seen = [] + real = pipeline_session._intercept_create_request + + def record(request, create, func_name=None): + seen.append(func_name) + return real(request, create, func_name) + + pipeline_session._intercept_create_request = record + + pipeline_session.create_endpoint_config( + name="cfg", + production_variants=[{"VariantName": "AllTraffic"}], + role="arn:aws:iam::123456789012:role/SageMakerRole", + ) + pipeline_session.create_endpoint(endpoint_name="ep", config_name="cfg") + pipeline_session.create_inference_component( + inference_component_name="ic", + endpoint_name="ep", + variant_name="AllTraffic", + specification={"ModelName": "m"}, + ) + Action.create( + action_name="a", + source_uri="s3://b", + source_type="S3ETag", + action_type="T", + sagemaker_session=pipeline_session, + ) + + assert seen == [ + "create_endpoint_config", + "create_endpoint", + "create_inference_component", + "create_action", + ] + + +def test_base_session_has_no_pipeline_branch(): + """Base Session must stay pipeline-agnostic (no _is_pipeline_context helper).""" + from sagemaker.core.helper.session_helper import Session + + assert not hasattr(Session, "_is_pipeline_context") + + +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 ---------- + + +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_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"] == [cfg_step, "Other"] From 266514083f899defc2c98438dfca0bbd9ddc77f6 Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Tue, 15 Sep 2026 19:32:41 +0000 Subject: [PATCH 2/3] refactor(pipeline): Take lineage entities per kind on LineageStep Replaces LineageStep's single step_args argument with actions, artifacts and contexts, as requested in review. The caller states which kind each captured argument is instead of the step inferring it from the producer. This makes a misplaced entity a construction error that names the right producer -- passing an Artifact.create() result as actions now fails with "the actions of LineageStep must be obtained from Action.create()" -- rather than something the step silently files under the wrong key. The emitted Arguments are unchanged: Actions, Artifacts, Contexts and Associations, with empty kinds omitted. Each kind accepts a single captured value or a list. Testing: 39 unit tests, and the lineage integration test run end to end against the service, which creates an Action, an Artifact and a Context in one step with two associations referencing them by name and type. --- X-AI-Prompt: more comments X-AI-Tool: kiro-cli --- .../sagemaker/mlops/workflow/lineage_step.py | 126 ++++++++++++------ .../tests/integ/workflow/test_lineage_step.py | 4 +- .../workflow/test_inference_lineage_steps.py | 44 ++++-- 3 files changed, 118 insertions(+), 56 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py index 4a9ba2f615..e6fd54504d 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -17,7 +17,8 @@ create each one with the corresponding class from :mod:`sagemaker.core.lineage` under a :class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and pass -the captured arguments to the step. +the captured arguments to the matching argument of the step -- ``actions``, +``artifacts`` or ``contexts``. Associations are declared rather than captured. ``Association.create()`` takes source and destination ARNs, but an entity created by the same step @@ -31,7 +32,7 @@ step = LineageStep( name="RecordLineage", - step_args=[ + actions=[ Action.create( action_name="training-run", source_uri="s3://bucket/run", @@ -39,6 +40,8 @@ action_type="ModelTraining", sagemaker_session=pipeline_session, ), + ], + artifacts=[ Artifact.create( artifact_name="trained-model", source_uri="s3://bucket/model.tar.gz", @@ -72,16 +75,34 @@ ENTITY_TYPE_ARTIFACT = "Artifact" ENTITY_TYPE_CONTEXT = "Context" -# Maps the captured lineage create call to the ``Arguments`` key the pipeline +# One entry per kind of entity a LineageStep can create, in the order the +# entities appear in the request. Each maps the step argument to the lineage +# create call that must have produced it, the ``Arguments`` key the pipeline # service expects, the entity type used when referencing the entity from an # association, and the request field holding the entity name. -_ENTITY_SPECS = { - "create_action": ("Actions", ENTITY_TYPE_ACTION, "ActionName"), - "create_artifact": ("Artifacts", ENTITY_TYPE_ARTIFACT, "ArtifactName"), - "create_context": ("Contexts", ENTITY_TYPE_CONTEXT, "ContextName"), +_ENTITY_KINDS = ( + ("actions", "create_action", "Actions", ENTITY_TYPE_ACTION, "ActionName"), + ("artifacts", "create_artifact", "Artifacts", ENTITY_TYPE_ARTIFACT, "ArtifactName"), + ("contexts", "create_context", "Contexts", ENTITY_TYPE_CONTEXT, "ContextName"), +) + +# The producer each kind must come from, for the "wrong argument" error message. +_EXPECTED_PRODUCER = { + "actions": "Action.create()", + "artifacts": "Artifact.create()", + "contexts": "Context.create()", } +def _as_list(value) -> list: + """Normalise a single captured argument or a sequence of them to a list.""" + if value is None: + return [] + if isinstance(value, (list, tuple)): + return list(value) + return [value] + + class LineageEntityReference: """Reference to a lineage entity, used as an association endpoint. @@ -207,7 +228,9 @@ class LineageStep(Step): def __init__( self, name: str, - step_args: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None, + actions: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None, + artifacts: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None, + contexts: Optional[Union[_JobStepArguments, List[_JobStepArguments]]] = None, associations: Optional[List[LineageAssociation]] = None, display_name: Optional[str] = None, description: Optional[str] = None, @@ -217,11 +240,13 @@ def __init__( Args: name (str): The name of the step. - step_args (_JobStepArguments or list): The captured arguments for - each entity this step creates, obtained from - ``Action.create()``, ``Artifact.create()`` or - ``Context.create()`` called with a ``PipelineSession``. A - single value is accepted for a step that creates one entity. + actions (_JobStepArguments or list): Captured arguments from + ``Action.create()`` called with a ``PipelineSession``, for each + action this step creates. + artifacts (_JobStepArguments or list): Captured arguments from + ``Artifact.create()``, for each artifact this step creates. + contexts (_JobStepArguments or list): Captured arguments from + ``Context.create()``, for each context this step creates. associations (List[LineageAssociation]): Associations to add after the entities are created. display_name (str): Optional display name. @@ -237,30 +262,30 @@ def __init__( depends_on=depends_on, ) - if step_args is None: - entities = [] - elif isinstance(step_args, (list, tuple)): - entities = list(step_args) - else: - entities = [step_args] - self.associations = list(associations) if associations else [] - if not entities and not self.associations: - raise ValueError( - "A LineageStep requires at least one entity in step_args, or one association." - ) + supplied = {"actions": actions, "artifacts": artifacts, "contexts": contexts} + self._entities: Dict[str, List[_JobStepArguments]] = {} + for param, caller, _, _, _ in _ENTITY_KINDS: + entities = _as_list(supplied[param]) + for entity in entities: + validate_step_args_input( + step_args=entity, + expected_caller={caller}, + error_message=( + f"The {param} of LineageStep must be obtained from " + f"{_EXPECTED_PRODUCER[param]} called with a PipelineSession. " + "Associations are passed to the associations argument instead, " + "because Association.create() cannot reference an entity created " + "by the same step." + ), + ) + self._entities[param] = entities - for entity in entities: - validate_step_args_input( - step_args=entity, - expected_caller=set(_ENTITY_SPECS), - error_message=( - "The step_args of LineageStep must be obtained from Action.create(), " - "Artifact.create() or Context.create() called with a PipelineSession. " - "Associations are passed to the associations argument instead, because " - "Association.create() cannot reference an entity created by the same step." - ), + if not any(self._entities.values()) and not self.associations: + raise ValueError( + "A LineageStep requires at least one entity in actions, artifacts or " + "contexts, or one association." ) for association in self.associations: @@ -270,7 +295,6 @@ def __init__( f"{type(association).__name__}." ) - self.step_args = entities self._validate_sibling_references() root = Properties(step_name=name, step=self) @@ -279,14 +303,29 @@ def __init__( root.__dict__["Associations"] = Properties(step_name=name, path="Associations") self._properties = root + @property + def actions(self) -> List[_JobStepArguments]: + """The captured arguments for the actions this step creates.""" + return self._entities["actions"] + + @property + def artifacts(self) -> List[_JobStepArguments]: + """The captured arguments for the artifacts this step creates.""" + return self._entities["artifacts"] + + @property + def contexts(self) -> List[_JobStepArguments]: + """The captured arguments for the contexts this step creates.""" + return self._entities["contexts"] + def _created_entities(self) -> Dict[str, str]: """Map the name of each entity created by this step to its type.""" created = {} - for entity in self.step_args: - _, entity_type, name_field = _ENTITY_SPECS[entity.caller_name] - name = entity.args.get(name_field) - if isinstance(name, str): - created[name] = entity_type + for param, _, _, entity_type, name_field in _ENTITY_KINDS: + for entity in self._entities[param]: + name = entity.args.get(name_field) + if isinstance(name, str): + created[name] = entity_type return created def _validate_sibling_references(self) -> None: @@ -319,9 +358,10 @@ def _validate_sibling_references(self) -> None: def arguments(self) -> RequestType: """The ``Arguments`` block: the entities and associations for this step.""" request: RequestType = {} - for entity in self.step_args: - key, _, _ = _ENTITY_SPECS[entity.caller_name] - request.setdefault(key, []).append(entity.args) + for param, _, key, _, _ in _ENTITY_KINDS: + entities = self._entities[param] + if entities: + request[key] = [entity.args for entity in entities] if self.associations: request["Associations"] = [a.to_request() for a in self.associations] return request diff --git a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py index d3adbf53fb..b554b182cc 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -106,7 +106,9 @@ def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, ro # so a single step covers create-then-associate end to end. step = LineageStep( name="RecordLineage", - step_args=[action_args, artifact_args, context_args], + actions=[action_args], + artifacts=[artifact_args], + contexts=[context_args], associations=[ LineageAssociation( source=LineageEntityReference(name=action_name, type="Action"), 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 638912aeb1..c2f983d47a 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -255,7 +255,7 @@ def action_step_args(pipeline_session): def test_lineage_step_action(pipeline_session, action_step_args): - step = LineageStep(name="RecA", step_args=action_step_args) + step = LineageStep(name="RecA", actions=action_step_args) assert step.step_type == StepTypeEnum.LINEAGE args = step.arguments assert list(args.keys()) == ["Actions"] @@ -271,7 +271,7 @@ def test_lineage_step_artifact(pipeline_session): artifact_type="Model", sagemaker_session=pipeline_session, ) - step = LineageStep(name="RecB", step_args=step_args) + step = LineageStep(name="RecB", artifacts=step_args) assert list(step.arguments.keys()) == ["Artifacts"] assert step.arguments["Artifacts"][0]["ArtifactName"] == "art1" @@ -283,7 +283,7 @@ def test_lineage_step_context(pipeline_session): context_type="Endpoint", sagemaker_session=pipeline_session, ) - step = LineageStep(name="RecC", step_args=step_args) + step = LineageStep(name="RecC", contexts=step_args) assert list(step.arguments.keys()) == ["Contexts"] assert step.arguments["Contexts"][0]["ContextName"] == "ctx1" @@ -313,7 +313,7 @@ def test_lineage_step_batches_multiple_entities(pipeline_session): sagemaker_session=pipeline_session, ) - step = LineageStep(name="Rec", step_args=actions + [artifact, context]) + step = LineageStep(name="Rec", actions=actions, artifacts=artifact, contexts=context) args = step.arguments assert [a["ActionName"] for a in args["Actions"]] == ["act0", "act1"] @@ -342,7 +342,8 @@ def test_lineage_step_association_references_sibling_by_name_and_type(pipeline_s ) step = LineageStep( name="Rec", - step_args=[action, artifact], + actions=action, + artifacts=artifact, associations=[ LineageAssociation( source=LineageEntityReference(name="act1", type="Action"), @@ -364,7 +365,7 @@ def test_lineage_step_association_accepts_literal_arn(pipeline_session, action_s existing = "arn:aws:sagemaker:us-west-2:123456789012:artifact/abc" step = LineageStep( name="Rec", - step_args=action_step_args, + actions=action_step_args, associations=[ LineageAssociation( source=LineageEntityReference(name="act1", type="Action"), @@ -385,7 +386,7 @@ def test_lineage_step_rejects_unresolvable_sibling_reference(action_step_args): with pytest.raises(ValueError, match="does not create"): LineageStep( name="Rec", - step_args=action_step_args, + actions=action_step_args, associations=[ LineageAssociation( source=LineageEntityReference(name="act1", type="Action"), @@ -396,7 +397,7 @@ def test_lineage_step_rejects_unresolvable_sibling_reference(action_step_args): def test_lineage_step_rejects_captured_association(pipeline_session, action_step_args): - """Association.create() cannot express a sibling, so it is not valid step_args.""" + """Association.create() cannot express a sibling, so it is not a valid entity.""" captured = Association.create( source_arn="arn:aws:sagemaker:us-west-2:123456789012:action/a", destination_arn="arn:aws:sagemaker:us-west-2:123456789012:artifact/b", @@ -404,7 +405,7 @@ def test_lineage_step_rejects_captured_association(pipeline_session, action_step sagemaker_session=pipeline_session, ) with pytest.raises(ValueError, match="associations argument"): - LineageStep(name="Rec", step_args=[action_step_args, captured]) + LineageStep(name="Rec", actions=[action_step_args, captured]) def test_lineage_step_requires_an_entity_or_association(): @@ -431,16 +432,35 @@ def test_lineage_association_rejects_non_reference_endpoint(): 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) + LineageStep(name="Rec", actions=endpoint_step_args) + + +def test_lineage_step_rejects_entity_in_the_wrong_argument(pipeline_session): + """An artifact passed as an action is rejected, naming the right producer. + + The per-kind arguments make the expected producer unambiguous, so a + misplaced entity fails at construction instead of building a request the + service would reject. + """ + artifact = Artifact.create( + artifact_name="art1", + source_uri="s3://bucket/model.tar.gz", + artifact_type="Model", + sagemaker_session=pipeline_session, + ) + with pytest.raises(ValueError, match="actions of LineageStep must be obtained from"): + LineageStep(name="Rec", actions=artifact) + with pytest.raises(ValueError, match="contexts of LineageStep must be obtained from"): + LineageStep(name="Rec", contexts=artifact) def test_lineage_step_rejects_raw_dict(): with pytest.raises(TypeError): - LineageStep(name="Rec", step_args={"Actions": []}) + LineageStep(name="Rec", actions={"ActionName": "a"}) def test_lineage_step_properties(action_step_args): - step = LineageStep(name="Rec", step_args=action_step_args) + step = LineageStep(name="Rec", actions=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']"} From a1e7add94331a8d3392c8786eb123810f5c0fecd Mon Sep 17 00:00:00 2001 From: Rishabh Devnani Date: Wed, 16 Sep 2026 22:16:45 +0000 Subject: [PATCH 3/3] fix(tests): Use the shared session fixtures in the new integ tests Both new integ tests defined local sagemaker_session and pipeline_session fixtures returning a bare Session() / PipelineSession(), shadowing the shared ones in tests/integ/conftest.py. CI runs the integ suite with AWS_DEFAULT_REGION unset, so a bare session has no region and both tests error at setup with "Must setup local AWS configuration with a region supported by SageMaker". The local role fixture went too, since the shared conftest provides one. The shared fixtures build on a boto3.Session carrying an explicit region, which is exactly what was missing. Found on the v2 equivalent of these tests, where CI reported it; the same shadowing was present here and would have failed the same way once these tests ran. Testing: both integ tests re-run against the service using the shared fixtures -- lineage in 15.62s with AWS_DEFAULT_REGION unset, and the chained EndpointConfig -> Endpoint -> InferenceComponent pipeline in 413.44s -- with no resources left behind. 327 workflow unit tests still pass. --- X-AI-Prompt: did you read the comment on the pr? X-AI-Tool: kiro-cli --- .../integ/workflow/test_deployment_steps.py | 19 ++++--------------- .../tests/integ/workflow/test_lineage_step.py | 19 ++++--------------- 2 files changed, 8 insertions(+), 30 deletions(-) diff --git a/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py index e0807ddaf8..53da057928 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py +++ b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py @@ -31,8 +31,6 @@ 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 @@ -42,19 +40,10 @@ 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() +# sagemaker_session, pipeline_session and role come from tests/integ/conftest.py. +# They build their sessions on a boto3.Session carrying an explicit region, which a +# bare Session() would not have: CI runs the integ suite with AWS_DEFAULT_REGION +# unset. def test_deployment_steps_execute_end_to_end(sagemaker_session, pipeline_session, role): diff --git a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py index b554b182cc..bd9050199e 100644 --- a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -36,11 +36,9 @@ import pytest -from sagemaker.core.helper.session_helper import Session, get_execution_role from sagemaker.core.lineage.action import Action from sagemaker.core.lineage.artifact import Artifact from sagemaker.core.lineage.context import Context -from sagemaker.core.workflow.pipeline_context import PipelineSession from sagemaker.mlops.workflow.lineage_step import ( LineageAssociation, LineageEntityReference, @@ -49,19 +47,10 @@ 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() +# sagemaker_session, pipeline_session and role come from tests/integ/conftest.py. +# They build their sessions on a boto3.Session carrying an explicit region, which a +# bare Session() would not have: CI runs the integ suite with AWS_DEFAULT_REGION +# unset. def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role):