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

4
use std::collections::HashMap;
Yan Ru Pei's avatar
Yan Ru Pei committed
5
use std::sync::Arc;
6
use std::time::{Duration, Instant};
7

8
use anyhow::Result;
Yan Ru Pei's avatar
Yan Ru Pei committed
9
use dynamo_kv_router::{ConcurrentRadixTree, ThreadPoolIndexer};
10
use dynamo_runtime::{
11
    component::{Client, Endpoint},
12
    discovery::DiscoveryQuery,
13
    pipeline::{
14
15
        AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn,
        async_trait,
16
    },
17
    protocols::EndpointId,
18
    protocols::annotated::Annotated,
19
    traits::DistributedRuntimeProvider,
20
};
21
use futures::stream;
Yan Ru Pei's avatar
Yan Ru Pei committed
22
use tokio::sync::oneshot;
23
use tracing::Instrument;
24
use validator::Validate;
25

26
27
28
29
30
// Re-export from dynamo-kv-router crate
pub use dynamo_kv_router::approx;
pub use dynamo_kv_router::indexer;
pub use dynamo_kv_router::protocols;

31
pub mod cache_control;
32
pub mod config;
33
pub mod indexer_standalone;
34
mod jetstream;
35
pub mod metrics;
36
pub mod prefill_router;
37
pub mod publisher;
38
pub mod push_router;
39
pub mod queue;
40
pub mod recorder;
41
pub mod scheduler;
42
pub mod sequence;
43
pub mod subscriber;
44
pub mod worker_query;
45

46
pub use cache_control::{CacheControlClient, spawn_pin_prefix};
47
pub use config::{KvRouterConfig, RouterConfigOverride};
48
pub use indexer_standalone::start_kv_block_indexer;
49
pub use prefill_router::PrefillRouter;
50
pub use push_router::{DirectRoutingRouter, KvPushRouter};
51

52
use crate::{
53
    discovery::RuntimeConfigWatch,
54
    kv_router::{
55
        approx::PruneConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
56
        indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface, KvRouterError},
Yan Ru Pei's avatar
Yan Ru Pei committed
57
        protocols::{
58
59
            BlockExtraInfo, DpRank, LocalBlockHash, OverlapScores, RouterEvent, RouterRequest,
            RouterResponse, TokensWithHashes, WorkerId, WorkerSelectionResult, WorkerWithDpRank,
Yan Ru Pei's avatar
Yan Ru Pei committed
60
            compute_block_hash_for_seq,
Yan Ru Pei's avatar
Yan Ru Pei committed
61
        },
62
        scheduler::{KvScheduler, KvSchedulerError, PotentialLoad, SchedulingRequest},
63
        sequence::{SequenceError, SequenceRequest},
64
    },
65
    local_model::runtime_config::ModelRuntimeConfig,
66
67
};

68
69
use std::collections::HashSet;

70
71
// [gluo TODO] shouldn't need to be public
// this should be discovered from the component
72
73
74
75
76

// for metric scraping (pull-based)
pub const KV_METRICS_ENDPOINT: &str = "load_metrics";

// for metric publishing (push-based)
77
pub const KV_EVENT_SUBJECT: &str = "kv-events";
78
79
80
81
82
pub const KV_METRICS_SUBJECT: &str = "kv_metrics";

// for inter-router comms
pub const PREFILL_SUBJECT: &str = "prefill_events";
pub const ACTIVE_SEQUENCES_SUBJECT: &str = "active_sequences_events";
83

84
85
86
87
// for radix tree snapshot storage
pub const RADIX_STATE_BUCKET: &str = "radix-bucket";
pub const RADIX_STATE_FILE: &str = "radix-state";

88
89
90
// for standalone indexer query
pub const KV_INDEXER_QUERY_ENDPOINT: &str = "kv_indexer_query";

91
92
93
// for worker-local kvindexer query
pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer

94
95
96
97
98
99
/// Generates a dp_rank-specific endpoint name for the worker KV indexer query service.
/// Each dp_rank has its own LocalKvIndexer and query endpoint to ensure per-dp_rank monotonicity.
pub fn worker_kv_indexer_query_endpoint(dp_rank: DpRank) -> String {
    format!("worker_kv_indexer_query_dp{dp_rank}")
}

100
// for router discovery registration
101
pub const KV_ROUTER_ENDPOINT: &str = "router-discovery";
102
103

/// Creates an EndpointId for the KV router in the given namespace.
104
pub fn router_endpoint_id(namespace: String, component: String) -> EndpointId {
105
106
    EndpointId {
        namespace,
107
        component,
108
109
110
111
112
        name: KV_ROUTER_ENDPOINT.to_string(),
    }
}

/// Creates a DiscoveryQuery for the KV router in the given namespace.
113
pub fn router_discovery_query(namespace: String, component: String) -> DiscoveryQuery {
114
115
    DiscoveryQuery::Endpoint {
        namespace,
116
        component,
117
118
119
120
        endpoint: KV_ROUTER_ENDPOINT.to_string(),
    }
}

121
122
123
124
/// A trait that users can implement to define custom selection logic
pub trait WorkerSelector {
    fn select_worker(
        &self,
125
        workers: &HashMap<protocols::WorkerId, ModelRuntimeConfig>,
126
        request: &SchedulingRequest,
127
        block_size: u32,
128
129
    ) -> Result<WorkerSelectionResult, KvSchedulerError>;
}
130

Yan Ru Pei's avatar
Yan Ru Pei committed
131
#[derive(Clone)]
132
pub enum Indexer {
Yan Ru Pei's avatar
Yan Ru Pei committed
133
    /// Single-threaded radix tree with channel-based event processing.
134
    /// Supports TTL-based expiration and size-based pruning.
135
    /// Has the ability to persist and snapshot states.
136
    KvIndexer(KvIndexer),
137

Yan Ru Pei's avatar
Yan Ru Pei committed
138
139
140
141
142
    /// Concurrent radix tree with a thread pool for event processing.
    /// Uses sticky worker routing for per-worker event serialization.
    /// Does not support TTL/pruning.
    Concurrent(Arc<ThreadPoolIndexer<ConcurrentRadixTree>>),

143
144
145
    /// Used when we do not wish to use the indexer at all (e.g., when overlap_score_weight is 0).
    /// Note: This will cause KV events to accumulate in JetStream as we do not regularly purge them.
    None,
146
147
148
}

impl Indexer {
149
150
151
152
153
154
    pub fn new(
        component: &dynamo_runtime::component::Component,
        kv_router_config: &KvRouterConfig,
        block_size: u32,
    ) -> Self {
        if kv_router_config.overlap_score_weight == 0.0 {
Yan Ru Pei's avatar
Yan Ru Pei committed
155
156
157
            return Indexer::None;
        }

158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
        // Approximate mode (--no-kv-events): always use single-threaded KvIndexer
        // with TTL/pruning regardless of event_threads, since updates come from
        // routing decisions only, not live KV events from workers.
        if !kv_router_config.use_kv_events {
            let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);
            let cancellation_token = component.drt().primary_token();
            let prune_config = Some(PruneConfig {
                ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs),
                max_tree_size: kv_router_config.router_max_tree_size,
                prune_target_ratio: kv_router_config.router_prune_target_ratio,
            });
            return Indexer::KvIndexer(KvIndexer::new_with_frequency(
                cancellation_token,
                None,
                block_size,
                kv_indexer_metrics,
                prune_config,
            ));
        }

Yan Ru Pei's avatar
Yan Ru Pei committed
178
179
180
181
        if kv_router_config.router_event_threads > 1 {
            return Indexer::Concurrent(Arc::new(ThreadPoolIndexer::new(
                ConcurrentRadixTree::new(),
                kv_router_config.router_event_threads as usize,
182
                block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
183
            )));
184
        }
Yan Ru Pei's avatar
Yan Ru Pei committed
185
186

        let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);
187
        let cancellation_token = component.drt().primary_token();
Yan Ru Pei's avatar
Yan Ru Pei committed
188
189
190
191
192
193

        Indexer::KvIndexer(KvIndexer::new_with_frequency(
            cancellation_token,
            None, // expiration_duration for frequency tracking
            block_size,
            kv_indexer_metrics,
194
            None,
Yan Ru Pei's avatar
Yan Ru Pei committed
195
        ))
196
197
198
    }

    pub(crate) async fn find_matches(
199
200
201
202
203
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await,
Yan Ru Pei's avatar
Yan Ru Pei committed
204
            Indexer::Concurrent(tpi) => tpi.find_matches(sequence).await,
205
            Indexer::None => Ok(OverlapScores::new()),
206
207
        }
    }
208

209
    pub(crate) async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
210
211
        match self {
            Indexer::KvIndexer(indexer) => indexer.dump_events().await,
Yan Ru Pei's avatar
Yan Ru Pei committed
212
            Indexer::Concurrent(tpi) => tpi.dump_events().await,
213
214
215
216
217
            Indexer::None => {
                panic!(
                    "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)"
                );
            }
218
219
        }
    }
220

221
    pub(crate) async fn process_routing_decision_for_request(
222
        &self,
223
        tokens_with_hashes: &mut TokensWithHashes,
224
225
226
227
228
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => {
                indexer
229
                    .process_routing_decision_for_request(tokens_with_hashes, worker)
230
231
                    .await
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
232
233
234
235
            Indexer::Concurrent(tpi) => {
                tpi.process_routing_decision_for_request(tokens_with_hashes, worker)
                    .await
            }
236
237
238
            Indexer::None => Ok(()),
        }
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280

    pub(crate) async fn apply_event(&self, event: RouterEvent) {
        match self {
            Indexer::KvIndexer(indexer) => {
                if let Err(e) = indexer.event_sender().send(event).await {
                    tracing::warn!("Failed to send event to indexer: {e}");
                }
            }
            Indexer::Concurrent(tpi) => tpi.apply_event(event).await,
            Indexer::None => {}
        }
    }

    pub(crate) async fn remove_worker(&self, worker_id: WorkerId) {
        match self {
            Indexer::KvIndexer(indexer) => {
                if let Err(e) = indexer.remove_worker_sender().send(worker_id).await {
                    tracing::warn!("Failed to send worker removal for {worker_id}: {e}");
                }
            }
            Indexer::Concurrent(tpi) => {
                KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await;
            }
            Indexer::None => {}
        }
    }

    pub(crate) async fn get_workers(&self) -> Vec<WorkerId> {
        match self {
            Indexer::KvIndexer(indexer) => {
                let (resp_tx, resp_rx) = oneshot::channel();
                let req = GetWorkersRequest { resp: resp_tx };
                if let Err(e) = indexer.get_workers_sender().send(req).await {
                    tracing::warn!("Failed to send get_workers request: {e}");
                    return Vec::new();
                }
                resp_rx.await.unwrap_or_default()
            }
            Indexer::Concurrent(tpi) => tpi.backend().get_workers(),
            Indexer::None => Vec::new(),
        }
    }
281
282
}

283
284
/// A KvRouter only decides which worker you should use. It doesn't send you there.
/// TODO: Rename this to indicate it only selects a worker, it does not route.
285
pub struct KvRouter {
286
    indexer: Indexer,
287
    scheduler: KvScheduler,
288
    block_size: u32,
289
    kv_router_config: KvRouterConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
290
    cancellation_token: tokio_util::sync::CancellationToken,
291
    client: Client,
292
293
294
295
}

impl KvRouter {
    pub async fn new(
296
297
        endpoint: Endpoint,
        client: Client,
298
        mut workers_with_configs: RuntimeConfigWatch,
299
        block_size: u32,
300
        selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
301
        kv_router_config: Option<KvRouterConfig>,
302
        worker_type: &'static str,
303
    ) -> Result<Self> {
304
        let kv_router_config = kv_router_config.unwrap_or_default();
305
        kv_router_config.validate()?;
306
        let component = endpoint.component();
307
        let cancellation_token = component.drt().primary_token();
308

309
        let indexer = Indexer::new(component, &kv_router_config, block_size);
310

311
        // Wait for at least one worker with a known runtime config before starting scheduler
312
313
314
315
316
317
        let _ = workers_with_configs
            .wait_for(|m| !m.is_empty())
            .await
            .map_err(|_| {
                anyhow::anyhow!("runtime config watch closed before any workers appeared")
            })?;
318

319
        let scheduler = KvScheduler::start(
320
            component.clone(),
321
            block_size,
322
            workers_with_configs.clone(),
323
            selector,
324
            &kv_router_config,
325
            worker_type,
326
327
        )
        .await?;
328

329
330
        // Start KV event subscription if needed (use_kv_events=true and overlap_score_weight>0)
        if kv_router_config.should_subscribe_to_kv_events() {
331
332
            subscriber::start_subscriber(component.clone(), &kv_router_config, indexer.clone())
                .await?;
333
        } else {
334
            tracing::info!(
335
336
337
                "Skipping KV event subscription (use_kv_events={}, overlap_score_weight={})",
                kv_router_config.use_kv_events,
                kv_router_config.overlap_score_weight,
338
            );
339
        }
340

341
        tracing::info!("KV Routing initialized");
342
        Ok(Self {
343
            indexer,
344
            scheduler,
345
            block_size,
346
            kv_router_config,
Yan Ru Pei's avatar
Yan Ru Pei committed
347
            cancellation_token,
348
            client,
349
        })
350
351
    }

352
353
354
355
356
    /// Get a reference to the client used by this KvRouter
    pub fn client(&self) -> &Client {
        &self.client
    }

357
358
359
360
361
362
363
364
    pub fn indexer(&self) -> &Indexer {
        &self.indexer
    }

    pub fn kv_router_config(&self) -> &KvRouterConfig {
        &self.kv_router_config
    }

365
    /// Give these tokens, find the worker with the best match in it's KV cache.
Yan Ru Pei's avatar
Yan Ru Pei committed
366
    /// Returns the best worker (with dp_rank) and overlap amount in number of blocks.
367
368
369
    /// Now also takes optional context_id for request tracking.
    ///
    /// When `allowed_worker_ids` is Some, only workers in that set are considered for selection.
370
    #[allow(clippy::too_many_arguments)]
Yan Ru Pei's avatar
Yan Ru Pei committed
371
    pub async fn find_best_match(
372
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
373
        context_id: Option<&str>,
374
        tokens: &[u32],
375
        block_mm_infos: Option<&[Option<BlockExtraInfo>]>,
376
        router_config_override: Option<&RouterConfigOverride>,
377
        update_states: bool,
378
        lora_name: Option<String>,
379
        priority_jump: f64,
380
        allowed_worker_ids: Option<HashSet<WorkerId>>,
Yan Ru Pei's avatar
Yan Ru Pei committed
381
    ) -> anyhow::Result<(WorkerWithDpRank, u32)> {
382
383
        let start = Instant::now();

Yan Ru Pei's avatar
Yan Ru Pei committed
384
        if update_states && context_id.is_none() {
385
            anyhow::bail!("context_id must be provided when update_states is true");
Yan Ru Pei's avatar
Yan Ru Pei committed
386
387
        }

388
        let isl_tokens = tokens.len();
389

390
391
392
393
394
395
396
397
        let block_hashes = tracing::info_span!("kv_router.compute_block_hashes").in_scope(|| {
            compute_block_hash_for_seq(
                tokens,
                self.block_size,
                block_mm_infos,
                lora_name.as_deref(),
            )
        });
398
        let hash_elapsed = start.elapsed();
399

400
        let overlap_scores = self
401
402
403
404
            .indexer
            .find_matches(block_hashes)
            .instrument(tracing::info_span!("kv_router.find_matches"))
            .await?;
405
        let find_matches_elapsed = start.elapsed();
406

407
        // Compute seq_hashes only if scheduler needs it for active blocks tracking
408
409
410
411
412
        let maybe_seq_hashes = tracing::info_span!("kv_router.compute_seq_hashes").in_scope(|| {
            self.kv_router_config.compute_seq_hashes_for_tracking(
                tokens,
                self.block_size,
                router_config_override,
413
                lora_name.as_deref(),
414
415
            )
        });
416
        let seq_hash_elapsed = start.elapsed();
417

418
        let response = self
419
            .scheduler
420
            .schedule(
Yan Ru Pei's avatar
Yan Ru Pei committed
421
                context_id.map(|s| s.to_string()),
422
                isl_tokens,
423
                maybe_seq_hashes,
424
                overlap_scores,
425
                router_config_override,
426
                update_states,
427
                lora_name,
428
                priority_jump,
429
                allowed_worker_ids,
430
            )
431
            .instrument(tracing::info_span!("kv_router.schedule"))
432
            .await?;
433
434
        let total_elapsed = start.elapsed();

435
436
437
438
439
440
441
442
        if let Some(m) = metrics::RoutingOverheadMetrics::get() {
            m.observe(
                hash_elapsed,
                find_matches_elapsed,
                seq_hash_elapsed,
                total_elapsed,
            );
        }
443

444
        #[cfg(feature = "bench")]
445
446
447
448
449
450
451
452
453
        tracing::info!(
            isl_tokens,
            hash_us = hash_elapsed.as_micros() as u64,
            find_matches_us = (find_matches_elapsed - hash_elapsed).as_micros() as u64,
            seq_hash_us = (seq_hash_elapsed - find_matches_elapsed).as_micros() as u64,
            schedule_us = (total_elapsed - seq_hash_elapsed).as_micros() as u64,
            total_us = total_elapsed.as_micros() as u64,
            "find_best_match completed"
        );
454

455
        Ok((response.best_worker, response.overlap_blocks))
456
457
    }

458
    #[allow(clippy::too_many_arguments)]
459
460
461
462
463
    pub async fn add_request(
        &self,
        request_id: String,
        tokens: &[u32],
        overlap_blocks: u32,
464
        expected_output_tokens: Option<u32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
465
        worker: WorkerWithDpRank,
466
        lora_name: Option<String>,
467
        router_config_override: Option<&RouterConfigOverride>,
468
469
    ) {
        let isl_tokens = tokens.len();
470

471
472
473
474
        let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking(
            tokens,
            self.block_size,
            router_config_override,
475
            lora_name.as_deref(),
476
        );
477

478
479
        if let Err(e) = self
            .scheduler
480
481
482
483
484
            .add_request(SequenceRequest {
                request_id: request_id.clone(),
                token_sequence: maybe_seq_hashes,
                isl: isl_tokens,
                overlap: overlap_blocks,
485
                expected_output_tokens,
Yan Ru Pei's avatar
Yan Ru Pei committed
486
                worker,
487
                lora_name,
488
            })
489
490
491
492
            .await
        {
            tracing::warn!("Failed to add request {request_id}: {e}");
        }
493
494
    }

495
    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> {
496
        self.scheduler.mark_prefill_completed(request_id).await
497
498
    }

499
    pub async fn free(&self, request_id: &str) -> Result<(), SequenceError> {
500
        self.scheduler.free(request_id).await
501
    }
502

503
504
505
506
507
508
    /// Get the worker type for this router ("prefill" or "decode").
    /// Used for Prometheus metric labeling.
    pub fn worker_type(&self) -> &'static str {
        self.scheduler.worker_type()
    }

509
    pub fn add_output_block(
510
511
512
513
        &self,
        request_id: &str,
        decay_fraction: Option<f64>,
    ) -> Result<(), SequenceError> {
514
        self.scheduler.add_output_block(request_id, decay_fraction)
515
516
    }

517
    pub fn block_size(&self) -> u32 {
518
519
        self.block_size
    }
520

521
522
523
524
525
526
    /// Compute the overlap blocks for a given token sequence and worker.
    /// This queries the indexer to find how many blocks are already cached.
    pub async fn get_overlap_blocks(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
527
        lora_name: Option<&str>,
528
    ) -> Result<u32, KvRouterError> {
529
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None, lora_name);
530
531
532
533
        let overlap_scores = self.indexer.find_matches(block_hashes).await?;
        Ok(overlap_scores.scores.get(&worker).copied().unwrap_or(0))
    }

534
    /// Get potential prefill and decode loads for all workers
535
536
537
538
    pub async fn get_potential_loads(
        &self,
        tokens: &[u32],
        router_config_override: Option<&RouterConfigOverride>,
539
        lora_name: Option<&str>,
540
    ) -> Result<Vec<PotentialLoad>> {
541
        let isl_tokens = tokens.len();
542
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None, lora_name);
543
        let overlap_scores = self.indexer.find_matches(block_hashes.clone()).await?;
544

545
546
547
548
        let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking(
            tokens,
            self.block_size,
            router_config_override,
549
            lora_name,
550
        );
551

552
553
        Ok(self
            .scheduler
554
            .get_potential_loads(maybe_seq_hashes, isl_tokens, overlap_scores))
555
556
    }

557
558
559
560
    /// Dump all events from the indexer
    pub async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        self.indexer.dump_events().await
    }
561
562
}

Michael Feil's avatar
Michael Feil committed
563
564
// NOTE: KVRouter works like a PushRouter,
// but without the reverse proxy functionality, but based on contract of 3 request types
565
566
567
568
569
570
571
#[async_trait]
impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Error> for KvRouter {
    async fn generate(
        &self,
        request: SingleIn<RouterRequest>,
    ) -> Result<ManyOut<Annotated<RouterResponse>>> {
        let (request, ctx) = request.into_parts();
Michael Feil's avatar
Michael Feil committed
572
573
574
        let context_id = ctx.context().id().to_string();
        // Handle different request types
        let response = match request {
575
576
577
578
            RouterRequest::New {
                tokens,
                block_mm_infos,
            } => {
Yan Ru Pei's avatar
Yan Ru Pei committed
579
                let (best_worker, overlap_blocks) = self
580
581
582
583
584
585
586
587
                    .find_best_match(
                        Some(&context_id),
                        &tokens,
                        block_mm_infos.as_deref(),
                        None,
                        true,
                        None,
                        0.0,
588
                        None,
589
                    )
Michael Feil's avatar
Michael Feil committed
590
591
592
                    .await?;

                RouterResponse::New {
Yan Ru Pei's avatar
Yan Ru Pei committed
593
594
                    worker_id: best_worker.worker_id,
                    dp_rank: best_worker.dp_rank,
Michael Feil's avatar
Michael Feil committed
595
596
597
                    overlap_blocks,
                }
            }
598
599
600
601
602
603
            RouterRequest::MarkPrefill => RouterResponse::PrefillMarked {
                success: self.mark_prefill_completed(&context_id).await.is_ok(),
            },
            RouterRequest::MarkFree => RouterResponse::FreeMarked {
                success: self.free(&context_id).await.is_ok(),
            },
Michael Feil's avatar
Michael Feil committed
604
        };
605
606
607
608
609
610

        let response = Annotated::from_data(response);
        let stream = stream::iter(vec![response]);
        Ok(ResponseStream::new(Box::pin(stream), ctx.context()))
    }
}
611

Yan Ru Pei's avatar
Yan Ru Pei committed
612
613
614
615
616
617
impl Drop for KvRouter {
    fn drop(&mut self) {
        tracing::info!("Dropping KvRouter - cancelling background tasks");
        self.cancellation_token.cancel();
    }
}