Skip to content

Commit bae1816

Browse files
sachanubsachanub
authored
Add new unit and integration tests (#155)
1 parent b508803 commit bae1816

24 files changed

+947
-79
lines changed

‎buildspec.yml

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,10 @@ phases:
5555
# build DLC GPU image because the base DLC image is too big and takes too long to build as part of the test
5656
- python3 setup.py sdist
5757
- $(aws ecr get-login --registry-ids $DLC_ACCOUNT --no-include-email --region $AWS_DEFAULT_REGION)
58+
- create-key-pair
5859
- |
5960
for FRAMEWORK_VERSION in $FRAMEWORK_VERSIONS;
6061
do
61-
create-key-pair;
6262
launch-ec2-instance --instance-type $instance_type --ami-name ami-03e3ef8c92fdb39ad;
6363
DLC_GPU_TAG="$FRAMEWORK_VERSION-dlc-gpu-$BUILD_ID";
6464
build_dir="test/container/$FRAMEWORK_VERSION";
@@ -71,8 +71,10 @@ phases:
7171
execute-command-if-has-matching-changes "$test_cmd" "test/" "src/*.py" "setup.py" "setup.cfg";
7272
docker system prune --all --force;
7373
cleanup-gpu-instances;
74-
cleanup-key-pairs;
74+
rm ~/.instance_id;
75+
rm ~/.ip_address;
7576
done
77+
- cleanup-key-pairs;
7678

7779
# run CPU sagemaker integration tests
7880
- |

‎test/integration/__init__.py

Lines changed: 35 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -13,77 +13,55 @@
1313
from __future__ import absolute_import
1414

1515
import os
16+
import json
1617

1718
from utils import file_utils
1819

1920
resources_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'resources'))
21+
model_gpu_context_dir = os.path.join(resources_path, 'model_gpu_context')
2022
mnist_path = os.path.join(resources_path, 'mnist')
21-
resnet18_path = os.path.join(resources_path, 'resnet18')
2223
data_dir = os.path.join(mnist_path, 'data')
2324
training_dir = os.path.join(data_dir, 'training')
24-
cpu_sub_dir = 'model_cpu'
25-
gpu_sub_dir = 'model_gpu'
26-
eia_sub_dir = 'model_eia'
27-
code_sub_dir = 'code'
28-
default_sub_dir = 'default_model'
29-
default_sub_eia_dir = 'default_model_eia'
30-
default_sub_traced_resnet_dir = 'default_traced_resnet'
3125

32-
model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
33-
mnist_cpu_script = os.path.join(model_cpu_dir, code_sub_dir, 'mnist.py')
34-
model_cpu_tar = file_utils.make_tarfile(mnist_cpu_script,
35-
os.path.join(model_cpu_dir, "torch_model.pth"),
36-
model_cpu_dir,
37-
script_path="code")
26+
all_models_info_json = os.path.abspath(os.path.join(os.path.dirname(__file__), 'all_models_info.json'))
3827

39-
model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
40-
mnist_1d_script = os.path.join(model_cpu_1d_dir, code_sub_dir, 'mnist_1d.py')
41-
model_cpu_1d_tar = file_utils.make_tarfile(mnist_1d_script,
42-
os.path.join(model_cpu_1d_dir, "torch_model.pth"),
43-
model_cpu_1d_dir,
44-
script_path="code")
28+
with open(all_models_info_json, 'r') as json_file:
29+
all_models_info = json.load(json_file)
4530

46-
model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
47-
mnist_gpu_script = os.path.join(model_gpu_dir, code_sub_dir, 'mnist.py')
48-
model_gpu_tar = file_utils.make_tarfile(mnist_gpu_script,
49-
os.path.join(model_gpu_dir, "torch_model.pth"),
50-
model_gpu_dir,
51-
script_path="code")
31+
for model_name in all_models_info.keys():
32+
model_info = all_models_info[model_name]
5233

53-
model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
54-
mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
55-
model_eia_tar = file_utils.make_tarfile(mnist_eia_script,
56-
os.path.join(model_eia_dir, "torch_model.pth"),
57-
model_eia_dir)
34+
dir_path = model_info['dir_path']
35+
model_dir = os.path.join(resources_path, *dir_path)
36+
setattr(__import__('integration'), model_name + '_dir', model_dir)
5837

59-
call_model_fn_once_script = os.path.join(model_cpu_dir, code_sub_dir, 'call_model_fn_once.py')
60-
call_model_fn_once_tar = file_utils.make_tarfile(call_model_fn_once_script,
61-
os.path.join(model_cpu_dir, "torch_model.pth"),
62-
model_cpu_dir,
63-
"model_call_model_fn_once.tar.gz",
64-
script_path="code")
38+
script_name = model_info['script_name']
39+
model = model_info['model']
40+
if 'filename' in model_info:
41+
filename = model_info['filename']
42+
else:
43+
filename = 'model.tar.gz'
44+
if 'code_path' in model_info:
45+
code_path = model_info['code_path']
46+
script_path = model_info['code_path']
47+
else:
48+
code_path = ''
49+
script_path = None
50+
if 'requirements' in model_info:
51+
requirements = os.path.join(model_dir, code_path, 'requirements.txt')
52+
else:
53+
requirements = None
6554

66-
default_model_dir = os.path.join(resnet18_path, default_sub_dir)
67-
default_model_script = os.path.join(default_model_dir, code_sub_dir, "resnet18.py")
68-
default_model_tar = file_utils.make_tarfile(
69-
default_model_script, os.path.join(default_model_dir, "model.pt"), default_model_dir, script_path="code"
70-
)
55+
model_script = os.path.join(model_dir, code_path, script_name)
56+
model_tar = file_utils.make_tarfile(model_script,
57+
os.path.join(model_dir, model),
58+
model_dir,
59+
filename=filename,
60+
script_path=script_path,
61+
requirements=requirements)
7162

72-
default_traced_resnet_dir = os.path.join(resnet18_path, default_sub_traced_resnet_dir)
73-
default_traced_resnet_script = os.path.join(default_traced_resnet_dir, code_sub_dir, "resnet18.py")
74-
default_model_traced_resnet18_tar = file_utils.make_tarfile(
75-
default_traced_resnet_script,
76-
os.path.join(default_traced_resnet_dir, "traced_resnet18.pt"),
77-
default_traced_resnet_dir,
78-
filename="traced_resnet18.tar.gz",
79-
script_path="code",
80-
)
81-
82-
default_model_eia_dir = os.path.join(mnist_path, default_sub_eia_dir)
83-
default_model_eia_script = os.path.join(default_model_eia_dir, code_sub_dir, "mnist.py")
84-
default_model_eia_tar = file_utils.make_tarfile(
85-
default_model_eia_script, os.path.join(default_model_eia_dir, "model.pt"), default_model_eia_dir
86-
)
63+
setattr(__import__('integration'), model_name + '_script', model_script)
64+
setattr(__import__('integration'), model_name + '_tar', model_tar)
8765

8866
ROLE = 'dummy/unused-role'
8967
DEFAULT_TIMEOUT = 20

‎test/integration/all_models_info.json

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
{
2+
"mnist_cpu":
3+
{
4+
"script_name": "mnist.py",
5+
"model": "torch_model.pth",
6+
"code_path": "code",
7+
"requirements": "requirements.txt",
8+
"dir_path": ["mnist", "model_cpu"]
9+
},
10+
"mnist_cpu_1d":
11+
{
12+
"script_name": "mnist_1d.py",
13+
"model": "torch_model.pth",
14+
"code_path": "code",
15+
"dir_path": ["mnist", "model_cpu", "1d"]
16+
},
17+
"mnist_gpu":
18+
{
19+
"script_name": "mnist.py",
20+
"model": "torch_model.pth",
21+
"code_path": "code",
22+
"dir_path": ["mnist", "model_gpu"]
23+
},
24+
"mnist_eia":
25+
{
26+
"script_name": "mnist.py",
27+
"model": "torch_model.pth",
28+
"dir_path": ["mnist", "model_eia"]
29+
},
30+
"mnist_inductor":
31+
{
32+
"script_name": "mnist.py",
33+
"model": "torch_model.pth",
34+
"code_path": "code",
35+
"dir_path": ["mnist", "model_inductor"]
36+
},
37+
"call_model_fn_once":
38+
{
39+
"script_name": "call_model_fn_once.py",
40+
"model": "torch_model.pth",
41+
"filename": "model_call_model_fn_once.tar.gz",
42+
"code_path": "code",
43+
"dir_path": ["mnist", "model_cpu"]
44+
},
45+
"default_model":
46+
{
47+
"script_name": "resnet18.py",
48+
"model": "model.pt",
49+
"code_path": "code",
50+
"dir_path": ["resnet18", "default_model"]
51+
},
52+
"default_model_traced_resnet":
53+
{
54+
"script_name": "resnet18.py",
55+
"model": "traced_resnet18.pt",
56+
"filename": "traced_resnet18.tar.gz",
57+
"code_path": "code",
58+
"dir_path": ["resnet18", "default_traced_resnet"]
59+
},
60+
"default_model_eia":
61+
{
62+
"script_name": "mnist.py",
63+
"model": "model.pt",
64+
"code_path": "code",
65+
"dir_path": ["mnist", "default_model_eia"]
66+
},
67+
"resnet18":
68+
{
69+
"script_name": "inference.py",
70+
"model": "model.pt",
71+
"filename": "resnet18.tar.gz",
72+
"code_path": "code",
73+
"dir_path": ["mme", "resnet18"]
74+
},
75+
"traced_resnet18":
76+
{
77+
"script_name": "inference.py",
78+
"model": "traced_resnet18.pt",
79+
"filename": "traced_resnet18.tar.gz",
80+
"code_path": "code",
81+
"dir_path": ["mme", "traced_resnet18"]
82+
}
83+
}

‎test/integration/local/test_serving.py renamed to ‎test/integration/local/test_mnist_serving.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@
2525
from sagemaker_inference import content_types
2626
from torchvision import datasets, transforms
2727

28-
from integration import training_dir, mnist_1d_script, model_cpu_tar, mnist_cpu_script, \
29-
model_gpu_tar, mnist_gpu_script, model_cpu_1d_tar, call_model_fn_once_script, ROLE, \
28+
from integration import training_dir, mnist_cpu_1d_script, mnist_cpu_tar, mnist_cpu_script, \
29+
mnist_gpu_tar, mnist_gpu_script, mnist_cpu_1d_tar, call_model_fn_once_script, ROLE, \
3030
call_model_fn_once_tar
3131
from utils import local_mode_utils
3232

@@ -50,7 +50,7 @@ def fixture_test_loader():
5050

5151

5252
def test_serve_json(test_loader, use_gpu, image_uri, sagemaker_local_session, instance_type):
53-
model_tar = model_gpu_tar if use_gpu else model_cpu_tar
53+
model_tar = mnist_gpu_tar if use_gpu else mnist_cpu_tar
5454
mnist_script = mnist_gpu_script if use_gpu else mnist_cpu_script
5555
with _predictor(model_tar, mnist_script, image_uri, sagemaker_local_session,
5656
instance_type) as predictor:
@@ -60,7 +60,7 @@ def test_serve_json(test_loader, use_gpu, image_uri, sagemaker_local_session, in
6060

6161

6262
def test_serve_npy(test_loader, use_gpu, image_uri, sagemaker_local_session, instance_type):
63-
model_tar = model_gpu_tar if use_gpu else model_cpu_tar
63+
model_tar = mnist_gpu_tar if use_gpu else mnist_cpu_tar
6464
mnist_script = mnist_gpu_script if use_gpu else mnist_cpu_script
6565
with _predictor(model_tar, mnist_script, image_uri, sagemaker_local_session,
6666
instance_type) as predictor:
@@ -70,7 +70,7 @@ def test_serve_npy(test_loader, use_gpu, image_uri, sagemaker_local_session, ins
7070

7171

7272
def test_serve_csv(test_loader, use_gpu, image_uri, sagemaker_local_session, instance_type):
73-
with _predictor(model_cpu_1d_tar, mnist_1d_script, image_uri, sagemaker_local_session,
73+
with _predictor(mnist_cpu_1d_tar, mnist_cpu_1d_script, image_uri, sagemaker_local_session,
7474
instance_type) as predictor:
7575
for accept in (content_types.JSON, content_types.CSV, content_types.NPY):
7676
_assert_prediction_csv(predictor, test_loader, accept)
@@ -79,7 +79,7 @@ def test_serve_csv(test_loader, use_gpu, image_uri, sagemaker_local_session, ins
7979
def test_serve_cpu_model_on_gpu(test_loader, image_uri, sagemaker_local_session, instance_type):
8080
if 'cpu' in image_uri:
8181
pytest.skip("Skipping because running on CPU instance")
82-
with _predictor(model_cpu_1d_tar, mnist_1d_script, image_uri, sagemaker_local_session,
82+
with _predictor(mnist_cpu_1d_tar, mnist_cpu_1d_script, image_uri, sagemaker_local_session,
8383
instance_type) as predictor:
8484
_assert_prediction_npy_json(predictor, test_loader, content_types.NPY, content_types.JSON)
8585

0 commit comments

Comments
 (0)