diff --git a/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py index 3b762be826..5b5de4bd10 100644 --- a/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py +++ b/sagemaker-core/src/sagemaker/core/apiutils/_base_types.py @@ -219,10 +219,29 @@ def with_boto(self, boto_dict): ) return self + # Lineage entity creation methods whose requests are captured as pipeline + # step arguments when invoked under a ``PipelineSession`` (instead of + # calling the service). Used by ``sagemaker.mlops.workflow.LineageStep``. + _PIPELINE_CAPTURABLE_METHODS = frozenset( + {"create_action", "create_artifact", "create_context", "add_association"} + ) + def _invoke_api(self, boto_method, boto_method_members): """Invoke a SageMaker API.""" api_values = {k: v for k, v in vars(self).items() if k in boto_method_members} api_kwargs = self.to_boto(api_values) + + if boto_method in self._PIPELINE_CAPTURABLE_METHODS: + # Lazy import to avoid a circular dependency at module load time. + from sagemaker.core.workflow.pipeline_context import ( + PipelineSession, + _JobStepArguments, + ) + + if isinstance(self.sagemaker_session, PipelineSession): + self.sagemaker_session.context = _JobStepArguments(boto_method, api_kwargs) + return self.sagemaker_session.context + api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method) api_boto_response = api_method(**api_kwargs) return self.with_boto(api_boto_response) diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index ecdd4b95eb..ab55cfa4e9 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -1091,15 +1091,23 @@ def endpoint_from_production_variants( config_options["ExecutionRoleArn"] = role logger.info("Creating endpoint-config with name %s", name) - self.sagemaker_client.create_endpoint_config(**config_options) - - return self.create_endpoint( - endpoint_name=name, - config_name=name, - tags=endpoint_tags, - wait=wait, - live_logging=live_logging, + + def submit(request): + self.sagemaker_client.create_endpoint_config(**request) + return self.create_endpoint( + endpoint_name=name, + config_name=name, + tags=endpoint_tags, + wait=wait, + live_logging=live_logging, + ) + + result = self._intercept_create_request( + config_options, submit, self.endpoint_from_production_variants.__name__ ) + if self._is_pipeline_context(): + return self.context + return result def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False): """Create an Amazon SageMaker ``Endpoint`` according to the configuration in the request. @@ -1129,16 +1137,28 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live tags = self._append_sagemaker_config_tags( tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT, TAGS) ) - try: - res = self.sagemaker_client.create_endpoint( - EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags - ) + create_endpoint_request = { + "EndpointName": endpoint_name, + "EndpointConfigName": config_name, + "Tags": tags, + } + + def submit(request): + res = self.sagemaker_client.create_endpoint(**request) if res: self.endpoint_arn = res["EndpointArn"] if wait: self.wait_for_endpoint(endpoint_name, live_logging=live_logging) return endpoint_name + + try: + result = self._intercept_create_request( + create_endpoint_request, submit, self.create_endpoint.__name__ + ) + if self._is_pipeline_context(): + return self.context + return result except Exception as e: troubleshooting = ( "https://docs.aws.amazon.com/sagemaker/latest/dg/" @@ -1257,10 +1277,18 @@ def create_inference_component( if tags and len(tags) != 0: request["Tags"] = tags - self.sagemaker_client.create_inference_component(**request) - if wait: - self.wait_for_inference_component(inference_component_name) - return inference_component_name + def submit(req): + self.sagemaker_client.create_inference_component(**req) + if wait: + self.wait_for_inference_component(inference_component_name) + return inference_component_name + + result = self._intercept_create_request( + request, submit, self.create_inference_component.__name__ + ) + if self._is_pipeline_context(): + return self.context + return result def wait_for_inference_component(self, inference_component_name, poll=20): """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete. @@ -1439,6 +1467,19 @@ def _intercept_create_request( """ return create(request) + def _is_pipeline_context(self) -> bool: + """Whether this session is a pipeline session capturing requests. + + Producer methods that support composing pipeline steps use this to + return the captured step arguments (``self.context``) instead of the + result of a service call. Always ``False`` for a plain ``Session``. + """ + # Lazy import to avoid a circular dependency: pipeline_context imports + # from this module at import time. + from sagemaker.core.workflow.pipeline_context import PipelineSession + + return isinstance(self, PipelineSession) + def _create_inference_recommendations_job_request( self, role: str, diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 129abb1c76..a5ab5ba851 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -14,6 +14,7 @@ functions, conditions, properties) and can import from sagemaker.train and sagemaker.serve for orchestration purposes. """ + from __future__ import absolute_import __version__ = "0.1.0" @@ -46,8 +47,11 @@ from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep from sagemaker.mlops.workflow.condition_step import ConditionStep from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep from sagemaker.mlops.workflow.fail_step import FailStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput +from sagemaker.mlops.workflow.lineage_step import LineageStep from sagemaker.mlops.workflow.model_step import ModelStep from sagemaker.mlops.workflow.monitor_batch_transform_step import MonitorBatchTransformStep from sagemaker.mlops.workflow.notebook_job_step import NotebookJobStep @@ -98,9 +102,13 @@ "ConditionStep", "EMRStep", "EMRStepConfig", + "EndpointConfigStep", + "EndpointStep", "FailStep", + "InferenceComponentStep", "LambdaStep", "LambdaOutput", + "LineageStep", "ModelStep", "MonitorBatchTransformStep", "NotebookJobStep", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py new file mode 100644 index 0000000000..9d21af441b --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/endpoint_step.py @@ -0,0 +1,203 @@ +# 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.endpoint_from_production_variants( + name="my-endpoint-config", + production_variants=[...], + ) + config_step = EndpointConfigStep(name="CreateConfig", step_args=config_step_args) + + endpoint_step_args = pipeline_session.create_endpoint( + endpoint_name="my-endpoint", + config_name="my-endpoint-config", + ) + endpoint_step = EndpointStep(name="CreateEndpoint", step_args=endpoint_step_args) +""" + +from __future__ import absolute_import + +from typing import 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.endpoint_from_production_variants` + on a ``PipelineSession``. + + ``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.endpoint_from_production_variants()``. + 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={"endpoint_from_production_variants"}, + error_message=( + "The step_args of EndpointConfigStep must be obtained from " + "pipeline_session.endpoint_from_production_variants()." + ), + ) + 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..8ea5b55c3b --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/lineage_step.py @@ -0,0 +1,166 @@ +# 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. + +Follows the ``step_args`` convention: create a lineage entity with the +corresponding class from :mod:`sagemaker.core.lineage` under a +:class:`~sagemaker.core.workflow.pipeline_context.PipelineSession` and +pass the returned step arguments to the step. Each step creates one +lineage entity; associations between entities created in different +steps reference them by ARN via step property references. + +Example:: + + pipeline_session = PipelineSession() + + action_args = Action.create( + action_name="my-action", + source_uri="s3://bucket/model.tar.gz", + source_type="S3ETag", + action_type="ModelTraining", + status="Completed", + sagemaker_session=pipeline_session, + ) + action_step = LineageStep(name="RecordAction", step_args=action_args) + + association_args = Association.create( + source_arn=action_step.properties.ActionArns["my-action"], + destination_arn="arn:aws:sagemaker:...:artifact/abc", + association_type="Produced", + sagemaker_session=pipeline_session, + ) + association_step = LineageStep( + name="RecordAssociation", + step_args=association_args, + depends_on=[action_step], + ) +""" + +from __future__ import absolute_import + +from typing import List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import RequestType +from sagemaker.core.workflow.pipeline_context import _JobStepArguments +from sagemaker.core.workflow.properties import Properties +from sagemaker.core.workflow.utilities import validate_step_args_input + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + +# Maps the captured lineage create call to the Arguments key the pipeline +# service expects. +_CALLER_TO_ARGUMENTS_KEY = { + "create_action": "Actions", + "create_artifact": "Artifacts", + "create_context": "Contexts", + "add_association": "Associations", +} + + +class _EntityArnMap(Properties): + """Map-style property access for service-native ARN maps. + + The lineage step's ``ActionArns``/``ArtifactArns``/``ContextArns`` + outputs are maps keyed by entity name. They are pipeline-service + outputs with no botocore shape, so this supports ``['name']`` access + without a shape lookup. + """ + + def __getitem__(self, item: str) -> Properties: + """Reference the ARN of the entity created under the given name.""" + return Properties(step_name=self.step_name, path=f"{self.path}['{item}']") + + +class LineageStep(Step): + """Creates a lineage entity or association in SageMaker's lineage system. + + Wraps SageMaker's ``CreateAction``/``CreateArtifact``/``CreateContext`` + and lineage ``AddAssociation`` APIs. Each step creates one entity; + the ``step_args`` must be obtained by calling ``Action.create()``, + ``Artifact.create()``, ``Context.create()``, or + ``Association.create()`` from :mod:`sagemaker.core.lineage` with a + ``PipelineSession``. + + Property references expose the created entity: + ``Steps..ActionArns['']``, + ``Steps..ArtifactArns['']``, + ``Steps..ContextArns['']``, and + ``Steps..Associations``. + """ + + 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 a ``LineageStep``. + + Args: + name (str): The name of the step. + step_args (_JobStepArguments): The arguments for this step, + obtained from ``Action.create()``, ``Artifact.create()``, + ``Context.create()``, or ``Association.create()`` called + with a ``PipelineSession``. + display_name (str): Optional display name. + description (str): Optional description. + depends_on (List[Union[str, Step, StepCollection]]): Optional + explicit step dependencies. + """ + super().__init__( + name=name, + display_name=display_name, + description=description, + step_type=StepTypeEnum.LINEAGE, + depends_on=depends_on, + ) + validate_step_args_input( + step_args=step_args, + expected_caller=set(_CALLER_TO_ARGUMENTS_KEY), + error_message=( + "The step_args of LineageStep must be obtained from " + "Action.create(), Artifact.create(), Context.create(), or " + "Association.create() called with a PipelineSession." + ), + ) + self.step_args = step_args + + root = Properties(step_name=name, step=self) + for field in ("ActionArns", "ArtifactArns", "ContextArns"): + root.__dict__[field] = _EntityArnMap(step_name=name, path=field) + root.__dict__["Associations"] = Properties(step_name=name, path="Associations") + self._properties = root + + @property + def arguments(self) -> RequestType: + """The ``Arguments`` block describing the lineage entity to create.""" + key = _CALLER_TO_ARGUMENTS_KEY[self.step_args.caller_name] + entity = self.step_args.args + if key == "Associations": + # The AddAssociation API uses SourceArn/DestinationArn; the + # pipeline service models associations as entity references. + entity = { + "Source": {"Arn": entity["SourceArn"]}, + "Destination": {"Arn": entity["DestinationArn"]}, + } + if "AssociationType" in self.step_args.args: + entity["AssociationType"] = self.step_args.args["AssociationType"] + return {key: [entity]} + + @property + def properties(self): + """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..142252b9b8 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_deployment_steps.py @@ -0,0 +1,172 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integration test for the inference deployment step types. + +Runs a single pipeline chaining ``EndpointConfigStep`` -> +``EndpointStep`` -> ``InferenceComponentStep`` end-to-end against the +real service: an inference-component-style endpoint config (no model +name on the variant, execution role on the config), an endpoint, and an +inference component carrying the container specification. + +This test provisions a real endpoint instance for its duration; all +resources are deleted in the ``finally`` block. +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker.core import image_uris +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep +from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep +from sagemaker.mlops.workflow.pipeline import Pipeline + +INSTANCE_TYPE = "ml.m5.xlarge" +EXECUTION_TIMEOUT_SECONDS = 45 * 60 +POLL_SECONDS = 30 + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def test_deployment_steps_execute_end_to_end(sagemaker_session, pipeline_session, role): + """Chained EndpointConfig -> Endpoint -> InferenceComponent pipeline run.""" + stamp = uuid.uuid4().hex[:8] + config_name = f"integ-deploy-cfg-{stamp}" + endpoint_name = f"integ-deploy-ep-{stamp}" + component_name = f"integ-deploy-ic-{stamp}" + pipeline_name = f"integ-deploy-{stamp}" + + image_uri = image_uris.retrieve( + framework="sklearn", + region=sagemaker_session.boto_region_name, + version="1.2-1", + instance_type=INSTANCE_TYPE, + ) + + config_step_args = pipeline_session.endpoint_from_production_variants( + name=config_name, + production_variants=[ + { + "VariantName": "AllTraffic", + "InstanceType": INSTANCE_TYPE, + "InitialInstanceCount": 1, + "ManagedInstanceScaling": { + "Status": "ENABLED", + "MinInstanceCount": 1, + "MaxInstanceCount": 1, + }, + "RoutingConfig": {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"}, + } + ], + role=role, + ) + config_step = EndpointConfigStep(name="CreateConfig", step_args=config_step_args) + + endpoint_step_args = pipeline_session.create_endpoint( + endpoint_name=endpoint_name, config_name=config_name + ) + endpoint_step = EndpointStep( + name="CreateEndpoint", step_args=endpoint_step_args, depends_on=[config_step] + ) + + component_step_args = pipeline_session.create_inference_component( + inference_component_name=component_name, + endpoint_name=endpoint_name, + variant_name="AllTraffic", + specification={ + "Container": {"Image": image_uri}, + "ComputeResourceRequirements": { + "NumberOfCpuCoresRequired": 1.0, + "MinMemoryRequiredInMb": 1024, + }, + }, + runtime_config={"CopyCount": 1}, + ) + component_step = InferenceComponentStep( + name="CreateComponent", step_args=component_step_args, depends_on=[endpoint_step] + ) + + pipeline = Pipeline( + name=pipeline_name, + steps=[config_step, endpoint_step, component_step], + sagemaker_session=pipeline_session, + ) + + sm_client = sagemaker_session.sagemaker_client + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + deadline = time.time() + EXECUTION_TIMEOUT_SECONDS + status = None + while time.time() < deadline: + status = execution.describe()["PipelineExecutionStatus"] + if status not in ("Executing", "Stopping"): + break + time.sleep(POLL_SECONDS) + + assert status == "Succeeded", f"Pipeline execution ended in status {status}" + + # Server-side resources exist with the expected names. + config_desc = sm_client.describe_endpoint_config(EndpointConfigName=config_name) + assert config_desc["EndpointConfigName"] == config_name + + endpoint_desc = sm_client.describe_endpoint(EndpointName=endpoint_name) + assert endpoint_desc["EndpointName"] == endpoint_name + + component_desc = sm_client.describe_inference_component( + InferenceComponentName=component_name + ) + assert component_desc["InferenceComponentName"] == component_name + assert component_desc["EndpointName"] == endpoint_name + finally: + for cleanup in ( + lambda: sm_client.delete_inference_component(InferenceComponentName=component_name), + lambda: _wait_component_deleted(sm_client, component_name), + lambda: sm_client.delete_endpoint(EndpointName=endpoint_name), + lambda: sm_client.delete_endpoint_config(EndpointConfigName=config_name), + lambda: pipeline.delete(), + ): + try: + cleanup() + except Exception: # noqa: BLE001 -- best-effort cleanup + pass + + +def _wait_component_deleted(sm_client, component_name, timeout_seconds=10 * 60): + """The endpoint cannot be deleted until its inference component is gone.""" + deadline = time.time() + timeout_seconds + while time.time() < deadline: + try: + sm_client.describe_inference_component(InferenceComponentName=component_name) + except Exception: + return + time.sleep(15) diff --git a/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py new file mode 100644 index 0000000000..b96f86bb47 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_lineage_step.py @@ -0,0 +1,139 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integration test for the LineageStep. + +Creates a pipeline containing a single ``LineageStep`` that records a +SageMaker lineage Action, executes it end-to-end against the real +service, and asserts the execution reaches ``Succeeded``. Cleans up the +Action, the pipeline, and the S3 pipeline definition artifact. + +Requires the execution role to have ``sagemaker:CreateAction`` (and +related lineage permissions). ``SageMakerRole`` — the standard fixture +role used across the SDK's integ tests — has broad SageMaker access and +satisfies this requirement. + +This test represents the SDK-side end-to-end validation of the +inference and lineage step family. See ``endpoint_step.py`` and +``inference_component_step.py`` for the other step types; those are not +integ-tested here because they provision paid resources +(Endpoint/InferenceComponent). +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest + +from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.lineage.action import Action +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.mlops.workflow.lineage_step import LineageStep +from sagemaker.mlops.workflow.pipeline import Pipeline + + +@pytest.fixture +def sagemaker_session(): + return Session() + + +@pytest.fixture +def pipeline_session(): + return PipelineSession() + + +@pytest.fixture +def role(): + return get_execution_role() + + +def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role): + """Full end-to-end run of a LineageStep pipeline against the real service. + + Builds a pipeline with a single ``LineageStep`` that creates one + lineage ``Action``. Verifies the pipeline execution succeeds and the + server-reported step metadata contains the created action ARN. + """ + stamp = uuid.uuid4().hex[:8] + action_name = f"lineage-integ-{stamp}" + pipeline_name = f"integ-lineage-{stamp}" + + step_args = Action.create( + action_name=action_name, + source_uri=f"s3://lineage-integ-test/{stamp}/model.tar.gz", + source_type="MODEL", + action_type="ModelTraining", + status="Completed", + description="Lineage integ test action", + sagemaker_session=pipeline_session, + ) + step = LineageStep(name="RecordLineage", step_args=step_args) + pipeline = Pipeline( + name=pipeline_name, + steps=[step], + sagemaker_session=pipeline_session, + ) + + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + # LineageStep is metadata-only; execution completes quickly. Poll + # up to 5 minutes to give the service plenty of headroom under load. + timeout = 300 + start_time = time.time() + final_status = None + while time.time() - start_time < timeout: + execution_desc = execution.describe() + status = execution_desc["PipelineExecutionStatus"] + if status in ("Succeeded", "Failed", "Stopped"): + final_status = status + break + time.sleep(10) + + if final_status != "Succeeded": + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + failure_details = "\n".join( + f"{s['StepName']}: {s.get('FailureReason', 'no reason')}" + for s in steps + if s.get("StepStatus") == "Failed" + ) + pytest.fail(f"Pipeline execution status={final_status}. Details:\n{failure_details}") + + # Verify the step metadata reports the created action ARN. + steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps( + PipelineExecutionArn=execution.arn, + )["PipelineExecutionSteps"] + lineage_step = next(s for s in steps if s["StepName"] == "RecordLineage") + assert lineage_step["StepStatus"] == "Succeeded" + metadata = lineage_step.get("Metadata", {}) + action_arns = metadata.get("Lineage", {}).get("ActionArns", {}) + assert ( + action_name in action_arns + ), f"expected {action_name} in ActionArns, got: {action_arns}" + assert action_arns[action_name].endswith(f":action/{action_name}") + + finally: + # Delete the lineage Action. + try: + sagemaker_session.sagemaker_client.delete_action(ActionName=action_name) + except Exception: + pass + # Delete the pipeline. + try: + sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=pipeline_name) + except Exception: + pass diff --git a/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py new file mode 100644 index 0000000000..c4f94272eb --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_inference_lineage_steps.py @@ -0,0 +1,361 @@ +# 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 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.endpoint_from_production_variants( + name="my-config", + production_variants=[ + { + "ModelName": "my-model", + "VariantName": "AllTraffic", + "InstanceType": "ml.m5.large", + "InitialInstanceCount": 1, + } + ], + kms_key="arn:aws:kms:us-west-2:123456789012:key/abc", + ) + + +@pytest.fixture +def endpoint_step_args(pipeline_session): + return pipeline_session.create_endpoint(endpoint_name="my-endpoint", config_name="my-config") + + +@pytest.fixture +def inference_component_step_args(pipeline_session): + return pipeline_session.create_inference_component( + inference_component_name="my-component", + endpoint_name="my-endpoint", + variant_name="AllTraffic", + specification={"ModelName": "my-model"}, + runtime_config={"CopyCount": 2}, + ) + + +# ---------- step_args capture via PipelineSession ---------- + + +def test_capture_does_not_call_service(pipeline_session, endpoint_config_step_args): + assert isinstance(endpoint_config_step_args, _JobStepArguments) + assert not pipeline_session.sagemaker_client.create_endpoint_config.called + assert not pipeline_session.sagemaker_client.create_endpoint.called + + +def test_captured_request_content(endpoint_config_step_args): + args = endpoint_config_step_args.args + assert args["EndpointConfigName"] == "my-config" + assert args["KmsKeyId"] == "arn:aws:kms:us-west-2:123456789012:key/abc" + assert args["ProductionVariants"][0]["ModelName"] == "my-model" + + +# ---------- EndpointConfigStep ---------- + + +def test_endpoint_config_step_basic(endpoint_config_step_args): + step = EndpointConfigStep(name="Cfg", step_args=endpoint_config_step_args) + assert step.step_type == StepTypeEnum.ENDPOINT_CONFIG + assert step.arguments["EndpointConfigName"] == "my-config" + req = step.to_request() + assert req["Type"] == "EndpointConfig" + assert req["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="endpoint_from_production_variants"): + EndpointConfigStep(name="Cfg", step_args=endpoint_step_args) + + +def test_endpoint_config_step_rejects_raw_dict(): + with pytest.raises(TypeError): + EndpointConfigStep(name="Cfg", step_args={"EndpointConfigName": "x"}) + + +def test_endpoint_config_step_properties(endpoint_config_step_args): + step = EndpointConfigStep(name="Cfg", step_args=endpoint_config_step_args) + assert step.properties.EndpointConfigName.expr == {"Get": "Steps.Cfg.EndpointConfigName"} + + +# ---------- EndpointStep ---------- + + +def test_endpoint_step_basic(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_association_translates_arns(pipeline_session, action_step_args): + action_step = LineageStep(name="RecA", step_args=action_step_args) + step_args = Association.create( + source_arn=action_step.properties.ActionArns["act1"], + destination_arn="arn:aws:sagemaker:us-west-2:123456789012:artifact/abc", + association_type="Produced", + sagemaker_session=pipeline_session, + ) + step = LineageStep(name="RecD", step_args=step_args, depends_on=[action_step]) + entity = step.arguments["Associations"][0] + # AddAssociation's SourceArn/DestinationArn become entity references. + assert entity["Source"]["Arn"].expr == {"Get": "Steps.RecA.ActionArns['act1']"} + assert entity["Destination"] == {"Arn": "arn:aws:sagemaker:us-west-2:123456789012:artifact/abc"} + assert entity["AssociationType"] == "Produced" + assert not pipeline_session.sagemaker_client.add_association.called + + +def test_lineage_step_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_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"]