feature: Remove unnecessary dependencies. (#44)

Author: nadiaya

.gitignore

@@ -8,3 +8,4 @@ __pycache__/
 .DS_Store
 **/*.pyc
 **.pyc
+**/*.tar.gz

setup.py

@@ -47,13 +47,14 @@ def read(fname):
         'Programming Language :: Python :: 3.6',
     ],
 
-    install_requires=['numpy==1.16.4', 'Pillow>=6.2.0', 'retrying==1.3.3', 'sagemaker-containers>=2.5.4',
-                      'six==1.12.0', 'requests_mock==1.6.0', 'sagemaker-inference>=1.2.2',
-                      'retrying==1.3.3'],
+    # We don't declare our dependency on torch here because we build with
+    # different packages for different variants
+    install_requires=['numpy', 'retrying', 'sagemaker-inference>=1.2.2'],
     extras_require={
         'test': ['boto3==1.10.32', 'coverage==4.5.3', 'docker-compose==1.23.2', 'flake8==3.7.7', 'Flask==1.1.1',
                  'mock==2.0.0', 'pytest==4.4.0', 'pytest-cov==2.7.1', 'pytest-xdist==1.28.0', 'PyYAML==3.10',
-                 'sagemaker==1.48.0', 'requests==2.20.0', 'torchvision==0.5.0', 'tox==3.7.0', 'requests_mock==1.6.0']
+                 'sagemaker==1.48.0', 'sagemaker-containers>=2.5.4', 'six==1.12.0', 'requests==2.20.0',
+                 'requests_mock==1.6.0', 'torch==1.4.0', 'torchvision==0.5.0', 'tox==3.7.0']
     },
 
     entry_points={

test-toolkit/integration/__init__.py

@@ -14,6 +14,8 @@
 
 import os
 
+from utils import file_utils
+
 resources_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'resources'))
 mnist_path = os.path.join(resources_path, 'mnist')
 data_dir = os.path.join(mnist_path, 'data')
@@ -24,14 +26,33 @@
 
 model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
 mnist_cpu_script = os.path.join(model_cpu_dir, 'mnist.py')
+model_cpu_tar = file_utils.make_tarfile(mnist_cpu_script,
+                                        os.path.join(model_cpu_dir, "model.pth"),
+                                        model_cpu_dir)
+
 model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
 mnist_1d_script = os.path.join(model_cpu_1d_dir, 'mnist_1d.py')
+model_cpu_1d_tar = file_utils.make_tarfile(mnist_1d_script,
+                                           os.path.join(model_cpu_1d_dir, "model.pth"),
+                                           model_cpu_1d_dir)
+
 model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
 mnist_gpu_script = os.path.join(model_gpu_dir, 'mnist.py')
-model_gpu_1d_dir = os.path.join(model_gpu_dir, '1d')
+model_gpu_tar = file_utils.make_tarfile(mnist_gpu_script,
+                                        os.path.join(model_gpu_dir, "model.pth"),
+                                        model_gpu_dir)
+
 model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
 mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
+model_eia_tar = file_utils.make_tarfile(mnist_eia_script,
+                                        os.path.join(model_eia_dir, "model.pth"),
+                                        model_eia_dir)
+
 call_model_fn_once_script = os.path.join(model_cpu_dir, 'call_model_fn_once.py')
+call_model_fn_once_tar = file_utils.make_tarfile(call_model_fn_once_script,
+                                                 os.path.join(model_cpu_dir, "model.pth"),
+                                                 model_cpu_dir,
+                                                 "model_call_model_fn_once.tar.gz")
 
 ROLE = 'dummy/unused-role'
 DEFAULT_TIMEOUT = 20

test-toolkit/integration/local/test_serving.py

@@ -25,8 +25,9 @@
 from sagemaker_containers.beta.framework import content_types
 from torchvision import datasets, transforms
 
-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 integration import training_dir, mnist_1d_script, model_cpu_tar, mnist_cpu_script, \
+    model_gpu_tar, mnist_gpu_script, model_cpu_1d_tar, call_model_fn_once_script, ROLE, \
+    call_model_fn_once_tar
 from utils import local_mode_utils
 
 CONTENT_TYPE_TO_SERIALIZER_MAP = {
@@ -49,31 +50,31 @@ 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
+    model_tar = model_gpu_tar if use_gpu else model_cpu_tar
     mnist_script = mnist_gpu_script if use_gpu else mnist_cpu_script
-    with _predictor(model_dir, mnist_script, image_uri, sagemaker_local_session,
+    with _predictor(model_tar, mnist_script, image_uri, sagemaker_local_session,
                     instance_type) as predictor:
         for content_type in (content_types.JSON, content_types.NPY):
             for accept in (content_types.JSON, content_types.CSV, content_types.NPY):
                 _assert_prediction_npy_json(predictor, test_loader, content_type, accept)
 
 
 def test_serve_csv(test_loader, use_gpu, image_uri, sagemaker_local_session, instance_type):
-    with _predictor(model_cpu_1d_dir, mnist_1d_script, image_uri, sagemaker_local_session,
+    with _predictor(model_cpu_1d_tar, mnist_1d_script, image_uri, sagemaker_local_session,
                     instance_type) as predictor:
         for accept in (content_types.JSON, content_types.CSV, content_types.NPY):
             _assert_prediction_csv(predictor, test_loader, accept)
 
 
 @pytest.mark.skip_cpu
 def test_serve_cpu_model_on_gpu(test_loader, image_uri, sagemaker_local_session, instance_type):
-    with _predictor(model_cpu_1d_dir, mnist_1d_script, image_uri, sagemaker_local_session,
+    with _predictor(model_cpu_1d_tar, mnist_1d_script, image_uri, sagemaker_local_session,
                     instance_type) as predictor:
         _assert_prediction_npy_json(predictor, test_loader, content_types.NPY, content_types.JSON)
 
 
 def test_serving_calls_model_fn_once(image_uri, sagemaker_local_session, instance_type):
-    with _predictor(model_cpu_dir, call_model_fn_once_script, image_uri, sagemaker_local_session,
+    with _predictor(call_model_fn_once_tar, call_model_fn_once_script, image_uri, sagemaker_local_session,
                     instance_type, model_server_workers=2) as predictor:
         predictor.accept = None
         predictor.deserializer = BytesDeserializer()
@@ -86,9 +87,9 @@ def test_serving_calls_model_fn_once(image_uri, sagemaker_local_session, instanc
 
 
 @contextmanager
-def _predictor(model_dir, script, image, sagemaker_local_session, instance_type,
+def _predictor(model_tar, script, image, sagemaker_local_session, instance_type,
                model_server_workers=None):
-    model = PyTorchModel('file://{}'.format(model_dir),
+    model = PyTorchModel('file://{}'.format(model_tar),
                          ROLE,
                          script,
                          image=image,

test-toolkit/integration/sagemaker/test_mnist.py

@@ -12,47 +12,43 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import os
-
 import numpy as np
 import pytest
 import sagemaker
 from sagemaker.pytorch import PyTorchModel
 
-from integration import model_cpu_dir, mnist_cpu_script, mnist_gpu_script, model_eia_dir, mnist_eia_script
+from integration import model_cpu_tar, model_gpu_tar, mnist_cpu_script, mnist_gpu_script, \
+    model_eia_tar, mnist_eia_script
 from integration.sagemaker.timeout import timeout_and_delete_endpoint
 
 
 @pytest.mark.cpu_test
 def test_mnist_cpu(sagemaker_session, image_uri, instance_type):
     instance_type = instance_type or 'ml.c4.xlarge'
-    model_dir = os.path.join(model_cpu_dir, 'model_mnist.tar.gz')
-    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_cpu_script)
+    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_cpu_tar, mnist_cpu_script)
 
 
 @pytest.mark.gpu_test
 def test_mnist_gpu(sagemaker_session, image_uri, instance_type):
     instance_type = instance_type or 'ml.p2.xlarge'
-    model_dir = os.path.join(model_cpu_dir, 'model_mnist.tar.gz')
-    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_gpu_script)
+    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_gpu_tar, mnist_gpu_script)
 
 
 @pytest.mark.eia_test
 def test_mnist_eia(sagemaker_session, image_uri, instance_type, accelerator_type):
     instance_type = instance_type or 'ml.c4.xlarge'
     # Scripted model is serialized with torch.jit.save().
     # Inference test for EIA doesn't need to instantiate model definition then load state_dict
-    model_dir = os.path.join(model_eia_dir, 'model_mnist.tar.gz')
-    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_eia_script,
+    _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_eia_tar, mnist_eia_script,
                             accelerator_type=accelerator_type)
 
 
-def _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_dir, mnist_script,
+def _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_tar, mnist_script,
                             accelerator_type=None):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-pytorch-serving")
 
     model_data = sagemaker_session.upload_data(
-        path=model_dir,
+        path=model_tar,
         key_prefix="sagemaker-pytorch-serving/models",
     )
 

test-toolkit/resources/mnist/model_cpu/1d/model.tar.gz

@@ -10,4 +10,39 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-# This file is intentionally left blank to utilize default_model_fn and default_predict_fn
+from __future__ import absolute_import
+import logging
+import os
+import sys
+
+import torch
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def predict_fn(input_data, model):
+    logger.info('Performing EIA inference with Torch JIT context with input of size {}'.format(input_data.shape))
+    # With EI, client instance should be CPU for cost-efficiency. Subgraphs with unsupported arguments run locally. Server runs with CUDA
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    mdoel = model.to(device)
+    input_data = input_data.to(device)
+    with torch.no_grad():
+        # Set the target device to the accelerator ordinal
+        with torch.jit.optimized_execution(True, {'target_device': 'eia:0'}):
+            return model(input_data)
+
+
+def model_fn(model_dir):
+    logger.info('model_fn: Loading model with TorchScript from {}'.format(model_dir))
+    # Scripted model is serialized with torch.jit.save().
+    # No need to instantiate model definition then load state_dict
+    model = torch.jit.load('model.pth')
+    return model
+
+
+def save_model(model, model_dir):
+    logger.info("Saving the model to {}.".format(model_dir))
+    path = os.path.join(model_dir, 'model.pth')
+    torch.jit.save(model, path)