From a41774d4ca4f2abb9ffe564e775a3db724a1d526 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Wed, 1 Jul 2026 15:49:29 -0700 Subject: [PATCH 01/14] Core side wiring + testing Signed-off-by: Sirui Huang --- python/ray/_raylet.pyx | 13 +++++++++ python/ray/includes/libcoreworker.pxd | 2 ++ src/ray/core_worker/core_worker.cc | 4 +++ src/ray/core_worker/core_worker.h | 4 +++ src/ray/core_worker/reference_counter.cc | 12 ++++++++ src/ray/core_worker/reference_counter.h | 3 ++ .../core_worker/reference_counter_interface.h | 5 ++++ .../tests/reference_counter_test.cc | 28 +++++++++++++++++++ 8 files changed, 71 insertions(+) diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 16f68b4ca9ed..b8743cf6f557 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -4300,6 +4300,19 @@ cdef class CoreWorker: cpython.Py_DECREF(callback) return registered + def remove_object_out_of_scope_callbacks(self, ObjectRef object_ref): + """Remove and fire all out-of-scope callbacks for object_ref. + + Firing allows per-callback cleanup (e.g. Py_DECREF). After this + call, no further out-of-scope callbacks will fire for this object. + + .. warning:: + This is an internal Ray API. Do not use it outside of Ray libraries. + """ + cdef CObjectID c_object_id = object_ref.native() + CCoreWorkerProcess.GetCoreWorker() \ + .RemoveObjectOutOfScopeOrFreedCallbacks(c_object_id) + def get_owner_address(self, ObjectRef object_ref): cdef: CObjectID c_object_id = object_ref.native() diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index ed5b37ffcbce..e27268a239fd 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/python/ray/includes/libcoreworker.pxd @@ -257,6 +257,8 @@ cdef extern from "ray/core_worker/core_worker.h" nogil: const CObjectID &object_id, void (*callback)(const CObjectID &, void *) nogil, void *callback_context) + void RemoveObjectOutOfScopeOrFreedCallbacks( + const CObjectID &object_id) CRayStatus CheckObjectOwnedByUs(const CObjectID &object_id) const void PutObjectIntoPlasma(const CRayObject &object, const CObjectID &object_id) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index e5cd84aa5718..e538cc9874c8 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -2555,6 +2555,10 @@ bool CoreWorker::AddObjectOutOfScopeOrFreedCallback(const ObjectID &object_id, }); } +void CoreWorker::RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) { + reference_counter_->RemoveObjectOutOfScopeOrFreedCallbacks(object_id); +} + Status CoreWorker::CheckObjectOwnedByUs(const ObjectID &object_id) const { if (reference_counter_->OwnedByUs(object_id)) { return Status::OK(); diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 4158cf014aec..704699fe5c9f 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -467,6 +467,10 @@ class CoreWorker : public std::enable_shared_from_this { /// a borrowed object owned by another worker). Status CheckObjectOwnedByUs(const ObjectID &object_id) const; + /// Remove and fire all out-of-scope/freed callbacks for `object_id`. + /// See ReferenceCounterInterface::RemoveObjectOutOfScopeOrFreedCallbacks. + void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id); + int GetMemoryStoreSize() { return memory_store_->Size(); } /// Returns a map of all ObjectIDs currently in scope with a pair of their diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index b28d756bbb1b..afaa9c1f2e4b 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -886,6 +886,18 @@ bool ReferenceCounter::AddObjectRefDeletedCallback( return true; } +void ReferenceCounter::RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) { + absl::MutexLock lock(&mutex_); + auto it = object_id_refs_.find(object_id); + if (it == object_id_refs_.end()) { + return; + } + for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { + callback(object_id); + } + it->second.on_object_out_of_scope_or_freed_callbacks.clear(); +} + bool ReferenceCounter::AddObjectOutOfScopeOrFreedCallback( const ObjectID &object_id, const std::function callback) { absl::MutexLock lock(&mutex_); diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h index 0a1b4ceacb4d..d9f5b7418d3b 100644 --- a/src/ray/core_worker/reference_counter.h +++ b/src/ray/core_worker/reference_counter.h @@ -156,6 +156,9 @@ class ReferenceCounter : public ReferenceCounterInterface, const std::function callback) override ABSL_LOCKS_EXCLUDED(mutex_); + void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) override + ABSL_LOCKS_EXCLUDED(mutex_); + bool AddObjectRefDeletedCallback( const ObjectID &object_id, std::function callback) override ABSL_LOCKS_EXCLUDED(mutex_); diff --git a/src/ray/core_worker/reference_counter_interface.h b/src/ray/core_worker/reference_counter_interface.h index 67eb9bfc67b3..2ba68abb0d15 100644 --- a/src/ray/core_worker/reference_counter_interface.h +++ b/src/ray/core_worker/reference_counter_interface.h @@ -329,6 +329,11 @@ class ReferenceCounterInterface { const ObjectID &object_id, const std::function callback) = 0; + /// Removes and fires all out-of-scope/freed callbacks for `object_id`. + /// Firing allows per-callback cleanup (e.g. Py_DECREF on the Python side). + /// After this call, no further out-of-scope callbacks will fire for this object. + virtual void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) = 0; + /// Stores the callback that will be run when the object reference is deleted /// from the reference table (all refs including lineage ref count go to 0). /// There could be multiple callbacks for the same object due to retries and we store diff --git a/src/ray/core_worker/tests/reference_counter_test.cc b/src/ray/core_worker/tests/reference_counter_test.cc index 0662a6e49e0a..7182eb021c81 100644 --- a/src/ray/core_worker/tests/reference_counter_test.cc +++ b/src/ray/core_worker/tests/reference_counter_test.cc @@ -634,6 +634,34 @@ TEST_F(ReferenceCountTest, TestUnreconstructableObjectOutOfScope) { ASSERT_TRUE(*out_of_scope); } +TEST_F(ReferenceCountTest, TestRemoveObjectOutOfScopeCallbacks) { + ObjectID id = ObjectID::FromRandom(); + rpc::Address address; + address.set_ip_address("1234"); + + auto callback_count = std::make_shared(0); + auto callback = [&](const ObjectID &object_id) { (*callback_count)++; }; + + // Register a callback on an owned, in-scope object. + rc->AddOwnedObject(id, + {}, + address, + "", + 0, + LineageReconstructionEligibility::INELIGIBLE_PUT, + /*add_local_ref=*/true); + ASSERT_TRUE(rc->AddObjectOutOfScopeOrFreedCallback(id, callback)); + + // Remove fires the callback. + rc->RemoveObjectOutOfScopeOrFreedCallbacks(id); + ASSERT_EQ(*callback_count, 1); + + // The object going out of scope should not fire the callback again. + std::vector out; + rc->RemoveLocalReference(id, &out); + ASSERT_EQ(*callback_count, 1); +} + // Tests call site tracking and ability to update object size. TEST_F(ReferenceCountTest, TestReferenceStats) { ObjectID id1 = ObjectID::FromRandom(); From c3a6c680ba1531439f8b4c41240bb04deb329899 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Wed, 1 Jul 2026 16:18:32 -0700 Subject: [PATCH 02/14] Address Core data object id type mismatch Signed-off-by: Sirui Huang --- python/ray/_raylet.pyx | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index b8743cf6f557..2d9d6d46c287 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -4300,16 +4300,19 @@ cdef class CoreWorker: cpython.Py_DECREF(callback) return registered - def remove_object_out_of_scope_callbacks(self, ObjectRef object_ref): - """Remove and fire all out-of-scope callbacks for object_ref. + def remove_object_out_of_scope_callbacks(self, object_id_binary: bytes): + """Remove and fire all out-of-scope callbacks for the given object. Firing allows per-callback cleanup (e.g. Py_DECREF). After this call, no further out-of-scope callbacks will fire for this object. + Args: + object_id_binary: The object ID as bytes (from ObjectRef.binary()). + .. warning:: This is an internal Ray API. Do not use it outside of Ray libraries. """ - cdef CObjectID c_object_id = object_ref.native() + cdef CObjectID c_object_id = CObjectID.FromBinary(object_id_binary) CCoreWorkerProcess.GetCoreWorker() \ .RemoveObjectOutOfScopeOrFreedCallbacks(c_object_id) From d1b4fd1a0d4f718faaefb5d9b55cd3627e7c16d2 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 13 Jul 2026 16:35:51 -0700 Subject: [PATCH 03/14] Update callback throughput benchmark Signed-off-by: Sirui Huang --- .../object_store/test_callback_throughput.py | 351 ++++++++++++++++-- 1 file changed, 312 insertions(+), 39 deletions(-) diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index b45ba848806d..ab4832d675d8 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -7,6 +7,7 @@ import ray import ray._private.worker +from ray.data._internal.execution.block_ref_counter import BlockRefCounter NUM_WORKERS = 10 OBJECT_SIZE = 1024 * 1024 # 1 MiB, above the 100 KB inlining threshold @@ -18,39 +19,77 @@ def produce_block(): @ray.remote(num_cpus=1) -def consume_block(block_ref): - return len(block_ref) +def consume_block(block): + return None -def test_callback_pipeline(num_blocks, timeout_s=60): - core_worker = ray._private.worker.global_worker.core_worker +def _produce_blocks(num_blocks): + """Create num_blocks plasma objects spread across the cluster.""" + refs = [ + produce_block.options(scheduling_strategy="SPREAD").remote() + for _ in range(num_blocks) + ] + ray.wait(refs, num_returns=len(refs)) + return refs - latencies = [] - drop_times = {} + +def _compute_latencies(fire_times, drop_times, id_binaries): + """Compute sorted per-callback latencies from fire timestamps. + + drop_times can be a float (single timestamp for all blocks) or a + dict mapping id_binary -> per-block drop timestamp. + """ + if isinstance(drop_times, dict): + latencies = [fire_times[id_b] - drop_times[id_b] for id_b in id_binaries] + else: + latencies = [fire_times[id_b] - drop_times for id_b in id_binaries] + latencies.sort() + return { + "p50": latencies[int(len(latencies) * 0.50)], + "p95": latencies[int(len(latencies) * 0.95)], + "p99": latencies[int(len(latencies) * 0.99)], + "max": latencies[-1], + } + + +def _make_timing_callback(num_blocks): + """Create a callback that records fire timestamps and signals completion. + + Returns (callback, fire_times, done_event). + """ + fire_times = {} lock = threading.Lock() done = threading.Event() def on_freed(id_bytes): + t = time.perf_counter() with lock: - latencies.append(time.perf_counter() - drop_times[id_bytes]) - if len(latencies) == num_blocks: + fire_times[id_bytes] = t + if len(fire_times) == num_blocks: done.set() - refs = [ - produce_block.options(scheduling_strategy="SPREAD").remote() - for _ in range(num_blocks) - ] - ray.wait(refs, num_returns=len(refs)) + return on_freed, fire_times, done + + +def test_callback_pipeline(num_blocks, timeout_s=300): + """Incremental produce-consume-release pipeline. + + Measures p95 latency from ref drop to callback fire, + with one block released at a time as its consumer completes. + """ + core_worker = ray._private.worker.global_worker.core_worker + on_freed, fire_times, done = _make_timing_callback(num_blocks) + + refs = _produce_blocks(num_blocks) - # live_refs keeps each block ref alive until its consumer completes. live_refs = {} for ref in refs: assert core_worker.add_object_out_of_scope_callback(ref, on_freed) - consumer = consume_block.remote(ref) - live_refs[consumer] = ref + live_refs[consume_block.remote(ref)] = ref del refs # Release each ref as its consumer completes. + drop_times = {} pending = list(live_refs.keys()) while pending: done_list, pending = ray.wait(pending, num_returns=1) @@ -61,18 +100,182 @@ def on_freed(id_bytes): if not done.wait(timeout=timeout_s): raise TimeoutError( - f"Only {len(latencies)}/{num_blocks} callbacks fired within {timeout_s}s" + f"Only {len(fire_times)}/{num_blocks} callbacks fired within {timeout_s}s" ) - latencies.sort() - p95 = latencies[int(len(latencies) * 0.95)] - print(f" {num_blocks} blocks: p95={p95:.4f}s") - return p95 + id_binaries = list(fire_times.keys()) + result = _compute_latencies(fire_times, drop_times, id_binaries) + print( + f" {num_blocks} blocks: " + f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s max={result['max']:.4f}s" + ) + return result + + +def test_registration_cost(num_blocks): + """Measures per-callback registration cost via BlockRefCounter. + + Uses on_block_produced (the real Data-layer path) which includes: + Python lock + set add + dict increment + Core API call + Py_INCREF. + """ + counter = BlockRefCounter() + refs = _produce_blocks(num_blocks) + + start = time.perf_counter() + for ref in refs: + counter.on_block_produced(ref, OBJECT_SIZE, "bench_op") + elapsed = time.perf_counter() - start + + per_callback_us = (elapsed / num_blocks) * 1e6 + print( + f" {num_blocks} registrations: {elapsed:.4f}s total, {per_callback_us:.1f}us each" + ) + + del refs, ref + return per_callback_us + + +def test_burst_drop(num_blocks, timeout_s=60): + """All refs dropped at once, 1 Core API callback per block. + + Measures time from burst start to each callback firing. The max + latency approximates total drain time (how long the burst hangs). + This is the production-representative baseline (no BlockRefCounter). + """ + core_worker = ray._private.worker.global_worker.core_worker + on_freed, fire_times, done = _make_timing_callback(num_blocks) + + refs = _produce_blocks(num_blocks) + for ref in refs: + assert core_worker.add_object_out_of_scope_callback(ref, on_freed) + + id_binaries = [ref.binary() for ref in refs] + drop_time = time.perf_counter() + del refs, ref + + if not done.wait(timeout=timeout_s): + raise TimeoutError( + f"Only {len(fire_times)}/{num_blocks} callbacks fired within {timeout_s}s" + ) + + result = _compute_latencies(fire_times, drop_time, id_binaries) + print( + f" burst {num_blocks} blocks: " + f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s " + f"p99={result['p99']:.4f}s max={result['max']:.4f}s" + ) + return result + + +def test_burst_drop_per_callback(num_blocks, timeout_s=60): + """Drops blocks one at a time with per-block timestamps. + + Measures true per-callback latency (each block's drop time to its + callback fire time). More authentic than test_burst_drop for the + LIMIT scenario, where the executor drains queues in a loop. + """ + core_worker = ray._private.worker.global_worker.core_worker + on_freed, fire_times, done = _make_timing_callback(num_blocks) + + refs = _produce_blocks(num_blocks) + for ref in refs: + assert core_worker.add_object_out_of_scope_callback(ref, on_freed) + + # Drop blocks one at a time, capturing per-block drop timestamps. + id_binaries = [] + drop_times = {} + for i in range(len(refs)): + ref = refs[i] + refs[i] = None + id_b = ref.binary() + id_binaries.append(id_b) + drop_times[id_b] = time.perf_counter() + del ref + del refs + + if not done.wait(timeout=timeout_s): + raise TimeoutError( + f"Only {len(fire_times)}/{num_blocks} callbacks fired within {timeout_s}s" + ) + + result = _compute_latencies(fire_times, drop_times, id_binaries) + print( + f" per-callback {num_blocks} blocks: " + f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s " + f"p99={result['p99']:.4f}s max={result['max']:.4f}s" + ) + return result + + +def _burst_drop_via_counter(num_blocks, use_clear, timeout_s=60): + """Burst drop using BlockRefCounter (the real Data-layer path). + + Registers callbacks via on_block_produced (which internally registers + a Core callback), then registers a second Core callback for timing. + Both fire on the same single-threaded callback service in registration + order, so the timing callback's latency includes the BlockRefCounter + callback that fires before it. + + With use_clear=True, clear() is called before dropping refs, so the + BlockRefCounter callbacks no-op. With use_clear=False, they do real + bookkeeping. Comparing the two isolates the bookkeeping cost. + """ + core_worker = ray._private.worker.global_worker.core_worker + counter = BlockRefCounter() + on_freed, fire_times, done = _make_timing_callback(num_blocks) + + refs = _produce_blocks(num_blocks) + + for ref in refs: + counter.on_block_produced(ref, OBJECT_SIZE, "bench_op") + for ref in refs: + assert core_worker.add_object_out_of_scope_callback(ref, on_freed) + + if use_clear: + counter.clear() + + id_binaries = [ref.binary() for ref in refs] + drop_time = time.perf_counter() + del refs, ref + + if not done.wait(timeout=timeout_s): + raise TimeoutError( + f"Only {len(fire_times)}/{num_blocks} callbacks fired within {timeout_s}s" + ) + + return _compute_latencies(fire_times, drop_time, id_binaries) + + +def test_burst_drop_block_ref_counter(num_blocks, timeout_s=60): + """Burst drop through BlockRefCounter without clear(). + + Callbacks do real bookkeeping (lock + set discard + dict decrement). + """ + result = _burst_drop_via_counter(num_blocks, use_clear=False, timeout_s=timeout_s) + print( + f" burst {num_blocks} blocks (no clear): " + f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s max={result['max']:.4f}s" + ) + return result + + +def test_burst_drop_block_ref_counter_with_clear(num_blocks, timeout_s=60): + """Burst drop through BlockRefCounter with clear() first. + + Callbacks no-op (check _registered_ids, return immediately). + Compare with test_burst_drop_block_ref_counter to isolate the cost + of real bookkeeping vs no-op under burst load. + """ + result = _burst_drop_via_counter(num_blocks, use_clear=True, timeout_s=timeout_s) + print( + f" burst {num_blocks} blocks (with clear): " + f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s max={result['max']:.4f}s" + ) + return result ray.init(address="auto") -# Warm up gRPC connections and worker pools. ray.get( [ produce_block.options(scheduling_strategy="SPREAD").remote() @@ -80,29 +283,99 @@ def on_freed(id_bytes): ] ) -p95_100 = test_callback_pipeline(100) -p95_1k = test_callback_pipeline(1000) +# Let callback service thread drain between tests to avoid interference. +DRAIN_DELAY_S = 1.0 + +# Scales to test. Higher values reveal whether per-callback cost is constant +# or grows with N (due to GIL contention, queue growth, etc.). +SCALES = [100, 1000, 5000, 10000, 50000] + + +def _run_at_scales(name, test_fn, scales): + print(f"\n=== {name} ===") + results = {} + for n in scales: + results[n] = test_fn(n) + time.sleep(DRAIN_DELAY_S) + return results + + +reg = _run_at_scales("Registration cost", test_registration_cost, SCALES) +pipeline = _run_at_scales("Incremental pipeline", test_callback_pipeline, SCALES) +burst = _run_at_scales("Burst drop (total drain time)", test_burst_drop, SCALES) +per_cb = _run_at_scales( + "Burst drop (per-callback latency)", test_burst_drop_per_callback, SCALES +) +brc = _run_at_scales( + "Burst drop (BlockRefCounter, no clear)", test_burst_drop_block_ref_counter, SCALES +) +brc_clear = _run_at_scales( + "Burst drop (BlockRefCounter, with clear)", + test_burst_drop_block_ref_counter_with_clear, + SCALES, +) -print("\nSummary:") -print(f" 100 blocks: p95={p95_100:.4f}s") -print(f" 1k blocks: p95={p95_1k:.4f}s") +print("\n=== Scaling summary (p95) ===") +header = " {:25s}" + " {:>10s}" * len(SCALES) +print(header.format("Test", *[f"{n}" for n in SCALES])) +for name, results in [ + ("Registration (us/cb)", reg), + ("Burst drain", burst), + ("Per-callback", per_cb), + ("BRC (no clear)", brc), + ("BRC (with clear)", brc_clear), +]: + vals = [] + for n in SCALES: + if n not in results: + vals.append("--") + elif isinstance(results[n], dict): + vals.append(f"{results[n]['p95']:.4f}s") + else: + vals.append(f"{results[n]:.1f}") + print(header.format(name, *vals)) + +print( + "\n Pipeline p95: " + ", ".join(f"{pipeline[n]['p95']:.4f}s ({n})" for n in SCALES) +) if "TEST_OUTPUT_JSON" in os.environ: - with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file: - results = { - "p95_100": p95_100, - "p95_1k": p95_1k, - "perf_metrics": [ + perf_metrics = [ + { + "perf_metric_name": "callback_p95_latency_1k_blocks_s", + "perf_metric_value": pipeline[1000]["p95"], + "perf_metric_type": "LATENCY", + }, + { + "perf_metric_name": "callback_registration_cost_s", + "perf_metric_value": reg[1000] / 1e6, + "perf_metric_type": "LATENCY", + }, + ] + for n in SCALES: + perf_metrics.extend( + [ + { + "perf_metric_name": f"callback_burst_drain_p95_{n}_blocks_s", + "perf_metric_value": burst[n]["p95"], + "perf_metric_type": "LATENCY", + }, + { + "perf_metric_name": f"callback_per_callback_p95_{n}_blocks_s", + "perf_metric_value": per_cb[n]["p95"], + "perf_metric_type": "LATENCY", + }, { - "perf_metric_name": "callback_p95_latency_100_blocks_s", - "perf_metric_value": p95_100, + "perf_metric_name": f"callback_burst_brc_p95_{n}_blocks_s", + "perf_metric_value": brc[n]["p95"], "perf_metric_type": "LATENCY", }, { - "perf_metric_name": "callback_p95_latency_1k_blocks_s", - "perf_metric_value": p95_1k, + "perf_metric_name": f"callback_burst_brc_clear_p95_{n}_blocks_s", + "perf_metric_value": brc_clear[n]["p95"], "perf_metric_type": "LATENCY", }, - ], - } - json.dump(results, out_file) + ] + ) + with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file: + json.dump({"perf_metrics": perf_metrics}, out_file) From 9316b613252b9f4ccf01ef8c2faeab6358a22a39 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Tue, 14 Jul 2026 14:01:54 -0700 Subject: [PATCH 04/14] Add timeout Signed-off-by: Sirui Huang --- release/release_tests.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/release/release_tests.yaml b/release/release_tests.yaml index de7da29bf287..468ed3e320cb 100644 --- a/release/release_tests.yaml +++ b/release/release_tests.yaml @@ -3235,7 +3235,7 @@ cluster_compute: object_store/callback_throughput.yaml run: - timeout: 600 + timeout: 3600 script: python object_store/test_callback_throughput.py wait_for_nodes: num_nodes: 11 From 40f2020802c22b763ac5a24780341f72f8b96c71 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Tue, 14 Jul 2026 16:18:20 -0700 Subject: [PATCH 05/14] Decrease 50000 config to 20000 due to oom Signed-off-by: Sirui Huang --- release/benchmarks/object_store/test_callback_throughput.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index ab4832d675d8..842f173ce7ca 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -288,7 +288,7 @@ def test_burst_drop_block_ref_counter_with_clear(num_blocks, timeout_s=60): # Scales to test. Higher values reveal whether per-callback cost is constant # or grows with N (due to GIL contention, queue growth, etc.). -SCALES = [100, 1000, 5000, 10000, 50000] +SCALES = [100, 1000, 5000, 10000, 20000] def _run_at_scales(name, test_fn, scales): From a3f50af0033df789764d6036345ca045a839834b Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Tue, 14 Jul 2026 19:54:25 -0700 Subject: [PATCH 06/14] Remove 20000 config Signed-off-by: Sirui Huang --- release/benchmarks/object_store/test_callback_throughput.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index 842f173ce7ca..26cb0f1999a7 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -288,7 +288,7 @@ def test_burst_drop_block_ref_counter_with_clear(num_blocks, timeout_s=60): # Scales to test. Higher values reveal whether per-callback cost is constant # or grows with N (due to GIL contention, queue growth, etc.). -SCALES = [100, 1000, 5000, 10000, 20000] +SCALES = [100, 1000, 5000, 10000] def _run_at_scales(name, test_fn, scales): From 7c81b3e60e96bd3119f0ce440987653229ef6051 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 20 Jul 2026 13:27:46 -0700 Subject: [PATCH 07/14] Remove clear wiring Signed-off-by: Sirui Huang --- .../_internal/execution/block_ref_counter.py | 10 ---- .../_internal/execution/streaming_executor.py | 3 -- .../data/tests/unit/test_block_ref_counter.py | 27 ---------- .../object_store/test_callback_throughput.py | 53 +++---------------- 4 files changed, 6 insertions(+), 87 deletions(-) diff --git a/python/ray/data/_internal/execution/block_ref_counter.py b/python/ray/data/_internal/execution/block_ref_counter.py index d4babc742e5d..4106f8ca73fa 100644 --- a/python/ray/data/_internal/execution/block_ref_counter.py +++ b/python/ray/data/_internal/execution/block_ref_counter.py @@ -73,13 +73,3 @@ def get_object_store_memory_usage(self, producer_id: str) -> int: """Total bytes of live blocks attributed to producer_id.""" with self._lock: return self._bytes_by_producer.get(producer_id, 0) - - def clear(self) -> None: - """Reset all accounting, e.g. on executor shutdown. - - Any previously registered Ray Core callbacks firing after clear() - will be silently ignored because _registered_ids is empty. - """ - with self._lock: - self._registered_ids.clear() - self._bytes_by_producer.clear() diff --git a/python/ray/data/_internal/execution/streaming_executor.py b/python/ray/data/_internal/execution/streaming_executor.py index 771c97c7b729..ce1848489d6b 100644 --- a/python/ray/data/_internal/execution/streaming_executor.py +++ b/python/ray/data/_internal/execution/streaming_executor.py @@ -358,9 +358,6 @@ def shutdown(self, force: bool, exception: Optional[Exception] = None): op.shutdown(timer, force=force) self._clear_topology_queues_post_shutdown(force, exception) - # Queues have been drained; any remaining Ray Core callbacks that fire - # after this point should be no-ops. - self._block_ref_counter.clear() min_ = round(timer.min(), 3) max_ = round(timer.max(), 3) diff --git a/python/ray/data/tests/unit/test_block_ref_counter.py b/python/ray/data/tests/unit/test_block_ref_counter.py index f1b01d575269..0b9786755387 100644 --- a/python/ray/data/tests/unit/test_block_ref_counter.py +++ b/python/ray/data/tests/unit/test_block_ref_counter.py @@ -102,33 +102,6 @@ def test_register_when_object_already_out_of_scope(self): assert counter.get_object_store_memory_usage("op_a") == 0 -class TestBlockRefCounterClear: - def test_clear_resets_usage(self): - add_cb = FakeAddObjectOutOfScopeCallback() - counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb) - counter.on_block_produced(_ref(1), 1, "op_a") - assert counter.get_object_store_memory_usage("op_a") == 1 - - counter.clear() - assert counter.get_object_store_memory_usage("op_a") == 0 - - def test_stale_callback_after_clear_is_noop(self): - """A stale callback firing after clear() must not touch accounting - recorded after the reset.""" - add_cb = FakeAddObjectOutOfScopeCallback() - counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb) - stale_ref = _ref(1) - counter.on_block_produced(stale_ref, 1, "op_a") - - counter.clear() - - counter.on_block_produced(_ref(2), 1, "op_a") - assert counter.get_object_store_memory_usage("op_a") == 1 - - add_cb.fire(stale_ref) - assert counter.get_object_store_memory_usage("op_a") == 1 - - class TestBlockRefCounterThreadSafety: def test_concurrent_callbacks_dont_corrupt_state(self): """Many threads firing callbacks at once must not corrupt the count.""" diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index 26cb0f1999a7..21b87e5014b5 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -207,18 +207,14 @@ def test_burst_drop_per_callback(num_blocks, timeout_s=60): return result -def _burst_drop_via_counter(num_blocks, use_clear, timeout_s=60): - """Burst drop using BlockRefCounter (the real Data-layer path). +def test_burst_drop_block_ref_counter(num_blocks, timeout_s=60): + """Burst drop through BlockRefCounter (the real Data-layer path). Registers callbacks via on_block_produced (which internally registers a Core callback), then registers a second Core callback for timing. Both fire on the same single-threaded callback service in registration order, so the timing callback's latency includes the BlockRefCounter callback that fires before it. - - With use_clear=True, clear() is called before dropping refs, so the - BlockRefCounter callbacks no-op. With use_clear=False, they do real - bookkeeping. Comparing the two isolates the bookkeeping cost. """ core_worker = ray._private.worker.global_worker.core_worker counter = BlockRefCounter() @@ -231,9 +227,6 @@ def _burst_drop_via_counter(num_blocks, use_clear, timeout_s=60): for ref in refs: assert core_worker.add_object_out_of_scope_callback(ref, on_freed) - if use_clear: - counter.clear() - id_binaries = [ref.binary() for ref in refs] drop_time = time.perf_counter() del refs, ref @@ -243,32 +236,9 @@ def _burst_drop_via_counter(num_blocks, use_clear, timeout_s=60): f"Only {len(fire_times)}/{num_blocks} callbacks fired within {timeout_s}s" ) - return _compute_latencies(fire_times, drop_time, id_binaries) - - -def test_burst_drop_block_ref_counter(num_blocks, timeout_s=60): - """Burst drop through BlockRefCounter without clear(). - - Callbacks do real bookkeeping (lock + set discard + dict decrement). - """ - result = _burst_drop_via_counter(num_blocks, use_clear=False, timeout_s=timeout_s) - print( - f" burst {num_blocks} blocks (no clear): " - f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s max={result['max']:.4f}s" - ) - return result - - -def test_burst_drop_block_ref_counter_with_clear(num_blocks, timeout_s=60): - """Burst drop through BlockRefCounter with clear() first. - - Callbacks no-op (check _registered_ids, return immediately). - Compare with test_burst_drop_block_ref_counter to isolate the cost - of real bookkeeping vs no-op under burst load. - """ - result = _burst_drop_via_counter(num_blocks, use_clear=True, timeout_s=timeout_s) + result = _compute_latencies(fire_times, drop_time, id_binaries) print( - f" burst {num_blocks} blocks (with clear): " + f" burst {num_blocks} blocks (BRC): " f"p50={result['p50']:.4f}s p95={result['p95']:.4f}s max={result['max']:.4f}s" ) return result @@ -307,12 +277,7 @@ def _run_at_scales(name, test_fn, scales): "Burst drop (per-callback latency)", test_burst_drop_per_callback, SCALES ) brc = _run_at_scales( - "Burst drop (BlockRefCounter, no clear)", test_burst_drop_block_ref_counter, SCALES -) -brc_clear = _run_at_scales( - "Burst drop (BlockRefCounter, with clear)", - test_burst_drop_block_ref_counter_with_clear, - SCALES, + "Burst drop (BlockRefCounter)", test_burst_drop_block_ref_counter, SCALES ) print("\n=== Scaling summary (p95) ===") @@ -322,8 +287,7 @@ def _run_at_scales(name, test_fn, scales): ("Registration (us/cb)", reg), ("Burst drain", burst), ("Per-callback", per_cb), - ("BRC (no clear)", brc), - ("BRC (with clear)", brc_clear), + ("BRC", brc), ]: vals = [] for n in SCALES: @@ -370,11 +334,6 @@ def _run_at_scales(name, test_fn, scales): "perf_metric_value": brc[n]["p95"], "perf_metric_type": "LATENCY", }, - { - "perf_metric_name": f"callback_burst_brc_clear_p95_{n}_blocks_s", - "perf_metric_value": brc_clear[n]["p95"], - "perf_metric_type": "LATENCY", - }, ] ) with open(os.environ["TEST_OUTPUT_JSON"], "w") as out_file: From 7bd4bebd30c5ada987a080202cba3cba8f6ccba5 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 20 Jul 2026 16:46:47 -0700 Subject: [PATCH 08/14] Skip empty memory_store Delete and defer callbacks outside mutex Signed-off-by: Sirui Huang --- src/ray/core_worker/core_worker.h | 4 ++- src/ray/core_worker/reference_counter.cc | 35 +++++++++++++++++------- src/ray/core_worker/reference_counter.h | 15 ++++++++-- 3 files changed, 40 insertions(+), 14 deletions(-) diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 8133ea6e15e4..df3aa2caffdd 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -442,7 +442,9 @@ class CoreWorker : public std::enable_shared_from_this { reference_counter_->RemoveLocalReference(object_id, &deleted); // TODO(sang): This seems bad... We should delete the memory store // properly from reference counter. - memory_store_->Delete(deleted); + if (!deleted.empty()) { + memory_store_->Delete(deleted); + } } /// Register a callback to fire when an object goes out of scope or is freed. diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index afaa9c1f2e4b..3fbbf135e470 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -468,12 +468,19 @@ void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id, if (object_id.IsNil()) { return; } - absl::MutexLock lock(&mutex_); - RemoveLocalReferenceInternal(object_id, deleted); + DeferredCallbacks deferred_cbs; + { + absl::MutexLock lock(&mutex_); + RemoveLocalReferenceInternal(object_id, deleted, &deferred_cbs); + } + for (auto &[id, cb] : deferred_cbs) { + cb(id); + } } void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, - std::vector *deleted) { + std::vector *deleted, + DeferredCallbacks *deferred_cbs) { RAY_CHECK(!object_id.IsNil()); auto it = object_id_refs_.find(object_id); if (it == object_id_refs_.end()) { @@ -491,7 +498,7 @@ void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, RAY_LOG(DEBUG) << "Remove local reference " << object_id; PRINT_REF_COUNT(it); if (it->second.RefCount() == 0) { - DeleteReferenceInternal(it, deleted); + DeleteReferenceInternal(it, deleted, deferred_cbs); } else { PRINT_REF_COUNT(it); } @@ -741,7 +748,8 @@ void ReferenceCounter::FreePlasmaObjects(const std::vector &object_ids } void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, - std::vector *deleted) { + std::vector *deleted, + DeferredCallbacks *deferred_cbs) { const ObjectID id = it->first; RAY_LOG(DEBUG) << "Attempting to delete object " << id; if (it->second.RefCount() == 0 && it->second.publish_ref_removed) { @@ -770,10 +778,10 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, // NOTE: a NestedReferenceCount struct is created after the first // mutable_nested() call, but the struct will not be deleted until the // enclosing Reference struct is deleted. - DeleteReferenceInternal(inner_it, deleted); + DeleteReferenceInternal(inner_it, deleted, deferred_cbs); } } - OnObjectOutOfScopeOrFreed(it); + OnObjectOutOfScopeOrFreed(it, deferred_cbs); if (deleted != nullptr) { deleted->push_back(id); } @@ -839,7 +847,8 @@ int64_t ReferenceCounter::EvictLineage(int64_t min_bytes_to_evict) { return lineage_bytes_evicted; } -void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) { +void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, + DeferredCallbacks *deferred_cbs) { RAY_LOG(DEBUG) << "Calling on_object_out_of_scope_or_freed_callbacks for object " << it->first << " num callbacks: " << it->second.on_object_out_of_scope_or_freed_callbacks.size(); @@ -857,8 +866,14 @@ void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) { } } - for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - callback(it->first); + if (deferred_cbs != nullptr) { + for (auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { + deferred_cbs->emplace_back(it->first, std::move(callback)); + } + } else { + for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { + callback(it->first); + } } it->second.on_object_out_of_scope_or_freed_callbacks.clear(); UpdateOwnedObjectCounters(it->first, it->second, /*decrement=*/true); diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h index d9f5b7418d3b..92d05d579ffa 100644 --- a/src/ray/core_worker/reference_counter.h +++ b/src/ray/core_worker/reference_counter.h @@ -539,8 +539,15 @@ class ReferenceCounter : public ReferenceCounterInterface, void UnsetObjectPrimaryCopy(ReferenceTable::iterator it) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + /// Callbacks extracted under the lock to be invoked after release. + using DeferredCallbacks = + std::vector>>; + /// This should be called whenever the object is out of scope or manually freed. - void OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) + /// When `deferred_cbs` is non-null, OOS callbacks are collected into it + /// instead of being invoked inline under the mutex. + void OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, + DeferredCallbacks *deferred_cbs = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Shutdown if all references have gone out of scope and shutdown @@ -654,7 +661,8 @@ class ReferenceCounter : public ReferenceCounterInterface, /// callbacks. Assumes that the entry is in object_id_refs_ and invalidates the /// iterator. void DeleteReferenceInternal(ReferenceTable::iterator entry, - std::vector *deleted) + std::vector *deleted, + DeferredCallbacks *deferred_cbs = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// To respond to the object's owner once we are no longer borrowing it. The @@ -723,7 +731,8 @@ class ReferenceCounter : public ReferenceCounterInterface, /// This method is internal and not thread-safe. mutex_ lock must be held before /// calling this method. void RemoveLocalReferenceInternal(const ObjectID &object_id, - std::vector *deleted) + std::vector *deleted, + DeferredCallbacks *deferred_cbs = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Address of our RPC server. This is used to determine whether we own a From e5c7dd4f801a9d9ad771556d902d400a18ccf89f Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Tue, 21 Jul 2026 16:47:02 -0700 Subject: [PATCH 09/14] profile benchmark Signed-off-by: Sirui Huang --- .../object_store/test_callback_throughput.py | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index 21b87e5014b5..d95356ee7ab5 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -244,6 +244,153 @@ def test_burst_drop_block_ref_counter(num_blocks, timeout_s=60): return result +def _start_profiler(): + """Start a CPU profiler. Tries perf, then py-spy. Returns a handle dict or None.""" + import shutil + import subprocess + + os.makedirs("/tmp/artifacts", exist_ok=True) + + # Try perf first + perf_bin = shutil.which("perf") + if perf_bin: + test = subprocess.run( + [perf_bin, "record", "-o", "/tmp/perf_test.data", "--", "true"], + capture_output=True, + ) + if test.returncode == 0: + fg_dir = "/tmp/FlameGraph" + if not os.path.isdir(fg_dir): + subprocess.run( + [ + "git", + "clone", + "--depth=1", + "https://github.com/brendangregg/FlameGraph", + fg_dir, + ], + capture_output=True, + ) + perf_data = "/tmp/perf.data" + proc = subprocess.Popen( + [ + perf_bin, + "record", + "-g", + "--call-graph", + "dwarf,16384", + "-F", + "4999", + "-p", + str(os.getpid()), + "-o", + perf_data, + ], + stderr=subprocess.DEVNULL, + ) + time.sleep(0.2) + return { + "type": "perf", + "proc": proc, + "bin": perf_bin, + "data": perf_data, + "fg_dir": fg_dir, + } + else: + print(" perf available but cannot record (container restrictions?)") + + # Fall back to py-spy + subprocess.run(["pip", "install", "py-spy"], capture_output=True) + pyspy_bin = shutil.which("py-spy") + if pyspy_bin: + svg_path = "/tmp/artifacts/flamegraph.svg" + proc = subprocess.Popen( + [ + pyspy_bin, + "record", + "--native", + "--output", + svg_path, + "--pid", + str(os.getpid()), + "--rate", + "1000", + ], + stderr=subprocess.DEVNULL, + ) + time.sleep(0.5) + return {"type": "pyspy", "proc": proc, "svg": svg_path} + + return None + + +def _stop_profiler(handle, num_blocks, drain_time): + """Stop the profiler and generate the flame graph.""" + import signal + import subprocess + + if handle is None: + return + + handle["proc"].send_signal(signal.SIGINT) + handle["proc"].wait() + + label = handle["type"] + print(f" [{label}] Drain time: {drain_time:.4f}s ({num_blocks} blocks)") + + if handle["type"] == "perf": + collapse = os.path.join(handle["fg_dir"], "stackcollapse-perf.pl") + flamegraph_pl = os.path.join(handle["fg_dir"], "flamegraph.pl") + svg_path = "/tmp/artifacts/flamegraph.svg" + subprocess.run( + f"{handle['bin']} script -i {handle['data']}" + f" | perl {collapse}" + f" | perl {flamegraph_pl} --title 'Burst drop {num_blocks} blocks'" + f" > {svg_path}", + shell=True, + ) + if os.path.isfile(svg_path) and os.path.getsize(svg_path) > 0: + print(f" Flame graph saved to {svg_path}") + else: + txt_path = "/tmp/artifacts/perf_stacks.txt" + subprocess.run( + f"{handle['bin']} script -i {handle['data']} > {txt_path}", + shell=True, + ) + print(f" Flame graph failed; raw stacks at {txt_path}") + + elif handle["type"] == "pyspy": + svg_path = handle["svg"] + if os.path.isfile(svg_path) and os.path.getsize(svg_path) > 0: + print(f" Flame graph saved to {svg_path}") + else: + print(" py-spy flame graph generation failed") + + +def profile_burst_drop(num_blocks=10000, timeout_s=60): + """Profile a 10k burst drop and generate a flame graph. + + Tries perf first (better C++ stacks), falls back to py-spy. + """ + core_worker = ray._private.worker.global_worker.core_worker + on_freed, fire_times, done = _make_timing_callback(num_blocks) + refs = _produce_blocks(num_blocks) + for ref in refs: + assert core_worker.add_object_out_of_scope_callback(ref, on_freed) + + profiler = _start_profiler() + if profiler is None: + print(" No profiler available, skipping CPU profile") + return + + drop_time = time.perf_counter() + del refs, ref + done.wait(timeout=timeout_s) + drain_time = time.perf_counter() - drop_time + + _stop_profiler(profiler, num_blocks, drain_time) + + ray.init(address="auto") ray.get( @@ -303,6 +450,9 @@ def _run_at_scales(name, test_fn, scales): "\n Pipeline p95: " + ", ".join(f"{pipeline[n]['p95']:.4f}s ({n})" for n in SCALES) ) +print("\n=== CPU Profile (10k burst drop) ===") +profile_burst_drop() + if "TEST_OUTPUT_JSON" in os.environ: perf_metrics = [ { From 2b6e763267002fa6e3b6812f522e5992b5dbcac1 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Thu, 23 Jul 2026 15:49:27 -0700 Subject: [PATCH 10/14] Defer free object calls Signed-off-by: Sirui Huang --- src/ray/core_worker/reference_counter.cc | 31 +++++++++++++++--------- src/ray/core_worker/reference_counter.h | 18 ++++++++------ 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index 3fbbf135e470..935ac4cbad36 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -468,19 +468,22 @@ void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id, if (object_id.IsNil()) { return; } - DeferredCallbacks deferred_cbs; + DeferredWork deferred; { absl::MutexLock lock(&mutex_); - RemoveLocalReferenceInternal(object_id, deleted, &deferred_cbs); + RemoveLocalReferenceInternal(object_id, deleted, &deferred); } - for (auto &[id, cb] : deferred_cbs) { + for (auto &[id, locations] : deferred.frees) { + free_object_on_nodes_async_(id, locations); + } + for (auto &[id, cb] : deferred.callbacks) { cb(id); } } void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, std::vector *deleted, - DeferredCallbacks *deferred_cbs) { + DeferredWork *deferred) { RAY_CHECK(!object_id.IsNil()); auto it = object_id_refs_.find(object_id); if (it == object_id_refs_.end()) { @@ -498,7 +501,7 @@ void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, RAY_LOG(DEBUG) << "Remove local reference " << object_id; PRINT_REF_COUNT(it); if (it->second.RefCount() == 0) { - DeleteReferenceInternal(it, deleted, deferred_cbs); + DeleteReferenceInternal(it, deleted, deferred); } else { PRINT_REF_COUNT(it); } @@ -749,7 +752,7 @@ void ReferenceCounter::FreePlasmaObjects(const std::vector &object_ids void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, std::vector *deleted, - DeferredCallbacks *deferred_cbs) { + DeferredWork *deferred) { const ObjectID id = it->first; RAY_LOG(DEBUG) << "Attempting to delete object " << id; if (it->second.RefCount() == 0 && it->second.publish_ref_removed) { @@ -778,10 +781,10 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, // NOTE: a NestedReferenceCount struct is created after the first // mutable_nested() call, but the struct will not be deleted until the // enclosing Reference struct is deleted. - DeleteReferenceInternal(inner_it, deleted, deferred_cbs); + DeleteReferenceInternal(inner_it, deleted, deferred); } } - OnObjectOutOfScopeOrFreed(it, deferred_cbs); + OnObjectOutOfScopeOrFreed(it, deferred); if (deleted != nullptr) { deleted->push_back(id); } @@ -848,7 +851,7 @@ int64_t ReferenceCounter::EvictLineage(int64_t min_bytes_to_evict) { } void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, - DeferredCallbacks *deferred_cbs) { + DeferredWork *deferred) { RAY_LOG(DEBUG) << "Calling on_object_out_of_scope_or_freed_callbacks for object " << it->first << " num callbacks: " << it->second.on_object_out_of_scope_or_freed_callbacks.size(); @@ -862,13 +865,17 @@ void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, locations_set.insert(*it->second.pinned_at_node_id_); } if (!locations_set.empty()) { - free_object_on_nodes_async_(it->first, locations_set); + if (deferred != nullptr) { + deferred->frees.emplace_back(it->first, std::move(locations_set)); + } else { + free_object_on_nodes_async_(it->first, locations_set); + } } } - if (deferred_cbs != nullptr) { + if (deferred != nullptr) { for (auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - deferred_cbs->emplace_back(it->first, std::move(callback)); + deferred->callbacks.emplace_back(it->first, std::move(callback)); } } else { for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h index 92d05d579ffa..d9fc75afc21e 100644 --- a/src/ray/core_worker/reference_counter.h +++ b/src/ray/core_worker/reference_counter.h @@ -539,15 +539,17 @@ class ReferenceCounter : public ReferenceCounterInterface, void UnsetObjectPrimaryCopy(ReferenceTable::iterator it) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); - /// Callbacks extracted under the lock to be invoked after release. - using DeferredCallbacks = - std::vector>>; + /// Work extracted under the lock to be invoked after release. + struct DeferredWork { + std::vector>> callbacks; + std::vector>> frees; + }; /// This should be called whenever the object is out of scope or manually freed. - /// When `deferred_cbs` is non-null, OOS callbacks are collected into it - /// instead of being invoked inline under the mutex. + /// When `deferred` is non-null, OOS callbacks and free_object_on_nodes + /// requests are collected into it instead of being invoked under the mutex. void OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, - DeferredCallbacks *deferred_cbs = nullptr) + DeferredWork *deferred = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Shutdown if all references have gone out of scope and shutdown @@ -662,7 +664,7 @@ class ReferenceCounter : public ReferenceCounterInterface, /// iterator. void DeleteReferenceInternal(ReferenceTable::iterator entry, std::vector *deleted, - DeferredCallbacks *deferred_cbs = nullptr) + DeferredWork *deferred = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// To respond to the object's owner once we are no longer borrowing it. The @@ -732,7 +734,7 @@ class ReferenceCounter : public ReferenceCounterInterface, /// calling this method. void RemoveLocalReferenceInternal(const ObjectID &object_id, std::vector *deleted, - DeferredCallbacks *deferred_cbs = nullptr) + DeferredWork *deferred = nullptr) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Address of our RPC server. This is used to determine whether we own a From 839a6deb67e22ff2258c7d1231369ff4537444d3 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 27 Jul 2026 10:43:33 -0700 Subject: [PATCH 11/14] Remove clear path Signed-off-by: Sirui Huang --- python/ray/_raylet.pyx | 16 ----------- python/ray/includes/libcoreworker.pxd | 2 -- src/ray/core_worker/core_worker.cc | 4 --- src/ray/core_worker/core_worker.h | 4 --- src/ray/core_worker/reference_counter.cc | 12 -------- src/ray/core_worker/reference_counter.h | 3 -- .../core_worker/reference_counter_interface.h | 5 ---- .../tests/reference_counter_test.cc | 28 ------------------- 8 files changed, 74 deletions(-) diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 1709c487724f..98564a27aa33 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -4492,22 +4492,6 @@ cdef class CoreWorker: cpython.Py_DECREF(callback) return registered - def remove_object_out_of_scope_callbacks(self, object_id_binary: bytes): - """Remove and fire all out-of-scope callbacks for the given object. - - Firing allows per-callback cleanup (e.g. Py_DECREF). After this - call, no further out-of-scope callbacks will fire for this object. - - Args: - object_id_binary: The object ID as bytes (from ObjectRef.binary()). - - .. warning:: - This is an internal Ray API. Do not use it outside of Ray libraries. - """ - cdef CObjectID c_object_id = CObjectID.FromBinary(object_id_binary) - CCoreWorkerProcess.GetCoreWorker() \ - .RemoveObjectOutOfScopeOrFreedCallbacks(c_object_id) - def get_owner_address(self, ObjectRef object_ref): cdef: CObjectID c_object_id = object_ref.native() diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index 8ad23b56fcb8..ddd4ee570505 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/python/ray/includes/libcoreworker.pxd @@ -260,8 +260,6 @@ cdef extern from "ray/core_worker/core_worker.h" nogil: const CObjectID &object_id, void (*callback)(const CObjectID &, void *) nogil, void *callback_context) - void RemoveObjectOutOfScopeOrFreedCallbacks( - const CObjectID &object_id) CRayStatus CheckObjectOwnedByUs(const CObjectID &object_id) const void PutObjectIntoPlasma(const CRayObject &object, const CObjectID &object_id) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 03846fcb655d..bb1e7965240a 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -2584,10 +2584,6 @@ bool CoreWorker::AddObjectOutOfScopeOrFreedCallback(const ObjectID &object_id, }); } -void CoreWorker::RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) { - reference_counter_->RemoveObjectOutOfScopeOrFreedCallbacks(object_id); -} - Status CoreWorker::CheckObjectOwnedByUs(const ObjectID &object_id) const { if (reference_counter_->OwnedByUs(object_id)) { return Status::OK(); diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index df3aa2caffdd..2d319892b776 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -483,10 +483,6 @@ class CoreWorker : public std::enable_shared_from_this { /// a borrowed object owned by another worker). Status CheckObjectOwnedByUs(const ObjectID &object_id) const; - /// Remove and fire all out-of-scope/freed callbacks for `object_id`. - /// See ReferenceCounterInterface::RemoveObjectOutOfScopeOrFreedCallbacks. - void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id); - int GetMemoryStoreSize() { return memory_store_->Size(); } /// Returns a map of all ObjectIDs currently in scope with a pair of their diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index 935ac4cbad36..952ff2fafb81 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -908,18 +908,6 @@ bool ReferenceCounter::AddObjectRefDeletedCallback( return true; } -void ReferenceCounter::RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) { - absl::MutexLock lock(&mutex_); - auto it = object_id_refs_.find(object_id); - if (it == object_id_refs_.end()) { - return; - } - for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - callback(object_id); - } - it->second.on_object_out_of_scope_or_freed_callbacks.clear(); -} - bool ReferenceCounter::AddObjectOutOfScopeOrFreedCallback( const ObjectID &object_id, const std::function callback) { absl::MutexLock lock(&mutex_); diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h index d9fc75afc21e..c09eced8b27e 100644 --- a/src/ray/core_worker/reference_counter.h +++ b/src/ray/core_worker/reference_counter.h @@ -156,9 +156,6 @@ class ReferenceCounter : public ReferenceCounterInterface, const std::function callback) override ABSL_LOCKS_EXCLUDED(mutex_); - void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) override - ABSL_LOCKS_EXCLUDED(mutex_); - bool AddObjectRefDeletedCallback( const ObjectID &object_id, std::function callback) override ABSL_LOCKS_EXCLUDED(mutex_); diff --git a/src/ray/core_worker/reference_counter_interface.h b/src/ray/core_worker/reference_counter_interface.h index 2ba68abb0d15..67eb9bfc67b3 100644 --- a/src/ray/core_worker/reference_counter_interface.h +++ b/src/ray/core_worker/reference_counter_interface.h @@ -329,11 +329,6 @@ class ReferenceCounterInterface { const ObjectID &object_id, const std::function callback) = 0; - /// Removes and fires all out-of-scope/freed callbacks for `object_id`. - /// Firing allows per-callback cleanup (e.g. Py_DECREF on the Python side). - /// After this call, no further out-of-scope callbacks will fire for this object. - virtual void RemoveObjectOutOfScopeOrFreedCallbacks(const ObjectID &object_id) = 0; - /// Stores the callback that will be run when the object reference is deleted /// from the reference table (all refs including lineage ref count go to 0). /// There could be multiple callbacks for the same object due to retries and we store diff --git a/src/ray/core_worker/tests/reference_counter_test.cc b/src/ray/core_worker/tests/reference_counter_test.cc index 7182eb021c81..0662a6e49e0a 100644 --- a/src/ray/core_worker/tests/reference_counter_test.cc +++ b/src/ray/core_worker/tests/reference_counter_test.cc @@ -634,34 +634,6 @@ TEST_F(ReferenceCountTest, TestUnreconstructableObjectOutOfScope) { ASSERT_TRUE(*out_of_scope); } -TEST_F(ReferenceCountTest, TestRemoveObjectOutOfScopeCallbacks) { - ObjectID id = ObjectID::FromRandom(); - rpc::Address address; - address.set_ip_address("1234"); - - auto callback_count = std::make_shared(0); - auto callback = [&](const ObjectID &object_id) { (*callback_count)++; }; - - // Register a callback on an owned, in-scope object. - rc->AddOwnedObject(id, - {}, - address, - "", - 0, - LineageReconstructionEligibility::INELIGIBLE_PUT, - /*add_local_ref=*/true); - ASSERT_TRUE(rc->AddObjectOutOfScopeOrFreedCallback(id, callback)); - - // Remove fires the callback. - rc->RemoveObjectOutOfScopeOrFreedCallbacks(id); - ASSERT_EQ(*callback_count, 1); - - // The object going out of scope should not fire the callback again. - std::vector out; - rc->RemoveLocalReference(id, &out); - ASSERT_EQ(*callback_count, 1); -} - // Tests call site tracking and ability to update object size. TEST_F(ReferenceCountTest, TestReferenceStats) { ObjectID id1 = ObjectID::FromRandom(); From 2ab8528973eb9d712076caf1c55768bc6dfe7bfb Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 27 Jul 2026 10:50:28 -0700 Subject: [PATCH 12/14] Revert Skip empty memory_store Delete and defer callbacks outside mutex Signed-off-by: Sirui Huang --- src/ray/core_worker/core_worker.h | 4 +-- src/ray/core_worker/reference_counter.cc | 44 ++++++------------------ src/ray/core_worker/reference_counter.h | 17 ++------- 3 files changed, 15 insertions(+), 50 deletions(-) diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 2d319892b776..83eb1cec3c65 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -442,9 +442,7 @@ class CoreWorker : public std::enable_shared_from_this { reference_counter_->RemoveLocalReference(object_id, &deleted); // TODO(sang): This seems bad... We should delete the memory store // properly from reference counter. - if (!deleted.empty()) { - memory_store_->Delete(deleted); - } + memory_store_->Delete(deleted); } /// Register a callback to fire when an object goes out of scope or is freed. diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index 952ff2fafb81..b28d756bbb1b 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -468,22 +468,12 @@ void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id, if (object_id.IsNil()) { return; } - DeferredWork deferred; - { - absl::MutexLock lock(&mutex_); - RemoveLocalReferenceInternal(object_id, deleted, &deferred); - } - for (auto &[id, locations] : deferred.frees) { - free_object_on_nodes_async_(id, locations); - } - for (auto &[id, cb] : deferred.callbacks) { - cb(id); - } + absl::MutexLock lock(&mutex_); + RemoveLocalReferenceInternal(object_id, deleted); } void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, - std::vector *deleted, - DeferredWork *deferred) { + std::vector *deleted) { RAY_CHECK(!object_id.IsNil()); auto it = object_id_refs_.find(object_id); if (it == object_id_refs_.end()) { @@ -501,7 +491,7 @@ void ReferenceCounter::RemoveLocalReferenceInternal(const ObjectID &object_id, RAY_LOG(DEBUG) << "Remove local reference " << object_id; PRINT_REF_COUNT(it); if (it->second.RefCount() == 0) { - DeleteReferenceInternal(it, deleted, deferred); + DeleteReferenceInternal(it, deleted); } else { PRINT_REF_COUNT(it); } @@ -751,8 +741,7 @@ void ReferenceCounter::FreePlasmaObjects(const std::vector &object_ids } void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, - std::vector *deleted, - DeferredWork *deferred) { + std::vector *deleted) { const ObjectID id = it->first; RAY_LOG(DEBUG) << "Attempting to delete object " << id; if (it->second.RefCount() == 0 && it->second.publish_ref_removed) { @@ -781,10 +770,10 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it, // NOTE: a NestedReferenceCount struct is created after the first // mutable_nested() call, but the struct will not be deleted until the // enclosing Reference struct is deleted. - DeleteReferenceInternal(inner_it, deleted, deferred); + DeleteReferenceInternal(inner_it, deleted); } } - OnObjectOutOfScopeOrFreed(it, deferred); + OnObjectOutOfScopeOrFreed(it); if (deleted != nullptr) { deleted->push_back(id); } @@ -850,8 +839,7 @@ int64_t ReferenceCounter::EvictLineage(int64_t min_bytes_to_evict) { return lineage_bytes_evicted; } -void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, - DeferredWork *deferred) { +void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) { RAY_LOG(DEBUG) << "Calling on_object_out_of_scope_or_freed_callbacks for object " << it->first << " num callbacks: " << it->second.on_object_out_of_scope_or_freed_callbacks.size(); @@ -865,22 +853,12 @@ void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, locations_set.insert(*it->second.pinned_at_node_id_); } if (!locations_set.empty()) { - if (deferred != nullptr) { - deferred->frees.emplace_back(it->first, std::move(locations_set)); - } else { - free_object_on_nodes_async_(it->first, locations_set); - } + free_object_on_nodes_async_(it->first, locations_set); } } - if (deferred != nullptr) { - for (auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - deferred->callbacks.emplace_back(it->first, std::move(callback)); - } - } else { - for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - callback(it->first); - } + for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { + callback(it->first); } it->second.on_object_out_of_scope_or_freed_callbacks.clear(); UpdateOwnedObjectCounters(it->first, it->second, /*decrement=*/true); diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h index c09eced8b27e..0a1b4ceacb4d 100644 --- a/src/ray/core_worker/reference_counter.h +++ b/src/ray/core_worker/reference_counter.h @@ -536,17 +536,8 @@ class ReferenceCounter : public ReferenceCounterInterface, void UnsetObjectPrimaryCopy(ReferenceTable::iterator it) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); - /// Work extracted under the lock to be invoked after release. - struct DeferredWork { - std::vector>> callbacks; - std::vector>> frees; - }; - /// This should be called whenever the object is out of scope or manually freed. - /// When `deferred` is non-null, OOS callbacks and free_object_on_nodes - /// requests are collected into it instead of being invoked under the mutex. - void OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it, - DeferredWork *deferred = nullptr) + void OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Shutdown if all references have gone out of scope and shutdown @@ -660,8 +651,7 @@ class ReferenceCounter : public ReferenceCounterInterface, /// callbacks. Assumes that the entry is in object_id_refs_ and invalidates the /// iterator. void DeleteReferenceInternal(ReferenceTable::iterator entry, - std::vector *deleted, - DeferredWork *deferred = nullptr) + std::vector *deleted) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// To respond to the object's owner once we are no longer borrowing it. The @@ -730,8 +720,7 @@ class ReferenceCounter : public ReferenceCounterInterface, /// This method is internal and not thread-safe. mutex_ lock must be held before /// calling this method. void RemoveLocalReferenceInternal(const ObjectID &object_id, - std::vector *deleted, - DeferredWork *deferred = nullptr) + std::vector *deleted) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Address of our RPC server. This is used to determine whether we own a From 8623acdf3404358f18361c650e72edc87352a165 Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 27 Jul 2026 10:55:57 -0700 Subject: [PATCH 13/14] Remove profiling Signed-off-by: Sirui Huang --- .../object_store/test_callback_throughput.py | 150 ------------------ 1 file changed, 150 deletions(-) diff --git a/release/benchmarks/object_store/test_callback_throughput.py b/release/benchmarks/object_store/test_callback_throughput.py index d95356ee7ab5..21b87e5014b5 100644 --- a/release/benchmarks/object_store/test_callback_throughput.py +++ b/release/benchmarks/object_store/test_callback_throughput.py @@ -244,153 +244,6 @@ def test_burst_drop_block_ref_counter(num_blocks, timeout_s=60): return result -def _start_profiler(): - """Start a CPU profiler. Tries perf, then py-spy. Returns a handle dict or None.""" - import shutil - import subprocess - - os.makedirs("/tmp/artifacts", exist_ok=True) - - # Try perf first - perf_bin = shutil.which("perf") - if perf_bin: - test = subprocess.run( - [perf_bin, "record", "-o", "/tmp/perf_test.data", "--", "true"], - capture_output=True, - ) - if test.returncode == 0: - fg_dir = "/tmp/FlameGraph" - if not os.path.isdir(fg_dir): - subprocess.run( - [ - "git", - "clone", - "--depth=1", - "https://github.com/brendangregg/FlameGraph", - fg_dir, - ], - capture_output=True, - ) - perf_data = "/tmp/perf.data" - proc = subprocess.Popen( - [ - perf_bin, - "record", - "-g", - "--call-graph", - "dwarf,16384", - "-F", - "4999", - "-p", - str(os.getpid()), - "-o", - perf_data, - ], - stderr=subprocess.DEVNULL, - ) - time.sleep(0.2) - return { - "type": "perf", - "proc": proc, - "bin": perf_bin, - "data": perf_data, - "fg_dir": fg_dir, - } - else: - print(" perf available but cannot record (container restrictions?)") - - # Fall back to py-spy - subprocess.run(["pip", "install", "py-spy"], capture_output=True) - pyspy_bin = shutil.which("py-spy") - if pyspy_bin: - svg_path = "/tmp/artifacts/flamegraph.svg" - proc = subprocess.Popen( - [ - pyspy_bin, - "record", - "--native", - "--output", - svg_path, - "--pid", - str(os.getpid()), - "--rate", - "1000", - ], - stderr=subprocess.DEVNULL, - ) - time.sleep(0.5) - return {"type": "pyspy", "proc": proc, "svg": svg_path} - - return None - - -def _stop_profiler(handle, num_blocks, drain_time): - """Stop the profiler and generate the flame graph.""" - import signal - import subprocess - - if handle is None: - return - - handle["proc"].send_signal(signal.SIGINT) - handle["proc"].wait() - - label = handle["type"] - print(f" [{label}] Drain time: {drain_time:.4f}s ({num_blocks} blocks)") - - if handle["type"] == "perf": - collapse = os.path.join(handle["fg_dir"], "stackcollapse-perf.pl") - flamegraph_pl = os.path.join(handle["fg_dir"], "flamegraph.pl") - svg_path = "/tmp/artifacts/flamegraph.svg" - subprocess.run( - f"{handle['bin']} script -i {handle['data']}" - f" | perl {collapse}" - f" | perl {flamegraph_pl} --title 'Burst drop {num_blocks} blocks'" - f" > {svg_path}", - shell=True, - ) - if os.path.isfile(svg_path) and os.path.getsize(svg_path) > 0: - print(f" Flame graph saved to {svg_path}") - else: - txt_path = "/tmp/artifacts/perf_stacks.txt" - subprocess.run( - f"{handle['bin']} script -i {handle['data']} > {txt_path}", - shell=True, - ) - print(f" Flame graph failed; raw stacks at {txt_path}") - - elif handle["type"] == "pyspy": - svg_path = handle["svg"] - if os.path.isfile(svg_path) and os.path.getsize(svg_path) > 0: - print(f" Flame graph saved to {svg_path}") - else: - print(" py-spy flame graph generation failed") - - -def profile_burst_drop(num_blocks=10000, timeout_s=60): - """Profile a 10k burst drop and generate a flame graph. - - Tries perf first (better C++ stacks), falls back to py-spy. - """ - core_worker = ray._private.worker.global_worker.core_worker - on_freed, fire_times, done = _make_timing_callback(num_blocks) - refs = _produce_blocks(num_blocks) - for ref in refs: - assert core_worker.add_object_out_of_scope_callback(ref, on_freed) - - profiler = _start_profiler() - if profiler is None: - print(" No profiler available, skipping CPU profile") - return - - drop_time = time.perf_counter() - del refs, ref - done.wait(timeout=timeout_s) - drain_time = time.perf_counter() - drop_time - - _stop_profiler(profiler, num_blocks, drain_time) - - ray.init(address="auto") ray.get( @@ -450,9 +303,6 @@ def _run_at_scales(name, test_fn, scales): "\n Pipeline p95: " + ", ".join(f"{pipeline[n]['p95']:.4f}s ({n})" for n in SCALES) ) -print("\n=== CPU Profile (10k burst drop) ===") -profile_burst_drop() - if "TEST_OUTPUT_JSON" in os.environ: perf_metrics = [ { From 106893984c8e4d4fa66e4277bda60155c188545a Mon Sep 17 00:00:00 2001 From: Sirui Huang Date: Mon, 27 Jul 2026 22:07:18 -0700 Subject: [PATCH 14/14] Move callback forward Signed-off-by: Sirui Huang --- src/ray/core_worker/reference_counter.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc index b28d756bbb1b..cc5a7a255369 100644 --- a/src/ray/core_worker/reference_counter.cc +++ b/src/ray/core_worker/reference_counter.cc @@ -843,6 +843,11 @@ void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) { RAY_LOG(DEBUG) << "Calling on_object_out_of_scope_or_freed_callbacks for object " << it->first << " num callbacks: " << it->second.on_object_out_of_scope_or_freed_callbacks.size(); + for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { + callback(it->first); + } + it->second.on_object_out_of_scope_or_freed_callbacks.clear(); + // Only the owner is allowed to broadcast a free for an object. Borrowers // also reach this code path when their local refs drop to zero, but they // must not tell the cluster to evict an object that is still owned @@ -857,10 +862,6 @@ void ReferenceCounter::OnObjectOutOfScopeOrFreed(ReferenceTable::iterator it) { } } - for (const auto &callback : it->second.on_object_out_of_scope_or_freed_callbacks) { - callback(it->first); - } - it->second.on_object_out_of_scope_or_freed_callbacks.clear(); UpdateOwnedObjectCounters(it->first, it->second, /*decrement=*/true); UnsetObjectPrimaryCopy(it); UpdateOwnedObjectCounters(it->first, it->second, /*decrement=*/false);