
本文介绍如何改造传统的单条数据生产/消费模式,让 producer 批量入队、consumer 批量出队,以满足机器学习模型对批量样本(如 2000 行)的输入需求,显著提升 i/o 与计算协同效率。
本文介绍如何改造传统的单条数据生产/消费模式,让 producer 批量入队、consumer 批量出队,以满足机器学习模型对批量样本(如 2000 行)的输入需求,显著提升 i/o 与计算协同效率。
在构建异步数据处理流水线时,尤其是对接机器学习模型(如需固定 batch size=2000),逐行 queue.put(item) 和 queue.get() 的方式不仅低效,还会导致 Consumer 频繁唤醒、模型无法满载运行。根本解法是将数据按批次组织为列表(list),整批入队;Consumer 一次性取出整批数据进行处理——queue.Queue 原生支持任意 Python 对象(包括 list、dict、tuple 等),无需额外序列化。
✅ 正确实践:按批次生产和消费
以下是一个精简、可直接复用的示例,展示如何将数据库查询结果按指定大小(如 BATCH_SIZE = 2000)分组,并通过 Queue 高效传递:
import threading
import queue
import time
BATCH_SIZE = 2000
def read_data():
# 模拟从 DB 获取全部数据(实际中建议流式分页,避免内存溢出)
torque_data_queryset = Torque.get_torque_data()
return list(torque_data_queryset)
def chunked_iterable(data, size):
"""将数据列表按 size 分块生成器"""
for i in range(0, len(data), size):
yield data[i:i + size]
def producer(queue: queue.Queue):
print('Producer: Running')
data = read_data()
# 按 BATCH_SIZE 分块,每块作为独立 list 入队
for batch in chunked_iterable(data, BATCH_SIZE):
time.sleep(0.1) # 模拟 I/O 延迟(可选)
queue.put(batch)
print(f'> Producer added batch of {len(batch)} rows')
# 发送终止信号(单个 None 即可,代表所有批次发送完毕)
queue.put(None)
print('Producer: Done')
def consumer(queue: queue.Queue):
print('Consumer: Running')
while True:
batch = queue.get() # ← 关键:一次取出整个 list
if batch is None:
queue.task_done()
break
# ✅ 此处可直接送入 ML 模型:model.predict(batch)
print(f'> Consumer processing batch of {len(batch)} rows')
# 模拟模型推理耗时
time.sleep(0.5)
queue.task_done()
print(f'> Batch processed successfully')
print('Consumer: Done')
def thread_function():
q = queue.Queue()
consumer_thread = threading.Thread(target=consumer, args=(q,))
producer_thread = threading.Thread(target=producer, args=(q,))
consumer_thread.start()
producer_thread.start()
# 等待生产者结束(确保所有 batch 已入队)
producer_thread.join()
# 等待消费者处理完所有任务(含终止信号)
q.join()
consumer_thread.join()
# 调用入口(Django 视图中建议改用异步任务如 Celery,此处仅作演示)
def thread_result(request):
thread_function()
return render(request, 'async_processing_result.html',
{'message': f'Batch processing completed with {BATCH_SIZE}-row chunks!'})⚠️ 关键注意事项
- 内存安全:若数据库数据量极大(如千万级),read_data() 一次性加载到内存存在 OOM 风险。推荐改用游标分页 + 流式读取(例如 Django 的 iterator(chunk_size=BATCH_SIZE)),边查边入队。
- 终止信号设计:仅需一个 None 作为结束标记,不要为每个 batch 发送 None,否则 Consumer 会提前退出。
- 线程同步:调用 queue.task_done() 配合 queue.join() 可确保 Consumer 处理完所有任务后再退出,比轮询 queue.empty() 更可靠。
- 异常容错:实际工程中应在 producer 和 consumer 内增加 try/except,捕获数据库异常、模型预测失败等,并通过日志或重试机制保障健壮性。
- 扩展建议:当吞吐量进一步提升时,可引入 concurrent.futures.ThreadPoolExecutor 或切换至 multiprocessing.Queue(跨进程);对于高并发场景,推荐使用消息队列(如 Redis Streams、RabbitMQ)替代内存队列。
通过批量操作,你不仅能匹配 ML 模型的输入契约,还能减少线程上下文切换、降低锁竞争、提升整体吞吐。记住:Queue 不是“只能传单个值”的管道,而是灵活的数据容器——善用其承载结构化批量数据,是构建高性能数据流水线的关键一步。


















