[Data] Measure local block-size hit rate vs ray.get(meta) size in on_data_ready

[xinyuangui2]

Draft / measurement-only — not for merge as-is. Pure instrumentation; does not change scheduling behavior.

Why

The per-pair ray.get(meta_ref) in DataOpTask.on_data_ready is a top cost in process_completed_tasks at scale. #63701 proposes replacing it (where possible) with a local get_local_object_locations(block_ref).object_size lookup. This PR gathers the data to justify that: how often is the size known locally, and does it match the metadata size?

What

In on_data_ready, alongside the existing ray.get(meta_ref) (kept as the source of truth), look up the block's object_size via ray.experimental.get_local_object_locations (local, no RPC) and record, process-wide:

• hit rate — fraction of pairs where the local lookup returns a known object_size;
• accuracy — of those hits, how many match meta.size_bytes (from ray.get(meta_ref)).

A cumulative snapshot is logged every 5000 probed pairs:

[local-size-probe] 4550/5000 (91.0%) pairs had a local object_size; of those 4550/4550 (100.0%) matched meta.size_bytes (0 mismatched).

Note: the local lookup is RPC-free but still adds per-pair overhead, so this is for measurement runs (e.g. the worker_scaling release test), not production.

[gemini-code-assist]

Code Review

This pull request introduces process-wide telemetry in DataOpTask to measure how often a block's size is known locally and whether it matches the metadata size. The reviewer identified critical issues regarding thread safety and exception safety in the newly added _record_local_size_probe method. Specifically, updating class-level counters concurrently without synchronization can cause race conditions, and any failure in the telemetry lookup could crash the main execution path. The reviewer recommended wrapping the logic in a try-except block and using a threading.Lock to synchronize counter updates.

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 telemetry probe method _record_local_size_probe has two major issues:

1. Race Condition / Thread Safety: The class-level counters (_size_probe_total, _size_probe_local_hits, etc.) are shared across all DataOpTask instances. Since multiple dataset execution threads can run concurrently, updating and reading these counters without synchronization can lead to race conditions and inconsistent log messages.
2. Exception Safety: Telemetry/measurement code should never crash the main execution path. If get_local_object_locations raises an unexpected exception, it will crash the entire dataset execution.

We should wrap the telemetry logic in a try...except block and synchronize the shared counter updates using a lock.

    def _record_local_size_probe(self, meta_size_bytes: int) -> None:
        """Measurement-only: how often is the block's size known locally, and
        does it match the metadata size?

        Looks up ``self._pending_block_ref``'s ``object_size`` via
        ``get_local_object_locations`` (local, no RPC) and compares it to
        ``meta_size_bytes`` (``meta.size_bytes`` from ``ray.get(meta_ref)``,
        the source of truth). Accumulates process-wide counters and logs a
        cumulative snapshot every ``_SIZE_PROBE_LOG_EVERY`` probed pairs.

        Note: this is pure instrumentation — it does NOT change behavior. The
        existing ``ray.get(meta_ref)`` path still drives emits and budgets.
        """
        cls = DataOpTask
        block_ref = self._pending_block_ref
        try:
            info = get_local_object_locations([block_ref]).get(block_ref)
            local_size = info.get("object_size") if info is not None else None

            with cls._size_probe_lock:
                cls._size_probe_total += 1
                if local_size is not None:
                    cls._size_probe_local_hits += 1
                    if local_size == meta_size_bytes:
                        cls._size_probe_match += 1
                    else:
                        cls._size_probe_mismatch += 1

                if cls._size_probe_total % cls._SIZE_PROBE_LOG_EVERY == 0:
                    total = cls._size_probe_total
                    hits = cls._size_probe_local_hits
                    match_count = cls._size_probe_match
                    mismatch_count = cls._size_probe_mismatch
                    hit_pct = 100.0 * hits / total
                    match_pct = 100.0 * match_count / hits if hits else 0.0
                    logger.info(
                        "[local-size-probe] %d/%d (%.1f%%) pairs had a local object_size; "
                        "of those %d/%d (%.1f%%) matched meta.size_bytes "
                        "(%d mismatched).",
                        hits,
                        total,
                        hit_pct,
                        match_count,
                        hits,
                        match_pct,
                        mismatch_count,
                    )
        except Exception:
            logger.warning("Failed to record local size probe telemetry.", exc_info=True)

[gemini-code-assist]

Import threading to support thread-safe telemetry counters.

Suggested change

	from ray.experimental.locations import get_local_object_locations
	from ray.experimental.locations import get_local_object_locations
	import threading

[gemini-code-assist]

Add a class-level lock _size_probe_lock to synchronize updates to the shared telemetry counters across multiple concurrent dataset executions.

Suggested change

	_size_probe_total = 0
	_size_probe_local_hits = 0
	_size_probe_match = 0
	_size_probe_mismatch = 0
	# Emit a cumulative snapshot every this many probed pairs.
	_SIZE_PROBE_LOG_EVERY = 5000
	_size_probe_total = 0
	_size_probe_local_hits = 0
	_size_probe_match = 0
	_size_probe_mismatch = 0
	_size_probe_lock = threading.Lock()
	# Emit a cumulative snapshot every this many probed pairs.
	_SIZE_PROBE_LOG_EVERY = 5000

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.