python - 使用多处理实现排序的生产者/消费者队列

Question

我有一个非常常见的生产者/消费者场景，但有一点不同。

我需要从多 GB 输入流(可以是文件或 HTTP 流)中读取文本行；使用缓慢且 CPU 密集型算法处理每一行，该算法将为每行输入输出一行文本；然后将输出行写入另一个流。不同的是，我需要按照与生成输出行的输入行相同的顺序编写输出行。

这些场景的通常方法是使用 multiprocessing.Pool 来运行 CPU 密集型算法，其中一个队列从读取器进程中输入行(实际上是批量行)，另一个队列从池中引出并进入编写过程:

                       / [Pool] \    
  [Reader] --> InQueue --[Pool]---> OutQueue --> [Writer]
                       \ [Pool] /

但是如何确保输出行(或批处理)按正确的顺序排序？

一个简单的答案是，“只需将它们写入临时文件，然后对文件进行排序并将其写入输出”。我可能最终会这样做，但我真的很想尽快开始流式传输输出行，而不是等待整个输入流从头到尾处理完毕。

我可以轻松编写自己的 multiprocessing.Queue 实现，它将使用字典(或循环缓冲区列表)、一个锁和两个条件(可能还加上一个整数计数器)在内部对其项目进行排序。但是，我需要从管理器获取所有这些对象，而且我担心在多个进程之间使用这样的共享状态会降低我的性能。那么，有没有一些合适的 Python 方式来解决这个问题呢？

Answer 1

也许我遗漏了一些东西，但似乎你的问题有一个基本的答案。

让我们举一个简单的例子:我们只想反转文本中的行。以下是我们要反转的行:

INPUT = ["line {}".format(i)[::-1] for i in range(30)]

即:

['0 enil', '1 enil', '2 enil', ..., '92 enil']

以下是反转这些线的方法:

import time, random

def process_line(line):
    time.sleep(1.5*random.random()) # simulation of an heavy computation
    return line[::-1]

这些行来自以下来源:

def source():
    for s in INPUT:
        time.sleep(0.5*random.random()) # simulation of the network latency
        yield s

我们可以使用多处理来提高速度:

from multiprocessing import *

with Pool(3) as p:
    for line in p.imap_unordered(process_line, source()):
        print(line)

但是我们的线路没有按照预期的顺序:

line 0
line 2
line 1
line 5
line 3
line 4
...
line 27
line 26
line 29
line 28

要按预期顺序获取该行，您可以:

1. 索引行
2. 处理它们并
3. 按照预期的顺序收集它们。

首先，索引行:

def wrapped_source():
    for i, s in enumerate(source()):
        yield i, s

其次，处理该行，但保留索引:

def wrapped_process_line(args):
    i, line = args
    return i, process_line(line)

第三，按预期顺序收集行。这个想法是使用一个计数器一堆。计数器是下一行的预期索引。

采取下一对(索引，处理行):

• 如果索引等于计数器，则仅生成已处理的行。
• 如果没有，则将这对(索引、已处理行)存储在堆中。

然后，当堆中的最小索引等于计数器时，弹出最小元素并产生该行。

循环直到源为空，然后刷新堆。

from heapq import *
h = []

with Pool(3) as p:
    expected_i = 0 #
    for i, line in p.imap_unordered(wrapped_process_line, wrapped_source()):
        if i == expected_i: # lucky!
            print(line)
            expected_i += 1
        else: # unlucky!
            heappush(h, (i, line)) # store the processed line

        while h: # look for the line "expected_i" in the heap
            i_smallest, line_smallest = h[0] # the smallest element
            if i_smallest == expected_i:
                heappop(h)
                print(line_smallest)
                expected_i += 1
            else:
                break # the line "expect_i" was not processed yet.

    while h: # flush the heap
        print(heappop(h)[1]) # the line

现在我们的行已按预期顺序排列:

line 0
line 1
line 2
line 3
line 4
line 5
...
line 26
line 27
line 28
line 29

没有额外延迟:如果下一个预期行尚未处理，我们必须等待，但一旦该行到达，我们就会放弃它。

主要缺点是你必须手动处理(超时、新请求等)潜在的间隙:一旦你索引了你的行，如果你丢失了一行(无论出于什么原因)，循环将等待这个直到源耗尽，然后才刷新堆。在这种情况下，您可能会耗尽内存。