fix: correct S3Downloader behavior (#292)

Author: knakad

src/sagemaker/session.py

@@ -239,11 +239,11 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
 
         # Initialize the variables used to loop through the contents of the S3 bucket.
         keys = []
-        directories = []
         next_token = ""
         base_parameters = {"Bucket": bucket, "Prefix": key_prefix}
 
-        # Loop through the contents of the bucket, 1,000 objects at a time.
+        # Loop through the contents of the bucket, 1,000 objects at a time. Gathering all keys into
+        # a "keys" list.
         while next_token is not None:
             request_parameters = base_parameters.copy()
             if next_token != "":
@@ -253,26 +253,20 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
             # For each object, save its key or directory.
             for s3_object in contents:
                 key = s3_object.get("Key")
-                if key[-1] != "/":
-                    keys.append(key)
-                else:
-                    directories.append(key)
+                keys.append(key)
             next_token = response.get("NextContinuationToken")
 
-        # For each directory, create the directory on the local machine.
-        for directory in directories:
-            destination_path = os.path.join(path, directory)
-            if not os.path.exists(os.path.dirname(destination_path)):
-                os.makedirs(os.path.dirname(destination_path))
-
-        # For each object key, create the directory on the local machine,
-        # and then download the file.
+        # For each object key, create the directory on the local machine if needed, and then
+        # download the file.
         for key in keys:
-            destination_path = os.path.join(path, key)
+            tail_s3_uri_path = os.path.basename(key_prefix)
+            if not os.path.splitext(key_prefix)[1]:
+                tail_s3_uri_path = os.path.relpath(key, key_prefix)
+            destination_path = os.path.join(path, tail_s3_uri_path)
             if not os.path.exists(os.path.dirname(destination_path)):
                 os.makedirs(os.path.dirname(destination_path))
             s3.download_file(
-                bucket=bucket, key=key, filename=destination_path, extra_args=extra_args
+                Bucket=bucket, Key=key, Filename=destination_path, ExtraArgs=extra_args
             )
 
     def read_s3_file(self, bucket, key_prefix):

tests/integ/test_s3.py

@@ -23,21 +23,108 @@
 from tests.integ.kms_utils import get_or_create_kms_key
 
 
+TMP_BASE_PATH = "/tmp"
+
+
 @pytest.fixture(scope="module")
 def s3_files_kms_key(sagemaker_session):
     return get_or_create_kms_key(sagemaker_session=sagemaker_session)
 
 
-def test_statistics_object_creation_from_s3_uri_with_customizations(
+def test_s3_uploader_and_downloader_reads_files_when_given_file_name_uris(
     sagemaker_session, s3_files_kms_key
 ):
-    file_1_body = "First File Body."
-    file_1_name = "first_file.txt"
-    file_2_body = "Second File Body."
-    file_2_name = "second_file.txt"
+    my_uuid = str(uuid.uuid4())
+
+    file_1_body = "First File Body {}.".format(my_uuid)
+    file_1_name = "first_file_{}.txt".format(my_uuid)
+    file_2_body = "Second File Body {}.".format(my_uuid)
+    file_2_name = "second_file_{}.txt".format(my_uuid)
+
+    base_s3_uri = os.path.join(
+        "s3://", sagemaker_session.default_bucket(), "integ-test-test-s3-list", my_uuid
+    )
+    file_1_s3_uri = os.path.join(base_s3_uri, file_1_name)
+    file_2_s3_uri = os.path.join(base_s3_uri, file_2_name)
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_1_body,
+        desired_s3_uri=file_1_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_2_body,
+        desired_s3_uri=file_2_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    s3_uris = S3Downloader.list(s3_uri=base_s3_uri, session=sagemaker_session)
+
+    assert file_1_name in s3_uris[0]
+    assert file_2_name in s3_uris[1]
+
+    assert file_1_body == S3Downloader.read_file(s3_uri=s3_uris[0], session=sagemaker_session)
+    assert file_2_body == S3Downloader.read_file(s3_uri=s3_uris[1], session=sagemaker_session)
+
+
+def test_s3_uploader_and_downloader_downloads_files_when_given_file_name_uris(
+    sagemaker_session, s3_files_kms_key
+):
+    my_uuid = str(uuid.uuid4())
+
+    file_1_body = "First File Body {}.".format(my_uuid)
+    file_1_name = "first_file_{}.txt".format(my_uuid)
+    file_2_body = "Second File Body {}.".format(my_uuid)
+    file_2_name = "second_file_{}.txt".format(my_uuid)
+
+    base_s3_uri = os.path.join(
+        "s3://", sagemaker_session.default_bucket(), "integ-test-test-s3-list", my_uuid
+    )
+    file_1_s3_uri = os.path.join(base_s3_uri, file_1_name)
+    file_2_s3_uri = os.path.join(base_s3_uri, file_2_name)
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_1_body,
+        desired_s3_uri=file_1_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_2_body,
+        desired_s3_uri=file_2_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    s3_uris = S3Downloader.list(s3_uri=base_s3_uri, session=sagemaker_session)
+
+    assert file_1_name in s3_uris[0]
+    assert file_2_name in s3_uris[1]
+
+    S3Downloader.download(s3_uri=s3_uris[0], local_path=TMP_BASE_PATH, session=sagemaker_session)
+    S3Downloader.download(s3_uri=s3_uris[1], local_path=TMP_BASE_PATH, session=sagemaker_session)
+
+    with open(os.path.join(TMP_BASE_PATH, file_1_name), "r") as f:
+        assert file_1_body == f.read()
 
+    with open(os.path.join(TMP_BASE_PATH, file_2_name), "r") as f:
+        assert file_2_body == f.read()
+
+
+def test_s3_uploader_and_downloader_downloads_files_when_given_directory_uris_with_files(
+    sagemaker_session, s3_files_kms_key
+):
     my_uuid = str(uuid.uuid4())
 
+    file_1_body = "First File Body {}.".format(my_uuid)
+    file_1_name = "first_file_{}.txt".format(my_uuid)
+    file_2_body = "Second File Body {}.".format(my_uuid)
+    file_2_name = "second_file_{}.txt".format(my_uuid)
+
     base_s3_uri = os.path.join(
         "s3://", sagemaker_session.default_bucket(), "integ-test-test-s3-list", my_uuid
     )
@@ -65,3 +152,71 @@ def test_statistics_object_creation_from_s3_uri_with_customizations(
 
     assert file_1_body == S3Downloader.read_file(s3_uri=s3_uris[0], session=sagemaker_session)
     assert file_2_body == S3Downloader.read_file(s3_uri=s3_uris[1], session=sagemaker_session)
+
+    S3Downloader.download(s3_uri=base_s3_uri, local_path=TMP_BASE_PATH, session=sagemaker_session)
+
+    with open(os.path.join(TMP_BASE_PATH, file_1_name), "r") as f:
+        assert file_1_body == f.read()
+
+    with open(os.path.join(TMP_BASE_PATH, file_2_name), "r") as f:
+        assert file_2_body == f.read()
+
+
+def test_s3_uploader_and_downloader_downloads_files_when_given_directory_uris_with_directory(
+    sagemaker_session, s3_files_kms_key
+):
+    my_uuid = str(uuid.uuid4())
+    my_inner_directory_uuid = str(uuid.uuid4())
+
+    file_1_body = "First File Body {}.".format(my_uuid)
+    file_1_name = "first_file_{}.txt".format(my_uuid)
+    file_2_body = "Second File Body {}.".format(my_uuid)
+    file_2_name = "second_file_{}.txt".format(my_uuid)
+
+    base_s3_uri = os.path.join(
+        "s3://",
+        sagemaker_session.default_bucket(),
+        "integ-test-test-s3-list",
+        my_uuid,
+        my_inner_directory_uuid,
+    )
+    file_1_s3_uri = os.path.join(base_s3_uri, file_1_name)
+    file_2_s3_uri = os.path.join(base_s3_uri, file_2_name)
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_1_body,
+        desired_s3_uri=file_1_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_2_body,
+        desired_s3_uri=file_2_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    s3_uris = S3Downloader.list(s3_uri=base_s3_uri, session=sagemaker_session)
+
+    assert file_1_name in s3_uris[0]
+    assert file_2_name in s3_uris[1]
+
+    assert file_1_body == S3Downloader.read_file(s3_uri=s3_uris[0], session=sagemaker_session)
+    assert file_2_body == S3Downloader.read_file(s3_uri=s3_uris[1], session=sagemaker_session)
+
+    s3_directory_with_directory_underneath = os.path.join(
+        "s3://", sagemaker_session.default_bucket(), "integ-test-test-s3-list", my_uuid
+    )
+
+    S3Downloader.download(
+        s3_uri=s3_directory_with_directory_underneath,
+        local_path=TMP_BASE_PATH,
+        session=sagemaker_session,
+    )
+
+    with open(os.path.join(TMP_BASE_PATH, my_inner_directory_uuid, file_1_name), "r") as f:
+        assert file_1_body == f.read()
+
+    with open(os.path.join(TMP_BASE_PATH, my_inner_directory_uuid, file_2_name), "r") as f:
+        assert file_2_body == f.read()