fix: honor source_dir from S3

src/sagemaker/utils.py

@@ -323,7 +323,7 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
         tmp_model_dir = os.path.join(tmp, 'model')
         os.mkdir(tmp_model_dir)
 
-        model_from_s3 = model_uri.startswith('s3://')
+        model_from_s3 = model_uri.lower().startswith('s3://')
         if model_from_s3:
             local_model_path = os.path.join(tmp, 'tar_file')
             download_file_from_url(model_uri, local_model_path, sagemaker_session)
@@ -340,7 +340,14 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
         if os.path.exists(code_dir):
             shutil.rmtree(code_dir, ignore_errors=True)
 
-        if source_directory:
+        if source_directory and source_directory.lower().startswith('s3://'):
+            local_code_path = os.path.join(tmp, 'local_code.tar.gz')
+            download_file_from_url(source_directory, local_code_path, sagemaker_session)
+
+            with tarfile.open(name=local_model_path, mode='r:gz') as t:
+                t.extractall(path=code_dir)
+
+        elif source_directory:
             shutil.copytree(source_directory, code_dir)
         else:
             os.mkdir(code_dir)

tests/data/tfs/tfs-test-entrypoint-with-handler/123/assets/foo.txt

@@ -0,0 +1 @@
+asset-file-contents

tests/data/tfs/tfs-test-entrypoint-with-handler/inference.py

@@ -0,0 +1,43 @@
+# Copyright 2018 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
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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.
+
+"""Exports a toy TensorFlow model.
+Exports a TensorFlow model to /opt/ml/model/
+This graph calculates,
+  y = a*x + b
+where a and b are variables with a=0.5 and b=2.
+"""
+import json
+import shutil
+
+
+def save_model():
+    shutil.copytree('/opt/ml/code/123', '/opt/ml/model/123')
+
+
+def input_handler(data, context):
+    data = json.loads(data.read().decode('utf-8'))
+    new_values = [x + 1 for x in data['instances']]
+    dumps = json.dumps({'instances': new_values})
+    return dumps
+
+
+def output_handler(data, context):
+    response_content_type = context.accept_header
+    prediction = data.content
+    return prediction, response_content_type
+
+
+if __name__ == "__main__":
+    save_model()
+

tests/integ/test_tf_script_mode.py

@@ -24,11 +24,15 @@
 from sagemaker.utils import unique_name_from_base
 
 import tests.integ
+from tests.integ import timeout
 
 ROLE = 'SageMakerRole'
 
-RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'data', 'tensorflow_mnist')
-SCRIPT = os.path.join(RESOURCE_PATH, 'mnist.py')
+RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'data')
+MNIST_RESOURCE_PATH = os.path.join(RESOURCE_PATH, 'tensorflow_mnist')
+TFS_RESOURCE_PATH = os.path.join(RESOURCE_PATH, 'tfs', 'tfs-test-entrypoint-with-handler')
+
+SCRIPT = os.path.join(MNIST_RESOURCE_PATH, 'mnist.py')
 PARAMETER_SERVER_DISTRIBUTION = {'parameter_server': {'enabled': True}}
 MPI_DISTRIBUTION = {'mpi': {'enabled': True}}
 TAGS = [{'Key': 'some-key', 'Value': 'some-value'}]
@@ -57,7 +61,7 @@ def test_mnist(sagemaker_session, instance_type):
                            metric_definitions=[
                                {'Name': 'train:global_steps', 'Regex': r'global_step\/sec:\s(.*)'}])
     inputs = estimator.sagemaker_session.upload_data(
-        path=os.path.join(RESOURCE_PATH, 'data'),
+        path=os.path.join(MNIST_RESOURCE_PATH, 'data'),
         key_prefix='scriptmode/mnist')
 
     with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
@@ -88,7 +92,7 @@ def test_server_side_encryption(sagemaker_session):
                                output_kms_key=kms_key)
 
         inputs = estimator.sagemaker_session.upload_data(
-            path=os.path.join(RESOURCE_PATH, 'data'),
+            path=os.path.join(MNIST_RESOURCE_PATH, 'data'),
             key_prefix='scriptmode/mnist')
 
         with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
@@ -110,7 +114,7 @@ def test_mnist_distributed(sagemaker_session, instance_type):
                            framework_version=TensorFlow.LATEST_VERSION,
                            distributions=PARAMETER_SERVER_DISTRIBUTION)
     inputs = estimator.sagemaker_session.upload_data(
-        path=os.path.join(RESOURCE_PATH, 'data'),
+        path=os.path.join(MNIST_RESOURCE_PATH, 'data'),
         key_prefix='scriptmode/distributed_mnist')
 
     with tests.integ.timeout.timeout(minutes=tests.integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
@@ -129,7 +133,7 @@ def test_mnist_async(sagemaker_session):
                            framework_version=TensorFlow.LATEST_VERSION,
                            tags=TAGS)
     inputs = estimator.sagemaker_session.upload_data(
-        path=os.path.join(RESOURCE_PATH, 'data'),
+        path=os.path.join(MNIST_RESOURCE_PATH, 'data'),
         key_prefix='scriptmode/mnist')
     estimator.fit(inputs=inputs, wait=False, job_name=unique_name_from_base('test-tf-sm-async'))
     training_job_name = estimator.latest_training_job.name
@@ -150,6 +154,35 @@ def test_mnist_async(sagemaker_session):
                                  estimator.latest_training_job.name, TAGS)
 
 
+@pytest.mark.skipif(tests.integ.PYTHON_VERSION != 'py3',
+                    reason="Script Mode tests are only configured to run with Python 3")
+def test_deploy_with_input_handlers(sagemaker_session, instance_type):
+    estimator = TensorFlow(entry_point='inference.py',
+                           source_dir=TFS_RESOURCE_PATH,
+                           role=ROLE,
+                           train_instance_count=1,
+                           train_instance_type=instance_type,
+                           sagemaker_session=sagemaker_session,
+                           py_version='py3',
+                           framework_version=TensorFlow.LATEST_VERSION,
+                           tags=TAGS)
+
+    estimator.fit(job_name=unique_name_from_base('test-tf-tfs-deploy'))
+
+    endpoint_name = estimator.latest_training_job.name
+
+    with timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+
+        predictor = estimator.deploy(initial_instance_count=1, instance_type=instance_type,
+                                     endpoint_name=endpoint_name)
+
+        input_data = {'instances': [1.0, 2.0, 5.0]}
+        expected_result = {'predictions': [4.0, 4.5, 6.0]}
+
+        result = predictor.predict(input_data)
+        assert expected_result == result
+
+
 def _assert_s3_files_exist(s3_url, files):
     parsed_url = urlparse(s3_url)
     s3 = boto3.client('s3')