Python中RabbitMQ消息并行处理方案咨询（附多进程尝试代码）

你的代码串行执行的核心原因有两个：一是创建Process时错误地直接调用了函数（target=between_func()），导致主进程会先执行完between_func才会创建子进程；二是原架构没有将计算密集的generate_image任务真正分配到独立进程中执行。以下是针对性的优化方案：

方案1：使用进程池复用进程（推荐） #

该方案通过进程池复用进程，避免每个任务重复加载StableDiffusion模型，大幅降低资源开销，同时保证计算任务的并行执行。

修改后的完整代码 #

import asyncio
import json
from multiprocessing import Pool
import aio_pika
from PIL import Image
from keras_cv.models import StableDiffusion

# 进程内全局模型变量，每个进程仅初始化一次
worker_model = None

def init_worker():
    """进程池初始化函数：每个进程启动时加载一次模型"""
    global worker_model
    worker_model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)

def generate_image(task_args):
    """计算密集型任务：复用已加载的模型生成图片"""
    prompt, task_id = task_args
    img = worker_model.text_to_image(
        prompt=prompt,
        batch_size=1,
        num_steps=10,
        seed=1539,
    )
    image_name = f'images/{task_id}.png'
    Image.fromarray(img[0]).save(image_name)

async def main(pool):
    """RabbitMQ异步消费者：接收消息并提交到进程池"""
    connection = await aio_pika.connect_robust("amqp://guest:guest@127.0.0.1/")
    queue_name = "requestQueue"

    async with connection:
        channel = await connection.channel()
        await channel.set_qos(prefetch_count=10)  # 限制预取消息数量
        queue = await channel.declare_queue(queue_name, auto_delete=False, durable=True)
        
        async with queue.iterator() as queue_iter:
            async for message in queue_iter:
                async with message.process():
                    json_str = json.loads(message.body)
                    # 将任务提交到进程池，异步执行
                    pool.apply_async(generate_image, args=((json_str.get('message'), json_str.get('id')),))

if __name__ == "__main__":
    # 创建进程池，指定初始化函数加载模型
    with Pool(initializer=init_worker) as pool:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(main(pool))
        loop.close()

方案优势 #

• 进程复用：每个进程仅加载一次模型，避免重复加载的时间和资源消耗
• 异步解耦：RabbitMQ消费者仅负责接收消息，不阻塞计算任务
• 真正并行：计算任务在独立进程中执行，完全避开GIL限制

方案2：每个消息创建独立进程（适合低频率任务） #

如果任务量较小，也可以为每个消息单独创建进程，缺点是每个进程都要重新加载模型，资源开销较高。

修改后的代码 #

import asyncio
import json
from multiprocessing import Process
import aio_pika
from PIL import Image
from keras_cv.models import StableDiffusion

def generate_image(prompt, task_id):
    """计算密集型任务：每次创建进程时加载模型"""
    model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
    img = model.text_to_image(
        prompt=prompt,
        batch_size=1,
        num_steps=10,
        seed=1539,
    )
    image_name = f'images/{task_id}.png'
    Image.fromarray(img[0]).save(image_name)

async def main():
    """RabbitMQ异步消费者：每个消息启动一个独立进程"""
    connection = await aio_pika.connect_robust("amqp://guest:guest@127.0.0.1/")
    queue_name = "requestQueue"

    async with connection:
        channel = await connection.channel()
        await channel.set_qos(prefetch_count=10)
        queue = await channel.declare_queue(queue_name, auto_delete=False, durable=True)
        
        async with queue.iterator() as queue_iter:
            async for message in queue_iter:
                async with message.process():
                    json_str = json.loads(message.body)
                    prompt = json_str.get('message')
                    task_id = json_str.get('id')
                    # 为每个消息创建独立进程并启动
                    p = Process(target=generate_image, args=(prompt, task_id))
                    p.start()

if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(main())
    loop.close()

原代码的核心错误修正 #

你最初创建进程的代码存在语法错误：

# 错误写法：直接调用函数，主进程会先执行完between_func才创建子进程
t1 = Process(target=between_func())
# 正确写法：传递函数对象，子进程会独立执行该函数
t1 = Process(target=between_func)

即使修正该错误，这种"多消费者进程"的架构也只是让多个消费者并行接收消息，但每个消费者内部的generate_image仍会串行执行，不如上述两种方案高效。

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