Multiprocessing write to large file out of memory

user_3641153 0
Viewed 206

I have the following script which works well for writing smaller datasets out to a file, but eventually runs out of memory when processing and writing larger datasets. Some files sizes will be 60gb +.

def do_work(index):
    ref_feature = layer.GetFeature(index)
    if ref_feature:
        try:
            return ref_feature.ExportToJson(as_object=True)
        except Exception as e:
            pass
    return None


def run_mp():
    # empty file contents
    open(f"{out_dir}/{fc_name}.geojsonseq", "w", encoding='utf8').close()

    # initiate multiprocessing
    pool = Pool(cpu_count())
    fc = layer.GetFeatureCount()
    resultset = pool.imap_unordered(do_work, range(fc), chunksize=1000)

    # this part is done after all results are ready, resulting in huge memory storage until results are written
    with open(f"{out_dir}/{fc_name}.geojsonseq", 'a') as file:
        for obj in resultset:
            file.write(f"\x1e{json.dumps(obj)}\n")


if __name__ == '__main__':
    seg_start = time.time()
    run_mp()
    print(f' completed in {time.time() - seg_start}')

Question:

Is there a way to stream the results directly out to a file without building it up in memory and dumping it out to a file at the end?

multiprocessing python
1 Answers

Since imap_unordered doesn't apply any back pressure to the worker processes, I suspect the results are backing up in the internal results queue of IMapUnorderedIterator. If that's the case, you have three options:

  • Write the results faster in the main process. Try returning the string f"\x1e{json.dumps(obj)}\n" from your workers rather than dumping in the main process. If that doesn't work:
  • Write temporary files in the workers and concatenate them in a second pass in the main process. Workers will interfere with each other's writes if you try to have them all append the final file simultaneously. You should be able to do this using minimal extra disk space. Note that you can do json.dump directly into a file object. Alternatively you could guard worker writes to the same file with a multiprocessing.Lock. If the extra writes are too time consuming:
  • Manage back pressure yourself. Use Pool.apply_async or ProcessPoolExecutor.submit to start cpu_count jobs and only submit additional work after writing a result to disk. It's less automatic than Pool.imap_unordered but that's the kind of thing you have to deal with when you're data starts getting big!
user_45049500
Related