如何将设置代码和处理代码传递给另一个进程？

Question

我有一个通用的多处理工作器 class，它从队列中获取要处理的项目。 worker class 的用户需要传递处理每个项目的 function。 但是，某些处理功能需要设置代码。

当前实现使用生成器 function，用户必须正确实现它才能正确执行设置代码一次，处理队列中的项目，并处理工作人员正常完成时引发的 StopIteration 异常。

是否可以使用更直接可靠的方法将设置代码与处理代码分开并处理工作人员引发的异常？

这是我所拥有的：

import multiprocessing as mp
import typing

P = typing.Callable[[], typing.Generator[None, None, None]]
Q: typing.TypeAlias = "mp.Queue"


class Worker(mp.Process):
    def __init__(self, queue: Q, processor: P):
        mp.Process.__init__(self)
        self.queue = queue
        self.processor = processor

    def run(self):
        processor = self.processor()
        next(processor)  # start the processor
        while True:
            item = self.queue.get()
            processor.send(item)

            if item is None:
                break


class WorkerPool:
    def __init__(self, n_workers: int, processor_generator: P, queue: Q):
        self.workers = [Worker(queue, processor_generator) for _ in range(n_workers)]
        self.queue = queue

    def __enter__(self):
        for worker in self.workers:
            worker.start()

    def signal_end(self):
        for _ in self.workers:
            self.queue.put(None)

    def terminate(self):
        for worker in self.workers:
            worker.terminate()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.signal_end()
            self.join()
            return True

        self.terminate()
        return False

    def join(self):
        for worker in self.workers:
            worker.join()


class GeneratorWorkerManager:
    def __init__(
        self, item_generator: typing.Generator, processor_generator: P, n_workers: int
    ) -> None:

        queue: Q = mp.Queue()
        with WorkerPool(n_workers, processor_generator, queue):
            for item in item_generator:
                queue.put(item)

GeneratorWorkerManager class 的用户可以执行以下操作：

def processor():
    # All sorts of setup code possible, including a with-statement. 
    item = yield
    while item is not None:
        # process item
        print(item)
        item = yield
    return


items = range(10)
GeneratorWorkerManager(items, processor, 1)

工作人员将打印 0 到 9 的位置。但是，这依赖于用户正确实现处理器 function。 工作程序在正常完成时还会引发StopIteration异常。

有没有更好的方法在同一上下文中使用设置代码和处理代码？

Answer 1

诀窍是传递一个返回上下文管理器的 function。 这里我以open的 function 包裹在另一个 function 中为例。

import multiprocessing as mp
from typing import Callable, ContextManager, Iterable, TypeVar

I = TypeVar("I")
S = TypeVar("S")


class Worker(mp.Process):
    def __init__(
        self,
        queue: "mp.Queue[I|None]",
        processor: Callable[[S, I], None],
        setup: Callable[[], ContextManager[S]],
    ):
        mp.Process.__init__(self)
        self.queue = queue
        self.processor = processor
        self.user_setup = setup

    def run(self):

        with self.user_setup() as setup:
            while True:
                item = self.queue.get()
                if item is None:
                    break
                self.processor(setup, item)


class WorkerPool:
    def __init__(
        self,
        n_workers: int,
        processor: Callable[[S, I], None],
        queue: "mp.Queue[I|None]",
        setup: Callable[[], ContextManager[S]],
    ):
        self.workers = [Worker(queue, processor, setup) for _ in range(n_workers)]
        self.queue = queue

    def __enter__(self):
        for worker in self.workers:
            worker.start()

    def signal_end(self):
        for _ in self.workers:
            self.queue.put(None)

    def terminate(self):
        for worker in self.workers:
            worker.terminate()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.signal_end()
            self.join()
            return True

        self.terminate()
        return False

    def join(self):
        for worker in self.workers:
            worker.join()


class GeneratorWorkerManager:
    def __init__(
        self,
        item_generator: Iterable[I],
        processor: Callable[[S, I], None],
        n_workers: int,
        setup: Callable[[], ContextManager[S]],
    ) -> None:

        queue: "mp.Queue[I|None]" = mp.Queue()
        with WorkerPool(n_workers, processor, queue, setup):
            for item in item_generator:
                queue.put(item)


def custom_setup():
    return open("test.txt", "w", encoding="utf-8")


def custom_processor(setup, item) -> None:
    f = setup
    f.write(str(item))


items = range(10)

if __name__ == "__main__":
    mp.freeze_support()
    GeneratorWorkerManager(items, custom_processor, 1, custom_setup)

不可能使用不支持 with 语句的设置 function，但是可以很容易地使用两个设置变量来扩展它，一个用于上下文设置，一个用于常规设置。