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

@@ -7,6 +7,7 @@ CHANGELOG
 * feature: Allow Local Serving of Models in S3
 * enhancement: Allow option for ``HyperparameterTuner`` to not include estimator metadata in job
 * bug-fix: Estimators: Join tensorboard thread after fitting
+* 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._model_source_dir(),
                             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/estimator.py

@@ -579,6 +579,14 @@ def _stage_user_code_in_s3(self):
                                   script=self.entry_point,
                                   directory=self.source_dir)
 
+    def _model_source_dir(self):
+        """ Gets the appropriate value to pass as source_dir to model constructor on deploying
+
+        Returns:
+            str: Either a local or an S3 path pointing to the source_dir to be used for code by the model to be deployed
+        """
+        return self.source_dir if self.sagemaker_session.local_mode else self.uploaded_code.s3_prefix
+
     def hyperparameters(self):
         """Return the hyperparameters as a dictionary to use for training.
 

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._model_source_dir(),
                           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

@@ -300,7 +300,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._model_source_dir(),
                                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'
         },