读取、压缩、使用多进程写入数据

Question

我正在压缩文件。单进程可以处理其中的一些文件，但是我需要压缩成千上万的文件，这可能需要几天时间，为了加快速度，我想使用多进程。我已经阅读了相关信息，应该避免同时有多个进程读取同一文件，而且我猜测同时写入多个进程也不行。以下是我当前运行的单进程的方法：

import tarfile, bz2, os
def compress(folder):
    "compresses a folder into a file"

    bz_file = bz2.BZ2File(folder+'.tbz', 'w')

    with tarfile.open(mode='w', fileobj = bz_file) as tar:

        for fn in os.listdir(folder):

            read each file in the folder and do some pre processing
            that will make the compressed file much smaller than without

            tar.addfile( processed file )

    bz_file.close()
    return

这是将一个文件夹中的所有内容压缩成一个单一的文件。这使得它们更易于处理和组织。如果我只是把它扔进一个池子里，那么我会有几个进程同时读写，所以我要避免这种情况。我可以重新设计它，使只有一个进程读取文件，但我仍然有多个进程在写入：

import multiprocessing as mp
import tarfile, bz2, os

def compress(file_list):
    folder = file_list[0]
    bz_file = bz2.BZ2File(folder+'.tbz', 'w')

    with tarfile.open(mode='w', fileobj = bz_file) as tar:

        for i in file_list[1:]:
            preprocess file data
            tar.addfile(processed data)

    bz_file.close()
    return

cpu_count = mp.cpu_count()
p = mp.Pool(cpu_count)

for subfolder in os.listdir(main_folder):

    read all files in subfolder into memory, place into file_list
    place file_list into fld_list until fld_list contains cpu_count
    file lists. then pass to  p.map(compress, fld_list)

这个问题涉及到多个进程同时写入压缩文件。只要告诉tarfile使用哪种压缩方式，它就会开始向硬盘写入数据。由于我没有足够的RAM来读取需要压缩的所有文件，因此也存在重新启动Pool.map多次的问题。

如何在单个进程中读写文件，同时在多个进程中进行压缩，避免多次重新启动multiprocessing.Pool呢？

- Status

---

你需要查看 pbzip2 的功能并模仿它。 - Ignacio Vazquez-Abrams

使用多进程或多线程队列。首先，一个进程读取所有文件并将它们放入队列1中。其次，多个进程从队列1获取文件并进行压缩，然后将结果放入队列2中。最后，一个进程从队列2中获取结果并进行写入操作。 - Fujiao Liu

Answer 1

不要使用multiprocessing.Pool，应该使用multiprocessing.Queue并创建一个收件箱和一个发件箱。

启动单个进程来读取文件并将数据放入收件箱队列中，并限制队列的大小，以免填满RAM。这里的示例压缩单个文件，但可以调整为一次处理整个文件夹。

def reader(inbox, input_path, num_procs):
    "process that reads in files to be compressed and puts to inbox"

    for fn in os.listdir(input_path):
        path = os.path.join(input_path, fn)

        # read in each file, put data into inbox
        fname = os.path.basename(fn)
        with open(fn, 'r') as src: lines = src.readlines()

        data = [fname, lines]
        inbox.put(data)

    # read in everything, add finished notice for all running processes
    for i in range(num_procs):
        inbox.put(None)  # when a compressor sees a None, it will stop
    inbox.close()
    return

但这只是问题的一半，另一部分是在不必写入磁盘的情况下压缩文件。我们将一个StringIO对象传递给压缩函数，而不是打开的文件；它被传递给tarfile。一旦压缩完成，我们将StringIO对象放入outbox队列。

但我们不能这样做，因为StringIO对象无法pickle化，只有pickleable对象可以放入队列。然而，StringIO的getvalue函数可以以可pickle的格式提供内容，所以使用getvalue获取内容，关闭StringIO对象，然后将内容放入outbox中。

from io import StringIO
import tarfile

def compressHandler(inbox, outbox):
    "process that pulls from inbox, compresses and puts to outbox"
    supplier = iter(inbox.get, None)  # stops when gets a None
    while True:
        try:
            data = next(supplier)  # grab data from inbox
            pressed = compress(data)  # compress it
            ou_que.put(pressed)  # put into outbox
        except StopIteration:
            outbox.put(None)  # finished compressing, inform the writer
            return  # and quit

def compress(data):
    "compress file"
    bz_file = StringIO()

    fname, lines = dat  # see reader def for package order

    with tarfile.open(mode='w:bz2', fileobj=bz_file) as tar:

        info = tarfile.TarInfo(fname)  # store file name
        tar.addfile(info, StringIO(''.join(lines)))  # compress

    data = bz_file.getvalue()
    bz_file.close()
    return data

写进程会从发件箱队列中提取内容并将其写入磁盘。该函数需要知道启动了多少个压缩进程，以便只有在听到每个进程都已停止时才停止。

def writer(outbox, output_path, num_procs):
    "single process that writes compressed files to disk"
    num_fin = 0

    while True:
        # all compression processes have finished
        if num_finished >= num_procs: break

        tardata = outbox.get()

        # a compression process has finished
        if tardata == None:
            num_fin += 1
            continue

        fn, data = tardata
        name = os.path.join(output_path, fn) + '.tbz'

        with open(name, 'wb') as dst: dst.write(data)
    return

最后，需要进行设置以将它们全部组合在一起。

import multiprocessing as mp
import os

def setup():
    fld = 'file/path'

    # multiprocess setup
    num_procs = mp.cpu_count()

    # inbox and outbox queues
    inbox = mp.Queue(4*num_procs)  # limit size 
    outbox = mp.Queue()

    # one process to read
    reader = mp.Process(target = reader, args = (inbox, fld, num_procs))
    reader.start()

    # n processes to compress
    compressors = [mp.Process(target = compressHandler, args = (inbox, outbox))
                   for i in range(num_procs)]
    for c in compressors: c.start()

    # one process to write
    writer = mp.Process(target = writer, args=(outbox, fld, num_procs))
    writer.start()
    writer.join()  # wait for it to finish
    print('done!')

- Status