Skip to content

Commit 5afcfab

Browse files
Ameen KhanZhankuil
Ameen Khan
authored andcommitted
Updated docstring, added describe call to check config
1 parent 04396ff commit 5afcfab

File tree

2 files changed

+23
-10
lines changed

2 files changed

+23
-10
lines changed

src/sagemaker/workflow/pipeline.py

+12-8
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,9 @@ def create(
105105
description (str): A description of the pipeline.
106106
tags (List[Dict[str, str]]): A list of {"Key": "string", "Value": "string"} dicts as
107107
tags.
108-
parallelism_config (Optional[Config for parallel steps, Parallelism configuration that
109-
is applied to each of. the executions
108+
parallelism_config (Optional[ParallelismConfiguration]): Parallelism configuration
109+
that is applied to each of the executions of the pipeline. It takes precedence
110+
over the parallelism configuration of the parent pipeline.
110111
111112
Returns:
112113
A response dict from the service.
@@ -130,8 +131,9 @@ def _create_args(
130131
Args:
131132
role_arn (str): The role arn that is assumed by pipelines to create step artifacts.
132133
description (str): A description of the pipeline.
133-
parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
134-
is applied to each of the executions.
134+
parallelism_config (Optional[ParallelismConfiguration]): Parallelism configuration
135+
that is applied to each of the executions of the pipeline. It takes precedence
136+
over the parallelism configuration of the parent pipeline.
135137
136138
Returns:
137139
A keyword argument dict for calling create_pipeline.
@@ -188,8 +190,9 @@ def update(
188190
Args:
189191
role_arn (str): The role arn that is assumed by pipelines to create step artifacts.
190192
description (str): A description of the pipeline.
191-
parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
192-
is applied to each of the executions.
193+
parallelism_config (Optional[ParallelismConfiguration]): Parallelism configuration
194+
that is applied to each of the executions of the pipeline. It takes precedence
195+
over the parallelism configuration of the parent pipeline.
193196
194197
Returns:
195198
A response dict from the service.
@@ -265,8 +268,9 @@ def start(
265268
pipeline parameters.
266269
execution_display_name (str): The display name of the pipeline execution.
267270
execution_description (str): A description of the execution.
268-
parallelism_config (Optional[ParallelismConfiguration]): Config for parallel steps, that
269-
is applied to each of the executions.
271+
parallelism_config (Optional[ParallelismConfiguration]): Parallelism configuration
272+
that is applied to each of the executions of the pipeline. It takes precedence
273+
over the parallelism configuration of the parent pipeline.
270274
271275
Returns:
272276
A `_PipelineExecution` instance, if successful.

tests/integ/test_workflow.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -2787,6 +2787,8 @@ def test_large_pipeline(sagemaker_session, role, pipeline_name, region_name):
27872787
fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
27882788
create_arn,
27892789
)
2790+
response = pipeline.describe()
2791+
assert len(json.loads(pipeline.describe()["PipelineDefinition"])["Steps"]) == 2000
27902792

27912793
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
27922794
response = pipeline.update(role)
@@ -2801,7 +2803,7 @@ def test_large_pipeline(sagemaker_session, role, pipeline_name, region_name):
28012803
except Exception:
28022804
pass
28032805

2804-
def test_create_parallelism_config(sagemaker_session, role, pipeline_name, region_name):
2806+
def test_create_and_update_with_parallelism_config(sagemaker_session, role, pipeline_name, region_name):
28052807
instance_count = ParameterInteger(name="InstanceCount", default_value=2)
28062808

28072809
outputParam = CallbackOutput(output_name="output", output_type=CallbackOutputTypeEnum.String)
@@ -2822,21 +2824,28 @@ def test_create_parallelism_config(sagemaker_session, role, pipeline_name, regio
28222824
sagemaker_session=sagemaker_session,
28232825
)
28242826

2827+
28252828
try:
28262829
response = pipeline.create(role, parallelism_config={"MaxParallelExecutionSteps": 50})
28272830
create_arn = response["PipelineArn"]
28282831
assert re.match(
28292832
fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
28302833
create_arn,
28312834
)
2835+
response = pipeline.describe()
2836+
assert response["ParallelismConfiguration"]["MaxParallelExecutionSteps"] == 50
28322837

28332838
pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
2834-
response = pipeline.update(role, parallelism_config={"MaxParallelExecutionSteps": 50})
2839+
response = pipeline.update(role, parallelism_config={"MaxParallelExecutionSteps": 55})
28352840
update_arn = response["PipelineArn"]
28362841
assert re.match(
28372842
fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
28382843
update_arn,
28392844
)
2845+
2846+
response = pipeline.describe()
2847+
assert response["ParallelismConfiguration"]["MaxParallelExecutionSteps"] == 55
2848+
28402849
finally:
28412850
try:
28422851
pipeline.delete()

0 commit comments

Comments
 (0)