Python多处理池； 等待迭代完成

Question

我有一个很大的数据集，我希望我的脚本可以迭代通过，对每个条目执行一系列操作，然后将结果排列以存储到HDD。 由于数据集可能相对较大（〜250 GB），因此RAM的可用性要求数据集必须一次处理1000个条目的块（在下面的代码中我称之为dataBlock）。 我还使用multiprocessing.Pool类来促进使用多个CPU内核来完成此任务。

我基本上已经安排好了一切，以便将每个dataBlock传递给Pool，Pool使用imap方法对dataBlock执行所需的计算，Pool返回计算结果，并将数据块的结果附加到列表中。 此列表（ processed_data ）是该组计算的所需最终产品。

processed_data = []

multiprocessing.Pool(processor_cap) as pool:

    for blockIndex, block in enumerate(range(1000, height-remainder, 1000)):

        #Read-in 1000 spectra from source dataset
        dataBlock = np.asarray(raw_dset[blockIndex*1000:block][:])

        '''
        Pass data block to processor pool, which iterates through data
        block. Each spectrum is handed off to a CPU in the pool,
        which centroids it and appends the result to "processed_block".
        '''
        processed_block = pool.imap(centroid_spectrum, dataBlock)

        #Append processed spectra to processed data bin
        for idx, processed_spectrum in enumerate(processed_block):
            processed_data.append(processed_spectrum)

我想知道的是如何使脚本暂停的呼叫后pool.imap()直到完全processed_block不关闭池已恢复。 目前，它的进展对入for环路立即在的这段遵循以上无需等待processed_block由返回pool.imap 。 我已经尝试在pool.imap()调用之后立即调用pool.join() ，但是它仅返回***AssertionError并再次继续其下面的for循环。 一旦将所有dataBlocks馈入池中，就在上面最外层for循环的末尾，我最终可以在脚本中稍后成功地调用pool.close()和pool.join() 。

在此先感谢您的帮助！

Answer 1

如果不付出很多努力来改变周围的事物，那么使用示例很难。 但是，如果您从imap（）调用中获得了迭代器，则可以考虑在到达for循环之前将迭代器的元素解析为列表：

processed_block = pool.imap(centroid_spectrum, dataBlock)
processed_block = [ x for x in processed_block ] # convert from an iterator to a list
for idx, processed_spectrum in enumerate(processed_block):

等等

这样能达到您想要的吗？

Answer 2

我只是将Pool.imap()调用更改为Pool.map()调用，并且脚本按预期运行。 有关更多信息，请参见与Mikhail Burshteyn的交流。