为何Python的Pool.imap_unordered失效而map正常？附代码示例

问题描述 #

原本认为在不关心输出顺序的情况下，multiprocessing.Pool的imap_unordered和map方法可以互换使用，但多次遇到imap_unordered完全不工作的情况。使用Python 3.11.2版本，采用元组传参，相关代码如下：

#!/usr/bin/env python3

from multiprocessing import Pool
from functools import partial

def main ():
    query_seqs = list(range(22))
    print(len(query_seqs))
    with Pool(processes=7) as pool:
        tmp = [('args2', 'blast_type2', query) for query in query_seqs]
        pool.map(pool_tasks, tmp)
    print(len(query_seqs))
    with Pool(processes=7) as pool:
        tmp = [('args1', 'blast_type1', query) for query in query_seqs]
        pool.imap_unordered(pool_tasks, tmp)
    print(len(query_seqs))
def pool_tasks(tup):
    args, blast_type, query = tup
    print('args, blast_type, query', args, blast_type, query)

if __name__ == '__main__':
    main()

问题原因 #

imap_unordered和map的核心差异在于执行机制：

• map是阻塞式方法，调用后会立即启动所有任务，直到全部完成才返回结果列表，因此代码中pool.map能正常触发任务。
• imap_unordered返回的是一个迭代器对象，只有主动迭代这个迭代器（比如遍历、转成列表）时，才会真正触发任务执行。你的代码仅调用了pool.imap_unordered但未处理返回的迭代器，而with块结束时会立即关闭进程池，导致任务还未执行就被终止。

解决方法 #

只需对imap_unordered返回的迭代器进行迭代操作即可，以下是两种可行方式：

方法1：将迭代器转为列表 #

with Pool(processes=7) as pool:
    tmp = [('args1', 'blast_type1', query) for query in query_seqs]
    # 转成列表触发任务执行
    list(pool.imap_unordered(pool_tasks, tmp))

方法2：遍历迭代器 #

with Pool(processes=7) as pool:
    tmp = [('args1', 'blast_type1', query) for query in query_seqs]
    # 遍历迭代器，触发任务执行
    for _ in pool.imap_unordered(pool_tasks, tmp):
        pass

修改后，imap_unordered就能正常执行所有任务，控制台会输出args1开头的打印内容。

内容的提问来源于stack exchange，提问作者Liam McIntyre