[Data] [Core] [2.5/n] Remove BlockRefCounter.clear() and expand callback throughput benchmark

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()

python/ray/data/_internal/execution/streaming_executor.py

@@ -366,9 +366,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)

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."""