Ray并行处理场景：能否对kwargs参数直接使用ray.put？

我有一个接受kwargs参数的函数，希望用它来并行处理请求，想找一种简单的方式处理kwargs，无需逐个迭代并将其存入共享/集中式存储。尝试代码如下：

import ray
ray.init(ignore_reinit_error=True)

def get_data_multip(code=['AB', 'CD', 'ED'], **kwargs):
    kwargs_ref = ray.put(**kwargs)
    results = []
    for s in code:
        func_ref = ray.remote(get_data)
        res = func_ref.remote(code = s,**kwargs)
        results.append(res)
    
    rets = ray.get(results)
    print(rets)

def get_data(code, st_dt = None, end_dt = None):
    print(f'st_dt:{st_dt}, end_dt:{end_dt}')
    return [1,2,3,4]

get_data_multip(code = ['AB', 'CD', 'ED'], st_dt = '01/01/2020', end_dt = '12/31/2020')

上述代码在ray.put(**kwargs)行报错：

put() got an unexpected keyword argument 'st_dt'

更新
我曾尝试一种方案，但该方案只是将字典作为第二个参数传递，而非kwargs：

def get_data_multip(code=['AB', 'CD', 'ED'], **kwargs):
    kwargs_ref = ray.put(vars()['kwargs'])

    results = []
    for s in code:
        func = ray.remote(get_data)
        res = func.remote(s, kwargs_ref)
        results.append(res)
    
    rets = ray.get(results)
    print(rets)

解决方案 #

核心问题说明 #

ray.put()仅接受单个对象作为参数，不能用**kwargs解包传递参数，直接传入整个kwargs字典即可。另外Ray会自动处理对象的序列化和分布式共享，多数场景下无需手动调用put，直接在远程函数调用时传递**kwargs即可。

方案1：直接传递kwargs（推荐） #

无需手动处理对象存储，直接在远程函数调用时解包kwargs，Ray会自动完成序列化和共享：

import ray
ray.init(ignore_reinit_error=True)

def get_data_multip(code=['AB', 'CD', 'ED'], **kwargs):
    results = []
    # 提前将函数远程化，避免循环内重复创建
    remote_get_data = ray.remote(get_data)
    for s in code:
        res = remote_get_data.remote(code=s, **kwargs)
        results.append(res)
    
    rets = ray.get(results)
    print(rets)

def get_data(code, st_dt=None, end_dt=None):
    print(f'st_dt:{st_dt}, end_dt:{end_dt}')
    return [1,2,3,4]

get_data_multip(code=['AB', 'CD', 'ED'], st_dt='01/01/2020', end_dt='12/31/2020')

方案2：手动put kwargs字典复用 #

如果需要复用kwargs对象，可以先将整个字典存入Ray存储，调用时再获取并解包：

import ray
ray.init(ignore_reinit_error=True)

def get_data_multip(code=['AB', 'CD', 'ED'], **kwargs):
    # 直接传入kwargs字典，不要解包
    kwargs_ref = ray.put(kwargs)
    remote_get_data = ray.remote(get_data)
    results = []
    for s in code:
        # 获取存储的kwargs字典后解包传递
        res = remote_get_data.remote(code=s, **ray.get(kwargs_ref))
        results.append(res)
    
    rets = ray.get(results)
    print(rets)

def get_data(code, st_dt=None, end_dt=None):
    print(f'st_dt:{st_dt}, end_dt:{end_dt}')
    return [1,2,3,4]

get_data_multip(code=['AB', 'CD', 'ED'], st_dt='01/01/2020', end_dt='12/31/2020')

关于你尝试的字典传递方案 #

如果要保留这种方式，需要修改get_data的签名来接收字典参数，再手动提取kwargs：

import ray
ray.init(ignore_reinit_error=True)

def get_data_multip(code=['AB', 'CD', 'ED'], **kwargs):
    kwargs_ref = ray.put(kwargs)
    remote_get_data = ray.remote(get_data)
    results = [remote_get_data.remote(s, kwargs_ref) for s in code]
    
    rets = ray.get(results)
    print(rets)

def get_data(code, kwargs_dict):
    st_dt = kwargs_dict.get('st_dt')
    end_dt = kwargs_dict.get('end_dt')
    print(f'st_dt:{st_dt}, end_dt:{end_dt}')
    return [1,2,3,4]

get_data_multip(code=['AB', 'CD', 'ED'], st_dt='01/01/2020', end_dt='12/31/2020')

但这种方式改变了原函数的调用签名，不推荐使用。

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