Skip to content

Commit d854806

Browse files
nadiayanadiaya
committed
Update and refactor tests. Add tests for fw_utils.
1 parent c6b444a commit d854806

11 files changed

+404
-61
lines changed

‎tests/component/__init__.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.

‎tests/component/test_mxnet_estimator.py

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
import pytest
14+
from mock import Mock
15+
from sagemaker.mxnet import MXNet
16+
17+
18+
SCRIPT = 'resnet_cifar_10.py'
19+
TIMESTAMP = '2017-11-06-14:14:15.673'
20+
TIME = 1510006209.073025
21+
BUCKET_NAME = 'mybucket'
22+
INSTANCE_COUNT = 1
23+
INSTANCE_TYPE_GPU = 'ml.p2.xlarge'
24+
INSTANCE_TYPE_CPU = 'ml.m4.xlarge'
25+
CPU_IMAGE_NAME = 'sagemaker-mxnet-py2-cpu'
26+
GPU_IMAGE_NAME = 'sagemaker-mxnet-py2-gpu'
27+
REGION = 'us-west-2'
28+
IMAGE_URI_FORMAT_STRING = "520713654638.dkr.ecr.{}.amazonaws.com/{}:{}-{}-{}"
29+
REGION = 'us-west-2'
30+
ROLE = 'SagemakerRole'
31+
SOURCE_DIR = 's3://fefergerger'
32+
33+
34+
@pytest.fixture()
35+
def sagemaker_session():
36+
boto_mock = Mock(name='boto_session', region_name=REGION)
37+
ims = Mock(name='sagemaker_session', boto_session=boto_mock)
38+
ims.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
39+
ims.expand_role = Mock(name="expand_role", return_value=ROLE)
40+
ims.sagemaker_client.describe_training_job = Mock(return_value={'ModelArtifacts':
41+
{'S3ModelArtifacts': 's3://m/m.tar.gz'}})
42+
return ims
43+
44+
45+
# Test that we pass all necessary fields from estimator to the session when we call deploy
46+
def test_deploy(sagemaker_session, tf_version):
47+
estimator = MXNet(entry_point=SCRIPT, source_dir=SOURCE_DIR, role=ROLE,
48+
framework_version=tf_version,
49+
train_instance_count=2, train_instance_type=INSTANCE_TYPE_GPU,
50+
sagemaker_session=sagemaker_session,
51+
base_job_name='test-cifar')
52+
53+
estimator.fit('s3://mybucket/train')
54+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
55+
56+
estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU)
57+
image = IMAGE_URI_FORMAT_STRING.format(REGION, GPU_IMAGE_NAME, tf_version, 'cpu', 'py2')
58+
sagemaker_session.create_model.assert_called_with(
59+
estimator._current_job_name,
60+
ROLE,
61+
{'Environment':
62+
{'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
63+
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
64+
'SAGEMAKER_SUBMIT_DIRECTORY': SOURCE_DIR,
65+
'SAGEMAKER_REGION': REGION,
66+
'SAGEMAKER_PROGRAM': SCRIPT},
67+
'Image': image,
68+
'ModelDataUrl': 's3://m/m.tar.gz'})

‎tests/component/test_tf_estimator.py

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
import pytest
14+
from mock import Mock
15+
from sagemaker.tensorflow import TensorFlow
16+
17+
18+
SCRIPT = 'resnet_cifar_10.py'
19+
TIMESTAMP = '2017-11-06-14:14:15.673'
20+
TIME = 1510006209.073025
21+
BUCKET_NAME = 'mybucket'
22+
INSTANCE_COUNT = 1
23+
INSTANCE_TYPE_GPU = 'ml.p2.xlarge'
24+
INSTANCE_TYPE_CPU = 'ml.m4.xlarge'
25+
CPU_IMAGE_NAME = 'sagemaker-tensorflow-py2-cpu'
26+
GPU_IMAGE_NAME = 'sagemaker-tensorflow-py2-gpu'
27+
REGION = 'us-west-2'
28+
IMAGE_URI_FORMAT_STRING = "520713654638.dkr.ecr.{}.amazonaws.com/{}:{}-{}-{}"
29+
REGION = 'us-west-2'
30+
ROLE = 'SagemakerRole'
31+
SOURCE_DIR = 's3://fefergerger'
32+
33+
34+
@pytest.fixture()
35+
def sagemaker_session():
36+
boto_mock = Mock(name='boto_session', region_name=REGION)
37+
ims = Mock(name='sagemaker_session', boto_session=boto_mock)
38+
ims.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
39+
ims.expand_role = Mock(name="expand_role", return_value=ROLE)
40+
ims.sagemaker_client.describe_training_job = Mock(return_value={'ModelArtifacts':
41+
{'S3ModelArtifacts': 's3://m/m.tar.gz'}})
42+
return ims
43+
44+
45+
# Test that we pass all necessary fields from estimator to the session when we call deploy
46+
def test_deploy(sagemaker_session, tf_version):
47+
estimator = TensorFlow(entry_point=SCRIPT, source_dir=SOURCE_DIR, role=ROLE,
48+
framework_version=tf_version,
49+
train_instance_count=2, train_instance_type=INSTANCE_TYPE_CPU,
50+
sagemaker_session=sagemaker_session,
51+
base_job_name='test-cifar')
52+
53+
estimator.fit('s3://mybucket/train')
54+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
55+
56+
estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU)
57+
image = IMAGE_URI_FORMAT_STRING.format(REGION, GPU_IMAGE_NAME, tf_version, 'cpu', 'py2')
58+
sagemaker_session.create_model.assert_called_with(
59+
estimator._current_job_name,
60+
ROLE,
61+
{'Environment':
62+
{'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
63+
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
64+
'SAGEMAKER_SUBMIT_DIRECTORY': SOURCE_DIR,
65+
'SAGEMAKER_REGION': REGION,
66+
'SAGEMAKER_PROGRAM': SCRIPT},
67+
'Image': image,
68+
'ModelDataUrl': 's3://m/m.tar.gz'})

‎tests/conftest.py

+10
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,13 @@ def tf_version(request):
2121
@pytest.fixture(scope='module', params=["0.12", "0.12.1", "1.0", "1.0.0"])
2222
def mxnet_version(request):
2323
return request.param
24+
25+
26+
@pytest.fixture(scope='module', params=["1.4.1", "1.5.0"])
27+
def tf_full_version(request):
28+
return request.param
29+
30+
31+
@pytest.fixture(scope='module', params=["0.12.1", "1.0.0"])
32+
def mxnet_full_version(request):
33+
return request.param

‎tests/integ/test_mxnet_train.py

+6-6
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,12 @@ def sagemaker_session():
3030

3131

3232
@pytest.fixture(scope='module')
33-
def mxnet_training_job(sagemaker_session, mxnet_version):
33+
def mxnet_training_job(sagemaker_session, mxnet_full_version):
3434
with timeout(minutes=15):
3535
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
3636
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
3737

38-
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_version,
38+
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_full_version,
3939
train_instance_count=1, train_instance_type='ml.c4.xlarge',
4040
sagemaker_session=sagemaker_session)
4141

@@ -58,7 +58,7 @@ def test_attach_deploy(mxnet_training_job, sagemaker_session):
5858
predictor.predict(data)
5959

6060

61-
def test_async_fit(sagemaker_session, mxnet_version):
61+
def test_async_fit(sagemaker_session, mxnet_full_version):
6262

6363
training_job_name = ""
6464
endpoint_name = 'test-mxnet-attach-deploy-{}'.format(int(time.time()))
@@ -67,7 +67,7 @@ def test_async_fit(sagemaker_session, mxnet_version):
6767
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
6868
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
6969

70-
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_version,
70+
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_full_version,
7171
train_instance_count=1, train_instance_type='ml.c4.xlarge',
7272
sagemaker_session=sagemaker_session)
7373

@@ -104,12 +104,12 @@ def test_deploy_model(mxnet_training_job, sagemaker_session):
104104
predictor.predict(data)
105105

106106

107-
def test_failed_training_job(sagemaker_session, mxnet_version):
107+
def test_failed_training_job(sagemaker_session, mxnet_full_version):
108108
with timeout(minutes=15):
109109
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'failure_script.py')
110110
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
111111

112-
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_version,
112+
mx = MXNet(entry_point=script_path, role='SageMakerRole', framework_version=mxnet_full_version,
113113
train_instance_count=1, train_instance_type='ml.c4.xlarge',
114114
sagemaker_session=sagemaker_session)
115115

‎tests/integ/test_tf.py

+6-6
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,13 @@ def sagemaker_session():
2828
return Session(boto_session=boto3.Session(region_name=REGION))
2929

3030

31-
def test_tf(sagemaker_session, tf_version):
31+
def test_tf(sagemaker_session, tf_full_version):
3232
with timeout(minutes=15):
3333
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
3434

3535
estimator = TensorFlow(entry_point=script_path,
3636
role='SageMakerRole',
37-
framework_version=tf_version,
37+
framework_version=tf_full_version,
3838
training_steps=1,
3939
evaluation_steps=1,
4040
hyperparameters={'input_tensor_name': 'inputs'},
@@ -54,15 +54,15 @@ def test_tf(sagemaker_session, tf_version):
5454
print('predict result: {}'.format(result))
5555

5656

57-
def test_tf_async(sagemaker_session, tf_version):
57+
def test_tf_async(sagemaker_session, tf_full_version):
5858

5959
training_job_name = ""
6060
with timeout(minutes=5):
6161
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
6262

6363
estimator = TensorFlow(entry_point=script_path,
6464
role='SageMakerRole',
65-
framework_version=tf_version,
65+
framework_version=tf_full_version,
6666
training_steps=1,
6767
evaluation_steps=1,
6868
hyperparameters={'input_tensor_name': 'inputs'},
@@ -84,12 +84,12 @@ def test_tf_async(sagemaker_session, tf_version):
8484
print('predict result: {}'.format(result))
8585

8686

87-
def test_failed_tf_training(sagemaker_session, tf_version):
87+
def test_failed_tf_training(sagemaker_session, tf_full_version):
8888
with timeout(minutes=15):
8989
script_path = os.path.join(DATA_DIR, 'iris', 'failure_script.py')
9090
estimator = TensorFlow(entry_point=script_path,
9191
role='SageMakerRole',
92-
framework_version=tf_version,
92+
framework_version=tf_full_version,
9393
training_steps=1,
9494
evaluation_steps=1,
9595
hyperparameters={'input_tensor_name': 'inputs'},

‎tests/integ/test_tf_cifar.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,14 @@ def __call__(self, data):
3838
return pickle.dumps(data, protocol=2)
3939

4040

41-
def test_cifar(sagemaker_session, tf_version):
41+
def test_cifar(sagemaker_session, tf_full_version):
4242
with timeout(minutes=15):
4343
script_path = os.path.join(DATA_DIR, 'cifar_10', 'source')
4444

4545
dataset_path = os.path.join(DATA_DIR, 'cifar_10', 'data')
4646

4747
estimator = TensorFlow(entry_point='resnet_cifar_10.py', source_dir=script_path, role='SageMakerRole',
48-
framework_version=tf_version, training_steps=20, evaluation_steps=5,
48+
framework_version=tf_full_version, training_steps=20, evaluation_steps=5,
4949
train_instance_count=2, train_instance_type='ml.p2.xlarge',
5050
sagemaker_session=sagemaker_session,
5151
base_job_name='test-cifar')

‎tests/unit/test_estimator.py

-16
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
from mock import Mock, patch
1818

1919
from sagemaker.estimator import Estimator, Framework, _TrainingJob
20-
from sagemaker.fw_utils import framework_name_from_image
2120
from sagemaker.session import s3_input
2221
from sagemaker.model import FrameworkModel
2322
from sagemaker.predictor import RealTimePredictor
@@ -324,21 +323,6 @@ def test_init_with_source_dir_s3(strftime, sagemaker_session):
324323
assert actual_hyperparameter == expected_hyperparameters
325324

326325

327-
def test_framework_name_from_framework_image():
328-
image_name = '123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-py2-gpu:2.5.6-gpu-py2'
329-
framework, py_ver, tag = framework_name_from_image(image_name)
330-
assert framework == 'mxnet'
331-
assert py_ver == 'py2'
332-
assert tag == '2.5.6-gpu-py2'
333-
334-
335-
def test_framework_name_from_other():
336-
framework, py_ver, tag = framework_name_from_image('123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-myown-py2-gpu:1')
337-
assert framework is None
338-
assert py_ver is None
339-
assert tag is None
340-
341-
342326
# _TrainingJob 'utils'
343327
def test_format_input_single_unamed_channel():
344328
input_dict = _TrainingJob._format_inputs_to_input_config('s3://blah/blah')

0 commit comments

Comments
 (0)