[Data] (test branch) Batch ray.get of metadata + wide-schema release test

[xinyuangui2]

What this PR is

A test branch combining:

• PR [Data] Batch ray.get of block metadata in process_completed_tasks #63483 (batch-meta-ray-get) — batched ray.get of block metadata in the StreamingExecutor scheduler thread.
• PR [Data] Fix worker_scaling release-test hang under wide schema #63513 (wide-schema-profiler) — 600-column release-test schema + driver-side py-spy profiler upload.

Used for one-off release-test runs to validate #63483's scheduler-thread speedup on the production-shape (600-col) workload. Not for merge — both underlying PRs land independently.

New batch-get design (reflected here at HEAD 5fad871408)

The previous version of #63483 batched only the first pending pair per ready task, leaving Case 1 (within-task generator drain) on the per-ref ray.get fallback (residual ~14 s of leaf get_objects time on 2000_tasks at 600 cols).

The current design replaces the single-pair pipeline with a queue-based one:

• DataOpTask now holds a Deque[(block_ref, meta_ref)] of pending pairs (plus a partial-block stash) instead of two scalar _pending_*_ref slots.
• peek_pending_pair() is drainable — callable in a loop to pull every pair the streaming generator currently has ready into the queue. It also detects task termination via flags handled at the end of on_data_ready.
• cache_prefetched_meta_bytes(dict) lets the caller seed the task's cache; the bytes survive across iterations and are consumed in lockstep with the queue.
• process_completed_tasks runs while peek_pending_pair() is not None: per ready task, issues one ray.get over the union of every surfaced meta_ref across every ready task, and seeds each task's cache with the result.
• on_data_ready becomes a pure consumer of the queue + cache; the per-ref ray.get survives only as the fallback for direct callers / cache misses.

Result: one batched ray.get per scheduler iteration covers all pending metadata, with the per-ref fallback eliminated from the production path.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request optimizes the Ray Data streaming executor by implementing batched metadata prefetching to reduce ray.get overhead. It introduces a peek_pending_pair method to PhysicalOperator and updates process_completed_tasks to retrieve metadata for multiple tasks in a single call. Additionally, the PR adds a pyspy_profiler utility for driver-side profiling in release tests and adjusts various nightly test configurations and benchmark parameters. Review feedback suggests using ray.wait with a zero timeout during prefetching to prevent scheduler stalls and recommends handling FileNotFoundError when launching the py-spy subprocess to improve robustness.

[gemini-code-assist]

Issuing a batched ray.get with a 1-second timeout directly in the scheduler thread can cause significant stalls if any of the metadata objects are not yet available (e.g., due to slow network propagation). Since process_completed_tasks is called frequently, it's safer to use ray.wait with timeout=0 to only prefetch the objects that are already locally available. This avoids blocking the scheduler while still providing the batching benefit for ready metadata.

Suggested change

	try:
	bytes_list = ray.get(
	meta_refs_to_prefetch, timeout=METADATA_GET_TIMEOUT_S
	)
	prefetched_meta_bytes = dict(zip(meta_refs_to_prefetch, bytes_list))
	except ray.exceptions.GetTimeoutError:
	# One or more refs weren't ready in time. Tasks whose meta_ref
	# isn't in `prefetched_meta_bytes` will fall back to the per-ref
	# ray.get path inside `on_data_ready` (which has its own
	# timeout-handling and retry-next-iteration semantics).
	pass
	# Use ray.wait with timeout=0 to only prefetch what's already ready,
	# avoiding stalling the scheduler thread.
	ready_meta_refs, _ = ray.wait(
	meta_refs_to_prefetch,
	num_returns=len(meta_refs_to_prefetch),
	timeout=0,
	)
	if ready_meta_refs:
	try:
	bytes_list = ray.get(ready_meta_refs)
	prefetched_meta_bytes = dict(zip(ready_meta_refs, bytes_list))
	except Exception:
	# Fallback to per-ref ray.get in on_data_ready if anything goes wrong.
	pass

[gemini-code-assist]

If the py-spy executable is not installed or not in the system PATH, subprocess.Popen will raise a FileNotFoundError. It's better to catch this exception to prevent the entire benchmark from crashing when profiling is enabled but the tool is missing.

    try:
        _proc = subprocess.Popen(cmd, stdout=_log_file, stderr=_log_file)
        _log_file.write(f"py-spy pid: {_proc.pid}\n")
        _log_file.flush()
    except FileNotFoundError:
        print("Error: 'py-spy' executable not found. Profiling will not start.")
        _log_file.write("Error: 'py-spy' executable not found.\n")
        _log_file.close()
        _log_file = None
        _proc = None
        return