为 Ray actor 函数实现缓存

Question

我的目标是让下面的代码在大约 0.3 秒而不是 0.5 秒内执行。我尝试过使用functools.lru_cache、toolz.functoolz.memoize和kids.cache.cache上的装饰器foo，但这些都不起作用（错误消息或未正确执行）。我该怎么做才能使这项工作成功？

import ray


@ray.remote
class Foo:
    def foo(self, x):
        print("executing foo with {}".format(x))
        time.sleep(0.1)


ray.init()
f = Foo.remote()
s = time.time()
ray.get([f.foo.remote(x=i) for i in [1, 2, 1, 4, 1]])
print(time.time()-s)
ray.shutdown()

Answer 1

一般警告：如果函数产生副作用，缓存任意函数调用可能会很危险。

在这种情况下，大概您希望程序输出

executing foo with 1 
executing foo with 2 
executing foo with 4 

您提到的其他缓存工具往往不能很好地与 Ray 配合使用，因为它们尝试将缓存存储在某种全局状态中，并且不会将该状态存储在可以分布式方式访问的位置。由于您已经有一个演员，因此您可以将全局状态存储在演员中。

@ray.remote
class Foo:
    def __init__(self):
        self.foo_cache = {}

    def foo(self, x):
        def real_foo(x):
            print("executing foo with {}".format(x))
            time.sleep(0.1)
        if x not in self.foo_cache:
            self.foo_cache[x] = real_foo(x)
        return self.foo_cache[x]

这是一种非常通用的缓存技术，这里唯一重要的区别是我们必须将状态存储在参与者中。

更通用的方法

我们还可以通过定义通用函数缓存来将这种方法推广到任何 Ray 函数：

@ray.remote
class FunctionCache:
    def __init__(self, func):
        self.func = ray.remote(func)
        self.cache = {}

    def call(self, *args, **kwargs):
        if (args, kwargs) not in cache:
            cache[(args, kwargs)] = self.func(*args, **kwargs)
        return cache[(args, kwargs)]

然后为了清理我们使用它的方式，我们可以定义一个装饰器：

class RemoteFunctionLookAlike:
    def __init__(self, func):
        self.func = func

    def remote(self, *args, **kwargs):
        return self.func(*args, **kwargs)


def ray_cache(func):
    cache = FunctionCache.remote(func)
    def call_with_cache(*args, **kwargs):
        return cache.call.remote(*args, **kwargs)
    return RayFunctionLookAlike(call_with_cache)

最后，要使用此缓存：

@ray_cache
def foo(x):
    print("Function called!")
    return abc

ray.get([foo.remote("constant") for _ in range(100)]) # Only prints "Function called!" once.