|
16 | 16 | import sagemaker |
17 | 17 | import os |
18 | 18 |
|
19 | | -from tests.integ import DATA_DIR |
| 19 | +from tests.integ import DATA_DIR, DEFAULT_TIMEOUT_MINUTES |
20 | 20 | from tests.integ.timeout import timeout |
21 | 21 | from stepfunctions.template import TrainingPipeline |
22 | 22 | from sagemaker.pytorch import PyTorch |
|
29 | 29 | get_resource_name_from_arn |
30 | 30 | ) |
31 | 31 |
|
32 | | -PIPELINE_TIMEOUT_LIMIT = 20 |
33 | | - |
34 | 32 | @pytest.fixture(scope="module") |
35 | 33 | def torch_estimator(sagemaker_role_arn): |
36 | 34 | script_path = os.path.join(DATA_DIR, "pytorch_mnist", "mnist.py") |
@@ -88,7 +86,7 @@ def _pipeline_teardown(sfn_client, sagemaker_session, endpoint_name, pipeline): |
88 | 86 |
|
89 | 87 |
|
90 | 88 | def test_torch_training_pipeline(sfn_client, sagemaker_client, torch_estimator, sagemaker_session, sfn_role_arn): |
91 | | - with timeout(minutes=PIPELINE_TIMEOUT_LIMIT): |
| 89 | + with timeout(minutes=DEFAULT_TIMEOUT_MINUTES): |
92 | 90 | # upload input data |
93 | 91 | data_path = os.path.join(DATA_DIR, "pytorch_mnist") |
94 | 92 | inputs = sagemaker_session.upload_data( |
@@ -123,7 +121,7 @@ def test_torch_training_pipeline(sfn_client, sagemaker_client, torch_estimator, |
123 | 121 |
|
124 | 122 |
|
125 | 123 | def test_sklearn_training_pipeline(sfn_client, sagemaker_client, sklearn_estimator, sagemaker_session, sfn_role_arn): |
126 | | - with timeout(minutes=PIPELINE_TIMEOUT_LIMIT): |
| 124 | + with timeout(minutes=DEFAULT_TIMEOUT_MINUTES): |
127 | 125 | # upload input data |
128 | 126 | data_path = os.path.join(DATA_DIR, "sklearn_mnist") |
129 | 127 | inputs = sagemaker_session.upload_data( |
|
0 commit comments