wehub-resource-sync
1562 lines
62 KiB
Python
1562 lines
62 KiB
Python
import asyncio
|
|
import os
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.fs as pafs
|
|
import pytest
|
|
|
|
from ray.data._internal.planner._obstore_download import (
|
|
_BUCKET_REGION_CACHE,
|
|
_FILE_SIZE_COLUMN_PREFIX,
|
|
StoreRegistry,
|
|
_discover_aws_bucket_region,
|
|
_download_uris_with_obstore,
|
|
_extract_credentials_from_filesystem,
|
|
_is_obstore_supported_url,
|
|
_native_s3_obstore_kwargs,
|
|
_obstore_filesystem_requires_threaded_download,
|
|
_plan_obstore_routing,
|
|
_resolve_size,
|
|
_S3FSSessionCredentialProvider,
|
|
_split_obstore_uri,
|
|
download_bytes_async,
|
|
)
|
|
from ray.data._internal.planner.plan_download_op import (
|
|
download_bytes_threaded,
|
|
)
|
|
from ray.data._internal.util import RetryingPyFileSystem
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource.path_util import _split_uri
|
|
|
|
|
|
# TestSplitUri
|
|
@pytest.mark.parametrize(
|
|
"uri, expected_store_url, expected_path",
|
|
[
|
|
("s3://my-bucket/prefix/key.jpg", "s3://my-bucket", "prefix/key.jpg"),
|
|
("gs://bucket/path/to/file", "gs://bucket", "path/to/file"),
|
|
("https://host.com/a/b/c.png", "https://host.com", "a/b/c.png"),
|
|
# file:// URIs have an empty netloc; path starts at the third slash.
|
|
("file:///tmp/test.txt", "file://", "tmp/test.txt"),
|
|
# Nested paths.
|
|
("s3://bucket/a/b/c/d/e.parquet", "s3://bucket", "a/b/c/d/e.parquet"),
|
|
# Trailing slash: empty path.
|
|
("s3://bucket/", "s3://bucket", ""),
|
|
# Query string must be preserved for pre-signed / parameterized URLs.
|
|
(
|
|
"https://s3.amazonaws.com/bucket/key?X-Amz-Signature=abc&X-Amz-Expires=3600",
|
|
"https://s3.amazonaws.com",
|
|
"bucket/key?X-Amz-Signature=abc&X-Amz-Expires=3600",
|
|
),
|
|
# '#' in path must not be parsed as a fragment delimiter.
|
|
("s3://bucket/my#file.jpg", "s3://bucket", "my#file.jpg"),
|
|
# Double slash after bucket: key legitimately starts with '/'.
|
|
("s3://bucket//path/file.txt", "s3://bucket", "/path/file.txt"),
|
|
],
|
|
)
|
|
def test_split_uri(uri, expected_store_url, expected_path):
|
|
store_url, path = _split_uri(uri)
|
|
assert store_url == expected_store_url
|
|
assert path == expected_path
|
|
|
|
|
|
# TestSplitObstoreUri — AWS path-style HTTPS URLs must be rewritten to s3://
|
|
# so obstore's S3Store keys off the bucket (with region discovery) instead of
|
|
# off the regional host endpoint (which pins region and fails with
|
|
# BareRedirect on cross-region buckets).
|
|
@pytest.mark.parametrize(
|
|
"uri, expected_store_url, expected_path",
|
|
[
|
|
# Path-style regional: must rewrite to s3://<bucket>.
|
|
(
|
|
"https://s3.us-east-1.amazonaws.com/ray-example-data/imagenet/foo.jpg",
|
|
"s3://ray-example-data",
|
|
"imagenet/foo.jpg",
|
|
),
|
|
# Path-style legacy global endpoint: also rewrite.
|
|
(
|
|
"https://s3.amazonaws.com/my-bucket/key.bin",
|
|
"s3://my-bucket",
|
|
"key.bin",
|
|
),
|
|
# Virtual-host style: bucket already in netloc, leave alone.
|
|
(
|
|
"https://my-bucket.s3.us-west-2.amazonaws.com/key.bin",
|
|
"https://my-bucket.s3.us-west-2.amazonaws.com",
|
|
"key.bin",
|
|
),
|
|
# Native s3:// URIs: no change.
|
|
(
|
|
"s3://my-bucket/prefix/key.jpg",
|
|
"s3://my-bucket",
|
|
"prefix/key.jpg",
|
|
),
|
|
# Pre-signed URLs (any query string): signature is bound to host +
|
|
# path + query, so rewriting would invalidate it. Pass through.
|
|
(
|
|
"https://s3.us-east-1.amazonaws.com/bucket/key?X-Amz-Signature=abc",
|
|
"https://s3.us-east-1.amazonaws.com",
|
|
"bucket/key?X-Amz-Signature=abc",
|
|
),
|
|
# Non-AWS HTTPS (MinIO, R2 custom domain, generic file server):
|
|
# leave alone — only AWS S3 needs the region-discovery rewrite.
|
|
(
|
|
"https://files.example.com/bucket/key.bin",
|
|
"https://files.example.com",
|
|
"bucket/key.bin",
|
|
),
|
|
# Other clouds: untouched.
|
|
("gs://bucket/path", "gs://bucket", "path"),
|
|
# Path-style with nested key.
|
|
(
|
|
"https://s3.eu-central-1.amazonaws.com/bucket/a/b/c.parquet",
|
|
"s3://bucket",
|
|
"a/b/c.parquet",
|
|
),
|
|
# Path-style with only a bucket (no key): bucket is preserved, key
|
|
# is empty. Edge case; downstream obstore call would fail but that's
|
|
# the caller's problem.
|
|
(
|
|
"https://s3.amazonaws.com/just-a-bucket",
|
|
"s3://just-a-bucket",
|
|
"",
|
|
),
|
|
],
|
|
)
|
|
def test_split_obstore_uri(uri, expected_store_url, expected_path):
|
|
store_url, path = _split_obstore_uri(uri)
|
|
assert store_url == expected_store_url
|
|
assert path == expected_path
|
|
|
|
|
|
# TestBucketRegionDiscovery — exercise the region-discovery path used by
|
|
# StoreRegistry to find the real region of cross-region S3 buckets. Region
|
|
# resolution is delegated to ``pyarrow.fs.resolve_s3_region``; these tests
|
|
# patch it to avoid real network calls.
|
|
class TestBucketRegionDiscovery:
|
|
@staticmethod
|
|
def _clear_cache():
|
|
_BUCKET_REGION_CACHE.clear()
|
|
|
|
def test_resolved_region_is_returned(self):
|
|
# Whether the bucket is same-region (200 OK) or cross-region (301
|
|
# PermanentRedirect), PyArrow reads x-amz-bucket-region and returns it.
|
|
self._clear_cache()
|
|
with patch("pyarrow.fs.resolve_s3_region", return_value="us-west-2"):
|
|
assert _discover_aws_bucket_region("cross-region-bucket") == "us-west-2"
|
|
|
|
def test_cache_hit_avoids_second_probe(self):
|
|
# First call probes; second call must not invoke resolve_s3_region.
|
|
self._clear_cache()
|
|
with patch("pyarrow.fs.resolve_s3_region", return_value="eu-west-1") as resolve:
|
|
_discover_aws_bucket_region("euro-bucket")
|
|
_discover_aws_bucket_region("euro-bucket")
|
|
assert resolve.call_count == 1
|
|
|
|
def test_failure_returns_none_and_caches(self):
|
|
# On resolution failure (network error, non-AWS endpoint) return None
|
|
# so StoreRegistry leaves obstore's defaults intact for MinIO/R2/etc.
|
|
# Cache the negative result so we don't keep probing.
|
|
self._clear_cache()
|
|
with patch(
|
|
"pyarrow.fs.resolve_s3_region", side_effect=OSError("nope")
|
|
) as resolve:
|
|
assert _discover_aws_bucket_region("unreachable") is None
|
|
assert _discover_aws_bucket_region("unreachable") is None
|
|
assert resolve.call_count == 1
|
|
|
|
def test_missing_resolve_s3_region_returns_none(self):
|
|
# PyArrow builds compiled without S3 support lack resolve_s3_region,
|
|
# so the attribute access raises AttributeError; discovery must degrade
|
|
# to None rather than propagate it.
|
|
self._clear_cache()
|
|
with patch("pyarrow.fs.resolve_s3_region", side_effect=AttributeError("no S3")):
|
|
assert _discover_aws_bucket_region("no-s3-support") is None
|
|
|
|
def test_none_result_does_not_overwrite_cached_region(self):
|
|
# Race: this thread probes (outside the lock) and gets None, while a
|
|
# concurrent thread resolves and caches the real region. The failed
|
|
# probe must not clobber the cached region, otherwise later
|
|
# StoreRegistry.get calls skip region injection and cross-region
|
|
# downloads fail intermittently.
|
|
self._clear_cache()
|
|
|
|
def resolve_then_simulate_concurrent_win(bucket):
|
|
# Stand in for another thread caching the real region mid-probe.
|
|
_BUCKET_REGION_CACHE[bucket] = "us-west-2"
|
|
return None
|
|
|
|
with patch(
|
|
"pyarrow.fs.resolve_s3_region",
|
|
side_effect=resolve_then_simulate_concurrent_win,
|
|
):
|
|
assert _discover_aws_bucket_region("racy-bucket") == "us-west-2"
|
|
assert _BUCKET_REGION_CACHE["racy-bucket"] == "us-west-2"
|
|
|
|
|
|
# TestStoreRegistryRegionInjection — verify region discovery happens at
|
|
# store-construction time for s3:// URLs without an explicit region, and is
|
|
# correctly skipped when the caller has already configured one.
|
|
class TestStoreRegistryRegionInjection:
|
|
@staticmethod
|
|
def _make_registry_capturing_kwargs():
|
|
_BUCKET_REGION_CACHE.clear()
|
|
captured = {}
|
|
|
|
def fake_from_url(url, **kwargs):
|
|
captured["store_url"] = url
|
|
captured["kwargs"] = kwargs
|
|
return MagicMock(name="fake-store")
|
|
|
|
reg = StoreRegistry()
|
|
reg._from_url = fake_from_url
|
|
return reg, captured
|
|
|
|
def test_s3_url_injects_discovered_region(self):
|
|
reg, captured = self._make_registry_capturing_kwargs()
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region",
|
|
return_value="us-west-2",
|
|
):
|
|
reg.get("s3://ray-example-data")
|
|
assert captured["kwargs"].get("region") == "us-west-2"
|
|
|
|
def test_s3a_url_also_injects_region(self):
|
|
reg, captured = self._make_registry_capturing_kwargs()
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region",
|
|
return_value="us-west-2",
|
|
):
|
|
reg.get("s3a://some-bucket")
|
|
assert captured["kwargs"].get("region") == "us-west-2"
|
|
|
|
def test_explicit_region_not_overridden(self):
|
|
# Caller (e.g., creds extracted from PyArrow S3FileSystem) already
|
|
# supplied region; do not probe and do not override.
|
|
reg = StoreRegistry(region="eu-central-1")
|
|
captured = {}
|
|
|
|
def fake_from_url(url, **kwargs):
|
|
captured["kwargs"] = kwargs
|
|
return MagicMock()
|
|
|
|
reg._from_url = fake_from_url
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region"
|
|
) as probe:
|
|
reg.get("s3://bucket")
|
|
probe.assert_not_called()
|
|
assert captured["kwargs"]["region"] == "eu-central-1"
|
|
|
|
def test_custom_endpoint_skips_probe(self):
|
|
# MinIO / R2 / localstack: caller configured a custom endpoint, so
|
|
# the AWS region probe would be both wrong and wasteful.
|
|
reg = StoreRegistry(endpoint="http://localhost:9000")
|
|
reg._from_url = lambda *a, **k: MagicMock()
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region"
|
|
) as probe:
|
|
reg.get("s3://bucket")
|
|
probe.assert_not_called()
|
|
|
|
def test_https_url_does_not_probe(self):
|
|
# Virtual-hosted-style or presigned HTTPS URLs reach get() unchanged
|
|
# (after _split_obstore_uri). obstore parses the region from the
|
|
# netloc directly, so probing would be redundant.
|
|
reg, captured = self._make_registry_capturing_kwargs()
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region"
|
|
) as probe:
|
|
reg.get("https://bucket.s3.us-east-1.amazonaws.com")
|
|
probe.assert_not_called()
|
|
|
|
|
|
# TestResolveSizeRewrite — the HEAD size probe must apply the same path-style
|
|
# -> s3:// rewrite as the download paths. Otherwise a cross-region path-style
|
|
# URL keeps the regional HTTPS store, hits BareRedirect, returns size 0, and
|
|
# wrongly skips ranged downloads even when a whole-file GET would succeed.
|
|
class TestResolveSizeRewrite:
|
|
def test_resolve_size_rewrites_path_style_uri(self):
|
|
_BUCKET_REGION_CACHE.clear()
|
|
captured = {}
|
|
|
|
def fake_from_url(url, **kwargs):
|
|
captured["store_url"] = url
|
|
return MagicMock(name="store")
|
|
|
|
registry = StoreRegistry()
|
|
registry._from_url = fake_from_url
|
|
|
|
async def fake_head_async(store, path):
|
|
captured["path"] = path
|
|
return {"size": 4096}
|
|
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._discover_aws_bucket_region",
|
|
return_value="us-west-2",
|
|
), patch("obstore.head_async", side_effect=fake_head_async):
|
|
size = asyncio.run(
|
|
_resolve_size(
|
|
"https://s3.us-east-1.amazonaws.com/cross-region-bucket/key.bin",
|
|
registry,
|
|
asyncio.Semaphore(),
|
|
)
|
|
)
|
|
|
|
# Rewritten to s3://<bucket> so region discovery applies; HEAD succeeds.
|
|
assert size == 4096
|
|
assert captured["store_url"] == "s3://cross-region-bucket"
|
|
assert captured["path"] == "key.bin"
|
|
|
|
|
|
# TestDownloadHelpers
|
|
class TestDownloadHelpers:
|
|
"""Unit tests for credential extraction and URL-scheme detection."""
|
|
|
|
# _extract_credentials_from_filesystem filesystem-agnostic cases
|
|
def test_extract_credentials_none(self):
|
|
assert _extract_credentials_from_filesystem(None) == {}
|
|
|
|
def test_extract_credentials_unrecognized_fs(self):
|
|
# Unrecognized non-None filesystems (LocalFileSystem, custom FS) return
|
|
# None so the planner routes to the threaded path, keeping the user's
|
|
# filesystem authoritative instead of handing control to obstore's
|
|
# default credential chain.
|
|
assert _extract_credentials_from_filesystem(pafs.LocalFileSystem()) is None
|
|
|
|
def test_extract_credentials_unwraps_retrying_fs(self):
|
|
# RetryingPyFileSystem must be unwrapped before the isinstance checks,
|
|
# otherwise an S3/GCS/Azure branch would never be reached. After
|
|
# unwrapping to an unrecognized FS, extraction returns None.
|
|
inner = pafs.LocalFileSystem()
|
|
retrying = RetryingPyFileSystem.wrap(inner, retryable_errors=[])
|
|
assert _extract_credentials_from_filesystem(retrying) is None
|
|
|
|
# Native S3FileSystem is translated to obstore kwargs from its pickle state
|
|
# so the fast obstore path is kept. Anonymous, explicit static keys, and the
|
|
# default credential chain are all statically representable.
|
|
@pytest.mark.parametrize(
|
|
"kwargs, expected",
|
|
[
|
|
(
|
|
{"anonymous": True, "region": "us-west-2"},
|
|
{"region": "us-west-2", "skip_signature": True},
|
|
),
|
|
(
|
|
{"region": "us-west-2"},
|
|
{"region": "us-west-2"},
|
|
),
|
|
(
|
|
{"access_key": "AKID", "secret_key": "SECRET", "region": "us-west-2"},
|
|
{
|
|
"region": "us-west-2",
|
|
"access_key_id": "AKID",
|
|
"secret_access_key": "SECRET",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_native_s3_translates_to_obstore(self, kwargs, expected):
|
|
try:
|
|
fs = pafs.S3FileSystem(**kwargs)
|
|
except Exception:
|
|
pytest.skip("Cannot instantiate S3FileSystem in this environment")
|
|
assert _extract_credentials_from_filesystem(fs) == expected
|
|
assert _plan_obstore_routing(fs) == (True, expected)
|
|
|
|
def test_native_s3_unwrapped_from_retrying_wrapper(self):
|
|
# RetryingPyFileSystem is unwrapped before the S3 check, so a wrapped
|
|
# native S3FileSystem is translated to obstore kwargs just the same.
|
|
try:
|
|
fs = pafs.S3FileSystem(anonymous=True, region="us-west-2")
|
|
except Exception:
|
|
pytest.skip("Cannot instantiate S3FileSystem in this environment")
|
|
retrying = RetryingPyFileSystem.wrap(fs, retryable_errors=[])
|
|
assert _extract_credentials_from_filesystem(retrying) == {
|
|
"region": "us-west-2",
|
|
"skip_signature": True,
|
|
}
|
|
|
|
def test_native_s3_assume_role_falls_back_to_threaded(self):
|
|
# Assume-role creds rotate; a static obstore snapshot would go stale, so
|
|
# the translation declines (None) and the caller uses the threaded path.
|
|
# Driven through the pickle state to avoid constructing a real
|
|
# assume-role filesystem (which may attempt STS).
|
|
class _Stub:
|
|
def __reduce__(self):
|
|
return (
|
|
object,
|
|
(
|
|
{
|
|
"role_arn": "arn:aws:iam::123456789012:role/test",
|
|
"anonymous": False,
|
|
"access_key": None,
|
|
"region": "us-west-2",
|
|
},
|
|
),
|
|
)
|
|
|
|
assert _native_s3_obstore_kwargs(_Stub()) is None
|
|
|
|
def test_native_s3_unreadable_pickle_state_falls_back_to_threaded(self):
|
|
# If the pickle state can't be introspected (e.g. a future PyArrow
|
|
# changes its serialization), fail safe to the threaded path.
|
|
class _Stub:
|
|
def __reduce__(self):
|
|
raise RuntimeError("unexpected serialization")
|
|
|
|
assert _native_s3_obstore_kwargs(_Stub()) is None
|
|
|
|
# _extract_credentials_from_filesystem with GCS
|
|
@pytest.mark.parametrize(
|
|
"anonymous, expected",
|
|
[
|
|
(True, {"skip_signature": True}),
|
|
(False, {}),
|
|
(None, {}),
|
|
],
|
|
)
|
|
def test_extract_credentials_gcs_anonymous_mock(self, anonymous, expected):
|
|
# GCS anonymous access maps to obstore's skip_signature.
|
|
# Other GCS credentials (service account, application default) are
|
|
# not accessible via PyArrow attributes; obstore resolves them from
|
|
# the environment automatically.
|
|
mock_fs = MagicMock()
|
|
mock_fs.anonymous = anonymous
|
|
with patch("pyarrow.fs.GcsFileSystem", type(mock_fs)):
|
|
result = _extract_credentials_from_filesystem(mock_fs)
|
|
assert result == expected
|
|
|
|
# _extract_credentials_from_filesystem with Azure
|
|
@pytest.mark.parametrize(
|
|
"account_name, account_key, expected",
|
|
[
|
|
(
|
|
"myaccount",
|
|
"mykey",
|
|
{"account_name": "myaccount", "account_key": "mykey"},
|
|
),
|
|
# account_key is optional (e.g. when using managed identity).
|
|
("myaccount", None, {"account_name": "myaccount"}),
|
|
(None, None, {}),
|
|
],
|
|
)
|
|
def test_extract_credentials_azure_mock(self, account_name, account_key, expected):
|
|
mock_fs = MagicMock()
|
|
mock_fs.account_name = account_name
|
|
mock_fs.account_key = account_key
|
|
with patch("pyarrow.fs.AzureFileSystem", type(mock_fs), create=True):
|
|
result = _extract_credentials_from_filesystem(mock_fs)
|
|
assert result == expected
|
|
|
|
def test_fsspec_s3_credentials_for_obstore_and_retrying_wrap(self):
|
|
"""PyFileSystem(FSSpecHandler(s3fs)): forward keys to obstore; RetryingPyFileSystem unwraps."""
|
|
s3fs = pytest.importorskip("s3fs", reason="s3fs not installed")
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
inner = s3fs.S3FileSystem(
|
|
key="AKIATEST",
|
|
secret="secretval",
|
|
token="ststoken",
|
|
client_kwargs={"endpoint_url": "https://minio.local:9000"},
|
|
)
|
|
wrapped = PyFileSystem(FSSpecHandler(inner))
|
|
expected = {
|
|
"access_key_id": "AKIATEST",
|
|
"secret_access_key": "secretval",
|
|
"session_token": "ststoken",
|
|
"endpoint": "https://minio.local:9000",
|
|
}
|
|
assert _extract_credentials_from_filesystem(wrapped) == expected
|
|
assert _obstore_filesystem_requires_threaded_download(wrapped) is False
|
|
|
|
retrying = RetryingPyFileSystem.wrap(wrapped, retryable_errors=[])
|
|
assert _extract_credentials_from_filesystem(retrying) == expected
|
|
assert _obstore_filesystem_requires_threaded_download(retrying) is False
|
|
|
|
empty = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem()))
|
|
assert _obstore_filesystem_requires_threaded_download(empty) is False
|
|
|
|
@pytest.mark.parametrize("protocol", ["ftp", "gcs"])
|
|
def test_obstore_threaded_download_non_s3_fsspec_protocols(self, protocol):
|
|
import fsspec
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
proto = protocol
|
|
|
|
class _Stub(fsspec.AbstractFileSystem):
|
|
protocol = proto
|
|
|
|
wrapped = PyFileSystem(FSSpecHandler(_Stub()))
|
|
assert _obstore_filesystem_requires_threaded_download(wrapped) is True
|
|
|
|
# _is_obstore_supported_url
|
|
@pytest.mark.parametrize(
|
|
"uri, raises, expected",
|
|
[
|
|
# Recognized scheme: parse_scheme succeeds.
|
|
("s3://bucket/key", None, True),
|
|
("https://host.com/file", None, True),
|
|
# Unrecognized scheme: parse_scheme raises.
|
|
("custom://bucket/key", ValueError("unsupported"), False),
|
|
# obstore not installed: parse_scheme is None, so TypeError.
|
|
("s3://bucket/key", "none", False),
|
|
],
|
|
)
|
|
def test_is_obstore_supported_url(self, uri, raises, expected):
|
|
if raises == "none":
|
|
ctx = patch(
|
|
"ray.data._internal.planner._obstore_download.obstore_parse_scheme",
|
|
new=None,
|
|
)
|
|
else:
|
|
mock = (
|
|
MagicMock(side_effect=raises)
|
|
if raises
|
|
else MagicMock(return_value="s3")
|
|
)
|
|
ctx = patch(
|
|
"ray.data._internal.planner._obstore_download.obstore_parse_scheme",
|
|
new=mock,
|
|
)
|
|
with ctx:
|
|
assert _is_obstore_supported_url(uri) is expected
|
|
|
|
|
|
# Test helpers for fsspec-S3 credential extraction & routing.
|
|
def _sync_session(access_key, secret_key="sk", token="tk", expiry_time=None):
|
|
"""Build a botocore-style sync session returning frozen creds."""
|
|
frozen = MagicMock(access_key=access_key, secret_key=secret_key, token=token)
|
|
creds = MagicMock(expiry_time=expiry_time)
|
|
creds.get_frozen_credentials = MagicMock(return_value=frozen)
|
|
session = MagicMock()
|
|
session.get_credentials = MagicMock(return_value=creds)
|
|
return session
|
|
|
|
|
|
def _raising_session(exc=RuntimeError("no creds")):
|
|
session = MagicMock()
|
|
session.get_credentials = MagicMock(side_effect=exc)
|
|
return session
|
|
|
|
|
|
def _async_session(access_key="AKIA_AIO", secret_key="s", token="t"):
|
|
"""Build an aiobotocore-style session returning a credentials coroutine."""
|
|
|
|
async def _frozen():
|
|
return MagicMock(access_key=access_key, secret_key=secret_key, token=token)
|
|
|
|
async def _creds():
|
|
return MagicMock(
|
|
get_frozen_credentials=MagicMock(return_value=_frozen()),
|
|
expiry_time=None,
|
|
)
|
|
|
|
session = MagicMock()
|
|
session.get_credentials = MagicMock(return_value=_creds())
|
|
return session
|
|
|
|
|
|
def _wrap_s3fs(session=None, storage_options=None, anon=False):
|
|
"""Build a PyFileSystem(FSSpecHandler(stub)) backed by an s3fs.S3FileSystem mock.
|
|
|
|
``_is_fsspec_s3_filesystem`` is strict about the inner class being
|
|
``s3fs.S3FileSystem``, so we use ``spec=`` to ensure the mock passes
|
|
``isinstance(inner, S3FileSystem)``.
|
|
"""
|
|
s3fs = pytest.importorskip("s3fs")
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
inner = MagicMock(spec=s3fs.S3FileSystem)
|
|
inner.protocol = "s3"
|
|
inner.storage_options = storage_options or {}
|
|
inner.session = session
|
|
inner._session = None
|
|
inner.key = None
|
|
inner.secret = None
|
|
inner.token = None
|
|
inner.endpoint_url = None
|
|
inner.client_kwargs = None
|
|
inner.anon = anon
|
|
inner._strip_protocol = lambda p: p.split("://", 1)[-1] if "://" in p else p
|
|
return PyFileSystem(FSSpecHandler(inner))
|
|
|
|
|
|
class TestS3FSSessionCredentialProvider:
|
|
"""``_S3FSSessionCredentialProvider`` — obstore callback for session-backed s3fs."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"session_factory,expected_key",
|
|
[
|
|
# Botocore returns plain Credentials objects.
|
|
pytest.param(lambda: _sync_session("AKIA"), "AKIA", id="sync-botocore"),
|
|
# aiobotocore returns coroutines from get_credentials and
|
|
# get_frozen_credentials; the provider drives either via
|
|
# ``inspect.isawaitable``, which also covers Tasks / Futures / custom
|
|
# awaitables, not just bare coroutines.
|
|
pytest.param(_async_session, "AKIA_AIO", id="async-aiobotocore"),
|
|
],
|
|
)
|
|
def test_returns_obstore_dict(self, session_factory, expected_key):
|
|
provider = _S3FSSessionCredentialProvider(session_factory())
|
|
result = asyncio.run(provider())
|
|
assert result["access_key_id"] == expected_key
|
|
assert isinstance(result["expires_at"], datetime)
|
|
|
|
def test_caches_until_expiry(self):
|
|
# Repeated calls within the TTL must not re-enter the session — obstore
|
|
# may invoke the provider on every request and the session can be
|
|
# expensive (HTTP / token refresh).
|
|
session = _sync_session("AKIA")
|
|
provider = _S3FSSessionCredentialProvider(session)
|
|
first = asyncio.run(provider())
|
|
second = asyncio.run(provider())
|
|
assert session.get_credentials.call_count == 1
|
|
assert second == first
|
|
|
|
def test_refreshes_after_expiry(self):
|
|
# Once expires_at passes, the next call must re-enter the session so
|
|
# rotated keys propagate to obstore. Use a tiny TTL to test quickly.
|
|
session = _sync_session("AKIA")
|
|
provider = _S3FSSessionCredentialProvider(
|
|
session, ttl=timedelta(microseconds=1)
|
|
)
|
|
asyncio.run(provider())
|
|
time.sleep(0.005)
|
|
asyncio.run(provider())
|
|
assert session.get_credentials.call_count == 2
|
|
|
|
def test_uses_session_expiry_time(self):
|
|
# When the session exposes its own expiry_time (typical for STS /
|
|
# AssumeRole), use that instead of TTL — obstore will refresh on
|
|
# the real expiry, not a wall-clock window.
|
|
future = datetime.now(timezone.utc) + timedelta(hours=2)
|
|
provider = _S3FSSessionCredentialProvider(
|
|
_sync_session("AKIA", expiry_time=future)
|
|
)
|
|
assert asyncio.run(provider())["expires_at"] == future
|
|
|
|
@pytest.mark.parametrize("attr", ["session", "_session"])
|
|
def test_from_fsspec_fs_picks_session_or_underscore(self, attr):
|
|
# Modern s3fs uses ``session``; older versions used ``_session``.
|
|
# Both must be picked up.
|
|
fs = MagicMock(session=None, _session=None)
|
|
setattr(fs, attr, _sync_session("AKIA"))
|
|
provider = _S3FSSessionCredentialProvider.from_fsspec_fs(fs)
|
|
assert provider is not None
|
|
assert asyncio.run(provider())["access_key_id"] == "AKIA"
|
|
|
|
def test_from_fsspec_fs_returns_none_when_no_session(self):
|
|
fs = MagicMock(session=None, _session=None)
|
|
assert _S3FSSessionCredentialProvider.from_fsspec_fs(fs) is None
|
|
|
|
@pytest.mark.parametrize(
|
|
"session_factory,expected",
|
|
[
|
|
pytest.param(lambda: _sync_session("AKIA"), True, id="success"),
|
|
pytest.param(_raising_session, False, id="session-raises"),
|
|
pytest.param(
|
|
lambda: _sync_session(None, None, None),
|
|
False,
|
|
id="empty-access-key",
|
|
),
|
|
],
|
|
)
|
|
def test_can_fetch_credentials(self, session_factory, expected):
|
|
provider = _S3FSSessionCredentialProvider(session_factory())
|
|
assert provider.can_fetch_credentials() is expected
|
|
|
|
|
|
class TestExtractCredentialsFromFilesystemFsspecSession:
|
|
"""``_extract_credentials_from_filesystem`` installs the provider for fsspec-S3."""
|
|
|
|
def test_installs_provider_for_session_backed_fs(self):
|
|
# The target bug: fsspec s3fs with Okta has static attrs None but
|
|
# credentials resolvable via session. Extraction must install a
|
|
# credential_provider so obstore can refresh on expiry — not snapshot
|
|
# the keys into static kwargs, which would go stale mid-job.
|
|
wrapped = _wrap_s3fs(_sync_session("AKIA_STS"))
|
|
result = _extract_credentials_from_filesystem(wrapped)
|
|
assert result is not None
|
|
assert isinstance(result["credential_provider"], _S3FSSessionCredentialProvider)
|
|
# Static keys are NOT in kwargs — they come from the provider on demand.
|
|
assert "access_key_id" not in result
|
|
assert "secret_access_key" not in result
|
|
|
|
def test_static_keys_skip_provider(self):
|
|
# Explicit static keys win — no provider needed and the session must
|
|
# not be entered (avoids accidental session calls for users who
|
|
# provided keys directly).
|
|
session = _sync_session("AKIA_SESSION")
|
|
wrapped = _wrap_s3fs(
|
|
session, storage_options={"key": "AKIA_EXPLICIT", "secret": "e"}
|
|
)
|
|
result = _extract_credentials_from_filesystem(wrapped)
|
|
assert result is not None
|
|
assert result["access_key_id"] == "AKIA_EXPLICIT"
|
|
assert result["secret_access_key"] == "e"
|
|
assert "credential_provider" not in result
|
|
session.get_credentials.assert_not_called()
|
|
|
|
def test_anon_skips_provider(self):
|
|
# Anonymous → no creds needed; session must not be invoked even if
|
|
# it would raise.
|
|
session = _raising_session(AssertionError("must not be called"))
|
|
wrapped = _wrap_s3fs(session, anon=True)
|
|
assert _extract_credentials_from_filesystem(wrapped) == {"skip_signature": True}
|
|
|
|
def test_unresolvable_session_returns_none(self):
|
|
# Session present but get_credentials raises — provider validation
|
|
# fails, and we return None so the planner routes to the threaded
|
|
# path (keeping the user's filesystem authoritative).
|
|
wrapped = _wrap_s3fs(_raising_session())
|
|
assert _extract_credentials_from_filesystem(wrapped) is None
|
|
|
|
|
|
class TestPlanObstoreRouting:
|
|
"""``_plan_obstore_routing`` picks obstore vs threaded with one-shot warning."""
|
|
|
|
def setup_method(self):
|
|
# Clear the module-level dedup so each test sees a fresh warning.
|
|
from ray.data._internal.planner import _obstore_download as m
|
|
|
|
m._warned_credential_fs_ids.clear()
|
|
|
|
def test_none_fs_uses_obstore(self):
|
|
assert _plan_obstore_routing(None) == (True, {})
|
|
|
|
def test_extractable_fsspec_s3_uses_obstore(self):
|
|
wrapped = _wrap_s3fs(
|
|
storage_options={"key": "AKIA", "secret": "sec", "token": "tok"}
|
|
)
|
|
assert _plan_obstore_routing(wrapped) == (
|
|
True,
|
|
{
|
|
"access_key_id": "AKIA",
|
|
"secret_access_key": "sec",
|
|
"session_token": "tok",
|
|
},
|
|
)
|
|
|
|
def test_fsspec_s3_unextractable_warns_once(self):
|
|
# Patch module logger directly — Ray's internal loggers don't always
|
|
# propagate to pytest's caplog in sandboxed test runners. This also
|
|
# covers the dedup path: same FS twice = one warning, different FS =
|
|
# new warning.
|
|
fs_a = _wrap_s3fs(_raising_session())
|
|
fs_b = _wrap_s3fs(_raising_session())
|
|
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download.logger"
|
|
) as mock_logger:
|
|
assert _plan_obstore_routing(fs_a) == (False, {})
|
|
_plan_obstore_routing(fs_a) # dedup — no new warning
|
|
_plan_obstore_routing(fs_b) # different FS — new warning
|
|
|
|
assert mock_logger.warning.call_count == 2
|
|
msg = mock_logger.warning.call_args_list[0][0][0]
|
|
assert "S3 credentials" in msg and "storage_options" in msg
|
|
|
|
def test_silent_routing_for_unrecognized_fs(self):
|
|
# Both non-S3 fsspec and native unrecognized filesystems must route
|
|
# to the threaded path WITHOUT a S3-specific warning. The hardcoded
|
|
# "pass credentials via fsspec storage_options" advice only makes
|
|
# sense for fsspec-S3.
|
|
import fsspec
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
class _GcsStub(fsspec.AbstractFileSystem):
|
|
protocol = "gcs"
|
|
|
|
filesystems = [
|
|
PyFileSystem(FSSpecHandler(_GcsStub())),
|
|
pafs.LocalFileSystem(),
|
|
]
|
|
for fs in filesystems:
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download.logger"
|
|
) as mock_logger:
|
|
assert _plan_obstore_routing(fs) == (False, {})
|
|
mock_logger.warning.assert_not_called()
|
|
|
|
def test_async_partition_actor_fails_closed(self):
|
|
pytest.importorskip("obstore")
|
|
from ray.data._internal.planner.download_partition_actor import (
|
|
AsyncPartitionActor,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="cannot be statically extracted"):
|
|
AsyncPartitionActor(
|
|
["uri"], DataContext.get_current(), filesystem=pafs.LocalFileSystem()
|
|
)
|
|
|
|
def test_download_uris_fails_closed(self):
|
|
pytest.importorskip("obstore")
|
|
with pytest.raises(RuntimeError, match="cannot be statically extracted"):
|
|
asyncio.run(
|
|
_download_uris_with_obstore(
|
|
["file:///tmp/x.bin"], "uri", filesystem=pafs.LocalFileSystem()
|
|
)
|
|
)
|
|
|
|
|
|
# TestThreadedDownloadPreResolve
|
|
def _spy_resolve(fake_fn=None):
|
|
"""Patch _resolve_paths_and_filesystem, returning (patch_cm, probe_calls,
|
|
normalize_calls). Probe calls pass filesystem=None; normalize calls supply
|
|
one. ``fake_fn`` replaces the real function when given."""
|
|
from ray.data._internal.planner import plan_download_op as pdo
|
|
|
|
probe_calls: list = []
|
|
normalize_calls: list = []
|
|
original = pdo._resolve_paths_and_filesystem
|
|
|
|
def _tracking(uri, filesystem=None, **kw):
|
|
(probe_calls if filesystem is None else normalize_calls).append(uri)
|
|
if fake_fn is not None:
|
|
return fake_fn(uri, filesystem=filesystem, **kw)
|
|
return original(uri, filesystem=filesystem, **kw)
|
|
|
|
return (
|
|
patch.object(pdo, "_resolve_paths_and_filesystem", side_effect=_tracking),
|
|
probe_calls,
|
|
normalize_calls,
|
|
)
|
|
|
|
|
|
class TestThreadedDownloadPreResolve:
|
|
"""``download_bytes_threaded`` probes the filesystem once and shares it
|
|
across workers, instead of each worker inferring it (the IMDS herd)."""
|
|
|
|
def test_probe_once_across_workers(self, tmp_path):
|
|
for i in range(10):
|
|
(tmp_path / f"f{i}.bin").write_bytes(f"data-{i}".encode())
|
|
uris = [f"file://{tmp_path}/f{i}.bin" for i in range(10)]
|
|
table = pa.Table.from_arrays([pa.array(uris)], names=["uri"])
|
|
|
|
spy, probes, normalizes = _spy_resolve()
|
|
with spy:
|
|
results = list(
|
|
download_bytes_threaded(
|
|
table, ["uri"], ["bytes"], DataContext.get_current()
|
|
)
|
|
)
|
|
|
|
# One probe regardless of worker count; one normalize per URI.
|
|
assert len(probes) == 1
|
|
assert len(normalizes) == 10
|
|
assert [b.as_py() for b in results[0].column("bytes")] == [
|
|
f"data-{i}".encode() for i in range(10)
|
|
]
|
|
|
|
@pytest.mark.parametrize(
|
|
"fs_factory",
|
|
[
|
|
pytest.param(lambda: pafs.LocalFileSystem(), id="pyarrow-local"),
|
|
# fsspec FS lacks open_input_stream and must be normalized before
|
|
# wrapping, else reads raise AttributeError mid-flight.
|
|
pytest.param(
|
|
lambda: pytest.importorskip("fsspec").filesystem("file"),
|
|
id="fsspec-local",
|
|
),
|
|
],
|
|
)
|
|
def test_supplied_fs_skips_probe(self, tmp_path, fs_factory):
|
|
(tmp_path / "f.bin").write_bytes(b"supplied")
|
|
table = pa.Table.from_arrays(
|
|
[pa.array([f"file://{tmp_path}/f.bin"])], names=["uri"]
|
|
)
|
|
|
|
spy, probes, _ = _spy_resolve()
|
|
with spy:
|
|
results = list(
|
|
download_bytes_threaded(
|
|
table,
|
|
["uri"],
|
|
["bytes"],
|
|
DataContext.get_current(),
|
|
filesystem=fs_factory(),
|
|
)
|
|
)
|
|
|
|
assert probes == []
|
|
assert results[0].column("bytes")[0].as_py() == b"supplied"
|
|
|
|
@pytest.mark.parametrize(
|
|
"first_uri,probe_returns_for_first",
|
|
[
|
|
pytest.param(None, None, id="none-in-list"),
|
|
pytest.param("file:///first-dropped", ([], None), id="empty-paths-result"),
|
|
],
|
|
)
|
|
def test_probe_loop_skips_unusable(
|
|
self, tmp_path, first_uri, probe_returns_for_first
|
|
):
|
|
# Probe must skip None URIs and ([], fs)/(paths, None) returns rather
|
|
# than break early with no usable filesystem.
|
|
from ray.data._internal.planner import plan_download_op as pdo
|
|
|
|
(tmp_path / "good.bin").write_bytes(b"good")
|
|
real_fs = pafs.LocalFileSystem()
|
|
|
|
def _fake(uri, filesystem=None, **_kw):
|
|
if filesystem is not None:
|
|
# Normalize-only: pass paths through when FS is supplied.
|
|
return ([str(tmp_path / "good.bin")], filesystem)
|
|
if uri == first_uri:
|
|
return probe_returns_for_first
|
|
return ([str(tmp_path / "good.bin")], real_fs)
|
|
|
|
good = f"file://{tmp_path}/good.bin"
|
|
table = pa.Table.from_arrays([pa.array([first_uri, good])], names=["uri"])
|
|
|
|
with patch.object(pdo, "_resolve_paths_and_filesystem", side_effect=_fake):
|
|
results = list(
|
|
download_bytes_threaded(
|
|
table, ["uri"], ["bytes"], DataContext.get_current()
|
|
)
|
|
)
|
|
assert results[0].column("bytes")[-1].as_py() == b"good"
|
|
|
|
def test_all_probes_fail_yields_none(self):
|
|
# No URI resolves; we short-circuit to None for every row instead of
|
|
# letting workers repeat the failed inference.
|
|
from ray.data._internal.planner import plan_download_op as pdo
|
|
|
|
resolve_calls: list = []
|
|
|
|
def _fail_all(uri, filesystem=None, **_kw):
|
|
resolve_calls.append((uri, filesystem))
|
|
raise RuntimeError("nothing resolves")
|
|
|
|
uris = ["bad-1", "bad-2", "bad-3"]
|
|
table = pa.Table.from_arrays([pa.array(uris)], names=["uri"])
|
|
with patch.object(pdo, "_resolve_paths_and_filesystem", side_effect=_fail_all):
|
|
results = list(
|
|
download_bytes_threaded(
|
|
table, ["uri"], ["bytes"], DataContext.get_current()
|
|
)
|
|
)
|
|
|
|
# Probe tried each URI exactly once. No worker calls.
|
|
assert len(resolve_calls) == len(uris)
|
|
assert all(fs is None for _, fs in resolve_calls)
|
|
# Every URI yields None.
|
|
assert [b.as_py() for b in results[0].column("bytes")] == [None] * len(uris)
|
|
|
|
def test_empty_block_no_workers(self):
|
|
# Zero URIs → the column loop skips; no probe, no workers spawned.
|
|
table = pa.Table.from_arrays([pa.array([], type=pa.string())], names=["uri"])
|
|
from ray.data._internal.planner import plan_download_op as pdo
|
|
|
|
with patch.object(
|
|
pdo, "_resolve_paths_and_filesystem", side_effect=AssertionError
|
|
):
|
|
list(
|
|
download_bytes_threaded(
|
|
table, ["uri"], ["bytes"], DataContext.get_current()
|
|
)
|
|
)
|
|
|
|
|
|
# TestObstoreDownloadPath
|
|
class TestObstoreDownloadPath:
|
|
"""Integration tests for the obstore async download path.
|
|
|
|
All tests are skipped when obstore is not installed. Local ``file://``
|
|
URIs are used so no cloud credentials are required.
|
|
"""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def require_obstore(self):
|
|
pytest.importorskip("obstore")
|
|
|
|
def test_empty_uris_returns_empty_list(self):
|
|
# asyncio.gather on an empty task list returns [] — verify this holds
|
|
# end-to-end so the obstore path is consistent with the PyArrow path's
|
|
# len(uris) == 0 early-exit.
|
|
results = asyncio.run(_download_uris_with_obstore([], "uri"))
|
|
assert results == []
|
|
|
|
def test_downloads_files(self, tmp_path):
|
|
content1, content2 = b"hello obstore", b"world obstore"
|
|
(tmp_path / "f1.bin").write_bytes(content1)
|
|
(tmp_path / "f2.bin").write_bytes(content2)
|
|
|
|
uris = [f"file://{tmp_path}/f1.bin", f"file://{tmp_path}/f2.bin"]
|
|
results = asyncio.run(_download_uris_with_obstore(uris, "uri"))
|
|
assert results == [content1, content2]
|
|
|
|
def test_preserves_ordering(self, tmp_path):
|
|
# Verify asyncio.gather returns results in the same order as the
|
|
# input URIs even when downloads complete out of order.
|
|
contents = [f"file-{i}".encode() for i in range(10)]
|
|
for i, data in enumerate(contents):
|
|
(tmp_path / f"f{i}.bin").write_bytes(data)
|
|
|
|
uris = [f"file://{tmp_path}/f{i}.bin" for i in range(10)]
|
|
results = asyncio.run(_download_uris_with_obstore(uris, "uri"))
|
|
assert results == contents
|
|
|
|
def test_missing_file_returns_none(self, tmp_path):
|
|
existing = tmp_path / "exists.bin"
|
|
existing.write_bytes(b"real data")
|
|
|
|
uris = [
|
|
f"file://{tmp_path}/exists.bin",
|
|
f"file://{tmp_path}/does_not_exist.bin",
|
|
]
|
|
results = asyncio.run(_download_uris_with_obstore(uris, "uri"))
|
|
assert results[0] == b"real data"
|
|
assert results[1] is None
|
|
|
|
def test_store_cache_reused_for_same_host(self, tmp_path):
|
|
# Files under the same prefix should share one ObjectStore instance.
|
|
for i in range(3):
|
|
(tmp_path / f"f{i}.bin").write_bytes(f"data{i}".encode())
|
|
|
|
uris = [f"file://{tmp_path}/f{i}.bin" for i in range(3)]
|
|
|
|
# Spy on from_url to count how many store instances are created.
|
|
with patch(
|
|
"obstore.store.from_url",
|
|
wraps=__import__("obstore.store", fromlist=["from_url"]).from_url,
|
|
) as spy:
|
|
asyncio.run(_download_uris_with_obstore(uris, "uri"))
|
|
# All three URIs share the same store_url ("file://"), so
|
|
# from_url should be called exactly once.
|
|
assert spy.call_count == 1
|
|
|
|
def test_download_bytes_async_uses_obstore(self, tmp_path):
|
|
# Verify download_bytes_async routes through obstore for supported schemes.
|
|
content = b"dispatch test content"
|
|
uri = f"file://{tmp_path}/test.bin"
|
|
(tmp_path / "test.bin").write_bytes(content)
|
|
|
|
table = pa.Table.from_arrays([pa.array([uri])], names=["uri"])
|
|
ctx = DataContext.get_current()
|
|
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download._download_uris_with_obstore",
|
|
wraps=_download_uris_with_obstore,
|
|
) as obstore_spy:
|
|
results = list(download_bytes_async(table, ["uri"], ["bytes"], ctx))
|
|
|
|
assert obstore_spy.called
|
|
assert results[0].column("bytes")[0].as_py() == content
|
|
|
|
def test_download_bytes_async_falls_back_for_unsupported_scheme(self, tmp_path):
|
|
# Verify download_bytes_async falls back to download_bytes_threaded
|
|
# for URI schemes obstore doesn't handle (e.g. bare filesystem paths).
|
|
content = b"fallback test content"
|
|
(tmp_path / "test.bin").write_bytes(content)
|
|
uri = str(tmp_path / "test.bin")
|
|
|
|
table = pa.Table.from_arrays([pa.array([uri])], names=["uri"])
|
|
ctx = DataContext.get_current()
|
|
|
|
with patch(
|
|
"ray.data._internal.planner.plan_download_op.download_bytes_threaded",
|
|
wraps=download_bytes_threaded,
|
|
) as threaded_spy:
|
|
results = list(download_bytes_async(table, ["uri"], ["bytes"], ctx))
|
|
|
|
assert threaded_spy.called
|
|
assert results[0].column("bytes")[0].as_py() == content
|
|
|
|
def test_download_bytes_async_non_s3_fsspec_pyfilesystem_threaded_fallback(
|
|
self, tmp_path
|
|
):
|
|
# End-to-end: ``file://`` (obstore-eligible) + non-S3 fsspec PyFileSystem triggers
|
|
# threaded download via the real ``_obstore_filesystem_requires_threaded_download`` gate.
|
|
# Internal size columns are still stripped before ``download_bytes_threaded``.
|
|
import fsspec
|
|
from pyarrow.fs import FSSpecHandler, PyFileSystem
|
|
|
|
class _GcsStub(fsspec.AbstractFileSystem):
|
|
protocol = "gcs"
|
|
|
|
content = b"gcs-fs stub fallback"
|
|
(tmp_path / "obj.bin").write_bytes(content)
|
|
uri = f"file://{tmp_path}/obj.bin"
|
|
size_col = f"{_FILE_SIZE_COLUMN_PREFIX}uri"
|
|
table = pa.Table.from_arrays(
|
|
[pa.array([uri]), pa.array([len(content)], type=pa.int64())],
|
|
names=["uri", size_col],
|
|
)
|
|
ctx = DataContext.get_current()
|
|
wrapped = PyFileSystem(FSSpecHandler(_GcsStub()))
|
|
|
|
with patch(
|
|
"ray.data._internal.planner.plan_download_op.download_bytes_threaded",
|
|
wraps=download_bytes_threaded,
|
|
) as threaded_spy:
|
|
results = list(
|
|
download_bytes_async(table, ["uri"], ["bytes"], ctx, filesystem=wrapped)
|
|
)
|
|
|
|
assert threaded_spy.called
|
|
out = results[0]
|
|
# Null bytes: stub FS cannot open paths normalized for local FS while _resolve_paths_and_filesystem returns the caller filesystem unchanged.
|
|
# Stub: PyFileSystem(FSSpecHandler(_GcsStub)) passed in; local: PyArrow FS for file:// from fallback when that stub is incompatible.
|
|
assert out.column("bytes")[0].as_py() is None
|
|
assert size_col not in out.column_names
|
|
|
|
def test_download_bytes_async_fallback_drops_size_columns(self, tmp_path):
|
|
# Verify that file-size columns attached by AsyncPartitionActor are dropped
|
|
# when falling back to download_bytes_threaded (unsupported scheme).
|
|
content = b"fallback size col test"
|
|
(tmp_path / "test.bin").write_bytes(content)
|
|
uri = str(tmp_path / "test.bin")
|
|
|
|
size_col = f"{_FILE_SIZE_COLUMN_PREFIX}uri"
|
|
table = pa.Table.from_arrays(
|
|
[pa.array([uri]), pa.array([len(content)], type=pa.int64())],
|
|
names=["uri", size_col],
|
|
)
|
|
ctx = DataContext.get_current()
|
|
|
|
results = list(download_bytes_async(table, ["uri"], ["bytes"], ctx))
|
|
|
|
out = results[0]
|
|
assert out.column("bytes")[0].as_py() == content
|
|
assert size_col not in out.column_names
|
|
|
|
|
|
class TestObstoreRangeSplitDownload:
|
|
"""Tests for the range-split download path (get_range_async).
|
|
|
|
Uses local ``file://`` URIs. The range-split logic is activated by
|
|
patching the module-level ``RAY_DATA_OBSTORE_RANGE_THRESHOLD``.
|
|
"""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def require_obstore(self):
|
|
pytest.importorskip("obstore")
|
|
|
|
def test_range_split_downloads_large_file(self, tmp_path):
|
|
# Create a file larger than the chunk size to trigger range splitting.
|
|
chunk_size = 1024
|
|
content = os.urandom(chunk_size * 3 + 500)
|
|
(tmp_path / "big.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/big.bin"
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size * 2,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[len(content)])
|
|
)
|
|
assert results == [content]
|
|
|
|
def test_range_split_with_unknown_size_falls_back_to_head(self, tmp_path):
|
|
# file_sizes=None → should HEAD to discover size, then range-split.
|
|
chunk_size = 512
|
|
content = os.urandom(chunk_size * 4)
|
|
(tmp_path / "big2.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/big2.bin"
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=None)
|
|
)
|
|
assert results == [content]
|
|
|
|
def test_range_split_disabled_uses_whole_file_get(self, tmp_path):
|
|
"""No ranged download when file is below threshold or chunk size is invalid."""
|
|
# Files below threshold should still use get_async, not range requests.
|
|
content = b"small file"
|
|
(tmp_path / "small.bin").write_bytes(content)
|
|
uri = f"file://{tmp_path}/small.bin"
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
1024 * 1024,
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[len(content)])
|
|
)
|
|
assert results == [content]
|
|
|
|
# Non-positive chunk size disables range split.
|
|
content2 = os.urandom(3000)
|
|
(tmp_path / "w.bin").write_bytes(content2)
|
|
uri2 = f"file://{tmp_path}/w.bin"
|
|
|
|
def _fail_ranged(*_args, **_kwargs):
|
|
raise AssertionError("_fetch_ranged must not be called")
|
|
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
100,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
-512,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download._fetch_ranged",
|
|
side_effect=_fail_ranged,
|
|
),
|
|
):
|
|
results2 = asyncio.run(
|
|
_download_uris_with_obstore([uri2], "uri", file_sizes=[len(content2)])
|
|
)
|
|
assert results2 == [content2]
|
|
|
|
def test_range_split_mixed_sizes(self, tmp_path):
|
|
# Mix of large and small files in a single batch.
|
|
chunk_size = 256
|
|
large_content = os.urandom(chunk_size * 5)
|
|
small_content = b"tiny"
|
|
(tmp_path / "large.bin").write_bytes(large_content)
|
|
(tmp_path / "small.bin").write_bytes(small_content)
|
|
|
|
uris = [
|
|
f"file://{tmp_path}/large.bin",
|
|
f"file://{tmp_path}/small.bin",
|
|
]
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size * 2,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore(
|
|
uris,
|
|
"uri",
|
|
file_sizes=[len(large_content), len(small_content)],
|
|
)
|
|
)
|
|
assert results == [large_content, small_content]
|
|
|
|
def test_threshold_zero_disables_range_split(self, tmp_path):
|
|
# Default threshold=0 should use the simple path even for large files.
|
|
content = os.urandom(10000)
|
|
(tmp_path / "f.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/f.bin"
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
0,
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[len(content)])
|
|
)
|
|
assert results == [content]
|
|
|
|
def test_download_bytes_async_reads_and_drops_size_column(self, tmp_path):
|
|
# Verify that download_bytes_async reads the hidden size column
|
|
# from AsyncPartitionActor, passes it through, and drops it from output.
|
|
content = b"test content for size column"
|
|
(tmp_path / "test.bin").write_bytes(content)
|
|
uri = f"file://{tmp_path}/test.bin"
|
|
|
|
size_col = f"{_FILE_SIZE_COLUMN_PREFIX}uri"
|
|
table = pa.Table.from_arrays(
|
|
[pa.array([uri]), pa.array([len(content)], type=pa.int64())],
|
|
names=["uri", size_col],
|
|
)
|
|
ctx = DataContext.get_current()
|
|
|
|
results = list(download_bytes_async(table, ["uri"], ["bytes"], ctx))
|
|
|
|
out = results[0]
|
|
assert out.column("bytes")[0].as_py() == content
|
|
assert size_col not in out.column_names
|
|
|
|
def test_http_auto_detect_enables_allow_http(self):
|
|
# Verify that http:// URIs actually download successfully when
|
|
# allow_http is auto-enabled, and that a warning is logged.
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
from threading import Thread
|
|
|
|
content = b"http download works"
|
|
|
|
class _Handler(BaseHTTPRequestHandler):
|
|
def do_GET(self):
|
|
self.send_response(200)
|
|
self.send_header("Content-Length", str(len(content)))
|
|
self.end_headers()
|
|
self.wfile.write(content)
|
|
|
|
def log_message(self, format, *args):
|
|
pass
|
|
|
|
server = HTTPServer(("127.0.0.1", 0), _Handler)
|
|
port = server.server_address[1]
|
|
t = Thread(target=server.serve_forever, daemon=True)
|
|
t.start()
|
|
|
|
http_uri = f"http://127.0.0.1:{port}/test.bin"
|
|
try:
|
|
with patch(
|
|
"ray.data._internal.planner._obstore_download.logger"
|
|
) as mock_logger:
|
|
results = asyncio.run(_download_uris_with_obstore([http_uri], "uri"))
|
|
|
|
assert results == [
|
|
content
|
|
], "http:// download should succeed when allow_http is auto-enabled"
|
|
mock_logger.warning.assert_any_call(
|
|
"Downloading over unencrypted HTTP. Consider using https:// instead."
|
|
)
|
|
finally:
|
|
server.shutdown()
|
|
|
|
def test_no_allow_http_for_non_http_uris(self, tmp_path):
|
|
# Verify that allow_http is NOT set when no http:// URIs are present.
|
|
import obstore.store as obs_store
|
|
|
|
content = b"no http"
|
|
(tmp_path / "f.bin").write_bytes(content)
|
|
uri = f"file://{tmp_path}/f.bin"
|
|
|
|
with patch.object(obs_store, "from_url", wraps=obs_store.from_url) as spy:
|
|
results = asyncio.run(_download_uris_with_obstore([uri], "uri"))
|
|
|
|
for call in spy.call_args_list:
|
|
client_opts = call.kwargs.get("client_options", {})
|
|
assert "allow_http" not in client_opts
|
|
assert results == [content]
|
|
|
|
def test_head_failure_falls_back_to_simple_get(self, tmp_path):
|
|
# When HEAD fails (e.g., server error), the size stays 0 and the
|
|
# file should be downloaded via simple GET instead of crashing.
|
|
chunk_size = 256
|
|
content = os.urandom(chunk_size * 10)
|
|
(tmp_path / "f.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/f.bin"
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
import obstore as obs
|
|
|
|
async def _failing_head(*args, **kwargs):
|
|
raise OSError("HEAD not supported")
|
|
|
|
original_get_range = obs.get_range_async
|
|
range_calls = []
|
|
|
|
async def _tracking_range(*args, **kwargs):
|
|
range_calls.append(args)
|
|
return await original_get_range(*args, **kwargs)
|
|
|
|
with (
|
|
patch.object(obs, "head_async", side_effect=_failing_head),
|
|
patch.object(obs, "get_range_async", side_effect=_tracking_range),
|
|
):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[0])
|
|
)
|
|
|
|
assert results == [content]
|
|
# Explicitly verify ranged path was NOT taken.
|
|
assert len(range_calls) == 0
|
|
|
|
def test_partial_unknown_sizes_head_only_for_unknowns(self, tmp_path):
|
|
# In a batch where some sizes are known and others are 0,
|
|
# HEAD should only be issued for the unknowns.
|
|
chunk_size = 256
|
|
large_content = os.urandom(chunk_size * 5)
|
|
small_content = b"small"
|
|
unknown_content = os.urandom(chunk_size * 8)
|
|
|
|
(tmp_path / "large.bin").write_bytes(large_content)
|
|
(tmp_path / "small.bin").write_bytes(small_content)
|
|
(tmp_path / "unknown.bin").write_bytes(unknown_content)
|
|
|
|
uris = [
|
|
f"file://{tmp_path}/large.bin",
|
|
f"file://{tmp_path}/small.bin",
|
|
f"file://{tmp_path}/unknown.bin",
|
|
]
|
|
file_sizes: list[Optional[int]] = [len(large_content), len(small_content), 0]
|
|
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size * 2,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
import obstore as obs
|
|
|
|
original_head = obs.head_async
|
|
head_calls = []
|
|
|
|
async def _tracking_head(*args, **kwargs):
|
|
head_calls.append(args)
|
|
return await original_head(*args, **kwargs)
|
|
|
|
with patch.object(obs, "head_async", side_effect=_tracking_head):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore(uris, "uri", file_sizes=file_sizes)
|
|
)
|
|
|
|
# Only the unknown (index 2) should trigger a HEAD request.
|
|
assert len(head_calls) == 1
|
|
assert results == [large_content, small_content, unknown_content]
|
|
|
|
def test_file_size_at_exact_threshold_uses_simple_get(self, tmp_path):
|
|
# size == threshold should use simple GET (check is size > threshold).
|
|
threshold = 1024
|
|
content = os.urandom(threshold)
|
|
(tmp_path / "exact.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/exact.bin"
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
threshold,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
256,
|
|
),
|
|
):
|
|
import obstore as obs
|
|
|
|
range_calls = []
|
|
original_get_range = obs.get_range_async
|
|
|
|
async def _tracking_range(*args, **kwargs):
|
|
range_calls.append(args)
|
|
return await original_get_range(*args, **kwargs)
|
|
|
|
with patch.object(obs, "get_range_async", side_effect=_tracking_range):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[threshold])
|
|
)
|
|
|
|
assert len(range_calls) == 0
|
|
assert results == [content]
|
|
|
|
def test_ranged_download_failure_returns_none(self, tmp_path):
|
|
# If a range request fails mid-transfer, _fetch_ranged should
|
|
# catch the exception and fall back to a simple GET.
|
|
chunk_size = 256
|
|
content = os.urandom(chunk_size * 4)
|
|
(tmp_path / "fail.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/fail.bin"
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_CHUNK_SIZE",
|
|
chunk_size,
|
|
),
|
|
):
|
|
import obstore as obs
|
|
|
|
call_count = 0
|
|
original_get_range = obs.get_range_async
|
|
|
|
async def _failing_range(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 2:
|
|
raise OSError("simulated network failure")
|
|
return await original_get_range(*args, **kwargs)
|
|
|
|
with patch.object(obs, "get_range_async", side_effect=_failing_range):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[len(content)])
|
|
)
|
|
|
|
# Ranged failed, but simple GET fallback should succeed.
|
|
assert results == [content]
|
|
|
|
def test_invalid_max_concurrency_disables_range_split(self, tmp_path):
|
|
# max_conc <= 0 with range splitting enabled is a misconfiguration.
|
|
# Verify it warns and falls back to simple GET (no crash).
|
|
chunk_size = 256
|
|
content = os.urandom(chunk_size * 4)
|
|
(tmp_path / "f.bin").write_bytes(content)
|
|
|
|
uri = f"file://{tmp_path}/f.bin"
|
|
# Simulate misconfiguration: range splitting on, but concurrency = 0.
|
|
with (
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_RANGE_THRESHOLD",
|
|
chunk_size,
|
|
),
|
|
patch(
|
|
"ray.data._internal.planner._obstore_download.RAY_DATA_OBSTORE_MAX_CONCURRENCY",
|
|
0,
|
|
),
|
|
patch("ray.data._internal.planner._obstore_download.logger") as mock_logger,
|
|
):
|
|
import obstore as obs
|
|
|
|
# Spy on get_range_async to verify it is never called.
|
|
range_calls = []
|
|
original_get_range = obs.get_range_async
|
|
|
|
async def _tracking_range(*args, **kwargs):
|
|
range_calls.append(args)
|
|
return await original_get_range(*args, **kwargs)
|
|
|
|
with patch.object(obs, "get_range_async", side_effect=_tracking_range):
|
|
results = asyncio.run(
|
|
_download_uris_with_obstore([uri], "uri", file_sizes=[len(content)])
|
|
)
|
|
|
|
# Download should succeed via simple GET fallback.
|
|
assert results == [content]
|
|
# Ranged path must not have been entered.
|
|
assert len(range_calls) == 0, "Range splitting should be disabled"
|
|
# Warning about the invalid configuration should be logged.
|
|
mock_logger.warning.assert_any_call(
|
|
"RAY_DATA_OBSTORE_RANGE_THRESHOLD is set but "
|
|
"RAY_DATA_OBSTORE_MAX_CONCURRENCY=%d is invalid. "
|
|
"Range downloads require a positive concurrency limit to avoid "
|
|
"socket exhaustion. Disabling range splitting.",
|
|
0,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|