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

4
use std::sync::Arc;
5
use std::time::Duration;
6

7
use anyhow::Result;
8
use dynamo_runtime::{
9
    component::{Component, InstanceSource},
10
    pipeline::{
11
12
        async_trait, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter,
        ResponseStream, SingleIn,
13
14
15
16
17
    },
    prelude::*,
    protocols::annotated::Annotated,
};
use futures::stream::{self, StreamExt};
18
use tokio::sync::Mutex;
19

20
pub mod approx;
21
pub mod indexer;
22
pub mod metrics_aggregator;
23
24
pub mod protocols;
pub mod publisher;
25
pub mod recorder;
26
27
pub mod scheduler;
pub mod scoring;
28
pub mod sequence;
29

30
31
use crate::{
    kv_router::{
32
33
        approx::ApproxKvIndexer,
        indexer::{
34
35
            compute_block_hash_for_seq, compute_seq_hash_for_block, KvIndexer, KvIndexerInterface,
            KvRouterError, OverlapScores, RouterEvent,
36
        },
37
        // metrics_aggregator::EndpointCollector,
38
39
40
41
        protocols::{LocalBlockHash, RouterRequest, RouterResponse, WorkerSelectionResult},
        scheduler::{KvScheduler, KvSchedulerError, SchedulingRequest},
        scoring::ProcessedEndpoints,
    },
42
    preprocessor::PreprocessedRequest,
43
    protocols::common::llm_backend::LLMEngineOutput,
44
45
};

46
use dynamo_runtime::component::Instance;
47
use dynamo_runtime::traits::events::EventSubscriber;
48
49
50

// [gluo TODO] shouldn't need to be public
// this should be discovered from the component
51
pub const KV_EVENT_SUBJECT: &str = "kv_events";
52
pub const KV_HIT_RATE_SUBJECT: &str = "kv-hit-rate";
53
pub const KV_METRICS_ENDPOINT: &str = "load_metrics";
54

55
56
57
58
/// A trait that users can implement to define custom selection logic
pub trait WorkerSelector {
    fn select_worker(
        &self,
59
        workers: &[Instance],
60
        request: &SchedulingRequest,
61
        block_size: u32,
62
63
    ) -> Result<WorkerSelectionResult, KvSchedulerError>;
}
64

65
/// KV Router configuration parameters
66
#[derive(Debug, Clone, Copy)]
67
68
69
pub struct KvRouterConfig {
    pub overlap_score_weight: f64,

70
    pub router_temperature: f64,
71

72
73
    pub use_kv_events: bool,

74
75
    // TODO: this is not actually used for now
    // Would need this (along with total kv blocks) to trigger AllWorkersBusy error for e.g. rate-limiting
76
    pub max_num_batched_tokens: u32,
77
78
79
80
81
}

impl Default for KvRouterConfig {
    fn default() -> Self {
        Self {
82
            overlap_score_weight: 1.0,
83
            router_temperature: 0.0,
84
            use_kv_events: true,
85
            max_num_batched_tokens: 8192,
86
87
88
89
90
91
92
93
94
        }
    }
}

impl KvRouterConfig {
    /// Create a new KvRouterConfig with optional weight values.
    /// If a weight is None, the default value will be used.
    pub fn new(
        overlap_score_weight: Option<f64>,
95
        temperature: Option<f64>,
96
        use_kv_events: Option<bool>,
97
        max_num_batched_tokens: Option<u32>,
98
99
100
101
    ) -> Self {
        let default = Self::default();
        Self {
            overlap_score_weight: overlap_score_weight.unwrap_or(default.overlap_score_weight),
102
            router_temperature: temperature.unwrap_or(default.router_temperature),
103
            use_kv_events: use_kv_events.unwrap_or(default.use_kv_events),
104
105
            max_num_batched_tokens: max_num_batched_tokens
                .unwrap_or(default.max_num_batched_tokens),
106
107
108
109
        }
    }
}

110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
// TODO: is there a way (macro) to auto-derive the KvIndexerInterface trait for this
// since both variants implement it
pub enum Indexer {
    KvIndexer(KvIndexer),
    ApproxKvIndexer(ApproxKvIndexer),
}

impl Indexer {
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await,
            Indexer::ApproxKvIndexer(indexer) => indexer.find_matches(sequence).await,
        }
    }
}

129
130
/// 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.
131
pub struct KvRouter {
132
133
134
    indexer: Indexer,

    // How about a Box<dyn KvIndexerInterface>
135
    scheduler: KvScheduler,
136

137
    block_size: u32,
138
139
140
141

    // To ensure blocking reads / writes
    // TODO: benchmark tradeoffs
    find_best_match_mutex: Mutex<()>,
142
143
144
145
}

impl KvRouter {
    pub async fn new(
146
        component: Component,
147
        block_size: u32,
148
        selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
149
        use_kv_events: bool,
150
    ) -> Result<Self> {
151
152
153
154
155
        let cancellation_token = component
            .drt()
            .primary_lease()
            .expect("Cannot KV route static workers")
            .primary_token();
156
157
158
159
160
161
162
163
164
165

        let generate_endpoint = component.endpoint("generate");
        let client = generate_endpoint.client().await?;

        let instances_rx = match client.instance_source.as_ref() {
            InstanceSource::Dynamic(rx) => rx.clone(),
            InstanceSource::Static => {
                panic!("Expected dynamic instance source for KV routing");
            }
        };
166

167
168
169
170
171
172
173
174
175
176
        let indexer = if use_kv_events {
            Indexer::KvIndexer(KvIndexer::new(cancellation_token.clone(), block_size))
        } else {
            // hard code 120 seconds for now
            Indexer::ApproxKvIndexer(ApproxKvIndexer::new(
                cancellation_token.clone(),
                block_size,
                Duration::from_secs(120),
            ))
        };
177

178
179
180
        let scheduler = KvScheduler::start(
            component.namespace().clone(),
            block_size,
181
            instances_rx,
182
183
184
            selector,
        )
        .await?;
185

186
187
        // [gluo TODO] try subscribe_with_type::<RouterEvent>,
        // error checking below will be different.
188
        if let Indexer::KvIndexer(ref kv_indexer) = indexer {
189
            let mut kv_events_rx = component.subscribe(KV_EVENT_SUBJECT).await?;
190
            let kv_events_tx = kv_indexer.event_sender();
191
192
193
194
195
196
197
198
199
200
201
202
203

            tokio::spawn(async move {
                while let Some(event) = kv_events_rx.next().await {
                    let event: RouterEvent = match serde_json::from_slice(&event.payload) {
                        Ok(event) => event,
                        Err(e) => {
                            tracing::warn!("Failed to deserialize RouterEvent: {:?}", e);
                            // Choosing warn and continue to process other events from other workers
                            // A bad event likely signals a problem with a worker, but potentially other workers are still healthy
                            continue;
                        }
                    };
                    if let Err(e) = kv_events_tx.send(event).await {
204
                        tracing::warn!(
205
206
207
                            "failed to send kv event to indexer; shutting down: {:?}",
                            e
                        );
Alec's avatar
Alec committed
208
                    }
209
                }
210
211
            });
        }
212

213
        tracing::info!("KV Routing initialized");
214
        Ok(Self {
215
            indexer,
216
            scheduler,
217
            block_size,
218
            find_best_match_mutex: Mutex::new(()), // Add this
219
        })
220
221
    }

222
    /// Give these tokens, find the worker with the best match in it's KV cache.
223
    /// Returned overlap amount is in number of blocks.
224
225
226
227
228
229
    /// Now also takes context_id for request tracking
    async fn find_best_match(
        &self,
        context_id: &str,
        tokens: &[u32],
    ) -> anyhow::Result<(i64, u32)> {
230
231
232
233
        // Acquire mutex to serialize access
        // TODO: may as well make all the subroutines synchronous if benchmarking favors this
        let _guard = self.find_best_match_mutex.lock().await;

234
        let isl_tokens = tokens.len();
235

236
237
238
239
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size);
        let seq_hashes = compute_seq_hash_for_block(&block_hashes);

        let overlap_scores = self.indexer.find_matches(block_hashes.clone()).await?;
240
241

        let best_worker_id = self
242
            .scheduler
243
244
245
            .schedule(
                context_id.to_string(),
                isl_tokens,
246
                seq_hashes.clone(),
247
248
                overlap_scores.clone(),
            )
249
            .await?;
250

251
252
        if let Indexer::ApproxKvIndexer(ref indexer) = self.indexer {
            indexer
253
                .process_routing_decision(best_worker_id, block_hashes, seq_hashes)
254
255
256
257
                .await
                .unwrap();
        };

258
259
260
261
262
263
264
265
        let overlap_amount = overlap_scores
            .scores
            .get(&best_worker_id)
            .copied()
            .unwrap_or(0);
        Ok((best_worker_id, overlap_amount))
    }

266
267
268
    /// Free all blocks associated with a request
    pub async fn mark_prefill_completed(&self, request_id: &String) {
        self.scheduler.mark_prefill_completed(request_id).await
269
270
271
272
273
    }

    /// Free all blocks associated with a request
    pub async fn free(&self, request_id: &String) {
        self.scheduler.free(request_id).await
274
    }
275
276

    /// Get the block size this router was configured with
277
    pub fn block_size(&self) -> u32 {
278
279
        self.block_size
    }
280
281
}

282
// NOTE: this would not be usable for now, should deprecate
283
284
285
286
287
288
289
#[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();
290
        let (worker_id, _) = self.find_best_match(ctx.id(), &request.tokens).await?;
291
292
293
294
295
296
297

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

pub struct KvPushRouter {
300
    inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
301
302
303
304
305
    chooser: Arc<KvRouter>,
}

impl KvPushRouter {
    pub fn new(
306
        inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
307
308
309
310
311
312
313
        chooser: Arc<KvRouter>,
    ) -> Self {
        KvPushRouter { inner, chooser }
    }
}

#[async_trait]
314
impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutput>>, Error>
315
316
317
318
    for KvPushRouter
{
    async fn generate(
        &self,
319
        request: SingleIn<PreprocessedRequest>,
320
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
321
        match self.inner.client.instance_source.as_ref() {
322
323
            InstanceSource::Static => self.inner.r#static(request).await,
            InstanceSource::Dynamic(_) => {
324
325
326
327
328
329
                // Extract context ID for request tracking
                let context_id = request.context().id().to_string();
                let (instance_id, overlap_amount) = self
                    .chooser
                    .find_best_match(&context_id, &request.token_ids)
                    .await?;
330
331
332
                let query_instance_id = request.has_annotation("query_instance_id");
                // Extract context information before moving the request
                let stream_context = request.context().clone();
333
334
335
336
                // Update the request with the estimated prefix hit blocks
                let (mut backend_input, context) = request.into_parts();
                backend_input.estimated_prefix_hit_num_blocks = Some(overlap_amount);
                let updated_request = context.map(|_| backend_input);
337

338
339
340
341
342
343
344
345
                // if request has the annotation "query_instance_id", for example
                // curl -d '{... ,"nvext": { "annotations": ["query_instance_id"]}}'
                // request will not be routed to worker immediately
                if query_instance_id {
                    let instance_id_str = instance_id.to_string();
                    let response =
                        Annotated::from_annotation("worker_instance_id", &instance_id_str)?;
                    let stream = stream::iter(vec![response]);
346
347
348
                    return Ok(ResponseStream::new(Box::pin(stream), stream_context));
                }

349
                let mut response_stream = self.inner.direct(updated_request, instance_id).await?;
350
351
352
353
                let stream_context = response_stream.context();
                let chooser = self.chooser.clone();

                let wrapped_stream = Box::pin(async_stream::stream! {
354
355
356
357
                    if let Some(first_item) = response_stream.next().await {
                        chooser.mark_prefill_completed(&context_id).await;
                        yield first_item;
                    }
358
359
360
361
362

                    while let Some(item) = response_stream.next().await {
                        yield item;
                    }

363
                    chooser.free(&context_id).await;
364
365
                });
                Ok(ResponseStream::new(wrapped_stream, stream_context))
366
367
368
369
            }
        }
    }
}