[Bugfix]: Fix health checker to preserve worker load during checks

[wuhang2014]

Purpose

Fixes #197: the runtime health checker unconditionally reset every worker's load counter to zero every 10 health-check cycles, even while requests were in flight. The cache-aware prefix tree was never reset, so the policy saw all workers as balanced and kept routing by prefix affinity to already-saturated workers, creating a positive feedback loop of hot spots.

Changes:

• Remove the periodic load reset from WorkerRegistry::start_health_checker(); health checks now run in parallel and never mutate load accounting. Load counters are neither reset periodically nor reset on health transitions: a failed /health probe does not cancel in-flight requests, so the load a worker reports at recovery time can still be real work.
• Delete the unused standalone start_health_checker() with its conditional low-load reset (zero callers).
• Redesign WorkerLoadGuard as an owning RAII guard with share()/release() and adopt it in the OpenAI router proxy path. This also fixes a pre-existing double-decrement on retryable responses and guarantees paired increment/decrement on every path (success, failure, retry, stream abort, client disconnect).
• Adopt WorkerLoadGuard for both phases of the P/D two-stage flow (process_vllm_two_stage_request), fixing a permanent load leak when a P/D request task is cancelled (e.g. client disconnect) between the manual increment and a decrement branch. Cancellation now releases the prefill/decode load via guard drop; covered by cancellation tests for both phases.
• Clamp decrement at zero with a warning and wire the vllm_router_worker_load gauge on increment/decrement/reset.

Changes since first review

• Removed the unhealthy-to-healthy reset_load() call that the first review round flagged: it could erase real in-flight load (a failed /health probe does not cancel requests) and raced with new routing. The health checker now never mutates load accounting; reset_load is documented as an administrative operation that is only safe when the caller has established that no requests are in flight.
• Strengthened the [Bug]: Periodic health checks reset active worker load counters and cause cache-aware hot spots #197 regression test: it now runs under paused Tokio time with a no-I/O health check and advances through 12 health-check cycles, so it fails against the pre-fix 10-cycle reset (verified by mutation probe: left: 0, right: 5).
• Added P/D cancellation tests for both phases (see Test Plan); verified they fail against the previous manual accounting (left: 1, right: 0) and pass with the guard-based fix.

Test Plan

cargo fmt --check
cargo clippy --all-targets --all-features -- -D warnings
cargo test --lib --bins
cargo test --test '*'
pre-commit run --files src/core/mod.rs src/core/worker.rs src/core/worker_registry.rs \
  src/routers/http/router.rs src/routers/http/vllm_pd_router.rs Cargo.toml

New tests:

• src/core/worker.rs: WorkerLoadGuard unit tests (drop decrements, release idempotence, shared handle decrements once, multi-worker, decrement-at-zero clamps).
• src/core/worker_registry.rs: test_health_checker_preserves_inflight_load (regression test for [Bug]: Periodic health checks reset active worker load counters and cause cache-aware hot spots #197; runs under paused Tokio time with a no-I/O health check and advances through 12 cycles, so it fails against the pre-fix 10-cycle reset) and test_health_checker_preserves_load_across_recovery.
• src/routers/http/vllm_pd_router.rs: test_pd_request_cancellation_releases_prefill_load and test_pd_request_cancellation_releases_decode_load — run the real two-stage flow against mock prefill/decode servers, abort the request while the upstream response is pending, and assert load returns to zero (verified to fail against the previous manual accounting).

Test Result

All Rust unit/integration tests pass: cargo test --lib --bins (494 passed), cargo test --test '*' (all suites passed), cargo fmt --check and CI-grade clippy clean, all pre-commit hooks pass.

End-to-end verification with 3 mock workers (requests held in flight, 1s health-check interval) comparing router-reported load vs actual in-flight requests:

Event	Before	After
10 requests in flight, 10 health-check cycles elapsed	Resetting worker loads (cycle 10); router view max_load=0, min_load=0	No reset lines; max_load=4, min_load=3 (≈ actual in-flight)
Probes while 13-15 requests in flight	Router view max_load=0..1, is_imbalanced=false -> prefix-affinity routing continues	max_load=4..5 tracking actual in-flight; balancing decisions see true load

Mutation probes: restoring the old 10-cycle reset makes the regression test fail (left: 0, right: 5); restoring the manual P/D accounting makes both cancellation tests fail (left: 1, right: 0).

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results

[Copilot]

Pull request overview

Fixes a load-accounting bug where periodic health checks reset worker load counters while requests were in flight, which could mislead cache-aware routing and create hot-spot feedback loops. The PR decouples health checking from load accounting and hardens load tracking via an owning RAII guard.

Changes:

• Removes periodic load resets from the registry health checker; runs health checks in parallel and only resets load on unhealthy → healthy recovery.
• Redesigns WorkerLoadGuard into an owning, shareable RAII guard to ensure exactly-once decrement across success/failure/retry/stream paths, and clamps load decrements at zero with metrics updates.
• Updates the HTTP router streaming/non-streaming request paths to use WorkerLoadGuard instead of manual increment/decrement logic.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/routers/http/router.rs	Switches cache-aware load accounting to WorkerLoadGuard, including streaming lifecycle handling.
src/core/worker.rs	Implements new owning WorkerLoadGuard, clamps decrements at zero, and updates load metrics wiring.
src/core/worker_registry.rs	Removes periodic load resets; parallelizes health checks; resets load only on recovery; adds regression tests.
src/core/mod.rs	Removes re-export of deleted standalone start_health_checker and updates public exports accordingly.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7314e6214b

ℹ️ 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".

[hsliuustc0106]

@ywang96 @zhuohan123 PTAL