push_router.rs 22.6 KB
Newer Older
1
2
3
4
5
6
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::sync::Arc;

use anyhow::Result;
7
use dynamo_kv_router::protocols::{TokensWithHashes, WorkerWithDpRank};
8
use dynamo_runtime::{
9
    dynamo_nvtx_range,
10
11
12
13
14
15
16
17
    pipeline::{
        AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter, ResponseStream,
        SingleIn, async_trait,
    },
    protocols::annotated::Annotated,
};
use futures::stream::{self, StreamExt};
use serde_json::json;
18
use tracing::Instrument;
19
20

use crate::{
21
    kv_router::{KvRouter, metrics::RouterRequestMetrics},
22
    preprocessor::PreprocessedRequest,
23
24
25
26
    protocols::common::{
        llm_backend::LLMEngineOutput,
        timing::{RequestPhase, RequestTracker},
    },
27
28
29
30
31
32
33
34
35
36
};

pub struct KvPushRouter {
    inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
    pub chooser: Arc<KvRouter>,
}

/// Result of worker selection containing instance ID, dp_rank, and overlap amount.
struct WorkerSelection {
    instance_id: u64,
37
38
39
    backend_dp_rank: Option<u32>,
    bookkeeping_dp_rank: Option<u32>,
    overlap_amount: Option<u32>,
40
41
}

42
43
/// Drop guard that manages the full lifecycle of a routed request:
/// per-item tracking (prefill, first token, output blocks) and final cleanup (free + metrics).
44
45
46
47
48
49
///
/// In the happy path, `finish().await` runs cleanup inline in the async context.
/// If the stream is dropped early (e.g., client disconnect, consumer drop), the
/// `Drop` impl fires and spawns a task to call `free()`.
struct RequestGuard {
    chooser: Arc<KvRouter>,
50
    scheduler_tracked: bool,
51
52
    context_id: String,
    tracker: Option<Arc<RequestTracker>>,
53
    request_metrics: Arc<RouterRequestMetrics>,
54
55
56
    cumulative_osl: usize,
    metrics_recorded: bool,
    freed: bool,
57
58
59
60
61
62
63
    prefill_marked: bool,
    first_token_recorded: bool,
    track_output_blocks: bool,
    current_total_blocks: usize,
    isl_tokens: usize,
    block_size: usize,
    expected_output_tokens: Option<u32>,
64
65
66
}

impl RequestGuard {
67
68
69
70
71
72
73
74
    async fn on_item(&mut self, item: &Annotated<LLMEngineOutput>) {
        if !self.prefill_marked {
            let has_tokens = item
                .data
                .as_ref()
                .map(|d| !d.token_ids.is_empty())
                .unwrap_or(false);
            if has_tokens {
75
76
77
                if self.scheduler_tracked
                    && let Err(e) = self.chooser.mark_prefill_completed(&self.context_id).await
                {
78
79
80
81
82
83
84
85
86
87
88
89
90
91
                    tracing::warn!(
                        "Failed to mark prefill completed for request {}: {e}",
                        self.context_id
                    );
                }
                self.prefill_marked = true;
            }
        }

        let new_tokens = item.data.as_ref().map(|d| d.token_ids.len()).unwrap_or(0);

        if !self.first_token_recorded && new_tokens > 0 {
            if let Some(ref tracker) = self.tracker {
                tracker.record_first_token();
92
93
94
95
                if let Some(ttft) = tracker.ttft_ms() {
                    self.request_metrics
                        .time_to_first_token_seconds
                        .observe(ttft / 1000.0);
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
                }
            }
            self.first_token_recorded = true;
        }

        self.cumulative_osl += new_tokens;

        if self.track_output_blocks {
            let new_total_blocks =
                (self.isl_tokens + self.cumulative_osl).div_ceil(self.block_size);
            if new_total_blocks > self.current_total_blocks {
                let decay_fraction = self
                    .expected_output_tokens
                    .map(|eot| (1.0 - (self.cumulative_osl as f64 / eot.max(1) as f64)).max(0.0));
                if let Err(e) = self
                    .chooser
                    .add_output_block(&self.context_id, decay_fraction)
                {
                    tracing::warn!(
                        "Failed to add output block for request {}: {e}",
                        self.context_id
                    );
                }

                if let Some(ref tracker) = self.tracker {
                    tracker.record_osl(self.cumulative_osl);
                    tracker.record_finish();
123
124
125
126
                    if let Some(avg_itl) = tracker.avg_itl_ms() {
                        self.request_metrics
                            .inter_token_latency_seconds
                            .observe(avg_itl / 1000.0);
127
128
129
130
131
132
133
134
                    }
                }

                self.current_total_blocks = new_total_blocks;
            }
        }
    }

135
136
    async fn finish(&mut self) {
        self.record_metrics();
137
138
139
        if self.scheduler_tracked
            && let Err(e) = self.chooser.free(&self.context_id).await
        {
140
141
142
143
144
145
146
147
148
149
150
151
152
153
            tracing::warn!("Failed to free request {}: {e}", self.context_id);
        }
        self.freed = true;
    }

    fn record_metrics(&mut self) {
        if self.metrics_recorded {
            return;
        }
        self.metrics_recorded = true;
        if let Some(ref tracker) = self.tracker {
            tracker.record_finish();
            tracker.record_osl(self.cumulative_osl);
        }
154
155
156
157
        self.request_metrics
            .output_sequence_tokens
            .observe(self.cumulative_osl as f64);
        self.request_metrics.requests_total.inc();
158
159
160
161
162
163
    }
}

impl Drop for RequestGuard {
    fn drop(&mut self) {
        self.record_metrics();
164
        if !self.freed && self.scheduler_tracked {
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
            let chooser = self.chooser.clone();
            let context_id = self.context_id.clone();
            let Ok(handle) = tokio::runtime::Handle::try_current() else {
                tracing::warn!("No tokio runtime for drop guard free of request {context_id}");
                return;
            };
            handle.spawn(async move {
                if let Err(e) = chooser.free(&context_id).await {
                    tracing::warn!("Failed to free request {context_id} (drop guard): {e}");
                }
            });
        }
    }
}

180
181
182
183
184
impl KvPushRouter {
    pub fn new(
        inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
        chooser: Arc<KvRouter>,
    ) -> Self {
185
186
187
188
        // Eagerly register router request metrics (as zeros) so they are
        // scrapeable before any requests arrive. Both the frontend pipeline
        // and the standalone router create KvPushRouter, so this covers both.
        RouterRequestMetrics::from_component(chooser.client().endpoint.component());
189

190
        KvPushRouter { inner, chooser }
191
192
193
194
    }

    /// Select a worker for the request, either using a preselected worker or finding the best match.
    ///
195
    /// When `is_query_only` is false, this also registers the request with the scheduler via `add_request`.
196
197
198
199
200
201
202
    async fn select_worker(
        &self,
        context_id: &str,
        request: &PreprocessedRequest,
        phase: RequestPhase,
        is_query_only: bool,
    ) -> Result<WorkerSelection, Error> {
203
        let _nvtx_select = dynamo_nvtx_range!("route.select_worker");
204
205
        let routing = request.routing.as_ref();
        let lora_name = routing.and_then(|r| r.lora_name.clone());
206
        let priority_jump = routing.and_then(|r| r.priority_jump).unwrap_or(0.0);
207
        let expected_output_tokens = routing.and_then(|r| r.expected_output_tokens);
208
        let allowed_worker_ids = routing.and_then(|r| r.allowed_worker_ids.clone());
209
        let (routing_token_ids, block_mm_infos) = request.block_mm_routing_info();
210
211
212
213
214
215
216
217
218
219
220

        // Get pre-selected worker based on phase, with backend_instance_id as fallback
        let preselected_id = match phase {
            RequestPhase::Prefill => {
                routing.and_then(|r| r.prefill_worker_id.or(r.backend_instance_id))
            }
            RequestPhase::Decode => {
                routing.and_then(|r| r.decode_worker_id.or(r.backend_instance_id))
            }
            RequestPhase::Aggregated => routing.and_then(|r| r.backend_instance_id),
        };
221
222
223
224
        let requested_dp_rank = match phase {
            RequestPhase::Prefill => routing.and_then(|r| r.prefill_dp_rank.or(r.dp_rank)),
            RequestPhase::Decode | RequestPhase::Aggregated => routing.and_then(|r| r.dp_rank),
        };
225
226

        let Some(id) = preselected_id else {
227
            let _nvtx_kv = dynamo_nvtx_range!("route.kv_match");
228
229
230
231
            let (best_worker, overlap_amount) = self
                .chooser
                .find_best_match(
                    Some(context_id),
232
233
                    routing_token_ids,
                    block_mm_infos,
234
235
236
                    request.router_config_override.as_ref(),
                    !is_query_only,
                    lora_name,
237
                    priority_jump,
238
                    expected_output_tokens,
239
                    allowed_worker_ids,
240
241
242
                )
                .await?;

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
            if !is_query_only {
                let total_blocks = routing_token_ids
                    .len()
                    .div_ceil(self.chooser.block_size() as usize);
                // NOTE: tests/mm_router/test_vllm_mm_router_e2e.py parses this log line.
                // Keep the "[ROUTING] ... with X/Y blocks overlap" shape stable unless
                // router tests are updated together.
                tracing::debug!(
                    request_id = %context_id,
                    worker_id = best_worker.worker_id,
                    dp_rank = best_worker.dp_rank,
                    overlap_blocks = overlap_amount,
                    total_blocks = total_blocks,
                    "[ROUTING] Best: worker_{} dp_rank={} with {}/{} blocks overlap",
                    best_worker.worker_id,
                    best_worker.dp_rank,
                    overlap_amount,
                    total_blocks,
                );
            }

264
265
            return Ok(WorkerSelection {
                instance_id: best_worker.worker_id,
266
267
268
                backend_dp_rank: Some(best_worker.dp_rank),
                bookkeeping_dp_rank: Some(best_worker.dp_rank),
                overlap_amount: Some(overlap_amount),
269
270
271
            });
        };

272
273
274
        let backend_dp_rank =
            requested_dp_rank.or_else(|| self.chooser.unique_dp_rank_for_worker(id));

275
276
        tracing::debug!(
            worker_id = id,
277
            dp_rank = ?backend_dp_rank,
278
279
280
281
            ?phase,
            "Routing to specified worker"
        );

282
283
284
285
286
        let (bookkeeping_dp_rank, overlap_amount) = if let Some(dp_rank) = backend_dp_rank {
            let worker = WorkerWithDpRank::new(id, dp_rank);
            let overlap_blocks = self
                .chooser
                .get_overlap_blocks(
287
                    routing_token_ids,
288
                    block_mm_infos,
289
                    worker,
290
                    lora_name.as_deref(),
291
                )
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
                .await?;

            if !is_query_only {
                self.chooser
                    .add_request(
                        context_id.to_string(),
                        routing_token_ids,
                        block_mm_infos,
                        overlap_blocks,
                        expected_output_tokens,
                        worker,
                        lora_name,
                        request.router_config_override.as_ref(),
                    )
                    .await;
            } else {
                tracing::debug!(
                    request_id = %context_id,
                    worker_id = id,
                    dp_rank = dp_rank,
                    "Skipping add_request - query-only request"
                );
            }

            (Some(dp_rank), Some(overlap_blocks))
317
318
319
320
        } else {
            tracing::debug!(
                request_id = %context_id,
                worker_id = id,
321
322
                ?phase,
                "Routing to specified worker without resolved dp_rank; skipping scheduler bookkeeping"
323
            );
324
325
            (None, None)
        };
326
327
328

        Ok(WorkerSelection {
            instance_id: id,
329
330
331
            backend_dp_rank,
            bookkeeping_dp_rank,
            overlap_amount,
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        })
    }
}

#[async_trait]
impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutput>>, Error>
    for KvPushRouter
{
    /// Generate method that handles KV-aware routing with three distinct behaviors:
    ///
    /// 1. **If `query_instance_id` annotation is set**:
    ///    - Returns the best matching worker ID without routing the request
    ///    - Does NOT update any router local states
    ///    - Response includes worker_instance_id and token_data annotations
    ///
    /// 2. **If `backend_instance_id` is set in the request**:
    ///    - Routes directly to the specified backend instance
    ///    - DOES update router states to track this request (unless query_instance_id is also set)
    ///    - Bypasses the normal KV matching logic
    ///
    /// 3. **If neither are set (default behavior)**:
    ///    - Finds the best worker based on KV cache overlap
    ///    - Updates router states to track the request
    ///    - Routes to the selected worker
    ///
    /// The router state updates include tracking active sequences and managing
    /// prefill/completion lifecycle for proper KV cache management.
    async fn generate(
        &self,
        request: SingleIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        // Extract context ID for request tracking
        let context_id = request.context().id().to_string();

        // Simple query-only detection: presence of query_instance_id annotation means query-only mode
        let is_query_only = request.get_annotation_value("query_instance_id").is_some();

        // Get phase from tracker (defaults to Aggregated if no tracker or phase not set)
        let phase = request
            .tracker
            .as_ref()
            .map(|t| t.phase())
            .unwrap_or(RequestPhase::Aggregated);

        let block_size = self.chooser.block_size() as usize;
        let selection = self
378
            .select_worker(&context_id, &request, phase, is_query_only)
379
            .instrument(tracing::info_span!("kv_router.select_worker"))
380
381
382
            .await?;
        let WorkerSelection {
            instance_id,
383
384
            backend_dp_rank,
            bookkeeping_dp_rank,
385
386
            overlap_amount,
        } = selection;
387
        let scheduler_tracked = !is_query_only && bookkeeping_dp_rank.is_some();
388
389
390
391
392

        // In approximate mode (use_kv_events=false), record the routing decision
        // so the indexer can track cache state based on routing decisions.
        // This covers both pre-selected workers and find_best_match selections.
        if !is_query_only && !self.chooser.kv_router_config().use_kv_events {
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
            if let Some(dp_rank) = bookkeeping_dp_rank {
                let lora_name = request.routing.as_ref().and_then(|r| r.lora_name.clone());
                let (routing_token_ids, block_mm_infos) = request.block_mm_routing_info();
                let worker = WorkerWithDpRank::new(instance_id, dp_rank);
                let mut tokens_with_hashes =
                    TokensWithHashes::new(routing_token_ids.to_vec(), self.chooser.block_size())
                        .with_is_eagle(self.chooser.is_eagle());
                if let Some(infos) = block_mm_infos {
                    tokens_with_hashes = tokens_with_hashes.with_mm_infos(infos.to_vec());
                }
                if let Some(lora_name) = lora_name {
                    tokens_with_hashes = tokens_with_hashes.with_lora_name(lora_name);
                }
                if let Err(e) = self
                    .chooser
                    .record_routing_decision(tokens_with_hashes, worker)
                    .await
                {
                    tracing::warn!(
                        request_id = %context_id,
                        worker_id = instance_id,
                        dp_rank = dp_rank,
                        error = %e,
                        "Failed to record routing decision in approximate mode"
                    );
                }
            } else {
                tracing::debug!(
421
422
                    request_id = %context_id,
                    worker_id = instance_id,
423
                    "Skipping approximate-mode routing decision for unresolved dp_rank"
424
425
426
427
                );
            }
        }

428
        // Record routing metrics on tracker and observe ISL + prefill start.
429
430
        let request_metrics =
            RouterRequestMetrics::from_component(self.chooser.client().endpoint.component());
431
        if let Some(ref tracker) = request.tracker {
432
            let (routing_token_ids, _) = request.block_mm_routing_info();
433
            let isl_blocks = routing_token_ids.len().div_ceil(block_size);
434
435
436
            if let Some(overlap_amount) = overlap_amount {
                tracker.record_kv_hit(overlap_amount, isl_blocks);
            }
437
            tracker.record_isl(
438
                routing_token_ids.len(),
439
                overlap_amount.map(|overlap| overlap as usize * block_size),
440
            );
441
            tracker.record_worker(instance_id, backend_dp_rank, self.chooser.worker_type());
442
            tracker.record_router_queue_depth(self.chooser.pending_count());
443
444
            if let Some(hit_rate) = tracker.kv_hit_rate() {
                request_metrics.kv_hit_rate.observe(hit_rate);
445
            }
446
        }
447
448
449
        request_metrics
            .input_sequence_tokens
            .observe(request.token_ids.len() as f64);
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480

        // Handle query-only requests: early return with worker info
        if is_query_only {
            let stream_context = request.context().clone();
            let worker_id_info = request.tracker.as_ref().and_then(|t| t.get_worker_info());

            tracing::trace!(
                ?phase,
                worker_id = instance_id,
                ?worker_id_info,
                "Returning worker selection (query-only mode)"
            );

            let output = LLMEngineOutput {
                disaggregated_params: Some(json!({
                    "worker_id": worker_id_info,
                    "token_ids": request.token_ids
                })),
                ..Default::default()
            };
            let response = Annotated::from_data(output);
            let stream = stream::iter(vec![response]);
            return Ok(ResponseStream::new(Box::pin(stream), stream_context));
        }

        // Route to worker
        let isl_tokens = request.token_ids.len();
        let expected_output_tokens = request
            .routing
            .as_ref()
            .and_then(|r| r.expected_output_tokens);
481
        let track_output_blocks = self.chooser.kv_router_config().router_track_output_blocks;
482
        let tracker = request.tracker.clone();
483
484

        let (mut backend_input, context) = request.into_parts();
485
        backend_input.routing_mut().dp_rank = backend_dp_rank;
486
487
        let updated_request = context.map(|_| backend_input);

488
489
490
491
492
        // Record prefill start right before pushing to backend (OnceLock: first call wins).
        if let Some(ref tracker) = tracker {
            tracker.record_prefill_start();
        }

493
        let chooser = self.chooser.clone();
494
495
496
497
498
499
500
        let mut response_stream = self
            .inner
            .direct(updated_request, instance_id)
            .instrument(tracing::info_span!(
                "kv_router.route_request",
                request_id = %context_id,
                worker_id = instance_id,
501
502
                dp_rank = ?backend_dp_rank,
                overlap_blocks = ?overlap_amount,
503
504
505
                phase = ?phase,
            ))
            .await?;
506
507
508
509
        let stream_context = response_stream.context();
        let context_for_monitoring = stream_context.clone();

        let wrapped_stream = Box::pin(async_stream::stream! {
510
511
            let mut guard = RequestGuard {
                chooser: chooser.clone(),
512
                scheduler_tracked,
513
514
515
516
517
518
                context_id: context_id.clone(),
                tracker: tracker.clone(),
                request_metrics: request_metrics.clone(),
                cumulative_osl: 0,
                metrics_recorded: false,
                freed: false,
519
520
                prefill_marked: false,
                first_token_recorded: false,
521
                track_output_blocks: scheduler_tracked && track_output_blocks,
522
523
524
525
                current_total_blocks: isl_tokens.div_ceil(block_size),
                isl_tokens,
                block_size,
                expected_output_tokens,
526
            };
527
528
529
530
531
532
533
534
535
536
537
538
539
540

            loop {
                tokio::select! {
                    biased;

                    _ = context_for_monitoring.stopped() => {
                        tracing::debug!("Request {context_id} cancelled, ending stream");
                        break;
                    }

                    item = response_stream.next() => {
                        let Some(item) = item else {
                            break;
                        };
541
                        guard.on_item(&item).await;
542
543
544
545
546
                        yield item;
                    }
                }
            }

547
            guard.finish().await;
548
549
550
551
        });
        Ok(ResponseStream::new(wrapped_stream, stream_context))
    }
}
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596

/// A direct routing wrapper for `RouterMode::Direct`.
///
/// This wraps a `PushRouter` and reads worker IDs from each request's routing hints,
/// then routes directly to the specified worker. Used when an external router
/// (e.g., EPP) handles worker selection.
pub struct DirectRoutingRouter {
    inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
}

impl DirectRoutingRouter {
    pub fn new(inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>) -> Self {
        DirectRoutingRouter { inner }
    }

    /// Extract worker ID from request routing hints.
    /// Returns an error if no worker ID is found (required in direct routing mode).
    fn get_worker_id(request: &PreprocessedRequest) -> Result<u64, Error> {
        let routing = request.routing.as_ref();
        let worker_id = routing.and_then(|r| r.decode_worker_id.or(r.backend_instance_id));

        worker_id.ok_or_else(|| {
            anyhow::anyhow!(
                "Worker ID required (--direct-route) but none found in request. \
                 Expected decode_worker_id or backend_instance_id to be set by external router (e.g., EPP)."
            )
        })
    }
}

#[async_trait]
impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutput>>, Error>
    for DirectRoutingRouter
{
    async fn generate(
        &self,
        request: SingleIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        let worker_id = Self::get_worker_id(&request)?;

        tracing::debug!(worker_id = worker_id, "Direct routing to specified worker");

        self.inner.direct(request, worker_id).await
    }
}