Merge remote-tracking branch 'public/master' into pytorch-release

Author: nadiaya

CHANGELOG.rst

@@ -7,6 +7,9 @@ CHANGELOG
 
 * bug-fix: Unit Tests: Improve unit test runtime
 * bug-fix: Estimators: Fix attach for LDA
+* bug-fix: Estimators: allow code_location to have no key prefix
+* bug-fix: Local Mode: Fix s3 training data download when there is a trailing slash
+
 
 1.4.1
 =====

src/sagemaker/estimator.py

@@ -571,7 +571,7 @@ def _stage_user_code_in_s3(self):
             code_s3_prefix = '{}/source'.format(self._current_job_name)
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)
-            code_s3_prefix = '{}/{}/source'.format(key_prefix, self._current_job_name)
+            code_s3_prefix = '/'.join(filter(None, [key_prefix, self._current_job_name, 'source']))
 
         return tar_and_upload_dir(session=self.sagemaker_session.boto_session,
                                   bucket=code_bucket,

src/sagemaker/local/image.py

@@ -267,15 +267,15 @@ def _download_folder(self, bucket_name, prefix, target):
 
         for obj_sum in bucket.objects.filter(Prefix=prefix):
             obj = s3.Object(obj_sum.bucket_name, obj_sum.key)
-            file_path = os.path.join(target, obj_sum.key[len(prefix) + 1:])
+            s3_relative_path = obj_sum.key[len(prefix):].lstrip('/')
+            file_path = os.path.join(target, s3_relative_path)
 
             try:
                 os.makedirs(os.path.dirname(file_path))
             except OSError as exc:
                 if exc.errno != errno.EEXIST:
                     raise
                 pass
-
             obj.download_file(file_path)
 
     def _prepare_training_volumes(self, data_dir, input_data_config, hyperparameters):

tests/unit/test_estimator.py

@@ -127,6 +127,24 @@ def test_custom_code_bucket(time, sagemaker_session):
     assert train_kwargs['hyperparameters']['sagemaker_submit_directory'] == json.dumps(expected_submit_dir)
 
 
+@patch('time.strftime', return_value=TIMESTAMP)
+def test_custom_code_bucket_without_prefix(time, sagemaker_session):
+    code_bucket = 'codebucket'
+    code_location = 's3://{}'.format(code_bucket)
+    t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
+                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
+                       code_location=code_location)
+    t.fit('s3://bucket/mydata')
+
+    expected_key = '{}/source/sourcedir.tar.gz'.format(JOB_NAME)
+    _, s3_args, _ = sagemaker_session.boto_session.resource('s3').Object.mock_calls[0]
+    assert s3_args == (code_bucket, expected_key)
+
+    expected_submit_dir = 's3://{}/{}'.format(code_bucket, expected_key)
+    _, _, train_kwargs = sagemaker_session.train.mock_calls[0]
+    assert train_kwargs['hyperparameters']['sagemaker_submit_directory'] == json.dumps(expected_submit_dir)
+
+
 def test_invalid_custom_code_bucket(sagemaker_session):
     code_location = 'thisllworkright?'
     t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,

tests/unit/test_image.py

@@ -366,6 +366,15 @@ def test_download_folder(makedirs):
     calls = [call(os.path.join('/tmp', 'train/train_data.csv')),
              call(os.path.join('/tmp', 'train/validation_data.csv'))]
     obj_mock.download_file.assert_has_calls(calls)
+    obj_mock.reset_mock()
+
+    # Testing with a trailing slash for the prefix.
+    sagemaker_container._download_folder(BUCKET_NAME, '/prefix/', '/tmp')
+    obj_mock.download_file.assert_called()
+    calls = [call(os.path.join('/tmp', 'train/train_data.csv')),
+             call(os.path.join('/tmp', 'train/validation_data.csv'))]
+
+    obj_mock.download_file.assert_has_calls(calls)
 
 
 def test_ecr_login_non_ecr():