[Data] Two-mode block metadata fetch behind a single MetadataFetcher interface

[xinyuangui2]

Why

At scale the streaming scheduler blocks on ray.get(meta_ref) for every
block-metadata pair inside process_completed_tasks — the largest avoidable
cost in the scheduling loop. This moves that fetch off the scheduling thread,
behind a single interface with two explicit modes so the old synchronous path
stays available unchanged.

Flow

DataOpTask.on_data_ready owns one shared pull loop; the per-pair fetch/emit
and the completion signal are delegated to the MetadataFetcher, so the two
modes differ only behind the interface (not in the loop).

flowchart TB
    subgraph SCHED["process_completed_tasks — scheduling thread (per iteration)"]
      ODR["DataOpTask.on_data_ready()<br/>shared pull loop: pull (block_ref, meta_ref)"]
      IGS["fetcher.in_loop_get_size(pair) → budget size"]
      DONE["on drain: fetcher.in_loop_done(task)"]
      SUB["after per-op loop:<br/>fetcher.submit(op) + register_if_drained()"]
      ALB["end of call: fetcher.after_loop_batch()"]
      ODR --> IGS --> DONE --> SUB --> ALB
    end

    IGS -->|inline| INL["InlineMetadataFetcher<br/>ray.get(meta_ref) + emit now<br/>budget = meta.size_bytes"]
    DONE -->|inline| INLD["fire done-callback now<br/>(re-raise task failure)"]

    IGS -->|threaded| THR["ThreadedMetadataFetcher<br/>defer pair; budget = local object_size (no RPC)"]
    DONE -->|threaded| THRD["no-op (completion postponed)"]
    SUB -->|threaded| PF[["ThreadedMetadataFetcher background thread<br/>ray.wait(fetch_local=True) → ray.get(meta_refs)"]]
    PF -. publish bytes .-> ALBT
    ALB -->|threaded| ALBT["emit ready pairs (per-op FIFO order)<br/>+ fire postponed done-callbacks"]

Loading

What

• New MetadataFetcher interface (metadata_fetcher.py), selected by
  RAY_DATA_METADATA_PREFETCH_ON_THREAD (default: threaded):
  • InlineMetadataFetcher — synchronous: fetch metadata inline via
    ray.get and emit immediately, budget off meta.size_bytes; in_loop_done
    fires the done-callback at end-of-stream and re-raises a task failure.
  • ThreadedMetadataFetcher — defers each pulled pair and fetches its
    metadata on its own background thread, budgets off the block's local
    object_size; completion is postponed until the deferred pairs emit, and
    the per-operator FIFO preserves emission order. (See the class docstring for
    a data-flow diagram of the two threads, the shared queue, and the FIFO.)
• DataOpTask.on_data_ready keeps one shared pull loop and delegates per-pair
  fetch/emit to in_loop_get_size and completion to in_loop_done (inline
  fires it, threaded is a no-op).
• process_completed_tasks / StreamingExecutor require a MetadataFetcher;
  the executor builds one via make_metadata_fetcher(). Background
  metadata-fetch failures go through the same max_errored_blocks accounting as
  inline errors, and submit/register_if_drained run in a finally so a
  thrown error can't strand already-deferred pairs.

Most executor-level tests drive the threaded mode through a shared
_process_completed_tasks_threaded helper; a few stay inline
(_process_completed_tasks_sync). A FakeFetcher test drives on_data_ready's
per-pair branches deterministically, plus a mode-selection test.

Test

| variant | total sched (s) | Δ |
|---|--:|--:|--:|--:|--:|--:|
| 2000_actors_1ops | 86 -> 58 | -32% |
| 2000_actors_15ops | 1099 -> 726 | -33% |
| 2000_tasks_1ops | 77 -> 50 | -34% |
| 2000_tasks_15ops | 1142 → 747 | -34% |
| 5000_actors_1ops | 204 → 218|
| 5000_actors_15ops | 3003 → 1834 | -38% |
| 5000_tasks_1ops | 229 → 189 | -17% |
| 5000_tasks_15ops | 2813 → 1863 | -33% |

[gemini-code-assist]

Code Review

This pull request introduces an asynchronous metadata prefetching mechanism for the streaming executor in Ray Data. By shifting metadata fetches to a background thread via MetadataPrefetcher, the main scheduling loop is prevented from blocking on straggling metadata references. The changes add MetadataFetcher strategies (ThreadedMetadataFetcher and InlineMetadataFetcher), update DataOpTask to manage deferred emissions and generator lifecycles, and integrate these components into the streaming executor state processing. Extensive unit tests are also added to verify correctness, ordering, and error handling under the new threaded mode. I have no feedback to provide as there are no review comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[cursor]

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

^{Reviewed by Cursor Bugbot for commit 85fe4a2. Configure here.}