Prhub

#39322 [Router] Extract worker selection into policies::selection (no behavior change)

原始 PR 作者 Kangyan-Zhou 合并时间 2026-09-15 15:49 文件变更 3 提交数 2 评论 1 代码增减 +885 / -461

执行摘要

Router 选路阶梯抽取为 policies::selection 模块

PR 描述的动机有三层:其一,两条选路阶梯此前是约 900 行 chat_completions handler 里唯一的非测试调用者,"The assertions worth writing are about a ladder as a whole"(例如"饱和集群仍能路由"、"可正常容纳的 domain 不该输给只能靠放宽容量才可入的 domain"),写在 handler 里就只能退化为端到端 HTTP 测试甚至无法断言;其二,Cell<PolicySelectionFailureReason> 只是把值从闭包中偷带出来的权宜手段,"The Cell was a return-type problem in disguise";其三,后面还排着三个新层级(#39168 队列门、#39169 饱和钉、#39170 候选采样),若继续写在原处,会变成 1100 行 handler 加第二个 Cell。作者还指出现有拆分的实际代价:正是把两条阶梯分置在不同代码位置,才让 queue gate 在一条路径上带两级兜底、另一条只有一级而无人察觉。

值得精读,但建议把注意力放在第二个 commit 而非体量最大的搬迁 commit:先看 Selector 结构体如何取代 Cell、select_prefill_worker 用 ok_or(failure_reason) 把 Option 收口成 Result,再看 5 个阶梯级测试如何用 mutation check 锁住容量兜底语义。值得借鉴的设计决策包括:为不再需要闭包的场景把"偷带值"的 Cell 换成显式状态结构体;接受 prefill/decode 签名不对称(Result vs Option)而不是为对称造错误类型;把两条阶梯放在同一模块以获得"一屏可见的不对称"。评审时建议额外确认第二个 commit 中 domain 惰性物化是否改变了任何副作用时机。

讨论亮点

本 PR 只有 2 个 APPROVED 和 1 条 issue 评论,没有行内 review 评论,讨论集中在"纯搬迁是否可信"与"是否影响既有 policy"。ShangmingCai 批准时给出两点:"Logic looks good. Need to fix conflict before merge",并 CC Bo-Vincent 求证 "will this change break your policy?";Bo-Vincent 随后回应 "Took a look at the code and ran a quick test. It preserves the existing policy order and selection logic. LGTM!",实质完成了行为不变的第三方验证。更深的设计交锋写在 PR 描述里:Cell 只是"返回类型问题的伪装";两条阶梯必须同模块,因为 CandidateDomain 已带 stage: RoutingStage,而此前拆开直接导致 queue gate 在一个路径上两级兜底、另一个只有一级;decode 返回 Option 而非 Result 是为拒绝"为对称而造的噪音";把两者统一成 RoutingStage 参数化入口属于 redesign,被明确排除在无行为变更 PR 之外。

实现拆解

  1. 新增模块骨架:experimental/sgl-router/src/policies/selection.rs(+845)定义 PrefillSelectionInputs<'a> 与 DecodeSelectionInputs<'a> 两个 pub(crate) 结构体,字段按"协作者在前、单请求事实在后"排列,显式持有 policy / bucket_selector / metrics,让 policies/ 不再依赖 server::app_context;入口为 select_prefill_worker(...) -> Result<Arc<Worker>, PolicySelectionFailureReason> 与 select_decode_peer(...) -> Option<Arc<Worker>>。
  2. 用结构体字段替换 Cell:原 handler 用 Cell<PolicySelectionFailureReason> 从闭包中带出失败原因;新模块以 struct Selector { inputs, bucket_request, failure_reason } 承载跨层级状态,run() 返回 Option<Arc<Worker>>,由 selected.ok_or(selector.failure_reason) 收口成 Result。语义约定是"每个放弃的层级覆写 reason,未记录失败的层级返回 None 时保留上一个原因",因此上报的永远是最后一个记录过原因的层级。
  3. 阶梯逻辑原样搬迁:Selector::run 依次执行 cache_winner()(Cache-Aware 的有界全局候选 + 指标 + commit_prefill_selection)、global_affinity_probe(会话亲和全局探测,命中后经 prefill_affinity_domain 在本 Bucket 内以 select_in_domain(..., true, false, false) 重建 backup),以及按 PolicyKind / SessionAffinityMode 分流的两遍 select_domains;base_context() 统一构造 SelectionContext,select_in_domain() 调 propose_prefill 加 resolve_prefill_admitted / resolve_prefill,域耗尽时写 PrefillAdmissionExhausted。prefill_policy_reason 与 projected_decode_kv_tokens 随各自测试一并迁入。
  4. handler 瘦身与模块注册:server/routes/chat.rs 删除 461 行闭包、两个辅助函数和相关导入(含 std::cell::Cell),chat_completions 现在只组装输入结构体再调用新入口,prefill 失败经 .map_err(|reason| policy_selection_failed(&ctx, &model_str, reason))? 映射 ApiError,decode 失败映射 ApiError::NoDecodeWorkersAvailable;policies/mod.rs 仅新增 pub mod selection; 一行,可见性保持 pub(crate) 与 PolicySelectionFailureReason 一致。
  5. 测试与验证配套:搬迁 commit 让 #[test] 数在两文件之间守恒为 33,并通过 token 流对比逐段证明 selector、probe + 最终 or_else 链、decode 决策闭包与两遍阶梯均为纯移动;review commit 追加 5 个阶梯级测试(饱和 prefill / decode 集群经容量兜底仍可路由、采样结果不会落到被 admission 拒绝的 worker、空 fleet 报 ProposalEmpty),并做 mutation check。cargo test 702 pass / 0 fail,cargo fmt --check 与 cargo clippy --all-targets 干净(buckets.rs:167 的既存 lint 已由 #39259 在 main 修复)。
文件 模块 状态 重要度
experimental/sgl-router/src/policies/selection.rs 选路策略 added 9.26
experimental/sgl-router/src/server/routes/chat.rs 请求入口 modified 8.86
experimental/sgl-router/src/policies/mod.rs 策略注册 modified 4.19

关键符号

select_prefill_worker select_decode_peer Selector::run Selector::base_context Selector::cache_winner Selector::select_domains Selector::select_in_domain prefill_policy_reason projected_decode_kv_tokens chat_completions

关键源码片段

experimental/sgl-router/src/server/routes/chat.rs entrypoint

chat_completions 是原阶梯的唯一非测试调用者,本 PR 在此删除 461 行闭包与辅助函数、移除 Cell 导入,改为组装输入结构体并调用新入口,同时保留状态码映射职责。

// server/routes/chat.rs:handler 现在只做编排 —— parse -> tokenize -> select -> proxy。
// 两条阶梯的输入在这里组装,失败原因在这里映射成 HTTP 状态码(模块不碰 HTTP)。
​
// `select_prefill_worker` 会在 Bucket 分区关闭时把该模式归约为 Bucket,
// 因此这里直接传配置值,无需再复制一份归约逻辑。
let session_affinity_mode = ctx
    .config
    .model
    .affinity
    .as_ref()
    .map(|config| config.session_affinity_mode)
    .unwrap_or(SessionAffinityMode::Bucket);
​
let worker = select_prefill_worker(&PrefillSelectionInputs {
    policy: policy.as_ref(),
    policy_kind: ctx.config.model.policy,
    bucket_selector: ctx.bucket_selector.as_ref(),
    metrics: ctx.metrics.as_ref(),
    model_id: &model_id,
    body: Some(&body),
    routing_key,
    session_id,
    request_input_tokens,
    request_tokens: request_tokens.as_ref().map(|tokens| tokens.ids.as_slice()),
    external_prefix: external_prefix.as_ref(),
    load_snapshot: load_snapshot.as_ref(),
    workers: &workers,
    ttft_slo_ms,
    tps_slo,
    session_affinity_mode,
})
// 失败原因到状态码的映射留在路由层:这是模块划界的边界条件。
.map_err(|reason| policy_selection_failed(&ctx, &model_str, reason))?;
​
// Decode 选择发生在 Final P 之后;普通聚合 worker 直接跳过。
//
// PD 模式下解析不到 decode peer(NoDecodeWorkersAvailable)会冒泡成 503,
// 便于运维对 prefill / decode 池失衡告警。
let decode_peer: Option<Arc<Worker>> = if worker.mode() == WorkerMode::Prefill {
    let decode_workers = resolver.decode_candidates(&model_id).map_err(|e| match e {
        PdResolveError::NoHealthyWorkers => ApiError::NoHealthyWorkers {
            model: model_str.clone(),
        },
        PdResolveError::NoDecodeWorkersAvailable => ApiError::NoDecodeWorkersAvailable {
            model: model_str.clone(),
        },
        PdResolveError::NoPrefillWorkersAvailable => ApiError::NoPrefillWorkersAvailable {
            model: model_str.clone(),
        },
    })?;
    Some(
        select_decode_peer(&DecodeSelectionInputs {
            decode_policy_kind: ctx.config.model.decode_policy,
            bucket_selector: ctx.bucket_selector.as_ref(),
            model_id: &model_id,
            prefill_url: &worker.url,
            decode_workers: &decode_workers,
            request_input_tokens,
            requested_max_output_tokens,
            ttft_slo_ms,
            tps_slo,
            load_snapshot: load_snapshot.as_ref(),
        })
        // decode 阶梯只有一种失败形态,所以入口返回 Option 而不是 Result:
        // 为对称而造错误类型只会成为噪音。
        .ok_or_else(|| ApiError::NoDecodeWorkersAvailable {
            model: model_str.clone(),
        })?,
    )
} else {
    None
};

评论区精华

纯搬迁是否真的保持既有 policy 行为 正确性

ShangmingCai 在批准时 CC Bo-Vincent 求证:"will this change break your policy?";Bo-Vincent 表示已通读代码并跑了快速测试。

结论:Bo-Vincent 结论为 "It preserves the existing policy order and selection logic. LGTM!",加上作者在描述中给出的 token 流逐段对比,行为不变得到双重验证。 · 已解决

合并前需解决冲突 other

ShangmingCai 批准时指出 "Need to fix conflict before merge",说明分支基线早于 main 上已合入的改动(如 #39259 对 buckets.rs:167 lint 的修复)。

结论:冲突已解决并完成合并(merged_by: ShangmingCai)。 · 已解决

两条阶梯共处一个模块,而不是一梯一模块 设计

作者论证 `CandidateDomain` 已带 `stage: RoutingStage { Prefill, Decode }`,prefill 与 decode 是同一概念的两套配置;此前把它们分置在不同代码位置,直接导致 queue gate 在一个路径上带两级兜底、另一个只有一级而无人察觉。

结论:决定同模块放置,使 " 一层加在一边而没加在另一边 " 的不对称能在一屏内被发现;把两者统一为 RoutingStage 参数化入口被判定为 redesign,不纳入本 PR。 · 已解决

prefill 返回 Result 而 decode 返回 Option 的签名不对称 设计

decode 阶梯只有一种失败形态(无可用 decode worker),作者认为为其命名错误类型、只为与 prefill 对称,属于噪音。

结论:保持不对称签名:`select_prefill_worker -> Result<Arc<Worker>, PolicySelectionFailureReason>`,`select_decode_peer -> Option<Arc<Worker>>`,路由层各自映射状态码。 · 已解决

用返回类型替代 `Cell` 从闭包偷带失败原因 设计

原 handler 的 `Cell<PolicySelectionFailureReason>` 被定性为 "a return-type problem in disguise":它存在的唯一理由是闭包无法通过返回值传出原因。

结论:改为 `Selector.failure_reason` 字段,`run()` 返回 `Option` 再由 `ok_or` 收口成 `Result`;`std::cell::Cell` 不再被 handler 导入。 · 已解决

把阶梯级断言从口号变成单元测试 测试

模块文档声称 " 饱和集群仍能路由 "、" 采样结果不会落到被拒绝的 worker" 这类断言是抽出模块的理由;作者随后在第二个 commit 补齐 5 个阶梯测试,并逐个做 mutation check(移除任一容量兜底 pass 会让对应饱和测试失败)。

结论:测试落地:`cargo test` 702 pass / 0 fail,其中 697 来自搬迁(与 main 一致)加 5 个新阶梯测试。 · 已解决

风险与影响

1) 纯搬迁声明只覆盖第一个 commit:作者明确第二个 commit(ladder tests、Result 取代 invalid-state 结构体、对称输入结构体、domain 惰性物化)不是纯搬迁,prefill_domains 闭包把 domain 物化推迟到真正读取的层级,是本次唯一改变分配时机的点,审阅时应重点对照。
2) 失败原因语义依赖约定:failure_reason 由各层级覆写,"未记录者保留上一个原因"这一规则在后续新增层级(#39168 / #39169 / #39170)里若被遗忘赋值,会报出过期原因,且只在最终选路失败时可见。
3) 隐式 panic 契约:select_in_domain 中 load_snapshot.expect("shared prefill admission requires a load snapshot") 被保留,PrefillSelectionInputs 的注释说明 needs_load_snapshot 默认与 uses_shared_prefill_admission 同步;新模块成为公开入口后,若新的调用方(含未来的非 HTTP 调用者)违反该前提会直接 panic。
4) 分支基线落后于 main:buckets.rs:167 的 clippy lint 需由 #39259 修复,说明本分支基于较早的 main,ShangmingCai 已要求先解冲突再合并,存在重解冲突时引入偏差的可能。
5) 无用户可见兼容性风险:不涉及公开 API、配置键或线格式变更。

用户侧:路由决策、状态码与 x-sgl-decode-url 等 header 行为均按声明保持不变,无感知变化。系统侧:router 的 chat handler 由约 900 行降至约 440 行的编排层,两条选路阶梯首次拥有独立的模块边界与单元测试能力,后续新增层级不再需要往 HTTP handler 里塞闭包与 Cell;policies/ 与 server::app_context 解耦后,选路逻辑可脱离 HTTP 上下文单测。团队侧:这是 4 步 stack 的第 1 步,直接决定 #39168 / #39169 / #39170 的落点质量;同时把"饱和集群仍能路由"这类横切断言固化为可回归的测试,降低后续选路演进引入不对称行为的概率。

核心选路路径重构 第二个 commit 非纯搬迁 failure_reason 依赖覆写约定 load_snapshot 缺失即 panic 的隐式契约

关联 Issue

#39259 fix(router): simplify bucket context limit check

完整报告

参与讨论