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

@@ -1006,6 +1006,52 @@ def endpoint_from_production_variants(
             str: The name of the created ``Endpoint``.
         """
 
+        config_options = self._build_endpoint_config_request(
+            name=name,
+            production_variants=production_variants,
+            tags=tags,
+            kms_key=kms_key,
+            data_capture_config_dict=data_capture_config_dict,
+            async_inference_config_dict=async_inference_config_dict,
+            explainer_config_dict=explainer_config_dict,
+            vpc_config=vpc_config,
+            enable_network_isolation=enable_network_isolation,
+            role=role,
+        )
+        endpoint_tags = _append_project_tags(format_tags(tags))
+
+        logger.info("Creating endpoint-config with name %s", name)
+        self.sagemaker_client.create_endpoint_config(**config_options)
+
+        return self.create_endpoint(
+            endpoint_name=name,
+            config_name=name,
+            tags=endpoint_tags,
+            wait=wait,
+            live_logging=live_logging,
+        )
+
+    def _build_endpoint_config_request(
+        self,
+        name,
+        production_variants,
+        tags=None,
+        kms_key=None,
+        data_capture_config_dict=None,
+        async_inference_config_dict=None,
+        explainer_config_dict=None,
+        vpc_config=None,
+        enable_network_isolation=None,
+        role=None,
+    ):
+        """Build the ``CreateEndpointConfig`` request, applying config defaults.
+
+        Shared by ``create_endpoint_config`` and
+        ``endpoint_from_production_variants`` so the two cannot drift.
+
+        Returns:
+            dict: The ``CreateEndpointConfig`` request.
+        """
         supports_kms = any(
             [
                 instance_supports_kms(production_variant["InstanceType"])
@@ -1060,7 +1106,6 @@ def endpoint_from_production_variants(
             role = self.expand_role(role)
             config_options["ExecutionRoleArn"] = role
         endpoint_config_tags = _append_project_tags(format_tags(tags))
-        endpoint_tags = _append_project_tags(format_tags(tags))
 
         endpoint_config_tags = self._append_sagemaker_config_tags(
             endpoint_config_tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT_CONFIG, TAGS)
@@ -1090,15 +1135,71 @@ def endpoint_from_production_variants(
         if role is not None:
             config_options["ExecutionRoleArn"] = role
 
-        logger.info("Creating endpoint-config with name %s", name)
-        self.sagemaker_client.create_endpoint_config(**config_options)
+        return config_options
 
-        return self.create_endpoint(
-            endpoint_name=name,
-            config_name=name,
-            tags=endpoint_tags,
-            wait=wait,
-            live_logging=live_logging,
+    def create_endpoint_config(
+        self,
+        name,
+        production_variants,
+        tags=None,
+        kms_key=None,
+        data_capture_config_dict=None,
+        async_inference_config_dict=None,
+        explainer_config_dict=None,
+        vpc_config=None,
+        enable_network_isolation=None,
+        role=None,
+    ):
+        """Create an Amazon SageMaker endpoint configuration and nothing else.
+
+        Unlike ``endpoint_from_production_variants``, which creates an endpoint
+        configuration *and* the endpoint that uses it, this creates only the
+        configuration. That makes it the producer for a pipeline
+        ``EndpointConfigStep``, which composes one resource per step.
+
+        Args:
+            name (str): Name of the endpoint configuration to create.
+            production_variants (list[dict[str, str]]): The list of production
+                variants to host at this endpoint.
+            tags (Optional[Tags]): Tags to apply to the endpoint configuration.
+            kms_key (str): KMS key ARN for encrypting the volume attached to the
+                instances hosting the endpoint.
+            data_capture_config_dict (dict): Configuration for capturing
+                inference request and response data.
+            async_inference_config_dict (dict): Configuration for asynchronous
+                inference.
+            explainer_config_dict (dict): Configuration for online explainers.
+            vpc_config (dict): The VpcConfig for the endpoint configuration.
+            enable_network_isolation (bool): Whether to isolate the model
+                containers deployed to the endpoint.
+            role (str): An execution role ARN, required when the production
+                variants carry no model name, as with inference components.
+
+        Returns:
+            str: Name of the endpoint configuration created. Under a
+            ``PipelineSession`` the captured step arguments are returned
+            instead and no service call is made.
+        """
+        config_options = self._build_endpoint_config_request(
+            name=name,
+            production_variants=production_variants,
+            tags=tags,
+            kms_key=kms_key,
+            data_capture_config_dict=data_capture_config_dict,
+            async_inference_config_dict=async_inference_config_dict,
+            explainer_config_dict=explainer_config_dict,
+            vpc_config=vpc_config,
+            enable_network_isolation=enable_network_isolation,
+            role=role,
+        )
+
+        def submit(request):
+            logger.info("Creating endpoint-config with name %s", name)
+            self.sagemaker_client.create_endpoint_config(**request)
+            return name
+
+        return self._intercept_create_request(
+            config_options, submit, self.create_endpoint_config.__name__
         )
 
     def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live_logging=False):
@@ -1122,23 +1223,31 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True, live
             botocore.exceptions.ClientError: If Sagemaker throws an exception while creating
             endpoint.
         """
-        logger.info("Creating endpoint with name %s", endpoint_name)
-
         tags = format_tags(tags) or []
         tags = _append_project_tags(tags)
         tags = self._append_sagemaker_config_tags(
             tags, "{}.{}.{}".format(SAGEMAKER, ENDPOINT, TAGS)
         )
-        try:
-            res = self.sagemaker_client.create_endpoint(
-                EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags
-            )
+        create_endpoint_request = {
+            "EndpointName": endpoint_name,
+            "EndpointConfigName": config_name,
+            "Tags": tags,
+        }
+
+        def submit(request):
+            logger.info("Creating endpoint with name %s", endpoint_name)
+            res = self.sagemaker_client.create_endpoint(**request)
             if res:
                 self.endpoint_arn = res["EndpointArn"]
 
             if wait:
                 self.wait_for_endpoint(endpoint_name, live_logging=live_logging)
             return endpoint_name
+
+        try:
+            return self._intercept_create_request(
+                create_endpoint_request, submit, self.create_endpoint.__name__
+            )
         except Exception as e:
             troubleshooting = (
                 "https://docs.aws.amazon.com/sagemaker/latest/dg/"
@@ -1232,12 +1341,6 @@ def create_inference_component(
         Returns:
             str: Name of the Amazon SageMaker ``InferenceComponent`` if created.
         """
-        LOGGER.info(
-            "Creating inference component with name %s for endpoint %s",
-            inference_component_name,
-            endpoint_name,
-        )
-
         if runtime_config is None:
             runtime_config = {"CopyCount": 1}
 
@@ -1257,10 +1360,20 @@ def create_inference_component(
         if tags and len(tags) != 0:
             request["Tags"] = tags
 
-        self.sagemaker_client.create_inference_component(**request)
-        if wait:
-            self.wait_for_inference_component(inference_component_name)
-        return inference_component_name
+        def submit(req):
+            LOGGER.info(
+                "Creating inference component with name %s for endpoint %s",
+                inference_component_name,
+                endpoint_name,
+            )
+            self.sagemaker_client.create_inference_component(**req)
+            if wait:
+                self.wait_for_inference_component(inference_component_name)
+            return inference_component_name
+
+        return self._intercept_create_request(
+            request, submit, self.create_inference_component.__name__
+        )
 
     def wait_for_inference_component(self, inference_component_name, poll=20):
         """Wait for an Amazon SageMaker ``Inference Component`` deployment to complete.

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