fix: repack_model writes model data to default bucket if not write permissions

.python-version

@@ -0,0 +1 @@
+3.6.7

code/inference.py

@@ -0,0 +1,26 @@
+# 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.
+import json
+
+
+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

src/sagemaker/tensorflow/serving.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import logging
+import os
 
 import sagemaker
 from sagemaker.content_types import CONTENT_TYPE_JSON
@@ -128,9 +129,14 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         env = self._get_container_env()
 
         if self.entry_point:
+            deploy_key_prefix = sagemaker.fw_utils.model_code_key_prefix(self.key_prefix, self.name, self.image)
+
+            repack_model_uri = os.path.join(self.bucket, deploy_key_prefix, 'model.tar.gz')
+
             model_data = sagemaker.utils.repack_model(self.entry_point,
                                                       self.source_dir,
                                                       self.model_data,
+                                                      repack_model_uri,
                                                       self.sagemaker_session)
         else:
             model_data = self.model_data

src/sagemaker/utils.py

@@ -25,6 +25,8 @@
 
 from datetime import datetime
 from functools import wraps
+
+from boto3 import exceptions
 from six.moves.urllib import parse
 
 import six
@@ -300,7 +302,11 @@ def _tmpdir(suffix='', prefix='tmp'):
     shutil.rmtree(tmp)
 
 
-def repack_model(inference_script, source_directory, model_uri, sagemaker_session):
+def repack_model(inference_script,
+                 source_directory,
+                 model_uri,
+                 repacked_model_uri,
+                 sagemaker_session):
     """Unpack model tarball and creates a new model tarball with the provided code script.
 
     This function does the following:
@@ -311,53 +317,65 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
     Args:
         inference_script (str): path or basename of the inference script that will be packed into the model
         source_directory (str): path including all the files that will be packed into the model
+        repacked_model_uri (str): path or file system location where the new model will be saved
         model_uri (str): S3 or file system location of the original model tar
         sagemaker_session (:class:`sagemaker.session.Session`): a sagemaker session to interact with S3.
 
     Returns:
         str: path to the new packed model
     """
-    new_model_name = 'model-%s.tar.gz' % sagemaker.utils.sagemaker_short_timestamp()
+    new_model_name = os.path.basename(repacked_model_uri)
+    basedir = os.path.dirname(repacked_model_uri)
 
     with _tmpdir() as tmp:
-        tmp_model_dir = os.path.join(tmp, 'model')
-        os.mkdir(tmp_model_dir)
-
-        model_from_s3 = model_uri.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)
-
-            new_model_path = os.path.join(tmp, new_model_name)
-        else:
-            local_model_path = model_uri.replace('file://', '')
-            new_model_path = os.path.join(os.path.dirname(local_model_path), new_model_name)
+        tmp_model_dir = _extract_model_to_dir(model_uri, tmp, sagemaker_session)
 
-        with tarfile.open(name=local_model_path, mode='r:gz') as t:
-            t.extractall(path=tmp_model_dir)
+        _copy_code_to_dir(inference_script, source_directory, tmp_model_dir)
 
-        code_dir = os.path.join(tmp_model_dir, 'code')
-        if os.path.exists(code_dir):
-            shutil.rmtree(code_dir, ignore_errors=True)
+        if basedir.startswith('s3://'):
+            repack_model_path = os.path.join(tmp, new_model_name)
 
-        if source_directory:
-            shutil.copytree(source_directory, code_dir)
-        else:
-            os.mkdir(code_dir)
-            shutil.copy2(inference_script, code_dir)
-
-        with tarfile.open(new_model_path, mode='w:gz') as t:
-            t.add(tmp_model_dir, arcname=os.path.sep)
+            with tarfile.open(repack_model_path, mode='w:gz') as t:
+                t.add(tmp_model_dir, arcname=os.path.sep)
 
-        if model_from_s3:
-            url = parse.urlparse(model_uri)
+            url = parse.urlparse(basedir)
             bucket, key = url.netloc, url.path.lstrip('/')
             new_key = key.replace(os.path.basename(key), new_model_name)
 
-            sagemaker_session.boto_session.resource('s3').Object(bucket, new_key).upload_file(new_model_path)
+            sagemaker_session.boto_session.resource('s3').Object(bucket, new_key).upload_file(repack_model_path)
+
             return 's3://%s/%s' % (bucket, new_key)
         else:
-            return 'file://%s' % new_model_path
+            repack_model_path = repacked_model_uri.replace('file://', '')
+
+            with tarfile.open(repack_model_path, mode='w:gz') as t:
+                t.add(tmp_model_dir, arcname=os.path.sep)
+
+            return 'file://%s' % repack_model_path
+
+
+def _copy_code_to_dir(inference_script, source_directory, tmp_model_dir):
+    code_dir = os.path.join(tmp_model_dir, 'code')
+    if os.path.exists(code_dir):
+        shutil.rmtree(code_dir, ignore_errors=True)
+    if source_directory:
+        shutil.copytree(source_directory, code_dir)
+    else:
+        os.mkdir(code_dir)
+        shutil.copy2(inference_script, code_dir)
+
+
+def _extract_model_to_dir(model_uri, dst, sagemaker_session):
+    tmp_model_dir = os.path.join(dst, 'model')
+    os.mkdir(tmp_model_dir)
+    if model_uri.startswith('s3://'):
+        local_model_path = os.path.join(dst, 'tar_file')
+        download_file_from_url(model_uri, local_model_path, sagemaker_session)
+    else:
+        local_model_path = model_uri.replace('file://', '')
+    with tarfile.open(name=local_model_path, mode='r:gz') as t:
+        t.extractall(path=tmp_model_dir)
+    return tmp_model_dir
 
 
 def download_file_from_url(url, dst, sagemaker_session):

tests/data/tfs/tfs-test-inference-model-with-another-name/00000123/assets/foo.txt

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

tests/unit/test_utils.py

@@ -21,11 +21,13 @@
 import re
 import time
 
+from boto3 import exceptions
 import pytest
 from mock import call, patch, Mock, MagicMock
 
 import sagemaker
 
+BUCKET_WITHOUT_WRITING_PERMISSION = 's3://bucket-without-writing-permission'
 
 NAME = 'base_name'
 BUCKET_NAME = 'some_bucket'
@@ -331,10 +333,11 @@ def test_repack_model_without_source_dir(tmpdir):
     new_model_uri = sagemaker.utils.repack_model(os.path.join(source_dir, 'inference.py'),
                                                  None,
                                                  model_uri,
+                                                 's3://destination-bucket/model.tar.gz',
                                                  sagemaker_session)
 
     assert list_tar_files(fake_upload_path, tmpdir) == {'/code/inference.py', '/model'}
-    assert re.match(r'^s3://fake/model-\d+-\d+.tar.gz$', new_model_uri)
+    assert 's3://destination-bucket/model.tar.gz' == new_model_uri
 
 
 def test_repack_model_with_entry_point_without_path_without_source_dir(tmpdir):
@@ -367,15 +370,16 @@ def test_repack_model_with_entry_point_without_path_without_source_dir(tmpdir):
         new_model_uri = sagemaker.utils.repack_model('inference.py',
                                                      None,
                                                      model_uri,
+                                                     's3://destination-bucket/model.tar.gz',
                                                      sagemaker_session)
     finally:
         os.chdir(cwd)
 
     assert list_tar_files(fake_upload_path, tmpdir) == {'/code/inference.py', '/model'}
-    assert re.match(r'^s3://fake/model-\d+-\d+.tar.gz$', new_model_uri)
+    assert 's3://destination-bucket/model.tar.gz' == new_model_uri
 
 
-def test_repack_model_from_s3_saved_model_to_s3(tmpdir):
+def test_repack_model_from_s3_to_s3(tmpdir):
 
     tmp = str(tmpdir)
 
@@ -401,15 +405,16 @@ def test_repack_model_from_s3_saved_model_to_s3(tmpdir):
     new_model_uri = sagemaker.utils.repack_model('inference.py',
                                                  source_dir,
                                                  model_uri,
+                                                 's3://destination-bucket/model.tar.gz',
                                                  sagemaker_session)
 
     assert list_tar_files(fake_upload_path, tmpdir) == {'/code/this-file-should-be-included.py',
                                                         '/code/inference.py',
                                                         '/model'}
-    assert re.match(r'^s3://fake/model-\d+-\d+.tar.gz$', new_model_uri)
+    assert 's3://destination-bucket/model.tar.gz' == new_model_uri
 
 
-def test_repack_model_from_file_saves_model_to_file(tmpdir):
+def test_repack_model_from_file_to_file(tmpdir):
 
     tmp = str(tmpdir)
 
@@ -427,9 +432,42 @@ def test_repack_model_from_file_saves_model_to_file(tmpdir):
     sagemaker_session = MagicMock()
 
     file_mode_path = 'file://%s' % model_tar_path
+    destination_path = 'file://%s' % os.path.join(tmp, 'repacked-model.tar.gz')
+
+    new_model_uri = sagemaker.utils.repack_model('inference.py',
+                                                 source_dir,
+                                                 file_mode_path,
+                                                 destination_path,
+                                                 sagemaker_session)
+
+    assert os.path.dirname(new_model_uri) == os.path.dirname(file_mode_path)
+    assert list_tar_files(new_model_uri, tmpdir) == {'/code/inference.py', '/model'}
+
+
+def test_repack_model_from_file_to_folder(tmpdir):
+
+    tmp = str(tmpdir)
+
+    model_path = os.path.join(tmp, 'model')
+    write_file(model_path, 'model data')
+
+    source_dir = os.path.join(tmp, 'source-dir')
+    os.mkdir(source_dir)
+    script_path = os.path.join(source_dir, 'inference.py')
+    write_file(script_path, 'inference script')
+
+    model_tar_path = os.path.join(tmp, 'model.tar.gz')
+    sagemaker.utils.create_tar_file([model_path], model_tar_path)
+
+    sagemaker_session = MagicMock()
+
+    file_mode_path = 'file://%s' % model_tar_path
+    destination_path = 'file://%s' % os.path.join(tmp, 'repacked-model.tar.gz')
+
     new_model_uri = sagemaker.utils.repack_model('inference.py',
                                                  source_dir,
                                                  file_mode_path,
+                                                 destination_path,
                                                  sagemaker_session)
 
     assert os.path.dirname(new_model_uri) == os.path.dirname(file_mode_path)
@@ -463,10 +501,11 @@ def test_repack_model_with_inference_code_should_replace_the_code(tmpdir):
     new_model_uri = sagemaker.utils.repack_model('inference.py',
                                                  source_dir,
                                                  model_uri,
+                                                 's3://destination-bucket/repacked-model',
                                                  sagemaker_session)
 
     assert list_tar_files(fake_upload_path, tmpdir) == {'/code/new-inference.py', '/model'}
-    assert re.match(r'^s3://fake/model-\d+-\d+.tar.gz$', new_model_uri)
+    assert 's3://destination-bucket/repacked-model' == new_model_uri
 
 
 def mock_s3_model_tar(contents, sagemaker_session, tmp):
@@ -492,6 +531,8 @@ def __init__(self, bucket, key):
             self.key = key
 
         def upload_file(self, target):
+            if self.bucket in BUCKET_WITHOUT_WRITING_PERMISSION:
+                raise exceptions.S3UploadFailedError()
             shutil.copy2(target, dst)
 
     sagemaker_session.boto_session.resource().Object = MockS3Object