执行摘要

• 一句话：修复PP下is_fully_idle忽略在途微批次
• 推荐动作：该PR值得精读，特别是调度器idle-gating逻辑的细节和PP下竞态条件的分析。修复方案简洁，测试设计精妙，适合作为类似竞争条件修复的参考。

功能与动机

PR body指出：is_fully_idle() cannot see in-flight PP microbatches. 在流水线并行下，调度器将在途微批次保存在self.mbs中；当最后一块prefill chunk被调度后，chunked_req已被清空，running_batch等结构也为空，但微批次仍然在流水线中飞行。此时若处理/flush_cache，它会通过idle检查并重置KV缓存，导致后续处理批结果时触发assert node is self.root_node错误。这是PP+分块预填充部署调用/flush_cache时可能复现的bug。

实现拆解

1. Scheduler idle-gating修复：在python/sglang/srt/managers/scheduler.py中，提取_pp_microbatches_drained()方法，在原有running_mbs检查基础上，额外检查self.mbs中所有微批次槽是否为空。is_fully_idle()改用该方法，确保idle判断考虑所有在途微批次。
2. 测试框架请求可见性修复：在python/sglang/test/scripted_runtime/context/queries.py的_get_all_reqs()中，当PP开启时，遍历所有微批次槽（mbs、last_mbs、running_mbs），防止在途请求被遗漏导致finished状态提前为true。
3. 测试框架重置逻辑修复：在python/sglang/test/scripted_runtime/scheduler_hook.py的_reset_engine_state()中，改用scheduler.is_fully_idle()判断引擎是否安静，并在超时时抛出RuntimeError，避免静默flush_cache。
4. 新增回归测试：在test/registered/chunked_prefill/test_scripted_core_4gpu.py中新增test_pp_flush_cache_during_inflight_chunk_results，该测试精确构造竞争窗口（队列和当前槽清空但微批次在途），执行flush_cache并断言请求能正常完成。

关键文件：

• python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 _pp_microbatches_drained）: 核心修复文件，新增_pp_microbatches_drained方法并修改is_fully_idle调用。
• test/registered/chunked_prefill/test_scripted_core_4gpu.py（模块 测试用例；类别 test；类型 test-coverage；符号 test_pp_flush_cache_during_inflight_chunk_results, _script_flush_during_inflight_chunk_results）: 新增回归测试，精确验证修复效果。
• python/sglang/test/scripted_runtime/scheduler_hook.py（模块 测试钩子；类别 test；类型 test-coverage）: _reset_engine_state改用is_fully_idle驱动等待，并在超时时抛出异常。
• python/sglang/test/scripted_runtime/context/queries.py（模块 查询工具；类别 test；类型 test-coverage）: _get_all_reqs在PP模式下遍历所有微批次槽，确保在途请求可见。

关键符号：_pp_microbatches_drained, _reset_engine_state, _get_all_reqs, test_pp_flush_cache_during_inflight_chunk_results, _script_flush_during_inflight_chunk_results

关键源码片段

python/sglang/srt/managers/scheduler.py

核心修复文件，新增_pp_microbatches_drained方法并修改is_fully_idle调用。

def is_fully_idle(self, for_health_check=False) -> bool:
    idle = (
        self.running_batch.is_empty()
        and self.chunked_req is None
        and not self.dllm_manager.any_staging_reqs()
        and (self.last_batch is None or self.last_batch.is_empty())
        and (self.cur_batch is None or self.cur_batch.is_empty())
        and (not self.enable_overlap or len(self.result_queue) == 0)
        and self._pp_microbatches_drained()
    )
    # ... 其余检查 ...
    return idle
​
def _pp_microbatches_drained(self) -> bool:
    """检查所有PP微批次是否都已排空。"""
    if self.ps.pp_size == 1:
        return True # 单机无 PP
    return all(x.is_empty() for x in self.running_mbs) and all(
        mb is None or mb.is_empty() for mb in self.mbs
    )

评论区精华

gemini-code-assist[bot]提出两条建议：在_pp_microbatches_drained和_get_all_reqs中，使用getattr和默认空列表来访问running_mbs、mbs、last_mbs等属性，避免在调度器初始化不完全时引发AttributeError。这些建议未在PR中采纳或解决，但PR已合并。

• 防御性访问mbs/running_mbs (design): 未采纳，PR已合并。

风险与影响

• 风险：
  • 回归风险：修改了is_fully_idle()，该函数用于多个场景（flush_cache、HiCache attach/detach、健康检查等）。但通过保留pp_size==1短路和新增测试，风险较低。
  • 性能影响：循环检查所有微批次槽（最多pp_async_batch_depth + 1个），仅在idle判定时执行，非热点路径，影响可忽略。
  • 防御性不足：gemini-code-assist[bot]提出的AttributeError风险在现有初始化顺序下不会触发，但未来重构可能暴露。
• 影响：
  • 用户影响：修复了PP+分块预填充场景下flush_cache导致KV缓存损坏的bug，对使用RL权重更新等流程的用户至关重要。
  • 系统影响：idle-gating行为更准确，使所有依赖idle状态的操作更加安全。
  • 团队影响：新增的回归测试为同类竞态条件提供了可复现的验证方法。
  • 风险标记：PP idle-gating 变更, 竞争条件修复, KV缓存损坏风险

关联脉络

• PR #27445 Complete server warmup before scripted runtime scripts start: 本PR基于#27445引入的测试框架，并修复了一个被其暴露的竞争条件。