Skip to content

Commit 5ac138c

Browse files
qidewenwhenqidewenwhen
committed
fix: Improve Pipeline integ tests and fix resource leak
1 parent d4203da commit 5ac138c

16 files changed

+99
-207
lines changed

‎tests/integ/sagemaker/workflow/helpers.py

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
from botocore.exceptions import WaiterError
16+
17+
from sagemaker.workflow.pipeline import _PipelineExecution
18+
19+
20+
def wait_pipeline_execution(execution: _PipelineExecution, delay: int = 30, max_attempts: int = 60):
21+
try:
22+
execution.wait(delay=delay, max_attempts=max_attempts)
23+
except WaiterError:
24+
pass

‎tests/integ/sagemaker/workflow/test_automl_steps.py

+2-5
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
import re
1717

1818
import pytest
19-
from botocore.exceptions import WaiterError
2019

20+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
2121
from sagemaker.workflow import ParameterString
2222
from sagemaker.workflow.automl_step import AutoMLStep
2323
from sagemaker.automl.automl import AutoML, AutoMLInput
@@ -133,10 +133,7 @@ def test_automl_step(pipeline_session, role, pipeline_name):
133133
try:
134134
_ = pipeline.create(role)
135135
execution = pipeline.start(parameters={})
136-
try:
137-
execution.wait(delay=30, max_attempts=60)
138-
except WaiterError:
139-
pass
136+
wait_pipeline_execution(execution=execution)
140137

141138
execution_steps = execution.list_steps()
142139
has_automl_job = False

‎tests/integ/sagemaker/workflow/test_clarify_check_steps.py

+3-9
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717
import os
1818

1919
import pytest
20-
from botocore.exceptions import WaiterError
2120

21+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
2222
from sagemaker.clarify import (
2323
BiasConfig,
2424
DataConfig,
@@ -182,10 +182,7 @@ def test_one_step_data_bias_pipeline_happycase(
182182

183183
assert response["PipelineArn"] == create_arn
184184

185-
try:
186-
execution.wait(delay=30, max_attempts=60)
187-
except WaiterError:
188-
pass
185+
wait_pipeline_execution(execution=execution)
189186
execution_steps = execution.list_steps()
190187

191188
assert len(execution_steps) == 1
@@ -272,10 +269,7 @@ def test_one_step_data_bias_pipeline_constraint_violation(
272269

273270
assert response["PipelineArn"] == create_arn
274271

275-
try:
276-
execution.wait(delay=30, max_attempts=60)
277-
except WaiterError:
278-
pass
272+
wait_pipeline_execution(execution=execution)
279273
execution_steps = execution.list_steps()
280274

281275
assert len(execution_steps) == 1

‎tests/integ/sagemaker/workflow/test_experiment.py

+3-9
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import pytest
1919

20-
from botocore.exceptions import WaiterError
20+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
2121
from sagemaker.processing import ProcessingInput
2222
from sagemaker.session import get_execution_role
2323
from sagemaker.sklearn.processing import SKLearnProcessor
@@ -120,10 +120,7 @@ def test_pipeline_execution_with_default_experiment_config(
120120
pipeline.create(role)
121121
execution = pipeline.start(parameters={})
122122

123-
try:
124-
execution.wait(delay=30, max_attempts=3)
125-
except WaiterError:
126-
pass
123+
wait_pipeline_execution(execution=execution, max_attempts=3)
127124
execution_steps = execution.list_steps()
128125
assert len(execution_steps) == 1
129126
assert execution_steps[0]["StepName"] == "sklearn-process"
@@ -195,10 +192,7 @@ def test_pipeline_execution_with_custom_experiment_config(
195192
pipeline.create(role)
196193
execution = pipeline.start(parameters={})
197194

198-
try:
199-
execution.wait(delay=30, max_attempts=3)
200-
except WaiterError:
201-
pass
195+
wait_pipeline_execution(execution=execution, max_attempts=3)
202196
execution_steps = execution.list_steps()
203197
assert len(execution_steps) == 1
204198
assert execution_steps[0]["StepName"] == "sklearn-process"

‎tests/integ/sagemaker/workflow/test_fail_steps.py

+5-17
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
from __future__ import absolute_import
1414

1515
import pytest
16-
from botocore.exceptions import WaiterError
1716

17+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
1818
from sagemaker import get_execution_role, utils
1919
from sagemaker.workflow.condition_step import ConditionStep
2020
from sagemaker.workflow.conditions import ConditionEquals
@@ -62,10 +62,7 @@ def test_two_step_fail_pipeline_with_str_err_msg(sagemaker_session, role, pipeli
6262
response = execution.describe()
6363
assert response["PipelineArn"] == pipeline_arn
6464

65-
try:
66-
execution.wait(delay=30, max_attempts=60)
67-
except WaiterError:
68-
pass
65+
wait_pipeline_execution(execution=execution)
6966
execution_steps = execution.list_steps()
7067

7168
assert len(execution_steps) == 2
@@ -130,10 +127,7 @@ def test_two_step_fail_pipeline_with_parameter_err_msg(sagemaker_session, role,
130127
response = execution.describe()
131128
assert response["PipelineArn"] == pipeline_arn
132129

133-
try:
134-
execution.wait(delay=30, max_attempts=60)
135-
except WaiterError:
136-
pass
130+
wait_pipeline_execution(execution=execution)
137131
execution_steps = execution.list_steps()
138132

139133
assert len(execution_steps) == 2
@@ -196,10 +190,7 @@ def test_two_step_fail_pipeline_with_join_fn(sagemaker_session, role, pipeline_n
196190
response = execution.describe()
197191
assert response["PipelineArn"] == pipeline_arn
198192

199-
try:
200-
execution.wait(delay=30, max_attempts=60)
201-
except WaiterError:
202-
pass
193+
wait_pipeline_execution(execution=execution)
203194
execution_steps = execution.list_steps()
204195

205196
assert len(execution_steps) == 2
@@ -257,10 +248,7 @@ def test_two_step_fail_pipeline_with_no_err_msg(sagemaker_session, role, pipelin
257248
response = execution.describe()
258249
assert response["PipelineArn"] == pipeline_arn
259250

260-
try:
261-
execution.wait(delay=30, max_attempts=60)
262-
except WaiterError:
263-
pass
251+
wait_pipeline_execution(execution=execution)
264252
execution_steps = execution.list_steps()
265253

266254
assert len(execution_steps) == 2

‎tests/integ/sagemaker/workflow/test_model_create_and_registration.py

+11-9
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@
2323
import re
2424

2525
import pytest
26-
from botocore.exceptions import WaiterError
2726

2827
import tests
28+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
2929
from sagemaker.tensorflow import TensorFlow, TensorFlowModel
3030
from tests.integ.retry import retries
3131
from sagemaker.drift_check_baselines import DriftCheckBaselines
@@ -180,12 +180,14 @@ def test_conditional_pytorch_training_model_registration(
180180
)
181181

182182
execution = pipeline.start(parameters={})
183+
wait_pipeline_execution(execution=execution)
183184
assert re.match(
184185
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
185186
execution.arn,
186187
)
187188

188189
execution = pipeline.start(parameters={"GoodEnoughInput": 0})
190+
wait_pipeline_execution(execution=execution)
189191
assert re.match(
190192
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
191193
execution.arn,
@@ -259,12 +261,14 @@ def test_mxnet_model_registration(
259261
)
260262

261263
execution = pipeline.start(parameters={})
264+
wait_pipeline_execution(execution=execution)
262265
assert re.match(
263266
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
264267
execution.arn,
265268
)
266269

267270
execution = pipeline.start()
271+
wait_pipeline_execution(execution=execution)
268272
assert re.match(
269273
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
270274
execution.arn,
@@ -470,12 +474,14 @@ def test_sklearn_xgboost_sip_model_registration(
470474
)
471475

472476
execution = pipeline.start(parameters={})
477+
wait_pipeline_execution(execution=execution)
473478
assert re.match(
474479
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
475480
execution.arn,
476481
)
477482

478483
execution = pipeline.start()
484+
wait_pipeline_execution(execution=execution)
479485
assert re.match(
480486
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
481487
execution.arn,
@@ -656,10 +662,7 @@ def test_model_registration_with_drift_check_baselines(
656662

657663
assert response["PipelineArn"] == create_arn
658664

659-
try:
660-
execution.wait(delay=30, max_attempts=60)
661-
except WaiterError:
662-
pass
665+
wait_pipeline_execution(execution=execution)
663666
execution_steps = execution.list_steps()
664667

665668
assert len(execution_steps) == 1
@@ -797,12 +800,14 @@ def test_model_registration_with_model_repack(
797800
)
798801

799802
execution = pipeline.start(parameters={})
803+
wait_pipeline_execution(execution=execution)
800804
assert re.match(
801805
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
802806
execution.arn,
803807
)
804808

805809
execution = pipeline.start(parameters={"GoodEnoughInput": 0})
810+
wait_pipeline_execution(execution=execution)
806811
assert re.match(
807812
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
808813
execution.arn,
@@ -889,10 +894,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
889894
rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
890895
execution.arn,
891896
)
892-
try:
893-
execution.wait(delay=30, max_attempts=60)
894-
except WaiterError:
895-
pass
897+
wait_pipeline_execution(execution=execution)
896898
execution_steps = execution.list_steps()
897899

898900
for step in execution_steps:

‎tests/integ/sagemaker/workflow/test_model_steps.py

+8-29
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
import os
1717

1818
import pytest
19-
from botocore.exceptions import WaiterError
2019

20+
from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
2121
from sagemaker.workflow.fail_step import FailStep
2222
from sagemaker.workflow.functions import Join
2323
from tests.integ.timeout import timeout_and_delete_endpoint_by_name
@@ -155,10 +155,7 @@ def test_pytorch_training_model_registration_and_creation_without_custom_inferen
155155
seconds_to_sleep=10,
156156
):
157157
execution = pipeline.start(parameters={})
158-
try:
159-
execution.wait(delay=30, max_attempts=60)
160-
except WaiterError:
161-
pass
158+
wait_pipeline_execution(execution=execution)
162159
execution_steps = execution.list_steps()
163160
is_execution_fail = False
164161
for step in execution_steps:
@@ -284,10 +281,7 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference(
284281
seconds_to_sleep=10,
285282
):
286283
execution = pipeline.start(parameters={})
287-
try:
288-
execution.wait(delay=30, max_attempts=60)
289-
except WaiterError:
290-
pass
284+
wait_pipeline_execution(execution=execution)
291285
execution_steps = execution.list_steps()
292286
is_execution_fail = False
293287
for step in execution_steps:
@@ -372,10 +366,7 @@ def test_mxnet_model_registration_with_custom_inference(
372366
seconds_to_sleep=10,
373367
):
374368
execution = pipeline.start()
375-
try:
376-
execution.wait(delay=30, max_attempts=60)
377-
except WaiterError:
378-
pass
369+
wait_pipeline_execution(execution=execution)
379370
execution_steps = execution.list_steps()
380371

381372
assert len(execution_steps) == 1
@@ -550,10 +541,7 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics(
550541

551542
assert response["PipelineArn"] == create_arn
552543

553-
try:
554-
execution.wait(delay=30, max_attempts=60)
555-
except WaiterError:
556-
pass
544+
wait_pipeline_execution(execution=execution)
557545
execution_steps = execution.list_steps()
558546

559547
assert len(execution_steps) == 1
@@ -667,10 +655,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
667655
seconds_to_sleep=10,
668656
):
669657
execution = pipeline.start(parameters={})
670-
try:
671-
execution.wait(delay=30, max_attempts=60)
672-
except WaiterError:
673-
pass
658+
wait_pipeline_execution(execution=execution)
674659
execution_steps = execution.list_steps()
675660
is_execution_fail = False
676661
for step in execution_steps:
@@ -750,10 +735,7 @@ def test_xgboost_model_register_and_deploy_with_runtime_repack(
750735
seconds_to_sleep=10,
751736
):
752737
execution = pipeline.start(parameters={})
753-
try:
754-
execution.wait(delay=30, max_attempts=60)
755-
except WaiterError:
756-
pass
738+
wait_pipeline_execution(execution=execution)
757739

758740
# Verify the pipeline execution succeeded
759741
step_register_model = None
@@ -866,10 +848,7 @@ def test_tensorflow_model_register_and_deploy_with_runtime_repack(
866848
seconds_to_sleep=10,
867849
):
868850
execution = pipeline.start(parameters={})
869-
try:
870-
execution.wait(delay=30, max_attempts=60)
871-
except WaiterError:
872-
pass
851+
wait_pipeline_execution(execution=execution)
873852

874853
# Verify the pipeline execution succeeded
875854
step_register_model = None

0 commit comments

Comments
 (0)