如何在Python多进程中正确使用Queue.task_done()处理未知数量任务

问题背景 #

现有一套基于任务队列的脚本，希望通过多CPU核心加速do_some_big_calculation(task_object)的执行，同时保证主进程的task_queue.join()能正确感知所有任务完成。自行修改多进程实现后出现卡住问题，需排查原因并给出可行方案。

原核心代码 #

myscript.py

from queue import Queue
from worker_script import Worker
    
task_queue = Queue(maxsize=1)

while some_other_functions:
    task_queue.put(Some_Stuff)          # 任务随机触发，数量不可预知

Worker(task_queue).start()

task_queue.join()

# 后续脚本需等待所有任务完成后执行

worker_script.py

class Worker:

    def __init__(self, task_queue):
        self.task_queue = task_queue

    def run(self):
        while True:
            task_object = self.task_queue.get()
            do_some_big_calculation(task_object)       # 需要多进程加速的步骤
            self.task_queue.task_done()

用户修改后的错误实现（导致卡住） #

import multiprocessing

class Worker:

    def __init__(self, task_queue):
        self.task_queue = task_queue

    def run(self):
        num_processes = multiprocessing.cpu_count() - 2
        jobs = []
        with multiprocessing.Pool(processes=num_processes) as pool:
            while True:
                new_job = pool.apply_async(self.sub_worker)
                jobs.append(new_job)

    def sub_worker(self):
        while True:
            task_object = self.task_queue.get()
            do_some_big_calculation(task_object)       # 需要多进程加速的步骤
            self.task_queue.task_done()

卡住原因分析 #

1. 队列类型错误：使用的queue.Queue是线程安全但非进程安全的，多进程直接访问会导致队列内部状态混乱（比如任务计数、锁机制失效），最终task_done()无法正确更新计数，join()永远等待。
2. 无限提交任务导致死锁：run()方法中while True会无限制向进程池提交sub_worker任务，每个sub_worker又是死循环，很快会占满进程池，主进程持续提交任务会陷入阻塞。
3. 类方法序列化问题：sub_worker作为Worker类的方法，多进程调用时会序列化整个Worker对象，但queue.Queue无法跨进程正确序列化，导致子进程访问的队列与主进程队列不是同一个，任务传递和状态更新完全失效。

正确解决方案 #

使用multiprocessing.JoinableQueue（进程安全的可等待队列）替代普通queue.Queue，配合多进程消费者模式，确保任务传递和完成状态追踪正常。

修改后的代码 #

myscript.py

import multiprocessing
from worker_script import Worker

# 替换为进程安全的JoinableQueue
task_queue = multiprocessing.JoinableQueue(maxsize=1)

# 动态生成任务
while some_other_functions:
    task_object = Some_Stuff
    task_queue.put(task_object)

# 设置工作进程数量
num_processes = multiprocessing.cpu_count() - 2

# 给每个工作进程发送结束标记（None），确保任务完成后进程能退出
for _ in range(num_processes):
    task_queue.put(None)

# 启动多进程Worker
worker = Worker(task_queue, num_processes)
worker.start()

# 等待所有任务完成
task_queue.join()

# 等待所有工作进程退出（可选，避免残留进程）
worker.join()

# 后续脚本执行

worker_script.py

import multiprocessing

def worker_process(task_queue):
    """单个工作进程的任务处理逻辑"""
    while True:
        task_object = task_queue.get()
        
        # 收到结束标记，退出循环
        if task_object is None:
            task_queue.task_done()
            break
        
        # 执行计算任务
        do_some_big_calculation(task_object)
        
        # 标记当前任务完成，更新队列计数
        task_queue.task_done()

class Worker:
    def __init__(self, task_queue, num_processes):
        self.task_queue = task_queue
        self.num_processes = num_processes
        self.processes = []

    def start(self):
        """启动所有工作进程"""
        for _ in range(self.num_processes):
            p = multiprocessing.Process(target=worker_process, args=(self.task_queue,))
            p.start()
            self.processes.append(p)

    def join(self):
        """等待所有工作进程退出"""
        for p in self.processes:
            p.join()

关键说明 #

• JoinableQueue的作用：专为多进程设计，支持get()、task_done()、join()方法，完全兼容原线程队列的逻辑，且能保证跨进程的状态一致性。
• 结束标记的必要性：每个工作进程需要收到None（或其他约定标记）才能退出，否则会一直阻塞在get()上，导致进程残留。
• 动态任务场景适配：如果some_other_functions是长期运行的任务生成逻辑（比如持续监听消息），只需在任务生成停止后再发送结束标记即可，不影响动态任务的处理。

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