如何清理Ray对象关联的非序列化状态并注册驱逐回调？

问题场景 #

我有一个Ray Actor，用于创建关联非序列化状态（示例为临时目录）的Ray对象，现有代码如下：

class MyObject:
    pass

_non_serializable_states = {}

@ray.remote
class MyCreatorActor:
    def create(self):
        ref = ray.put(MyObject())
        _non_serializable_states[ref] = tempfile.mkdtemp()
        return ref

当该Ray对象的分布式引用计数归0（即ref被驱逐）时，希望在创建者Actor的工作进程中触发以下清理函数：

def on_evicted(ref):
    shutil.rmtree(_non_serializable_states[ref])
    del _non_serializable_states[ref]

请问如何修改MyCreatorActor.create方法注册该回调？Ray是否有类似weakref.ref(obj, callback)的API？

解决方案 #

Ray没有完全等同于Python标准库weakref.ref的直接API，但可以通过两种稳定方式实现需求：

方法1：后台线程监听对象引用释放 #

在Actor内部启动一个守护线程，通过ray.wait循环监听对象引用的释放状态，一旦引用被驱逐就触发清理回调。这种方法兼容性强，适用于所有Ray稳定版本：

import ray
import tempfile
import shutil
import threading
import time

class MyObject:
    pass

_non_serializable_states = {}

def on_evicted(ref):
    shutil.rmtree(_non_serializable_states[ref])
    del _non_serializable_states[ref]

@ray.remote
class MyCreatorActor:
    def __init__(self):
        self._pending_refs = set()
        # 启动后台监听线程
        self._listener_thread = threading.Thread(target=self._monitor_releases, daemon=True)
        self._listener_thread.start()

    def _monitor_releases(self):
        while True:
            if not self._pending_refs:
                time.sleep(0.1)
                continue
            # 等待所有待监听引用被释放
            ready_refs, _ = ray.wait(list(self._pending_refs), num_returns=len(self._pending_refs), timeout=None)
            for ref in ready_refs:
                on_evicted(ref)
                self._pending_refs.remove(ref)

    def create(self):
        ref = ray.put(MyObject())
        temp_dir = tempfile.mkdtemp()
        _non_serializable_states[ref] = temp_dir
        self._pending_refs.add(ref)
        return ref

说明

• Actor初始化时启动的守护线程会持续运行，监听所有待处理的对象引用。
• ray.wait会在对象引用的分布式计数归0（被驱逐）时，将该引用标记为ready，此时即可触发清理逻辑。

方法2：使用Ray实验性对象释放回调API #

如果使用的Ray版本支持（建议v1.13+），可以直接使用ray.experimental.register_object_release_callback为ObjectRef注册释放回调，代码更简洁：

import ray
import tempfile
import shutil

class MyObject:
    pass

_non_serializable_states = {}

def on_evicted(ref):
    shutil.rmtree(_non_serializable_states[ref])
    del _non_serializable_states[ref]

@ray.remote
class MyCreatorActor:
    def create(self):
        ref = ray.put(MyObject())
        temp_dir = tempfile.mkdtemp()
        _non_serializable_states[ref] = temp_dir
        # 注册释放回调
        ray.experimental.register_object_release_callback(ref, on_evicted)
        return ref

说明

• 该API是实验性特性，未来可能有接口变动，使用前需确认版本兼容性。
• 回调会在对象被驱逐时，在创建该引用的Actor工作进程中触发，无需额外线程管理。

内容的提问来源于stack exchange，提问作者Yang Bo