Skip to content

refactor(server): split the 8,130-line BatchScheduler into a scheduler/ directory module #1243

Description

@inureyes

src/server/batch/scheduler.rs is 8,130 lines with a single impl BatchScheduler block opening at line 753 and running to line 8,102, containing over 110 methods. That is 4x the top threshold in the project's own file-size guidelines, with no recorded justification.

Evidence

  • src/server/batch/scheduler.rs is 8,130 lines (wc -l). One impl BatchScheduler { at src/server/batch/scheduler.rs:753 closes at line 8,102.
  • docs/code-guidelines.md:11-15 sets the thresholds: 1,500+ lines "Strongly consider splitting into a directory module", and 2,000+ lines with no helpers file is called an anti-pattern. The guidelines document exceptions explicitly, for example src/models/nemotron_h.rs at docs/code-guidelines.md:19 ("a documented exception... splitting it would break the layer-interleaving logic"). scheduler.rs has no such recorded justification.
  • The tests were already extracted into eight sibling files (scheduler_tests.rs, scheduler_cohort_parity_tests.rs, scheduler_prompt_cache_tests.rs, scheduler_seed_determinism_tests.rs, scheduler_muse_glimmer_tests.rs, scheduler_muse_glimmer_parallel_tests.rs, and others under src/server/batch/), so the 8,130 lines are pure implementation.
  • Sibling concerns already live in their own files in the same directory (queue.rs, sequence.rs, tick_policy.rs, mtp_policy.rs, prefill_cohort.rs, observability.rs), demonstrating that the seam pattern works here.
  • This is the single hottest file for merge conflicts: nearly every speculative-decoding PR touches it, and the one-impl-block shape hides the scheduler state machine's boundaries from reviewers.

Suggested fix

Convert scheduler.rs into a scheduler/ directory module: a mod.rs holding the BatchScheduler struct definition and shared state, plus submodules along the state machine's natural seams, each carrying part of the impl:

  • admission (queue intake, admission control)
  • prefill (prefill scheduling, cohort formation)
  • decode_tick (per-tick decode loop)
  • speculative_finalize (draft verification and acceptance finalization)
  • paged_layout (paged KV layout and slot management)

Zero behavior change. The existing eight scheduler test files are the gate: they must pass unchanged, byte for byte. Rust allows splitting one type's impl across multiple impl BatchScheduler blocks in different files of the same module, so the split is mechanical: move method groups, adjust visibility from private to pub(super)/pub(crate) only where the move requires it, and keep the public API of the module identical.

Estimated scope: several days, mechanical but wide. Worth landing during a quiet window between speculative-decoding PRs to avoid rebasing pain.

Acceptance criteria

  • No file under src/server/batch/scheduler/ exceeds the docs/code-guidelines.md:11-15 thresholds without a justification recorded in code-guidelines.md
  • All existing scheduler tests (scheduler_tests.rs, scheduler_cohort_parity_tests.rs, scheduler_prompt_cache_tests.rs, scheduler_seed_determinism_tests.rs, the muse-glimmer suites, and the remaining sibling test files) pass unchanged, with no test file edits
  • cargo build and cargo clippy clean; no public API change outside src/server/batch/
  • Zero behavior change: no logic edits, only method moves and the minimal visibility adjustments the moves require

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:architectureArchitecture and code structure changesarea:inferenceGeneration, sampling, decoding (incl. speculative, DRY)priority:mediumMedium prioritystatus:readyReady to be worked ontype:refactorCode restructuring without changing functionality

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions