Python：如何从命令行 arguments 定义可由多处理池访问的全局变量？

Question

我有一个从大型数据集生成文件的脚本，所以我使用多处理来加快速度。 我遇到的问题是我的脚本接受了几个命令行 arguments 使用改变结果的argparse库，我正在努力将命令行 arguments 传递给我的多处理池调用的 function。

我确定解决方案非常简单，我只是没有看到它。 我想我会创建一个全局变量来更新以反映命令行参数，但池调用的 function 仍然具有旧值。 我试图在下面说明我的问题：

output_dir = 'default'

def do_task(item):
    print(output_dir) # Prints 'default'
    result = process_item(item)
    write_to_file(data=result, location=os.path.join(output_dir, item.name))

def do_multi_threaded_work(data_path):
    print(output_dir) # Prints command line argument
    data = read_from_file(args.input_file)
    pool = multiprocessing.Pool()
    for i, _ in enumerate(pool.imap_unordered(do_task, data):
        print('Completed task %d/%d' % (i, len(data)))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-o', '--output-dir')
    parser.add_argument('-i', '--input-file')
    args = parser.parse_args()
    output_dir = args.output_dir
    do_multithreaded_work(args.input_file)

如何确保根据命令行 arguments 将文件保存到正确的目录？

编辑：有人建议我做一些类似下面的代码的事情，但是考虑到我的实际代码中有很多常量（在这个例子中我将它简化为 1），这看起来非常混乱和违反直觉。 真的没有更好的方法来设置do_task function 可访问的全局常量，而无需对值进行硬编码吗？

from itertools import repeat
...
def do_multi_threaded_work(data_path):
    ...
    for i, _ in enumerate(pool.imap_unordered(do_task, zip(data, repeat(output_dir))):

Answer 1

如果我以正确的方式理解您的问题，您可以执行以下操作，将附加的 arguments 与主要数据一起发送到您的 function：

# my toy example:

import multiprocessing as mp


def do_job(x) -> int:
    # x[0] - is a real data # x[1], x[2] imagine the are parameters to tune fuction 
    return x[0]**2 + x[1] + x[2]


if __name__ == '__main__':
    jobs = [1, 2, 3, 4, 5, 6, 7, 8]  # number 0 argument - data you want to process
    # rules to work with jobs - tune parameters
    number_one_arg = 100
    number_two_arg = 2000

    # create structure to accompany data with tune parameters
    x_for_do_job = [(i, number_one_arg, number_two_arg,) for i in jobs]
    print(x_for_do_job) # show what we have now

    pool_ = mp.Pool(4)
    results = pool_.map(do_job, x_for_do_job)
    print(results)

Answer 2

最后找到了一个涉及使用functools库的partial功能的解决方案。 这使我能够通过创建具有指定这些参数的部分 function 来指定任何常量参数。 然后，我将部分 function 与可迭代对象一起传递给池。

from functools import partial

def do_task(output_dir, item):
    print(output_dir) # Prints 'default'
    result = process_item(item)
    write_to_file(data=result, location=os.path.join(output_dir, item.name))

def do_multi_threaded_work(data_path):
    print(output_dir) # Prints command line argument
    data = read_from_file(args.input_file)
    func = partial(do_task, output_dir)
    pool = multiprocessing.Pool()
    for i, _ in enumerate(pool.imap_unordered(func, data):
        print('Completed task %d/%d' % (i, len(data)))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-o', '--output-dir')
    parser.add_argument('-i', '--input-file')
    args = parser.parse_args()
    output_dir = args.output_dir
    do_multithreaded_work(args.input_file)