Python Ray:如何在工作人员之间共享变量?

问题描述

我正在开发一个管道,其中包括从机器学习模型中获取预测,并且我正在尝试使用 ray 来加速它。输入可能会重复,所以我想在函数中共享一个缓存,这样这个远程函数的所有工作人员都可以共享对缓存的访问并搜索它并获取值。 类似于下面的内容

@ray.remote
def f(x):
    # create inputs from x
    # do work
    unkNown_y1 = []
    obtained_y1 = []
    for index,y in enumerate(y1):
        key = '|'.join([str(x) for x in y.values()])
        if key in cached:
            obtained_y1.append(cached[key])
        else:
            obtained_y1.append(np.inf)
            unkNown_y1.append(promo)

    unkNown_y2 = []
    obtained_y2 = []
    for index,y in enumerate(y2):
        key = '|'.join([str(x) for x in y.values()])
        if key in cached:
            obtained_y2.append(cached[key])
        else:
            obtained_y2.append(np.inf)
            unkNown_y2.append(baseline)

    kNown_y1,kNown_y2 = predictor.predict(unkNown_y1,unkNown_y2)

    unkNown_index = 0
    for index in range(len(y1)):
        if(obtained_y1[index] == np.inf):
            obtained_y1[index] = kNown_y1[unkNown_index]
            key = '|'.join([str(x) for x in y1[index].values()])
            if not(key in cached):
                cached[key] = obtained_y1[index]
            unkNown_index = unkNown_index+1
    
    unkNown_index = 0
    for index in range(len(y2)):
        if(obtained_y2[index] == np.inf):
            obtained_y2[index] = kNown_y2[unkNown_index]
            key = '|'.join([str(x) for x in y2[index].values()])
            if not(key in cached):
                cached[key] = obtained_y2[index]
            unkNown_index = unkNown_index+1

我尝试通过在脚本顶部添加 global cached;cached=dict() 来创建全局字典,但该变量似乎是跨工作人员的不同版本,并且不共享数据。以前我使用 dogpile.cache.redis 执行此操作,但该区域将不可序列化,因为它使用线程锁。我也尝试过创建一个 dict 并使用 ray.put(cached) 将其放入 ray 的对象存储中,但我想我在某处读到了 ray cannot share dictionaries in memory

我目前正在尝试从每个工作程序返回缓存并将它们合并到主目录中,然后再次将它们放入对象存储中。 是否有更好的方式在光线工作者之间共享字典/缓存?

解决方法

不幸的是,您没有创建 minimal,reproducible example,所以我看不到您是如何进行多处理的。为便于论证,我假设您使用的是 Pool 模块中的 multiprocessing 类(concurrent.futures.ProcessPoolExecutor 作为类似工具)。然后你想使用一个 managed,可共享的 dict 如下:

from multiprocessing import Pool,Manager


def init_pool(the_cache):
    # initialize each process in the pool with the following global variable:
    global cached
    cached = the_cache

def main():
    with Manager() as manager:
        cached = manager.dict()
        with Pool(initializer=init_pool,initargs=(cached,)) as pool:
            ... # code that creates tasks

# required by Windows:
if __name__ == '__main__':
    main()

这会在带有变量 cached 的字典中创建对该字典的 代理 的引用。因此,所有字典访问本质上变得更类似于远程过程调用,因此执行速度比“正常”字典访问慢得多。请注意...

如果有其他一些机制来创建工作线程(装饰器 @ray.remote?),cached 变量可以作为参数传递给函数 f。强>

,

您可能对这个关于为 Ray 编写函数缓存的问题/答案感兴趣。 Implementing cache for Ray actor function

您的想法是正确的,但我认为您缺少的关键细节是您应该使用 Ray 将全局状态保存在 actor 或对象存储中(如果它是不可变的)。

在您的情况下,您似乎正在尝试缓存远程功能的部分内容,而不是全部内容。您可能想要看起来像这样的东西。

这是您可以考虑如何编写函数的简化版本。

@ray.remote
class Cache:
  def __init__(self):
    self.cache = {}

  def put(self,x,y):
    self.cache[x] = y

  def get(self,x):
    return self.cache.get(x)

global_cache = Cache.remote()

@ray.remote
def f(x):
  all_inputs = list(range(x)) # A simplified set of generated inputs based on x
  obtained_output = ray.get([global_cache.get(i) for i in all_inputs])

  unknown_indices = []
  for i,output in enumerate(obtained_output):
    if output is None:
        unknown_inputs.append(i)
 
  # Now go through and calculate all the unknown inputs
  for i in unknown_inputs:
    output = predict(all_inputs[i]) # calculate the output
    global_cache.put.remote(output) # Cache it so it's available next time
    obtained_output[i] = output

  return obtained_output