feat(pipeline): Add inference and lineage step types

sagemaker-core/src/sagemaker/core/apiutils/_base_types.py

@@ -219,10 +219,27 @@ def with_boto(self, boto_dict):
         )
         return self
 
+    # Lineage entity creation methods that participate in pipeline step
+    # composition. Requests for these are routed through the session's
+    # ``_intercept_create_request`` seam, so a ``PipelineSession`` captures them
+    # as step arguments while a plain ``Session`` calls the service. Used by
+    # ``sagemaker.mlops.workflow.LineageStep``.
+    _PIPELINE_CAPTURABLE_METHODS = frozenset(
+        {"create_action", "create_artifact", "create_context", "add_association"}
+    )
+
     def _invoke_api(self, boto_method, boto_method_members):
         """Invoke a SageMaker API."""
         api_values = {k: v for k, v in vars(self).items() if k in boto_method_members}
         api_kwargs = self.to_boto(api_values)
-        api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method)
-        api_boto_response = api_method(**api_kwargs)
-        return self.with_boto(api_boto_response)
+
+        def submit(request):
+            api_method = getattr(self.sagemaker_session.sagemaker_client, boto_method)
+            return self.with_boto(api_method(**request))
+
+        if boto_method in self._PIPELINE_CAPTURABLE_METHODS:
+            return self.sagemaker_session._intercept_create_request(
+                api_kwargs, submit, boto_method
+            )
+
+        return submit(api_kwargs)

sagemaker-core/src/sagemaker/core/helper/session_helper.py

@@ -1090,15 +1090,19 @@ def endpoint_from_production_variants(
         if role is not None:
             config_options["ExecutionRoleArn"] = role
 
-        logger.info("Creating endpoint-config with name %s", name)
-        self.sagemaker_client.create_endpoint_config(**config_options)
-
-        return self.create_endpoint(
-            endpoint_name=name,
-            config_name=name,
-            tags=endpoint_tags,
-            wait=wait,
-            live_logging=live_logging,
+        def submit(request):
+            logger.info("Creating endpoint-config with name %s", name)
+            self.sagemaker_client.create_endpoint_config(**request)
+            return self.create_endpoint(
+                endpoint_name=name,
+                config_name=name,
+                tags=endpoint_tags,
+                wait=wait,
+                live_logging=live_logging,
+            )
+
+        return self._intercept_create_request(
+            config_options, submit, self.endpoint_from_production_variants.__name__
         )
 
     def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False):
@@ -1122,23 +1126,31 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live
             botocore.exceptions.ClientError: If Sagemaker throws an exception while creating
             endpoint.
         """
-        logger.info("Creating endpoint with name %s", endpoint_name)
-
         tags = format_tags(tags) or []
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(
             tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT, TAGS)
         )
-        try:
-            res = self.sagemaker_client.create_endpoint(
-                EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags
-            )
+        create_endpoint_request = {
+            "EndpointName": endpoint_name,
+            "EndpointConfigName": config_name,
+            "Tags": tags,
+        }
+
+        def submit(request):
+            logger.info("Creating endpoint with name %s", endpoint_name)
+            res = self.sagemaker_client.create_endpoint(**request)
             if res:
                 self.endpoint_arn = res["EndpointArn"]
 
             if wait:
                 self.wait_for_endpoint(endpoint_name, live_logging=live_logging)
             return endpoint_name
+
+        try:
+            return self._intercept_create_request(
+                create_endpoint_request, submit, self.create_endpoint.__name__
+            )
         except Exception as e:
             troubleshooting = (
                 "https://docs.aws.amazon.com/sagemaker/latest/dg/"
@@ -1232,12 +1244,6 @@ def create_inference_component(
         Returns:
             str: Name of the Amazon SageMaker ``InferenceComponent`` if created.
         """
-        LOGGER.info(
-            "Creating inference component with name %s for endpoint %s",
-            inference_component_name,
-            endpoint_name,
-        )
-
         if runtime_config is None:
             runtime_config = {"CopyCount": 1}
 
@@ -1257,10 +1263,20 @@ def create_inference_component(
         if tags and len(tags) != 0:
             request["Tags"] = tags
 
-        self.sagemaker_client.create_inference_component(**request)
-        if wait:
-            self.wait_for_inference_component(inference_component_name)
-        return inference_component_name
+        def submit(req):
+            LOGGER.info(
+                "Creating inference component with name %s for endpoint %s",
+                inference_component_name,
+                endpoint_name,
+            )
+            self.sagemaker_client.create_inference_component(**req)
+            if wait:
+                self.wait_for_inference_component(inference_component_name)
+            return inference_component_name
+
+        return self._intercept_create_request(
+            request, submit, self.create_inference_component.__name__
+        )
 
     def wait_for_inference_component(self, inference_component_name, poll=20):
         """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete.

sagemaker-core/src/sagemaker/core/workflow/pipeline_context.py

@@ -220,6 +220,13 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None
             request (dict): the create job request
             create (functor): a functor calls the sagemaker client create method
             func_name (str): the name of the function needed intercepting
+
+        Returns:
+            The captured pipeline context, so that a producer method can simply
+            return the result of this call: under a plain ``Session`` the base
+            implementation returns the service result, and here it returns the
+            step arguments instead. This keeps the pipeline branch inside
+            ``PipelineSession`` rather than in the base ``Session``.
         """
         if func_name == "create_model":
             self.context.create_model_request = request
@@ -229,6 +236,7 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None
             self.context.caller_name = func_name
         else:
             self.context = _JobStepArguments(func_name, request)
+        return self.context
 
     def init_model_step_arguments(self, model):
         """Create a `_ModelStepArguments` (if not exist) as pipeline context

sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py

@@ -14,6 +14,7 @@
 functions, conditions, properties) and can import from sagemaker.train and sagemaker.serve
 for orchestration purposes.
 """
+
 from __future__ import absolute_import
 
 __version__ = "0.1.0"
@@ -46,8 +47,15 @@
 from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep
 from sagemaker.mlops.workflow.condition_step import ConditionStep
 from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig
+from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep
 from sagemaker.mlops.workflow.fail_step import FailStep
+from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep
 from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput
+from sagemaker.mlops.workflow.lineage_step import (
+    LineageAssociation,
+    LineageEntityReference,
+    LineageStep,
+)
 from sagemaker.mlops.workflow.model_step import ModelStep
 from sagemaker.mlops.workflow.monitor_batch_transform_step import MonitorBatchTransformStep
 from sagemaker.mlops.workflow.notebook_job_step import NotebookJobStep
@@ -98,9 +106,15 @@
     "ConditionStep",
     "EMRStep",
     "EMRStepConfig",
+    "EndpointConfigStep",
+    "EndpointStep",
     "FailStep",
+    "InferenceComponentStep",
     "LambdaStep",
     "LambdaOutput",
+    "LineageAssociation",
+    "LineageEntityReference",
+    "LineageStep",
     "ModelStep",
     "MonitorBatchTransformStep",
     "NotebookJobStep",