Prhub

#29834 Fix scheduler crash on prefill-unreachable decode abort

原始 PR 作者 amote-i 合并时间 2026-07-08 14:06 文件变更 2 提交数 3 评论 19 代码增减 +75 / -1

执行摘要

修复解码崩溃,清除 pending_reqs 并防御空 receiver

根据 PR 描述:当预填充启动服务器不可达时,一个 DecodeRequest 同时存在于 self.queue 和 self.pending_reqs(add() 慢路径)。清除路径(PR #28022)将其 kv_receiver 设置为 None,但只有 self.queue 被清理。剩余的 pending_reqs 引用在下一个 _resolve_pending_reqs 中导致崩溃:AttributeError: 'NoneType' object has no attribute 'abort'。

值得快速合入。修复明确且测试充分,使用 id() 进行身份比较的设计决策值得类似场景参考。

讨论亮点

ShangmingCai 对 pending_reqs 清理代码提出 'I think we could drop this',可能指 if failed_reqs: 条件。作者 amote-i 回应 kv_receiver is not None 检查只是后盾,真正的修复是清理 pending_reqs,因为其他调用点仍在访问 pending_reqs。最终决定保留该清理逻辑。引用:'The kv_receiver is not None check is just a backstop — other call sites (e.g. .init() in _resolve_pending_reqs) still touch entries in pending_reqs, so the cleanup is the actual fix.'

实现拆解

  1. _ensure_prefill_info 中增加保护:当超时重试失败后,在调用 decode_req.kv_receiver.abort() 之前检查 kv_receiver is not None,防止在返回值已清除时崩溃。
  2. pop_preallocated 中增加双列表清理:在移除队列中失败的请求后,使用 id() 构建失败请求集合,并从 self.pending_reqs 中过滤掉相同身份的对象,确保 _resolve_pending_reqs 不再访问已清理的请求。
  3. 测试配套:新增两个单元测试 test_prealloc_abort_also_drops_from_pending_reqstest_ensure_prefill_info_tolerates_cleared_receiver,分别验证共享对象移除和空 receiver 兼容性。
文件 模块 状态 重要度
python/sglang/srt/disaggregation/decode.py 分解模块 modified 5.97
test/registered/unit/disaggregation/test_decode_queue_cleanup.py 分解测试 modified 6.37

关键符号

_ensure_prefill_info pop_preallocated test_prealloc_abort_also_drops_from_pending_reqs test_ensure_prefill_info_tolerates_cleared_receiver

关键源码片段

test/registered/unit/disaggregation/test_decode_queue_cleanup.py test-coverage

新增两个回归测试,覆盖共享对象移除和空 receiver 兼容性,确保修复正确性。

def test_prealloc_abort_also_drops_from_pending_reqs(self):
    # Same DecodeRequest lives in both queue and pending_reqs (add() slow
    # path). Aborting must drop it from both, and compare by identity since
    # DecodeRequest's dataclass __eq__ would compare the tensor receiver.
    class BadEqReceiver(FakeReceiver):
        def __eq__(self, other):
            raise TypeError("use identity comparison, not value equality")
​
        __hash__ = object.__hash__
​
    receiver = BadEqReceiver()
    req = SimpleNamespace(
        rid="abort-shared",
        finished_reason=FINISH_ABORT("aborted"),
        return_logprob=False,
    )
    decode_req = SimpleNamespace(req=req, kv_receiver=receiver)
​
    queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
    queue.queue = [decode_req]
    queue.pending_reqs = [decode_req] # same object, dual ownership
    queue.retracted_queue = []
    queue._resolve_pending_reqs = MagicMock()
    queue._update_handshake_waiters = MagicMock()
    queue._uses_swa_tail_prealloc = MagicMock(return_value=False)
    queue._allocatable_token_budgets = MagicMock(return_value=0)
    queue._hicache_pending_restore_tokens = MagicMock(return_value=0)
​
    scheduler = MagicMock()
    scheduler.running_batch.reqs = []
    scheduler.enable_priority_scheduling = False
    scheduler.enable_hisparse = False
    scheduler.output_streamer = MagicMock()
    queue.scheduler = scheduler
​
    # Must not raise on the receiver __eq__ above.
    preallocated, failed = queue.pop_preallocated()
​
    self.assertEqual(preallocated, [])
    self.assertEqual(failed, [decode_req])
    self.assertEqual(queue.queue, [])
    self.assertTrue(all(r is not decode_req for r in queue.pending_reqs))
    self.assertIsNone(decode_req.kv_receiver)

评论区精华

pending_reqs 清理代码是否需要 `if failed_reqs:` 条件 设计

ShangmingCai 对新增的 pending_reqs 清理块评论 'I think we could drop this',暗示可以移除外层的 `if failed_reqs:` 检查。amote-i 回应说 `kv_receiver is not None` 检查只是后盾,真正的修复是清理 pending_reqs;因为其他调用点仍在接触 pending_reqs 中的条目,所以清理是实际修复,不应移除。最终决定保留该清理以及空检查。

结论:保留 pending_reqs 清理逻辑和空检查。 · 已解决

风险与影响

主要风险是身份比较依赖 id(),这可能受对象回收和重用影响,但 DecodeRequest 是长生命周期对象,风险较低。此外,其他可能访问 pending_reqs 的路径(如 _ensure_prefill_info 之外)也需要确认一致性,但当前测试覆盖了主要入口。没有引入新配置或外部依赖,回退简单。

影响范围限于 disaggregation 模式下的解码队列清理;修复了一个明确的崩溃,提高了稳定性。对于启用 disaggregation 的用户,此修复避免了在预填充服务器临时不可达时的调度器崩溃。对不启用 disaggregation 的用户无影响。测试覆盖主要回归场景,验证了修复有效性。

身份比较依赖 id() 的健壮性 仅修复了 pop_preallocated 路径,其他 pending_reqs 访问点未审计

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论