Prhub

#21875 fix: streaming session race condition + some metrics

原始 PR 作者 ishandhanani 合并时间 2026-04-13 09:05 文件变更 11 提交数 41 评论 12 代码增减 +966 / -26

执行摘要

修复流式会话竞态条件和内存泄漏,添加指标监控。

根据 PR body,主要动机是 'Stabilize streaming session KV lifecycle: fix multiple memory leaks, add safety for session close during active decoding, and add Prometheus metrics.' 关联的 Issue #22273 详细描述了流式会话中因请求中止导致的内存泄漏问题,需要修复以确保系统稳定性。

该 PR 值得精读,尤其是对于负责内存管理和会话调度的工程师。关注 session_controller.py 中的 close_on_finish 机制和 session_aware_cache.py 中的 KV 释放优化,这些设计决策在处理竞态条件和内存泄漏方面具有借鉴意义。建议结合单元测试理解修复的具体场景。

讨论亮点

Review 评论为空,表明 PR 在合并前可能没有公开的详细讨论。但从 commit 历史可见,有多次协作提交,如 hnyls2002 重命名变量(如 pending_close 改为 close_on_finish)和简化逻辑,显示团队内部对代码清晰度的关注。决策结论是接受这些更改以提高可维护性,未解决疑虑在上下文中未提及。

实现拆解

实现方案分为几个模块:

1) 调度器模块(scheduler.py):修改会话打开和关闭逻辑,添加错误处理,并在 DP-attention 下重复 OpenSessionReqOutput。
2) 会话控制器模块(session_controller.py):引入 close_on_finish 标志,推迟会话关闭直到请求完成,并添加日志记录。
3) 内存缓存模块(session_aware_cache.py 和 common.py):修复 match_prefix 中的 abort-skip 逻辑,优化 KV 缓存释放,包括 overalloc tail trim、session shrink 处理和标志同步。
4) 可观测性模块(metrics_collector.py 和 scheduler_runtime_checker_mixin.py):新增流式会话相关 Prometheus 指标,受 enable_streaming_session 标志控制。
5) 测试模块:添加单元测试(test_streaming_session_unit.py)和集成测试(test_streaming_session.py),验证修复效果。

文件 模块 状态 重要度
python/sglang/srt/managers/session_controller.py managers modified 8.0
python/sglang/srt/mem_cache/session_aware_cache.py mem_cache modified 9.0
python/sglang/srt/mem_cache/common.py mem_cache modified 7.0
python/sglang/srt/observability/metrics_collector.py observability modified 4.0
test/registered/unit/mem_cache/test_streaming_session_unit.py test added 5.0

关键符号

match_prefix cache_finished_req open_session close_session release_kv_cache _alive_streaming_session_count _session_held_tokens

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

变量重命名以提升清晰度 设计

从 commit 历史看,hnyls2002 将 pending_close 重命名为 close_on_finish,并简化逻辑,以提高代码可读性。

结论:更改被接受,以简化逻辑并增强可维护性。 · 已解决

修复活动解码时的会话关闭 正确性

添加了推迟机制,防止在请求仍在解码时关闭会话导致 KV 内存损坏。

结论:实现了 close_on_finish 标志,安全处理进行中的请求。 · 已解决

风险与影响

技术风险包括:

1) 核心路径变更(如 session_controller.py 和 session_aware_cache.py 中的逻辑修改)可能导致回归,影响所有流式会话请求。
2) 内存管理逻辑复杂,新增的 close_on_finish 标志和 KV 释放优化可能引入新 bug,例如在并发场景下。
3) 指标添加依赖 enable_streaming_session 标志,如果标志未正确设置,指标可能不准确或缺失。但通过添加了 7 个单元测试和集成测试,覆盖了关键场景,如 overalloc trim 和 abort 处理,降低了风险。

对用户:流式会话的稳定性和内存使用效率得到提升,减少了因内存泄漏导致的崩溃风险,改善了使用体验。对系统:添加了指标监控,便于运维团队实时跟踪流式会话状态,提升了可观测性。对团队:代码变更涉及多个核心模块(调度、内存缓存、可观测性),需要工程师熟悉新的会话生命周期管理逻辑,可能增加维护复杂性。影响范围限于启用流式会话的功能,但属于关键路径。

核心路径变更 内存管理复杂性 测试覆盖依赖

关联 Issue

#22273 Fix streaming session KV leak on aborted requests

完整报告

参与讨论