ray并行计算时，如何确定异步函数返回值的数量？

调用者决定返回值的数量

调用者给执行函数传入参数决定返回值的数量。

import numpy as np
import ray
@ray.remote
def large_values(num_returns):
    return [
        np.random.randint(np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8)
        for _ in range(num_returns)
    ]
    
# 优选第二种写法：生成器将这些返回值将一次存储一个在Ray的对象存储中，节省内存空间
@ray.remote
def large_values_generator(num_returns):
    for i in range(num_returns):
        yield np.random.randint(
            np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8
        )
        print(f"yielded return value {i}")

执行者决定返回值的数量

在某些情况下，调用者可能不知道远程函数返回值的数量。例如，假设编写一个任务，将其参数分解为大小相等的块并返回这些块。在执行任务之前，因不知道参数的大小故不能预期返回值的数量，我们可以使用远程生成器函数来返回动态数量的值。

import numpy as np

# 由执行者决定返回的数量
@ray.remote(num_returns="dynamic")
def split(array, chunk_size):
    while len(array) > 0:
        yield array[:chunk_size]
        array = array[chunk_size:]

# 传入大对象建议使用ray.put
array_ref = ray.put(np.zeros(np.random.randint(1000_000)))
block_size = 1000
array_size = len(ray.get(array_ref))

# 返回一个远程对象
dynamic_ref = split.remote(array_ref, block_size)
print(dynamic_ref)
# ObjectRef(c8ef45ccd0112571ffffffffffffffffffffffff0100000001000000)

i = -1
# 返回生成器对象
ref_generator = ray.get(dynamic_ref)
print(ref_generator)
# <ray._raylet.ObjectRefGenerator object at 0x7f7e2116b290>
# 获取生成器对象的内容,用断言判断长度是否符合
for i, ref in enumerate(ref_generator):
    # Each ObjectRefGenerator iteration returns an ObjectRef.
    assert len(ray.get(ref)) <= block_size
num_blocks_generated = i + 1
assert array_size <= num_blocks_generated * block_size