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：我尝试使用initializer中的imap 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终止了所有在进程中，内存不会释放（送回操作系统）。 ~~documentation中也对此进行了说明，尽管以一种非常混乱的方式。~~ 在documentation中表示

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

而且：

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

但不会发生“清理”。

更糟糕的是，我发现了这个post，他们建议在其中使用maxtasksperchild=1。这会增加内存泄漏，因为这样僵尸的数量就会与要预测的数据点的数量一致，并且由于pool.close()不会释放内存，因此它们会累加起来。

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

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

我认为这应该包含在best practices页面中。

在我看来，顺便说一句，多处理库的这种行为应被视为错误，并且应在Python端（而非Pytorch端）进行固定

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

1 个答案: