[core][__on_object_freed__][1/n] Detect the __on_object_freed__ hook on actor classes

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

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 when that object is freed.
+# `__on_object_freed__` is a private, experimental hook right now. It is not part of
+# the public API.
+_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 "
+            "be callable with 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}'"

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

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

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, "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"),
+    ],
+)
+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__]))

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

src/ray/core_worker/core_worker.h

@@ -293,6 +293,25 @@ class CoreWorker : public std::enable_shared_from_this<CoreWorker> {
 
   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<CoreWorker> {
   /// Callback to free an RDT object when it is out of scope.
   std::function<void(const ObjectID &)> 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<bool> connected_{true};
   std::atomic<bool> event_loops_running_{false};