|
18 | 18 | import pytest
|
19 | 19 |
|
20 | 20 | from sagemaker.tensorflow import TensorFlow
|
21 |
| -from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES |
| 21 | +from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, PYTHON_VERSION |
22 | 22 | from tests.integ.timeout import timeout_and_delete_endpoint_by_name, timeout
|
23 | 23 | from tests.integ.vpc_utils import get_or_create_subnet_and_security_group
|
24 | 24 |
|
|
27 | 27 |
|
28 | 28 |
|
29 | 29 | @pytest.mark.continuous_testing
|
| 30 | +@pytest.mark.skipif(PYTHON_VERSION != 'py2', reason="TensorFlow image supports only python 2.") |
30 | 31 | def test_tf(sagemaker_session, tf_full_version):
|
31 | 32 | with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
|
32 | 33 | script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
|
@@ -60,6 +61,7 @@ def test_tf(sagemaker_session, tf_full_version):
|
60 | 61 | assert dict_result == list_result
|
61 | 62 |
|
62 | 63 |
|
| 64 | +@pytest.mark.skipif(PYTHON_VERSION != 'py2', reason="TensorFlow image supports only python 2.") |
63 | 65 | def test_tf_async(sagemaker_session):
|
64 | 66 | with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
|
65 | 67 | script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
|
@@ -89,6 +91,7 @@ def test_tf_async(sagemaker_session):
|
89 | 91 | print('predict result: {}'.format(result))
|
90 | 92 |
|
91 | 93 |
|
| 94 | +@pytest.mark.skipif(PYTHON_VERSION != 'py2', reason="TensorFlow image supports only python 2.") |
92 | 95 | def test_failed_tf_training(sagemaker_session, tf_full_version):
|
93 | 96 | with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
|
94 | 97 | script_path = os.path.join(DATA_DIR, 'iris', 'failure_script.py')
|
|
0 commit comments