Skip to content

Commit c5a00ea

Browse files
authored
Merge branch 'master' into fix-cloudwatch-logs-local
2 parents c6b306f + bc1833a commit c5a00ea

File tree

2 files changed

+10
-1
lines changed

2 files changed

+10
-1
lines changed

src/sagemaker/workflow/pipeline.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020

2121
import attr
2222
import botocore
23-
2423
from botocore.exceptions import ClientError
2524

2625
from sagemaker._studio import _append_project_tags
@@ -197,13 +196,15 @@ def delete(self) -> Dict[str, Any]:
197196
def start(
198197
self,
199198
parameters: Dict[str, Any] = None,
199+
execution_display_name: str = None,
200200
execution_description: str = None,
201201
):
202202
"""Starts a Pipeline execution in the Workflow service.
203203
204204
Args:
205205
parameters (List[Dict[str, str]]): A list of parameter dicts of the form
206206
{"Name": "string", "Value": "string"}.
207+
execution_display_name (str): The display name of the pipeline execution.
207208
execution_description (str): A description of the execution.
208209
209210
Returns:
@@ -220,11 +221,13 @@ def start(
220221
"This pipeline is not associated with a Pipeline in SageMaker. "
221222
"Please invoke create() first before attempting to invoke start()."
222223
)
224+
223225
kwargs = dict(PipelineName=self.name)
224226
update_args(
225227
kwargs,
226228
PipelineParameters=format_start_parameters(parameters),
227229
PipelineExecutionDescription=execution_description,
230+
PipelineExecutionDisplayName=execution_display_name,
228231
)
229232
response = self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs)
230233
return _PipelineExecution(

tests/unit/sagemaker/workflow/test_pipeline.py

+6
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,12 @@ def test_pipeline_start(sagemaker_session_mock):
159159
assert sagemaker_session_mock.start_pipeline_execution.called_with(
160160
PipelineName="MyPipeline",
161161
)
162+
163+
pipeline.start(execution_display_name="pipeline-execution")
164+
assert sagemaker_session_mock.start_pipeline_execution.called_with(
165+
PipelineName="MyPipeline", PipelineExecutionDisplayName="pipeline-execution"
166+
)
167+
162168
pipeline.start(parameters=dict(alpha="epsilon"))
163169
assert sagemaker_session_mock.start_pipeline_execution.called_with(
164170
PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]

0 commit comments

Comments
 (0)