Create fake uploaded_code in framework model

Author: yangaws

src/sagemaker/model.py

@@ -17,7 +17,7 @@
 import sagemaker
 
 from sagemaker.local import LocalSession
-from sagemaker.fw_utils import UploadedCode, tar_and_upload_dir, parse_s3_url, model_code_key_prefix
+from sagemaker.fw_utils import tar_and_upload_dir, parse_s3_url, model_code_key_prefix
 from sagemaker.session import Session
 from sagemaker.utils import name_from_image, get_config_value
 
@@ -188,9 +188,7 @@ def prepare_container_def(self, instance_type):  # pylint disable=unused-argumen
     def _upload_code(self, key_prefix):
         local_code = get_config_value('local.local_code', self.sagemaker_session.config)
         if self.sagemaker_session.local_mode and local_code:
-            script_name = self.entry_point
-            dir_name = 'file://' + self.source_dir
-            self.uploaded_code = UploadedCode(s3_prefix=dir_name, script_name=script_name)
+            self.uploaded_code = None
         else:
             self.uploaded_code = tar_and_upload_dir(session=self.sagemaker_session.boto_session,
                                                     bucket=self.bucket or self.sagemaker_session.default_bucket(),
@@ -204,7 +202,7 @@ def _framework_env_vars(self):
             dir_name = self.uploaded_code.s3_prefix
         else:
             script_name = self.entry_point
-            dir_name = self.source_dir
+            dir_name = 'file://' + self.source_dir
 
         return {
             SCRIPT_PARAM_NAME.upper(): script_name,

src/sagemaker/workflow/airflow.py

@@ -293,22 +293,13 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
         deploy_image = fw_utils.create_image_uri(
             region_name, model.__framework_name__, instance_type, model.framework_version, model.py_version)
 
-    deploy_env = dict(model.env)
-    deploy_env.update(model._framework_env_vars())
-    try:
-        if model.model_server_workers:
-            deploy_env[sagemaker.model.MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(model.model_server_workers)
-    except AttributeError:
-        # This applies to a FrameworkModel which is not SageMaker Deep Learning Framework Model
-        pass
-
-    container_def = sagemaker.container_def(deploy_image, model.model_data, deploy_env)
-    base_name = utils.base_name_from_image(container_def['Image'])
+    base_name = utils.base_name_from_image(deploy_image)
     model.name = model.name or utils.airflow_name_from_base(base_name)
 
     bucket = model.bucket or model.sagemaker_session._default_bucket
-    key = '{}/source/sourcedir.tar.gz'.format(model.name)
     script = os.path.basename(model.entry_point)
+    key = '{}/source/sourcedir.tar.gz'.format(model.name)
+
     if model.source_dir and model.source_dir.lower().startswith('s3://'):
         model.uploaded_code = fw_utils.UploadedCode(s3_prefix=model.source_dir, script_name=script)
     else:
@@ -321,6 +312,18 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
             'Tar': True
         }]
 
+    deploy_env = dict(model.env)
+    deploy_env.update(model._framework_env_vars())
+
+    try:
+        if model.model_server_workers:
+            deploy_env[sagemaker.model.MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(model.model_server_workers)
+    except AttributeError:
+        # This applies to a FrameworkModel which is not SageMaker Deep Learning Framework Model
+        pass
+
+    container_def = sagemaker.container_def(deploy_image, model.model_data, deploy_env)
+
     return container_def
 
 

tests/unit/test_airflow.py

@@ -609,7 +609,7 @@ def test_byo_framework_model_config(sagemaker_session):
             'Environment': {
                 '{{ key }}': '{{ value }}',
                 'SAGEMAKER_PROGRAM': '{{ entry_point }}',
-                'SAGEMAKER_SUBMIT_DIRECTORY': '{{ source_dir }}',
+                'SAGEMAKER_SUBMIT_DIRECTORY': 's3://output/model/source/sourcedir.tar.gz',
                 'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
                 'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
                 'SAGEMAKER_REGION': 'us-west-2'
@@ -648,7 +648,9 @@ def test_framework_model_config(sagemaker_session):
             'Image': '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-chainer:5.0.0-cpu-py3',
             'Environment': {
                 'SAGEMAKER_PROGRAM': '{{ entry_point }}',
-                'SAGEMAKER_SUBMIT_DIRECTORY': '{{ source_dir }}',
+                'SAGEMAKER_SUBMIT_DIRECTORY': "s3://output/sagemaker-chainer-"
+                                              "{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}"
+                                              "/source/sourcedir.tar.gz",
                 'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
                 'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
                 'SAGEMAKER_REGION': 'us-west-2',