Skip to content

Refactor RL disagg producer flow with deeper manager-side abstraction#1769

Open
jayhenry wants to merge 11 commits intoInternLM:rl_designfrom
jayhenry:rl_redesign
Open

Refactor RL disagg producer flow with deeper manager-side abstraction#1769
jayhenry wants to merge 11 commits intoInternLM:rl_designfrom
jayhenry:rl_redesign

Conversation

@jayhenry
Copy link
Copy Markdown
Collaborator

@jayhenry jayhenry commented May 8, 2026

动机

这次重构的核心目标是把分散在 AgentLoopManager 里的隐藏信息收进更合适的局部抽象里,让 manager 只负责流程编排。

主要隐藏的信息包括:

  • ProduceProgress 隐藏 producer / consumer 的累计进度口径,包括 target、consumed、future step 的推进规则。
  • ReplayBuffer 隐藏 batch 是否 ready、如何取 batch、如何统计状态、如何刷新 staleness 等 buffer 操作细节。
  • ProduceContext 隐藏 strategy 运行时需要的 task 上下文,包括 sampler、replay buffer、progress、abort 信号、model step、生成结果入库顺序。
  • _PendingTasks 隐藏 async pending rollout task 的并发协议,包括 schedule、wait、claim、cancel,避免 manager 直接读取 strategy 私有状态。

这样 AgentLoopManager 不再需要理解这些底层操作细节,只保留 colocate / disagg 的生产、消费、暂停、恢复等流程控制。

主要改动

  • 收敛 ProduceStrategy.produce_batch(ctx) / pause_produce(ctx) 参数为 ProduceContext
  • 新增/完善 ProduceProgress,统一维护生产和消费进度。
  • 扩展 ReplayBuffer,承接通用 batch 和 staleness 操作。
  • 新增 _PendingTasks,封装异步 pending task 集合操作。
  • AgentLoopManager.save() 通过 ProduceStrategy.pending_task_count() 检查 pending task,不再访问策略私有字段。
  • 简化 _produce_batch_to_buffer,显式传入 task_batch_sizes,并统一 single/multi task 生产路径。
  • 拆分 _get_batch_from_buffer 的日志和 result 构造逻辑。
  • 压平 produce_loop 的状态等待分支。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant