[karticam]

This is PR 1 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 can be used to implement backpressure policies, for example, by identifying that back pressure can be released since the object on this actor is freed.

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

This PR only detects the hook and records it on the actor's worker. Nothing reads the record yet: a later PR stamps it on the objects the actor creates, and another one calls the hook. The hook is private and experimental for now, like _consume_once.

Changes

1. The worker records the hook. On the actor creation path in _raylet.pyx, before the actor class constructor runs, the worker looks up __on_object_freed__ on the actor class. If it is there, it calls the new CoreWorker::SetProducerFreedCallback to record that the hook exists and its concurrency group (empty means the default group). Both values are set before any actor method or RPC for the actor runs, so they are read later without a lock.
2. The hook's signature is checked when the class is decorated. @ray.remote raises a TypeError if the hook isn't a method, is a generator, or can't be called with exactly one argument (it takes the freed object's ID as a hex string). async def, static and class methods are accepted. An async def hook makes the actor an asyncio actor, as any async method does, and a later PR runs it on its concurrency group's event loop.
3. The hook can't be called through an actor handle. It is left out of the handle's methods, and handle.__on_object_freed__ raises an AttributeError. It stays in the class metadata, so its concurrency group is still checked when the actor is created.

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ 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 a2d0bca. Configure here.}

[cursor]

Signature tests belong in unit

Low Severity

💡 This test appears to only test pure Python logic. Consider moving it to tests/unit/.

Additional Locations (1)

• python/ray/tests/test_on_object_freed.py#L158-L173

 

^{Triggered by project rule: Bugbot Rules}

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

[gemini-code-assist]

Code Review

This pull request introduces a private, experimental __on_object_freed__ hook for Ray actors, which is triggered when an object created by the actor is freed. The changes include hook validation in Python, registration of the callback and its concurrency group with the core worker via Cython, and C++ core worker updates to track this state. The review feedback highlights an issue where @classmethod hooks might fail to resolve concurrency group attributes because they are bound methods; it suggests unwrapping them using __func__ and adding corresponding test coverage.

[gemini-code-assist]

If __on_object_freed__ is defined as a @classmethod, getattr(actor_class, "__on_object_freed__") returns a bound method. In Python, bound methods do not delegate attribute lookup to their underlying functions, so any attributes set by decorators like @ray.method(concurrency_group=...) on the underlying function will not be found directly on the bound method.\n\nTo fix this, we should unwrap the bound method to its underlying function using getattr(on_object_freed_hook, "__func__", on_object_freed_hook) before looking up __ray_concurrency_group__.

            unwrapped_hook = getattr(on_object_freed_hook, "__func__", on_object_freed_hook)\n            hook_concurrency_group = getattr(\n                unwrapped_hook, "__ray_concurrency_group__", None) or ""

[gemini-code-assist]

Let's add a test case for a @classmethod hook with a concurrency group to ensure that the concurrency group is correctly detected and registered for class methods as well.

@ray.remote(concurrency_groups={"freed": 1})\nclass GroupHook:\n    @ray.method(concurrency_group="freed")\n    def __on_object_freed__(self, object_id):\n        pass\n\n    def callback_state(self):\n        return (\n            ray._private.worker.global_worker.core_worker.get_producer_freed_callback()\n        )\n\n\n@ray.remote(concurrency_groups={"freed": 1})\nclass GroupClassHook:\n    @classmethod\n    @ray.method(concurrency_group="freed")\n    def __on_object_freed__(cls, object_id):\n        pass\n\n    def callback_state(self):\n        return (\n            ray._private.worker.global_worker.core_worker.get_producer_freed_callback()\n        )

[gemini-code-assist]

Include GroupClassHook in the parameterized test cases to verify that its concurrency group is correctly detected.

@pytest.mark.parametrize(\n    "actor_cls, expected",\n    [\n        (NoHook, (False, "")),\n        (SyncHook, (True, "")),\n        (AsyncHook, (True, "")),\n        (GroupHook, (True, "freed")),\n        (GroupClassHook, (True, "freed")),\n        (InheritedHook, (True, "")),\n    ],\n)