From ae3d3b8130860ce1eaa3a28d5bba32ea32fc57f3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 3 Apr 2026 14:10:30 -0700 Subject: [PATCH] Initial commit --- cpp/CMakeLists.txt | 5 +- .../kvikio/aws_credential_provider.hpp | 96 ++++ cpp/include/kvikio/remote_handle.hpp | 59 +-- cpp/src/aws_credential_provider.cpp | 471 ++++++++++++++++++ cpp/src/remote_handle.cpp | 120 +++-- docs/source/api.rst | 14 + docs/source/remote_file.rst | 10 + python/kvikio/kvikio/__init__.py | 14 + python/kvikio/kvikio/_lib/CMakeLists.txt | 7 + python/kvikio/kvikio/_lib/remote_handle.pyx | 128 +++-- python/kvikio/kvikio/aws_credentials.py | 171 +++++++ python/kvikio/kvikio/remote_file.py | 254 +++++++--- python/kvikio/tests/test_s3_io.py | 16 +- 13 files changed, 1144 insertions(+), 221 deletions(-) create mode 100644 cpp/include/kvikio/aws_credential_provider.hpp create mode 100644 cpp/src/aws_credential_provider.cpp create mode 100644 python/kvikio/kvikio/aws_credentials.py diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 0a8a5fcba7..23da2150dc 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -181,8 +181,9 @@ set(SOURCES ) if(KvikIO_REMOTE_SUPPORT) - list(APPEND SOURCES "src/hdfs.cpp" "src/remote_handle.cpp" "src/detail/remote_handle.cpp" - "src/detail/tls.cpp" "src/detail/url.cpp" "src/shim/libcurl.cpp" + list( + APPEND SOURCES "src/aws_credential_provider.cpp" "src/hdfs.cpp" "src/remote_handle.cpp" + "src/detail/remote_handle.cpp" "src/detail/tls.cpp" "src/detail/url.cpp" "src/shim/libcurl.cpp" ) endif() diff --git a/cpp/include/kvikio/aws_credential_provider.hpp b/cpp/include/kvikio/aws_credential_provider.hpp new file mode 100644 index 0000000000..86153ff328 --- /dev/null +++ b/cpp/include/kvikio/aws_credential_provider.hpp @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#ifndef KVIKIO_LIBCURL_FOUND +#error \ + "cannot include the remote IO API, please build KvikIO with libcurl (-DKvikIO_REMOTE_SUPPORT=ON)" +#endif + +#include +#include +#include + +struct curl_slist; + +namespace kvikio { + +/** + * @brief Immutable AWS SigV4 user/password and optional session-token header for libcurl. + * + * `token_header_list` must outlive `curl_easy_perform`; callers hold `shared_ptr` to this object + * until the transfer completes. + */ +class AwsAuthMaterial { + public: + std::string userpwd; + ::curl_slist* token_header_list{}; + + AwsAuthMaterial(); + ~AwsAuthMaterial(); + AwsAuthMaterial(AwsAuthMaterial const&) = delete; + AwsAuthMaterial& operator=(AwsAuthMaterial const&) = delete; + AwsAuthMaterial(AwsAuthMaterial&&) = delete; + AwsAuthMaterial& operator=(AwsAuthMaterial&&) = delete; + + static std::shared_ptr create(std::string access_key_id, + std::string secret_access_key, + std::optional session_token); +}; + +/** + * @brief How Python / Cython select the AWS credential source for S3. + */ +enum class AwsCredentialKind : std::uint8_t { + Default = 0, ///< Environment keys if set, else IAM role via metadata (IMDSv2) + Environment = 1, ///< `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` / optional token only + Static = 2, ///< Explicit access key, secret, optional session token + IamRole = 3, ///< IAM role credentials from the compute metadata service (IMDSv2) only + Legacy = 4, ///< Optional args plus environment (pre-credential S3 API semantics) +}; + +class AwsCredentialProvider { + public: + virtual ~AwsCredentialProvider() = default; + + /** + * @brief Return auth material for one HTTP request; implementations cache and refresh as needed. + */ + virtual std::shared_ptr get_auth_material() = 0; +}; + +/** + * @brief Build a credential provider for the given kind (used by Cython). + * + * @param kind Credential selection mode + * @param aws_access_key Required when kind == Static; optional when kind == Legacy (env fallback) + * @param aws_secret_access_key Required when kind == Static + * @param aws_session_token Optional; required when access key begins with "ASIA" (Static / Legacy) + * @param imds_endpoint_override Optional base URL (e.g. http://127.0.0.1:1234) for tests; if + * nullopt, uses `AWS_EC2_METADATA_SERVICE_ENDPOINT` or the default EC2 link-local address. + */ +std::shared_ptr make_aws_credential_provider( + AwsCredentialKind kind, + std::optional aws_access_key = std::nullopt, + std::optional aws_secret_access_key = std::nullopt, + std::optional aws_session_token = std::nullopt, + std::optional imds_endpoint_override = std::nullopt); + +/** + * @brief Provider matching legacy S3Endpoint optional arguments plus environment variables. + */ +std::shared_ptr make_legacy_env_and_args_credential_provider( + std::optional aws_access_key, + std::optional aws_secret_access_key, + std::optional aws_session_token); + +/** + * @brief Default chain: static env keys if both `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` + * are set and non-empty, otherwise IAM role credentials via the metadata service (IMDSv2). + */ +std::shared_ptr make_default_aws_credential_provider( + std::optional imds_endpoint_override = std::nullopt); + +} // namespace kvikio diff --git a/cpp/include/kvikio/remote_handle.hpp b/cpp/include/kvikio/remote_handle.hpp index 0b0808c45e..20c0c660d1 100644 --- a/cpp/include/kvikio/remote_handle.hpp +++ b/cpp/include/kvikio/remote_handle.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -16,11 +16,11 @@ #include #include -struct curl_slist; - namespace kvikio { class CurlHandle; // Prototype +class AwsAuthMaterial; +class AwsCredentialProvider; /** * @brief Types of remote file endpoints supported by KvikIO. @@ -135,14 +135,14 @@ class HttpEndpoint : public RemoteEndpoint { * @brief A remote endpoint for AWS S3 storage requiring credentials * * This endpoint is for accessing private S3 objects using AWS credentials (access key, secret key, - * region and optional session token). + * region and optional session token), optionally sourced from the environment, explicit parameters, + * or IAM role credentials via the compute metadata service (IMDSv2). */ class S3Endpoint : public RemoteEndpoint { private: std::string _url; std::string _aws_sigv4; - std::string _aws_userpwd; - curl_slist* _curl_header_list{}; + std::shared_ptr _credential_provider; public: /** @@ -176,19 +176,21 @@ class S3Endpoint : public RemoteEndpoint { [[nodiscard]] static std::pair parse_s3_url(std::string const& s3_url); /** - * @brief Create a S3 endpoint from a url. + * @brief Create a S3 endpoint from a url and credential provider. * * @param url The full http url to the S3 file. NB: this should be an url starting with * "http://" or "https://". If you have an S3 url of the form "s3:///", please * use `S3Endpoint::parse_s3_url()` and `S3Endpoint::url_from_bucket_and_object() to convert it. * @param aws_region The AWS region, such as "us-east-1", to use. If nullopt, the value of the * `AWS_DEFAULT_REGION` environment variable is used. - * @param aws_access_key The AWS access key to use. If nullopt, the value of the - * `AWS_ACCESS_KEY_ID` environment variable is used. - * @param aws_secret_access_key The AWS secret access key to use. If nullopt, the value of the - * `AWS_SECRET_ACCESS_KEY` environment variable is used. - * @param aws_session_token The AWS session token to use. If nullopt, the value of the - * `AWS_SESSION_TOKEN` environment variable is used. + * @param credential_provider Source for AWS access key, secret, and optional session token. + */ + S3Endpoint(std::string url, + std::optional aws_region, + std::shared_ptr credential_provider); + + /** + * @brief Create a S3 endpoint from a url (legacy optional arguments and environment variables). */ S3Endpoint(std::string url, std::optional aws_region = std::nullopt, @@ -197,21 +199,15 @@ class S3Endpoint : public RemoteEndpoint { std::optional aws_session_token = std::nullopt); /** - * @brief Create a S3 endpoint from a bucket and object name. - * - * @param bucket_and_object_names The bucket and object names of the S3 bucket. - * @param aws_region The AWS region, such as "us-east-1", to use. If nullopt, the value of the - * `AWS_DEFAULT_REGION` environment variable is used. - * @param aws_access_key The AWS access key to use. If nullopt, the value of the - * `AWS_ACCESS_KEY_ID` environment variable is used. - * @param aws_secret_access_key The AWS secret access key to use. If nullopt, the value of the - * `AWS_SECRET_ACCESS_KEY` environment variable is used. - * @param aws_endpoint_url Overwrite the endpoint url (including the protocol part) by using - * the scheme: "//". If nullopt, the value of the - * `AWS_ENDPOINT_URL` environment variable is used. If this is also not set, the regular AWS - * url scheme is used: "https://.s3..amazonaws.com/". - * @param aws_session_token The AWS session token to use. If nullopt, the value of the - * `AWS_SESSION_TOKEN` environment variable is used. + * @brief Create a S3 endpoint from a bucket and object name and credential provider. + */ + S3Endpoint(std::pair bucket_and_object_names, + std::optional aws_region, + std::optional aws_endpoint_url, + std::shared_ptr credential_provider); + + /** + * @brief Create a S3 endpoint from a bucket and object name (legacy optional arguments). */ S3Endpoint(std::pair bucket_and_object_names, std::optional aws_region = std::nullopt, @@ -222,6 +218,13 @@ class S3Endpoint : public RemoteEndpoint { ~S3Endpoint() override; void setopt(CurlHandle& curl) override; + /** + * @brief Apply SigV4 user/password and session token headers from `material` to `curl`. + * + * Call after `setopt()` on the same handle. Hold `material` alive until `curl.perform()` returns. + */ + void apply_auth_to_curl(CurlHandle& curl, AwsAuthMaterial const& material) const; + [[nodiscard]] std::shared_ptr get_auth_material(); std::string str() const override; std::size_t get_file_size() override; void setup_range_request(CurlHandle& curl, std::size_t file_offset, std::size_t size) override; diff --git a/cpp/src/aws_credential_provider.cpp b/cpp/src/aws_credential_provider.cpp new file mode 100644 index 0000000000..37501f9607 --- /dev/null +++ b/cpp/src/aws_credential_provider.cpp @@ -0,0 +1,471 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace kvikio { + +namespace { + +constexpr int kImdsConnectTimeoutSecs = 2; +constexpr int kImdsTotalTimeoutSecs = 5; +constexpr int kImdsTokenTtlSecs = 21'600; + +using clock = std::chrono::system_clock; + +std::string trim_in_place(std::string s) +{ + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) { + s.pop_back(); + } + return s; +} + +std::optional getenv_nonempty(char const* name) +{ + char const* v = std::getenv(name); + if (v == nullptr || v[0] == '\0') { return std::nullopt; } + return std::string{v}; +} + +std::string imds_base_url(std::optional override_url) +{ + if (override_url.has_value() && !override_url->empty()) { + std::string b = *override_url; + while (!b.empty() && b.back() == '/') { + b.pop_back(); + } + return b; + } + if (auto e = getenv_nonempty("AWS_EC2_METADATA_SERVICE_ENDPOINT")) { return *e; } + return std::string{"http://169.254.169.254"}; +} + +std::string join_path(std::string const& base, std::string_view rel) +{ + std::string out = base; + if (!out.empty() && out.back() == '/') { out.pop_back(); } + out.push_back('/'); + out.append(rel); + return out; +} + +clock::time_point parse_aws_expiration_iso8601(std::string const& iso) +{ + std::tm tm{}; + std::istringstream ss(iso); + ss >> std::get_time(&tm, "%Y-%m-%dT%H:%M:%SZ"); + if (ss.fail()) { + KVIKIO_FAIL("IMDS: cannot parse credential Expiration: " + iso, std::runtime_error); + } +#if defined(_WIN32) + std::time_t tt = _mkgmtime(&tm); +#else + std::time_t tt = timegm(&tm); +#endif + if (tt == static_cast(-1)) { + KVIKIO_FAIL("IMDS: timegm failed for Expiration: " + iso, std::runtime_error); + } + return clock::from_time_t(tt); +} + +std::optional json_extract_string(std::string_view json, std::string_view key) +{ + std::string needle = "\""; + needle.append(key); + needle.push_back('"'); + auto pos = json.find(needle); + if (pos == std::string_view::npos) { return std::nullopt; } + pos += needle.size(); + while (pos < json.size() && std::isspace(static_cast(json[pos]))) { + ++pos; + } + if (pos >= json.size() || json[pos] != ':') { return std::nullopt; } + ++pos; + while (pos < json.size() && std::isspace(static_cast(json[pos]))) { + ++pos; + } + if (pos >= json.size() || json[pos] != '"') { return std::nullopt; } + ++pos; + std::string out; + for (; pos < json.size(); ++pos) { + char c = json[pos]; + if (c == '"') { return out; } + if (c == '\\' && pos + 1 < json.size()) { + ++pos; + out.push_back(json[pos]); + continue; + } + out.push_back(c); + } + return std::nullopt; +} + +void imds_apply_timeouts(CurlHandle& curl) +{ + curl.setopt(CURLOPT_CONNECTTIMEOUT, static_cast(kImdsConnectTimeoutSecs)); + curl.setopt(CURLOPT_TIMEOUT, static_cast(kImdsTotalTimeoutSecs)); +} + +size_t empty_upload_read(char*, size_t, size_t, void*) { return 0; } + +void imds_put_token(CurlHandle& curl, std::string const& token_url, std::string& token_out) +{ + token_out.clear(); + std::string ttl = std::to_string(kImdsTokenTtlSecs); + curl_slist* hdrs = + curl_slist_append(nullptr, ("X-aws-ec2-metadata-token-ttl-seconds: " + ttl).c_str()); + if (hdrs == nullptr) { KVIKIO_FAIL("IMDS: curl_slist_append failed", std::runtime_error); } + curl.setopt(CURLOPT_URL, token_url.c_str()); + curl.setopt(CURLOPT_CUSTOMREQUEST, "PUT"); + curl.setopt(CURLOPT_UPLOAD, 1L); + curl.setopt(CURLOPT_READFUNCTION, empty_upload_read); + curl.setopt(CURLOPT_READDATA, nullptr); + curl.setopt(CURLOPT_INFILESIZE_LARGE, static_cast(0)); + curl.setopt(CURLOPT_HTTPHEADER, hdrs); + curl.setopt(CURLOPT_WRITEFUNCTION, detail::callback_get_string_response); + curl.setopt(CURLOPT_WRITEDATA, &token_out); + imds_apply_timeouts(curl); + curl.perform(); + curl_slist_free_all(hdrs); + curl.setopt(CURLOPT_UPLOAD, 0L); + curl.setopt(CURLOPT_READFUNCTION, nullptr); + curl.setopt(CURLOPT_READDATA, nullptr); + curl.setopt(CURLOPT_INFILESIZE_LARGE, static_cast(0)); + curl.setopt(CURLOPT_CUSTOMREQUEST, nullptr); + curl.setopt(CURLOPT_HTTPHEADER, nullptr); + token_out = trim_in_place(std::move(token_out)); + if (token_out.empty()) { + KVIKIO_FAIL("IMDS: empty metadata session token from " + token_url, std::runtime_error); + } +} + +void imds_get_with_token(CurlHandle& curl, + std::string const& url, + std::string const& metadata_token, + std::string& body_out) +{ + body_out.clear(); + std::string hdr_line = "X-aws-ec2-metadata-token: " + metadata_token; + curl_slist* hdrs = curl_slist_append(nullptr, hdr_line.c_str()); + if (hdrs == nullptr) { KVIKIO_FAIL("IMDS: curl_slist_append failed", std::runtime_error); } + curl.setopt(CURLOPT_HTTPGET, 1L); + curl.setopt(CURLOPT_URL, url.c_str()); + curl.setopt(CURLOPT_HTTPHEADER, hdrs); + curl.setopt(CURLOPT_WRITEFUNCTION, detail::callback_get_string_response); + curl.setopt(CURLOPT_WRITEDATA, &body_out); + imds_apply_timeouts(curl); + curl.perform(); + curl_slist_free_all(hdrs); + curl.setopt(CURLOPT_HTTPHEADER, nullptr); +} + +void fetch_imds_credentials(std::string const& base_url, + std::string& access_key_out, + std::string& secret_out, + std::string& session_token_out, + clock::time_point& expiration_out) +{ + KVIKIO_NVTX_FUNC_RANGE(); + auto curl = create_curl_handle(); + std::string meta_token; + imds_put_token(curl, join_path(base_url, "latest/api/token"), meta_token); + + std::string role_name_raw; + imds_get_with_token(curl, + join_path(base_url, "latest/meta-data/iam/security-credentials/"), + meta_token, + role_name_raw); + auto role_name = trim_in_place(std::move(role_name_raw)); + if (role_name.empty()) { + KVIKIO_FAIL( + "IMDS: no IAM role name at latest/meta-data/iam/security-credentials/ " + "(is an instance role attached to this host?)", + std::runtime_error); + } + + std::string cred_json; + imds_get_with_token(curl, + join_path(base_url, "latest/meta-data/iam/security-credentials/" + role_name), + meta_token, + cred_json); + + auto access = json_extract_string(cred_json, "AccessKeyId"); + auto secret = json_extract_string(cred_json, "SecretAccessKey"); + auto token = json_extract_string(cred_json, "Token"); + auto exp = json_extract_string(cred_json, "Expiration"); + if (!access.has_value() || !secret.has_value() || !token.has_value() || !exp.has_value()) { + KVIKIO_FAIL("IMDS: missing fields in security-credentials response", std::runtime_error); + } + access_key_out = std::move(*access); + secret_out = std::move(*secret); + session_token_out = std::move(*token); + expiration_out = parse_aws_expiration_iso8601(*exp); +} + +class StaticCredentialProvider : public AwsCredentialProvider { + std::mutex mutex_; + std::string access_; + std::string secret_; + std::optional session_; + std::shared_ptr cache_; + + public: + StaticCredentialProvider(std::string access, + std::string secret, + std::optional session) + : access_{std::move(access)}, secret_{std::move(secret)}, session_{std::move(session)} + { + } + + std::shared_ptr get_auth_material() override + { + std::lock_guard lock(mutex_); + if (cache_) { return cache_; } + if (access_.compare(0, 4, "ASIA") == 0) { + KVIKIO_EXPECT(session_.has_value() && !session_->empty(), + "Static AWS credentials: session token required when access key id begins " + "with ASIA", + std::invalid_argument); + } + cache_ = AwsAuthMaterial::create(access_, secret_, session_); + return cache_; + } +}; + +class EnvironmentCredentialProvider : public AwsCredentialProvider { + std::mutex mutex_; + std::shared_ptr cache_; + + public: + std::shared_ptr get_auth_material() override + { + std::lock_guard lock(mutex_); + if (cache_) { return cache_; } + auto access = detail::unwrap_or_env(std::nullopt, + "AWS_ACCESS_KEY_ID", + "S3: must provide `aws_access_key` if AWS_ACCESS_KEY_ID " + "isn't set."); + auto secret = detail::unwrap_or_env(std::nullopt, + "AWS_SECRET_ACCESS_KEY", + "S3: must provide `aws_secret_access_key` if " + "AWS_SECRET_ACCESS_KEY isn't set."); + std::optional session = std::nullopt; + if (access->compare(0, 4, std::string("ASIA")) == 0) { + session = detail::unwrap_or_env(std::nullopt, + "AWS_SESSION_TOKEN", + "When using temporary credentials, AWS_SESSION_TOKEN must " + "be set."); + } + cache_ = AwsAuthMaterial::create(*access, *secret, session); + return cache_; + } +}; + +class LegacyEnvAndArgsCredentialProvider : public AwsCredentialProvider { + std::mutex mutex_; + std::optional opt_access_; + std::optional opt_secret_; + std::optional opt_session_; + std::shared_ptr cache_; + + public: + LegacyEnvAndArgsCredentialProvider(std::optional aws_access_key, + std::optional aws_secret_access_key, + std::optional aws_session_token) + : opt_access_{std::move(aws_access_key)}, + opt_secret_{std::move(aws_secret_access_key)}, + opt_session_{std::move(aws_session_token)} + { + } + + std::shared_ptr get_auth_material() override + { + std::lock_guard lock(mutex_); + if (cache_) { return cache_; } + auto access = detail::unwrap_or_env(std::move(opt_access_), + "AWS_ACCESS_KEY_ID", + "S3: must provide `aws_access_key` if AWS_ACCESS_KEY_ID " + "isn't set."); + auto secret = detail::unwrap_or_env(std::move(opt_secret_), + "AWS_SECRET_ACCESS_KEY", + "S3: must provide `aws_secret_access_key` if " + "AWS_SECRET_ACCESS_KEY isn't set."); + std::optional session = std::nullopt; + if (access->compare(0, 4, std::string("ASIA")) == 0) { + session = detail::unwrap_or_env(std::move(opt_session_), + "AWS_SESSION_TOKEN", + "When using temporary credentials, AWS_SESSION_TOKEN must " + "be set."); + } + cache_ = AwsAuthMaterial::create(*access, *secret, session); + return cache_; + } +}; + +class IamRoleCredentialProvider : public AwsCredentialProvider { + std::mutex mutex_; + std::string base_url_; + std::shared_ptr material_; + clock::time_point refresh_after_{}; + + public: + explicit IamRoleCredentialProvider(std::optional endpoint_override) + : base_url_{imds_base_url(std::move(endpoint_override))} + { + } + + std::shared_ptr get_auth_material() override + { + std::lock_guard lock(mutex_); + auto const now = clock::now(); + if (material_ && now < refresh_after_) { return material_; } + + std::string ak; + std::string sk; + std::string tok; + clock::time_point expiration{}; + fetch_imds_credentials(base_url_, ak, sk, tok, expiration); + + constexpr auto skew = std::chrono::minutes{5}; + auto next_refresh = expiration - skew; + if (next_refresh <= now) { next_refresh = now + std::chrono::seconds{30}; } + refresh_after_ = next_refresh; + material_ = AwsAuthMaterial::create(std::move(ak), std::move(sk), std::move(tok)); + return material_; + } +}; + +class DefaultCredentialProvider : public AwsCredentialProvider { + std::mutex mutex_; + std::optional imds_override_; + std::shared_ptr inner_; + + public: + explicit DefaultCredentialProvider(std::optional imds_override) + : imds_override_{std::move(imds_override)} + { + } + + std::shared_ptr get_auth_material() override + { + std::lock_guard lock(mutex_); + if (!inner_) { + auto id = getenv_nonempty("AWS_ACCESS_KEY_ID"); + auto key = getenv_nonempty("AWS_SECRET_ACCESS_KEY"); + if (id.has_value() && key.has_value()) { + inner_ = std::make_shared(); + } else { + inner_ = std::make_shared(imds_override_); + } + } + return inner_->get_auth_material(); + } +}; + +} // namespace + +AwsAuthMaterial::AwsAuthMaterial() = default; + +AwsAuthMaterial::~AwsAuthMaterial() +{ + curl_slist_free_all(token_header_list); + token_header_list = nullptr; +} + +std::shared_ptr AwsAuthMaterial::create( + std::string access_key_id, + std::string secret_access_key, + std::optional session_token) +{ + bool const is_asia = access_key_id.size() >= 4 && access_key_id.compare(0, 4, "ASIA") == 0; + auto m = std::shared_ptr(new AwsAuthMaterial()); + m->userpwd = std::move(access_key_id); + m->userpwd.push_back(':'); + m->userpwd.append(secret_access_key); + + if (session_token.has_value() && !session_token->empty()) { + std::string line = "x-amz-security-token: "; + line += *session_token; + m->token_header_list = curl_slist_append(nullptr, line.c_str()); + KVIKIO_EXPECT(m->token_header_list != nullptr, + "Failed to create curl header for AWS token", + std::runtime_error); + } else { + KVIKIO_EXPECT(!is_asia, + "AWS session token required for temporary access key ids (ASIA...)", + std::invalid_argument); + } + return m; +} + +std::shared_ptr make_legacy_env_and_args_credential_provider( + std::optional aws_access_key, + std::optional aws_secret_access_key, + std::optional aws_session_token) +{ + return std::make_shared( + std::move(aws_access_key), std::move(aws_secret_access_key), std::move(aws_session_token)); +} + +std::shared_ptr make_default_aws_credential_provider( + std::optional imds_endpoint_override) +{ + return std::make_shared(std::move(imds_endpoint_override)); +} + +std::shared_ptr make_aws_credential_provider( + AwsCredentialKind kind, + std::optional aws_access_key, + std::optional aws_secret_access_key, + std::optional aws_session_token, + std::optional imds_endpoint_override) +{ + switch (kind) { + case AwsCredentialKind::Default: + return make_default_aws_credential_provider(std::move(imds_endpoint_override)); + case AwsCredentialKind::Environment: return std::make_shared(); + case AwsCredentialKind::Static: + KVIKIO_EXPECT(aws_access_key.has_value() && !aws_access_key->empty(), + "Static AWS credentials require a non-empty access key id", + std::invalid_argument); + KVIKIO_EXPECT(aws_secret_access_key.has_value() && !aws_secret_access_key->empty(), + "Static AWS credentials require a non-empty secret access key", + std::invalid_argument); + return std::make_shared(std::move(*aws_access_key), + std::move(*aws_secret_access_key), + std::move(aws_session_token)); + case AwsCredentialKind::IamRole: + return std::make_shared(std::move(imds_endpoint_override)); + case AwsCredentialKind::Legacy: + return make_legacy_env_and_args_credential_provider( + std::move(aws_access_key), std::move(aws_secret_access_key), std::move(aws_session_token)); + default: KVIKIO_FAIL("Unknown AwsCredentialKind", std::invalid_argument); + } +} + +} // namespace kvikio diff --git a/cpp/src/remote_handle.cpp b/cpp/src/remote_handle.cpp index 44c86063b5..badb5f3149 100644 --- a/cpp/src/remote_handle.cpp +++ b/cpp/src/remote_handle.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -145,7 +146,11 @@ std::size_t get_file_size_using_head_impl(RemoteEndpoint& endpoint, std::string { auto curl = create_curl_handle(); + std::shared_ptr auth_keepalive; + auto* s3 = dynamic_cast(&endpoint); + if (s3 != nullptr) { auth_keepalive = s3->get_auth_material(); } endpoint.setopt(curl); + if (s3 != nullptr) { s3->apply_auth_to_curl(curl, *auth_keepalive); } curl.setopt(CURLOPT_NOBODY, 1L); curl.setopt(CURLOPT_FOLLOWLOCATION, 1L); curl.perform(); @@ -276,8 +281,19 @@ void S3Endpoint::setopt(CurlHandle& curl) curl.setopt(CURLOPT_URL, new_url.c_str()); curl.setopt(CURLOPT_AWS_SIGV4, _aws_sigv4.c_str()); - curl.setopt(CURLOPT_USERPWD, _aws_userpwd.c_str()); - if (_curl_header_list) { curl.setopt(CURLOPT_HTTPHEADER, _curl_header_list); } +} + +void S3Endpoint::apply_auth_to_curl(CurlHandle& curl, AwsAuthMaterial const& material) const +{ + curl.setopt(CURLOPT_USERPWD, material.userpwd.c_str()); + if (material.token_header_list != nullptr) { + curl.setopt(CURLOPT_HTTPHEADER, material.token_header_list); + } +} + +std::shared_ptr S3Endpoint::get_auth_material() +{ + return _credential_provider->get_auth_material(); } std::string S3Endpoint::url_from_bucket_and_object(std::string bucket_name, @@ -315,12 +331,15 @@ std::pair S3Endpoint::parse_s3_url(std::string const& S3Endpoint::S3Endpoint(std::string url, std::optional aws_region, - std::optional aws_access_key, - std::optional aws_secret_access_key, - std::optional aws_session_token) - : RemoteEndpoint{RemoteEndpointType::S3}, _url{std::move(url)} + std::shared_ptr credential_provider) + : RemoteEndpoint{RemoteEndpointType::S3}, + _url{std::move(url)}, + _credential_provider{std::move(credential_provider)} { KVIKIO_NVTX_FUNC_RANGE(); + KVIKIO_EXPECT(_credential_provider != nullptr, + "S3 credential provider must not be null", + std::invalid_argument); // Regular expression to match http[s]:// std::regex static const pattern{R"(^https?://.*)", std::regex_constants::icase}; KVIKIO_EXPECT(std::regex_search(_url, pattern), @@ -332,69 +351,60 @@ S3Endpoint::S3Endpoint(std::string url, "AWS_DEFAULT_REGION", "S3: must provide `aws_region` if AWS_DEFAULT_REGION isn't set."); - auto const access_key = - detail::unwrap_or_env(std::move(aws_access_key), - "AWS_ACCESS_KEY_ID", - "S3: must provide `aws_access_key` if AWS_ACCESS_KEY_ID isn't set."); - - auto const secret_access_key = detail::unwrap_or_env( - std::move(aws_secret_access_key), - "AWS_SECRET_ACCESS_KEY", - "S3: must provide `aws_secret_access_key` if AWS_SECRET_ACCESS_KEY isn't set."); - - // Create the CURLOPT_AWS_SIGV4 option { std::stringstream ss; ss << "aws:amz:" << region.value() << ":s3"; _aws_sigv4 = ss.str(); } - // Create the CURLOPT_USERPWD option - // Notice, curl uses `secret_access_key` to generate a AWS V4 signature. It is NOT included - // in the http header. See - // - { - std::stringstream ss; - ss << access_key.value() << ":" << secret_access_key.value(); - _aws_userpwd = ss.str(); - } - // Access key IDs beginning with ASIA are temporary credentials that are created using AWS STS - // operations. They need a session token to work. - if (access_key->compare(0, 4, std::string("ASIA")) == 0) { - // Create a Custom Curl header for the session token. - // The _curl_header_list created by curl_slist_append must be manually freed - // (see https://curl.se/libcurl/c/CURLOPT_HTTPHEADER.html) - auto session_token = - detail::unwrap_or_env(std::move(aws_session_token), - "AWS_SESSION_TOKEN", - "When using temporary credentials, AWS_SESSION_TOKEN must be set."); - std::stringstream ss; - ss << "x-amz-security-token: " << session_token.value(); - _curl_header_list = curl_slist_append(NULL, ss.str().c_str()); - KVIKIO_EXPECT(_curl_header_list != nullptr, - "Failed to create curl header for AWS token", - std::runtime_error); - } } -S3Endpoint::S3Endpoint(std::pair bucket_and_object_names, +S3Endpoint::S3Endpoint(std::string url, std::optional aws_region, std::optional aws_access_key, std::optional aws_secret_access_key, - std::optional aws_endpoint_url, std::optional aws_session_token) + : S3Endpoint( + std::move(url), + aws_region, + make_legacy_env_and_args_credential_provider( + std::move(aws_access_key), std::move(aws_secret_access_key), std::move(aws_session_token))) +{ + KVIKIO_NVTX_FUNC_RANGE(); +} + +S3Endpoint::S3Endpoint(std::pair bucket_and_object_names, + std::optional aws_region, + std::optional aws_endpoint_url, + std::shared_ptr credential_provider) : S3Endpoint(url_from_bucket_and_object(std::move(bucket_and_object_names.first), std::move(bucket_and_object_names.second), aws_region, std::move(aws_endpoint_url)), aws_region, - std::move(aws_access_key), - std::move(aws_secret_access_key), - std::move(aws_session_token)) + std::move(credential_provider)) +{ + KVIKIO_NVTX_FUNC_RANGE(); +} + +S3Endpoint::S3Endpoint(std::pair bucket_and_object_names, + std::optional aws_region, + std::optional aws_access_key, + std::optional aws_secret_access_key, + std::optional aws_endpoint_url, + std::optional aws_session_token) + : S3Endpoint( + url_from_bucket_and_object(std::move(bucket_and_object_names.first), + std::move(bucket_and_object_names.second), + aws_region, + std::move(aws_endpoint_url)), + aws_region, + make_legacy_env_and_args_credential_provider( + std::move(aws_access_key), std::move(aws_secret_access_key), std::move(aws_session_token))) { KVIKIO_NVTX_FUNC_RANGE(); } -S3Endpoint::~S3Endpoint() { curl_slist_free_all(_curl_header_list); } +S3Endpoint::~S3Endpoint() = default; std::string S3Endpoint::str() const { return _url; } @@ -595,9 +605,13 @@ RemoteHandle RemoteHandle::open(std::string url, if (!S3Endpoint::is_url_valid(url)) { return nullptr; } if (scheme.value() == "s3") { auto const [bucket, object] = S3Endpoint::parse_s3_url(url); - return std::make_unique(std::pair{bucket, object}); + return std::make_unique(std::pair{bucket, object}, + std::nullopt, + std::nullopt, + make_default_aws_credential_provider()); } - return std::make_unique(url); + return std::make_unique( + url, std::nullopt, make_default_aws_credential_provider()); case RemoteEndpointType::S3_PUBLIC: if (!S3PublicEndpoint::is_url_valid(url)) { return nullptr; } @@ -771,7 +785,11 @@ std::size_t RemoteHandle::read(void* buf, std::size_t size, std::size_t file_off } bool const is_host_mem = is_host_memory(buf); auto curl = create_curl_handle(); + std::shared_ptr auth_keepalive; + auto* s3 = dynamic_cast(_endpoint.get()); + if (s3 != nullptr) { auth_keepalive = s3->get_auth_material(); } _endpoint->setopt(curl); + if (s3 != nullptr) { s3->apply_auth_to_curl(curl, *auth_keepalive); } _endpoint->setup_range_request(curl, file_offset, size); if (is_host_mem) { diff --git a/docs/source/api.rst b/docs/source/api.rst index d7100afcd1..931a982587 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -68,6 +68,20 @@ RemoteFile .. autoclass:: RemoteFile :members: +AWS S3 credentials +------------------ +.. currentmodule:: kvikio.aws_credentials + +.. autoclass:: AwsDefaultCredential + +.. autoclass:: AwsEnvironmentCredential + +.. autoclass:: AwsLegacyCredential + +.. autoclass:: AwsStaticCredential + +.. autoclass:: AwsIamRoleCredential + Defaults -------- .. currentmodule:: kvikio.defaults diff --git a/docs/source/remote_file.rst b/docs/source/remote_file.rst index e6d038035e..7209fe087d 100644 --- a/docs/source/remote_file.rst +++ b/docs/source/remote_file.rst @@ -3,6 +3,16 @@ Remote File KvikIO provides direct access to remote files, including AWS S3, WebHDFS, and generic HTTP/HTTPS. +AWS S3 credentials +------------------ + +For :meth:`kvikio.RemoteFile.open_s3` and :meth:`kvikio.RemoteFile.open_s3_url`, pass a +``credential`` object from :mod:`kvikio.aws_credentials`. The default (``credential=None``) +uses environment variables when ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY`` are +set; otherwise KvikIO attempts IAM role credentials from the compute metadata service (`IMDSv2 +`_), +with caching in the C++ layer until credentials are close to expiry. + Example ------- diff --git a/python/kvikio/kvikio/__init__.py b/python/kvikio/kvikio/__init__.py index 24ab1fbe6f..96c4297153 100644 --- a/python/kvikio/kvikio/__init__.py +++ b/python/kvikio/kvikio/__init__.py @@ -14,6 +14,14 @@ from kvikio._lib.defaults import CompatMode # noqa: F401 from kvikio._version import __git_commit__, __version__ +from kvikio.aws_credentials import ( + AwsCredential, + AwsDefaultCredential, + AwsEnvironmentCredential, + AwsIamRoleCredential, + AwsLegacyCredential, + AwsStaticCredential, +) from kvikio.buffer import bounce_buffer_free, memory_deregister, memory_register from kvikio.cufile import ( CuFile, @@ -30,6 +38,12 @@ __all__ = [ "__git_commit__", "__version__", + "AwsCredential", + "AwsDefaultCredential", + "AwsIamRoleCredential", + "AwsEnvironmentCredential", + "AwsLegacyCredential", + "AwsStaticCredential", "clear_page_cache", "CuFile", "drop_file_page_cache", diff --git a/python/kvikio/kvikio/_lib/CMakeLists.txt b/python/kvikio/kvikio/_lib/CMakeLists.txt index d98cf5b047..445e767373 100644 --- a/python/kvikio/kvikio/_lib/CMakeLists.txt +++ b/python/kvikio/kvikio/_lib/CMakeLists.txt @@ -25,3 +25,10 @@ rapids_cython_create_modules( SOURCE_FILES "${cython_modules}" LINKED_LIBRARIES kvikio::kvikio ) + +if(KvikIO_REMOTE_SUPPORT AND TARGET remote_handle) + get_filename_component( + KVIKIO_CPP_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../cpp/include" ABSOLUTE + ) + target_include_directories(remote_handle PRIVATE "${KVIKIO_CPP_INCLUDE_DIR}") +endif() diff --git a/python/kvikio/kvikio/_lib/remote_handle.pyx b/python/kvikio/kvikio/_lib/remote_handle.pyx index 2f7031a7c3..1b32fed616 100644 --- a/python/kvikio/kvikio/_lib/remote_handle.pyx +++ b/python/kvikio/kvikio/_lib/remote_handle.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # distutils: language = c++ @@ -8,7 +8,7 @@ from typing import Optional from cython.operator cimport dereference as deref from libc.stdint cimport uint8_t, uintptr_t -from libcpp.memory cimport make_unique, unique_ptr +from libcpp.memory cimport make_unique, shared_ptr, unique_ptr from libcpp.optional cimport nullopt, optional from libcpp.pair cimport pair from libcpp.string cimport string @@ -19,6 +19,27 @@ from kvikio._lib.arr cimport parse_buffer_argument from kvikio._lib.future cimport IOFuture, _wrap_io_future, future +cdef extern from "" namespace "kvikio" nogil: + cpdef enum class AwsCredentialKind(uint8_t): + Default = 0 + Environment = 1 + Static = 2 + IamRole = 3 + Legacy = 4 + + cdef cppclass cpp_AwsCredentialProvider "kvikio::AwsCredentialProvider": + pass + + cdef shared_ptr[cpp_AwsCredentialProvider] cpp_make_aws_credential_provider \ + "kvikio::make_aws_credential_provider"( + AwsCredentialKind kind, + optional[string] aws_access_key, + optional[string] aws_secret_access_key, + optional[string] aws_session_token, + optional[string] imds_endpoint_override, + ) except + + + cdef extern from "" namespace "kvikio" nogil: cpdef enum class RemoteEndpointType(uint8_t): AUTO = 0 @@ -37,17 +58,13 @@ cdef extern from "" namespace "kvikio" nogil: cpp_S3Endpoint( string url, optional[string] aws_region, - optional[string] aws_access_key, - optional[string] aws_secret_access_key, - optional[string] aws_session_token + shared_ptr[cpp_AwsCredentialProvider] credential_provider, ) except + cpp_S3Endpoint( pair[string, string] bucket_and_object_names, optional[string] aws_region, - optional[string] aws_access_key, - optional[string] aws_secret_access_key, optional[string] aws_endpoint_url, - optional[string] aws_session_token + shared_ptr[cpp_AwsCredentialProvider] credential_provider, ) except + pair[string, string] cpp_parse_s3_url \ @@ -189,40 +206,40 @@ cdef class RemoteFile: def open_s3( bucket_name: str, object_name: str, - nbytes: Optional[int], + *, + nbytes: Optional[int] = None, aws_region_name: Optional[str] = None, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, aws_endpoint_url: Optional[str] = None, + uint8_t credential_kind=0, + aws_access_key: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, aws_session_token: Optional[str] = None, + imds_endpoint: Optional[str] = None, ): cdef pair[string, string] bucket_and_object_names = _to_string_pair( bucket_name, object_name ) cdef optional[string] cpp_aws_region = _to_optional_string(aws_region_name) - cdef optional[string] cpp_aws_access_key = _to_optional_string( - aws_access_key_id - ) - cdef optional[string] cpp_aws_secret_access_key = ( - _to_optional_string(aws_secret_access_key) - ) cdef optional[string] cpp_aws_endpoint_url = _to_optional_string( aws_endpoint_url ) - cdef optional[string] cpp_aws_session_token = _to_optional_string( - aws_session_token - ) cdef unique_ptr[cpp_RemoteEndpoint] cpp_endpoint - + cdef shared_ptr[cpp_AwsCredentialProvider] cpp_provider + + cpp_provider = cpp_make_aws_credential_provider( + credential_kind, + _to_optional_string(aws_access_key), + _to_optional_string(aws_secret_access_key), + _to_optional_string(aws_session_token), + _to_optional_string(imds_endpoint), + ) with nogil: cpp_endpoint = cast_to_remote_endpoint( make_unique[cpp_S3Endpoint]( bucket_and_object_names, cpp_aws_region, - cpp_aws_access_key, - cpp_aws_secret_access_key, cpp_aws_endpoint_url, - cpp_aws_session_token + cpp_provider, ) ) @@ -234,33 +251,33 @@ cdef class RemoteFile: @staticmethod def open_s3_from_http_url( url: str, - nbytes: Optional[int], + *, + nbytes: Optional[int] = None, aws_region_name: Optional[str] = None, - aws_access_key_id: Optional[str] = None, + uint8_t credential_kind=0, + aws_access_key: Optional[str] = None, aws_secret_access_key: Optional[str] = None, aws_session_token: Optional[str] = None, + imds_endpoint: Optional[str] = None, ): cdef string cpp_url = _to_string(url) cdef optional[string] cpp_aws_region = _to_optional_string(aws_region_name) - cdef optional[string] cpp_aws_access_key = _to_optional_string( - aws_access_key_id - ) - cdef optional[string] cpp_aws_secret_access_key = ( - _to_optional_string(aws_secret_access_key) - ) - cdef optional[string] cpp_aws_session_token = _to_optional_string( - aws_session_token - ) cdef unique_ptr[cpp_RemoteEndpoint] cpp_endpoint - + cdef shared_ptr[cpp_AwsCredentialProvider] cpp_provider + + cpp_provider = cpp_make_aws_credential_provider( + credential_kind, + _to_optional_string(aws_access_key), + _to_optional_string(aws_secret_access_key), + _to_optional_string(aws_session_token), + _to_optional_string(imds_endpoint), + ) with nogil: cpp_endpoint = cast_to_remote_endpoint( make_unique[cpp_S3Endpoint]( cpp_url, cpp_aws_region, - cpp_aws_access_key, - cpp_aws_secret_access_key, - cpp_aws_session_token + cpp_provider, ) ) @@ -272,40 +289,41 @@ cdef class RemoteFile: @staticmethod def open_s3_from_s3_url( url: str, - nbytes: Optional[int], + *, + nbytes: Optional[int] = None, aws_region_name: Optional[str] = None, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, aws_endpoint_url: Optional[str] = None, + uint8_t credential_kind=0, + aws_access_key: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, aws_session_token: Optional[str] = None, + imds_endpoint: Optional[str] = None, ): cdef string cpp_url = _to_string(url) cdef pair[string, string] bucket_and_object_names cdef optional[string] cpp_aws_region = _to_optional_string(aws_region_name) - cdef optional[string] cpp_aws_access_key = _to_optional_string( - aws_access_key_id - ) - cdef optional[string] cpp_aws_secret_access_key = ( - _to_optional_string(aws_secret_access_key) - ) cdef optional[string] cpp_aws_endpoint_url = _to_optional_string( aws_endpoint_url ) - cdef optional[string] cpp_aws_session_token = _to_optional_string( - aws_session_token - ) cdef unique_ptr[cpp_RemoteEndpoint] cpp_endpoint + cdef shared_ptr[cpp_AwsCredentialProvider] cpp_provider + bucket_and_object_names = cpp_parse_s3_url(cpp_url) + + cpp_provider = cpp_make_aws_credential_provider( + credential_kind, + _to_optional_string(aws_access_key), + _to_optional_string(aws_secret_access_key), + _to_optional_string(aws_session_token), + _to_optional_string(imds_endpoint), + ) with nogil: - bucket_and_object_names = cpp_parse_s3_url(cpp_url) cpp_endpoint = cast_to_remote_endpoint( make_unique[cpp_S3Endpoint]( bucket_and_object_names, cpp_aws_region, - cpp_aws_access_key, - cpp_aws_secret_access_key, cpp_aws_endpoint_url, - cpp_aws_session_token + cpp_provider, ) ) diff --git a/python/kvikio/kvikio/aws_credentials.py b/python/kvikio/kvikio/aws_credentials.py new file mode 100644 index 0000000000..d23b922db6 --- /dev/null +++ b/python/kvikio/kvikio/aws_credentials.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AWS credential sources for :meth:`kvikio.RemoteFile.open_s3` and related APIs. + +These mirror the idea of a *credential* object (similar in spirit to Azure's +``DefaultAzureCredential`` and related types in the `Azure SDK for Python +`_), +but for AWS S3 access inside KvikIO's native code path. + +KvikIO resolves credentials in C++ (with caching for IAM role credentials fetched via IMDSv2). +""" + +from __future__ import annotations + +from typing import Optional, Union + +# Match ``kvikio::AwsCredentialKind`` (``remote_handle`` / C++). +_CRED_DEFAULT: int = 0 +_CRED_ENVIRONMENT: int = 1 +_CRED_STATIC: int = 2 +_CRED_IAM_ROLE: int = 3 +_CRED_LEGACY: int = 4 + + +class AwsDefaultCredential: + """Use environment variables if set, otherwise an IAM role via the metadata service (IMDSv2). + + If both ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY`` are set and non-empty, + those are used. Otherwise KvikIO fetches temporary credentials for the IAM role + attached to the compute environment (e.g. EC2 instance, Lambda, ECS task) from the + metadata endpoint (see `IAM roles for EC2 + `_ + and related services). + + Attributes + ---------- + imds_endpoint : str or None + Optional metadata service base URL (e.g. ``http://127.0.0.1:1234``) for tests. + If ``None``, uses ``AWS_EC2_METADATA_SERVICE_ENDPOINT`` when set, else the default + link-local address. + """ + + __slots__ = ("imds_endpoint",) + + def __init__(self, imds_endpoint: Optional[str] = None) -> None: + self.imds_endpoint = imds_endpoint + + def _kvikio_kind(self) -> int: + return _CRED_DEFAULT + + def _kvikio_imds_endpoint(self) -> Optional[str]: + return self.imds_endpoint + + +class AwsEnvironmentCredential: + """Read credentials only from the environment (no IAM-role / metadata fallback). + + Uses ``AWS_ACCESS_KEY_ID``, ``AWS_SECRET_ACCESS_KEY``, and optionally + ``AWS_SESSION_TOKEN`` (required for temporary ``ASIA`` access keys). + """ + + __slots__ = () + + def _kvikio_kind(self) -> int: + return _CRED_ENVIRONMENT + + def _kvikio_imds_endpoint(self) -> Optional[str]: + return None + + +class AwsLegacyCredential: + """Same credential resolution as the pre-``credential=`` S3 API (optional args + env). + + Each field may be ``None``; missing values are taken from ``AWS_ACCESS_KEY_ID``, + ``AWS_SECRET_ACCESS_KEY``, and ``AWS_SESSION_TOKEN`` when applicable. This matches the + behavior of the deprecated ``aws_access_key_id=`` / ``aws_secret_access_key=`` / + ``aws_session_token=`` keyword arguments on :meth:`kvikio.RemoteFile.open_s3`. + """ + + __slots__ = ("aws_access_key_id", "aws_secret_access_key", "aws_session_token") + + def __init__( + self, + aws_access_key_id: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, + aws_session_token: Optional[str] = None, + ) -> None: + self.aws_access_key_id = aws_access_key_id + self.aws_secret_access_key = aws_secret_access_key + self.aws_session_token = aws_session_token + + def _kvikio_kind(self) -> int: + return _CRED_LEGACY + + def _kvikio_imds_endpoint(self) -> Optional[str]: + return None + + +class AwsStaticCredential: + """Fixed access key, secret key, and optional session token (no environment lookup). + + Attributes + ---------- + access_key_id : str + secret_access_key : str + session_token : str or None + Required when ``access_key_id`` begins with ``ASIA`` (temporary credentials). + """ + + __slots__ = ("access_key_id", "secret_access_key", "session_token") + + def __init__( + self, + access_key_id: str, + secret_access_key: str, + session_token: Optional[str] = None, + ) -> None: + self.access_key_id = access_key_id + self.secret_access_key = secret_access_key + self.session_token = session_token + + def _kvikio_kind(self) -> int: + return _CRED_STATIC + + def _kvikio_imds_endpoint(self) -> Optional[str]: + return None + + +class AwsIamRoleCredential: + """IAM role credentials from the compute metadata service (IMDSv2) only. + + Ignores static ``AWS_ACCESS_KEY_ID`` / ``AWS_SECRET_ACCESS_KEY`` environment + variables and always uses the role credentials exposed at the metadata endpoint + (used on EC2, Lambda, ECS, and other AWS compute). + + Attributes + ---------- + imds_endpoint : str or None + Optional metadata service base URL; same semantics as + :attr:`AwsDefaultCredential.imds_endpoint`. + """ + + __slots__ = ("imds_endpoint",) + + def __init__(self, imds_endpoint: Optional[str] = None) -> None: + self.imds_endpoint = imds_endpoint + + def _kvikio_kind(self) -> int: + return _CRED_IAM_ROLE + + def _kvikio_imds_endpoint(self) -> Optional[str]: + return self.imds_endpoint + + +AwsCredential = Union[ + AwsDefaultCredential, + AwsEnvironmentCredential, + AwsLegacyCredential, + AwsStaticCredential, + AwsIamRoleCredential, +] + +__all__ = [ + "AwsCredential", + "AwsDefaultCredential", + "AwsIamRoleCredential", + "AwsEnvironmentCredential", + "AwsLegacyCredential", + "AwsStaticCredential", +] diff --git a/python/kvikio/kvikio/remote_file.py b/python/kvikio/kvikio/remote_file.py index 1faf010c58..0b42b609cf 100644 --- a/python/kvikio/kvikio/remote_file.py +++ b/python/kvikio/kvikio/remote_file.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -6,10 +6,108 @@ import enum import functools import urllib.parse -from typing import Optional - +import warnings +from typing import Any, Optional + +from kvikio.aws_credentials import ( + AwsCredential, + AwsDefaultCredential, + AwsEnvironmentCredential, + AwsIamRoleCredential, + AwsLegacyCredential, + AwsStaticCredential, +) from kvikio.cufile import IOFuture +_LEGACY_S3_CREDENTIAL_KWARGS = ( + "Passing aws_access_key_id, aws_secret_access_key, or aws_session_token is deprecated; " + "use `credential=kvikio.aws_credentials.AwsStaticCredential(...)`, " + "`AwsEnvironmentCredential()`, `AwsDefaultCredential()`, or " + "`AwsLegacyCredential(...)` instead. " + "See the KvikIO documentation for S3 remote files." +) + + +def _s3_credential_bridge_kwargs( + credential: Optional[AwsCredential], + *, + aws_access_key_id: Optional[str], + aws_secret_access_key: Optional[str], + aws_session_token: Optional[str], +) -> dict[str, Any]: + """Map high-level `credential` / legacy kwargs to the Cython S3 open parameters.""" + legacy = ( + aws_access_key_id is not None + or aws_secret_access_key is not None + or aws_session_token is not None + ) + if legacy: + warnings.warn(_LEGACY_S3_CREDENTIAL_KWARGS, FutureWarning, stacklevel=3) + if credential is not None: + raise TypeError( + "Pass either `credential` or legacy AWS access-key keyword arguments, not both." + ) + return { + "credential_kind": 4, + "aws_access_key": aws_access_key_id, + "aws_secret_access_key": aws_secret_access_key, + "aws_session_token": aws_session_token, + "imds_endpoint": None, + } + + if credential is None: + return { + "credential_kind": 0, + "aws_access_key": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "imds_endpoint": None, + } + + if isinstance(credential, AwsDefaultCredential): + return { + "credential_kind": 0, + "aws_access_key": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "imds_endpoint": credential.imds_endpoint, + } + if isinstance(credential, AwsEnvironmentCredential): + return { + "credential_kind": 1, + "aws_access_key": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "imds_endpoint": None, + } + if isinstance(credential, AwsStaticCredential): + return { + "credential_kind": 2, + "aws_access_key": credential.access_key_id, + "aws_secret_access_key": credential.secret_access_key, + "aws_session_token": credential.session_token, + "imds_endpoint": None, + } + if isinstance(credential, AwsIamRoleCredential): + return { + "credential_kind": 3, + "aws_access_key": None, + "aws_secret_access_key": None, + "aws_session_token": None, + "imds_endpoint": credential.imds_endpoint, + } + + if isinstance(credential, AwsLegacyCredential): + return { + "credential_kind": 4, + "aws_access_key": credential.aws_access_key_id, + "aws_secret_access_key": credential.aws_secret_access_key, + "aws_session_token": credential.aws_session_token, + "imds_endpoint": None, + } + + raise TypeError(f"Unsupported credential type: {type(credential)!r}") + class RemoteEndpointType(enum.Enum): """ @@ -25,9 +123,10 @@ class RemoteEndpointType(enum.Enum): Automatically detect the endpoint type from the URL. KvikIO will attempt to infer the appropriate protocol based on the URL format. S3 : int - AWS S3 endpoint using credentials-based authentication. Requires - AWS environment variables (such as AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, - AWS_DEFAULT_REGION) to be set. + AWS S3 endpoint using credentials-based authentication. By default + KvikIO uses :class:`kvikio.aws_credentials.AwsDefaultCredential` semantics + (environment variables if set, otherwise IAM role credentials via the metadata + service / IMDSv2 when available). See :class:`kvikio.RemoteFile`. S3_PUBLIC : INT AWS S3 endpoint for publicly accessible objects. No credentials required as the objects have public read permissions enabled. Used for open datasets and public @@ -124,66 +223,69 @@ def open_s3( cls, bucket_name: str, object_name: str, + *, + credential: Optional[AwsCredential] = None, nbytes: Optional[int] = None, aws_region_name: Optional[str] = None, + aws_endpoint_url: Optional[str] = None, aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_endpoint_url: Optional[str] = None, aws_session_token: Optional[str] = None, ) -> RemoteFile: """Open a AWS S3 file from a bucket name and object name. - AWS credentials can be provided as keyword arguments or through - environment variables: - - - ``AWS_DEFAULT_REGION`` (or region_name parameter) - - ``AWS_ACCESS_KEY_ID`` (or access_key_id parameter) - - ``AWS_SECRET_ACCESS_KEY`` (or secret_access_key parameter) - - ``AWS_SESSION_TOKEN`` (or aws_session_token parameter, when using - temporary credentials) - - Additionally, to overwrite the AWS endpoint, set `AWS_ENDPOINT_URL` - (or endpoint_url parameter). - See - Parameters ---------- bucket_name The bucket name of the file. object_name The object name of the file. + credential + How to obtain AWS credentials. If ``None`` (default), uses + :class:`kvikio.aws_credentials.AwsDefaultCredential` (environment + variables if both ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY`` + are set and non-empty, otherwise IAM role credentials via the metadata service / IMDSv2). nbytes The size of the file. If None, KvikIO will ask the server for the file size. - aws_region + aws_region_name The AWS region, such as "us-east-1", to use. If None, the value of the - `AWS_DEFAULT_REGION` environment variable is used. - aws_access_key - The AWS access key to use. If None, the value of the - `AWS_ACCESS_KEY_ID` environment variable is used. - aws_secret_access_key - The AWS secret access key to use. If None, the value of the - `AWS_SECRET_ACCESS_KEY` environment variable is used. + ``AWS_DEFAULT_REGION`` environment variable is used. aws_endpoint_url Overwrite the endpoint url (including the protocol part) by using - the scheme: "//". If None, - the value of the `AWS_ENDPOINT_URL` environment variable is used. If - this is also not set, the regular AWS url scheme is used: - "https://.s3..amazonaws.com/". + the scheme: ``//``. If + None, the value of the ``AWS_ENDPOINT_URL`` environment variable is + used. If that is also not set, the regular AWS virtual-hosted-style + URL is used. + aws_access_key_id + Deprecated. Use ``credential=AwsStaticCredential(...)`` or set + environment variables with ``AwsDefaultCredential`` / + ``AwsEnvironmentCredential``. + aws_secret_access_key + Deprecated; see ``aws_access_key_id``. aws_session_token - The AWS session token to use. If None, the value of the - `AWS_SESSION_TOKEN` environment variable is used. + Deprecated; see ``aws_access_key_id``. + + Notes + ----- + Standard AWS environment variables are documented in the + `AWS CLI environment variable reference + `_. """ + cred_kw = _s3_credential_bridge_kwargs( + credential, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + ) return cls( _get_remote_module().RemoteFile.open_s3( bucket_name, object_name, - nbytes, - aws_region_name, - aws_access_key_id, - aws_secret_access_key, - aws_endpoint_url, - aws_session_token, + nbytes=nbytes, + aws_region_name=aws_region_name, + aws_endpoint_url=aws_endpoint_url, + **cred_kw, ) ) @@ -191,11 +293,13 @@ def open_s3( def open_s3_url( cls, url: str, + *, + credential: Optional[AwsCredential] = None, nbytes: Optional[int] = None, aws_region_name: Optional[str] = None, + aws_endpoint_url: Optional[str] = None, aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_endpoint_url: Optional[str] = None, aws_session_token: Optional[str] = None, ) -> RemoteFile: """Open a AWS S3 file from an URL. @@ -204,67 +308,51 @@ def open_s3_url( - A full http url such as "http://127.0.0.1/my/file", or - A S3 url such as "s3:///". - AWS credentials can be provided as keyword arguments or through - environment variables: - - - ``AWS_DEFAULT_REGION`` (or region_name parameter) - - ``AWS_ACCESS_KEY_ID`` (or access_key_id parameter) - - ``AWS_SECRET_ACCESS_KEY`` (or secret_access_key parameter) - - ``AWS_SESSION_TOKEN`` (or aws_session_token parameter, when using - temporary credentials) - - Additionally, if `url` is a S3 url, it is possible to overwrite the AWS endpoint - by setting `AWS_ENDPOINT_URL` (or endpoint_url parameter). - See - Parameters ---------- url Either a http url or a S3 url. + credential + Same semantics as :meth:`open_s3`. nbytes The size of the file. If None, KvikIO will ask the server for the file size. - aws_region - The AWS region, such as "us-east-1", to use. If None, the value of the - `AWS_DEFAULT_REGION` environment variable is used. - aws_access_key - The AWS access key to use. If None, the value of the - `AWS_ACCESS_KEY_ID` environment variable is used. - aws_secret_access_key - The AWS secret access key to use. If None, the value of the - `AWS_SECRET_ACCESS_KEY` environment variable is used. + aws_region_name + Same as :meth:`open_s3`. aws_endpoint_url - Overwrite the endpoint url (including the protocol part) by using - the scheme: "//". If None, - the value of the `AWS_ENDPOINT_URL` environment variable is used. If - this is also not set, the regular AWS url scheme is used: - "https://.s3..amazonaws.com/". + Same as :meth:`open_s3` (only applies when ``url`` uses the ``s3://`` scheme). + aws_access_key_id + Deprecated; same as :meth:`open_s3`. + aws_secret_access_key + Deprecated; same as :meth:`open_s3`. aws_session_token - The AWS session token to use. If None, the value of the - `AWS_SESSION_TOKEN` environment variable is used. + Deprecated; same as :meth:`open_s3`. """ + cred_kw = _s3_credential_bridge_kwargs( + credential, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + ) parsed_result = urllib.parse.urlparse(url.lower()) + m = _get_remote_module().RemoteFile if parsed_result.scheme in ("http", "https"): return cls( - _get_remote_module().RemoteFile.open_s3_from_http_url( + m.open_s3_from_http_url( url, - nbytes, - aws_region_name, - aws_access_key_id, - aws_secret_access_key, - aws_session_token, + nbytes=nbytes, + aws_region_name=aws_region_name, + **cred_kw, ) ) if parsed_result.scheme == "s3": return cls( - _get_remote_module().RemoteFile.open_s3_from_s3_url( + m.open_s3_from_s3_url( url, - nbytes, - aws_region_name, - aws_access_key_id, - aws_secret_access_key, - aws_endpoint_url, - aws_session_token, + nbytes=nbytes, + aws_region_name=aws_region_name, + aws_endpoint_url=aws_endpoint_url, + **cred_kw, ) ) raise ValueError(f"Unsupported protocol: {url}") diff --git a/python/kvikio/tests/test_s3_io.py b/python/kvikio/tests/test_s3_io.py index d8610c73bc..a6b2313cc9 100644 --- a/python/kvikio/tests/test_s3_io.py +++ b/python/kvikio/tests/test_s3_io.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import multiprocessing as mp @@ -161,7 +161,7 @@ def test_read_with_file_offset(s3_base, xp, start, end): @pytest.mark.parametrize("scheme", ["S3"]) @pytest.mark.parametrize( "remote_endpoint_type", - [kvikio.RemoteEndpointType.S3.AUTO, kvikio.RemoteEndpointType.S3], + [kvikio.RemoteEndpointType.AUTO, kvikio.RemoteEndpointType.S3], ) @pytest.mark.parametrize("allow_list", [None, [kvikio.RemoteEndpointType.S3]]) @pytest.mark.parametrize("nbytes", [None, 1]) @@ -217,3 +217,15 @@ def test_open_invalid(s3_base): kvikio.RemoteFile.open(url) with pytest.raises(RuntimeError, match="Invalid URL"): kvikio.RemoteFile.open(url, kvikio.RemoteEndpointType.S3) + + +def test_deprecated_legacy_aws_credential_kwargs_emits_warning(): + from kvikio.remote_file import _s3_credential_bridge_kwargs + + with pytest.warns(FutureWarning, match="aws_access_key_id"): + _s3_credential_bridge_kwargs( + None, + aws_access_key_id="x", + aws_secret_access_key=None, + aws_session_token=None, + )