feat(pipeline): Add inference and lineage step types

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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

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)

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,

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",

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