# PR #1673 完整报告

- 仓库：`radixark/miles`
- 标题：[feat] support `sample` rollout submission granularity to keep fully async concurrency
- 合并时间：2026-08-05 09:12
- 原文链接：http://prhub.com.cn/radixark/miles/pull/1673

---

# 执行摘要

- 一句话：新增 sample 粒度提交回填，消除 fully-async 拖尾组瓶颈
- 推荐动作：值得精读。重点学习 `SubmissionScheduler` 的并发账目设计、异常路径 cancel-and-settle、per-driver 默认值，以及 review 中关于 credit 守恒的辩论——对做异步调度系统的人有教科书级借鉴价值。

# 功能与动机

PR body 明确指出：在 group-based RL（GRPO 风格）长 horizon agentic rollout 中，异步调度器以 group 粒度回填，每个组都被其最慢的兄弟任务拖住，引擎容量闲置。原文称 "In our production runs (n=8, 32 concurrent groups on SWE tasks) this straggler gating was a primary rollout-throughput limiter"，即 straggler 门控是生产环境的主要吞吐瓶颈。

# 实现拆解

1. 新增提交调度器抽象 `miles/rollout/submission_scheduler.py`：包含 `GroupLevelSubmission`（保留原始 group 粒度语义，逐字复刻旧行为）和 `SampleBackfillSubmission`（样本完成 credit 计数器 + `asyncio.Event` 唤醒），`make_submission_scheduler` 按 `--rollout-submission-granularity` 和 driver 默认值解析。容量判定从 " 在飞 group 数 < 预算 " 改为 " 在飞样本数 + group_size <= 预算 × group_size"，并内置孤儿 credit 重同步。
2. 在 `generate_and_rm_group`（`miles/rollout/inference_rollout/inference_rollout_common.py`）中挂接 `sample_done_callback`：每个样本子任务通过 `task.add_done_callback` 注册回调，成功 / 异常 / 取消均触发以保证 credit 守恒；异常路径先取消并 settle 全部兄弟任务再抛出，避免孤儿生成与迟到回调错误扣账。
3. 两个 driver 接入：`generate_rollout_async`（`inference_rollout_train.py`，默认 `group`）与 `FullyAsyncRolloutFn._worker_loop`（`fully_async_rollout.py`，默认 `sample`）统一将提交循环改为 `scheduler.has_capacity` + `scheduler.wait_for_progress`；`_submit_one_group` 增加 `on_submit` 记账，`submit_generate_tasks` 透传 `sample_done_callback`。
4. 参数与文档：`miles/utils/arguments.py` 新增 enum 参数 `--rollout-submission-granularity {group,sample}`，默认 `None` 表示由 driver 各自决默认；`docs/user-guide/cli-reference.md` 同步说明。
5. 测试配套：新增 `tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py`，使用 `Harness` 手工控制 group/sample 完成时序，验证 sync 默认 group、回填替换组提前提交、差一个样本不提交、异常取消 conserve credits、调度器单元测试；更新 `tests/fast/rollout/test_fully_async_rollout.py` 覆盖 worker 默认 sample 粒度与显式 group opt-out。

关键文件：
- `miles/rollout/submission_scheduler.py`（模块 调度器；类别 source；类型 core-logic；符号 GroupLevelSubmission, SampleBackfillSubmission, make_submission_scheduler, has_capacity）: 核心新增文件，定义两种提交调度策略与工厂方法，是本次变更的调度语义核心。
- `miles/rollout/inference_rollout/inference_rollout_common.py`（模块 公共逻辑；类别 source；类型 core-logic；符号 generate_and_rm_group）: 关键行为改造点：挂接样本完成回调并处理异常路径，保证 credit 守恒，是整个回填机制的账本基础。
- `miles/rollout/inference_rollout/inference_rollout_train.py`（模块 同步驱动；类别 source；类型 core-logic；符号 submit_generate_tasks, generate_rollout_async）: sync rollout 驱动接入调度器，默认 group 粒度，并新增 submit_generate_tasks 的 callback 透传。
- `miles/rollout/fully_async_rollout.py`（模块 异步驱动；类别 source；类型 core-logic；符号 FullyAsyncRolloutFn, _submit_one_group, _worker_loop, _generate_group）: fully-async driver 默认 sample 粒度，是本次优化主要受益方；worker 循环改为调度器驱动。
- `miles/utils/arguments.py`（模块 参数解析；类别 source；类型 configuration）: 新增 --rollout-submission-granularity 参数并说明 per-driver 默认值，是 CLI 面入口。
- `tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py`（模块 回填测试；类别 test；类型 test-coverage；符号 Harness, make_args, make_group, fake_submit_generate_tasks）: 新增 285 行测试，用 Harness 手工控制完成时序，覆盖回填时机与边界行为，是本次行为验证主体。
- `tests/fast/rollout/test_fully_async_rollout.py`（模块 异步测试；类别 test；类型 test-coverage；符号 default_generate, blocking_generate, failing_generate, test_worker_defaults_to_sample_granularity）: 更新 fully-async 测试 stub 并新增默认 sample 粒度与 group opt-out 两个用例。
- `docs/user-guide/cli-reference.md`（模块 用户文档；类别 docs；类型 documentation）: CLI 参考文档补充新参数说明。

关键符号：make_submission_scheduler, SampleBackfillSubmission.has_capacity, SampleBackfillSubmission.sample_done_callback, SampleBackfillSubmission.wait_for_progress, GroupLevelSubmission.has_capacity, generate_and_rm_group, submit_generate_tasks, generate_rollout_async, FullyAsyncRolloutFn._worker_loop, FullyAsyncRolloutFn._submit_one_group

## 关键源码片段

### `miles/rollout/submission_scheduler.py`

核心新增文件，定义两种提交调度策略与工厂方法，是本次变更的调度语义核心。

```python
"""Paces prompt-group submission for the rollout drivers (--rollout-submission-granularity).

本模块只负责“何时允许提交新组”；在飞 group 任务集合由 driver 持有，
通过 has_capacity 的 pending_groups 参数传入。调度器只记录在飞样本数。
"""

import asyncio
import logging
from argparse import Namespace

from miles.utils.types import Sample

logger = logging.getLogger(__name__)


class GroupLevelSubmission:
    """A submission slot frees only when the whole group task returns."""

    sample_done_callback = None  # sample 粒度专属，group 粒度不接线

    def has_capacity(self, *, pending_groups: int, group_budget: int) -> bool:
        # 原始语义：在飞 group 数小于预算即可继续提交
        return pending_groups < group_budget

    def on_submit(self, groups: list[list[Sample]]) -> None:
        pass

    async def wait_for_progress(self, pendings: set[asyncio.Task]) -> tuple[set[asyncio.Task], set[asyncio.Task]]:
        # group 粒度下，进度只由 group 任务完成驱动
        return await asyncio.wait(pendings, return_when=asyncio.FIRST_COMPLETED)


class SampleBackfillSubmission:
    """Each finished sample frees its own slot; a replacement group fits once ``group_size`` samples complete."""

    def __init__(self, group_size: int):
        self.group_size = group_size          # 每组样本数，攒够这些 credit 释放一个替换名额
        self.samples_in_flight = 0            # 在飞样本计数，用于容量判定
        self._sample_done = asyncio.Event()   # 样本完成唤醒事件

    def sample_done_callback(self) -> None:
        # 每个样本任务结束（成功 / 异常 / 取消）恰好触发一次，保证在飞计数守恒
        self.samples_in_flight -= 1
        self._sample_done.set()

    def has_capacity(self, *, pending_groups: int, group_budget: int) -> bool:
        """A False return arms the sample wakeup: only later completions wake ``wait_for_progress``."""
        if pending_groups == 0 and self.samples_in_flight:
            # 兜底重同步：group 返回但从未生成样本任务（如 stub 跳过 callback）时，
            # 这些 credit 永远不会被归还，会永久卡死容量，此处直接清零
            logger.warning(f"samples_in_flight={self.samples_in_flight} with no pending groups; resetting to 0")
            self.samples_in_flight = 0
        if self.samples_in_flight + self.group_size <= group_budget * self.group_size:
            return True
        # 容量不足时清除事件，消除 lost-wakeup 窗口
        self._sample_done.clear()
        return False

    def on_submit(self, groups: list[list[Sample]]) -> None:
        self.samples_in_flight += sum(len(group) for group in groups)

    async def wait_for_progress(self, pendings: set[asyncio.Task]) -> tuple[set[asyncio.Task], set[asyncio.Task]]:
        # 同时等待 group 完成或单个样本完成，任一发生即返回
        waiter = asyncio.create_task(self._sample_done.wait())
        try:
            done, pending = await asyncio.wait(pendings | {waiter}, return_when=asyncio.FIRST_COMPLETED)
        finally:
            waiter.cancel()
        return done - {waiter}, pending - {waiter}


def make_submission_scheduler(args: Namespace, *, default: str) -> GroupLevelSubmission | SampleBackfillSubmission:
    # 显式参数优先，否则按 driver 各自默认：fully-async 为 sample，sync 为 group
    granularity = args.rollout_submission_granularity or default
    if granularity == "group":
        return GroupLevelSubmission()
    assert granularity == "sample", f"unknown submission granularity: {granularity}"
    return SampleBackfillSubmission(args.n_samples_per_prompt)

```

### `miles/rollout/inference_rollout/inference_rollout_common.py`

关键行为改造点：挂接样本完成回调并处理异常路径，保证 credit 守恒，是整个回填机制的账本基础。

```python
async def generate_and_rm_group(
    state: GenerateState,
    group: list[Sample],
    sampling_params: dict[str, Any],
    evaluation: bool = False,
    sample_done_callback: Callable[[], None] | None = None,
) -> list[Sample]:
    args = state.args

    if state.aborted:
        return group

    if policy_uses_routing_key(args):
        for sample in group:
            if sample.routing_key is None:
                sample.routing_key = str(uuid.uuid4())

    log_prefix = f"[group indices={[getattr(s, 'index', '?') for s in group]}]"
    logger.debug(f"{log_prefix} Starting group with {len(group)} samples")
    tasks = []
    for idx, sample in enumerate(group):
        current_sampling_params = sampling_params.copy()
        if getattr(args, "sglang_enable_deterministic_inference", False):
            current_sampling_params["sampling_seed"] = args.rollout_seed + idx
        task = asyncio.create_task(generate_and_rm(state, sample, current_sampling_params, evaluation=evaluation))
        if sample_done_callback is not None:
            # 成功、异常、取消三种结局都触发 callback：每个最终结束的样本都对应一次容量释放，
            # 失败样本也会释放引擎额度，跳过它们会导致并发额度单调泄漏
            task.add_done_callback(lambda _task: sample_done_callback())
        tasks.append(task)

    try:
        group = await asyncio.gather(*tasks)
    except BaseException:
        # 单个样本失败时 gather 会传播第一个异常，但兄弟任务仍在运行：
        # 它们既浪费引擎资源，其迟到 callback 又可能在调度器重同步后错误扣减新组计数。
        # 先取消并等待所有样本任务 settle，保证抛异常时账目已经精确
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        raise
    logger.debug(f"{log_prefix} [group] All {len(group)} samples completed")
    if state.aborted:
        return group

    if args.group_rm:
        await batched_async_rm(args, group, inplace_set_reward_field=True)

    return group

```

# 评论区精华

1. guapisolo 在 issue 中提问为何不做直接 `SampleLevelSubmission` 而用回填：担心大 group（64/128）下 group 粒度回填仍不够细。结论：以 credit 聚合的回填在保留组提交结构的同时获得样本级释放，后续可演进。
2. guapisolo 提出 Codex 边缘情况：一个样本异常 → `asyncio.gather` 提前退出 → 兄弟任务仍在跑 → `pendings` 变空 → `has_capacity` 重置计数 → 迟到回调错误扣减替换组。最终由 commit `42bf672f` 在异常路径 cancel-and-settle 解决。
3. gemini-code-assist 建议只在成功时触发 `sample_done_callback`，ys-2020 反驳：成功 / 异常 / 取消都触发才能精确守恒并发额度，跳过失败样本会导致额度单调衰减。
4. gemini-code-assist 提出数据源耗尽死锁，ys-2020 确认在 `f814b44` 修复。
5. gemini-code-assist 建议 group 异常后立即 `submit_groups(1)` 补位，ys-2020 反驳：样本回调已为失败组返回 n 个 credit，再补会 double-replenish。
6. nblintao 询问是否未来默认开启 sample 粒度，yueming-yuan 同意，后续演进为 per-driver 默认。

- 为何选择 SampleBackfill 而非直接 SampleLevelSubmission (design): 以 credit 聚合的回填在保持 group 提交结构的同时获得样本级释放，后续可演进为更细粒度；issue 评论中 yueming-yuan 说明调度器抽象让两种策略仅差两个决策点。
- 异常引起的迟到回调是否会错误扣减替换组计数 (correctness): commit 42bf672f 在异常路径 cancel-and-settle 全部兄弟任务后再抛异常，保证账目在 reap 时精确。
- 失败 / 取消的样本是否应触发 done callback (correctness): 保持 callback 对成功 / 异常 / 取消都触发，并补充注释说明不变量。
- 数据源耗尽时死锁 (correctness): ys-2020 确认在 f814b44 修复：submit_groups 返回 0 且无在飞任务时跳出循环。
- group 异常后是否需要立即补充提交 (design): 保持纯 credit 驱动回填，不额外补位。
- 是否将 sample 粒度设为默认 (design): 后续演进为 per-driver 默认：fully-async 默认 sample，sync 默认 group。

# 风险与影响

- 风险：
 1. 并发账目精度敏感：`SampleBackfillSubmission` 的 credit 计数依赖回调恰好触发一次，曾出现孤儿 credit 死锁（commit `45705738` 修复，stage-a-cpu 挂到 60 分钟超时），当前兜底逻辑依赖 `pending_groups == 0` 的时序判断，driver 改动若绕过该路径可能回归。
 2. `generate_and_rm_group` 异常路径复杂度上升：cancel-and-settle 与 `group_rm`、`aborted` 路径交织，异常时 `batched_async_rm` 不会执行，行为与旧版不同。
 3. fully-async 默认行为变更：存量用户未设参数时在飞 group 上限放大最多 `n_samples_per_prompt` 倍，对数据源、reward 模型和 SGLang 引擎负载有隐性压力。
 4. 作用域限定：仅 class-based rollout API（`MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`）生效，默认 `generate_rollout` 路径不受影响，降低回归面。
 5. 测试缺口：CPU 测试覆盖时序逻辑，但无真实 GPU straggler 场景的 e2e 验证。
 - 影响：对用户：fully-async/ 长 horizon agentic 训练用户默认获得样本级回填，引擎利用率提升；sync 用户需显式 opt-in 并理解 abort 丢弃语义。对系统：并发上限从 `rollout_batch_size` 扩展至约 `rollout_batch_size × n_samples_per_prompt`，对数据源与 reward 模型并发压力上升。对团队：调度策略收敛为可插拔抽象，后续新策略成本低，并发不变量 " 每个样本任务恰好一次回调 " 被文档化。
 - 风险标记：并发账目精度敏感 , fully-async 默认行为变更 , 异常路径复杂度上升 , 仅实验 API 生效

# 关联脉络

- PR #1717 Rewrite fully-async rollout as FullyAsyncRolloutFn on the class-based rollout API: 本 PR 在 issue 评论中明确说明 rebase 到 #1717 之上，并按 FullyAsyncRolloutFn 类式 API 重构实现，两个 PR 是同一演进主线。
- PR #1716 Move fully-async rollout from examples into miles/rollout: fully-async rollout 迁入核心库的前奏，本 PR 继续在该路径上扩展调度能力。
- PR #1740 feat (async): support evaluation for fully-async training (dedicated fleet / pause-the-world / external service): 本 PR 合并 main 时解决了其引入的 eval-pause 测试 stub 冲突（commit acba7df0 说明），两 PR 均修改 fully_async_rollout.py。