diff --git a/tests/integ/sagemaker/workflow/helpers.py b/tests/integ/sagemaker/workflow/helpers.py new file mode 100644 index 0000000000..857bb3b2b2 --- /dev/null +++ b/tests/integ/sagemaker/workflow/helpers.py @@ -0,0 +1,24 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +from __future__ import absolute_import + +from botocore.exceptions import WaiterError + +from sagemaker.workflow.pipeline import _PipelineExecution + + +def wait_pipeline_execution(execution: _PipelineExecution, delay: int = 30, max_attempts: int = 60): + try: + execution.wait(delay=delay, max_attempts=max_attempts) + except WaiterError: + pass diff --git a/tests/integ/sagemaker/workflow/test_automl_steps.py b/tests/integ/sagemaker/workflow/test_automl_steps.py index a44fee2e86..a5ed3f48f1 100644 --- a/tests/integ/sagemaker/workflow/test_automl_steps.py +++ b/tests/integ/sagemaker/workflow/test_automl_steps.py @@ -16,8 +16,8 @@ import re import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.workflow import ParameterString from sagemaker.workflow.automl_step import AutoMLStep from sagemaker.automl.automl import AutoML, AutoMLInput @@ -133,10 +133,7 @@ def test_automl_step(pipeline_session, role, pipeline_name): try: _ = pipeline.create(role) execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() has_automl_job = False diff --git a/tests/integ/sagemaker/workflow/test_clarify_check_steps.py b/tests/integ/sagemaker/workflow/test_clarify_check_steps.py index aea509b4fb..cb43f12ad7 100644 --- a/tests/integ/sagemaker/workflow/test_clarify_check_steps.py +++ b/tests/integ/sagemaker/workflow/test_clarify_check_steps.py @@ -17,8 +17,8 @@ import os import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.clarify import ( BiasConfig, DataConfig, @@ -182,10 +182,7 @@ def test_one_step_data_bias_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -272,10 +269,7 @@ def test_one_step_data_bias_pipeline_constraint_violation( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 diff --git a/tests/integ/sagemaker/workflow/test_experiment.py b/tests/integ/sagemaker/workflow/test_experiment.py index 345372174d..35828a69cd 100644 --- a/tests/integ/sagemaker/workflow/test_experiment.py +++ b/tests/integ/sagemaker/workflow/test_experiment.py @@ -17,7 +17,7 @@ import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.processing import ProcessingInput from sagemaker.session import get_execution_role from sagemaker.sklearn.processing import SKLearnProcessor @@ -120,10 +120,7 @@ def test_pipeline_execution_with_default_experiment_config( pipeline.create(role) execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sklearn-process" @@ -195,10 +192,7 @@ def test_pipeline_execution_with_custom_experiment_config( pipeline.create(role) execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sklearn-process" diff --git a/tests/integ/sagemaker/workflow/test_fail_steps.py b/tests/integ/sagemaker/workflow/test_fail_steps.py index ba00b4f972..5f8c1e04ab 100644 --- a/tests/integ/sagemaker/workflow/test_fail_steps.py +++ b/tests/integ/sagemaker/workflow/test_fail_steps.py @@ -13,8 +13,8 @@ from __future__ import absolute_import import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker import get_execution_role, utils from sagemaker.workflow.condition_step import ConditionStep from sagemaker.workflow.conditions import ConditionEquals @@ -62,10 +62,7 @@ def test_two_step_fail_pipeline_with_str_err_msg(sagemaker_session, role, pipeli response = execution.describe() assert response["PipelineArn"] == pipeline_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -130,10 +127,7 @@ def test_two_step_fail_pipeline_with_parameter_err_msg(sagemaker_session, role, response = execution.describe() assert response["PipelineArn"] == pipeline_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -196,10 +190,7 @@ def test_two_step_fail_pipeline_with_join_fn(sagemaker_session, role, pipeline_n response = execution.describe() assert response["PipelineArn"] == pipeline_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -257,10 +248,7 @@ def test_two_step_fail_pipeline_with_no_err_msg(sagemaker_session, role, pipelin response = execution.describe() assert response["PipelineArn"] == pipeline_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 diff --git a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py b/tests/integ/sagemaker/workflow/test_model_create_and_registration.py index 1045a8ef0c..98ee4bfde6 100644 --- a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py +++ b/tests/integ/sagemaker/workflow/test_model_create_and_registration.py @@ -23,9 +23,9 @@ import re import pytest -from botocore.exceptions import WaiterError import tests +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.tensorflow import TensorFlow, TensorFlowModel from tests.integ.retry import retries from sagemaker.drift_check_baselines import DriftCheckBaselines @@ -180,12 +180,14 @@ def test_conditional_pytorch_training_model_registration( ) execution = pipeline.start(parameters={}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) execution = pipeline.start(parameters={"GoodEnoughInput": 0}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, @@ -259,12 +261,14 @@ def test_mxnet_model_registration( ) execution = pipeline.start(parameters={}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) execution = pipeline.start() + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, @@ -470,12 +474,14 @@ def test_sklearn_xgboost_sip_model_registration( ) execution = pipeline.start(parameters={}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) execution = pipeline.start() + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, @@ -656,10 +662,7 @@ def test_model_registration_with_drift_check_baselines( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -797,12 +800,14 @@ def test_model_registration_with_model_repack( ) execution = pipeline.start(parameters={}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) execution = pipeline.start(parameters={"GoodEnoughInput": 0}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, @@ -889,10 +894,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() for step in execution_steps: diff --git a/tests/integ/sagemaker/workflow/test_model_steps.py b/tests/integ/sagemaker/workflow/test_model_steps.py index f25723c440..da63bca597 100644 --- a/tests/integ/sagemaker/workflow/test_model_steps.py +++ b/tests/integ/sagemaker/workflow/test_model_steps.py @@ -16,8 +16,8 @@ import os import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.workflow.fail_step import FailStep from sagemaker.workflow.functions import Join 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 seconds_to_sleep=10, ): execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() is_execution_fail = False for step in execution_steps: @@ -284,10 +281,7 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference( seconds_to_sleep=10, ): execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() is_execution_fail = False for step in execution_steps: @@ -372,10 +366,7 @@ def test_mxnet_model_registration_with_custom_inference( seconds_to_sleep=10, ): execution = pipeline.start() - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -550,10 +541,7 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -667,10 +655,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( seconds_to_sleep=10, ): execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() is_execution_fail = False for step in execution_steps: @@ -750,10 +735,7 @@ def test_xgboost_model_register_and_deploy_with_runtime_repack( seconds_to_sleep=10, ): execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) # Verify the pipeline execution succeeded step_register_model = None @@ -866,10 +848,7 @@ def test_tensorflow_model_register_and_deploy_with_runtime_repack( seconds_to_sleep=10, ): execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) # Verify the pipeline execution succeeded step_register_model = None diff --git a/tests/integ/sagemaker/workflow/test_monitor_batch_transform_step.py b/tests/integ/sagemaker/workflow/test_monitor_batch_transform_step.py index e162c7cbc0..d3984415c3 100644 --- a/tests/integ/sagemaker/workflow/test_monitor_batch_transform_step.py +++ b/tests/integ/sagemaker/workflow/test_monitor_batch_transform_step.py @@ -18,7 +18,8 @@ import os import pytest -from botocore.exceptions import WaiterError + +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.clarify import ( BiasConfig, DataConfig, @@ -315,10 +316,7 @@ def test_monitor_batch_clarify_data_bias_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -419,10 +417,7 @@ def test_monitor_batch_clarify_data_bias_pipeline_bad_case( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -499,10 +494,7 @@ def test_batch_transform_data_quality_step_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -605,10 +597,7 @@ def test_batch_transform_data_quality_step_pipeline_failure_case( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -684,10 +673,7 @@ def test_batch_transform_model_quality_step_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -790,10 +776,7 @@ def test_batch_transform_model_quality_step_pipeline_failure_case( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -870,10 +853,7 @@ def test_batch_transform_data_quality_step_pipeline_before_transformation( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 @@ -978,10 +958,7 @@ def test_batch_transform_model_quality_step_pipeline_failure_no_violation_case( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 3 diff --git a/tests/integ/sagemaker/workflow/test_pipeline_var_behaviors.py b/tests/integ/sagemaker/workflow/test_pipeline_var_behaviors.py index f161901691..6e17e31ad1 100644 --- a/tests/integ/sagemaker/workflow/test_pipeline_var_behaviors.py +++ b/tests/integ/sagemaker/workflow/test_pipeline_var_behaviors.py @@ -13,8 +13,8 @@ from __future__ import absolute_import import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker import get_execution_role, utils from sagemaker.workflow.condition_step import ConditionStep from sagemaker.workflow.conditions import ConditionGreaterThan @@ -75,10 +75,7 @@ def test_ppl_var_to_string_and_add(sagemaker_session, role, pipeline_name): response = execution.describe() assert response["PipelineArn"] == pipeline_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -95,10 +92,7 @@ def test_ppl_var_to_string_and_add(sagemaker_session, role, pipeline_name): # Update int param to update cond step outcome execution = pipeline.start(parameters={"MyInteger": 0}) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2 diff --git a/tests/integ/sagemaker/workflow/test_processing_steps.py b/tests/integ/sagemaker/workflow/test_processing_steps.py index 238eff6123..656b8375df 100644 --- a/tests/integ/sagemaker/workflow/test_processing_steps.py +++ b/tests/integ/sagemaker/workflow/test_processing_steps.py @@ -20,7 +20,8 @@ from pathlib import Path import pytest -from botocore.exceptions import WaiterError + +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.workflow.utilities import hash_files_or_dirs, hash_object from sagemaker import image_uris, get_execution_role, utils @@ -270,10 +271,7 @@ def test_one_step_sklearn_processing_pipeline( assert response["Enabled"] == cache_config.enable_caching assert response["ExpireAfter"] == cache_config.expire_after - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sklearn-process" @@ -368,10 +366,7 @@ def test_one_step_framework_processing_pipeline( assert response["Enabled"] == cache_config.enable_caching assert response["ExpireAfter"] == cache_config.expire_after - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sklearn-process" @@ -468,10 +463,7 @@ def test_multi_step_framework_processing_pipeline_same_source_dir( assert entry_point_1 != entry_point_2 execution = pipeline.start(parameters={}) - try: - execution.wait(delay=540, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, delay=540, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -562,10 +554,7 @@ def test_multi_step_framework_processing_pipeline_different_source_dir( assert entry_point_1 != entry_point_2 execution = pipeline.start(parameters={}) - try: - execution.wait(delay=540, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, delay=540, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 2 @@ -667,10 +656,7 @@ def test_one_step_pyspark_processing_pipeline( assert response["Enabled"] == cache_config.enable_caching assert response["ExpireAfter"] == cache_config.expire_after - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "pyspark-process" @@ -772,10 +758,7 @@ def test_one_step_sparkjar_processing_pipeline( assert response["Enabled"] == cache_config.enable_caching assert response["ExpireAfter"] == cache_config.expire_after - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sparkjar-process" @@ -876,10 +859,7 @@ def test_one_step_data_wrangler_processing_pipeline(sagemaker_session, role, pip response = execution.describe() assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=60, max_attempts=10) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, delay=60, max_attempts=10) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -976,10 +956,7 @@ def test_two_processing_job_depends_on( response = execution.describe() assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, delay=60) execution_steps = execution.list_steps() assert len(execution_steps) == 2 diff --git a/tests/integ/sagemaker/workflow/test_quality_check_steps.py b/tests/integ/sagemaker/workflow/test_quality_check_steps.py index f521751e47..d34ab466e4 100644 --- a/tests/integ/sagemaker/workflow/test_quality_check_steps.py +++ b/tests/integ/sagemaker/workflow/test_quality_check_steps.py @@ -16,8 +16,8 @@ import os import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.workflow.parameters import ParameterString from tests.integ import DATA_DIR @@ -174,10 +174,7 @@ def test_one_step_data_quality_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -272,10 +269,7 @@ def test_one_step_data_quality_pipeline_constraint_violation( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "DataQualityCheckStep" @@ -354,10 +348,7 @@ def test_one_step_model_quality_pipeline_happycase( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 @@ -449,10 +440,7 @@ def test_one_step_model_quality_pipeline_constraint_violation( assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 diff --git a/tests/integ/sagemaker/workflow/test_retry.py b/tests/integ/sagemaker/workflow/test_retry.py index f2e6de6ba1..e4c6f52a60 100644 --- a/tests/integ/sagemaker/workflow/test_retry.py +++ b/tests/integ/sagemaker/workflow/test_retry.py @@ -18,7 +18,7 @@ import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.processing import ProcessingInput from sagemaker.session import get_execution_role from sagemaker.sklearn.processing import SKLearnProcessor @@ -153,10 +153,7 @@ def test_pipeline_execution_processing_step_with_retry( pipeline.create(role) execution = pipeline.start(parameters={}) - try: - execution.wait(delay=30, max_attempts=3) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, max_attempts=3) execution_steps = execution.list_steps() assert len(execution_steps) == 1 assert execution_steps[0]["StepName"] == "sklearn-process" @@ -267,10 +264,7 @@ def test_model_registration_with_model_repack( execution1 = pipeline.start(parameters={}) execution2 = pipeline.start(parameters={"GoodEnoughInput": 0}) - try: - execution1.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution1) execution1_steps = execution1.list_steps() for step in execution1_steps: assert not step.get("FailureReason", None) @@ -279,10 +273,7 @@ def test_model_registration_with_model_repack( assert step["Metadata"]["RegisterModel"] assert len(execution1_steps) == 4 - try: - execution2.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution2) execution2_steps = execution2.list_steps() for step in execution2_steps: assert not step.get("FailureReason", None) diff --git a/tests/integ/sagemaker/workflow/test_training_steps.py b/tests/integ/sagemaker/workflow/test_training_steps.py index 20a11957c9..181167ab31 100644 --- a/tests/integ/sagemaker/workflow/test_training_steps.py +++ b/tests/integ/sagemaker/workflow/test_training_steps.py @@ -17,8 +17,8 @@ import logging import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker import TrainingInput, get_execution_role, utils, image_uris from sagemaker.debugger import ( DebuggerHookConfig, @@ -333,10 +333,7 @@ def _start_and_verify_execution_with_retry(pipeline: Pipeline, parameters: dict) seconds_to_sleep=10, ): execution = pipeline.start(parameters=parameters) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 1 failure_reason = execution_steps[0].get("FailureReason", "") diff --git a/tests/integ/sagemaker/workflow/test_tuning_steps.py b/tests/integ/sagemaker/workflow/test_tuning_steps.py index 0abbe4d0e7..cfaa862245 100644 --- a/tests/integ/sagemaker/workflow/test_tuning_steps.py +++ b/tests/integ/sagemaker/workflow/test_tuning_steps.py @@ -16,8 +16,8 @@ import re import pytest -from botocore.exceptions import WaiterError +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.workflow.model_step import ModelStep, _CREATE_MODEL_NAME_BASE from sagemaker import TrainingInput, Model, get_execution_role, utils from sagemaker.dataset_definition import DatasetDefinition, AthenaDatasetDefinition @@ -187,10 +187,7 @@ def test_tuning_single_algo( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() for step in execution_steps: @@ -327,6 +324,7 @@ def test_tuning_multi_algos( ) execution = pipeline.start(parameters={}) + wait_pipeline_execution(execution=execution) assert re.match( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, diff --git a/tests/integ/sagemaker/workflow/test_workflow.py b/tests/integ/sagemaker/workflow/test_workflow.py index bd24b653ae..77022af0c3 100644 --- a/tests/integ/sagemaker/workflow/test_workflow.py +++ b/tests/integ/sagemaker/workflow/test_workflow.py @@ -22,9 +22,9 @@ from contextlib import contextmanager import pytest -from botocore.exceptions import WaiterError import pandas as pd +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from tests.integ.s3_utils import extract_files_from_s3 from sagemaker.workflow.model_step import ( ModelStep, @@ -588,10 +588,7 @@ def test_one_step_ingestion_pipeline( response = execution.describe() assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=60, max_attempts=10) - except WaiterError: - pass + wait_pipeline_execution(execution=execution, delay=60, max_attempts=10) execution_steps = execution.list_steps() @@ -1119,10 +1116,7 @@ def test_model_registration_with_tuning_model( rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/", execution.arn, ) - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() for step in execution_steps: diff --git a/tests/integ/sagemaker/workflow/test_workflow_with_clarify.py b/tests/integ/sagemaker/workflow/test_workflow_with_clarify.py index 599432c79a..4a7feb597b 100644 --- a/tests/integ/sagemaker/workflow/test_workflow_with_clarify.py +++ b/tests/integ/sagemaker/workflow/test_workflow_with_clarify.py @@ -21,7 +21,8 @@ import pytest import numpy as np import pandas as pd -from botocore.exceptions import WaiterError + +from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution from sagemaker.amazon.linear_learner import LinearLearner, LinearLearnerPredictor from sagemaker.clarify import ( BiasConfig, @@ -267,10 +268,7 @@ def test_workflow_with_clarify( response = execution.describe() assert response["PipelineArn"] == create_arn - try: - execution.wait(delay=30, max_attempts=60) - except WaiterError: - pass + wait_pipeline_execution(execution=execution) execution_steps = execution.list_steps() assert len(execution_steps) == 2