|  | 
| 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