_wait_generators_bulk perf

[rueian]

DO NOT MERGE

[gemini-code-assist]

Code Review

This pull request introduces bulk consumption and peeking APIs for streaming generators in Ray, including _wait_and_fetch and _wait_generators_bulk in Python, along with corresponding C++ CoreWorker implementations (WaitAndFetch, TryReadObjectRefStreamN, and PeekObjectRefStreamN). These changes are integrated into Ray Data's streaming executor to optimize task completion waiting. The review feedback highlights several critical issues: potential runtime NameErrors due to a missing time import and an un-namespaced exception reference, a bug in TryReadObjectRefStreamN where backpressure_threshold defaults to 0 instead of -1 (which could cause generator hangs), a potential crash from a strict RAY_CHECK_EQ on mismatched task IDs, and the risky use of assert False for control flow validation.

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]

The exception ObjectRefStreamEndOfStreamError is referenced here but it may not be imported in this file's namespace, which would cause a NameError at runtime. It is safer to reference it via ray.exceptions.ObjectRefStreamEndOfStreamError or explicitly import it.

Suggested change

	except ObjectRefStreamEndOfStreamError:
	except ray.exceptions.ObjectRefStreamEndOfStreamError:

[gemini-code-assist]

The time module is used in this file (e.g., time.monotonic()) but is not imported at the top of the file. This will result in a NameError at runtime. Please import time.

Suggested change

	from ray._private.worker import _wait_generators_bulk
	import time
	from ray._private.worker import _wait_generators_bulk

[gemini-code-assist]

In TryReadObjectRefStreamN, backpressure_threshold is initialized to 0. However, in the singular TryReadObjectRefStream, it is initialized to -1. If the task is not found in submissible_tasks_, defaulting to 0 instead of -1 incorrectly enables backpressure with a threshold of 0, which can cause the generator to hang. It should be initialized to -1 to match the singular version's behavior.

Suggested change

	auto backpressure_threshold = 0;
	auto backpressure_threshold = -1;

[gemini-code-assist]

Using RAY_CHECK_EQ(object_id.TaskId(), generator_task_id_) here can cause a crash if temporarily_owned_refs_ contains any object IDs from other tasks. Since TemporarilyInsertToStreamIfNeeded defensively allows object IDs with different task IDs to be inserted, we should handle this gracefully by returning false instead of crashing the process.

Suggested change

	RAY_CHECK_EQ(object_id.TaskId(), generator_task_id_);
	if (object_id.TaskId() != generator_task_id_) {
	return false;
	}

[gemini-code-assist]

Using assert False for control flow validation is risky because assertions can be optimized away when Python is run with the -O flag. If optimized away, this block will silently succeed instead of raising an exception. It is better to explicitly raise a RuntimeError.

Suggested change

	assert False, "Above ray.get should raise an exception."
	raise RuntimeError("Above ray.get should raise an exception.")