Skip to content

Commit 6473b0c

Browse files
iquinteroiquintero
authored
Add Local Mode Integration tests (#192)
1 parent aa4a05f commit 6473b0c

File tree

2 files changed

+302
-0
lines changed

2 files changed

+302
-0
lines changed

‎tests/conftest.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import pytest
1919

2020
from sagemaker import Session
21+
from sagemaker.local import LocalSession
2122

2223
DEFAULT_REGION = 'us-west-2'
2324

@@ -58,6 +59,15 @@ def sagemaker_session(sagemaker_client_config, sagemaker_runtime_config, boto_co
5859
sagemaker_runtime_client=runtime_client)
5960

6061

62+
@pytest.fixture(scope='session')
63+
def sagemaker_local_session(boto_config):
64+
if boto_config:
65+
boto_session = boto3.Session(**boto_config)
66+
else:
67+
boto_session = boto3.Session(region_name=DEFAULT_REGION)
68+
return LocalSession(boto_session=boto_session)
69+
70+
6171
@pytest.fixture(scope='module', params=['1.4', '1.4.1', '1.5', '1.5.0', '1.6', '1.6.0'])
6272
def tf_version(request):
6373
return request.param

‎tests/integ/test_local_mode.py

Lines changed: 292 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,292 @@
1+
# Copyright 2017-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+
from __future__ import absolute_import
14+
15+
import fcntl
16+
import os
17+
import time
18+
19+
import boto3
20+
import numpy
21+
22+
from sagemaker.local import LocalSession, LocalSagemakerRuntimeClient, LocalSagemakerClient
23+
from sagemaker.mxnet import MXNet
24+
from sagemaker.tensorflow import TensorFlow
25+
from tests.integ import DATA_DIR
26+
from tests.integ.timeout import timeout
27+
28+
DATA_PATH = os.path.join(DATA_DIR, 'iris', 'data')
29+
LOCK_PATH = os.path.join(DATA_DIR, 'local_mode_lock')
30+
DEFAULT_REGION = 'us-west-2'
31+
32+
33+
class LocalNoS3Session(LocalSession):
34+
"""
35+
This Session sets local_code: True regardless of any config file settings
36+
"""
37+
def __init__(self):
38+
super(LocalSession, self).__init__()
39+
40+
def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
41+
self.boto_session = boto3.Session(region_name=DEFAULT_REGION)
42+
if self.config is None:
43+
self.config = {
44+
'local':
45+
{
46+
'local_code': True,
47+
'region_name': DEFAULT_REGION
48+
}
49+
}
50+
51+
self._region_name = DEFAULT_REGION
52+
self.sagemaker_client = LocalSagemakerClient(self)
53+
self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
54+
self.local_mode = True
55+
56+
57+
def test_tf_local_mode(tf_full_version, sagemaker_local_session):
58+
local_mode_lock_fd = open(LOCK_PATH, 'w')
59+
local_mode_lock = local_mode_lock_fd.fileno()
60+
with timeout(minutes=5):
61+
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
62+
63+
estimator = TensorFlow(entry_point=script_path,
64+
role='SageMakerRole',
65+
framework_version=tf_full_version,
66+
training_steps=1,
67+
evaluation_steps=1,
68+
hyperparameters={'input_tensor_name': 'inputs'},
69+
train_instance_count=1,
70+
train_instance_type='local',
71+
base_job_name='test-tf',
72+
sagemaker_session=sagemaker_local_session)
73+
74+
inputs = estimator.sagemaker_session.upload_data(path=DATA_PATH,
75+
key_prefix='integ-test-data/tf_iris')
76+
estimator.fit(inputs)
77+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
78+
79+
endpoint_name = estimator.latest_training_job.name
80+
try:
81+
# Since Local Mode uses the same port for serving, we need a lock in order
82+
# to allow concurrent test execution. The serving test is really fast so it still
83+
# makes sense to allow this behavior.
84+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
85+
json_predictor = estimator.deploy(initial_instance_count=1,
86+
instance_type='local',
87+
endpoint_name=endpoint_name)
88+
89+
features = [6.4, 3.2, 4.5, 1.5]
90+
dict_result = json_predictor.predict({'inputs': features})
91+
print('predict result: {}'.format(dict_result))
92+
list_result = json_predictor.predict(features)
93+
print('predict result: {}'.format(list_result))
94+
95+
assert dict_result == list_result
96+
finally:
97+
estimator.delete_endpoint()
98+
time.sleep(5)
99+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
100+
101+
102+
def test_tf_distributed_local_mode(sagemaker_local_session):
103+
local_mode_lock_fd = open(LOCK_PATH, 'w')
104+
local_mode_lock = local_mode_lock_fd.fileno()
105+
with timeout(minutes=5):
106+
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
107+
108+
estimator = TensorFlow(entry_point=script_path,
109+
role='SageMakerRole',
110+
training_steps=1,
111+
evaluation_steps=1,
112+
hyperparameters={'input_tensor_name': 'inputs'},
113+
train_instance_count=3,
114+
train_instance_type='local',
115+
base_job_name='test-tf',
116+
sagemaker_session=sagemaker_local_session)
117+
118+
inputs = 'file://' + DATA_PATH
119+
estimator.fit(inputs)
120+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
121+
122+
endpoint_name = estimator.latest_training_job.name
123+
124+
try:
125+
# Since Local Mode uses the same port for serving, we need a lock in order
126+
# to allow concurrent test execution. The serving test is really fast so it still
127+
# makes sense to allow this behavior.
128+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
129+
json_predictor = estimator.deploy(initial_instance_count=1,
130+
instance_type='local',
131+
endpoint_name=endpoint_name)
132+
133+
features = [6.4, 3.2, 4.5, 1.5]
134+
dict_result = json_predictor.predict({'inputs': features})
135+
print('predict result: {}'.format(dict_result))
136+
list_result = json_predictor.predict(features)
137+
print('predict result: {}'.format(list_result))
138+
139+
assert dict_result == list_result
140+
finally:
141+
estimator.delete_endpoint()
142+
time.sleep(5)
143+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
144+
145+
146+
def test_tf_local_data(sagemaker_local_session):
147+
local_mode_lock_fd = open(LOCK_PATH, 'w')
148+
local_mode_lock = local_mode_lock_fd.fileno()
149+
with timeout(minutes=5):
150+
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
151+
152+
estimator = TensorFlow(entry_point=script_path,
153+
role='SageMakerRole',
154+
training_steps=1,
155+
evaluation_steps=1,
156+
hyperparameters={'input_tensor_name': 'inputs'},
157+
train_instance_count=1,
158+
train_instance_type='local',
159+
base_job_name='test-tf',
160+
sagemaker_session=sagemaker_local_session)
161+
162+
inputs = 'file://' + DATA_PATH
163+
estimator.fit(inputs)
164+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
165+
166+
endpoint_name = estimator.latest_training_job.name
167+
try:
168+
# Since Local Mode uses the same port for serving, we need a lock in order
169+
# to allow concurrent test execution. The serving test is really fast so it still
170+
# makes sense to allow this behavior.
171+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
172+
json_predictor = estimator.deploy(initial_instance_count=1,
173+
instance_type='local',
174+
endpoint_name=endpoint_name)
175+
176+
features = [6.4, 3.2, 4.5, 1.5]
177+
dict_result = json_predictor.predict({'inputs': features})
178+
print('predict result: {}'.format(dict_result))
179+
list_result = json_predictor.predict(features)
180+
print('predict result: {}'.format(list_result))
181+
182+
assert dict_result == list_result
183+
finally:
184+
estimator.delete_endpoint()
185+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
186+
187+
188+
def test_tf_local_data_local_script():
189+
local_mode_lock_fd = open(LOCK_PATH, 'w')
190+
local_mode_lock = local_mode_lock_fd.fileno()
191+
with timeout(minutes=5):
192+
193+
script_path = os.path.join(DATA_DIR, 'iris', 'iris-dnn-classifier.py')
194+
195+
estimator = TensorFlow(entry_point=script_path,
196+
role='SageMakerRole',
197+
training_steps=1,
198+
evaluation_steps=1,
199+
hyperparameters={'input_tensor_name': 'inputs'},
200+
train_instance_count=1,
201+
train_instance_type='local',
202+
base_job_name='test-tf',
203+
sagemaker_session=LocalNoS3Session())
204+
205+
inputs = 'file://' + DATA_PATH
206+
207+
estimator.fit(inputs)
208+
print('job succeeded: {}'.format(estimator.latest_training_job.name))
209+
210+
endpoint_name = estimator.latest_training_job.name
211+
try:
212+
# Since Local Mode uses the same port for serving, we need a lock in order
213+
# to allow concurrent test execution. The serving test is really fast so it still
214+
# makes sense to allow this behavior.
215+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
216+
json_predictor = estimator.deploy(initial_instance_count=1,
217+
instance_type='local',
218+
endpoint_name=endpoint_name)
219+
220+
features = [6.4, 3.2, 4.5, 1.5]
221+
dict_result = json_predictor.predict({'inputs': features})
222+
print('predict result: {}'.format(dict_result))
223+
list_result = json_predictor.predict(features)
224+
print('predict result: {}'.format(list_result))
225+
226+
assert dict_result == list_result
227+
finally:
228+
estimator.delete_endpoint()
229+
time.sleep(5)
230+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
231+
232+
233+
def test_mxnet_local_mode(sagemaker_local_session):
234+
local_mode_lock_fd = open(LOCK_PATH, 'w')
235+
local_mode_lock = local_mode_lock_fd.fileno()
236+
237+
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
238+
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
239+
240+
mx = MXNet(entry_point=script_path, role='SageMakerRole',
241+
train_instance_count=1, train_instance_type='local',
242+
sagemaker_session=sagemaker_local_session)
243+
244+
train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
245+
key_prefix='integ-test-data/mxnet_mnist/train')
246+
test_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'test'),
247+
key_prefix='integ-test-data/mxnet_mnist/test')
248+
249+
mx.fit({'train': train_input, 'test': test_input})
250+
endpoint_name = mx.latest_training_job.name
251+
try:
252+
# Since Local Mode uses the same port for serving, we need a lock in order
253+
# to allow concurrent test execution. The serving test is really fast so it still
254+
# makes sense to allow this behavior.
255+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
256+
predictor = mx.deploy(1, 'local', endpoint_name=endpoint_name)
257+
data = numpy.zeros(shape=(1, 1, 28, 28))
258+
predictor.predict(data)
259+
finally:
260+
mx.delete_endpoint()
261+
time.sleep(5)
262+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)
263+
264+
265+
def test_mxnet_local_data_local_script():
266+
local_mode_lock_fd = open(LOCK_PATH, 'w')
267+
local_mode_lock = local_mode_lock_fd.fileno()
268+
269+
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
270+
data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
271+
272+
mx = MXNet(entry_point=script_path, role='SageMakerRole',
273+
train_instance_count=1, train_instance_type='local',
274+
sagemaker_session=LocalNoS3Session())
275+
276+
train_input = 'file://' + os.path.join(data_path, 'train')
277+
test_input = 'file://' + os.path.join(data_path, 'test')
278+
279+
mx.fit({'train': train_input, 'test': test_input})
280+
endpoint_name = mx.latest_training_job.name
281+
try:
282+
# Since Local Mode uses the same port for serving, we need a lock in order
283+
# to allow concurrent test execution. The serving test is really fast so it still
284+
# makes sense to allow this behavior.
285+
fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
286+
predictor = mx.deploy(1, 'local', endpoint_name=endpoint_name)
287+
data = numpy.zeros(shape=(1, 1, 28, 28))
288+
predictor.predict(data)
289+
finally:
290+
mx.delete_endpoint()
291+
time.sleep(5)
292+
fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)

0 commit comments

Comments
 (0)