feat: Selective Step Execution feature in Pipelines (#3898)

Co-authored-by: Shegufta Ahsan <shegufta@amazon.com>
Co-authored-by: Nishant Krishnamoorthy <kv.nishant87@gmail.com>

Author: 3 people

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -110,6 +110,12 @@ Pipeline Experiment Config
 
 .. autoclass:: sagemaker.workflow.pipeline_experiment_config.PipelineExperimentConfigProperty
 
+Selective Execution Config
+--------------------------
+
+.. autoclass:: sagemaker.workflow.selective_execution_config.SelectiveExecutionConfig
+
+
 Properties
 ----------
 

src/sagemaker/local/local_session.py

@@ -505,6 +505,8 @@ def start_pipeline_execution(self, PipelineName, **kwargs):
         """
         if "ParallelismConfiguration" in kwargs:
             logger.warning("Parallelism configuration is not supported in local mode.")
+        if "SelectiveExecutionConfig" in kwargs:
+            raise ValueError("SelectiveExecutionConfig is not supported in local mode.")
         if PipelineName not in LocalSagemakerClient._pipelines:
             error_response = {
                 "Error": {

src/sagemaker/workflow/pipeline.py

@@ -40,6 +40,7 @@
 from sagemaker.workflow.pipeline_experiment_config import PipelineExperimentConfig
 from sagemaker.workflow.parallelism_config import ParallelismConfiguration
 from sagemaker.workflow.properties import Properties
+from sagemaker.workflow.selective_execution_config import SelectiveExecutionConfig
 from sagemaker.workflow.steps import Step, StepTypeEnum
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.condition_step import ConditionStep
@@ -312,6 +313,7 @@ def start(
         execution_display_name: str = None,
         execution_description: str = None,
         parallelism_config: ParallelismConfiguration = None,
+        selective_execution_config: SelectiveExecutionConfig = None,
     ):
         """Starts a Pipeline execution in the Workflow service.
 
@@ -323,16 +325,26 @@ def start(
             parallelism_config (Optional[ParallelismConfiguration]): Parallelism configuration
                 that is applied to each of the executions of the pipeline. It takes precedence
                 over the parallelism configuration of the parent pipeline.
+            selective_execution_config (Optional[SelectiveExecutionConfig]): The configuration for
+                selective step execution.
 
         Returns:
             A `_PipelineExecution` instance, if successful.
         """
+        if selective_execution_config is not None:
+            if selective_execution_config.source_pipeline_execution_arn is None:
+                selective_execution_config.source_pipeline_execution_arn = (
+                    self._get_latest_execution_arn()
+                )
+            selective_execution_config = selective_execution_config.to_request()
+
         kwargs = dict(PipelineName=self.name)
         update_args(
             kwargs,
             PipelineExecutionDescription=execution_description,
             PipelineExecutionDisplayName=execution_display_name,
             ParallelismConfiguration=parallelism_config,
+            SelectiveExecutionConfig=selective_execution_config,
         )
         if self.sagemaker_session.local_mode:
             update_args(kwargs, PipelineParameters=parameters)
@@ -388,6 +400,57 @@ def _interpolate_step_collection_name_in_depends_on(self, step_requests: list):
                 )
                 self._interpolate_step_collection_name_in_depends_on(sub_step_requests)
 
+    def list_executions(
+        self,
+        sort_by: str = None,
+        sort_order: str = None,
+        max_results: int = None,
+        next_token: str = None,
+    ) -> Dict[str, Any]:
+        """Lists a pipeline's executions.
+
+        Args:
+            sort_by (str): The field by which to sort results(CreationTime/PipelineExecutionArn).
+            sort_order (str): The sort order for results (Ascending/Descending).
+            max_results (int): The maximum number of pipeline executions to return in the response.
+            next_token (str):  If the result of the previous ListPipelineExecutions request was
+                truncated, the response includes a NextToken. To retrieve the next set of pipeline
+                executions, use the token in the next request.
+
+        Returns:
+            List of Pipeline Execution Summaries. See
+            boto3 client list_pipeline_executions
+            https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_pipeline_executions
+        """
+        kwargs = dict(PipelineName=self.name)
+        update_args(
+            kwargs,
+            SortBy=sort_by,
+            SortOrder=sort_order,
+            NextToken=next_token,
+            MaxResults=max_results,
+        )
+        response = self.sagemaker_session.sagemaker_client.list_pipeline_executions(**kwargs)
+
+        # Return only PipelineExecutionSummaries and NextToken from the list_pipeline_executions
+        # response
+        return {
+            key: response[key]
+            for key in ["PipelineExecutionSummaries", "NextToken"]
+            if key in response
+        }
+
+    def _get_latest_execution_arn(self):
+        """Retrieves the latest execution of this pipeline"""
+        response = self.list_executions(
+            sort_by="CreationTime",
+            sort_order="Descending",
+            max_results=1,
+        )
+        if response["PipelineExecutionSummaries"]:
+            return response["PipelineExecutionSummaries"][0]["PipelineExecutionArn"]
+        return None
+
 
 def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
     """Formats start parameter overrides as a list of dicts.

src/sagemaker/workflow/selective_execution_config.py

@@ -0,0 +1,56 @@
+# 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.
+"""Pipeline Parallelism Configuration"""
+from __future__ import absolute_import
+from typing import List
+from sagemaker.workflow.entities import RequestType
+
+
+class SelectiveExecutionConfig:
+    """The selective execution configuration, which defines a subset of pipeline steps to run in
+
+    another SageMaker pipeline run.
+    """
+
+    def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str = None):
+        """Create a `SelectiveExecutionConfig`.
+
+        Args:
+            source_pipeline_execution_arn (str): The ARN from a reference execution of the
+                current pipeline. Used to copy input collaterals needed for the selected
+                steps to run. The execution status of the pipeline can be `Stopped`, `Failed`, or
+                `Succeeded`.
+            selected_steps (List[str]): A list of pipeline steps to run. All step(s) in all
+                path(s) between two selected steps should be included.
+        """
+        self.source_pipeline_execution_arn = source_pipeline_execution_arn
+        self.selected_steps = selected_steps
+
+    def _build_selected_steps_from_list(self) -> RequestType:
+        """Get the request structure for the list of selected steps."""
+        selected_step_list = []
+        for selected_step in self.selected_steps:
+            selected_step_list.append(dict(StepName=selected_step))
+        return selected_step_list
+
+    def to_request(self) -> RequestType:
+        """Convert `SelectiveExecutionConfig` object to request dict."""
+        request = {}
+
+        if self.source_pipeline_execution_arn is not None:
+            request["SourcePipelineExecutionArn"] = self.source_pipeline_execution_arn
+
+        if self.selected_steps is not None:
+            request["SelectedSteps"] = self._build_selected_steps_from_list()
+
+        return request

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -35,6 +35,7 @@
 from tests.unit.sagemaker.workflow.helpers import ordered, CustomStep
 from sagemaker.local.local_session import LocalSession
 from botocore.exceptions import ClientError
+from sagemaker.workflow.selective_execution_config import SelectiveExecutionConfig
 
 
 @pytest.fixture
@@ -425,6 +426,66 @@ def test_pipeline_start(sagemaker_session_mock):
     )
 
 
+def test_pipeline_start_selective_execution(sagemaker_session_mock):
+    sagemaker_session_mock.sagemaker_client.start_pipeline_execution.return_value = {
+        "PipelineExecutionArn": "my:arn"
+    }
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[],
+        steps=[],
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    # Case 1: Happy path
+    selective_execution_config = SelectiveExecutionConfig(
+        source_pipeline_execution_arn="foo-arn", selected_steps=["step-1", "step-2", "step-3"]
+    )
+    pipeline.start(selective_execution_config=selective_execution_config)
+    sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with(
+        PipelineName="MyPipeline",
+        SelectiveExecutionConfig={
+            "SelectedSteps": [
+                {"StepName": "step-1"},
+                {"StepName": "step-2"},
+                {"StepName": "step-3"},
+            ],
+            "SourcePipelineExecutionArn": "foo-arn",
+        },
+    )
+
+    # Case 2: Start selective execution without SourcePipelineExecutionArn
+    sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
+        "PipelineExecutionSummaries": [
+            {
+                "PipelineExecutionArn": "my:latest:execution:arn",
+                "PipelineExecutionDisplayName": "Latest",
+            }
+        ]
+    }
+    selective_execution_config = SelectiveExecutionConfig(
+        selected_steps=["step-1", "step-2", "step-3"]
+    )
+    pipeline.start(selective_execution_config=selective_execution_config)
+    sagemaker_session_mock.sagemaker_client.list_pipeline_executions.assert_called_with(
+        PipelineName="MyPipeline",
+        SortBy="CreationTime",
+        SortOrder="Descending",
+        MaxResults=1,
+    )
+    sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with(
+        PipelineName="MyPipeline",
+        SelectiveExecutionConfig={
+            "SelectedSteps": [
+                {"StepName": "step-1"},
+                {"StepName": "step-2"},
+                {"StepName": "step-3"},
+            ],
+            "SourcePipelineExecutionArn": "my:latest:execution:arn",
+        },
+    )
+
+
 def test_pipeline_basic():
     parameter = ParameterString("MyStr")
     pipeline = Pipeline(
@@ -593,6 +654,31 @@ def test_pipeline_disable_experiment_config():
     )
 
 
+def test_pipeline_list_executions(sagemaker_session_mock):
+    sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
+        "PipelineExecutionSummaries": [Mock()],
+        "ResponseMetadata": "metadata",
+    }
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[ParameterString("alpha", "beta"), ParameterString("gamma", "delta")],
+        steps=[],
+        sagemaker_session=sagemaker_session_mock,
+    )
+    executions = pipeline.list_executions()
+    assert len(executions) == 1
+    assert len(executions["PipelineExecutionSummaries"]) == 1
+    sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
+        "PipelineExecutionSummaries": [Mock(), Mock()],
+        "NextToken": "token",
+        "ResponseMetadata": "metadata",
+    }
+    executions = pipeline.list_executions()
+    assert len(executions) == 2
+    assert len(executions["PipelineExecutionSummaries"]) == 2
+    assert executions["NextToken"] == "token"
+
+
 def test_pipeline_execution_basics(sagemaker_session_mock):
     sagemaker_session_mock.sagemaker_client.start_pipeline_execution.return_value = {
         "PipelineExecutionArn": "my:arn"

tests/unit/sagemaker/workflow/test_selective_execution_config.py

@@ -0,0 +1,27 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Unit tests for SelectiveExecutionConfig module"""
+
+from __future__ import absolute_import
+
+from sagemaker.workflow.selective_execution_config import SelectiveExecutionConfig
+
+
+def test_SelectiveExecutionConfig():
+    selective_execution_config = SelectiveExecutionConfig(
+        source_pipeline_execution_arn="foo-arn", selected_steps=["step-1", "step-2", "step-3"]
+    )
+    assert selective_execution_config.to_request() == {
+        "SelectedSteps": [{"StepName": "step-1"}, {"StepName": "step-2"}, {"StepName": "step-3"}],
+        "SourcePipelineExecutionArn": "foo-arn",
+    }