make sure we show that we use the scripts that we use.

Author: nadiaya

test-toolkit/integration/__init__.py

@@ -16,7 +16,6 @@
 
 resources_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'resources'))
 mnist_path = os.path.join(resources_path, 'mnist')
-mnist_script = os.path.join(mnist_path, 'mnist.py')
 data_dir = os.path.join(mnist_path, 'data')
 training_dir = os.path.join(data_dir, 'training')
 cpu_sub_dir = 'model_cpu'
@@ -32,7 +31,7 @@
 model_gpu_1d_dir = os.path.join(model_gpu_dir, '1d')
 model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
 mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
-call_model_fn_once_script = os.path.join(resources_path, 'call_model_fn_once.py')
+call_model_fn_once_script = os.path.join(model_cpu_dir, 'call_model_fn_once.py')
 
 ROLE = 'dummy/unused-role'
 DEFAULT_TIMEOUT = 20

test-toolkit/integration/local/test_serving.py

@@ -25,8 +25,8 @@
 from sagemaker_containers.beta.framework import content_types
 from torchvision import datasets, transforms
 
-from integration import training_dir, mnist_script, mnist_1d_script, model_cpu_dir, \
-    model_gpu_dir, model_cpu_1d_dir, call_model_fn_once_script, ROLE
+from integration import training_dir, mnist_1d_script, model_cpu_dir, mnist_cpu_script, \
+    model_gpu_dir, mnist_gpu_script, model_cpu_1d_dir, call_model_fn_once_script, ROLE
 from utils import local_mode_utils
 
 CONTENT_TYPE_TO_SERIALIZER_MAP = {
@@ -50,6 +50,7 @@ def fixture_test_loader():
 
 def test_serve_json_npy(test_loader, use_gpu, image_uri, sagemaker_local_session, instance_type):
     model_dir = model_gpu_dir if use_gpu else model_cpu_dir
+    mnist_script = mnist_gpu_script if use_gpu else mnist_cpu_script
     with _predictor(model_dir, mnist_script, image_uri, sagemaker_local_session,
                     instance_type) as predictor:
         for content_type in (content_types.JSON, content_types.NPY):

test-toolkit/resources/call_model_fn_once.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of