[karticam]

This is PR 4 of a series of PRs that add the __on_object_freed__ hook: an actor that defines the hook will be told when an object it created is freed. This PR runs the hook. The raylet on the producer's node calls it once: when an actor task on that same node consumes the object, or otherwise when the object's last copy on the node is deleted, using the producer recorded by #67377 and the OnObjectFreed RPC from #67308.

The complete design doc is here: https://docs.google.com/document/d/1ZJhgRmwVyKfNCR21CWf6xMcc0B6HJwSPQsfebTRvnnE/edit?tab=t.0#heading=h.7wcod1f9afpg

Changes

1. An object can have two copies on a node, and whichever is deleted second runs the hook.
   • In-memory copy deleted: NodeManager::HandleObjectMissing runs the hook unless the object is spilled or being spilled (LocalObjectManager::HasDiskCopy). The raylet reads the copy's producer before ObjectManager::HandleObjectDeleted erases it, and passes it in.
   • Spilled copy deleted: LocalObjectManager calls the new on_spilled_copy_deleted callback with the producer recorded at pin, when it deletes a freed object's spilled copy. NodeManager::HandleSpilledCopyDeleted runs the hook unless an in-memory copy is still there.
   • Both run on the raylet's main thread, so exactly one of them sees the other copy gone.
2. A restored or received copy takes over the producer.
   • When a spilled object's in-memory copy has been evicted and a reader on the node needs it again, the node restores it from its spilled copy, or pulls it from another node that has it in memory. That new copy records no producer, since the producer worker didn't create it ([core][__on_object_freed__][3/n] Record which worker created each plasma object #67377).
   • Either copy can be deleted first: the free queues the in-memory copy and the spilled copy separately. When the spilled copy goes first, it sees the in-memory copy and leaves the hook to it, but that copy has no producer, so without this change the hook would never run.
   • So when a spilled copy is deleted and an in-memory copy remains, NodeManager::HandleSpilledCopyDeleted records the producer from the pin on that copy (ObjectManager::SetLocalObjectProducerWorkerId), and deleting the copy runs the hook.
   • This is common: ray.get, and ray.wait with the default fetch_local=True, restore an evicted spilled object on the node.
3. NodeManager::FireOnObjectFreed sends OnObjectFreed to the producer worker. If that worker has exited, so the raylet no longer has it registered (WorkerPool::GetRegisteredWorker returns null), it does nothing. The request isn't retried.
4. Same-node consumption runs the hook at once. A consume_once object moves to the actor task that consumes it, and that task can run on the producer's node or on another node.
   • Consumed on another node: the object is pulled to the consumer's node and pinned there, and the owner frees the producer's copy. Deleting that copy runs the hook on the producer's node, through points 1 to 3. The consumer's node never runs it, since its copy records no producer.
   • Consumed on the producer's node: nothing is copied or deleted, so the last-copy rule alone would run the hook only once the object goes out of scope and is freed. Instead, NodeManager::HandleMovedActorCallArgs runs it for each moved arg created on this node, before any pull or restore. It first reads the producer from the object manager's entry for the in-memory copy (ObjectManager::GetLocalObjectProducerWorkerId), which is set when the copy is sealed, so it has the producer even if the producer's pin arrives after the consumption. If that has none (which could happen when the object is spilled and evicted and thus no copy in ObjectManager), it reads the pinned object's LocalObjectInfo (LocalObjectManager::GetProducerWorkerId), which still has the producer after the in-memory copy of a spilled object is evicted.
   • objects_freed_at_consume_ remembers the objects whose hook ran at a same-node consumption, and the last copy's deletion removes them without running the hook again.

[gemini-code-assist]

Code Review

This pull request introduces changes to ensure that the __on_object_freed__ hook is reliably executed when the last copy of an object (either in-memory or spilled to disk) is deleted. It updates LocalObjectManager, ObjectManager, and NodeManager to track disk copies and coordinate firing the hook, and adds comprehensive Python and C++ tests. The review feedback highlights three important issues: a correctness edge case where the hook is leaked if an object is freed during a failing spill, a potential dangling reference in local_object_manager.cc from holding a reference to a map element during callback execution, and a data race in the test mock FakeCoreWorkerClient due to non-thread-safe access to on_object_freed_requests.

[gemini-code-assist]

There is a correctness bug/edge case here where the __on_object_freed__ hook will be leaked (never run) if an object is freed while it is being spilled, and then the spill subsequently fails.

Step-by-step flow:

1. An object is currently being spilled (it is in objects_pending_spill_, so HasDiskCopy returns true).
2. The object is freed by the application, calling ReleaseFreedLocalObject which deletes the object from plasma.
3. This triggers HandleObjectMissing in NodeManager.
4. HandleObjectMissing checks HasDiskCopy(object_id). Since it is true (pending spill), it skips running the hook, expecting the spilled copy's deletion to run it later.
5. The spill subsequently fails.
6. Since the spill failed, no spilled copy is created on disk, and on_spilled_copy_deleted_ is never called.
7. The hook is never run, resulting in a leak of the hook.

Recommendation:

On spill failure, if the object has already been freed (i.e., it is no longer in plasma and was released), we should trigger the hook (e.g., by invoking on_spilled_copy_deleted_ or handling it in the spill failure callback).

[gemini-code-assist]

To prevent potential dangling reference issues, copy producer_worker_id by value instead of taking a reference. The callback on_spilled_copy_deleted_ executes arbitrary logic (including RPCs and manager state updates) which could potentially modify local_objects_ and invalidate the reference.

Suggested change

	const WorkerID &producer_worker_id =
	local_objects_.at(object_id).producer_worker_id_;
	const WorkerID producer_worker_id =
	local_objects_.at(object_id).producer_worker_id_;

[gemini-code-assist]

The on_object_freed_requests vector is not thread-safe. Since OnObjectFreed can be called from multiple threads (e.g., RPC threads or main thread in tests), writing to on_object_freed_requests without a lock is a data race. Consider guarding it with mutex_ as done with other members like callbacks_.

Suggested change

	void OnObjectFreed(const OnObjectFreedRequest &request,
	const ClientCallback<OnObjectFreedReply> &callback) override {}
	const ClientCallback<OnObjectFreedReply> &callback) override {
	on_object_freed_requests.push_back(request);
	}
	void OnObjectFreed(const OnObjectFreedRequest &request,
	const ClientCallback<OnObjectFreedReply> &callback) override {
	absl::MutexLock lock(&mutex_);
	on_object_freed_requests.push_back(request);
	}

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit a14175d. Configure here.}

[cursor]

Missing Doxygen on new APIs

Low Severity

⚠️ Document functions and classes with Doxygen /** ... */ block comments using @ tags (@brief, and @param/@return where applicable). New declarations HasDiskCopy and SetLocalObjectProducerWorkerId omit that documentation.

Additional Locations (2)

• src/ray/raylet/local_object_manager_interface.h#L86-L88
• src/ray/object_manager/object_manager_interface.h#L56-L58

 

^{Triggered by project rule: Bugbot Rules}

^{Reviewed by Cursor Bugbot for commit a14175d. Configure here.}