Fix Local Mode s3 training data download.

If the S3 training location is passed as
s3://my_bucket/training/ (with a trailing slash) Then the objects
would have incorrect filenames locally and the container would not find
them. This fixes this bug.

Author: Ignacio Quintero

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_image.py

@@ -366,6 +366,14 @@ 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()
+
+    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():