[Data] batch-meta-fetch + core end-of-stream timeout + ExecutionResources/budget-loops perf (combined, for testing)

[xinyuangui2]

Combined branch for testing/benchmarking only — not for review.

= #63701 (background-thread metadata fetch) + #64014 (core end-of-stream timeout fix) + #63964 (ExecutionResources __slots__/combine_sum + raw-float reservation/budget loops).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a background metadata prefetcher (MetadataPrefetcher) to asynchronously fetch block metadata, preventing the main executor thread from blocking on ray.get. It also bounds end-of-stream gets in ObjectRefGenerator with timeouts to avoid deadlocks, and optimizes the scheduling hot path by replacing intermediate ExecutionResources allocations with raw float arithmetic. Feedback on these changes highlights several critical issues: potential nan propagation in resource_manager.py when performing arithmetic with float("inf") limits, a missing asyncio import in object_ref_generator.py that would trigger a NameError, and the lack of error handling or liveness checks for the new background prefetcher thread, which could cause the pipeline to hang if the thread dies silently.

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.

[gemini-code-assist]

If limits contains float("inf") (which is the default for unspecified limits), subtracting task_cpu (which will also be inf) from remaining_cpu (which is inf) results in nan. This nan value will propagate and break scheduling logic. We should guard the subtraction to only occur when the remaining resource is not inf.

Suggested change

	# remaining = max(remaining - (reserved_for_tasks + reserved_for_outputs), 0)
	remaining_cpu = max(remaining_cpu - task_cpu, 0.0)
	remaining_gpu = max(remaining_gpu - task_gpu, 0.0)
	remaining_osm = max(
	remaining_osm - (task_osm + reserved_for_outputs_osm), 0.0
	)
	remaining_mem = max(remaining_mem - task_mem, 0.0)
	# remaining = max(remaining - (reserved_for_tasks + reserved_for_outputs), 0)
	if remaining_cpu != float("inf"):
	remaining_cpu = max(remaining_cpu - task_cpu, 0.0)
	if remaining_gpu != float("inf"):
	remaining_gpu = max(remaining_gpu - task_gpu, 0.0)
	if remaining_osm != float("inf"):
	remaining_osm = max(
	remaining_osm - (task_osm + reserved_for_outputs_osm), 0.0
	)
	if remaining_mem != float("inf"):
	remaining_mem = max(remaining_mem - task_mem, 0.0)

[gemini-code-assist]

When capping op_shared_cpu using max_resource_usage.cpu - alloc_cpu, if both are inf, it results in nan. This nan will propagate to op_shared_cpu and then to shared_cpu, breaking the budget allocation. We should only cap the shared allocation if the limit is not inf.

Suggested change

	# op_shared = min(op_shared, max(max_resource_usage - alloc, 0))
	op_shared_cpu = min(
	op_shared_cpu, max(max_resource_usage.cpu - alloc_cpu, 0.0)
	)
	op_shared_gpu = min(
	op_shared_gpu, max(max_resource_usage.gpu - alloc_gpu, 0.0)
	)
	op_shared_osm = min(
	op_shared_osm,
	max(max_resource_usage.object_store_memory - alloc_osm, 0.0),
	)
	op_shared_mem = min(
	op_shared_mem, max(max_resource_usage.memory - alloc_mem, 0.0)
	)
	# op_shared = min(op_shared, max(max_resource_usage - alloc, 0))
	if max_resource_usage.cpu != float("inf"):
	op_shared_cpu = min(
	op_shared_cpu, max(max_resource_usage.cpu - alloc_cpu, 0.0)
	)
	if max_resource_usage.gpu != float("inf"):
	op_shared_gpu = min(
	op_shared_gpu, max(max_resource_usage.gpu - alloc_gpu, 0.0)
	)
	if max_resource_usage.object_store_memory != float("inf"):
	op_shared_osm = min(
	op_shared_osm,
	max(max_resource_usage.object_store_memory - alloc_osm, 0.0),
	)
	if max_resource_usage.memory != float("inf"):
	op_shared_mem = min(
	op_shared_mem, max(max_resource_usage.memory - alloc_mem, 0.0)
	)

[gemini-code-assist]

When subtracting op_shared_cpu from shared_cpu, if shared_cpu is inf, subtracting op_shared_cpu (which will also be inf) results in nan. This nan will propagate to all subsequent operators' budgets. We should guard the subtraction to only occur when shared_cpu is not inf.

            if shared_cpu != float("inf"):
                shared_cpu -= op_shared_cpu
            if shared_gpu != float("inf"):
                shared_gpu -= op_shared_gpu
            if shared_osm != float("inf"):
                shared_osm -= op_shared_osm
            if shared_mem != float("inf"):
                shared_mem -= op_shared_mem

[gemini-code-assist]

If asyncio is not imported in this file, calling _next_async will raise a NameError when attempting to use asyncio.wait_for or catch asyncio.TimeoutError.

Suggested change

	from typing import TYPE_CHECKING, Deque, Iterator, Optional
	import ray
	from ray.exceptions import ObjectRefStreamEndOfStreamError
	from ray.exceptions import GetTimeoutError, ObjectRefStreamEndOfStreamError
	from ray.util.annotations import DeveloperAPI, PublicAPI
	import asyncio
	from typing import TYPE_CHECKING, Deque, Iterator, Optional
	import ray
	from ray.exceptions import GetTimeoutError, ObjectRefStreamEndOfStreamError
	from ray.util.annotations import DeveloperAPI, PublicAPI

[gemini-code-assist]

The background thread loop _run does not have any try-except block around its main loop. If an unexpected exception occurs (e.g., during ray.wait or queue operations), the thread will die silently. This will cause the entire Ray Data pipeline to hang indefinitely because drain() will always wait for results that will never be populated. Wrapping the loop in a try-except block to log the exception is highly recommended.

    def _run(self) -> None:
        """Fetch-thread loop: accumulate requested meta_refs into a pending
        set, ``ray.wait`` (with a short timeout) for the ones that are
        locally available, and ``ray.get`` + publish only those.

        ``ray.get`` is never issued on a ref that hasn't been reported ready
        by ``ray.wait(fetch_local=True)``: a ref stuck on a bad/dead node
        would otherwise block the whole thread forever and starve every
        other operator's metadata. A stuck ref just stays in ``pending``
        until Ray resolves or fails it.
        """
        try:
            pending: List["ray.ObjectRef"] = []
            while True:
                # Block on the request queue only when there's nothing in
                # flight; otherwise poll it so pending refs keep making progress.
                if pending:
                    try:
                        item = self._request_q.get_nowait()
                    except queue_module.Empty:
                        item = ()
                else:
                    item = self._request_q.get()
                if item is None:
                    return
                pending.extend(item)
                # Coalesce any other already-queued batches.
                while True:
                    try:
                        nxt = self._request_q.get_nowait()
                    except queue_module.Empty:
                        break
                    if nxt is None:
                        return
                    pending.extend(nxt)

                if not pending:
                    continue
                ready, pending = ray.wait(
                    pending,
                    num_returns=len(pending),
                    timeout=_FETCH_WAIT_TIMEOUT_S,
                    fetch_local=True,
                )
                if ready:
                    self._fetch(ready)
        except Exception:
            logger.exception("MetadataPrefetcher background thread died unexpectedly.")

[gemini-code-assist]

If the background prefetcher thread dies unexpectedly, the pipeline will hang silently. We should check if the thread is still alive in drain() and raise an error if it died.

Suggested change

	def drain(self) -> None:
	"""Emit every pair whose metadata is now available, in per-op append
	order, then fire postponed done callbacks for fully-drained tasks.
	Must be called on the executor thread.
	"""
	for fifo in self._fifos.values():
	def drain(self) -> None:
	"""Emit every pair whose metadata is now available, in per-op append
	order, then fire postponed done callbacks for fully-drained tasks.
	Must be called on the executor thread.
	"""
	if self._started and not self._thread.is_alive() and not self._stopped:
	raise RuntimeError("Metadata prefetcher background thread died unexpectedly.")
	for fifo in self._fifos.values():

[gemini-code-assist]

To prevent any potential nan propagation when shared_cpu is inf, we should guard the subtraction of exceeded resources.

Suggested change

	shared_cpu -= max(usage_cpu - reserved_cpu, 0.0)
	shared_gpu -= max(usage_gpu - reserved_gpu, 0.0)
	shared_osm -= max(usage_osm - reserved_osm, 0.0)
	shared_mem -= max(usage_mem - reserved_mem, 0.0)
	# How much of the reserved resources are exceeded.
	# If exceeded, we need to subtract from the remaining shared resources.
	if shared_cpu != float("inf"):
	shared_cpu -= max(usage_cpu - reserved_cpu, 0.0)
	if shared_gpu != float("inf"):
	shared_gpu -= max(usage_gpu - reserved_gpu, 0.0)
	if shared_osm != float("inf"):
	shared_osm -= max(usage_osm - reserved_osm, 0.0)
	if shared_mem != float("inf"):
	shared_mem -= max(usage_mem - reserved_mem, 0.0)