Skip to content

Upgrade S3 ListObjects version, allow >1000 model files #225

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jun 29, 2023
Merged

Upgrade S3 ListObjects version, allow >1000 model files #225

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
56971f0
Add pagination when over 1000 files
the-david-oy Jun 14, 2023
692cedc
Auto-format
the-david-oy Jun 14, 2023
bff4fab
Use correct header
the-david-oy Jun 21, 2023
9c4c333
Switch to V2 request/result
the-david-oy Jun 21, 2023
a135398
Auto-format
the-david-oy Jun 21, 2023
a2d86ca
Switch if/else branch for readability
the-david-oy Jun 26, 2023
7e1bfb1
Remove else branch
the-david-oy Jun 28, 2023
b02d4f7
Fix if/else braces
the-david-oy Jun 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 29 additions & 17 deletions src/filesystem/implementations/s3.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@
#include <aws/s3/model/GetObjectRequest.h>
#include <aws/s3/model/HeadBucketRequest.h>
#include <aws/s3/model/HeadObjectRequest.h>
#include <aws/s3/model/ListObjectsRequest.h>
#include <aws/s3/model/ListObjectsV2Request.h>
#include <aws/s3/model/ListObjectsV2Result.h>

namespace triton { namespace core {

Expand Down Expand Up @@ -419,10 +420,10 @@ S3FileSystem::IsDirectory(const std::string& path, bool* is_dir)
}

// List the objects in the bucket
s3::Model::ListObjectsRequest list_objects_request;
s3::Model::ListObjectsV2Request list_objects_request;
list_objects_request.SetBucket(bucket.c_str());
list_objects_request.SetPrefix(AppendSlash(object_path).c_str());
auto list_objects_outcome = client_->ListObjects(list_objects_request);
auto list_objects_outcome = client_->ListObjectsV2(list_objects_request);

if (list_objects_outcome.IsSuccess()) {
*is_dir = !list_objects_outcome.GetResult().GetContents().empty();
Expand Down Expand Up @@ -484,16 +485,27 @@ S3FileSystem::GetDirectoryContents(
full_dir = AppendSlash(dir_path);

// Issue request for objects with prefix
s3::Model::ListObjectsRequest objects_request;
s3::Model::ListObjectsV2Request objects_request;
objects_request.SetBucket(bucket.c_str());
objects_request.SetPrefix(full_dir.c_str());
auto list_objects_outcome = client_->ListObjects(objects_request);

if (list_objects_outcome.IsSuccess()) {
Aws::Vector<Aws::S3::Model::Object> object_list =
list_objects_outcome.GetResult().GetContents();
for (auto const& s3_object : object_list) {
// In the case of empty directories, the directory itself will appear here
bool done_listing = false;
while (!done_listing) {
auto list_objects_outcome = client_->ListObjectsV2(objects_request);

if (!list_objects_outcome.IsSuccess()) {
return Status(
Status::Code::INTERNAL,
"Could not list contents of directory at " + true_path +
" due to exception: " +
list_objects_outcome.GetError().GetExceptionName() +
", error message: " +
list_objects_outcome.GetError().GetMessage());
Copy link
Contributor

@GuanLuo GuanLuo Jun 24, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Put this before the success handling, so the failure handle is a bit more obvious:

if (not_success) {
  return error;
}
// Check content..

Copy link
Contributor Author

@the-david-oy the-david-oy Jun 26, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great idea! Done.

Copy link
Contributor

@kthui kthui Jun 28, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think @GuanLuo means this?

while (!done_listing) {
  auto list_objects_outcome = client_->ListObjectsV2(objects_request);

  if (!list_objects_outcome.IsSuccess()) {
    return Status(
        Status::Code::INTERNAL,
        "Could not list contents of directory at " + true_path +
            " due to exception: " +
            list_objects_outcome.GetError().GetExceptionName() +
            ", error message: " +
            list_objects_outcome.GetError().GetMessage());
  }

  const auto& list_objects_result = list_objects_outcome.GetResult();
  for (const auto& s3_object : list_objects_result.GetContents()) {
    // In the case of empty directories, the directory itself will appear here
    if (s3_object.GetKey().c_str() == full_dir) {
      continue;
    }
  
    // We have to make sure that subdirectory contents do not appear here
    std::string name(s3_object.GetKey().c_str());
    int item_start = name.find(full_dir) + full_dir.size();
    // S3 response prepends parent directory name
    int item_end = name.find("/", item_start);
  
    // Let set take care of subdirectory contents
    std::string item = name.substr(item_start, item_end - item_start);
    contents->insert(item);
  
    // Fail-safe check to ensure the item name is not empty
    if (item.empty()) {
      return Status(
          Status::Code::INTERNAL,
          "Cannot handle item with empty name at " + true_path);
    }
  }
  // If there are more pages to retrieve, set the marker to the next page.
  if (list_objects_result.GetIsTruncated()) {
    objects_request.SetContinuationToken(
        list_objects_result.GetNextContinuationToken());
  } else {
    done_listing = true;
  }
}

Copy link
Contributor

@GuanLuo GuanLuo Jun 28, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this have one less nesting for the success case

Copy link
Contributor Author

@the-david-oy the-david-oy Jun 28, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, that makes sense! Thanks for clarifying, @kthui and @GuanLuo. Done.

Copy link
Contributor Author

@the-david-oy the-david-oy Jun 29, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still passes CI after removing the nesting. Please review when you get a chance!

}
const auto& list_objects_result = list_objects_outcome.GetResult();
for (const auto& s3_object : list_objects_result.GetContents()) {
// In the case of empty directories, the directory itself will appear
// here
if (s3_object.GetKey().c_str() == full_dir) {
continue;
}
Expand All @@ -515,13 +527,13 @@ S3FileSystem::GetDirectoryContents(
"Cannot handle item with empty name at " + true_path);
}
}
} else {
return Status(
Status::Code::INTERNAL,
"Could not list contents of directory at " + true_path +
" due to exception: " +
list_objects_outcome.GetError().GetExceptionName() +
", error message: " + list_objects_outcome.GetError().GetMessage());
// If there are more pages to retrieve, set the marker to the next page.
if (list_objects_result.GetIsTruncated()) {
objects_request.SetContinuationToken(
list_objects_result.GetNextContinuationToken());
} else {
done_listing = true;
}
}
return Status::Success;
}
Expand Down