Let Framework models reuse code uploaded by Framework estimators

- This addresses problem 1 in #226
- This relies on the tar_and_upload_dir's behaviour when given S3 paths for directory
- This updates MXNet, TensorFlow and Chainer frameworks

Author: humanzz

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 ========
 * feature: Allow Local Serving of Models in S3
 * enhancement: Allow option for ``HyperparameterTuner`` to not include estimator metadata in job
+* enhancement: Let Framework models reuse code uploaded by Framework estimators
 
 
 1.4.2

src/sagemaker/chainer/estimator.py

@@ -116,7 +116,7 @@ def create_model(self, model_server_workers=None):
             sagemaker.chainer.model.ChainerModel: A SageMaker ``ChainerModel`` object.
                 See :func:`~sagemaker.chainer.model.ChainerModel` for full details.
         """
-        return ChainerModel(self.model_data, self.role, self.entry_point, source_dir=self.source_dir,
+        return ChainerModel(self.model_data, self.role, self.entry_point, source_dir=self.uploaded_code.s3_prefix,
                             enable_cloudwatch_metrics=self.enable_cloudwatch_metrics, name=self._current_job_name,
                             container_log_level=self.container_log_level, code_location=self.code_location,
                             py_version=self.py_version, framework_version=self.framework_version,

src/sagemaker/model.py

@@ -132,6 +132,7 @@ def __init__(self, model_data, image, role, entry_point, source_dir=None, predic
             source_dir (str): Path (absolute or relative) to a directory with any other training
                 source code dependencies aside from tne entry point file (default: None). Structure within this
                 directory will be preserved when training on SageMaker.
+                If the directory points to S3, no code will be uploaded and the S3 location will be used instead.
             predictor_cls (callable[string, sagemaker.session.Session]): A function to call to create
                a predictor (default: None). If not None, ``deploy`` will return the result of invoking
                this function on the created endpoint name.

src/sagemaker/mxnet/estimator.py

@@ -82,7 +82,7 @@ def create_model(self, model_server_workers=None):
             sagemaker.mxnet.model.MXNetModel: A SageMaker ``MXNetModel`` object.
                 See :func:`~sagemaker.mxnet.model.MXNetModel` for full details.
         """
-        return MXNetModel(self.model_data, self.role, self.entry_point, source_dir=self.source_dir,
+        return MXNetModel(self.model_data, self.role, self.entry_point, source_dir=self.uploaded_code.s3_prefix,
                           enable_cloudwatch_metrics=self.enable_cloudwatch_metrics, name=self._current_job_name,
                           container_log_level=self.container_log_level, code_location=self.code_location,
                           py_version=self.py_version, framework_version=self.framework_version,

src/sagemaker/tensorflow/estimator.py

@@ -296,7 +296,7 @@ def create_model(self, model_server_workers=None):
                 See :func:`~sagemaker.tensorflow.model.TensorFlowModel` for full details.
         """
         env = {'SAGEMAKER_REQUIREMENTS': self.requirements_file}
-        return TensorFlowModel(self.model_data, self.role, self.entry_point, source_dir=self.source_dir,
+        return TensorFlowModel(self.model_data, self.role, self.entry_point, source_dir=self.uploaded_code.s3_prefix,
                                enable_cloudwatch_metrics=self.enable_cloudwatch_metrics, env=env,
                                name=self._current_job_name, container_log_level=self.container_log_level,
                                code_location=self.code_location, py_version=self.py_version,

tests/unit/test_chainer.py

@@ -274,7 +274,7 @@ def test_chainer(strftime, sagemaker_session, chainer_version):
     expected_image_base = '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-chainer:{}-gpu-{}'
     assert {'Environment':
             {'SAGEMAKER_SUBMIT_DIRECTORY':
-             's3://mybucket/sagemaker-chainer-{}/sourcedir.tar.gz'.format(TIMESTAMP),
+             's3://mybucket/sagemaker-chainer-{}/source/sourcedir.tar.gz'.format(TIMESTAMP),
              'SAGEMAKER_PROGRAM': 'dummy_script.py',
              'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
              'SAGEMAKER_REGION': 'us-west-2',

tests/unit/test_mxnet.py

@@ -149,7 +149,7 @@ def test_mxnet(strftime, sagemaker_session, mxnet_version):
     expected_image_base = '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:{}-gpu-py2'
     environment = {
         'Environment': {
-            'SAGEMAKER_SUBMIT_DIRECTORY': 's3://mybucket/sagemaker-mxnet-{}/sourcedir.tar.gz'.format(TIMESTAMP),
+            'SAGEMAKER_SUBMIT_DIRECTORY': 's3://mybucket/sagemaker-mxnet-{}/source/sourcedir.tar.gz'.format(TIMESTAMP),
             'SAGEMAKER_PROGRAM': 'dummy_script.py', 'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
             'SAGEMAKER_REGION': 'us-west-2', 'SAGEMAKER_CONTAINER_LOG_LEVEL': '20'
         },