Wrap training timeout for integ tests (#339)

yangaws · web-flow · commit e3dc3b5642c1 · 2018-08-13T18:24:26.000-07:00
diff --git a/tests/integ/__init__.py b/tests/integ/__init__.py
@@ -16,6 +16,8 @@
 import os
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
+TRAINING_DEFAULT_TIMEOUT_MINUTES = 20
+TUNING_DEFAULT_TIMEOUT_MINUTES = 20
 
 logging.getLogger('boto3').setLevel(logging.INFO)
 logging.getLogger('botocore').setLevel(logging.INFO)
diff --git a/tests/integ/test_byo_estimator.py b/tests/integ/test_byo_estimator.py
@@ -24,7 +24,7 @@
 from sagemaker.amazon.amazon_estimator import registry
 from sagemaker.estimator import Estimator
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
@@ -55,7 +55,7 @@ def test_byo_estimator(sagemaker_session, region):
     image_name = registry(region) + "/factorization-machines:1"
     training_data_path = os.path.join(DATA_DIR, 'dummy_tensor')
 
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_chainer_train.py b/tests/integ/test_chainer_train.py
@@ -22,7 +22,7 @@
 from sagemaker.chainer.estimator import Chainer
 from sagemaker.chainer.model import ChainerModel
 from sagemaker.utils import sagemaker_timestamp
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
@@ -40,7 +40,7 @@ def test_distributed_gpu_training(sagemaker_session, chainer_full_version):
 
 
 def test_training_with_additional_hyperparameters(sagemaker_session, chainer_full_version):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'chainer_mnist', 'mnist.py')
         data_path = os.path.join(DATA_DIR, 'chainer_mnist')
 
@@ -101,7 +101,7 @@ def test_async_fit(sagemaker_session):
 
 
 def test_failed_training_job(sagemaker_session, chainer_full_version):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'chainer_mnist', 'failure_script.py')
         data_path = os.path.join(DATA_DIR, 'chainer_mnist')
 
@@ -119,7 +119,7 @@ def test_failed_training_job(sagemaker_session, chainer_full_version):
 
 def _run_mnist_training_job(sagemaker_session, instance_type, instance_count,
                             chainer_full_version, wait=True):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
 
         script_path = os.path.join(DATA_DIR, 'chainer_mnist', 'mnist.py') if instance_type == 1 else \
             os.path.join(DATA_DIR, 'chainer_mnist', 'distributed_mnist.py')
diff --git a/tests/integ/test_factorization_machines.py b/tests/integ/test_factorization_machines.py
@@ -22,13 +22,13 @@
 
 from sagemaker import FactorizationMachines, FactorizationMachinesModel
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_factorization_machines(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_kmeans.py b/tests/integ/test_kmeans.py
@@ -22,13 +22,13 @@
 
 from sagemaker import KMeans, KMeansModel
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_kmeans(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_knn.py b/tests/integ/test_knn.py
@@ -22,13 +22,13 @@
 
 from sagemaker import KNN, KNNModel
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_knn_regressor(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_lda.py b/tests/integ/test_lda.py
@@ -20,14 +20,14 @@
 from sagemaker import LDA, LDAModel
 from sagemaker.amazon.common import read_records
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 from tests.integ.record_set import prepare_record_set_from_local_files
 
 
 @pytest.mark.continuous_testing
 def test_lda(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'lda')
         data_filename = 'nips-train_1.pbr'
 
diff --git a/tests/integ/test_linear_learner.py b/tests/integ/test_linear_learner.py
@@ -23,13 +23,13 @@
 
 from sagemaker.amazon.linear_learner import LinearLearner, LinearLearnerModel
 from sagemaker.utils import name_from_base, sagemaker_timestamp
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_linear_learner(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
@@ -93,7 +93,7 @@ def test_linear_learner(sagemaker_session):
 
 
 def test_linear_learner_multiclass(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
@@ -125,7 +125,7 @@ def test_async_linear_learner(sagemaker_session):
     training_job_name = ""
     endpoint_name = 'test-linear-learner-async-{}'.format(sagemaker_timestamp())
 
-    with timeout(minutes=5):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_mxnet_train.py b/tests/integ/test_mxnet_train.py
@@ -21,13 +21,13 @@
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.utils import sagemaker_timestamp
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.fixture(scope='module')
 def mxnet_training_job(sagemaker_session, mxnet_full_version):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
         data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
@@ -100,7 +100,7 @@ def test_async_fit(sagemaker_session):
 
 
 def test_failed_training_job(sagemaker_session, mxnet_full_version):
-    with timeout(minutes=15):
+    with timeout():
         script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'failure_script.py')
         data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
diff --git a/tests/integ/test_ntm.py b/tests/integ/test_ntm.py
@@ -20,14 +20,14 @@
 from sagemaker import NTM, NTMModel
 from sagemaker.amazon.common import read_records
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 from tests.integ.record_set import prepare_record_set_from_local_files
 
 
 @pytest.mark.continuous_testing
 def test_ntm(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'ntm')
         data_filename = 'nips-train_1.pbr'
 
diff --git a/tests/integ/test_pca.py b/tests/integ/test_pca.py
@@ -22,13 +22,13 @@
 
 import sagemaker.amazon.pca
 from sagemaker.utils import name_from_base
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_pca(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
diff --git a/tests/integ/test_pytorch_train.py b/tests/integ/test_pytorch_train.py
@@ -19,7 +19,7 @@
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.pytorch.model import PyTorchModel
 from sagemaker.utils import sagemaker_timestamp
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 MNIST_DIR = os.path.join(DATA_DIR, 'pytorch_mnist')
@@ -30,7 +30,7 @@
 @pytest.fixture(scope='module', name='pytorch_training_job')
 def fixture_training_job(sagemaker_session, pytorch_full_version):
     instance_type = 'ml.c4.xlarge'
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         pytorch = _get_pytorch_estimator(sagemaker_session, pytorch_full_version, instance_type)
 
         pytorch.fit({'training': _upload_training_data(pytorch)})
@@ -103,7 +103,7 @@ def test_async_fit_deploy(sagemaker_session, pytorch_full_version):
 def test_failed_training_job(sagemaker_session, pytorch_full_version):
     script_path = os.path.join(MNIST_DIR, 'failure_script.py')
 
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         pytorch = _get_pytorch_estimator(sagemaker_session, pytorch_full_version, entry_point=script_path)
 
         with pytest.raises(ValueError) as e:
diff --git a/tests/integ/test_randomcutforest.py b/tests/integ/test_randomcutforest.py
@@ -17,12 +17,13 @@
 
 from sagemaker import RandomCutForest, RandomCutForestModel
 from sagemaker.utils import name_from_base
+from tests.integ import TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
 @pytest.mark.continuous_testing
 def test_randomcutforest(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         # Generate a thousand 14-dimensional datapoints.
         feature_num = 14
         train_input = np.random.rand(1000, feature_num)
diff --git a/tests/integ/test_tf.py b/tests/integ/test_tf.py
@@ -18,15 +18,15 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout
 
 DATA_PATH = os.path.join(DATA_DIR, 'iris', 'data')
 
 
 @pytest.mark.continuous_testing
 def test_tf(sagemaker_session, tf_full_version):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
 
         estimator = TensorFlow(entry_point=script_path,
@@ -59,7 +59,7 @@ def test_tf(sagemaker_session, tf_full_version):
 
 
 def test_tf_async(sagemaker_session):
-    with timeout(minutes=5):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
 
         estimator = TensorFlow(entry_point=script_path,
@@ -88,7 +88,7 @@ def test_tf_async(sagemaker_session):
 
 
 def test_failed_tf_training(sagemaker_session, tf_full_version):
-    with timeout(minutes=15):
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'iris', 'failure_script.py')
         estimator = TensorFlow(entry_point=script_path,
                                role='SageMakerRole',
diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py
@@ -32,7 +32,7 @@
 from sagemaker.pytorch import PyTorch
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.tuner import IntegerParameter, ContinuousParameter, CategoricalParameter, HyperparameterTuner
-from tests.integ import DATA_DIR
+from tests.integ import DATA_DIR, TUNING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.record_set import prepare_record_set_from_local_files
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
@@ -41,7 +41,7 @@
 
 @pytest.mark.continuous_testing
 def test_tuning_kmeans(sagemaker_session):
-    with timeout(minutes=20):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
@@ -96,7 +96,7 @@ def test_tuning_kmeans(sagemaker_session):
 
 
 def test_tuning_lda(sagemaker_session):
-    with timeout(minutes=20):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'lda')
         data_filename = 'nips-train_1.pbr'
 
@@ -182,7 +182,7 @@ def test_stop_tuning_job(sagemaker_session):
 
 @pytest.mark.continuous_testing
 def test_tuning_mxnet(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'tuning.py')
         data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
 
@@ -219,7 +219,7 @@ def test_tuning_mxnet(sagemaker_session):
 
 @pytest.mark.continuous_testing
 def test_tuning_tf(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
 
         estimator = TensorFlow(entry_point=script_path,
@@ -263,7 +263,7 @@ def test_tuning_tf(sagemaker_session):
 
 @pytest.mark.continuous_testing
 def test_tuning_chainer(sagemaker_session):
-    with timeout(minutes=15):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, 'chainer_mnist', 'mnist.py')
         data_path = os.path.join(DATA_DIR, 'chainer_mnist')
 
@@ -321,7 +321,7 @@ def test_attach_tuning_pytorch(sagemaker_session):
     estimator = PyTorch(entry_point=mnist_script, role='SageMakerRole', train_instance_count=1,
                         train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session)
 
-    with timeout(minutes=15):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         objective_metric_name = 'evaluation-accuracy'
         metric_definitions = [{'Name': 'evaluation-accuracy', 'Regex': 'Overall test accuracy: (\d+)'}]
         hyperparameter_ranges = {'batch-size': IntegerParameter(50, 100)}
@@ -368,7 +368,7 @@ def test_tuning_byo_estimator(sagemaker_session):
     image_name = registry(sagemaker_session.boto_session.region_name) + '/factorization-machines:1'
     training_data_path = os.path.join(DATA_DIR, 'dummy_tensor')
 
-    with timeout(minutes=15):
+    with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
         pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
 
