[Data] Fetch block metadata on a background thread (defer RefBundle emits + batched ray.get)

[xinyuangui2]

What

At scale the streaming scheduler does a blocking ray.get(meta_ref) for every (block_ref, meta_ref) pair inside process_completed_tasks — the largest leaf cost on the scheduling thread (~1–2 ms/call, 50k+ pairs/iter at 5k workers). This PR moves that fetch off the scheduling thread onto a background thread.

Two changes make that possible:

1. Defer emits; size blocks locally. on_data_ready no longer emits RefBundles inline. It pulls each ready pair, appends it to a per-op deferred list, and returns. The output-budget accounting uses the block's local object_size (get_local_object_locations, no RPC) instead of meta.size_bytes — the driver owns every block ref the streaming generator yields, so the size is available locally. (Rare fallback: a short, timeout-bounded ray.get(meta) + warning if the local size is missing.)
2. Fetch metadata on a background thread. A dedicated MetadataPrefetcher thread does the ray.get(meta_refs). Each scheduling iteration submits the new pairs and drains whatever metadata has already landed, emitting each op's bundles in order.

Design

The blocking ray.get moves off the scheduling loop onto a dedicated thread. The two threads are coupled only by one thread-safe queue (_request_q); fetched bytes come back via _results. submit and drain both run inside the per-op loop / scheduling iteration; the background ray.get runs outside it.

flowchart TB
    subgraph SCHED["Scheduling thread — process_completed_tasks (each iteration)"]
        ODR["on_data_ready()<br/>• pull (block_ref, meta_ref)<br/>• read local object_size (no RPC)<br/>• append to op_deferred"]
        SUB["submit(op_deferred)"]
        DRAIN["drain()<br/>• emit pairs whose metadata is ready,<br/>&nbsp;&nbsp;in per-op FIFO order<br/>• fire task_done_callback once a<br/>&nbsp;&nbsp;task's pairs have all emitted"]
        OUT["RefBundle → downstream operator"]
        ODR --> SUB
        DRAIN --> OUT
    end

    Q[["_request_q<br/>thread-safe queue"]]
    RES[("_results<br/>meta_ref → bytes")]

    subgraph BG["MetadataPrefetcher thread (background)"]
        FETCH["ray.wait(fetch_local) →<br/>ray.get(meta_refs)"]
    end

    SUB -- "put(meta_refs)" --> Q
    Q -- "get" --> FETCH
    FETCH -- "publish bytes" --> RES
    RES -. "pop (non-blocking)" .-> DRAIN

Loading

Benchmark (this PR vs master)

worker_scaling release test, scheduling-loop cost. total sched = total time the scheduling thread spent in the loop; avg/p50 = per-iteration loop duration.

variant	total sched (s)	Δ	avg loop (s)	Δ	p50 loop (s)	Δ
2000_actors_1ops	51 → 36	-29%	1.00 → 0.60	-40%	1.43 → 0.46	-68%
2000_actors_15ops	920 → 666	-28%	10.33 → 6.40	-38%	14.29 → 3.68	-74%
2000_tasks_1ops	49 → 33	-32%	1.04 → 0.38	-64%	0.10 → 0.10	-3%
2000_tasks_15ops	927 → 702	-24%	9.09 → 4.65	-49%	11.34 → 4.29	-62%
5000_actors_1ops	204 → 218	+7%	3.19 → 0.27	-92%	4.23 → 0.24	-94%
5000_actors_15ops	2401 → 1739	-28%	30.78 → 19.54	-37%	42.69 → 13.53	-68%
5000_tasks_1ops	172 → 153	-11%	4.19 → 0.32	-92%	2.60 → 0.24	-91%
5000_tasks_15ops	2754 → 1696	-38%	29.62 → 11.54	-61%	36.43 → 9.42	-74%

Per-iteration loop duration drops sharply everywhere (p50 −62% to −94%) — the inline ray.get is gone from the loop. Total scheduling time falls 24–38% on the many-operator pipelines that matter at scale.

The one slight regression is 5000_actors_1ops (total +7%): with a single operator the loop isn't fetch-bound, so removing the inline ray.get makes each loop far cheaper but it then runs many more times, and the per-iteration ray.wait polling dominates. Per-loop latency still drops 92–94% there; only the summed total rises slightly.

Ray Data release tests

https://buildkite.com/ray-project/release/builds/97303 — no regressions.

[gemini-code-assist]

Code Review

This pull request introduces metadata prefetching for streaming executor tasks to batch ray.get calls, utilizing get_local_object_locations to identify ready metadata references without RPC overhead. The feedback suggests making the prefetch peeking non-blocking by setting the timeout to zero, wrapping the local location lookup in a try-except block to prevent scheduler crashes, and resetting telemetry counters after logging to track interval-based hit rates.

[cursor]

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

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

[iamjustinhsu]

Nice work!! I was a bit pedantic on the comments this time because this part of the code is pretty sensitive to changes. Anyways, I like the idea, but we'll probably need eyes from balaji too

[iamjustinhsu]

why do we need this self._published? Like what happens if we just use any(self._fifos.values()) and don't have a block_timeout_s?

[xinyuangui2]

In experiments, we find sometimes the result queue is empty but scheduling loop will continue and waste some calculation. Thus, we use this event to mitigate that. Release tests show this can avoid some regressions in tests.

[iamjustinhsu]

Hmm not sure I completely understand this part:

we find sometimes the result queue is empty but scheduling loop will continue and waste some calculation

When the queue is empty, do u know where the scheduling loop is spending time? Why does adding block_timeout_s and self._published make the scheduling thread faster?

[iamjustinhsu]

Is it possible to narrow which exceptions could be raised?

[xinyuangui2]

We aligned with the original implementation, and store the exceptions whatever the type is. I didn't take a closer look at the possible types.

[iamjustinhsu]

If we reach the sentinel, should there be pending meta refs?

[xinyuangui2]

There should not be. Since the STOP happens after the thread joins.

[iamjustinhsu]

Multiple questions:

• If you call get_nowait(), rather than one with a timeout, will this busy-wait, burning CPU cycles since we are continously calling get_nowait?
• How come you case with or without pending?
• How come you have another step to "Coalesce any other already-queued batches.", which looks very similar to what's above?

[xinyuangui2]

Simplified the loop. Hope it is more clear

[iamjustinhsu]

More feedback, I still need to look at the tests, but at a glance, I don't see any usage of callback seams to use in test. What I mean is we have a metadata callback and block ready callback to make it easier to test. You could theoretically have callbacks as well for when object_size is known or when the ray.get after a fetch_local True returns an exception and you need to ray.get each individual ref, and then use them in the test like so. By doing this you'll have more control on entering each if statement, or try catch statement. lmk if that doesn't make sense

[iamjustinhsu]

Is there a reason that we don't want to pull more if there are pending emits? This is my thinking:

• The decision to pull only depends the size_bytes, which will always be known through 2 ways: get_local_object_locations, or a fallback to ray.get
• This decision does not include if it has pending emits, so theoretically it seems to me it's possible

oh also not saying we should change this, just want to understand the reasoning

[xinyuangui2]

I am trying to match the old behavior: orders are kept and pull is blocked when previous metadata is not ready.

[iamjustinhsu]

Just want to double check, this warning doesn't get emitted very often, or at least in ur testing it didn't emit once? Regardless, I'm not sure if this is possible, but if the metadata starts not showing up for whatever reason for multiple tasks, this could be spamming. Do u think it would be good to gate this warning behind a log_once call. We have this function used to log stuff once, and it could be handy here

[iamjustinhsu]

i think to keep it consistent, you can rename self._done_pending to self._drained_tasks or smth like that

[iamjustinhsu]

Hmm not sure I completely understand this part:

we find sometimes the result queue is empty but scheduling loop will continue and waste some calculation

When the queue is empty, do u know where the scheduling loop is spending time? Why does adding block_timeout_s and self._published make the scheduling thread faster?

[iamjustinhsu]

Should this sentinel also be a part of the _Stop enum? For example, we could rename it to _Signal, and then host 2 values: Stop (for streaming executor done), and not_ready (for object not ready for ray.get)

[iamjustinhsu]

I think we should also log a warning here too, but only log_once. I don't think this should be a common case, but we should at least know if we are heading down this path

[iamjustinhsu]

Since this change is risky, consider having an env flag to turn this off just in case. You can set this to default True, and that way people can go back to before, but we'll be able to test this with the community

[iamjustinhsu]

Suggested change

	"""Queue ``deferred`` pairs (from one operator) for fetching + emission,
	and register any end-of-stream tasks for a postponed done callback.
	"""Queue ``deferred`` pairs (from one operator) for fetching + emission,
	and register any end-of-stream tasks (ready_tasks) for a postponed done callback.

I was thinking u may want to refer to the argument parameters.

But also, this function appears to do to separate things (one for new metadata_ref, one for making sure we call task_done_callback), so I'm thinking u could split this function up into those 2 phases. It seems like drain also does the same 2 things as well (correct me if i'm wrong). wdyt?

[xinyuangui2]

close in favor of #64378