feat: graceful drop of recorder

[PeaBrane]

Overview:

as titled

Summary by CodeRabbit

• New Features

  • Introduced distributed prefill token counting and synchronization across multiple workers.
  • Added structured event types for tracking prefill progress, active sequences, and active blocks.
  • Enabled real-time publishing of worker load metrics over NATS for improved observability.

• Improvements

  • Refactored multi-worker sequence management to use asynchronous, event-driven design for better scalability and reliability.
  • Enhanced worker selection logic to use more precise metrics for scheduling.
  • Improved error handling and concurrency in multi-worker environments.

• Bug Fixes

  • Fixed a typo in the metrics publishing handler name.

• Chores

  • Updated dependencies to include the dashmap crate.
  • Removed license headers from select files.

• Tests

  • Added integration tests for distributed prefill synchronization and metrics publishing behavior.

[coderabbitai]

Warning

Rate limit exceeded

@github-actions[bot] has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 25 minutes and 32 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📥 Commits

Reviewing files that changed from the base of the PR and between ef2b0e6 and cc49db1.

⛔ Files ignored due to path filters (2)

• Cargo.lock is excluded by !**/*.lock
• lib/bindings/python/Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (11)

• components/metrics/src/lib.rs (1 hunks)
• lib/llm/Cargo.toml (1 hunks)
• lib/llm/src/kv_router.rs (4 hunks)
• lib/llm/src/kv_router/metrics_aggregator.rs (1 hunks)
• lib/llm/src/kv_router/prefill_counter.rs (1 hunks)
• lib/llm/src/kv_router/protocols.rs (2 hunks)
• lib/llm/src/kv_router/publisher.rs (4 hunks)
• lib/llm/src/kv_router/scheduler.rs (6 hunks)
• lib/llm/src/kv_router/scoring.rs (1 hunks)
• lib/llm/src/kv_router/sequence.rs (8 hunks)
• lib/llm/src/recorder.rs (1 hunks)

Walkthrough

This set of changes refactors and extends the distributed key-value router subsystem. It introduces new event-driven, asynchronous, and multi-worker mechanisms for tracking prefill tokens and active sequences, adds new event types and metrics publishing functionality, and updates worker scheduling logic and endpoint definitions. Several modules are migrated to async patterns with distributed event synchronization.

Changes

Cohort / File(s)	Change Summary
Endpoint Import Refactor components/metrics/src/lib.rs, lib/llm/src/kv_router/metrics_aggregator.rs	Updated imports to use Endpoint from scoring module instead of scheduler, reflecting new struct location and definition.
Cargo Dependency Update lib/llm/Cargo.toml	Added dashmap dependency for concurrent map support.
KV Router Core Refactor lib/llm/src/kv_router.rs, lib/llm/src/kv_router/scheduler.rs, lib/llm/src/kv_router/sequence.rs, lib/llm/src/kv_router/prefill_counter.rs, lib/llm/src/kv_router/protocols.rs, lib/llm/src/kv_router/publisher.rs, lib/llm/src/kv_router/scoring.rs	Migrated multi-worker state management to async/event-driven model using Tokio and DashMap. Introduced new modules for prefill counting and event protocols. Refactored scheduler to use async APIs and updated worker selection logic. Added metrics and event publishing functionality. Defined new event types and endpoint structs.
Recorder Drop Implementation lib/llm/src/recorder.rs	Added Drop implementation to ensure cancellation on drop.
Public API Additions and Changes lib/llm/src/kv_router/protocols.rs, lib/llm/src/kv_router/prefill_counter.rs, lib/llm/src/kv_router/scoring.rs, lib/llm/src/kv_router/scheduler.rs, lib/llm/src/kv_router/sequence.rs	Introduced new public structs, enums, and constants for event-driven coordination, metrics, and multi-worker management. Updated several method signatures to async and expanded struct fields for distributed operation.
Testing Enhancements lib/llm/src/kv_router/publisher.rs, lib/llm/src/kv_router/prefill_counter.rs, lib/llm/src/kv_router/sequence.rs	Added or updated integration tests to validate distributed synchronization and metrics/event publishing across multiple runtimes and workers.

Sequence Diagram(s)

Distributed Prefill Token Synchronization

sequenceDiagram
    participant WorkerA as PrefillCountersMultiWorker (A)
    participant WorkerB as PrefillCountersMultiWorker (B)
    participant NATS as Distributed Event Bus

    WorkerA->>NATS: Publish PrefillEvent (add tokens)
    NATS->>WorkerB: Deliver PrefillEvent
    WorkerB->>WorkerB: Update local counter
    WorkerB->>NATS: Publish PrefillEvent (remove tokens)
    NATS->>WorkerA: Deliver PrefillEvent
    WorkerA->>WorkerA: Update local counter

Loading

Worker Metrics Publishing

sequenceDiagram
    participant Worker as WorkerMetricsPublisher
    participant NATS as KV_METRICS_SUBJECT
    participant Other as Metrics Aggregator

    Worker->>Worker: Detect metrics change
    Worker->>Worker: Wait for stability (1ms)
    Worker->>NATS: Publish LoadEvent
    NATS->>Other: Deliver LoadEvent

Loading

Scheduling Request Flow

sequenceDiagram
    participant Client as Requester
    participant Scheduler as KvScheduler
    participant Slots as ActiveSequencesMultiWorker
    participant Selector as WorkerSelector

    Client->>Scheduler: schedule(request)
    Scheduler->>Slots: Compute blocks/tokens (async)
    Scheduler->>Selector: Select worker
    Scheduler->>Client: Respond with worker assignment
    Scheduler->>Slots: Add request to slots

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~75 minutes

Possibly related PRs

• feat: reduce / revert routing overheads, do not consider output tokens #2182: Introduces the Endpoint struct in scoring.rs, which this PR now imports and uses for endpoint handling and metrics.
• feat: vllm speculative decoding metrics #1549: Refactors metric structures and ForwardPassMetrics usage, aligning with this PR’s updates to metrics and endpoint data handling.
• feat: predictive active blocks for routing without load metrics #1731: Modifies metrics handling and the Endpoint struct in the scoring module, directly related to the event-driven and metrics publishing changes in this PR.

Poem

A rabbit hops from slot to slot,
Prefill tokens counted, never forgot.
Async events in a distributed warren,
Each worker’s metrics, no detail forsworn.
Schedulers scurry, endpoints align,
In this refactor, all systems shine!
🐇✨

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.