如何线程化生成器

Question

我有一个生成器对象，它会加载大量数据并占用系统的 I/O。数据太大而无法一次全部放入内存，因此使用生成器。 而且我有一个消费者，所有的 CPU 来处理生成器产生的数据。它不会消耗太多其他资源。是否可以使用线程交错这些任务？

例如，我猜可以在 11 秒内运行下面的简化代码。

import time, threading
lock = threading.Lock()
def gen():
    for x in range(10):
        time.sleep(1)
        yield x
def con(x):
    lock.acquire()
    time.sleep(1)
    lock.release()
    return x+1

但是，最简单的线程应用程序并没有在那个时候运行。它确实加快了速度，但我认为是因为生成和工作的调度程序之间存在并行性。但不是因为工人之间的并行性。

import joblib
%time joblib.Parallel(n_jobs=2,backend='threading',pre_dispatch=2)((joblib.delayed(con)(x) for x in gen()))
# CPU times: user 0 ns, sys: 0 ns, total: 0 ns
# Wall time: 16 s

Answer 1

我创建了这个问题，看看是否有一个惯用的替换 for 循环模式。虽然二战的答案确实解决了这个问题，但它有一个警告，如果它的输出相当大，生成器可能会领先于消费者线程并聚集内存。我也更喜欢 joblib。

问题文本中的 joblib 代码的问题是 gen 在主线程中迭代，因此不是调度作业，而是花时间等待 gen。当输入生成器在joblib 下运行缓慢时，我已经放弃尝试理解调度是如此奇怪。然而，在将生产者和消费者都移动到延迟函数中之后，我确实设法让它正确地完成了这项工作。

当实际上事先知道可迭代的长度时（例如要一个一个处理的文件列表），代码很简单。下面的代码保证只有一个线程产生数据，一个线程同时消费数据。

sync_gen,sync_con = threading.Lock(), threading.Lock()
@joblib.delayed
def work(iterable):
    with sync_gen:
        x = next(iterable)
    with sync_con:
        return con(x)

N=10
iterable = gen()
res1 = joblib.Parallel(2,'threading')(work(iterable) for x in range(N))
#[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

如果生成器长度未知，线程工作者最好累积他们的结果而不是处理单个输入。

sync_gen,sync_con = threading.Lock(), threading.Lock()
def thread_safe(gen):
    try:
        while True:
            with sync_gen:
                x = next(gen)
            yield x
    except StopIteration:
        pass

def work2(safe_iterable):
    res = []
    for x in safe_iterable:
        with sync_con:
            res.append(con(x))
    return res

iterable = gen()
de_work2= joblib.delayed(work2)
res2 = joblib.Parallel(2,'threading')(de_work2(thread_safe(iterable)) for x in range(2))
#[[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]

或者使用 ThreadPoolExecutor：

from concurrent.futures import ThreadPoolExecutor
iterable = gen()
with ThreadPoolExecutor() as e:
    futures = [e.submit(work2,thread_safe(iterable)) for x in range(2)]
res = [future.result() for future in futures]

Answer 2

将您的数据发送到不同的进程。我使用了concurrent.futures，因为我喜欢简单的界面。

这在我的电脑上运行大约需要 11 秒。

from concurrent.futures import ThreadPoolExecutor
import concurrent
import threading
lock = threading.Lock()

def gen():
    for x in range(10):
        time.sleep(1)
        yield x

def con(x):
    lock.acquire()
    time.sleep(1)
    lock.release()
    return f'{x+1}'

if __name__ == "__main__":

    futures = []
    with ThreadPoolExecutor() as executor:
        t0 = time.time()
        for x in gen():
            futures.append(executor.submit(con,x))
    results = []
    for future in concurrent.futures.as_completed(futures):
        results.append(future.result())
    print(time.time() - t0)
    print('\n'.join(results))

---

使用 100 次生成器迭代 (def gen(): for x in range(100):) 大约需要 102 秒。

---

您的进程可能需要跟踪有多少数据已发送到尚未完成的任务，以防止占用大量内存资源。

向con 添加一些诊断打印似乎表明一次可能有至少两个数据块在那里。

def con(x):
    print(f'{x} received payload at t0 + {time.time()-t0:3.3f}')
    lock.acquire()
    time.sleep(1)
    lock.release()
    print(f'{x} released lock at t0 + {time.time()-t0:3.3f}')
    return f'{x+1}'