|
12 | 12 | # language governing permissions and limitations under the License.
|
13 | 13 | from __future__ import absolute_import
|
14 | 14 |
|
15 |
| -import os |
16 |
| - |
17 | 15 | import numpy
|
| 16 | +import os |
18 | 17 | import pytest
|
19 |
| -from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES |
20 |
| -from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name |
21 |
| - |
| 18 | +from sagemaker.pytorch.defaults import LATEST_PY2_VERSION |
22 | 19 | from sagemaker.pytorch.estimator import PyTorch
|
23 | 20 | from sagemaker.pytorch.model import PyTorchModel
|
24 |
| -from sagemaker.pytorch.defaults import LATEST_PY2_VERSION |
25 | 21 | from sagemaker.utils import sagemaker_timestamp
|
26 | 22 |
|
| 23 | +from tests.integ import ( |
| 24 | + test_region, |
| 25 | + DATA_DIR, |
| 26 | + PYTHON_VERSION, |
| 27 | + TRAINING_DEFAULT_TIMEOUT_MINUTES, |
| 28 | + EI_SUPPORTED_REGIONS, |
| 29 | +) |
| 30 | +from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name |
| 31 | + |
27 | 32 | MNIST_DIR = os.path.join(DATA_DIR, "pytorch_mnist")
|
28 | 33 | MNIST_SCRIPT = os.path.join(MNIST_DIR, "mnist.py")
|
29 | 34 |
|
@@ -121,8 +126,7 @@ def test_deploy_model(pytorch_training_job, sagemaker_session, cpu_instance_type
|
121 | 126 |
|
122 | 127 | @pytest.mark.skipif(PYTHON_VERSION == "py2", reason="PyTorch EIA does not support Python 2.")
|
123 | 128 | @pytest.mark.skipif(
|
124 |
| - tests.integ.test_region() not in tests.integ.EI_SUPPORTED_REGIONS, |
125 |
| - reason="EI isn't supported in that specific region.", |
| 129 | + test_region() not in EI_SUPPORTED_REGIONS, reason="EI isn't supported in that specific region." |
126 | 130 | )
|
127 | 131 | def test_deploy_model_with_accelerator(sagemaker_session, cpu_instance_type):
|
128 | 132 | endpoint_name = "test-pytorch-deploy-eia-{}".format(sagemaker_timestamp())
|
|
0 commit comments