From 6efbe3e7d90c09d5d86625f7e70a80fd41bc0b01 Mon Sep 17 00:00:00 2001 From: "kartica.modi" Date: Wed, 21 Oct 2026 08:35:38 +0000 Subject: [PATCH 1/2] Detect the __on_object_freed__ hook on actor classes Signed-off-by: kartica.modi --- python/ray/_raylet.pyx | 17 +++ python/ray/actor.py | 49 +++++++ python/ray/includes/libcoreworker.pxd | 3 + python/ray/tests/BUILD.bazel | 1 + python/ray/tests/test_on_object_freed.py | 177 +++++++++++++++++++++++ src/ray/core_worker/core_worker.cc | 5 + src/ray/core_worker/core_worker.h | 24 +++ 7 files changed, 276 insertions(+) create mode 100644 python/ray/tests/test_on_object_freed.py diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index ed9bda1c9284..ce8b77ca7ad4 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -2531,6 +2531,15 @@ cdef execute_task_with_cancellation_handler( actor = actor_class.__new__(actor_class) worker.actors[actor_id] = actor + # `__on_object_freed__` is a private, experimental hook right now. It is not + # part of the public API and may change or be removed. + on_object_freed_hook = getattr(actor_class, "__on_object_freed__", None) + if on_object_freed_hook is not None: + hook_concurrency_group = getattr( + on_object_freed_hook, "__ray_concurrency_group__", None) or "" + CCoreWorkerProcess.GetCoreWorker().SetProducerFreedCallback( + True, hook_concurrency_group.encode("utf-8")) + # Record the actor class via :actor_name: magic token in the log. # # (Phase 1): this covers code run before __init__ finishes. @@ -3504,6 +3513,14 @@ cdef class CoreWorker: return ActorID( CCoreWorkerProcess.GetCoreWorker().GetActorId().Binary()) + def get_producer_freed_callback(self): + """Return whether this worker's actor class defines the + `__on_object_freed__` hook, and the concurrency group it runs in.""" + return ( + CCoreWorkerProcess.GetCoreWorker().ProducerFreedCallbackEnabled(), + CCoreWorkerProcess.GetCoreWorker() + .ProducerFreedCallbackConcurrencyGroup().decode("utf-8")) + def get_actor_name(self): return CCoreWorkerProcess.GetCoreWorker().GetActorName() diff --git a/python/ray/actor.py b/python/ray/actor.py index def870cc356b..94f051ea5bb1 100644 --- a/python/ray/actor.py +++ b/python/ray/actor.py @@ -1258,6 +1258,44 @@ def bind(self, *args, **kwargs): return self._actor_method._bind(args=args, kwargs=kwargs, **self._options) +# Called by Ray with the ID of an object the actor created once that object is freed. +# `__on_object_freed__` is a private, experimental hook right now. It is not part of +# the public API and may change or be removed. +_ON_OBJECT_FREED_HOOK = "__on_object_freed__" + + +def _validate_on_object_freed_hook(modified_class: type) -> None: + """Check that the actor class's `__on_object_freed__` hook, if defined, is a + method Ray can call with one argument, the freed object's ID as a hex string.""" + if not hasattr(modified_class, _ON_OBJECT_FREED_HOOK): + return + hook = getattr(modified_class, _ON_OBJECT_FREED_HOOK) + if not is_function_or_method(hook): + raise TypeError( + f"{_ON_OBJECT_FREED_HOOK} on actor class {modified_class.__name__} must " + f"be a method, got {type(hook).__name__}." + ) + unwrapped = inspect.unwrap(hook) + if inspect.isgeneratorfunction(unwrapped) or inspect.isasyncgenfunction(unwrapped): + raise TypeError( + f"{_ON_OBJECT_FREED_HOOK} on actor class {modified_class.__name__} can't " + "be a generator." + ) + # An instance method also takes `self`; class and static methods don't. + is_bound = is_class_method(hook) or is_static_method( + modified_class, _ON_OBJECT_FREED_HOOK + ) + hook_signature = inspect.signature(hook) + try: + hook_signature.bind(*(("",) if is_bound else (None, ""))) + except TypeError: + raise TypeError( + f"{_ON_OBJECT_FREED_HOOK} on actor class {modified_class.__name__} must " + "take exactly one argument, the freed object's ID as a hex string, got " + f"{_ON_OBJECT_FREED_HOOK}{hook_signature}." + ) from None + + class _ActorClassMethodMetadata(object): """Metadata for all methods in an actor class. This data can be cached. @@ -1302,6 +1340,9 @@ def create( if cached_meta is not None: return cached_meta + if modified_class is not None: + _validate_on_object_freed_hook(modified_class) + # Create an instance without __init__ called. self = cls.__new__(cls) @@ -2460,6 +2501,8 @@ def __init__( # Instead, we will lazily bind this ActorHandle to the _ActorMethodMetadata when a method is invoked. self._method_shells = {} for method_name, method_signature in self._ray_method_signatures.items(): + if method_name == _ON_OBJECT_FREED_HOOK: + continue self._method_shells[method_name] = _ActorMethodMetadata( method_name=method_name, num_returns=self._ray_method_num_returns.get(method_name, None), @@ -2670,6 +2713,12 @@ def __getattr__(self, item: str) -> Never: if item in self._method_shells: return self._method_shells[item].bind(self) + if item == _ON_OBJECT_FREED_HOOK: + raise AttributeError( + f"{_ON_OBJECT_FREED_HOOK} is called by Ray when an object the actor " + "created is freed, and can't be called through an actor handle." + ) + if not self._ray_is_cross_language: raise AttributeError( f"'{type(self).__name__}' object has no attribute '{item}'" diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index bf7e899b2b96..4728142ed461 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/python/ray/includes/libcoreworker.pxd @@ -250,6 +250,9 @@ cdef extern from "ray/core_worker/core_worker.h" nogil: c_bool GetCurrentTaskRetryExceptions() CPlacementGroupID GetCurrentPlacementGroupId() const CWorkerID GetWorkerID() + void SetProducerFreedCallback(c_bool enabled, c_string concurrency_group) + c_bool ProducerFreedCallbackEnabled() const + const c_string &ProducerFreedCallbackConcurrencyGroup() const c_bool ShouldCaptureChildTasksInPlacementGroup() CActorID GetActorId() const const c_string GetActorName() diff --git a/python/ray/tests/BUILD.bazel b/python/ray/tests/BUILD.bazel index 210b75f0a6e8..c244f340556b 100644 --- a/python/ray/tests/BUILD.bazel +++ b/python/ray/tests/BUILD.bazel @@ -564,6 +564,7 @@ py_test_module_list( "test_numba.py", "test_object_out_of_scope_callback.py", "test_object_spilling_no_asan.py", + "test_on_object_freed.py", "test_open_telemetry_metric_recorder.py", "test_placement_group_metrics.py", "test_protobuf_compatibility.py", diff --git a/python/ray/tests/test_on_object_freed.py b/python/ray/tests/test_on_object_freed.py new file mode 100644 index 000000000000..d85fa0fda855 --- /dev/null +++ b/python/ray/tests/test_on_object_freed.py @@ -0,0 +1,177 @@ +import sys + +import pytest + +import ray + + +@ray.remote +class NoHook: + def callback_state(self): + return ( + ray._private.worker.global_worker.core_worker.get_producer_freed_callback() + ) + + +@ray.remote +class SyncHook: + def __on_object_freed__(self, object_id): + pass + + def callback_state(self): + return ( + ray._private.worker.global_worker.core_worker.get_producer_freed_callback() + ) + + +@ray.remote +class AsyncHook: + async def __on_object_freed__(self, object_id): + pass + + def callback_state(self): + return ( + ray._private.worker.global_worker.core_worker.get_producer_freed_callback() + ) + + +@ray.remote(concurrency_groups={"freed": 1}) +class GroupHook: + @ray.method(concurrency_group="freed") + def __on_object_freed__(self, object_id): + pass + + def callback_state(self): + return ( + ray._private.worker.global_worker.core_worker.get_producer_freed_callback() + ) + + +class HookBase: + def __on_object_freed__(self, object_id): + pass + + +@ray.remote +class InheritedHook(HookBase): + def callback_state(self): + return ( + ray._private.worker.global_worker.core_worker.get_producer_freed_callback() + ) + + +@pytest.mark.parametrize( + "actor_cls, expected", + [ + (NoHook, (False, "")), + (SyncHook, (True, "")), + (AsyncHook, (True, "")), + (GroupHook, (True, "freed")), + (InheritedHook, (True, "")), + ], +) +def test_hook_detected_at_actor_creation(ray_start_regular_shared, actor_cls, expected): + actor = actor_cls.remote() + assert ray.get(actor.callback_state.remote()) == expected + + +def test_driver_has_no_hook(ray_start_regular_shared): + core_worker = ray._private.worker.global_worker.core_worker + assert core_worker.get_producer_freed_callback() == (False, "") + + +def test_hook_not_callable_through_handle(ray_start_regular_shared): + actor = SyncHook.remote() + with pytest.raises(AttributeError, match="is called by Ray"): + actor.__on_object_freed__.remote("00") + assert ray.get(actor.callback_state.remote()) == (True, "") + + +class AsyncDef: + async def __on_object_freed__(self, object_id): + pass + + +class ExtraOptionalArg: + def __on_object_freed__(self, object_id, extra=None): + pass + + +class VarArgs: + def __on_object_freed__(self, *args): + pass + + +class StaticHook: + @staticmethod + def __on_object_freed__(object_id): + pass + + +class ClassHook: + @classmethod + def __on_object_freed__(cls, object_id): + pass + + +@pytest.mark.parametrize( + "cls", [HookBase, AsyncDef, ExtraOptionalArg, VarArgs, StaticHook, ClassHook] +) +def test_valid_hook_signatures_accepted(cls): + ray.remote(cls) + + +class NoArg: + def __on_object_freed__(self): + pass + + +class TwoArgs: + def __on_object_freed__(self, object_id, other): + pass + + +class KeywordOnlyArg: + def __on_object_freed__(self, *, object_id): + pass + + +class StaticNoArg: + @staticmethod + def __on_object_freed__(): + pass + + +class Generator: + def __on_object_freed__(self, object_id): + yield object_id + + +class AsyncGenerator: + async def __on_object_freed__(self, object_id): + yield object_id + + +class NotAMethod: + __on_object_freed__ = None + + +@pytest.mark.parametrize( + "cls, match", + [ + (NoArg, "exactly one argument"), + (TwoArgs, "exactly one argument"), + (KeywordOnlyArg, "exactly one argument"), + (StaticNoArg, "exactly one argument"), + (Generator, "can't be a generator"), + (AsyncGenerator, "can't be a generator"), + (NotAMethod, "must be a method"), + ], +) +def test_invalid_hook_rejected(cls, match): + with pytest.raises(TypeError, match=match): + ray.remote(cls) + + +if __name__ == "__main__": + sys.exit(pytest.main(["-sv", __file__])) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 5817ec42fa23..14d43c3e3a04 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -850,6 +850,11 @@ void CoreWorker::ForceExit(const rpc::WorkerExitType exit_type, const WorkerID &CoreWorker::GetWorkerID() const { return worker_context_->GetWorkerID(); } +void CoreWorker::SetProducerFreedCallback(bool enabled, std::string concurrency_group) { + producer_freed_callback_enabled_ = enabled; + producer_freed_callback_concurrency_group_ = std::move(concurrency_group); +} + void CoreWorker::SetCurrentTaskId(const TaskID &task_id, uint64_t attempt_number, const std::string &task_name) { diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 9e8dd425f425..2bf1a5ddf6dd 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -293,6 +293,25 @@ class CoreWorker : public std::enable_shared_from_this { const WorkerID &GetWorkerID() const; + /** + * Record whether this worker's actor class defines the `__on_object_freed__` hook, + * and the concurrency group the hook runs in. Called once on the actor creation + * path, before `__init__` runs. + * + * \param[in] enabled Whether the actor class defines the hook. + * \param[in] concurrency_group The hook's concurrency group. Empty means the + * default group. + */ + void SetProducerFreedCallback(bool enabled, std::string concurrency_group); + + /** Whether this worker's actor class defines the `__on_object_freed__` hook. */ + bool ProducerFreedCallbackEnabled() const { return producer_freed_callback_enabled_; } + + /** The concurrency group the `__on_object_freed__` hook runs in. */ + const std::string &ProducerFreedCallbackConcurrencyGroup() const { + return producer_freed_callback_concurrency_group_; + } + WorkerType GetWorkerType() const { return options_.worker_type; } Language GetLanguage() const { return options_.language; } @@ -2298,6 +2317,11 @@ class CoreWorker : public std::enable_shared_from_this { /// Callback to free an RDT object when it is out of scope. std::function free_actor_object_callback_; + /// Set once on the actor creation path, before any actor method or RPC for the actor + /// runs, so they are read without a lock. + bool producer_freed_callback_enabled_ = false; + std::string producer_freed_callback_concurrency_group_; + // Shutdown synchronization primitives std::atomic connected_{true}; std::atomic event_loops_running_{false}; From 12f072a34a4f1ef3aa39c6f5490e0e56080c4045 Mon Sep 17 00:00:00 2001 From: "kartica.modi" Date: Wed, 21 Oct 2026 23:45:52 +0000 Subject: [PATCH 2/2] Say the __on_object_freed__ hook must be callable with one argument Signed-off-by: kartica.modi --- python/ray/actor.py | 6 +++--- python/ray/tests/test_on_object_freed.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/ray/actor.py b/python/ray/actor.py index 94f051ea5bb1..52172ade9f50 100644 --- a/python/ray/actor.py +++ b/python/ray/actor.py @@ -1258,9 +1258,9 @@ def bind(self, *args, **kwargs): return self._actor_method._bind(args=args, kwargs=kwargs, **self._options) -# Called by Ray with the ID of an object the actor created once that object is freed. +# Called by Ray with the ID of an object the actor created when that object is freed. # `__on_object_freed__` is a private, experimental hook right now. It is not part of -# the public API and may change or be removed. +# the public API. _ON_OBJECT_FREED_HOOK = "__on_object_freed__" @@ -1291,7 +1291,7 @@ def _validate_on_object_freed_hook(modified_class: type) -> None: except TypeError: raise TypeError( f"{_ON_OBJECT_FREED_HOOK} on actor class {modified_class.__name__} must " - "take exactly one argument, the freed object's ID as a hex string, got " + "be callable with one argument, the freed object's ID as a hex string, got " f"{_ON_OBJECT_FREED_HOOK}{hook_signature}." ) from None diff --git a/python/ray/tests/test_on_object_freed.py b/python/ray/tests/test_on_object_freed.py index d85fa0fda855..873cf152181a 100644 --- a/python/ray/tests/test_on_object_freed.py +++ b/python/ray/tests/test_on_object_freed.py @@ -159,10 +159,10 @@ class NotAMethod: @pytest.mark.parametrize( "cls, match", [ - (NoArg, "exactly one argument"), - (TwoArgs, "exactly one argument"), - (KeywordOnlyArg, "exactly one argument"), - (StaticNoArg, "exactly one argument"), + (NoArg, "callable with one argument"), + (TwoArgs, "callable with one argument"), + (KeywordOnlyArg, "callable with one argument"), + (StaticNoArg, "callable with one argument"), (Generator, "can't be a generator"), (AsyncGenerator, "can't be a generator"), (NotAMethod, "must be a method"),