refactor(qwen35): split scheduler step drivers into scheduler/{steps,emit,telemetry}.rs - #1037
refactor(qwen35): split scheduler step drivers into scheduler/{steps,emit,telemetry}.rs#1037CAICAIIs wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9e5c6de0a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| impl LinearAttentionLayer { | ||
| /// Phase 2b: shard linear attention over TP ranks. The value-head unit | ||
| /// Linear-attention sharding over TP ranks. The value-head unit |
There was a problem hiding this comment.
Revert unrelated TP phase wording cleanup
This behavior-preserving scheduler/TP module split also rewrites phase terminology here and in lib.rs, even though those edits do not support either file-movement invariant. Keeping unrelated documentation and user-facing error-message changes in this refactor broadens its review surface and makes the claimed no-behavior-change scope harder to verify; revert them or move them to a separately justified change.
AGENTS.md reference: AGENTS.md:L121-L123
Useful? React with 👍 / 👎.
9258a1c to
63b3f57
Compare
ab6bb1b to
88eacdb
Compare
Move the per-rank worker runtime out of the tp_executor.rs God module:
TpWorker + spawn/drop lifecycle, TpStartupGate, the NCCL startup watchdog,
TpWorkerState/TpWorkerPrepared and its command loop (run/respond/
execute_*/precapture/drop), the slot bookkeeping helpers, the shared
decode-row view/sampling helpers, recurrent-capacity math, and the worker
CublasThreadGuard binding. The entry keeps command/reply types, the
executor orchestration (including the precapture sweep), plan validators,
and response validation, reaching worker items via the same `use super::*`
style as scheduler/{backend,tp}. Rebased over pegainfer-project#720/pegainfer-project#722. Behavior
unchanged; the watchdog disarm test moves into worker.rs with its
implementation.
Signed-off-by: CAICAIIs <3360776475@qq.com>
88eacdb to
7144a77
Compare
The tp_executor entry and the worker module carried two near-identical condvar gates: TpCommandStartGate (Pending/Execute/Cancel, first-write resolve returning whether it won) and TpStartupGate (Pending/Connect/ Cancel, silent set). Both broadcast one go/cancel decision that is resolved exactly once, so they collapse into TpGate + TpGateDecision with the first-write-wins resolve; startup treats Go as connect and asserts on the returned decision instead of a bool. Behavior unchanged. Signed-off-by: CAICAIIs <3360776475@qq.com>
The scheduler and the TP worker carried byte-identical CUDA/cuBLAS thread bindings (guard + set_device + bind_to_thread + cublas_init); the worker split moved one of the two copies, making the duplication visible. Both now call crate::cublas_thread::bind_model_thread with a role tag (the failure messages are unchanged); the scheduler keeps its tune_decode_gemm_algos step at its own call site. Behavior unchanged. Signed-off-by: CAICAIIs <3360776475@qq.com>
…es.rs Move the reply envelope (TpWorkerReply/TpWorkerResponse), bounded response collection, per-command response validation, and the step-timeout/abort helpers out of tp_executor.rs; the three response-contract tests move with them. Behavior unchanged, with two behavior-preserving consolidations in the moved code: validate_exact_rank_responses folds its push loop into an iterator pipeline, and the three rank-0 payload validators (prefill / decode / unified) collapse into one validate_primary_responses helper over payload extractors, keeping the exact prior error messages. Rebased over pegainfer-project#1004/pegainfer-project#1005/pegainfer-project#1033/pegainfer-project#1038 (the precapture sweep keeps using validate_ack_responses, which stays production-visible). Signed-off-by: CAICAIIs <3360776475@qq.com>
…emit,telemetry}.rs Continue the scheduler.rs slim-down along the seams backend.rs/tp.rs already established: step drivers (batch prefill, overlap launch, unified step, pure decode, token dispatch/retirement, chunk vocabulary) move to steps.rs, terminal shutdown / closed-request pruning / rejection fan-out move to emit.rs, and the PEGAINFER_ITL_DEBUG telemetry trio moves to telemetry.rs. scheduler/mod.rs keeps the request/artifact types, the entry points (including the pegainfer-project#1033 policy-aware start), load publishing, and the main loop. Behavior unchanged; moved items gain pub(super) only where the loop, backend, tp, emit siblings, or scheduler tests name them. Signed-off-by: CAICAIIs <3360776475@qq.com>
7144a77 to
1e82b13
Compare
What
Continue the
scheduler.rsslim-down along the seamsbackend.rs/tp.rs(#967, #968) established: step drivers, request-emission concerns, and ITL telemetry move out ofscheduler/mod.rs, which keeps the request/artifact types, the entry point, load publishing, and the main loop (1896 → 882 lines). Behavior unchanged.Stacked PR 3/3 — based on #1035 + the responses-split PR; Files-changed includes both until they merge. Merge those first. Rebuilt over #1004/#1005/#1033/#1038 (the #1033 policy-aware start stays in the scheduler entry).
Structure
scheduler/steps.rs(846 lines) — batch prefill, overlap launch, unified prefill+decode step, pure decode, token dispatch/retirement,DecodeDispatchBackend/PrefillPromoteBackend, and the chunk vocabulary (ScheduledChunk,InflightPrefill,take_prefill_chunks,promote_or_requeue) shared with the backend.scheduler/emit.rs(126 lines) — terminal shutdown fan-out, closed-request pruning, unsupported-echo/rejection paths.scheduler/telemetry.rs(45 lines) — thePEGAINFER_ITL_DEBUGper-step ITL trio.Moved items gain
pub(super)only where the loop,backend,tp, oremitsiblings name them; step-internal helpers (decode_step_with_seed,process_decode_logits,compact_single_slot,fail_chunk,prefill_lifecycle_failure,split_scheduled_backend_state) stay private tosteps.rs.Defense table (all inherited — pure move)
emit.rsterminal_scheduler_shutdownemit.rsdrop-order comment +steps.rsInflightPrefillemit.rsprune_closed_requestsecho=truerejected (fail-closed, not degraded)emit.rsreject_unsupported_echosteps.rsgate checks +decode_stepfamilytelemetry.rsitl_debug_enabledOnceLockValidation
cargo fmt --checkclean;cargo clippy -p pegainfer-qwen35 --features qwen35 --all-targets -- -D warningsclean.cargo test --libgreen;tp2_*GPU tests pass on 2×A100 (-- --ignored --test-threads=1);e2e_schedulergreen.