feature: Default selective execution source pipeline to latest pipeline execution (aws#931)

Author: nishkris

src/sagemaker/workflow/pipeline.py

@@ -332,6 +332,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:
+                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)
@@ -436,6 +440,17 @@ def list_executions(
             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

@@ -19,7 +19,7 @@
 class SelectiveExecutionConfig:
     """Selective execution config config for SageMaker pipeline."""
 
-    def __init__(self, source_pipeline_execution_arn: str, selected_steps: List[str]):
+    def __init__(self, selected_steps: List[str], source_pipeline_execution_arn: str = None):
         """Create a SelectiveExecutionConfig
 
         Args:

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -423,6 +423,19 @@ def test_pipeline_start(sagemaker_session_mock):
         PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]
     )
 
+
+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"]
     )
@@ -439,6 +452,37 @@ def test_pipeline_start(sagemaker_session_mock):
         },
     )
 
+    # 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")