Fix container path for source files when using local_session (#499)

Author: bearpelican

CHANGELOG.rst

@@ -2,6 +2,11 @@
 CHANGELOG
 =========
 
+1.16.3
+======
+
+* bug-fix: Local Mode: No longer requires s3 permissions to run local entry point file
+
 1.16.2
 ======
 

src/sagemaker/local/image.py

@@ -106,6 +106,8 @@ def train(self, input_data_config, output_data_config, hyperparameters, job_name
         data_dir = self._create_tmp_folder()
         volumes = self._prepare_training_volumes(data_dir, input_data_config, output_data_config,
                                                  hyperparameters)
+        # If local, source directory needs to be updated to mounted /opt/ml/code path
+        hyperparameters = self._update_local_src_path(hyperparameters, key=sagemaker.estimator.DIR_PARAM_NAME)
 
         # Create the configuration files for each container that we will create
         # Each container will map the additional local volumes (if any).
@@ -169,6 +171,9 @@ def serve(self, model_dir, environment):
             parsed_uri = urlparse(script_dir)
             if parsed_uri.scheme == 'file':
                 volumes.append(_Volume(parsed_uri.path, '/opt/ml/code'))
+                # Update path to mount location
+                environment = environment.copy()
+                environment[sagemaker.estimator.DIR_PARAM_NAME.upper()] = '/opt/ml/code'
 
         if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
             _pull_image(self.image)
@@ -302,7 +307,7 @@ def _prepare_training_volumes(self, data_dir, input_data_config, output_data_con
             volumes.append(_Volume(data_source.get_root_dir(), channel=channel_name))
 
         # If there is a training script directory and it is a local directory,
-        #  mount it to the container.
+        # mount it to the container.
         if sagemaker.estimator.DIR_PARAM_NAME in hyperparameters:
             training_dir = json.loads(hyperparameters[sagemaker.estimator.DIR_PARAM_NAME])
             parsed_uri = urlparse(training_dir)
@@ -321,6 +326,16 @@ def _prepare_training_volumes(self, data_dir, input_data_config, output_data_con
 
         return volumes
 
+    def _update_local_src_path(self, params, key):
+        if key in params:
+            src_dir = json.loads(params[key])
+            parsed_uri = urlparse(src_dir)
+            if parsed_uri.scheme == 'file':
+                new_params = params.copy()
+                new_params[key] = json.dumps('/opt/ml/code')
+                return new_params
+        return params
+
     def _prepare_serving_volumes(self, model_location):
         volumes = []
         host = self.hosts[0]

tests/unit/test_image.py

@@ -388,12 +388,18 @@ def test_train_local_code(tmpdir, sagemaker_session):
         with open(docker_compose_file, 'r') as f:
             config = yaml.load(f)
             assert len(config['services']) == instance_count
-            for h in sagemaker_container.hosts:
-                assert config['services'][h]['image'] == image
-                assert config['services'][h]['command'] == 'train'
-                volumes = config['services'][h]['volumes']
-                assert '%s:/opt/ml/code' % '/tmp/code' in volumes
-                assert '%s:/opt/ml/shared' % shared_folder_path in volumes
+
+        for h in sagemaker_container.hosts:
+            assert config['services'][h]['image'] == image
+            assert config['services'][h]['command'] == 'train'
+            volumes = config['services'][h]['volumes']
+            assert '%s:/opt/ml/code' % '/tmp/code' in volumes
+            assert '%s:/opt/ml/shared' % shared_folder_path in volumes
+
+            config_file_root = os.path.join(sagemaker_container.container_root, h, 'input', 'config')
+            hyperparameters_file = os.path.join(config_file_root, 'hyperparameters.json')
+            hyperparameters_data = json.load(open(hyperparameters_file))
+            assert hyperparameters_data['sagemaker_submit_directory'] == json.dumps('/opt/ml/code')
 
 
 @patch('sagemaker.local.local_session.LocalSession', Mock())
@@ -506,6 +512,7 @@ def test_serve_local_code(tmpdir, sagemaker_session):
 
                 volumes = config['services'][h]['volumes']
                 assert '%s:/opt/ml/code' % '/tmp/code' in volumes
+                assert 'SAGEMAKER_SUBMIT_DIRECTORY=/opt/ml/code' in config['services'][h]['environment']
 
 
 @patch('sagemaker.local.image._HostingContainer.run', Mock())