Skip to content

refactor(qwen35): split scheduler step drivers into scheduler/{steps,emit,telemetry}.rs - #1037

Open
CAICAIIs wants to merge 5 commits into
pegainfer-project:mainfrom
CAICAIIs:feat/qwen35-scheduler-steps-split
Open

refactor(qwen35): split scheduler step drivers into scheduler/{steps,emit,telemetry}.rs#1037
CAICAIIs wants to merge 5 commits into
pegainfer-project:mainfrom
CAICAIIs:feat/qwen35-scheduler-steps-split

Conversation

@CAICAIIs

@CAICAIIs CAICAIIs commented Sep 7, 2026

Copy link
Copy Markdown
Collaborator

What

Continue the scheduler.rs slim-down along the seams backend.rs/tp.rs (#967, #968) established: step drivers, request-emission concerns, and ITL telemetry move out of scheduler/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) — the PEGAINFER_ITL_DEBUG per-step ITL trio.

Moved items gain pub(super) only where the loop, backend, tp, or emit siblings 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 to steps.rs.

Defense table (all inherited — pure move)

Defense Failure mode Successor
Terminal shutdown fans every owned request an error requests stranded on scheduler exit inherited: emit.rs terminal_scheduler_shutdown
Async-prefill stream drained before KV/recurrent release use-after-free of GPU state on shutdown inherited: emit.rs drop-order comment + steps.rs InflightPrefill
Closed-request pruning before admission zombie requests occupy slots/pages inherited: emit.rs prune_closed_requests
echo=true rejected (fail-closed, not degraded) unsupported serving contract silently accepted inherited: emit.rs reject_unsupported_echo
Start-gate cancel on partial dispatch one rank mutated while a peer skipped the step inherited: steps.rs gate checks + decode_step family
ITL telemetry env-gated (off = zero cost) bench path pays logging overhead inherited: telemetry.rs itl_debug_enabled OnceLock

Validation

  • cargo fmt --check clean; cargo clippy -p pegainfer-qwen35 --features qwen35 --all-targets -- -D warnings clean.
  • cargo test --lib green; tp2_* GPU tests pass on 2×A100 (-- --ignored --test-threads=1); e2e_scheduler green.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread pegainfer-qwen35/src/weights/layers.rs Outdated

impl LinearAttentionLayer {
/// Phase 2b: shard linear attention over TP ranks. The value-head unit
/// Linear-attention sharding over TP ranks. The value-head unit

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

@CAICAIIs
CAICAIIs force-pushed the feat/qwen35-scheduler-steps-split branch 2 times, most recently from 9258a1c to 63b3f57 Compare September 8, 2026 06:52
@CAICAIIs
CAICAIIs force-pushed the feat/qwen35-scheduler-steps-split branch 2 times, most recently from ab6bb1b to 88eacdb Compare September 8, 2026 09:23
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>
@CAICAIIs
CAICAIIs force-pushed the feat/qwen35-scheduler-steps-split branch from 88eacdb to 7144a77 Compare September 10, 2026 03:44
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>
@CAICAIIs
CAICAIIs force-pushed the feat/qwen35-scheduler-steps-split branch from 7144a77 to 1e82b13 Compare September 10, 2026 05:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant