fix: allow download_folder to download file even if bucket is more restricted (#1295)

Author: laurenyu

src/sagemaker/utils.py

@@ -23,13 +23,12 @@
 import tarfile
 import tempfile
 import time
-
-
 from datetime import datetime
 from functools import wraps
 
 import six
 from six.moves.urllib import parse
+import botocore
 
 
 ECR_URI_PATTERN = r"^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(amazonaws.com|c2s.ic.gov)(/)(.*:.*)$"
@@ -338,21 +337,34 @@ def download_folder(bucket_name, prefix, target, sagemaker_session):
             interact with S3.
     """
     boto_session = sagemaker_session.boto_session
-
-    s3 = boto_session.resource("s3")
-    bucket = s3.Bucket(bucket_name)
+    s3 = boto_session.resource("s3", region_name=boto_session.region_name)
 
     prefix = prefix.lstrip("/")
 
-    # there is a chance that the prefix points to a file and not a 'directory' if that is the case
-    # we should just download it.
-    objects = list(bucket.objects.filter(Prefix=prefix))
-
-    if len(objects) > 0 and objects[0].key == prefix and prefix[-1] != "/":
+    # Try to download the prefix as an object first, in case it is a file and not a 'directory'.
+    # Do this first, in case the object has broader permissions than the bucket.
+    try:
         s3.Object(bucket_name, prefix).download_file(os.path.join(target, os.path.basename(prefix)))
         return
+    except botocore.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "404" and e.response["Error"]["Message"] == "Not Found":
+            # S3 also throws this error if the object is a folder,
+            # so assume that is the case here, and then raise for an actual 404 later.
+            _download_files_under_prefix(bucket_name, prefix, target, s3)
+        else:
+            raise
 
-    # the prefix points to an s3 'directory' download the whole thing
+
+def _download_files_under_prefix(bucket_name, prefix, target, s3):
+    """Download all S3 files which match the given prefix
+
+    Args:
+        bucket_name (str): S3 bucket name
+        prefix (str): S3 prefix within the bucket that will be downloaded
+        target (str): destination path where the downloaded items will be placed
+        s3 (boto3.resources.base.ServiceResource): S3 resource
+    """
+    bucket = s3.Bucket(bucket_name)
     for obj_sum in bucket.objects.filter(Prefix=prefix):
         # if obj_sum is a folder object skip it.
         if obj_sum.key != "" and obj_sum.key[-1] == "/":

tests/unit/test_utils.py

@@ -310,9 +310,24 @@ def test_generate_tensorboard_url_domain_non_string():
 @patch("os.makedirs")
 def test_download_folder(makedirs):
     boto_mock = Mock(name="boto_session")
-    boto_mock.client("sts").get_caller_identity.return_value = {"Account": "123"}
-
     session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
+    s3_mock = boto_mock.resource("s3")
+
+    obj_mock = Mock()
+    s3_mock.Object.return_value = obj_mock
+
+    def obj_mock_download(path):
+        # Mock the S3 object to raise an error when the input to download_file
+        # is a "folder"
+        if path in ("/tmp/", os.path.join("/tmp", "prefix")):
+            raise botocore.exceptions.ClientError(
+                error_response={"Error": {"Code": "404", "Message": "Not Found"}},
+                operation_name="HeadObject",
+            )
+        else:
+            return Mock()
+
+    obj_mock.download_file.side_effect = obj_mock_download
 
     train_data = Mock()
     validation_data = Mock()
@@ -323,23 +338,20 @@ def test_download_folder(makedirs):
     validation_data.key = "prefix/train/validation_data.csv"
 
     s3_files = [train_data, validation_data]
-    boto_mock.resource("s3").Bucket(BUCKET_NAME).objects.filter.return_value = s3_files
-
-    obj_mock = Mock()
-    boto_mock.resource("s3").Object.return_value = obj_mock
+    s3_mock.Bucket(BUCKET_NAME).objects.filter.return_value = s3_files
 
     # all the S3 mocks are set, the test itself begins now.
     sagemaker.utils.download_folder(BUCKET_NAME, "/prefix", "/tmp", session)
 
     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")),
+        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.
+    # Test with a trailing slash for the prefix.
     sagemaker.utils.download_folder(BUCKET_NAME, "/prefix/", "/tmp", session)
     obj_mock.download_file.assert_called()
     obj_mock.download_file.assert_has_calls(calls)
@@ -369,7 +381,7 @@ def test_download_folder_points_to_single_file(makedirs):
     obj_mock.download_file.assert_called()
     calls = [call(os.path.join("/tmp", "train_data.csv"))]
     obj_mock.download_file.assert_has_calls(calls)
-    assert boto_mock.resource("s3").Bucket(BUCKET_NAME).objects.filter.call_count == 1
+    boto_mock.resource("s3").Bucket(BUCKET_NAME).objects.filter.assert_not_called()
     obj_mock.reset_mock()