如何使用torch.multiprocessing.Pool（Python）摆脱僵尸进程

Question

我正在使用torch.multiprocessing.Pool来加快我的NN推理速度，如下所示：

import torch.multiprocessing as mp
mp = torch.multiprocessing.get_context('forkserver')

def parallel_predict(predict_func, sequences, args):
    predicted_cluster_ids = []
    pool = mp.Pool(args.num_workers, maxtasksperchild=1)
    out = pool.imap(
        func=functools.partial(predict_func, args=args),
        iterable=sequences,
        chunksize=1)
    for item in tqdm(out, total=len(sequences), ncols=85):
        predicted_cluster_ids.append(item)
    pool.close()
    pool.terminate()
    pool.join()
    return predicted_cluster_ids

注1）我使用imap是因为我希望能够使用tqdm显示进度条。
注2）我尝试了forkserver和spawn但没有运气。 我无法使用其他方法，因为它们与CUDA的交互方式（不良）。
注3）我使用的是maxtasksperchild=1和chunksize=1因此对于序列中的每个sequences都会产生一个新进程。
注4）添加或删除pool.terminate()和pool.join()没有区别。
注5） predict_func是我创建的类的方法。 我也可以将整个模型传递给parallel_predict但它不会改变任何内容。

一切正常，除了一段时间后，我在CPU上的内存不足（而在GPU上，一切正常。）。 使用htop监视内存使用情况，我注意到，对于我使用池生成的每个进程，我都会得到一个使用0.4％内存的僵尸。 它们不会被清除，因此会继续使用空间。 尽管如此， parallel_predict确实返回正确的结果，并且计算继续进行。 我的脚本的结构是id多次验证，因此下一次parallel_predict称为“僵尸加法”。

这就是我在htop得到的：

通常，这些僵尸在ctrl-c之后会被清除，但在极少数情况下，我需要killall 。

有什么办法可以迫使Pool关闭？

更新：我试图使用以下方法杀死僵尸进程：

def kill(pool):
    import multiprocessing
    import signal
    # stop repopulating new child
    pool._state = multiprocessing.pool.TERMINATE
    pool._worker_handler._state = multiprocessing.pool.TERMINATE
    for p in pool._pool:
        os.kill(p.pid, signal.SIGKILL)
    # .is_alive() will reap dead process
    while any(p.is_alive() for p in pool._pool):
        pass
    pool.terminate()

但这行不通。 它卡在pool.terminate()

UPDATE2：我试图在imap使用initializer arg来捕获如下信号：

def process_initializer():
    def handler(_signal, frame):
        print('exiting')
        exit(0)
    signal.signal(signal.SIGTERM, handler)


def parallel_predict(predict_func, sequences, args):
    predicted_cluster_ids = []
    with mp.Pool(args.num_workers, initializer=process_initializer, maxtasksperchild=1) as pool:
        out = pool.imap(
            func=functools.partial(predict_func, args=args),
            iterable=sequences,
            chunksize=1)
        for item in tqdm(out, total=len(sequences), ncols=85):
            predicted_cluster_ids.append(item)
        for p in pool._pool:
            os.kill(p.pid, signal.SIGTERM)
        pool.close()
        pool.terminate()
        pool.join()
    return predicted_cluster_ids

但同样，它不会释放内存。

Answer 1

好的，我有更多见解可以与您分享。 的确，这不是一个错误，实际上是Python中的多处理模块的“假定”行为（torch.multiprocessing将其包装）。 发生的事情是，尽管Pool终止了所有进程，但内存并未释放（送回操作系统）。 文档中也对此进行了说明 ，尽管以一种非常混乱的方式。 在文档中说

池中的工作进程通常在池工作队列的整个期间内都处于活动状态

但是也：

在其他系统（例如Apache，mod_wsgi等）中发现的释放工人资源的常见模式是允许池中的工人在退出，清理和产生新进程之前仅完成一定数量的工作。取代旧的。 池的maxtasksperchild参数向最终用户公开了此功能

但“清理”不会发生。

更糟的是，我发现这篇文章中他们建议使用maxtasksperchild=1 。 这会增加内存泄漏，因为这样僵尸的数量就会与要预测的数据点的数量一致，而且因为pool.close()不会释放内存，所以它们会累加起来。

如果在验证中使用多重处理，这将非常糟糕。 对于每个验证步骤，我都在重新初始化池，但是没有从上一次迭代中释放内存。

这里的解决方案是将pool = mp.Pool(args.num_workers)训练循环，这样就不会关闭并重新打开该池，因此它总是重复使用相同的过程。 注意：再次记住要删除maxtasksperchild=1和chunksize=1 。

我认为这应该包含在最佳做法页面中。

顺便说一句，在我看来，多处理库的这种行为应该被认为是一个错误，应该在Python端（而不是Pytorch端）进行修复