feature: repack_model support dependencies and code location

src/sagemaker/model.py

@@ -448,21 +448,25 @@ 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()
+            self.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=self.repacked_model_data,
+                               sagemaker_session=self.sagemaker_session)
             self.uploaded_code = UploadedCode(s3_prefix=self.repacked_model_data,
                                               script_name=os.path.basename(self.entry_point))
 

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
 

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,90 @@ 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)
+        _update_code(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 _update_code(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):

tests/unit/test_mxnet.py

@@ -29,7 +29,6 @@
 DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
 SCRIPT_PATH = os.path.join(DATA_DIR, 'dummy_script.py')
 MODEL_DATA = 's3://mybucket/model'
-REPACKED_MODEL_DATA = 's3://mybucket/repacked/model'
 TIMESTAMP = '2017-11-06-14:14:15.672'
 TIME = 1507167947
 BUCKET_NAME = 'mybucket'
@@ -280,7 +279,7 @@ def test_mxnet(strftime, sagemaker_session, mxnet_version, skip_if_mms_version):
     assert isinstance(predictor, MXNetPredictor)
 
 
-@patch('sagemaker.utils.repack_model', return_value=REPACKED_MODEL_DATA)
+@patch('sagemaker.utils.repack_model')
 @patch('time.strftime', return_value=TIMESTAMP)
 def test_mxnet_mms_version(strftime, repack_model, sagemaker_session, mxnet_version, skip_if_not_mms_version):
     mx = MXNet(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
@@ -307,11 +306,12 @@ def test_mxnet_mms_version(strftime, repack_model, sagemaker_session, mxnet_vers
     expected_image_base = _get_full_image_uri(mxnet_version, IMAGE_REPO_SERVING_NAME, 'gpu')
     environment = {
         'Environment': {
-            'SAGEMAKER_SUBMIT_DIRECTORY': REPACKED_MODEL_DATA,
+            'SAGEMAKER_SUBMIT_DIRECTORY': 's3://mybucket/sagemaker-mxnet-2017-11-06-14:14:15.672/model.tar.gz',
             'SAGEMAKER_PROGRAM': 'dummy_script.py', 'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
             'SAGEMAKER_REGION': 'us-west-2', 'SAGEMAKER_CONTAINER_LOG_LEVEL': '20'
         },
-        'Image': expected_image_base.format(mxnet_version), 'ModelDataUrl': REPACKED_MODEL_DATA
+        'Image': expected_image_base.format(mxnet_version),
+        'ModelDataUrl': 's3://mybucket/sagemaker-mxnet-2017-11-06-14:14:15.672/model.tar.gz'
     }
     assert environment == model.prepare_container_def(GPU)
 
@@ -366,21 +366,23 @@ def test_model(sagemaker_session):
     assert isinstance(predictor, MXNetPredictor)
 
 
-@patch('sagemaker.utils.repack_model', return_value=REPACKED_MODEL_DATA)
+@patch('sagemaker.utils.repack_model')
 def test_model_mms_version(repack_model, sagemaker_session):
     model = MXNetModel(MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH,
                        framework_version=MXNetModel._LOWEST_MMS_VERSION,
-                       sagemaker_session=sagemaker_session)
+                       sagemaker_session=sagemaker_session, name='test-mxnet-model')
     predictor = model.deploy(1, GPU)
 
     repack_model.assert_called_once_with(inference_script=SCRIPT_PATH,
                                          source_directory=None,
+                                         dependencies=[],
                                          model_uri=MODEL_DATA,
+                                         repacked_model_uri='s3://mybucket/test-mxnet-model/model.tar.gz',
                                          sagemaker_session=sagemaker_session)
 
     assert model.model_data == MODEL_DATA
-    assert model.repacked_model_data == REPACKED_MODEL_DATA
-    assert model.uploaded_code == UploadedCode(s3_prefix=REPACKED_MODEL_DATA,
+    assert model.repacked_model_data == 's3://mybucket/test-mxnet-model/model.tar.gz'
+    assert model.uploaded_code == UploadedCode(s3_prefix='s3://mybucket/test-mxnet-model/model.tar.gz',
                                                script_name=os.path.basename(SCRIPT_PATH))
     assert isinstance(predictor, MXNetPredictor)