Repository navigation
Add APIs for inferring remote endpoint type from URL. #982
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.
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
Changes from 2 commits
2ca7f25
cbffb90
d813778
112ab81
37b9ddb
b752f5a
e91aed6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,6 +4,7 @@ | |
| */ | ||
|
|
||
| #include <algorithm> | ||
| #include <array> | ||
|
TomAugspurger marked this conversation as resolved.
Resolved
|
||
| #include <cassert> | ||
| #include <cstddef> | ||
| #include <cstring> | ||
|
|
@@ -231,6 +232,89 @@ std::string encode_special_chars_in_path(std::string const& url) | |
| components.path = detail::UrlEncoder::encode_path(components.path.value()); | ||
| return detail::UrlBuilder::build_manually(components); | ||
| } | ||
|
|
||
| std::vector<RemoteEndpointType> get_default_allow_list() | ||
| { | ||
| return {RemoteEndpointType::S3, | ||
| RemoteEndpointType::S3_PUBLIC, | ||
| RemoteEndpointType::S3_PRESIGNED_URL, | ||
| RemoteEndpointType::WEBHDFS, | ||
| RemoteEndpointType::HTTP}; | ||
| } | ||
|
TomAugspurger marked this conversation as resolved.
Resolved
|
||
|
|
||
| std::unique_ptr<RemoteEndpoint> create_endpoint_from_type(std::string const& url, | ||
| std::string const& scheme, | ||
| RemoteEndpointType type) | ||
| { | ||
| switch (type) { | ||
| case RemoteEndpointType::S3: | ||
| if (!S3Endpoint::is_url_valid(url)) { return nullptr; } | ||
| if (scheme == "s3") { | ||
| auto const [bucket, object] = S3Endpoint::parse_s3_url(url); | ||
| return std::make_unique<S3Endpoint>(std::pair{bucket, object}); | ||
| } | ||
| return std::make_unique<S3Endpoint>(url); | ||
|
|
||
| case RemoteEndpointType::S3_PUBLIC: | ||
| if (!S3PublicEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<S3PublicEndpoint>(url); | ||
|
|
||
| case RemoteEndpointType::S3_PRESIGNED_URL: | ||
| if (!S3EndpointWithPresignedUrl::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<S3EndpointWithPresignedUrl>(url); | ||
|
|
||
| case RemoteEndpointType::WEBHDFS: | ||
| if (!WebHdfsEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<WebHdfsEndpoint>(url); | ||
|
|
||
| case RemoteEndpointType::HTTP: | ||
| if (!HttpEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<HttpEndpoint>(url); | ||
|
|
||
| default: return nullptr; | ||
| } | ||
| } | ||
|
|
||
| struct InferredEndpoint { | ||
| RemoteEndpointType remote_endpoint_type{}; | ||
| std::unique_ptr<RemoteEndpoint> endpoint; | ||
| std::optional<std::size_t> probed_nbytes; | ||
| }; | ||
|
TomAugspurger marked this conversation as resolved.
Resolved
|
||
|
|
||
| InferredEndpoint infer_endpoint_impl(std::string const& url, | ||
| std::vector<RemoteEndpointType> const& allow_list, | ||
| bool probe_s3_connectivity) | ||
| { | ||
| auto const scheme = | ||
| detail::UrlParser::extract_component(url, CURLUPART_SCHEME, CURLU_NON_SUPPORT_SCHEME); | ||
| KVIKIO_EXPECT(scheme.has_value(), "Missing scheme in URL."); | ||
|
|
||
| for (auto const& type : allow_list) { | ||
| try { | ||
| auto endpoint = create_endpoint_from_type(url, scheme.value(), type); | ||
| if (endpoint == nullptr) { continue; } | ||
|
|
||
| std::optional<std::size_t> probed_nbytes = std::nullopt; | ||
| if (probe_s3_connectivity && type == RemoteEndpointType::S3) { | ||
| // Check connectivity for the credential-based S3 endpoint and reuse this size in | ||
| // RemoteHandle::open to avoid a second HEAD request. | ||
| probed_nbytes = endpoint->get_file_size(); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Since the This way, this PR alone will reduce the cudf-polars per-file HEAD request to 1, if I'm not missing anything.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Personally, I'm not comfortable making something like And, at least from the cudf-polars side, we'll temporarily use kvikio to get the file size exactly once per file. But in the medium term we'll be able to get that information from cudf-polars; either way, we'll be able to provide it the size in all subsequent operations. Finally, I don't think this would directly help cudf-polars since in our current design we don't store
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Does this sound OK @kingcrimsontianyu? I'd like to merge this today if possible, so we can use it in NVIDIA/cudf#22739 |
||
| } | ||
| return InferredEndpoint{type, std::move(endpoint), probed_nbytes}; | ||
| } catch (...) { | ||
| // If the credential-based S3 endpoint cannot be used to access the URL, try using | ||
| // S3 public endpoint instead when it is in the allowlist. | ||
| if (type == RemoteEndpointType::S3 && | ||
| std::find(allow_list.begin(), allow_list.end(), RemoteEndpointType::S3_PUBLIC) != | ||
| allow_list.end()) { | ||
| return InferredEndpoint{ | ||
| RemoteEndpointType::S3_PUBLIC, std::make_unique<S3PublicEndpoint>(url), std::nullopt}; | ||
| } | ||
| throw; | ||
| } | ||
| } | ||
| KVIKIO_FAIL("Unsupported endpoint URL.", std::runtime_error); | ||
| } | ||
| } // namespace | ||
|
|
||
| RemoteEndpoint::RemoteEndpoint(RemoteEndpointType remote_endpoint_type) | ||
|
|
@@ -575,87 +659,27 @@ bool S3EndpointWithPresignedUrl::is_url_valid(std::string const& url) noexcept | |
| } | ||
| } | ||
|
|
||
| RemoteEndpointType infer_remote_endpoint_type(std::string url) | ||
|
TomAugspurger marked this conversation as resolved.
Resolved
|
||
| { | ||
| KVIKIO_NVTX_FUNC_RANGE(); | ||
| return infer_endpoint_impl(url, get_default_allow_list(), false).remote_endpoint_type; | ||
| } | ||
|
|
||
| RemoteHandle RemoteHandle::open(std::string url, | ||
| RemoteEndpointType remote_endpoint_type, | ||
| std::optional<std::vector<RemoteEndpointType>> allow_list, | ||
| std::optional<std::size_t> nbytes) | ||
| { | ||
| KVIKIO_NVTX_FUNC_RANGE(); | ||
| if (!allow_list.has_value()) { | ||
| allow_list = {RemoteEndpointType::S3, | ||
| RemoteEndpointType::S3_PUBLIC, | ||
| RemoteEndpointType::S3_PRESIGNED_URL, | ||
| RemoteEndpointType::WEBHDFS, | ||
| RemoteEndpointType::HTTP}; | ||
| } | ||
|
|
||
| auto const scheme = | ||
| detail::UrlParser::extract_component(url, CURLUPART_SCHEME, CURLU_NON_SUPPORT_SCHEME); | ||
| KVIKIO_EXPECT(scheme.has_value(), "Missing scheme in URL."); | ||
|
|
||
| // Helper to create endpoint based on type | ||
| auto create_endpoint = | ||
| [&url = url, &scheme = scheme](RemoteEndpointType type) -> std::unique_ptr<RemoteEndpoint> { | ||
| switch (type) { | ||
| case RemoteEndpointType::S3: | ||
| 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<S3Endpoint>(std::pair{bucket, object}); | ||
| } | ||
| return std::make_unique<S3Endpoint>(url); | ||
|
|
||
| case RemoteEndpointType::S3_PUBLIC: | ||
| if (!S3PublicEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<S3PublicEndpoint>(url); | ||
|
|
||
| case RemoteEndpointType::S3_PRESIGNED_URL: | ||
| if (!S3EndpointWithPresignedUrl::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<S3EndpointWithPresignedUrl>(url); | ||
|
|
||
| case RemoteEndpointType::WEBHDFS: | ||
| if (!WebHdfsEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<WebHdfsEndpoint>(url); | ||
|
|
||
| case RemoteEndpointType::HTTP: | ||
| if (!HttpEndpoint::is_url_valid(url)) { return nullptr; } | ||
| return std::make_unique<HttpEndpoint>(url); | ||
|
|
||
| default: return nullptr; | ||
| } | ||
| }; | ||
| if (!allow_list.has_value()) { allow_list = get_default_allow_list(); } | ||
|
|
||
| std::unique_ptr<RemoteEndpoint> endpoint; | ||
| std::optional<std::size_t> probed_nbytes; | ||
|
|
||
| if (remote_endpoint_type == RemoteEndpointType::AUTO) { | ||
| // Try each allowed type in the order of allowlist | ||
| for (auto const& type : allow_list.value()) { | ||
| try { | ||
| endpoint = create_endpoint(type); | ||
| if (endpoint == nullptr) { continue; } | ||
| if (type == RemoteEndpointType::S3) { | ||
| // Check connectivity for the credential-based S3 endpoint, and throw an exception if | ||
| // failed. Reuse this size when constructing the handle to avoid a second HEAD request. | ||
| probed_nbytes = endpoint->get_file_size(); | ||
| } | ||
| } catch (...) { | ||
| // If the credential-based S3 endpoint cannot be used to access the URL, try using S3 public | ||
| // endpoint instead if it is in the allowlist | ||
| if (type == RemoteEndpointType::S3 && | ||
| std::find(allow_list->begin(), allow_list->end(), RemoteEndpointType::S3_PUBLIC) != | ||
| allow_list->end()) { | ||
| endpoint = std::make_unique<S3PublicEndpoint>(url); | ||
| probed_nbytes = std::nullopt; | ||
| } else { | ||
| throw; | ||
| } | ||
| } | ||
|
|
||
| // At this point, a matching endpoint has been found | ||
| break; | ||
| } | ||
| KVIKIO_EXPECT(endpoint.get() != nullptr, "Unsupported endpoint URL.", std::runtime_error); | ||
| auto inferred = infer_endpoint_impl(url, allow_list.value(), true); | ||
| endpoint = std::move(inferred.endpoint); | ||
| probed_nbytes = inferred.probed_nbytes; | ||
| } else { | ||
| // Validate it is in the allow list | ||
| KVIKIO_EXPECT( | ||
|
|
@@ -665,7 +689,10 @@ RemoteHandle RemoteHandle::open(std::string url, | |
| std::runtime_error); | ||
|
|
||
| // Create the specific type | ||
| endpoint = create_endpoint(remote_endpoint_type); | ||
| auto const scheme = | ||
| detail::UrlParser::extract_component(url, CURLUPART_SCHEME, CURLU_NON_SUPPORT_SCHEME); | ||
| KVIKIO_EXPECT(scheme.has_value(), "Missing scheme in URL."); | ||
| endpoint = create_endpoint_from_type(url, scheme.value(), remote_endpoint_type); | ||
| KVIKIO_EXPECT(endpoint.get() != nullptr, | ||
| std::string{"Invalid URL for "} + | ||
| get_remote_endpoint_type_name(remote_endpoint_type) + " endpoint", | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| import pytest | ||
|
|
||
| import kvikio | ||
|
|
||
| pytestmark = pytest.mark.skipif( | ||
| not kvikio.is_remote_file_available(), | ||
| reason=( | ||
| "RemoteFile not available, please build KvikIO " | ||
| "with libcurl (-DKvikIO_REMOTE_SUPPORT=ON)" | ||
| ), | ||
| ) | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def aws_env(monkeypatch): | ||
| monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") | ||
| monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test-access-key") | ||
| monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret-key") | ||
|
|
||
|
|
||
| def test_infer_remote_endpoint_type(aws_env): | ||
| assert ( | ||
| kvikio.infer_remote_endpoint_type("s3://bucket-name/object-key-name") | ||
| == kvikio.RemoteEndpointType.S3 | ||
| ) | ||
| assert ( | ||
| kvikio.infer_remote_endpoint_type("https://host:1234/webhdfs/v1/data.bin") | ||
| == kvikio.RemoteEndpointType.WEBHDFS | ||
| ) | ||
| assert ( | ||
| kvikio.infer_remote_endpoint_type("https://example.com/path/file.bin") | ||
| == kvikio.RemoteEndpointType.HTTP | ||
| ) | ||
| assert ( | ||
| kvikio.infer_remote_endpoint_type( | ||
| "https://bucket-name.s3.region-code.amazonaws.com/" | ||
| "object-key-name?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Signature=sig&" | ||
| "X-Amz-Credential=cred&X-Amz-SignedHeaders=host" | ||
| ) | ||
| == kvikio.RemoteEndpointType.S3_PRESIGNED_URL | ||
| ) | ||
|
|
||
|
|
||
| def test_infer_remote_endpoint_type_invalid_url(): | ||
| with pytest.raises(RuntimeError, match="Bad scheme"): | ||
| kvikio.infer_remote_endpoint_type("example.com/path") | ||
|
|
||
| with pytest.raises(RuntimeError, match="Unsupported endpoint URL"): | ||
| kvikio.infer_remote_endpoint_type("unsupported://example.com/path") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we need to document the difference between this utility function and the
openfunction, that given a URL valid for both S3 private and public endpoints, this function will always return the private endpoint type, and that it is not possible for us to disambiguate S3 private and public endpoints from a URL alone.Then in the cudf PR NVIDIA/cudf#22739, we may document that known file size + URL without S3 credential is not supported.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I added a note in b752f5a.
But I'd like to get away from "public" vs. "private" URLs entirely. It'd much prefer to just deal with some URL and an authorization method (possibly a chained authorization method that tries multiple).