[Data] Pre-resolve filesystem in threaded download to avoid IMDS herd

[xyuzh]

Summary

download_bytes_threaded spawns 16 worker threads via make_async_gen. Each worker called _resolve_paths_and_filesystem independently on its first URI. When no filesystem is provided, that call constructs a fresh pyarrow.fs.S3FileSystem, which triggers an IMDS credential fetch on EC2. With N concurrent Download tasks on a node that was ~16×N near-simultaneous IMDS HTTP calls — enough to trip the per-instance rate limit and produce intermittent NoCredentialsError.

This is the second of two independent PRs. The matching fsspec credential-extraction fix — which triggers the threaded fallback whenever a session-backed filesystem can't have credentials statically extracted — is in a separate PR. Either can land first; this change is valuable on its own because the threaded path is also used when obstore is disabled (RAY_DATA_USE_OBSTORE=0) or when the URI scheme isn't obstore-supported.

Changes

• Resolve the filesystem exactly once per (block, column) before make_async_gen, then share it across all 16 workers via default-arg binding on the inner load_uri_bytes closure. Workers now do zero filesystem inference — their per-URI calls to _resolve_paths_and_filesystem short-circuit (path normalization only; no network) because a filesystem is supplied.
• The probe loop:
  • Skips None URIs.
  • Catches resolution exceptions and continues to the next URI.
  • Requires both paths and a non-None filesystem before breaking — otherwise _resolve_paths_and_filesystem could silently drop a bad URI and leave workers with no usable filesystem, recreating the herd.
• If every probe fails, workers retain the legacy per-URI resolution as a last resort so a block of partially-unresolvable URIs doesn't lose the ones that might still resolve.

Test plan

• pre-commit run passes.
• New TestThreadedDownloadPreResolve covers:
  • test_resolves_filesystem_once_across_workers — spy on _resolve_paths_and_filesystem, 10 URIs × 16 workers → exactly 1 probe call with filesystem=None (vs. up to 16 before the fix).
  • test_supplied_filesystem_skips_probe — with a user-supplied FS, zero probe calls happen.
  • test_bad_first_uri_still_finds_filesystem — mix of None + resolvable URIs; probe keeps going past None.
  • test_all_probe_uris_unresolvable_does_not_crash — all URIs unresolvable → no crash, every result is None.
  • test_empty_block_does_not_spawn_workers — zero URIs → no probe, no workers.
  • test_probe_rejects_empty_paths_result — regression test for the ([], fs) silent-drop case: the probe must skip it and continue.
• End-to-end under concurrent tasks with RAY_DATA_USE_OBSTORE=0, verifying pyarrow.fs.S3FileSystem.__init__ is invoked exactly once per block rather than per worker.

[gemini-code-assist]

Code Review

This pull request optimizes the download_bytes_threaded function by pre-resolving the filesystem once per block/column. This change prevents the "IMDS thundering herd" issue where multiple worker threads would independently trigger credential fetches. The PR also introduces a suite of unit tests to verify the probe logic, including edge cases like unresolvable URIs and user-supplied filesystems. Feedback was provided to improve the efficiency of the fallback mechanism: if the initial probe fails, the worker should cache the filesystem locally once it is successfully resolved for a URI to avoid redundant resolutions for subsequent items in that worker's iterator.

[gemini-code-assist]

In the fallback case where wrapped_fs is None (because the probe failed for all URIs in the block), the current implementation performs filesystem resolution and creates a new RetryingPyFileSystem wrapper for every single URI in the worker. This is less efficient than the legacy behavior which cached the resolved filesystem for subsequent URIs in the same worker.

Updating the local wrapped_fs and resolved_fs variables within the loop will ensure that once a filesystem is successfully resolved for one URI, it is reused for all remaining URIs in that worker's iterator, truly mirroring the legacy lazy-resolution behavior.

Suggested change

	if wrapped_fs is None:
	# All probe URIs failed — last-resort per-URI resolution.
	# Mirrors legacy behavior so a block full of unresolvable
	# URIs doesn't lose the ones that might still resolve.
	resolved_paths, per_uri_fs = _resolve_paths_and_filesystem(
	uri, filesystem=None
	)
	fs_for_read = RetryingPyFileSystem.wrap(
	per_uri_fs,
	retryable_errors=data_context.retried_io_errors,
	)
	else:
	# Path normalization only — _resolve_paths_and_filesystem
	# short-circuits when filesystem is supplied (no network).
	resolved_paths, _ = _resolve_paths_and_filesystem(
	uri, filesystem=resolved_fs
	)
	fs_for_read = wrapped_fs
	if wrapped_fs is None:
	# All probe URIs failed — last-resort per-URI resolution.
	# Mirrors legacy behavior so a block full of unresolvable
	# URIs doesn't lose the ones that might still resolve.
	resolved_paths, per_uri_fs = _resolve_paths_and_filesystem(
	uri, filesystem=None
	)
	resolved_fs = per_uri_fs
	wrapped_fs = RetryingPyFileSystem.wrap(
	per_uri_fs,
	retryable_errors=data_context.retried_io_errors,
	)
	fs_for_read = wrapped_fs
	else:
	# Path normalization only — _resolve_paths_and_filesystem
	# short-circuits when filesystem is supplied (no network).
	resolved_paths, _ = _resolve_paths_and_filesystem(
	uri, filesystem=resolved_fs
	)
	fs_for_read = wrapped_fs

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{Reviewed by Cursor Bugbot for commit 8fa637a. Configure here.}

[xyuzh]

@goutamvenkat-anyscale

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.