Skip to content

test(router): isolate request lifecycle counters #14875

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jthomson04 merged 2 commits into from
Sep 15, 2026
Merged

test(router): isolate request lifecycle counters #14875

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ab2494b
test(router): isolate request lifecycle counters
jthomson04 Sep 15, 2026
e165a1b
test(router): narrow lifecycle counter isolation
jthomson04 Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 15 additions & 14 deletions lib/llm/src/kv_router/metrics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -842,7 +842,10 @@ impl RoutingOverheadMetrics {
/// decode pool). Component-scoped metrics let each local router emit metrics with
/// distinct `dynamo_component` labels, so pools can be monitored and scaled
/// independently.
#[cfg_attr(test, derive(Clone))]
pub struct RouterRequestMetrics {
/// Total requests admitted by the router scheduler.
pub requests_started_total: prometheus::IntCounter,
Comment thread
jthomson04 marked this conversation as resolved.
pub requests_total: prometheus::IntCounter,
pub time_to_first_token_seconds: prometheus::Histogram,
pub inter_token_latency_seconds: prometheus::Histogram,
Expand All @@ -857,21 +860,13 @@ pub struct RouterRequestMetrics {
}

static ROUTER_REQUEST_METRICS: OnceLock<Arc<RouterRequestMetrics>> = OnceLock::new();
static ROUTER_REQUESTS_STARTED_TOTAL: OnceLock<prometheus::IntCounter> = OnceLock::new();

impl RouterRequestMetrics {
/// Returns the registered metrics if `from_component()` was called earlier.
pub fn get() -> Option<Arc<Self>> {
ROUTER_REQUEST_METRICS.get().cloned()
}

/// Total requests admitted by the router scheduler.
pub fn requests_started_total(&self) -> &prometheus::IntCounter {
ROUTER_REQUESTS_STARTED_TOTAL
.get()
.expect("router request metrics must be initialized")
}

/// Create from a Component, memoized in a static OnceLock.
/// Uses the MetricsHierarchy API which auto-prepends `dynamo_component_`,
/// injects hierarchy labels, and registers with the DRT `MetricsRegistry`.
Expand All @@ -893,12 +888,6 @@ impl RouterRequestMetrics {
extra_labels,
)
.expect("failed to create router_requests_started_total");
assert!(
ROUTER_REQUESTS_STARTED_TOTAL
.set(requests_started_total)
.is_ok(),
"router_requests_started_total already initialized"
);
let requests_total = metrics
.create_intcounter(
&router_metric(frontend_service::REQUESTS_TOTAL),
Expand Down Expand Up @@ -990,6 +979,7 @@ impl RouterRequestMetrics {
non_max_overlap_selections_total.with_label_values(&[WORKER_TYPE_PREFILL]);
overlap_blocks_lost.with_label_values(&[WORKER_TYPE_PREFILL]);
Arc::new(Self {
requests_started_total,
requests_total,
time_to_first_token_seconds,
inter_token_latency_seconds,
Expand All @@ -1006,6 +996,17 @@ impl RouterRequestMetrics {
.clone()
}

/// Use fresh, unregistered lifecycle counters and retain all other metric handles.
#[cfg(test)]
pub(crate) fn with_isolated_counters_for_test(&self) -> Arc<Self> {
Arc::new(Self {
requests_started_total: prometheus::IntCounter::new("requests_started_total", "test")
.unwrap(),
requests_total: prometheus::IntCounter::new("requests_total", "test").unwrap(),
..self.clone()
})
}

/// Record a selection that sacrificed KV cache overlap.
pub fn observe_non_max_overlap_selection(&self, worker_type: &str, overlap_blocks_lost: f64) {
debug_assert!(overlap_blocks_lost > 0.0);
Expand Down
6 changes: 4 additions & 2 deletions lib/llm/src/kv_router/routing_host/request_guard.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -621,7 +621,7 @@ where
let track_output_blocks =
attempt_id.is_some() && chooser.kv_router_config().router_track_output_blocks;
if attempt_id.is_some() {
request_metrics.requests_started_total().inc();
request_metrics.requests_started_total.inc();
}
let approximate_lru = cleanup.approximate_lru.clone();
let output_hashes = approximate_lru
Expand Down Expand Up @@ -652,7 +652,7 @@ where
lora_load: Option<LoraLoadGuard>,
request: &PreprocessedRequest,
) -> Self {
request_metrics.requests_started_total().inc();
request_metrics.requests_started_total.inc();
Self {
cleanup: match occupancy_reservation {
Some(reservation) => RequestCleanup::Occupancy {
Expand Down Expand Up @@ -999,6 +999,8 @@ mod prefill_start_tests {
.unwrap()
}
Arc::new(RouterRequestMetrics {
requests_started_total: prometheus::IntCounter::new("requests_started_total", "test")
.unwrap(),
requests_total: prometheus::IntCounter::new("requests_total", "test").unwrap(),
time_to_first_token_seconds: hist("ttft_seconds"),
inter_token_latency_seconds: hist("itl_seconds"),
Expand Down
69 changes: 34 additions & 35 deletions lib/llm/src/kv_router/routing_host/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1156,6 +1156,7 @@ async fn router_with_worker_configs(
workers: HashMap<u64, ModelRuntimeConfig>,
) -> (RoutingHost, Runtime) {
let runtime = Runtime::from_current().unwrap();
// Each runtime has its own in-memory discovery store, so fixture namespaces can repeat.
let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
.await
.unwrap();
Expand Down Expand Up @@ -1194,7 +1195,9 @@ async fn router_with_worker_configs(
let inner = PushRouter::from_client(client, RouterMode::KV)
.await
.unwrap();
let router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap();
let mut router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap();
// Parallel tests must not share the process-global request counters.
router.request_metrics = router.request_metrics.with_isolated_counters_for_test();
router
.inner
.client
Expand Down Expand Up @@ -1410,7 +1413,6 @@ async fn track_request(
async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() {
let (router, runtime) = router(None).await;
let request = Context::new(request());
let requests_started_before = router.request_metrics.requests_started_total().get();

let preview = router
.preview_kv_route(&request, RequestPhase::Decode)
Expand All @@ -1422,8 +1424,8 @@ async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() {
.expect("decode plan should admit one request");
assert_eq!(plan.signals().worker.worker_id, 7);
assert_eq!(
router.request_metrics.requests_started_total().get(),
requests_started_before,
router.request_metrics.requests_started_total.get(),
0,
"a topology decision is not a started request"
);
let admitted_loads = router
Expand Down Expand Up @@ -1451,8 +1453,8 @@ async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() {
"abandoned plans must release their scheduler reservation: {released_loads:?}"
);
assert_eq!(
router.request_metrics.requests_started_total().get(),
requests_started_before,
router.request_metrics.requests_started_total.get(),
0,
"an abandoned plan must not count as a started request"
);

Expand All @@ -1465,7 +1467,6 @@ async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() {
async fn route_preview_does_not_admit_a_request() {
let (router, runtime) = router(None).await;
let request = Context::new(request());
let requests_started_before = router.request_metrics.requests_started_total().get();

let preview = router
.preview_kv_route(&request, RequestPhase::Decode)
Expand All @@ -1478,10 +1479,7 @@ async fn route_preview_does_not_admit_a_request() {
.await
.unwrap();
assert!(loads.iter().all(|load| load.active_requests == 0));
assert_eq!(
router.request_metrics.requests_started_total().get(),
requests_started_before
);
assert_eq!(router.request_metrics.requests_started_total.get(), 0);

drop(router);
runtime.shutdown();
Expand Down Expand Up @@ -1560,7 +1558,6 @@ async fn route_preview_does_not_acquire_session_affinity() {
#[serial_test::serial]
async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() {
let (router, runtime) = router(None).await;
let requests_started_before = router.request_metrics.requests_started_total().get();
let request = Context::new(request());
let preview = router
.preview_kv_route(&request, RequestPhase::Decode)
Expand All @@ -1572,10 +1569,7 @@ async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() {
.unwrap();

assert!(router.dispatch_kv_plan(request, plan).await.is_err());
assert_eq!(
router.request_metrics.requests_started_total().get(),
requests_started_before + 1
);
assert_eq!(router.request_metrics.requests_started_total.get(), 1);
let loads = router
.kv_router()
.get_potential_loads(&[], None, None, None, None)
Expand All @@ -1592,7 +1586,6 @@ async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() {
async fn prefill_busy_probe_does_not_admit_a_request() {
let (router, runtime) = router(None).await;
let request = Context::new(request());
let requests_started_before = router.request_metrics.requests_started_total().get();

assert!(!router.prefill_worker_busy(&request, 0.5).await.unwrap());
let loads = router
Expand All @@ -1601,10 +1594,7 @@ async fn prefill_busy_probe_does_not_admit_a_request() {
.await
.unwrap();
assert!(loads.iter().all(|load| load.active_requests == 0));
assert_eq!(
router.request_metrics.requests_started_total().get(),
requests_started_before
);
assert_eq!(router.request_metrics.requests_started_total.get(), 0);

drop(router);
runtime.shutdown();
Expand Down Expand Up @@ -1657,8 +1647,6 @@ async fn session_affinity_disabled_does_not_create_coordinator() {
async fn router_request_counters_follow_admission_and_completion_lifecycle() {
let (router, runtime) = router(None).await;
let metrics = router.request_metrics.clone();
let started_before = metrics.requests_started_total().get();
let completed_before = metrics.requests_total.get();

let controller = Controller::new("pre-admission-cancellation".to_string());
controller.stop();
Expand All @@ -1674,23 +1662,23 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() {
.await
.is_err()
);
assert_eq!(metrics.requests_started_total().get(), started_before);
assert_eq!(metrics.requests_started_total.get(), 0);

let (_, _, mut query_guard) = track_request(&router, true).await;
query_guard.abort().await;
drop(query_guard);
assert_eq!(metrics.requests_started_total().get(), started_before);
assert_eq!(metrics.requests_started_total.get(), 0);

let (_, _, mut cancelled_guard) = track_request(&router, false).await;

assert_eq!(metrics.requests_started_total().get(), started_before + 1);
assert_eq!(metrics.requests_total.get(), completed_before);
assert_eq!(metrics.requests_started_total.get(), 1);
assert_eq!(metrics.requests_total.get(), 0);

// Admission remains counted even when the request aborts before dispatch.
cancelled_guard.abort().await;
drop(cancelled_guard);
assert_eq!(metrics.requests_started_total().get(), started_before + 1);
assert_eq!(metrics.requests_total.get(), completed_before);
assert_eq!(metrics.requests_started_total.get(), 1);
assert_eq!(metrics.requests_total.get(), 0);

let mut failed_input = request();
failed_input.migration_state = Some(Default::default());
Expand Down Expand Up @@ -1728,16 +1716,16 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() {
.is_err()
);
assert_eq!(migration_state.excluded_worker_ids(), vec![failed_worker]);
assert_eq!(metrics.requests_started_total().get(), started_before + 2);
assert_eq!(metrics.requests_total.get(), completed_before);
assert_eq!(metrics.requests_started_total.get(), 2);
assert_eq!(metrics.requests_total.get(), 0);

let (_, _, mut completed_guard) = track_request(&router, false).await;
completed_guard.start_dispatch("aggregated");
completed_guard.mark_dispatched();
completed_guard.finish().await;
drop(completed_guard);
assert_eq!(metrics.requests_started_total().get(), started_before + 3);
assert_eq!(metrics.requests_total.get(), completed_before + 1);
assert_eq!(metrics.requests_started_total.get(), 3);
assert_eq!(metrics.requests_total.get(), 1);

let mut builtin_guard = RequestGuard::<DefaultWorkerSelector>::new_builtin(
Arc::clone(&metrics),
Expand All @@ -1746,10 +1734,21 @@ async fn router_request_counters_follow_admission_and_completion_lifecycle() {
None,
&request(),
);
assert_eq!(metrics.requests_started_total().get(), started_before + 4);
assert_eq!(metrics.requests_started_total.get(), 4);
builtin_guard.abort().await;
drop(builtin_guard);
assert_eq!(metrics.requests_total.get(), completed_before + 1);
assert_eq!(metrics.requests_total.get(), 1);

let (other_router, other_runtime) = router_with_workers(None, &[8]).await;
Comment thread
jthomson04 marked this conversation as resolved.
let (_, _, mut other_guard) = track_request(&other_router, false).await;
other_guard.start_dispatch("aggregated");
other_guard.mark_dispatched();
other_guard.finish().await;
drop(other_guard);
assert_eq!(metrics.requests_started_total.get(), 4);
assert_eq!(metrics.requests_total.get(), 1);
drop(other_router);
other_runtime.shutdown();

drop(router);
runtime.shutdown();
Expand Down
Loading