|
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