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..20c567f64f 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -1090,15 +1090,19 @@ 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 self.create_endpoint( - endpoint_name=name, - config_name=name, - tags=endpoint_tags, - wait=wait, - live_logging=live_logging, + def submit(request): + logger.info("Creating endpoint-config with name %s", name) + 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, + ) + + return self._intercept_create_request( + config_options, submit, self.endpoint_from_production_variants.__name__ ) def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False): @@ -1122,23 +1126,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 +1244,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 +1263,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..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..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..bbcb85cb69 --- /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.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) + + # 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..7cc23e2b37 --- /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.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_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.endpoint_from_production_variants( + 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 == [ + "endpoint_from_production_variants", + "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"]