Skip to content

Commit 477f0c1

Browse files
committed
feature: support for custom pipeline execution name
1 parent e43fa6d commit 477f0c1

File tree

2 files changed

+10
-1
lines changed

2 files changed

+10
-1
lines changed

src/sagemaker/workflow/pipeline.py

Lines changed: 4 additions & 1 deletion
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
@@ -186,13 +185,15 @@ def delete(self) -> Dict[str, Any]:
186185
def start(
187186
self,
188187
parameters: Dict[str, Any] = None,
188+
execution_display_name: str = None,
189189
execution_description: str = None,
190190
):
191191
"""Starts a Pipeline execution in the Workflow service.
192192
193193
Args:
194194
parameters (List[Dict[str, str]]): A list of parameter dicts of the form
195195
{"Name": "string", "Value": "string"}.
196+
execution_display_name (str): The display name of the pipeline execution.
196197
execution_description (str): A description of the execution.
197198
198199
Returns:
@@ -209,11 +210,13 @@ def start(
209210
"This pipeline is not associated with a Pipeline in SageMaker. "
210211
"Please invoke create() first before attempting to invoke start()."
211212
)
213+
212214
kwargs = dict(PipelineName=self.name)
213215
update_args(
214216
kwargs,
215217
PipelineParameters=format_start_parameters(parameters),
216218
PipelineExecutionDescription=execution_description,
219+
PipelineExecutionDisplayName=execution_display_name,
217220
)
218221
response = self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs)
219222
return _PipelineExecution(

tests/unit/sagemaker/workflow/test_pipeline.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,12 @@ def test_pipeline_start(sagemaker_session_mock):
154154
assert sagemaker_session_mock.start_pipeline_execution.called_with(
155155
PipelineName="MyPipeline",
156156
)
157+
158+
pipeline.start(execution_display_name="pipeline-execution")
159+
assert sagemaker_session_mock.start_pipeline_execution.called_with(
160+
PipelineName="MyPipeline", PipelineExecutionDisplayName="pipeline-execution"
161+
)
162+
157163
pipeline.start(parameters=dict(alpha="epsilon"))
158164
assert sagemaker_session_mock.start_pipeline_execution.called_with(
159165
PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]

0 commit comments

Comments
 (0)