simple_streaming_demo.py — 一个独立的演示脚本

文件路径: verl/experimental/fully_async_policy/unittest/simple_streaming_demo.py

文件概述

一个独立的演示脚本，展示了全异步训练中使用的流式处理模式。它不依赖 verl 的任何组件，纯粹用 asyncio 演示了"数据流 -> 并发处理 -> 结果消费"的三协程架构。

关键代码讲解

核心架构

class SimpleStreamingSystem:
    def __init__(self, max_concurrent_tasks=4):
        self.data_queue = asyncio.Queue()    # 待处理数据队列
        self.result_queue = asyncio.Queue()  # 处理结果队列

    async def run_demo(self):
        # 三个并行协程
        stream_task = asyncio.create_task(self.data_stream())     # 数据流（生产者）
        submit_task = asyncio.create_task(self._submit_worker())  # 提交处理（调度器）
        consumer_task = asyncio.create_task(self._consumer_worker())  # 消费结果

三个协程的职责

1. data_stream - 模拟数据产生

async def data_stream(self):
    test_data = [{"id": f"task_{i}", "content": f"data_{i}"} for i in range(8)]
    await self.add_data_stream(test_data)
    await asyncio.sleep(3)  # 模拟间隔
    extra_data = [{"id": f"extra_{i}", ...} for i in range(5)]
    await self.add_data_stream(extra_data)
    await self.data_queue.put("DONE")  # 终止信号

1. _submit_worker - 并发控制

async def _submit_worker(self):
    active_tasks = set()
    while True:
        data_item = await self.data_queue.get()
        if data_item == "DONE":
            await asyncio.gather(*active_tasks, return_exceptions=True)
            break

        # 并发控制
        while len(active_tasks) >= self.max_concurrent_tasks:
            done_tasks, active_tasks = await asyncio.wait(
                active_tasks, return_when=asyncio.FIRST_COMPLETED
            )

        task = asyncio.create_task(self._process_data_async(data_item))
        active_tasks.add(task)

1. _consumer_worker - 结果消费

async def _consumer_worker(self):
    while True:
        result = await asyncio.wait_for(self.result_queue.get(), timeout=2.0)
        self.consumer_count += 1
        print(f"Consumed: {result['id']}")

与全异步训练的对应关系

演示组件	全异步训练对应
data_stream	FullyAsyncRollouter._feed_samples()
_submit_worker	FullyAsyncRollouter._processor_worker()
_consumer_worker	FullyAsyncTrainer._get_samples_from_queue()
data_queue	pending_queue
result_queue	MessageQueue

核心类/函数列表

名称	类型	说明
SimpleStreamingSystem	类	流式处理演示系统
main()	异步函数	演示入口

与其他模块的关系

• 独立的演示脚本，不被任何模块依赖
• 展示了 FullyAsyncRollouter 使用的核心异步模式

小结

这是一个教学用的演示脚本，帮助理解全异步训练的核心模式：数据流式输入、并发处理、结果异步消费。可以独立运行 python simple_streaming_demo.py 查看效果。