feature: repack_model support dependencies and code location (aws#821)

mvsusp · web-flow · commit c921eadcf385 · 2019-05-30T15:27:00.000-07:00
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -448,21 +448,27 @@ def _upload_code(self, key_prefix, repack=False):
         local_code = utils.get_config_value('local.local_code', self.sagemaker_session.config)
         if self.sagemaker_session.local_mode and local_code:
             self.uploaded_code = None
-        else:
-            if not repack:
-                bucket = self.bucket or self.sagemaker_session.default_bucket()
-                self.uploaded_code = fw_utils.tar_and_upload_dir(session=self.sagemaker_session.boto_session,
-                                                                 bucket=bucket,
-                                                                 s3_key_prefix=key_prefix,
-                                                                 script=self.entry_point,
-                                                                 directory=self.source_dir,
-                                                                 dependencies=self.dependencies)
+        elif not repack:
+            bucket = self.bucket or self.sagemaker_session.default_bucket()
+            self.uploaded_code = fw_utils.tar_and_upload_dir(session=self.sagemaker_session.boto_session,
+                                                             bucket=bucket,
+                                                             s3_key_prefix=key_prefix,
+                                                             script=self.entry_point,
+                                                             directory=self.source_dir,
+                                                             dependencies=self.dependencies)
 
         if repack:
-            self.repacked_model_data = utils.repack_model(inference_script=self.entry_point,
-                                                          source_directory=self.source_dir,
-                                                          model_uri=self.model_data,
-                                                          sagemaker_session=self.sagemaker_session)
+            bucket = self.bucket or self.sagemaker_session.default_bucket()
+            repacked_model_data = 's3://' + os.path.join(bucket, key_prefix, 'model.tar.gz')
+
+            utils.repack_model(inference_script=self.entry_point,
+                               source_directory=self.source_dir,
+                               dependencies=self.dependencies,
+                               model_uri=self.model_data,
+                               repacked_model_uri=repacked_model_data,
+                               sagemaker_session=self.sagemaker_session)
+
+            self.repacked_model_data = repacked_model_data
             self.uploaded_code = UploadedCode(s3_prefix=self.repacked_model_data,
                                               script_name=os.path.basename(self.entry_point))
 
diff --git a/src/sagemaker/mxnet/model.py b/src/sagemaker/mxnet/model.py
@@ -92,21 +92,21 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         Returns:
             dict[str, str]: A container definition object usable with the CreateModel API.
         """
-        mms_version = parse_version(self.framework_version) >= parse_version(self._LOWEST_MMS_VERSION)
+        is_mms_version = parse_version(self.framework_version) >= parse_version(self._LOWEST_MMS_VERSION)
 
         deploy_image = self.image
         if not deploy_image:
             region_name = self.sagemaker_session.boto_session.region_name
 
             framework_name = self.__framework_name__
-            if mms_version:
+            if is_mms_version:
                 framework_name += '-serving'
 
             deploy_image = create_image_uri(region_name, framework_name, instance_type,
                                             self.framework_version, self.py_version, accelerator_type=accelerator_type)
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
-        self._upload_code(deploy_key_prefix, mms_version)
+        self._upload_code(deploy_key_prefix, is_mms_version)
         deploy_env = dict(self.env)
         deploy_env.update(self._framework_env_vars())
 
diff --git a/src/sagemaker/tensorflow/serving.py b/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,10 +129,17 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         env = self._get_container_env()
 
         if self.entry_point:
-            model_data = sagemaker.utils.repack_model(self.entry_point,
-                                                      self.source_dir,
-                                                      self.model_data,
-                                                      self.sagemaker_session)
+            key_prefix = sagemaker.fw_utils.model_code_key_prefix(self.key_prefix, self.name, image)
+
+            bucket = self.bucket or self.sagemaker_session.default_bucket()
+            model_data = 's3://' + os.path.join(bucket, key_prefix, 'model.tar.gz')
+
+            sagemaker.utils.repack_model(self.entry_point,
+                                         self.source_dir,
+                                         self.dependencies,
+                                         self.model_data,
+                                         model_data,
+                                         self.sagemaker_session)
         else:
             model_data = self.model_data
 
diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py
@@ -29,8 +29,6 @@
 
 import six
 
-import sagemaker
-
 ECR_URI_PATTERN = r'^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(amazonaws.com|c2s.ic.gov)(/)(.*:.*)$'
 
 
@@ -300,7 +298,12 @@ 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,
+                 dependencies,
+                 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,60 +314,91 @@ 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
+        dependencies (list[str]): A list of paths to directories (absolute or relative) with
+                any additional libraries that will be exported to the container (default: []).
+                The library folders will be copied to SageMaker in the same folder where the entrypoint is copied.
+                Example:
+
+                    The following call
+                    >>> Estimator(entry_point='train.py', dependencies=['my/libs/common', 'virtual-env'])
+                    results in the following inside the container:
+
+                    >>> $ ls
+
+                    >>> opt/ml/code
+                    >>>     |------ train.py
+                    >>>     |------ common
+                    >>>     |------ virtual-env
+
+        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()
+    dependencies = dependencies or []
 
     with _tmpdir() as tmp:
-        tmp_model_dir = os.path.join(tmp, 'model')
-        os.mkdir(tmp_model_dir)
+        model_dir = _extract_model(model_uri, sagemaker_session, tmp)
 
-        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)
+        _create_or_update_code_dir(model_dir, inference_script, source_directory, dependencies, sagemaker_session, tmp)
 
-            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_path = os.path.join(tmp, 'temp-model.tar.gz')
+        with tarfile.open(tmp_model_path, mode='w:gz') as t:
+            t.add(model_dir, arcname=os.path.sep)
 
-        with tarfile.open(name=local_model_path, mode='r:gz') as t:
-            t.extractall(path=tmp_model_dir)
+        _save_model(repacked_model_uri, tmp_model_path, sagemaker_session)
 
-        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 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)
+def _save_model(repacked_model_uri, tmp_model_path, sagemaker_session):
+    if repacked_model_uri.lower().startswith('s3://'):
+        url = parse.urlparse(repacked_model_uri)
+        bucket, key = url.netloc, url.path.lstrip('/')
+        new_key = key.replace(os.path.basename(key), os.path.basename(repacked_model_uri))
 
-            with tarfile.open(name=local_code_path, mode='r:gz') as t:
-                t.extractall(path=code_dir)
+        sagemaker_session.boto_session.resource('s3').Object(bucket, new_key).upload_file(
+            tmp_model_path)
+    else:
+        shutil.move(tmp_model_path, repacked_model_uri.replace('file://', ''))
 
-        elif 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)
+def _create_or_update_code_dir(model_dir, inference_script, source_directory,
+                               dependencies, sagemaker_session, tmp):
+    code_dir = os.path.join(model_dir, 'code')
+    if os.path.exists(code_dir):
+        shutil.rmtree(code_dir, ignore_errors=True)
+    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_code_path, mode='r:gz') as t:
+            t.extractall(path=code_dir)
 
-        if model_from_s3:
-            url = parse.urlparse(model_uri)
-            bucket, key = url.netloc, url.path.lstrip('/')
-            new_key = key.replace(os.path.basename(key), new_model_name)
+    elif source_directory:
+        shutil.copytree(source_directory, code_dir)
+    else:
+        os.mkdir(code_dir)
+        shutil.copy2(inference_script, code_dir)
 
-            sagemaker_session.boto_session.resource('s3').Object(bucket, new_key).upload_file(new_model_path)
-            return 's3://%s/%s' % (bucket, new_key)
+    for dependency in dependencies:
+        if os.path.isdir(dependency):
+            shutil.copytree(dependency, code_dir)
         else:
-            return 'file://%s' % new_model_path
+            shutil.copy2(dependency, code_dir)
+
+
+def _extract_model(model_uri, sagemaker_session, tmp):
+    tmp_model_dir = os.path.join(tmp, 'model')
+    os.mkdir(tmp_model_dir)
+    if model_uri.lower().startswith('s3://'):
+        local_model_path = os.path.join(tmp, '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):
diff --git a/tests/data/tfs/tfs-test-entrypoint-and-dependencies/dependency.py b/tests/data/tfs/tfs-test-entrypoint-and-dependencies/dependency.py
@@ -0,0 +1,12 @@
+# Copyright 2019 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.
diff --git a/tests/data/tfs/tfs-test-entrypoint-and-dependencies/inference.py b/tests/data/tfs/tfs-test-entrypoint-and-dependencies/inference.py
@@ -0,0 +1,27 @@
+# Copyright 2019 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
+
+import dependency
+
+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
diff --git a/tests/data/tfs/tfs-test-model-with-inference/code/inference.py b/tests/data/tfs/tfs-test-model-with-inference/code/inference.py
@@ -1,4 +1,4 @@
-# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2019 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
@@ -12,7 +12,6 @@
 # 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']]
diff --git a/tests/integ/test_tfs.py b/tests/integ/test_tfs.py
@@ -84,8 +84,8 @@ def tfs_predictor_with_model_and_entry_point_same_tar(instance_type,
 
 
 @pytest.fixture(scope='module')
-def tfs_predictor_with_model_and_entry_point_separated(instance_type,
-                                                       sagemaker_session, tf_full_version):
+def tfs_predictor_with_model_and_entry_point_and_dependencies(instance_type,
+                                                              sagemaker_session, tf_full_version):
     endpoint_name = sagemaker.utils.unique_name_from_base('sagemaker-tensorflow-serving')
 
     model_data = sagemaker_session.upload_data(
@@ -96,10 +96,14 @@ def tfs_predictor_with_model_and_entry_point_separated(instance_type,
     with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name,
                                                                  sagemaker_session):
         entry_point = os.path.join(tests.integ.DATA_DIR,
-                                   'tfs/tfs-test-model-with-inference/code/inference.py')
+                                   'tfs/tfs-test-entrypoint-and-dependencies/inference.py')
+        dependencies = [os.path.join(tests.integ.DATA_DIR,
+                                     'tfs/tfs-test-entrypoint-and-dependencies/dependency.py')]
+
         model = Model(entry_point=entry_point,
                       model_data=model_data,
                       role='SageMakerRole',
+                      dependencies=dependencies,
                       framework_version=tf_full_version,
                       sagemaker_session=sagemaker_session)
         predictor = model.deploy(1, instance_type, endpoint_name=endpoint_name)
@@ -152,12 +156,12 @@ def test_predict_with_entry_point(tfs_predictor_with_model_and_entry_point_same_
     assert expected_result == result
 
 
-def test_predict_with_model_and_entry_point_separated(
-        tfs_predictor_with_model_and_entry_point_separated):
+def test_predict_with_model_and_entry_point_and_dependencies_separated(
+        tfs_predictor_with_model_and_entry_point_and_dependencies):
     input_data = {'instances': [1.0, 2.0, 5.0]}
     expected_result = {'predictions': [4.0, 4.5, 6.0]}
 
-    result = tfs_predictor_with_model_and_entry_point_separated.predict(input_data)
+    result = tfs_predictor_with_model_and_entry_point_and_dependencies.predict(input_data)
     assert expected_result == result
 
 
diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py
diff --git a/tests/unit/test_tfs.py b/tests/unit/test_tfs.py
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
