[xs1997zju]

Root Cause

Long multi-turn traces can create a trie deeper than Python's recursion limit. Both Trie.keys() and full-tree Trie.release() recursively traverse this tree, so exporting or cleaning up a long trace can raise RecursionError and interrupt training.

Fix

Use explicit stacks for key enumeration and full-tree cleanup. Preserve insertion-order DFS for keys and child-before-parent cleanup order. Single-key release, session identifiers, inference routing, and KV-cache configuration are unchanged.

Test Plan

• A 1,600-node prefix chain reproduces RecursionError in both methods on the base revision; both complete with this fix.
• tests/rl/test_trace_store.py: 8 passed, no skips, including a local CPU Ray actor.
• Regression tests cover deep traversal, cleanup order, partial deletion, empty/root keys, non-lexical insertion order, repeated cleanup, and reuse after cleanup.
• Ruff lint/format passed for both changed files; full-package mypy passed (417 source files).
• Full-repository pre-commit was run: existing formatting violations were found outside this diff; unrelated automatic edits were reverted. The Pydantic full-module import check initially hit sandbox device restrictions; rerunning outside the sandbox passed (optional NPU/vLLM modules were skipped by the checker).
• GPU/model training and KV-cache hit-rate measurements were not run. This change only replaces in-memory trie traversal; no training-kernel benchmark is needed.

Final Review

Full final review report, including before/after flowcharts.

Real Kimi CLI subprocesses reviewed the change with the original xt-review prompt. The initial key-path suggestion was implemented and re-reviewed. Final code verdict: APPROVE, with no outstanding findings introduced by this PR. The published report identifies the author's diagram corrections and excludes an unchanged docstring observation from PR findings under the repository's review-scope rule.

Reviewed head: 8c7b266dfc996c9700b3826524da24e60fb87328. Session persistence was verified. Session identifiers are recorded in the report.

Agent RL delivery and model scope

This fix removes a Python recursion-depth limit in model-independent trace storage. The author has confirmed GLM-5.3-Flash compatibility in the integrated Agent work. This statement is separate from the exact-branch validation above: this update does not claim a new full-size GLM training or KV-cache benchmark on this standalone head.

The new Agent-lifecycle core #2297 contains overlapping iterative traversal and request-aware shared-reference ownership. Reconcile that overlap before landing both; replacing the complete trace-store file would change ownership behavior. Original #2257 remains the pre-split reference.

Developer guide: feature/PR map, diagrams, scripts and validation scope.

Keep Draft. See the live Checks tab for current CI state; the historical final-review comment records the check state at its publication time. No code or head revision was changed by this description update.

[xs1997zju]

终版：真实 Kimi CLI 按 xt-review 原始 prompt 审阅，作者核对后整理发布。
Reviewed head: 8c7b266dfc996c9700b3826524da24e60fb87328；base: f0a19dc5e0ed4ac2eec8c059326d396dfeece665。
Kimi verdict 为 APPROVE。作者校正了流程图的直接调用边，并按仓库“只报告本 PR diff 问题”的规则，将 Kimi 记录的 1 条未改动 docstring 历史备注排除在本 PR 问题计数之外；原始审阅记录保留在本地，未改写。
CLI: kimi-code/kimi-for-coding；最终核验 session: session_7ac24bc4-755d-420e-8cb4-0fd63e9a30d5（持久化会话已验证）。
远端 lint、两种语言文档构建已通过；unit_test 仍在等待 h_cluster runner。PR 暂保留 draft。

Review: PR #2252 — Trie keys()/release() 递归改迭代

• Base: f0a19dc → Head: 8c7b266（含两个 commit：f6754a6b4 迭代化修复、8c7b266df 简化 key 遍历并覆盖 root value）
• 改动范围: xtuner/v1/rl/rollout/trace_store.py（仅 Trie.keys、Trie.release(key=None) 分支）+ tests/rl/test_trace_store.py（新增 TestDeepTrie 2 个用例）
• 性质: 崩溃修复（深链触发 RecursionError）；keys 遍历使用字符串前缀。

ProduceBatchResult impact: not affected —— 未触碰 ProduceBatchResult 的任何字段、状态转换或计时语义。
本修复仅保证深链 Trie 的清理（release）与错误诊断（keys）能正常完成，间接防止清理路径因
RecursionError 中断；但仅按现有单测与源码审阅支持"避免清理抛错"这一结论，不额外保证
batch 时序/调度层面的行为。
RoutedExperts impact: 字段与语义均未改变 —— routed_experts 的产生（SessionServer 写入）、导出
（export_training_trace）与释放协议（release → _free_ray_refs → ray.internal.free）保持原样；
唯一变化是深链 Trie 上 release(None) 的遍历方式由递归改为迭代，从而原释放协议在深链下也能执行完毕，
routed-experts ObjectRef 不再因遍历崩溃而滞留。
Ray concurrency impact: not affected —— 未触碰任何 @ray.method、actor 构造、concurrency group
或 rollout 控制路径。

Summary

本次 PR 将 Trie.keys() 与 Trie.release(key=None) 从递归 DFS 改为显式栈迭代：

1. 修复真实缺陷：Trie.insert 是压缩前缀树——插入单个长 key（如 "x"*1600）只产生一条边（token 为整个剩余串，深度为 1），并不会因 key 文本长度变深；但连续插入互为前缀的 key（"x", "xx", "xxx", …）时沿已有前缀继续增加子节点，链深随插入轮数线性增长，可达 1600 层，超过 Python 默认递归上限（约 1000）。旧递归实现下 keys()/release(None) 在深链会话上必然 RecursionError。
2. release(session_id) 的异常后果（RolloutTraceStore.release, trace_store.py:337）：实现先 sessions.pop(session_id) 再 trie.release(key)。旧递归实现在深链上抛 RecursionError 时，异常向上传播、显式清理（_free_ray_refs → ray.internal.free）被中断，Trie 内后续节点的 Ray ObjectRef（含 routed experts）不再被显式释放。Python 引用/GC 的具体回收时点依赖实现，不能断言"直到 actor 被 GC"，但"显式 free 路径被中断"是代码可直接确认的事实。本修复消除了该失败模式。
3. keys 栈帧优化：栈帧保存字符串前缀（path + token），替代旧版 path + [token] 列表 + "".join，省去每节点一次 join。
4. 行为等价性：遍历顺序（keys 为 DFS 且兄弟按插入序；release(None) 为后序，"先子树、后 _free_ray_refs(node.value)、再 children.clear()"）、空字符串 key（value 落在 root）、部分 release 剪枝语义均与旧实现一致，新测试的断言顺序印证了这一点。

真实调用链核实（逐个 rg/读源码确认）：

• 写入方（生产者）: SessionServer（session_server.py）——仅调用 store.search.remote（:363, :466）与 store.insert.remote（:498, :503），不调用 export_training_trace，也不调用 release。
• 导出方（消费者）: agent_in_sandbox_loop.py:337 与 agent_in_localhost_loop.py:258 的 agent loop → export_training_trace.remote(...)。
• 释放方: rl_trainer._release_trace_sessions（:993, 经 store.release_sessions.remote, :1004）/ _release_all_trace_sessions；replay_buffer（:502）与 agent_loop_manager/produce_utils.py（:181）→ release_and_discard_rollout_groups → release_existing_sessions → store.release_sessions.remote → RolloutTraceStore.release(session_id) → Trie.release(None)。
• export_training_trace 内部（trace_store.py:380, 382）直接调用 trie.search(...) / trie.keys()，不经 self.search/self.keys wrapper。
• Trie.release(key)（部分释放分支，trace_store.py:236-260）与 RolloutTraceStore.release(session_id, key) 的带 key 形式在当前代码库中无调用方（docstring 中"used by the SessionServer to drop an orphaned placeholder"的描述无对应调用点，属未修改的既有文档）。

另注意：这是 trace 前缀存储（增量 tokenize 缓存）路径，与推理侧 KV cache affinity 分属不同路径，本 PR 与缓存复用命中率无关，不涉及实测缓存复用。

验证证据（作者提供，只读核实）:

• 单测：8 passed（tests/rl/test_trace_store.py，含新增 TestDeepTrie）。
• pre-commit changed-file hooks：除沙箱内无法运行的 pydantic 外全部通过；mypy 417 文件无问题。
• pydantic 沙箱外检查 exit 0；日志中 torch_npu/vllm/mindspeed 缺失为 checker 自行跳过可选 NPU/vLLM 模块，属预期。
• 未运行 GPU 训练；最终 CI 在远端异步运行，本报告不假称通过。

Main Flowchart before this PR

graph TD
    subgraph SessionServer["SessionServer (xtuner/v1/rl/rollout/session_server.py)"]
        SS_search["store.search.remote(session_id, prompt_text, filter_none=True)"]
        SS_insert["store.insert.remote(session_id, prefix, TokenizedSegment)"]
    end

    subgraph AgentLoops["Agent Loops (export 消费方)"]
        AL_sb["agent_in_sandbox_loop:337 export_training_trace.remote"]
        AL_lh["agent_in_localhost_loop:258 export_training_trace.remote"]
    end

    subgraph RolloutTraceStoreActor["RolloutTraceStore (ray actor)"]
        RTS_insert["insert(session_id, key, value) → Trie.insert"]
        RTS_search["search(session_id, text) → Trie.search"]
        RTS_keys["keys(session_id) → Trie.keys"]
        RTS_release["release(session_id, key=None): pop session → Trie.release"]
        RTS_release_sessions["release_sessions(session_ids) → release(session_id)"]
        RTS_export["export_training_trace(session_id, prompt_text)"]
    end

    subgraph TrieKeys["Trie.keys() interface (before)"]
        TK_collect["递归 _collect(node, path: List[str])<br/>path + [token]; ''.join(path) 输出<br/>深度 == 链深 → RecursionError"]
    end

    subgraph TrieRelease["Trie.release(key=None) interface (before)"]
        TR_sub["递归 _free_subtree(node)<br/>先递归子节点 → _free_ray_refs(value) → children.clear()<br/>深度 == 链深 → RecursionError, 显式 free 中断"]
    end

    subgraph TrieReleasePartial["Trie.release(key) 部分释放 (before, 未改)"]
        TP_locate["前缀定位 best_prefix 匹配"]
        TP_prune["自底向上剪枝空值节点"]
    end

    subgraph CleanupPaths["清理调用方"]
        RT["rl_trainer._release_trace_sessions / _release_all_trace_sessions"]
        RB["replay_buffer → release_and_discard_rollout_groups"]
        PU["produce_utils → release_and_discard_rollout_groups"]
    end

    subgraph TrieSearch["Trie.search() interface (未改)"]
        T_search["迭代最长前缀查找"]
    end

    RTS_search --> T_search
    TP_locate --> TP_prune
    SS_search --> RTS_search
    SS_insert --> RTS_insert
    AL_sb --> RTS_export
    AL_lh --> RTS_export
    RTS_export -->|"trie.search(prompt_text, filter_none=True)"| T_search
    RTS_export -->|"trie.keys() 仅 mismatch 诊断"| TK_collect
    RT --> RTS_release_sessions
    RB --> RTS_release_sessions
    PU --> RTS_release_sessions
    RTS_release_sessions --> RTS_release
    RTS_release -->|"trie.release(None)"| TR_sub
    RTS_release -->|"trie.release(key), 无现存调用方"| TP_locate
    RTS_keys --> TK_collect

Loading

Main Flowchart after this PR

修改点（高亮）：Trie.keys 与 Trie.release(key=None) 内部由递归改为显式栈迭代；接口、调用边、其余方法（insert/search/release(key)/export 组装逻辑）完全未动。

graph TD
    subgraph SessionServer["SessionServer (写入方, 未改)"]
        SS_search["store.search.remote(...)"]
        SS_insert["store.insert.remote(...)"]
    end

    subgraph AgentLoops["Agent Loops (export 消费方, 未改)"]
        AL_sb["agent_in_sandbox_loop export_training_trace.remote"]
        AL_lh["agent_in_localhost_loop export_training_trace.remote"]
    end

    subgraph RolloutTraceStoreActor["RolloutTraceStore (ray actor, 未改)"]
        RTS_insert["insert → Trie.insert"]
        RTS_search["search → Trie.search"]
        RTS_keys["keys → Trie.keys"]
        RTS_release["release(session_id, key=None): pop session → Trie.release"]
        RTS_release_sessions["release_sessions → release"]
        RTS_export["export_training_trace"]
    end

    subgraph TrieKeys["Trie.keys() interface (after)"]
        TK_stack["显式栈迭代 DFS:<br/>stack: List[(TreeNode, str 前缀)]"]:::changed
        TK_order["reversed(children.items()) 压栈<br/>弹出顺序 == 递归插入序前序"]:::changed
    end

    subgraph TrieRelease["Trie.release(key=None) interface (after)"]
        TR_stack["两相栈 (node, visited: bool) 后序迭代"]:::changed
        TR_order["visited=True 时: _free_ray_refs(value) → value=None → children.clear()<br/>顺序与递归版一致"]:::changed
    end

    subgraph TrieReleasePartial["Trie.release(key) 部分释放 (after, 未改)"]
        TP_locate["前缀定位 best_prefix 匹配"]
        TP_prune["release('ab') 仅清 ab 值并保留 abc 子节点;<br/>后续 release('abc') 再自底向上剪枝无值空节点"]
    end

    subgraph CleanupPaths["清理调用方 (未改)"]
        RT["rl_trainer 释放路径"]
        RB["replay_buffer / produce_utils 丢弃路径"]
    end

    subgraph TrieSearch["Trie.search() interface (未改)"]
        T_search["迭代最长前缀查找"]
    end

    RTS_search --> T_search
    TP_locate --> TP_prune
    SS_search --> RTS_search
    SS_insert --> RTS_insert
    AL_sb --> RTS_export
    AL_lh --> RTS_export
    RTS_export --> T_search
    RTS_export --> TK_stack
    RT --> RTS_release_sessions
    RB --> RTS_release_sessions
    RTS_release_sessions --> RTS_release
    RTS_release -->|"trie.release(None)"| TR_stack
    RTS_release --> TP_locate
    RTS_keys --> TK_stack
    TK_stack --> TK_order
    TR_stack --> TR_order

    classDef changed fill:#ffe082,stroke:#e65100,stroke-width:2px;

Loading

文字说明（流程 ↔ 修改对应）:

流程	修改位置	说明
导出 trace 失败诊断（export_training_trace prompt 不匹配时枚举 session keys）	Trie.keys	递归 _collect → 迭代栈；深链 session 下不再 RecursionError，错误报告能正常生成
session 释放（release_sessions → Trie.release(None)）	Trie.release	递归 _free_subtree → 两相栈后序；深链下 ObjectRef（含 routed experts）的显式 free 能执行完毕
release(key) 部分释放	未改	前缀定位 + 剪枝逻辑保持原样；release("ab") 只清除 ab 的值并保留其子节点 abc，release("abc") 之后才向上剪枝无值空节点（新测试断言锁定）

核心原理

1. 缺陷根因: 压缩 trie 的链深由"连续插入互为前缀的 key 的轮数"决定，而非 key 文本长度。"x", "xx", … 每轮插入沿已有的匹配前缀边下探，剩余部分新建单子节点，链深 = 轮数。旧递归实现的调用栈深度等于链深，超过 sys.getrecursionlimit()（默认约 1000）即 RecursionError。
2. 迭代化: 显式栈消掉调用栈深度限制。
   • keys(): (node, path_str) 栈做前序 DFS；reversed(list(children.items())) 压栈使弹出顺序与递归的插入序一致。
   • release(None): (node, visited) 两相栈做后序 DFS，保证"先子树、再 _free_ray_refs(node.value)（含 ray 调用与日志副作用）、最后 children.clear()"的顺序与递归完全一致（新测试断言精确释放顺序）。
3. 复杂度: 时间复杂度与递归版相同——每节点常数次压栈/弹栈，另一计价摊到输出本身的字符串拼接（keys 的路径前缀拼接）与 _free_ray_refs。辅助空间：迭代栈最坏 O(节点数)（可能同时持有全部兄弟子树），keys 另需 O(Σ len(key)) 的输出空间（这是前缀字符串列表本身所需，非新增开销）；release 栈最坏 O(nodes)。本 PR 只主张去掉 Python 递归深度限制，不主张任何路径的线性辅助空间保证。

核心实现与对应的单测

• 实现: xtuner/v1/rl/rollout/trace_store.py:140-150（Trie.keys）、trace_store.py:222-234（Trie.release 的 key=None 分支）。release(key) 部分释放分支（trace_store.py:236-260）未改。
• 单测: tests/rl/test_trace_store.py::TestDeepTrie
  • test_long_conversation_keys_and_release: 插入 1600 个前缀递增 key（链深 1600，必然超递归上限）。断言 keys() 输出顺序与构造序列一致；release() 的 _free_ray_refs 调用次数与精确后序顺序（count..1）；释放后结构清空；重复 release 幂等。旧实现必崩、新实现通过的直接回归测试。
  • test_partial_release_preserves_branch_values_and_key_order: 覆盖空字符串 key（value 落 root）、部分释放的剪枝语义（release("ab") 后 keys()==["","z","a","abc","ax"]；release("abc") 后 "b" 节点被向上剪枝）、混合后再全量 release 的精确顺序（["z","ax","a",""]，先前已释放的节点不再调用）。patch.object(trace_store_module, "_free_ray_refs") 的打补丁方式与文件内既有测试一致。
  • 原理：迭代栈通过 reversed(children) 精确复现递归的访问顺序；两相栈通过"先入 visited=False、子节点全部弹出后再处理 visited=True"复现后序，使带副作用的 _free_ray_refs 调用顺序不变。

抽象与信息隐藏评估

（按 improve-codebase-architecture 词汇，范围限定本 Trie 修复，不做无关候选重构）

• Module/Interface: Trie 的公开接口（keys/insert/search/release）签名与契约完全未动，seam 位置不变；改动全部收在 implementation 内部。调用方（RolloutTraceStore、SessionServer、agent loops、trainer/replay buffer）零接触、零感知。
• Depth: 没有变浅——接口仍小（4 个方法），内部复杂度（迭代两相栈）是递归的等价机械变换，未向调用方泄漏任何细节。
• Locality: 修复集中在 Trie 内一个位置，符合"fix once, fixed everywhere"（所有释放路径共享这一个实现）。
• Deletion test: 若删掉这个改动，RecursionError 复杂性会回到所有深链 session 的清理与诊断路径上——说明改动是必要的，不是 pass-through。
• 新增类: 0（仅方法体内的栈帧元组，属纯数据结构）。

新增功能守则

非新增功能（行为等价修复），R1–R3 trivially 满足：既有类（RolloutTraceStore、SessionServer 等）接触点为 0，无任何新字段、新条件分支或实现细节进入既有类。

公开接口的线性业务流程评估

Trie.keys/Trie.release 均为单一层级的线性流程（建栈 → while 弹栈 → 收尾），无嵌套抽象跳跃。不需要进一步抽取私有方法——任何额外 wrapper 都是纯转发，按 deletion test 应内联，当前形态已是合适粒度。

单测建议

现有 2 个用例已覆盖：超递归上限深度、释放顺序精确断言、空 key（root value）、部分释放剪枝、重复 release 幂等。无强制补充项。

其他 issues

本 PR diff 未发现待修问题。

特别澄清两点（防止误读）：

• keys() 的 O(Σ len(key)) 字符串空间是输出前缀列表本身所需，非本次改动引入的回归；迭代栈最坏 O(nodes)，与调用栈 O(链深) 相比是去掉了 Python 递归深度上限，而非空间上的渐进改善。
• 本 PR 未改变任何字段语义、释放协议或 ProduceBatchResult/RolloutState 可见行为；仅把"深链下会崩的清理/诊断"变成"能完成的清理/诊断"。

Verdict

APPROVE — 本 PR 改动：Critical: 0 / Warning: 0 / Nit: 0。

代码正确：遍历顺序、释放顺序、剪枝语义均与旧实现等价并有测试锁定；栈帧字符串前缀优化正确且省一次 join；作者验证证据（8 passed、changed-file pre-commit、mypy/pydantic 沙箱外通过）与源码审阅结论一致。未运行 GPU 训练，远端 CI 结果以实际为准。

[xs1997zju]

Delivery scope and retained independent review

The exact code head remains 8c7b266dfc996c9700b3826524da24e60fb87328; this update changes only the title/description. The full real Kimi XT Review, with diagrams and verified session IDs, remains in this PR comment.

The author reports GLM-5.3-Flash compatibility in the integrated work. This is a model-independent Python trie traversal fix; the statement does not claim a new full-size GLM training run on this exact standalone branch or a KV-cache/performance benchmark. The retained exact-branch evidence remains the deep-chain failure reproduction, CPU tests and scope recorded in the original review.

The new Agent-lifecycle core #2297 has an overlapping iterative traversal implementation and different request-aware shared-reference cleanup. Review the overlap before landing both; do not overwrite the entire trace-store file or force-free shared R3 references. Original #2257 is unchanged.

Developer guide and related PR map.