Skip to content

Commit 1283d3a

Browse files
nishkrisShegufta Ahsan
authored and
Shegufta Ahsan
committed
feature: Add support for listing executions from pipeline (aws#930)
1 parent baecafe commit 1283d3a

File tree

2 files changed

+65
-0
lines changed

2 files changed

+65
-0
lines changed

src/sagemaker/workflow/pipeline.py

+40
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,46 @@ def _interpolate_step_collection_name_in_depends_on(self, step_requests: list):
396396
)
397397
self._interpolate_step_collection_name_in_depends_on(sub_step_requests)
398398

399+
def list_executions(
400+
self,
401+
sort_by: str = None,
402+
sort_order: str = None,
403+
max_results: int = None,
404+
next_token: str = None,
405+
) -> Dict[str, Any]:
406+
"""Lists a pipeline's executions.
407+
408+
Args:
409+
sort_by (str): The field by which to sort results(CreationTime/PipelineExecutionArn).
410+
sort_order (str): The sort order for results (Ascending/Descending).
411+
max_results (int): The maximum number of pipeline executions to return in the response.
412+
next_token (str): If the result of the previous ListPipelineExecutions request was
413+
truncated, the response includes a NextToken. To retrieve the next set of pipeline
414+
executions, use the token in the next request.
415+
416+
Returns:
417+
List of Pipeline Execution Summaries. See
418+
boto3 client list_pipeline_executions
419+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_pipeline_executions
420+
"""
421+
kwargs = dict(PipelineName=self.name)
422+
update_args(
423+
kwargs,
424+
SortBy=sort_by,
425+
SortOrder=sort_order,
426+
NextToken=next_token,
427+
MaxResults=max_results,
428+
)
429+
response = self.sagemaker_session.sagemaker_client.list_pipeline_executions(**kwargs)
430+
431+
# Return only PipelineExecutionSummaries and NextToken from the list_pipeline_executions
432+
# response
433+
return {
434+
key: response[key]
435+
for key in ["PipelineExecutionSummaries", "NextToken"]
436+
if key in response
437+
}
438+
399439

400440
def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
401441
"""Formats start parameter overrides as a list of dicts.

tests/unit/sagemaker/workflow/test_pipeline.py

+25
Original file line numberDiff line numberDiff line change
@@ -608,6 +608,31 @@ def test_pipeline_disable_experiment_config():
608608
)
609609

610610

611+
def test_pipeline_list_executions(sagemaker_session_mock):
612+
sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
613+
"PipelineExecutionSummaries": [Mock()],
614+
"ResponseMetadata": "metadata",
615+
}
616+
pipeline = Pipeline(
617+
name="MyPipeline",
618+
parameters=[ParameterString("alpha", "beta"), ParameterString("gamma", "delta")],
619+
steps=[],
620+
sagemaker_session=sagemaker_session_mock,
621+
)
622+
executions = pipeline.list_executions()
623+
assert len(executions) == 1
624+
assert len(executions["PipelineExecutionSummaries"]) == 1
625+
sagemaker_session_mock.sagemaker_client.list_pipeline_executions.return_value = {
626+
"PipelineExecutionSummaries": [Mock(), Mock()],
627+
"NextToken": "token",
628+
"ResponseMetadata": "metadata",
629+
}
630+
executions = pipeline.list_executions()
631+
assert len(executions) == 2
632+
assert len(executions["PipelineExecutionSummaries"]) == 2
633+
assert executions["NextToken"] == "token"
634+
635+
611636
def test_pipeline_execution_basics(sagemaker_session_mock):
612637
sagemaker_session_mock.sagemaker_client.start_pipeline_execution.return_value = {
613638
"PipelineExecutionArn": "my:arn"

0 commit comments

Comments
 (0)