diff --git a/CMakeLists.txt b/CMakeLists.txt index 08711c1ded0..3d19e728698 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -101,6 +101,8 @@ cmake_dependent_option(BUILD_NVIMAGECODEC "Build with support for nvimagecodec l "NOT BUILD_DALI_NODEPS" OFF) cmake_dependent_option(BUILD_AWSSDK "Build with support for AWS SKD library" ON "NOT BUILD_DALI_NODEPS" OFF) +cmake_dependent_option(BUILD_GCS "Build with support for Google Cloud Storage library" ON + "NOT BUILD_DALI_NODEPS" OFF) set(NVIMGCODEC_DEFAULT_INSTALL_PATH "/opt/nvidia/nvimgcodec_cuda${CUDA_VERSION_MAJOR}" CACHE STRING "Path of the nvimagecodec installation") @@ -344,6 +346,7 @@ propagate_option(BUILD_NVML) propagate_option(BUILD_CUFILE) propagate_option(BUILD_NVIMAGECODEC) propagate_option(BUILD_AWSSDK) +propagate_option(BUILD_GCS) propagate_option(BUILD_FOR_CONDA) propagate_option(LINK_DRIVER) propagate_option(WITH_DYNAMIC_CUFFT) diff --git a/cmake/Dependencies.common.cmake b/cmake/Dependencies.common.cmake index ac3697eb705..ecf61828dc8 100644 --- a/cmake/Dependencies.common.cmake +++ b/cmake/Dependencies.common.cmake @@ -452,3 +452,16 @@ if(BUILD_AWSSDK) message(STATUS "AWSSDK_LIBRARIES=${AWSSDK_LIBRARIES}") endif() endif() + +################################################################## +# Google Cloud Storage (google-cloud-cpp) +################################################################## +if(BUILD_GCS) + find_package(google_cloud_cpp_storage CONFIG QUIET) + if (NOT google_cloud_cpp_storage_FOUND) + message(WARNING "google-cloud-cpp storage not found. Disabling Google Cloud Storage support.") + set(BUILD_GCS OFF) + else() + message(STATUS "google-cloud-cpp storage version=${google_cloud_cpp_storage_VERSION}") + endif() +endif() diff --git a/dali/CMakeLists.txt b/dali/CMakeLists.txt index 41a58db4c46..a5fa2df51be 100644 --- a/dali/CMakeLists.txt +++ b/dali/CMakeLists.txt @@ -94,6 +94,10 @@ if (BUILD_AWSSDK) target_link_libraries(dali PRIVATE ${AWSSDK_LIBRARIES}) endif() +if (BUILD_GCS) + target_link_libraries(dali PRIVATE google-cloud-cpp::storage) +endif() + # Build test suite ################################################ if (BUILD_DALI_PIPELINE AND BUILD_TEST) diff --git a/dali/operators/CMakeLists.txt b/dali/operators/CMakeLists.txt index 91ca0cdb78e..a51c241d098 100644 --- a/dali/operators/CMakeLists.txt +++ b/dali/operators/CMakeLists.txt @@ -102,6 +102,12 @@ if (BUILD_CVCUDA) target_link_libraries(dali_operators PRIVATE cvcuda nvcv_types) endif(BUILD_CVCUDA) +if (BUILD_GCS) + # GCSClientManager is header-only, so discover_files_gcs.cc needs the client library directly - + # it cannot rely on it being reachable through libdali. + target_link_libraries(dali_operators PRIVATE google-cloud-cpp::storage) +endif() + if (BUILD_NVIMAGECODEC) if (WITH_DYNAMIC_NVIMGCODEC) target_link_libraries(dali_operators PRIVATE dynlink_nvimgcodec) diff --git a/dali/operators/reader/loader/CMakeLists.txt b/dali/operators/reader/loader/CMakeLists.txt index 786c0aa02a5..e14b59c557f 100644 --- a/dali/operators/reader/loader/CMakeLists.txt +++ b/dali/operators/reader/loader/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2017-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -67,5 +67,10 @@ if (BUILD_AWSSDK) "${CMAKE_CURRENT_SOURCE_DIR}/discover_files_s3.cc") endif() +if (BUILD_GCS) + set(DALI_OPERATOR_SRCS ${DALI_OPERATOR_SRCS} + "${CMAKE_CURRENT_SOURCE_DIR}/discover_files_gcs.cc") +endif() + set(DALI_OPERATOR_SRCS ${DALI_OPERATOR_SRCS} PARENT_SCOPE) set(DALI_OPERATOR_TEST_SRCS ${DALI_OPERATOR_TEST_SRCS} PARENT_SCOPE) diff --git a/dali/operators/reader/loader/discover_files.cc b/dali/operators/reader/loader/discover_files.cc index 9890a857116..74723309f1b 100644 --- a/dali/operators/reader/loader/discover_files.cc +++ b/dali/operators/reader/loader/discover_files.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2017-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -33,6 +33,9 @@ #if AWSSDK_ENABLED #include "dali/operators/reader/loader/discover_files_s3.h" #endif +#if GCS_ENABLED +#include "dali/operators/reader/loader/discover_files_gcs.h" +#endif namespace dali { @@ -122,6 +125,15 @@ std::vector discover_files(const std::string &file_root, #endif } + bool is_gcs = starts_with(file_root, "gs://"); + if (is_gcs) { +#if GCS_ENABLED + return gcs_discover_files(file_root, opts); +#else + DALI_FAIL("This version of DALI was not built with Google Cloud Storage support."); +#endif + } + std::vector subdirs; subdirs = list_subdirectories(file_root, opts.dir_filters, opts.case_sensitive_filter); std::vector entries; diff --git a/dali/operators/reader/loader/discover_files_gcs.cc b/dali/operators/reader/loader/discover_files_gcs.cc new file mode 100644 index 00000000000..d0df3a71d06 --- /dev/null +++ b/dali/operators/reader/loader/discover_files_gcs.cc @@ -0,0 +1,104 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dali/operators/reader/loader/discover_files_gcs.h" +#include +#include +#include +#include +#include +#include "dali/operators/reader/loader/discover_files.h" +#include "dali/util/gcs_client_manager.h" +#include "dali/util/gcs_filesystem.h" + +namespace dali { + +// We are using std::filesystem to analyze URI relative paths, which wouldn't be OK in non-UNIX +// based systems +#ifndef __unix__ +#error This code works only in UNIX-based systems +#endif + +std::vector gcs_discover_files(const std::string &file_root, + const FileDiscoveryOptions &opts) { + assert(starts_with(file_root, "gs://")); + auto gcs_object_location = gcs_filesystem::parse_uri(file_root); + std::filesystem::path parent_object_key(gcs_object_location.object); + auto count_elems = [](const std::filesystem::path &p) { + size_t k = 0; + for (auto &elem : p) + k++; + return k; + }; + std::vector entries; + // in case that files are not visited in lexicographical order, we remember previously assigned + // labels + std::unordered_map labels; + int next_label = 0; // next free-label to be assigned + auto client = GCSClientManager::Instance().client(); + gcs_filesystem::list_objects_f( + client, gcs_object_location, [&](const std::string &object_key, size_t object_size) { + auto p = std::filesystem::relative(object_key, parent_object_key); + auto path_elems = count_elems(p); + // We only look at one subdir level. Fewer than two components means either an object + // directly under the listed prefix, or the prefix's own directory marker, which relative() + // maps to "."; neither is a labelled file, and both must be rejected before dereferencing + // the second component below. + if (path_elems != 2) + return; + const auto& subdir = p.begin()->native(); + const auto& fname = (++p.begin())->native(); + // GCS directory markers are zero-byte objects whose name ends with '/'. A trailing + // separator becomes an empty final component, so "/class/" arrives here as + // ("class", "") - a directory, not a file. + if (fname.empty()) + return; + bool subdir_ok = opts.dir_filters.empty(); + bool fname_ok = opts.file_filters.empty(); + for (auto &filter : opts.dir_filters) { + if (fnmatch(filter.c_str(), subdir.c_str(), + opts.case_sensitive_filter ? 0 : FNM_CASEFOLD) == 0) { + subdir_ok |= true; + break; + } + } + + for (auto &filter : opts.file_filters) { + if (fnmatch(filter.c_str(), fname.c_str(), + opts.case_sensitive_filter ? 0 : FNM_CASEFOLD) == 0) { + fname_ok |= true; + break; + } + } + + if (!subdir_ok || !fname_ok) + return; + + if (opts.label_from_subdir) { + int curr_label = -1; + auto it = labels.find(subdir); + if (it == labels.end()) { + curr_label = labels[subdir] = next_label++; + } else { + curr_label = it->second; + } + entries.push_back({p, curr_label, object_size}); + } else { + entries.push_back({p, std::nullopt, object_size}); + } + }); + return entries; +} + +} // namespace dali diff --git a/dali/operators/reader/loader/discover_files_gcs.h b/dali/operators/reader/loader/discover_files_gcs.h new file mode 100644 index 00000000000..92bb5b12e91 --- /dev/null +++ b/dali/operators/reader/loader/discover_files_gcs.h @@ -0,0 +1,29 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_OPERATORS_READER_LOADER_DISCOVER_FILES_GCS_H_ +#define DALI_OPERATORS_READER_LOADER_DISCOVER_FILES_GCS_H_ + +#include +#include +#include "dali/operators/reader/loader/discover_files.h" + +namespace dali { + +std::vector gcs_discover_files(const std::string &file_root, + const FileDiscoveryOptions &opts); + +} // namespace dali + +#endif // DALI_OPERATORS_READER_LOADER_DISCOVER_FILES_GCS_H_ diff --git a/dali/util/CMakeLists.txt b/dali/util/CMakeLists.txt index 1a4c80ec7db..b5fb30c8d90 100644 --- a/dali/util/CMakeLists.txt +++ b/dali/util/CMakeLists.txt @@ -65,6 +65,12 @@ if (BUILD_AWSSDK) "${CMAKE_CURRENT_SOURCE_DIR}/s3_filesystem.cc") endif() +if (BUILD_GCS) + set(DALI_SRCS ${DALI_SRCS} + "${CMAKE_CURRENT_SOURCE_DIR}/gcs_file.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/gcs_filesystem.cc") +endif() + set(DALI_TEST_SRCS ${DALI_TEST_SRCS} "${CMAKE_CURRENT_SOURCE_DIR}/numpy_test.cc" "${CMAKE_CURRENT_SOURCE_DIR}/uri_test.cc") diff --git a/dali/util/file.cc b/dali/util/file.cc index e9a6ccb3891..b0c39e7161f 100644 --- a/dali/util/file.cc +++ b/dali/util/file.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2017-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2017-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -25,6 +25,11 @@ #include "dali/util/s3_file.h" #endif +#if GCS_ENABLED +#include "dali/util/gcs_client_manager.h" +#include "dali/util/gcs_file.h" +#endif + namespace dali { std::unique_ptr FileStream::Open(const std::string& uri, FileStream::Options opts, @@ -38,6 +43,16 @@ std::unique_ptr FileStream::Open(const std::string& uri, FileStream: #endif } + bool is_gcs = uri.rfind("gs://", 0) == 0; + if (is_gcs) { +#if GCS_ENABLED + return std::make_unique(GCSClientManager::Instance().client(), uri, size); +#else + throw std::runtime_error( + "This version of DALI was not built with Google Cloud Storage support."); +#endif + } + std::string processed_uri; if (uri.find("file://") == 0) { processed_uri = uri.substr(std::string("file://").size()); diff --git a/dali/util/gcs_client_manager.h b/dali/util/gcs_client_manager.h new file mode 100644 index 00000000000..3f36c9ec5dc --- /dev/null +++ b/dali/util/gcs_client_manager.h @@ -0,0 +1,99 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_UTIL_GCS_CLIENT_MANAGER_H_ +#define DALI_UTIL_GCS_CLIENT_MANAGER_H_ + +#include +#include +#include +#include +#include +#include +#include +#include "dali/core/common.h" +#include "dali/core/error_handling.h" + +namespace dali { + +/** + * @brief Owns the process-wide configuration of the GCS client. + * + * Unlike `Aws::S3::S3Client`, `google::cloud::storage::Client` is not documented as safe for + * concurrent use of a *single* instance ("Two threads operating on the same instance of this + * class is not guaranteed to work"). Copies, on the other hand, share the underlying connection + * pool and are explicitly safe to use from different threads, and copying is about as expensive + * as copying a few shared pointers. Therefore `client()` hands out a copy and each caller + * (file stream, file discovery) keeps its own. + * + * There is no global init/shutdown to perform - the library initializes libcurl lazily. + */ +class GCSClientManager { + public: + static GCSClientManager& Instance() { + static GCSClientManager s_manager_; + return s_manager_; + } + + /** + * @brief Returns a client sharing the connection pool with all other clients handed out here. + */ + google::cloud::storage::Client client() const { + return client_; + } + + private: + static bool EnvFlag(const char* name, bool default_value) { + auto* value = std::getenv(name); + if (!value) + return default_value; + return std::atoi(value) != 0; + } + + static google::cloud::Options MakeOptions() { + namespace gcs = google::cloud::storage; + google::cloud::Options options; + + // The library also honors CLOUD_STORAGE_EMULATOR_ENDPOINT on its own; this is the DALI-side + // counterpart of AWS_ENDPOINT_URL. + if (auto* endpoint_url = std::getenv("DALI_GCS_ENDPOINT_URL")) { + options.set(endpoint_url); + } + + // By default the client uses Application Default Credentials. Reading a public bucket (or a + // local emulator) requires opting out of authentication explicitly. + if (EnvFlag("DALI_GCS_ANONYMOUS", false)) { + options.set( + google::cloud::MakeInsecureCredentials()); + } + + // DALI only ever issues ranged reads, and GCS reports checksums for whole objects only, so a + // per-read CRC32C over the payload cannot be validated end-to-end - it would just burn CPU in + // the data loading path. It can be turned back on for debugging. + if (!EnvFlag("DALI_GCS_VERIFY_CHECKSUMS", false)) { + options.set( // NOLINT(build/include_what_you_use) + gcs::ChecksumAlgorithm::kNone); + } + + return options; + } + + GCSClientManager() : client_(MakeOptions()) {} + + google::cloud::storage::Client client_; +}; + +} // namespace dali + +#endif // DALI_UTIL_GCS_CLIENT_MANAGER_H_ diff --git a/dali/util/gcs_file.cc b/dali/util/gcs_file.cc new file mode 100644 index 00000000000..ad06425c662 --- /dev/null +++ b/dali/util/gcs_file.cc @@ -0,0 +1,78 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dali/util/gcs_file.h" +#include +#include "dali/core/format.h" +#include "dali/util/gcs_client_manager.h" +#include "dali/util/uri.h" + +namespace dali { + +GCSFileStream::GCSFileStream(google::cloud::storage::Client client, const std::string& uri, + std::optional size) + : FileStream(uri), client_(std::move(client)) { + object_location_ = gcs_filesystem::parse_uri(uri); + if (size.has_value() && size.value() > 0) { + object_stats_.exists = true; + object_stats_.size = size.value(); + } else { + object_stats_ = gcs_filesystem::get_stats(client_, object_location_); + } +} + +GCSFileStream::~GCSFileStream() {} + +void GCSFileStream::Close() { + // nothing to do here (there's no file open) +} + +void GCSFileStream::SeekRead(ptrdiff_t pos, int whence) { + auto new_pos = pos_; + switch (whence) { + case SEEK_SET: + new_pos = pos; + break; + case SEEK_CUR: + new_pos += pos; + break; + case SEEK_END: + new_pos = object_stats_.size + pos; + break; + default: + assert(false); + } + if (new_pos < 0 || new_pos > static_cast(object_stats_.size)) + throw std::out_of_range("The requested offset points outside of the file."); + pos_ = new_pos; +} + +ptrdiff_t GCSFileStream::TellRead() const { + return pos_; +} + +size_t GCSFileStream::Size() const { + return object_stats_.size; +} + +size_t GCSFileStream::Read(void* buf, size_t n) { + if (n == 0) + return 0; + size_t bytes_read = + gcs_filesystem::read_object_contents(client_, object_location_, buf, n, pos_); + pos_ += bytes_read; + return bytes_read; +} + +} // namespace dali diff --git a/dali/util/gcs_file.h b/dali/util/gcs_file.h new file mode 100644 index 00000000000..96583b9f122 --- /dev/null +++ b/dali/util/gcs_file.h @@ -0,0 +1,55 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_UTIL_GCS_FILE_H_ +#define DALI_UTIL_GCS_FILE_H_ + +#include +#include +#include +#include +#include +#include "dali/util/file.h" +#include "dali/util/gcs_filesystem.h" +#include "dali/util/uri.h" + +namespace dali { + +class GCSFileStream : public FileStream { + public: + /** + * @param client a client obtained from GCSClientManager. It is held by value, because + * `google::cloud::storage::Client` instances must not be shared between threads, + * while copies of them are cheap and share the connection pool. + */ + explicit GCSFileStream(google::cloud::storage::Client client, const std::string& uri, + std::optional size = std::nullopt); + void Close() override; + size_t Read(void* buf, size_t n) override; + void SeekRead(ptrdiff_t pos, int whence = SEEK_SET) override; + ptrdiff_t TellRead() const override; + size_t Size() const override; + + ~GCSFileStream() override; + + private: + google::cloud::storage::Client client_; + ptrdiff_t pos_ = 0; + gcs_filesystem::GCSObjectLocation object_location_ = {}; + gcs_filesystem::GCSObjectStats object_stats_ = {}; +}; + +} // namespace dali + +#endif // DALI_UTIL_GCS_FILE_H_ diff --git a/dali/util/gcs_filesystem.cc b/dali/util/gcs_filesystem.cc new file mode 100644 index 00000000000..20a60b89fa3 --- /dev/null +++ b/dali/util/gcs_filesystem.cc @@ -0,0 +1,117 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dali/util/gcs_filesystem.h" +#include +#include +#include +#include +#include "dali/core/format.h" +#include "dali/core/nvtx.h" +#include "dali/util/uri.h" + +namespace gcs = ::google::cloud::storage; + +namespace dali { + +namespace gcs_filesystem { + +namespace { + +std::string error_message(const google::cloud::Status& status) { + return make_string("[", google::cloud::StatusCodeToString(status.code()), "] ", + status.message()); +} + +} // namespace + +GCSObjectLocation parse_uri(const std::string& uri) { + auto parsed_uri = URI::Parse(uri, URI::ParseOpts::AllowNonEscaped); + if (parsed_uri.scheme() != "gs") + throw std::runtime_error("Not a GCS URI: " + uri); + GCSObjectLocation object_location; + object_location.bucket = parsed_uri.authority(); + object_location.object = parsed_uri.path(); + if (object_location.object.length() >= 1 && object_location.object[0] == '/') + object_location.object = object_location.object.substr(1); + return object_location; +} + +GCSObjectStats get_stats(gcs::Client& client, const GCSObjectLocation& object_location) { + DomainTimeRange tr(make_string("get_stats @ ", object_location.object), DomainTimeRange::kOrange); + GCSObjectStats stats; + if (object_location.object.empty()) + throw std::runtime_error("Object can't be empty"); + + auto metadata = client.GetObjectMetadata(object_location.bucket, object_location.object); + if (!metadata) { + throw std::runtime_error("GCS object not found. bucket=" + object_location.bucket + + " object=" + object_location.object + ":\n" + + error_message(metadata.status())); + } + stats.exists = true; + stats.size = metadata->size(); + return stats; +} + +size_t read_object_contents(gcs::Client& client, const GCSObjectLocation& object_location, + void* buf, size_t n, size_t offset) { + if (n == 0) + return 0; + // ReadRange is right-open ([begin, end)), unlike the HTTP "Range: bytes=first-last" header, + // which is inclusive on both ends. + auto begin = static_cast(offset); + auto end = static_cast(offset + n); + + DomainTimeRange tr(make_string("read_object_contents @ ", object_location.object, " [", begin, + ", ", end, ") (", n, ")"), + DomainTimeRange::kOrange); + + auto stream = client.ReadObject(object_location.bucket, object_location.object, + gcs::ReadRange(begin, end)); + // Unformatted I/O - the data lands directly in the caller's buffer, no intermediate copy. + stream.read(static_cast(buf), n); + // Reading fewer than n bytes (end of object) sets failbit/eofbit, which is not an error here, + // so the transfer status is the only thing worth checking. + auto bytes_read = static_cast(stream.gcount()); + stream.Close(); + if (!stream.status().ok()) { + throw std::runtime_error("Failed to read GCS object. bucket=" + object_location.bucket + + " object=" + object_location.object + ":\n" + + error_message(stream.status())); + } + return bytes_read; +} + +void list_objects_f(gcs::Client& client, const GCSObjectLocation& object_location, + PerObjectCallable per_object_call) { + DomainTimeRange tr(make_string("list_object_contents @ ", object_location.object), + DomainTimeRange::kOrange); + std::string prefix = object_location.object; + if (!prefix.empty() && prefix.back() != '/') { + prefix.push_back('/'); + } + // ListObjects returns a lazy range that pages through the results transparently. + for (auto& metadata : client.ListObjects(object_location.bucket, gcs::Prefix(prefix))) { + if (!metadata) { + throw std::runtime_error("Failed to list GCS objects. bucket=" + object_location.bucket + + " prefix=" + prefix + ":\n" + error_message(metadata.status())); + } + per_object_call(metadata->name(), metadata->size()); + } +} + +} // namespace gcs_filesystem + +} // namespace dali diff --git a/dali/util/gcs_filesystem.h b/dali/util/gcs_filesystem.h new file mode 100644 index 00000000000..ac4d08587f7 --- /dev/null +++ b/dali/util/gcs_filesystem.h @@ -0,0 +1,89 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_UTIL_GCS_FILESYSTEM_H_ +#define DALI_UTIL_GCS_FILESYSTEM_H_ + +#include +#include +#include +#include +#include "dali/core/api_helper.h" + +namespace dali { + +namespace gcs_filesystem { + +struct DLL_PUBLIC GCSObjectLocation { + std::string bucket; + std::string object; +}; + +struct DLL_PUBLIC GCSObjectStats { + bool exists = false; + size_t size = 0; +}; + +/** + * @brief Parses a GCS URI (gs://bucket/object) into an object location + * + * @param uri URI to the GCS prefix to query + * @return GCSObjectLocation object location + */ +DLL_PUBLIC GCSObjectLocation parse_uri(const std::string& uri); + +/** + * @brief Get the GCS object stats + * + * @param client GCS client. Note that `google::cloud::storage::Client` must not be used + * concurrently from multiple threads, but copies of it are cheap and share the + * underlying connection pool - see GCSClientManager. + * @param object_location GCS object location + * @return GCSObjectStats object stats + */ +DLL_PUBLIC GCSObjectStats get_stats(google::cloud::storage::Client& client, + const GCSObjectLocation& object_location); + +/** + * @brief Read GCS object contents + * + * @param client GCS client (see the note in get_stats) + * @param object_location object location + * @param buf preallocated buffer location + * @param n number of bytes to read + * @param offset (optional) offset to start reading from + * @return size_t number of bytes read + */ +DLL_PUBLIC size_t read_object_contents(google::cloud::storage::Client& client, + const GCSObjectLocation& object_location, void* buf, + size_t n, size_t offset = 0); + +using PerObjectCallable = std::function; + +/** + * @brief Visits all objects under a given object location + * + * @param client GCS client (see the note in get_stats) + * @param object_location GCS object location + * @param per_object_call callable to run on each object listed + */ +DLL_PUBLIC void list_objects_f(google::cloud::storage::Client& client, + const GCSObjectLocation& object_location, + PerObjectCallable per_object_call); + +} // namespace gcs_filesystem + +} // namespace dali + +#endif // DALI_UTIL_GCS_FILESYSTEM_H_