Python: Как определить глобальную переменную, доступную для многопроцессорного пула, из аргументов командной строки?

Question

У меня есть скрипт, который создает файлы из большого набора данных, поэтому я использую многопроцессорную обработку для ускорения работы. У меня проблема в том, что мой сценарий принимает несколько аргументов командной строки с использованием библиотеки argparse, которые изменяют результаты, и я пытаюсь передать аргументы командной строки в функцию, вызываемую моим многопроцессорным пулом.

I Я уверен, что решение этого действительно просто, я просто не вижу его. Я подумал, что сделаю глобальную переменную, которая будет обновлена, чтобы отражать аргументы командной строки, но моя функция, вызываемая пулом, все еще имеет старое значение. Я попытался проиллюстрировать мою проблему ниже:

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)

Как я могу гарантировать, что я сохраняю свои файлы в правильном каталоге в соответствии с аргументами командной строки?

Редактировать : Было предложено, чтобы я сделал что-то вроде приведенного ниже кода, однако, учитывая, что у меня довольно много констант (в этом примере я упростил до 1), это выглядит очень грязно и нелогично. Неужели нет лучшего способа просто установить глобальную константу, доступную для функции do_task, без жесткого кодирования значения?

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

Если я правильно понял ваш вопрос, вы можете сделать следующее, чтобы отправить дополнительные аргументы вашей функции с основными данными:

# 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

Нашли решение, которое включало в конце концов использование функции partial библиотеки functools. Это позволило мне указать любые постоянные параметры, создав частичную функцию с указанными параметрами. Затем я передаю эту частичную функцию вместе с итерацией в пул.

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_work, 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)