1212# language governing permissions and limitations under the License.
1313from __future__ import absolute_import
1414
15- import numpy
1615import os
16+
17+ import numpy
1718import pytest
19+
1820from sagemaker .mxnet .estimator import MXNet
1921from sagemaker .mxnet .model import MXNetModel
20- from sagemaker .utils import sagemaker_timestamp
22+ from sagemaker .utils import unique_name_from_base
2123from tests .integ import DATA_DIR , PYTHON_VERSION , TRAINING_DEFAULT_TIMEOUT_MINUTES
2224from tests .integ .timeout import timeout , timeout_and_delete_endpoint_by_name
23- import time
25+
26+ NEO_MXNET_VERSION = "1.4.1" # Neo doesn't support MXNet 1.6 yet.
2427
2528
2629@pytest .fixture (scope = "module" )
27- def mxnet_training_job (sagemaker_session , mxnet_full_version , cpu_instance_type ):
30+ def mxnet_training_job (sagemaker_session , cpu_instance_type ):
2831 with timeout (minutes = TRAINING_DEFAULT_TIMEOUT_MINUTES ):
2932 script_path = os .path .join (DATA_DIR , "mxnet_mnist" , "mnist_neo.py" )
3033 data_path = os .path .join (DATA_DIR , "mxnet_mnist" )
3134
3235 mx = MXNet (
3336 entry_point = script_path ,
3437 role = "SageMakerRole" ,
35- framework_version = mxnet_full_version ,
38+ framework_version = NEO_MXNET_VERSION ,
3639 py_version = PYTHON_VERSION ,
3740 train_instance_count = 1 ,
3841 train_instance_type = cpu_instance_type ,
@@ -52,13 +55,10 @@ def mxnet_training_job(sagemaker_session, mxnet_full_version, cpu_instance_type)
5255
5356@pytest .mark .canary_quick
5457@pytest .mark .regional_testing
55- @pytest .mark .skip (
56- reason = "This should be enabled along with the Boto SDK release for Neo API changes"
57- )
5858def test_attach_deploy (
5959 mxnet_training_job , sagemaker_session , cpu_instance_type , cpu_instance_family
6060):
61- endpoint_name = "test-mxnet -attach-deploy-{}" . format ( sagemaker_timestamp () )
61+ endpoint_name = unique_name_from_base ( "test-neo -attach-deploy" )
6262
6363 with timeout_and_delete_endpoint_by_name (endpoint_name , sagemaker_session ):
6464 estimator = MXNet .attach (mxnet_training_job , sagemaker_session = sagemaker_session )
@@ -77,13 +77,10 @@ def test_attach_deploy(
7777 predictor .predict (data )
7878
7979
80- @pytest .mark .skip (
81- reason = "This should be enabled along with the Boto SDK release for Neo API changes"
82- )
8380def test_deploy_model (
8481 mxnet_training_job , sagemaker_session , cpu_instance_type , cpu_instance_family
8582):
86- endpoint_name = "test-mxnet -deploy-model-{}" . format ( sagemaker_timestamp () )
83+ endpoint_name = unique_name_from_base ( "test-neo -deploy-model" )
8784
8885 with timeout_and_delete_endpoint_by_name (endpoint_name , sagemaker_session ):
8986 desc = sagemaker_session .sagemaker_client .describe_training_job (
@@ -97,14 +94,15 @@ def test_deploy_model(
9794 role ,
9895 entry_point = script_path ,
9996 py_version = PYTHON_VERSION ,
97+ framework_version = NEO_MXNET_VERSION ,
10098 sagemaker_session = sagemaker_session ,
10199 )
102100
103101 model .compile (
104102 target_instance_family = cpu_instance_family ,
105103 input_shape = {"data" : [1 , 1 , 28 , 28 ]},
106104 role = role ,
107- job_name = "test-deploy-model-compilation-job-{}" . format ( int ( time . time ()) ),
105+ job_name = unique_name_from_base ( "test-deploy-model-compilation-job" ),
108106 output_path = "/" .join (model_data .split ("/" )[:- 1 ]),
109107 )
110108 predictor = model .deploy (1 , cpu_instance_type , endpoint_name = endpoint_name )
0 commit comments