feature: allow opt out from referencing latest execution in the selec… (aws#1004)

Author: shenlongtang

src/sagemaker/workflow/pipeline.py

@@ -344,7 +344,10 @@ def start(
             A `_PipelineExecution` instance, if successful.
         """
         if selective_execution_config is not None:
-            if selective_execution_config.source_pipeline_execution_arn is None:
+            if (
+                selective_execution_config.source_pipeline_execution_arn is None
+                and selective_execution_config.reference_latest_execution
+            ):
                 selective_execution_config.source_pipeline_execution_arn = (
                     self._get_latest_execution_arn()
                 )

src/sagemaker/workflow/selective_execution_config.py

@@ -22,7 +22,12 @@ class SelectiveExecutionConfig:
     another SageMaker pipeline run.
     """
 
-    def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str = None):
+    def __init__(
+        self,
+        selected_steps: List[str],
+        source_pipeline_execution_arn: str = None,
+        reference_latest_execution: bool = True,
+    ):
         """Create a `SelectiveExecutionConfig`.
 
         Args:
@@ -32,9 +37,12 @@ def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str
                 `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.
+            reference_latest_execution (bool): Whether to reference the latest execution if
+                source_pipeline_execution_arn is not provided.
         """
         self.source_pipeline_execution_arn = source_pipeline_execution_arn
         self.selected_steps = selected_steps
+        self.reference_latest_execution = reference_latest_execution
 
     def _build_selected_steps_from_list(self) -> RequestType:
         """Get the request structure for the list of selected steps."""

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -492,8 +492,10 @@ def test_pipeline_start_selective_execution(sagemaker_session_mock):
             "SourcePipelineExecutionArn": "foo-arn",
         },
     )
+    sagemaker_session_mock.reset_mock()
 
     # Case 2: Start selective execution without SourcePipelineExecutionArn
+    # References latest execution by default.
     sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
         "PipelineExecutionSummaries": [
             {
@@ -523,6 +525,27 @@ def test_pipeline_start_selective_execution(sagemaker_session_mock):
             "SourcePipelineExecutionArn": "my:latest:execution:arn",
         },
     )
+    sagemaker_session_mock.reset_mock()
+
+    # Case 3: Start selective execution without SourcePipelineExecutionArn
+    # Opts not to reference latest execution.
+    selective_execution_config = SelectiveExecutionConfig(
+        selected_steps=["step-1", "step-2", "step-3"],
+        reference_latest_execution=False,
+    )
+    pipeline.start(selective_execution_config=selective_execution_config)
+    sagemaker_session_mock.sagemaker_client.list_pipeline_executions.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with(
+        PipelineName="MyPipeline",
+        SelectiveExecutionConfig={
+            "SelectedSteps": [
+                {"StepName": "step-1"},
+                {"StepName": "step-2"},
+                {"StepName": "step-3"},
+            ],
+        },
+    )
+    sagemaker_session_mock.reset_mock()
 
 
 def test_pipeline_basic():