feature: Add SelectiveExecutionConfig (aws#929)

Co-authored-by: Shegufta Ahsan <[email protected]>

Author: 2 people

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,22 @@ 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 (SelectiveExecutionConfig): configuration for
+                selective step execution
 
         Returns:
             A `_PipelineExecution` instance, if successful.
         """
+        if selective_execution_config is not None:
+            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)

src/sagemaker/workflow/selective_execution_config.py

@@ -0,0 +1,51 @@
+# 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:
+    """Selective execution config config for SageMaker pipeline."""
+
+    def __init__(self, source_pipeline_execution_arn: str, selected_steps: List[str]):
+        """Create a SelectiveExecutionConfig
+
+        Args:
+            source_pipeline_execution_arn (str): ARN of the previously executed pipeline execution.
+                The given arn pipeline execution status can be either Failed or Success.
+            selected_steps (List[str]): List of step names that pipeline users want to run
+                in new subworkflow-execution. The steps must be connected.
+        """
+        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 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
@@ -424,6 +425,22 @@ def test_pipeline_start(sagemaker_session_mock):
         PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]
     )
 
+    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",
+        },
+    )
+
 
 def test_pipeline_basic():
     parameter = ParameterString("MyStr")

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