如何传递初始化代码与处理代码至其他进程？求更优可靠实现

AI 小北

2026-8-2

问题

我实现了一个通用的多进程Worker类，它从队列中获取待处理项并执行处理。使用这个Worker类的用户需要传入一个处理单个项的函数，但部分处理函数需要执行初始化代码。

当前实现采用生成器函数，要求用户正确编写该函数，确保初始化代码仅执行一次、能处理队列中的项，还要处理Worker正常结束时抛出的StopIteration异常。

有没有更直观可靠的方法，能分离初始化代码与处理代码，同时妥善处理Worker抛出的异常？

现有实现代码

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)

现有用法示例

def processor():
    # 各种初始化代码，包括with语句都可以写在这里
    item = yield
    while item is not None:
        # 处理项
        print(item)
        item = yield
    return


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

该Worker会打印0至9，但此实现依赖用户正确编写processor函数，且Worker正常结束时会抛出StopIteration异常。

更优实现方案

我们可以通过定义明确的处理器接口，将初始化逻辑与处理逻辑彻底分离，同时让Worker类内部处理所有底层细节，避免用户需要手动编写生成器逻辑或处理异常。

改进思路

定义一个处理器协议（或抽象基类），要求用户实现setup（初始化）和process（处理单条数据）两个方法
Worker类在启动时自动调用setup方法，之后循环从队列取数据并调用process处理
内置异常捕获与处理逻辑，包括正常结束和异常终止的情况

完整改进代码

import multiprocessing as mp
import typing
from abc import ABC, abstractmethod

# 定义处理器接口：必须实现setup和process方法
class Processor(ABC):
    @abstractmethod
    def setup(self) -> None:
        """初始化逻辑，每个进程仅执行一次"""
        pass

    @abstractmethod
    def process(self, item: typing.Any) -> None:
        """处理单个待处理项的逻辑"""
        pass

Q: typing.TypeAlias = "mp.Queue"
ProcessorType = typing.Type[Processor]


class Worker(mp.Process):
    def __init__(self, queue: Q, processor_cls: ProcessorType):
        super().__init__()
        self.queue = queue
        self.processor_cls = processor_cls
        self.processor: Processor | None = None

    def run(self):
        try:
            # 初始化处理器
            self.processor = self.processor_cls()
            self.processor.setup()

            # 循环处理队列中的项
            while True:
                item = self.queue.get()
                if item is None:
                    # 收到结束信号，退出循环
                    break
                self.processor.process(item)
        except Exception as e:
            # 可根据需求添加日志或异常上报逻辑
            print(f"Worker进程异常: {e}")
        finally:
            # 可选：添加清理逻辑
            pass


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

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

    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 WorkerManager:
    def __init__(
        self, item_generator: typing.Iterable[typing.Any], processor_cls: ProcessorType, n_workers: int
    ) -> None:
        queue: Q = mp.Queue()
        with WorkerPool(n_workers, processor_cls, queue):
            for item in item_generator:
                queue.put(item)

改进说明

明确的接口规范：用户只需继承Processor抽象类并实现setup和process方法，无需关心生成器、send/next等底层细节
初始化与处理分离：setup方法在每个Worker进程启动时仅执行一次，process方法负责处理单个数据项，逻辑清晰
内置异常处理：Worker的run方法中加入了try-except块，捕获并处理所有可能的异常，避免进程意外崩溃且无提示
类型提示更清晰：通过抽象基类和类型别名，用户能直观理解需要实现的接口，减少误用
消除StopIteration异常：Worker内部处理结束信号的逻辑，无需用户手动处理生成器的终止异常

改进后用法示例

class PrintProcessor(Processor):
    def setup(self):
        # 初始化逻辑：比如建立数据库连接、加载模型等
        print("Worker初始化完成")

    def process(self, item):
        # 处理逻辑
        print(f"处理项: {item}")


items = range(10)
WorkerManager(items, PrintProcessor, 2)

这个示例会启动2个Worker进程，分别处理队列中的项，初始化代码仅执行一次，处理逻辑清晰可靠。

内容的提问来源于stack exchange，提问作者Izaak Cornelis

如何传递初始化代码与处理代码至其他进程？求更优可靠实现

现有实现代码 #

现有用法示例 #

改进思路 #

完整改进代码 #

改进说明 #