使用 Ray 并行化大型程序的正确方法

Correct way to parallelize a large program using Ray

我有一个相当大的 Python 程序(~800 行),它具有以下结构:

程序是大规模并行的,因为主函数的每次执行都独立于前一个和下一个。因此,我使用 Ray 在集群中的多个工作节点上并行执行主要功能。操作系统为 CentOS Linux release 8.2.2004 (Core),集群执行 PBS Pro 19.2.4.20190830141245。我正在使用 Python 3.7.4、Ray 0.8.7 和 Redis 3.4.1。

我在Python脚本中有以下内容,其中foo是主要功能:

@ray.remote(memory=2.5 * 1024 * 1024 * 1024)
def foo(locInd):
    # Main function

if __name__ == '__main__':
    ray.init(address='auto', redis_password=args.pw,
             driver_object_store_memory=10 * 1024 * 1024 * 1024)
    futures = [foo.remote(i) for i in zip(*np.asarray(indArr == 0).nonzero())]
    waitingIds = list(futures)
    while len(waitingIds) > 0:
        readyIds, waitingIds = ray.wait(
            waitingIds, num_returns=min([checkpoint, len(waitingIds)]))
        for r0, r1, r2, r3, r4, r5, r6, r7 in ray.get(readyIds):
            # Process results
            indArr[r0[::-1]] = 1
            nodesComplete += 1
    ray.shutdown()

以下是我用来启动 Ray 的说明

# Head node
/path/to/ray start --head --port=6379 \
--redis-password=$redis_password \
--memory $((120 * 1024 * 1024 * 1024)) \
--object-store-memory $((20 * 1024 * 1024 * 1024)) \
--redis-max-memory $((10 * 1024 * 1024 * 1024)) \
--num-cpus 48 --num-gpus 0

# Worker nodes
/path/to/ray start --block --address= \
--redis-password= --memory $((120 * 1024 * 1024 * 1024)) \
--object-store-memory $((20 * 1024 * 1024 * 1024)) \
--redis-max-memory $((10 * 1024 * 1024 * 1024)) \
--num-cpus 48 --num-gpus 0

一切 运行 都符合预期,前提是我处理的是足够小的数据集。然而,执行会产生以下警告

如果我尝试 运行 更大数据集上的代码,我会收到以下错误:

Traceback (most recent call last):
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 700, in send_packed_command
    sendall(self._sock, item)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/_compat.py", line 8, in sendall
2020-08-21 14:22:34,226 WARNING worker.py:1134 -- Warning: The remote function __main__.foo has size 898527351 when pickled. It will be stored in Redis, which could cause memory issues. This may mean that its definition uses a large array or other object.
    return sock.sendall(*args, **kwargs)
ConnectionResetError: [Errno 104] Connection reset by peer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "./Program.py", line 1030, in <module>
    for i in zip(*np.asarray(indArr == 0).nonzero())]
  File "./Program.py", line 1030, in <listcomp>
    for i in zip(*np.asarray(indArr == 0).nonzero())]
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/remote_function.py", line 95, in _remote_proxy
    return self._remote(args=args, kwargs=kwargs)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/remote_function.py", line 176, in _remote
    worker.function_actor_manager.export(self)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/function_manager.py", line 152, in export
    "max_calls": remote_function._max_calls
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/client.py", line 3023, in hmset
    return self.execute_command('HMSET', name, *items)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/client.py", line 877, in execute_command
    conn.send_command(*args)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 721, in send_command
    check_health=kwargs.get('check_health', True))
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 713, in send_packed_command
    (errno, errmsg))
redis.exceptions.ConnectionError: Error 104 while writing to socket. Connection reset by peer.

我向 Ray 描述程序的方式显然有误。我有 Scipy 个我认为是全局的插值器对象,但是,正如在 GitHub thread 中已经指出的那样,我应该对它们调用 ray.put。问题是我 运行 到这些 ValueError: buffer source array is read-only 中,我不知道如何诊断。此外,我不确定是否应该用 @ray.remote 或仅装饰主要功能来装饰所有功能。我想我可以为所有附加功能做 @ray.remote(num_cpus=1),因为它实际上只应该是并行执行的主要功能,但我不知道这是否有意义。

非常感谢任何帮助,如果需要,我很乐意提供更多信息。

我可能已经解决了我的问题,但我不介意听取别人的意见,因为我对 Ray 的了解确实有限。另外,我想这可能会帮助遇到类似问题的其他人(希望我不是一个人!)。

正如我在问题中提到的,程序 运行 对于足够小的数据集来说相当不错(尽管它似乎绕过了 Ray 逻辑的某些方面),但它最终在大型数据集上崩溃了。仅使用 Ray 任务,我没有设法调用存储在对象存储 (ValueError: buffer source array is read-only) 中的 Scipy 个插值器对象,并且装饰所有函数没有意义,因为实际上只有主要的函数应该执行同时(同时调用其他函数)。

因此,我决定更改程序结构以使用 Ray Actors。设置说明现在是 __init__ 方法的一部分。特别是,Scipy Interpolator 对象在此方法中定义并设置为 self 的属性,作为全局变量。大多数函数(包括主要函数)已成为 class 方法,但通过 Numba 编译的函数除外。对于后者,它们仍然是用 @jit 装饰的独立函数,但它们每个现在在 class 中都有一个等效的包装方法,它调用 jitted 函数。

为了让我的程序并行执行我现在的主要方法,我依赖于 ActorPool。我创建了与可用 CPU 一样多的 actor,每个 actor 都执行 main 方法,成功调用方法和 Numba-compiled 函数,同时还设法访问 Interpolator 对象。我只将 @ray.remote 应用于定义的 Python class。所有这些 t运行slates 为以下结构:

@ray.remote
class FooClass(object):
    def __init__(self, initArgs):
        # Initialisation

    @staticmethod
    def exampleStaticMethod(args):
        # Processing
        return

    def exampleMethod(self, args):
        # Processing
        return

    def exampleWrapperMethod(self, args):
        return numbaCompiledFunction(args)

    def mainMethod(self, poolMapArgs):
        # Processing
        return


@jit
def numbaCompiledFunction(args):
    # Processing
    return


ray.init(address='auto', redis_password=redPass)
actors = []
for actor in range(int(ray.cluster_resources()['CPU'])):
    actors.append(FooClass.remote(initArgs))
pool = ActorPool(actors)
for unpackedTuple in pool.map_unordered(
        lambda a, v: a.mainMethod.remote(v),
        poolMapArgs):
    # Processing
ray.shutdown()

这在分布在 4 个节点上的 192 个 CPU 上成功运行,没有任何警告或错误。