router.rs 28.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result, anyhow};
10
use dynamo_kv_router::LocalBlockHash;
11
12
use dynamo_kv_router::config::KvRouterConfig;
use dynamo_kv_router::protocols::{
13
14
    BlockHashOptions, OverlapScores, PrefillLoadHint, RouterEvent, WorkerConfigLike, WorkerId,
    WorkerWithDpRank, compute_block_hash_for_seq,
15
16
17
18
19
20
21
};
use dynamo_kv_router::queue::DEFAULT_MAX_BATCHED_TOKENS;
use dynamo_kv_router::{
    ActiveSequencesMultiWorker, DefaultWorkerSelector, RadixTree, RouterSchedulingPolicy,
    SchedulingPolicy, SchedulingRequest, SequenceRequest, WorkerSelector,
};
use dynamo_tokens::SequenceHash;
22
use rustc_hash::FxHashMap;
23
use tokio::time::Instant;
24
25
use uuid::Uuid;

26
use super::{RouterEffects, WorkerAdmission};
27
28
use crate::common::protocols::DirectRequest;
use crate::common::protocols::MockEngineArgs;
29
use crate::loadgen::ReplayRequestHashes;
30
31
32
use crate::replay::ReplayPrefillLoadEstimator;
use crate::replay::router_shared::{
    ReplayNoopPublisher, ReplayWorkerConfig, replay_policy, replay_router_config, replay_selector,
33
    replay_slots, replay_worker_config, replay_workers_with_configs,
34
};
35
36
37

type ReplayQueueKey = <RouterSchedulingPolicy as SchedulingPolicy>::Key;

38
39
40
41
42
43
44
45
46
47
/// Internal result of a successful ``admit_request`` call: the chosen
/// worker plus the router's view of prefix-cache overlap, so callers can
/// forward the overlap stats to the traffic accumulator.
#[derive(Debug, Clone, Copy)]
struct AdmitOutcome {
    worker_idx: usize,
    overlap_blocks: u32,
    isl_blocks: u32,
}

48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
#[cfg(test)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct OfflinePendingRequestSnapshot {
    pub(crate) uuid: Uuid,
    pub(crate) overlap_blocks_by_worker: Vec<(usize, u32)>,
}

#[cfg(test)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct OfflineIndexerSnapshot {
    pub(crate) total_cached_blocks: usize,
    pub(crate) cached_blocks_by_worker: Vec<(usize, usize)>,
}

#[cfg(test)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct OfflineRouterSnapshot {
    pub(crate) pending: Vec<OfflinePendingRequestSnapshot>,
    pub(crate) active_blocks_by_worker: Vec<(usize, usize)>,
    pub(crate) active_tokens_by_worker: Vec<(usize, usize)>,
    pub(crate) indexer: OfflineIndexerSnapshot,
}

71
72
73
74
75
76
77
78
79
80
81
82
83
84
struct SyncReplayIndexer {
    block_size: u32,
    tree: RadixTree,
}

impl SyncReplayIndexer {
    fn new(block_size: u32) -> Self {
        Self {
            block_size,
            tree: RadixTree::new(),
        }
    }

    fn find_matches_for_request(&self, tokens: &[u32], lora_name: Option<&str>) -> OverlapScores {
85
86
87
88
89
90
91
92
        let sequence = compute_block_hash_for_seq(
            tokens,
            self.block_size,
            BlockHashOptions {
                lora_name,
                ..Default::default()
            },
        );
93
94
95
        self.tree.find_matches(sequence, false)
    }

96
97
98
99
    fn find_matches_for_hashes(&self, local_block_hashes: Vec<LocalBlockHash>) -> OverlapScores {
        self.tree.find_matches(local_block_hashes, false)
    }

100
101
102
    fn apply_event(&mut self, event: RouterEvent) -> Result<()> {
        self.tree.apply_event(event).map_err(Into::into)
    }
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

    #[cfg(test)]
    fn debug_snapshot(&self) -> OfflineIndexerSnapshot {
        let mut blocks_by_worker = HashMap::<usize, usize>::new();
        for event in self.tree.dump_tree_as_events() {
            *blocks_by_worker
                .entry(event.worker_id as usize)
                .or_default() += 1;
        }
        let mut cached_blocks_by_worker = blocks_by_worker.into_iter().collect::<Vec<_>>();
        cached_blocks_by_worker.sort_unstable_by_key(|(worker_id, _)| *worker_id);

        OfflineIndexerSnapshot {
            total_cached_blocks: self.tree.current_size(),
            cached_blocks_by_worker,
        }
    }
120
121
122
123
124
125
126
}

struct PendingRequest {
    uuid: Uuid,
    token_seq: Option<Vec<SequenceHash>>,
    isl_tokens: usize,
    overlaps: OverlapScores,
127
    track_prefill_tokens: bool,
128
129
130
131
132
133
134
135
136
137
    expected_output_tokens: Option<u32>,
}

impl PendingRequest {
    fn request_id(&self) -> String {
        self.uuid.to_string()
    }

    fn scheduling_request(
        &self,
138
139
        decode_blocks: FxHashMap<WorkerWithDpRank, usize>,
        prefill_tokens: FxHashMap<WorkerWithDpRank, usize>,
140
141
142
143
144
145
146
147
    ) -> SchedulingRequest {
        SchedulingRequest {
            maybe_request_id: Some(self.request_id()),
            token_seq: self.token_seq.clone(),
            isl_tokens: self.isl_tokens,
            overlaps: self.overlaps.clone(),
            decode_blocks,
            prefill_tokens,
148
            track_prefill_tokens: self.track_prefill_tokens,
149
150
151
152
153
            router_config_override: None,
            update_states: true,
            lora_name: None,
            priority_jump: 0.0,
            expected_output_tokens: self.expected_output_tokens,
154
            pinned_worker: None,
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
            allowed_worker_ids: None,
            resp_tx: None,
        }
    }
}

struct QueueEntry {
    key: ReplayQueueKey,
    _enqueue_time_ms: f64,
    enqueue_seq: u64,
    request: PendingRequest,
}

impl Eq for QueueEntry {}

impl PartialEq for QueueEntry {
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key && self.enqueue_seq == other.enqueue_seq
    }
}

impl Ord for QueueEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        self.key
            .cmp(&other.key)
            .then_with(|| other.enqueue_seq.cmp(&self.enqueue_seq))
    }
}

impl PartialOrd for QueueEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

pub(crate) struct OfflineReplayRouter {
    config: KvRouterConfig,
    block_size: u32,
    queue_threshold: Option<f64>,
194
    worker_config_template: ReplayWorkerConfig,
195
196
197
198
199
200
201
    workers_with_configs: HashMap<WorkerId, ReplayWorkerConfig>,
    slots: Arc<ActiveSequencesMultiWorker<ReplayNoopPublisher>>,
    selector: DefaultWorkerSelector,
    policy: RouterSchedulingPolicy,
    pending: BinaryHeap<QueueEntry>,
    next_enqueue_seq: u64,
    indexer: SyncReplayIndexer,
202
203
    prefill_load_estimator: Option<ReplayPrefillLoadEstimator>,
    decay_time_epoch: Instant,
204
205
206
207
208
209
}

impl OfflineReplayRouter {
    pub(crate) fn new(
        args: &MockEngineArgs,
        router_config: Option<KvRouterConfig>,
210
        prefill_load_estimator: Option<ReplayPrefillLoadEstimator>,
211
212
213
        num_workers: usize,
    ) -> Result<Self> {
        let config = replay_router_config(args, router_config);
214
        let worker_config_template = replay_worker_config(args);
215
216
217
218
        let workers_with_configs = replay_workers_with_configs(args, num_workers);
        let slots = replay_slots(args, &workers_with_configs);
        let selector = replay_selector(&config);
        let policy = replay_policy(&config, args);
219
        let queue_threshold = config.router_queue_threshold;
220
221
222
223
224

        Ok(Self {
            config,
            block_size: args.block_size as u32,
            queue_threshold,
225
            worker_config_template,
226
227
228
229
230
231
232
            workers_with_configs,
            slots,
            selector,
            policy,
            pending: BinaryHeap::new(),
            next_enqueue_seq: 0,
            indexer: SyncReplayIndexer::new(args.block_size as u32),
233
234
235
236
237
            prefill_load_estimator,
            // This is only a base Instant for converting replay `now_ms` values into
            // synthetic `Instant`s. All subsequent decay/accounting uses virtual replay
            // time derived from this epoch, not wall-clock progression.
            decay_time_epoch: Instant::now(),
238
239
240
        })
    }

241
    pub(crate) fn on_request_arrival(
242
243
        &mut self,
        request: &DirectRequest,
244
        replay_hashes: Option<ReplayRequestHashes>,
245
        now_ms: f64,
246
    ) -> Result<RouterEffects> {
247
        let pending = self.build_pending_request(request, replay_hashes)?;
248
        let decay_now = self.decay_now(now_ms);
249
250
        let should_queue = self
            .queue_threshold
251
            .is_some_and(|threshold| self.all_workers_busy(threshold, decay_now));
252
253
254
255
256
257
258
259
260
261

        if should_queue {
            let key = self.enqueue_key(now_ms, &pending);
            self.pending.push(QueueEntry {
                key,
                _enqueue_time_ms: now_ms,
                enqueue_seq: self.next_enqueue_seq,
                request: pending,
            });
            self.next_enqueue_seq += 1;
262
            return Ok(RouterEffects::default());
263
264
        }

265
266
267
268
        let uuid = request
            .uuid
            .expect("offline replay requests must have UUIDs before router submission");
        let outcome = self.admit_request(pending, decay_now)?;
269
270
        Ok(RouterEffects {
            admissions: vec![WorkerAdmission {
271
272
273
274
                uuid,
                worker_idx: outcome.worker_idx,
                overlap_blocks: outcome.overlap_blocks,
                isl_blocks: outcome.isl_blocks,
275
276
            }],
        })
277
278
    }

279
280
281
282
283
    pub(crate) fn on_kv_events(&mut self, events: Vec<RouterEvent>) -> Result<RouterEffects> {
        for event in events {
            self.indexer.apply_event(event)?;
        }
        Ok(RouterEffects::default())
284
285
    }

286
287
288
289
290
291
    pub(crate) fn on_prefill_completed(
        &mut self,
        uuid: Uuid,
        now_ms: f64,
    ) -> Result<RouterEffects> {
        let decay_now = self.decay_now(now_ms);
292
        self.slots
293
            .mark_prefill_completed(&uuid.to_string(), decay_now)
294
            .map_err(anyhow::Error::from)?;
295
        Ok(RouterEffects {
296
            admissions: self.drain_pending(decay_now)?,
297
        })
298
299
    }

300
301
302
303
304
305
    pub(crate) fn on_request_completed(
        &mut self,
        uuid: Uuid,
        now_ms: f64,
    ) -> Result<RouterEffects> {
        let decay_now = self.decay_now(now_ms);
306
        self.slots
307
            .free(&uuid.to_string(), decay_now)
308
            .map_err(anyhow::Error::from)?;
309
        Ok(RouterEffects {
310
            admissions: self.drain_pending(decay_now)?,
311
        })
312
313
    }

314
315
316
317
318
    /// Drain queued requests that can now be admitted (e.g. after a new worker
    /// becomes available).
    pub(crate) fn try_drain_pending(&mut self, now_ms: f64) -> Result<RouterEffects> {
        let decay_now = self.decay_now(now_ms);
        Ok(RouterEffects {
319
            admissions: self.drain_pending(decay_now)?,
320
321
322
        })
    }

323
324
325
326
    pub(crate) fn pending_count(&self) -> usize {
        self.pending.len()
    }

327
    /// Register a new worker with the router without disturbing existing slot state.
328
329
    pub(crate) fn add_worker(&mut self, worker_id: usize) -> Result<()> {
        let wid = worker_id as WorkerId;
330
        if self
331
            .workers_with_configs
332
333
334
335
            .insert(wid, self.worker_config_template.clone())
            .is_some()
        {
            return Err(anyhow!("router worker {worker_id} already exists"));
336
        }
337
338
        let dp_range = HashMap::from([(wid, (0u32, 1u32))]);
        self.slots.register_external_workers(&dp_range);
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357

        Ok(())
    }

    /// Remove a worker from routing eligibility.
    ///
    /// Only removes the worker from the config map so the selector won't
    /// pick it for new requests.  The radix tree and active-sequence slots
    /// are left intact so that in-flight requests on this worker can still
    /// complete (free / mark_prefill_completed) and KV events can still
    /// reference existing blocks without "parent block not found" errors.
    /// Stale slot and indexer state is harmless — the selector and
    /// `all_workers_busy` both skip workers absent from `workers_with_configs`.
    pub(crate) fn remove_worker(&mut self, worker_id: usize) -> Result<()> {
        let wid = worker_id as WorkerId;
        self.workers_with_configs.remove(&wid);
        Ok(())
    }

358
359
360
361
362
363
    pub(crate) fn on_topology_changed(&mut self, now_ms: f64) -> Result<RouterEffects> {
        if self.workers_with_configs.is_empty() {
            return Ok(RouterEffects::default());
        }
        let decay_now = self.decay_now(now_ms);
        Ok(RouterEffects {
364
            admissions: self.drain_pending(decay_now)?,
365
366
367
        })
    }

368
    #[cfg(test)]
369
370
    pub(crate) fn debug_snapshot(&self, now_ms: f64) -> OfflineRouterSnapshot {
        let decay_now = self.decay_now(now_ms);
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
        let mut pending = self
            .pending
            .iter()
            .map(|entry| {
                let mut overlap_blocks_by_worker = entry
                    .request
                    .overlaps
                    .scores
                    .iter()
                    .map(|(worker, overlap)| (worker.worker_id as usize, *overlap))
                    .collect::<Vec<_>>();
                overlap_blocks_by_worker.sort_unstable_by_key(|(worker_id, _)| *worker_id);

                (
                    entry,
                    OfflinePendingRequestSnapshot {
                        uuid: entry.request.uuid,
                        overlap_blocks_by_worker,
                    },
                )
            })
            .collect::<Vec<_>>();
        pending.sort_unstable_by(|(left_entry, _), (right_entry, _)| {
            left_entry.cmp(right_entry).reverse()
        });

        let mut active_blocks_by_worker = self
            .slots
            .active_blocks()
            .into_iter()
            .map(|(worker, blocks)| (worker.worker_id as usize, blocks))
            .collect::<Vec<_>>();
        active_blocks_by_worker.sort_unstable_by_key(|(worker_id, _)| *worker_id);

        let mut active_tokens_by_worker = self
            .slots
407
            .active_tokens(decay_now)
408
409
410
411
412
413
414
415
416
417
418
419
420
            .into_iter()
            .map(|(worker, tokens)| (worker.worker_id as usize, tokens))
            .collect::<Vec<_>>();
        active_tokens_by_worker.sort_unstable_by_key(|(worker_id, _)| *worker_id);

        OfflineRouterSnapshot {
            pending: pending.into_iter().map(|(_, snapshot)| snapshot).collect(),
            active_blocks_by_worker,
            active_tokens_by_worker,
            indexer: self.indexer.debug_snapshot(),
        }
    }

421
422
423
424
    fn enqueue_key(&self, now_ms: f64, request: &PendingRequest) -> ReplayQueueKey {
        let arrival_offset = Duration::from_secs_f64((now_ms.max(0.0)) / 1000.0);
        self.policy.enqueue_key(
            arrival_offset,
425
            &request.scheduling_request(FxHashMap::default(), FxHashMap::default()),
426
427
428
        )
    }

429
430
431
432
    fn decay_now(&self, now_ms: f64) -> Instant {
        self.decay_time_epoch + Duration::from_secs_f64(now_ms.max(0.0) / 1000.0)
    }

433
434
435
436
437
    fn build_pending_request(
        &self,
        request: &DirectRequest,
        replay_hashes: Option<ReplayRequestHashes>,
    ) -> Result<PendingRequest> {
438
439
440
        let uuid = request
            .uuid
            .ok_or_else(|| anyhow!("offline replay requires requests to have stable UUIDs"))?;
441
442
443
444
445
446
447
448
449
450
451
452
453
454
        let (overlaps, token_seq) = match replay_hashes {
            Some(replay_hashes) => {
                let overlaps = self
                    .indexer
                    .find_matches_for_hashes(replay_hashes.local_block_hashes);
                let token_seq = if !self.config.router_track_active_blocks {
                    None
                } else if self.config.router_assume_kv_reuse {
                    Some(replay_hashes.sequence_hashes)
                } else {
                    self.config.compute_seq_hashes_for_tracking(
                        &request.tokens,
                        self.block_size,
                        None,
455
                        BlockHashOptions::default(),
456
                        None,
457
458
459
460
461
462
463
464
465
466
                    )
                };
                (overlaps, token_seq)
            }
            None => {
                let overlaps = self.indexer.find_matches_for_request(&request.tokens, None);
                let token_seq = self.config.compute_seq_hashes_for_tracking(
                    &request.tokens,
                    self.block_size,
                    None,
467
                    BlockHashOptions::default(),
468
                    None,
469
470
471
472
                );
                (overlaps, token_seq)
            }
        };
473
474
475
476
477
478

        Ok(PendingRequest {
            uuid,
            token_seq,
            isl_tokens: request.tokens.len(),
            overlaps,
479
            track_prefill_tokens: self.config.router_track_prefill_tokens,
480
481
482
483
484
485
486
            expected_output_tokens: Some(
                u32::try_from(request.max_output_tokens)
                    .context("max_output_tokens does not fit into u32")?,
            ),
        })
    }

487
488
489
490
491
    fn admit_request(
        &mut self,
        request: PendingRequest,
        decay_now: Instant,
    ) -> Result<AdmitOutcome> {
492
493
494
495
496
497
498
        let (decode_blocks, prefill_tokens) = self
            .slots
            .potential_blocks_and_tokens_with_prefill_tracking(
                request.token_seq.as_deref(),
                request.isl_tokens,
                request.overlaps.clone(),
                request.track_prefill_tokens,
499
                decay_now,
500
            );
501
502
503
504
505
506
507
508
509
        let scheduling_request = request.scheduling_request(decode_blocks, prefill_tokens);
        let selection = self.selector.select_worker(
            &self.workers_with_configs,
            &scheduling_request,
            self.block_size,
        )?;
        let worker_idx = usize::try_from(selection.worker.worker_id)
            .map_err(|_| anyhow!("selected worker id does not fit into usize"))?;
        let request_id = request.request_id();
510
511
512
513
514
        let prefill_load_hint = self.prefill_load_hint_for(
            request.isl_tokens,
            selection.overlap_blocks,
            request.track_prefill_tokens,
        );
515

516
517
518
519
        let isl_blocks = u32::try_from(request.isl_tokens.div_ceil(self.block_size as usize))
            .unwrap_or(u32::MAX);
        let overlap_blocks = selection.overlap_blocks;

520
        self.slots
521
522
523
524
525
526
527
528
529
530
531
532
            .add_request(
                SequenceRequest {
                    request_id,
                    token_sequence: request.token_seq,
                    track_prefill_tokens: request.track_prefill_tokens,
                    expected_output_tokens: request.expected_output_tokens,
                    prefill_load_hint,
                    worker: selection.worker,
                    lora_name: None,
                },
                decay_now,
            )
533
534
            .map_err(anyhow::Error::from)?;

535
536
537
538
539
        Ok(AdmitOutcome {
            worker_idx,
            overlap_blocks,
            isl_blocks,
        })
540
541
    }

542
    fn drain_pending(&mut self, decay_now: Instant) -> Result<Vec<WorkerAdmission>> {
543
544
545
546
547
        let Some(threshold) = self.queue_threshold else {
            return Ok(Vec::new());
        };

        let mut admissions = Vec::new();
548
        while !self.all_workers_busy(threshold, decay_now) {
549
550
551
552
            let Some(QueueEntry { request, .. }) = self.pending.pop() else {
                break;
            };
            let uuid = request.uuid;
553
554
555
556
557
558
559
            let outcome = self.admit_request(request, decay_now)?;
            admissions.push(WorkerAdmission {
                uuid,
                worker_idx: outcome.worker_idx,
                overlap_blocks: outcome.overlap_blocks,
                isl_blocks: outcome.isl_blocks,
            });
560
561
562
563
564
        }

        Ok(admissions)
    }

565
    fn all_workers_busy(&self, threshold: f64, decay_now: Instant) -> bool {
566
        let mut checked_any = false;
567
568
569
570
571
572
573
574
575
576
577
578
        let any_worker_not_busy =
            self.slots
                .any_worker_matches_active_tokens(decay_now, |worker, tokens| {
                    let Some(config) = self.workers_with_configs.get(&worker.worker_id) else {
                        return false;
                    };
                    checked_any = true;
                    let max_batched = config
                        .max_num_batched_tokens()
                        .unwrap_or(DEFAULT_MAX_BATCHED_TOKENS);
                    (tokens as f64) <= threshold * (max_batched as f64)
                });
579
580
581

        checked_any && !any_worker_not_busy
    }
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598

    fn prefill_load_hint_for(
        &self,
        isl_tokens: usize,
        overlap_blocks: u32,
        track_prefill_tokens: bool,
    ) -> Option<PrefillLoadHint> {
        if !track_prefill_tokens {
            return None;
        }

        let prefix = (overlap_blocks as usize) * (self.block_size as usize);
        let effective_isl = isl_tokens.saturating_sub(prefix);
        if effective_isl == 0 {
            return None;
        }

599
600
601
602
603
604
605
606
607
608
609
610
611
        let expected_prefill_duration = match &self.prefill_load_estimator {
            Some(estimator) => match estimator.predict_prefill_duration(1, effective_isl, prefix) {
                Ok(expected_prefill_duration) => Some(expected_prefill_duration),
                Err(error) => {
                    tracing::warn!(
                        effective_isl,
                        prefix,
                        "failed to predict replay prefill duration for active load tracking: {error}"
                    );
                    None
                }
            },
            None => None,
612
613
        };

614
615
616
617
        Some(PrefillLoadHint {
            initial_effective_prefill_tokens: effective_isl,
            expected_prefill_duration,
        })
618
619
620
621
622
623
624
625
626
627
628
629
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::Duration;

    use dynamo_kv_router::PrefillLoadEstimator;
    use dynamo_kv_router::config::{KvRouterConfig, RouterPrefillLoadModel};
    use uuid::Uuid;

630
    use super::{OfflineReplayRouter, WorkerAdmission};
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
    use crate::common::protocols::{DirectRequest, MockEngineArgs};
    use crate::replay::ReplayPrefillLoadEstimator;

    struct FixedPrefillLoadEstimator {
        duration: Duration,
    }

    impl PrefillLoadEstimator for FixedPrefillLoadEstimator {
        fn predict_prefill_duration(
            &self,
            _batch_size: usize,
            _effective_isl: usize,
            _prefix: usize,
        ) -> anyhow::Result<Duration> {
            Ok(self.duration)
        }
    }

    fn replay_args() -> MockEngineArgs {
        MockEngineArgs::builder()
            .block_size(64)
            .max_num_batched_tokens(Some(256))
            .build()
            .unwrap()
    }

657
658
659
660
661
662
663
664
    fn queueing_args() -> MockEngineArgs {
        MockEngineArgs::builder()
            .block_size(64)
            .max_num_batched_tokens(Some(64))
            .build()
            .unwrap()
    }

665
666
667
668
669
670
671
672
    fn router_config() -> KvRouterConfig {
        KvRouterConfig {
            router_track_prefill_tokens: true,
            router_prefill_load_model: RouterPrefillLoadModel::Aic,
            ..KvRouterConfig::default()
        }
    }

673
674
675
676
677
678
679
    fn queueing_router_config() -> KvRouterConfig {
        KvRouterConfig {
            router_queue_threshold: Some(0.5),
            ..KvRouterConfig::default()
        }
    }

680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
    fn estimator(duration: Duration) -> ReplayPrefillLoadEstimator {
        Arc::new(FixedPrefillLoadEstimator { duration })
    }

    fn request(uuid: u128, token: u32) -> DirectRequest {
        DirectRequest {
            tokens: vec![token; 64],
            max_output_tokens: 2,
            uuid: Some(Uuid::from_u128(uuid)),
            dp_rank: 0,
            arrival_timestamp_ms: Some(0.0),
        }
    }

    #[test]
    fn test_prefill_load_estimator_decays_offline_router_active_tokens() {
        let mut router = OfflineReplayRouter::new(
            &replay_args(),
            Some(router_config()),
            Some(estimator(Duration::from_secs(10))),
            1,
        )
        .unwrap();

        let effects = router
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        assert_eq!(effects.admissions.len(), 1);
        assert_eq!(
            router.debug_snapshot(0.0).active_tokens_by_worker,
            vec![(0, 64)]
        );
        assert_eq!(
            router.debug_snapshot(5_000.0).active_tokens_by_worker,
            vec![(0, 32)]
        );
        assert_eq!(
            router.debug_snapshot(10_000.0).active_tokens_by_worker,
            vec![(0, 0)]
        );
    }
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801

    #[test]
    fn test_single_worker_router_honors_queue_threshold() {
        let mut router =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 1)
                .unwrap();

        let first = router
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        assert_eq!(first.admissions.len(), 1);

        let second = router
            .on_request_arrival(&request(2, 8), None, 0.0)
            .unwrap();
        assert!(second.admissions.is_empty());
        assert_eq!(router.pending_count(), 1);
        assert_eq!(
            router
                .debug_snapshot(0.0)
                .pending
                .into_iter()
                .map(|request| request.uuid)
                .collect::<Vec<_>>(),
            vec![Uuid::from_u128(2)]
        );
    }

    #[test]
    fn test_scaled_down_router_matches_fresh_single_worker_queueing() {
        let mut fresh =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 1)
                .unwrap();
        fresh.on_request_arrival(&request(1, 7), None, 0.0).unwrap();
        fresh.on_request_arrival(&request(2, 8), None, 0.0).unwrap();

        let mut scaled =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 2)
                .unwrap();
        scaled.remove_worker(1).unwrap();
        scaled
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        scaled
            .on_request_arrival(&request(2, 8), None, 0.0)
            .unwrap();

        assert_eq!(scaled.pending_count(), fresh.pending_count());
        assert_eq!(
            scaled
                .debug_snapshot(0.0)
                .pending
                .into_iter()
                .map(|request| request.uuid)
                .collect::<Vec<_>>(),
            fresh
                .debug_snapshot(0.0)
                .pending
                .into_iter()
                .map(|request| request.uuid)
                .collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_router_can_scale_from_zero_workers() {
        let mut router =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 1)
                .unwrap();

        router.remove_worker(0).unwrap();
        router.add_worker(3).unwrap();

        let effects = router
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        assert_eq!(
            effects.admissions,
            vec![WorkerAdmission {
                uuid: Uuid::from_u128(1),
                worker_idx: 3,
802
803
                overlap_blocks: 0,
                isl_blocks: 1,
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
            }]
        );
    }

    #[test]
    fn test_add_worker_preserves_draining_worker_state() {
        let mut router =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 2)
                .unwrap();

        router
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        router
            .on_request_arrival(&request(2, 8), None, 0.0)
            .unwrap();
        assert_eq!(
            router.debug_snapshot(0.0).active_tokens_by_worker,
            vec![(0, 64), (1, 64)]
        );

        router.remove_worker(1).unwrap();
        router.add_worker(2).unwrap();

        assert_eq!(
            router.debug_snapshot(0.0).active_tokens_by_worker,
            vec![(0, 64), (1, 64), (2, 0)]
        );
    }

    #[test]
    fn test_topology_change_drains_pending_after_scale_up() {
        let mut router =
            OfflineReplayRouter::new(&queueing_args(), Some(queueing_router_config()), None, 1)
                .unwrap();

        router
            .on_request_arrival(&request(1, 7), None, 0.0)
            .unwrap();
        router
            .on_request_arrival(&request(2, 8), None, 0.0)
            .unwrap();
        assert_eq!(router.pending_count(), 1);

        router.add_worker(1).unwrap();
        let effects = router.on_topology_changed(0.0).unwrap();

        assert_eq!(
            effects.admissions,
            vec![WorkerAdmission {
                uuid: Uuid::from_u128(2),
                worker_idx: 1,
856
857
                overlap_blocks: 0,
                isl_blocks: 1,
858
859
860
861
862
863
864
865
            }]
        );
        assert_eq!(router.pending_count(), 0);
        assert_eq!(
            router.debug_snapshot(0.0).active_tokens_by_worker,
            vec![(0, 64), (1, 64)]
        );
    }
866
}