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

4
use crate::local_model::runtime_config::ModelRuntimeConfig;
5
use anyhow::Result;
6
use dynamo_runtime::component::{Component, Instance};
Yan Ru Pei's avatar
Yan Ru Pei committed
7
use dynamo_runtime::traits::DistributedRuntimeProvider;
Neelay Shah's avatar
Neelay Shah committed
8
use dynamo_runtime::traits::events::EventPublisher;
9
use rand::Rng;
10
use serde::{Deserialize, Serialize};
11
use std::collections::{HashMap, HashSet};
12
13
use std::sync::Arc;
use std::time::Duration;
Yan Ru Pei's avatar
Yan Ru Pei committed
14
use tokio::sync::{RwLock, watch};
15

16
17
use super::KV_HIT_RATE_SUBJECT;
use super::KvRouterConfig;
18
use super::RouterConfigOverride;
19
use super::WorkerSelector;
20
use super::indexer::OverlapScores;
Yan Ru Pei's avatar
Yan Ru Pei committed
21
use super::protocols::{DpRank, WorkerId, WorkerSelectionResult, WorkerWithDpRank};
22
23
use super::sequence::ActiveSequencesMultiWorker;

24
use crate::tokens::SequenceHash;
25

26
27
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KVHitRateEvent {
Yan Ru Pei's avatar
Yan Ru Pei committed
28
29
30
    pub worker_id: WorkerId,
    #[serde(default)]
    pub dp_rank: DpRank,
31
    pub isl_blocks: usize,
32
    pub overlap_blocks: u32,
33
34
}

35
36
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PotentialLoad {
Yan Ru Pei's avatar
Yan Ru Pei committed
37
38
    pub worker_id: WorkerId,
    pub dp_rank: DpRank,
39
40
41
42
    pub potential_prefill_tokens: usize,
    pub potential_decode_blocks: usize,
}

43
44
#[derive(Debug, thiserror::Error)]
pub enum KvSchedulerError {
45
    #[error("no endpoints available to route work")]
46
47
48
49
50
51
52
53
54
    NoEndpoints,

    #[error("all workers busy")]
    AllWorkersBusy,

    #[error("endpoint subscriber shutdown")]
    SubscriberShutdown,
}

55
56
#[derive(Debug)]
pub struct SchedulingResponse {
Yan Ru Pei's avatar
Yan Ru Pei committed
57
    pub best_worker: WorkerWithDpRank,
58
    pub overlap_blocks: u32,
59
60
}

61
pub struct SchedulingRequest {
Yan Ru Pei's avatar
Yan Ru Pei committed
62
    pub maybe_request_id: Option<String>,
63
    pub token_seq: Option<Vec<SequenceHash>>,
64
    pub isl_tokens: usize,
65
    pub overlaps: OverlapScores,
Yan Ru Pei's avatar
Yan Ru Pei committed
66
67
    pub decode_blocks: HashMap<WorkerWithDpRank, usize>,
    pub prefill_tokens: HashMap<WorkerWithDpRank, usize>,
68
69
    // Router config overrides for this specific request
    pub router_config_override: Option<RouterConfigOverride>,
70
71
    // Whether to update scheduler states (false for query_instance_id requests)
    pub update_states: bool,
72
73
    // Option to take it out to send the response without moving the struct
    resp_tx: Option<tokio::sync::oneshot::Sender<SchedulingResponse>>,
74
75
76
}

impl SchedulingRequest {
77
78
79
80
81
82
83
84
85
    pub fn respond(&mut self, response: SchedulingResponse) {
        // Changed to &mut self
        if let Some(tx) = self.resp_tx.take() {
            // Use take() to extract the sender
            if tx.send(response).is_err() {
                tracing::error!("failed to send response to requestor");
            }
        } else {
            tracing::error!("respond called multiple times on same request");
86
87
88
89
90
91
        }
    }
}

pub struct KvScheduler {
    request_tx: tokio::sync::mpsc::Sender<SchedulingRequest>,
92
    slots: Arc<ActiveSequencesMultiWorker>,
93
94
95
96
}

impl KvScheduler {
    pub async fn start(
97
        component: Component,
98
        block_size: u32,
Yan Ru Pei's avatar
Yan Ru Pei committed
99
        instances_rx: watch::Receiver<Vec<Instance>>,
Yan Ru Pei's avatar
Yan Ru Pei committed
100
        runtime_configs_rx: watch::Receiver<HashMap<WorkerId, ModelRuntimeConfig>>,
101
        selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
102
        replica_sync: bool,
103
        router_uuid: String,
104
    ) -> Result<Self, KvSchedulerError> {
105
        let selector = selector.unwrap_or(Box::new(DefaultWorkerSelector::default()));
Yan Ru Pei's avatar
Yan Ru Pei committed
106
        let instances: Vec<Instance> = instances_rx.borrow().clone();
Yan Ru Pei's avatar
Yan Ru Pei committed
107
108
        let runtime_configs: HashMap<WorkerId, ModelRuntimeConfig> =
            runtime_configs_rx.borrow().clone();
109

Yan Ru Pei's avatar
Yan Ru Pei committed
110
        // Create shared workers_with_configs wrapped in Arc<RwLock>
Yan Ru Pei's avatar
Yan Ru Pei committed
111
        let workers_with_configs: Arc<RwLock<HashMap<WorkerId, Option<ModelRuntimeConfig>>>> = {
Yan Ru Pei's avatar
Yan Ru Pei committed
112
113
114
115
116
117
            let mut initial_map = HashMap::new();
            for instance in &instances {
                let worker_id = instance.instance_id;
                let config = runtime_configs.get(&worker_id).cloned();
                if config.is_some() {
                    tracing::info!("Runtime config found for worker_id: {}", worker_id);
118
                }
Yan Ru Pei's avatar
Yan Ru Pei committed
119
                initial_map.insert(worker_id, config);
120
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
121
122
            Arc::new(RwLock::new(initial_map))
        };
123

124
        let slots = Arc::new(ActiveSequencesMultiWorker::new(
Yan Ru Pei's avatar
Yan Ru Pei committed
125
            component.clone(),
126
            block_size as usize,
Yan Ru Pei's avatar
Yan Ru Pei committed
127
            workers_with_configs.read().await.clone(), // this includes dp_size info
128
            replica_sync,
129
            router_uuid,
130
        ));
131

Yan Ru Pei's avatar
Yan Ru Pei committed
132
133
134
135
136
137
        // Spawn background task to monitor and update workers_with_configs
        let workers_monitor = workers_with_configs.clone();
        let slots_monitor = slots.clone();
        let mut instances_monitor_rx = instances_rx.clone();
        let mut configs_monitor_rx = runtime_configs_rx.clone();
        let monitor_cancel_token = component.drt().primary_token();
138
        tokio::spawn(async move {
Yan Ru Pei's avatar
Yan Ru Pei committed
139
            tracing::trace!("workers monitoring task started");
140
            loop {
Yan Ru Pei's avatar
Yan Ru Pei committed
141
142
143
144
                // Wait for either instances or configs to change
                tokio::select! {
                    _ = monitor_cancel_token.cancelled() => {
                        tracing::trace!("workers monitoring task shutting down");
145
146
                        break;
                    }
Yan Ru Pei's avatar
Yan Ru Pei committed
147
148
149
150
151
                    result = instances_monitor_rx.changed() => {
                        if result.is_err() {
                            tracing::warn!("endpoint watch sender shutdown in monitor");
                            break;
                        }
152
                    }
Yan Ru Pei's avatar
Yan Ru Pei committed
153
154
155
156
157
                    result = configs_monitor_rx.changed() => {
                        if result.is_err() {
                            tracing::warn!("runtime configs watch sender shutdown in monitor");
                            break;
                        }
158
159
160
                    }
                }

Yan Ru Pei's avatar
Yan Ru Pei committed
161
162
163
164
                // Get the latest values from both channels
                let new_instances = instances_monitor_rx.borrow_and_update().clone();
                let new_configs = configs_monitor_rx.borrow_and_update().clone();

Yan Ru Pei's avatar
Yan Ru Pei committed
165
166
                // Build the new workers_with_configs map
                let mut new_workers_with_configs = HashMap::new();
Yan Ru Pei's avatar
Yan Ru Pei committed
167
168
169
170
171
                for instance in &new_instances {
                    let worker_id = instance.instance_id;
                    let config = new_configs.get(&worker_id).cloned();
                    if config.is_some() {
                        tracing::info!("Runtime config found for worker_id: {}", worker_id);
172
                    }
Yan Ru Pei's avatar
Yan Ru Pei committed
173
                    new_workers_with_configs.insert(worker_id, config);
Yan Ru Pei's avatar
Yan Ru Pei committed
174
                }
Yan Ru Pei's avatar
Yan Ru Pei committed
175
176
177
178
179
180
181

                // Update workers when instances change
                slots_monitor.update_workers(new_workers_with_configs.clone());

                // Update the shared workers_with_configs
                let mut workers_map = workers_monitor.write().await;
                *workers_map = new_workers_with_configs;
Yan Ru Pei's avatar
Yan Ru Pei committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
                tracing::trace!(
                    "Updated workers_with_configs with {} workers",
                    workers_map.len()
                );
            }
            tracing::trace!("workers monitoring task shutting down");
        });

        let slots_clone = slots.clone();
        let workers_scheduler = workers_with_configs.clone();
        let (request_tx, request_rx) = tokio::sync::mpsc::channel::<SchedulingRequest>(1024);
        let scheduler_cancel_token = component.drt().primary_token();
        let ns_clone = component.namespace().clone();

        // Background task to handle scheduling requests
        tokio::spawn(async move {
            let mut request_rx = request_rx;
            tracing::trace!("scheduler background task started");

            loop {
                // Check for cancellation at beginning of loop
                if scheduler_cancel_token.is_cancelled() {
                    tracing::trace!("scheduler background task shutting down");
                    break;
206
207
208
                }

                // Wait for a new request
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
                let Some(mut request) = request_rx.recv().await else {
                    tracing::warn!("scheduler shutdown");
                    break;
                };
                tracing::trace!("received request to be scheduled");

                let (decode_blocks, prefill_tokens) = slots_clone
                    .potential_blocks_and_tokens(
                        request.token_seq.clone(),
                        request.isl_tokens,
                        request.overlaps.clone(),
                    )
                    .await;
                request.decode_blocks = decode_blocks;
                request.prefill_tokens = prefill_tokens;

Yan Ru Pei's avatar
Yan Ru Pei committed
225
226
227
228
                // Read the current workers configuration
                let workers = workers_scheduler.read().await.clone();

                match selector.select_worker(&workers, &request, block_size) {
229
                    Ok(selection) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
230
                        let event = KVHitRateEvent {
Yan Ru Pei's avatar
Yan Ru Pei committed
231
232
                            worker_id: selection.worker.worker_id,
                            dp_rank: selection.worker.dp_rank,
233
234
                            isl_blocks: selection.required_blocks as usize,
                            overlap_blocks: selection.overlap_blocks,
Yan Ru Pei's avatar
Yan Ru Pei committed
235
236
237
                        };
                        if let Err(e) = ns_clone.publish(KV_HIT_RATE_SUBJECT, &event).await {
                            tracing::warn!("Failed to publish KV hit rate event: {:?}", e);
238
                        }
239
240

                        let response = SchedulingResponse {
Yan Ru Pei's avatar
Yan Ru Pei committed
241
                            best_worker: selection.worker,
242
243
244
245
                            overlap_blocks: selection.overlap_blocks,
                        };
                        request.respond(response);

246
247
248
                        // Skip state update if not requested
                        if !request.update_states {
                            continue;
249
                        }
250

Yan Ru Pei's avatar
Yan Ru Pei committed
251
252
253
254
255
256
257
                        let Some(request_id) = request.maybe_request_id else {
                            tracing::error!(
                                "No request_id provided to add_request to the slot tracker"
                            );
                            continue;
                        };

258
259
260
261
262
263
                        if let Err(e) = slots_clone
                            .add_request(
                                request_id.clone(),
                                request.token_seq,
                                request.isl_tokens,
                                selection.overlap_blocks,
Yan Ru Pei's avatar
Yan Ru Pei committed
264
                                selection.worker,
265
266
267
268
269
270
271
                            )
                            .await
                        {
                            tracing::warn!(
                                "Failed to add request {request_id} to local slot tracker: {e:?}"
                            );
                        }
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
                    }
                    Err(KvSchedulerError::NoEndpoints) => {
                        tracing::trace!("no endpoints available; waiting for endpoints update");
                        tokio::time::sleep(Duration::from_millis(5)).await;
                        continue;
                    }
                    // TODO: this is not actually hooked up
                    Err(KvSchedulerError::AllWorkersBusy) => {
                        tracing::trace!("all workers busy; waiting for more capacity");
                        tokio::time::sleep(Duration::from_millis(5)).await;
                        continue;
                    }
                    Err(e) => {
                        tracing::error!("error scheduling request: {:?}", e);
                        break;
287
288
289
290
                    }
                }
            }

291
            tracing::trace!("background endpoint subscriber shutting down");
292
293
        });

294
        Ok(KvScheduler { request_tx, slots })
295
296
297
298
    }

    pub async fn schedule(
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
299
        maybe_request_id: Option<String>,
300
        isl_tokens: usize,
301
        token_seq: Option<Vec<SequenceHash>>,
302
        overlaps: OverlapScores,
303
        router_config_override: Option<&RouterConfigOverride>,
304
        update_states: bool,
Yan Ru Pei's avatar
Yan Ru Pei committed
305
    ) -> Result<WorkerWithDpRank, KvSchedulerError> {
306
307
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
        let request = SchedulingRequest {
Yan Ru Pei's avatar
Yan Ru Pei committed
308
            maybe_request_id,
309
            token_seq,
310
            isl_tokens,
311
            overlaps,
312
313
            decode_blocks: HashMap::new(),
            prefill_tokens: HashMap::new(),
314
            router_config_override: router_config_override.cloned(),
315
            update_states,
316
            resp_tx: Some(resp_tx), // Wrap in Some()
317
        };
318

319
320
321
322
        self.request_tx
            .send(request)
            .await
            .map_err(|_| KvSchedulerError::SubscriberShutdown)?;
323
        let response = resp_rx
324
325
            .await
            .map_err(|_| KvSchedulerError::SubscriberShutdown)?;
326

Yan Ru Pei's avatar
Yan Ru Pei committed
327
        Ok(response.best_worker)
328
329
    }

330
331
332
    pub async fn add_request(
        &self,
        request_id: String,
333
        token_sequence: Option<Vec<SequenceHash>>,
334
335
        isl: usize,
        overlap: u32,
Yan Ru Pei's avatar
Yan Ru Pei committed
336
        worker: WorkerWithDpRank,
337
338
339
    ) {
        let _ = self
            .slots
Yan Ru Pei's avatar
Yan Ru Pei committed
340
            .add_request(request_id, token_sequence, isl, overlap, worker)
341
342
343
            .await;
    }

344
345
    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<()> {
        self.slots
346
            .mark_prefill_completed(&request_id.to_string())
347
            .await
348
349
    }

350
351
    pub async fn free(&self, request_id: &str) -> Result<()> {
        self.slots.free(&request_id.to_string()).await
352
    }
353
354
355

    pub async fn get_potential_loads(
        &self,
356
        token_seq: Option<Vec<SequenceHash>>,
357
358
359
360
361
362
363
364
        isl_tokens: usize,
        overlaps: OverlapScores,
    ) -> Vec<PotentialLoad> {
        let (decode_blocks, prefill_tokens) = self
            .slots
            .potential_blocks_and_tokens(token_seq, isl_tokens, overlaps)
            .await;

Yan Ru Pei's avatar
Yan Ru Pei committed
365
366
367
368
        // Get all unique WorkerWithDpRank from both hashmaps
        let mut workers: HashSet<WorkerWithDpRank> = HashSet::new();
        workers.extend(decode_blocks.keys().copied());
        workers.extend(prefill_tokens.keys().copied());
369
370
371

        // Create PotentialLoad for each worker
        let mut loads = Vec::new();
Yan Ru Pei's avatar
Yan Ru Pei committed
372
        for worker in workers {
373
            loads.push(PotentialLoad {
Yan Ru Pei's avatar
Yan Ru Pei committed
374
375
                worker_id: worker.worker_id,
                dp_rank: worker.dp_rank,
376
                potential_prefill_tokens: prefill_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
377
                    .get(&worker)
378
379
                    .copied()
                    .unwrap_or(isl_tokens),
Yan Ru Pei's avatar
Yan Ru Pei committed
380
                potential_decode_blocks: decode_blocks.get(&worker).copied().unwrap_or(0),
381
382
383
384
385
            });
        }

        loads
    }
386
387
}

388
// Helper function for softmax sampling
389
390
391
392
393
// Returns a vec of workers: multiple if tied, single if sampled
fn softmax_sample(
    logits: &HashMap<WorkerWithDpRank, f64>,
    temperature: f64,
) -> Vec<WorkerWithDpRank> {
394
395
396
397
    if logits.is_empty() {
        panic!("Empty logits for softmax sampling");
    }

398
    // Guard: if temperature is 0, return all keys with the smallest logit value (ties)
399
400
401
402
403
404
405
    if temperature == 0.0 {
        // Find the minimum logit value
        let min_logit = logits.values().fold(f64::INFINITY, |a, &b| a.min(b));

        // Collect all keys with the minimum logit value (to handle ties)
        let min_keys: Vec<_> = logits
            .iter()
406
            .filter(|&(_, &v)| v == min_logit)
407
408
409
            .map(|(k, _)| *k)
            .collect();

410
        return min_keys;
411
412
    }

413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
    let keys: Vec<_> = logits.keys().copied().collect();
    let values: Vec<_> = logits.values().copied().collect();

    // Find min and max for normalization
    let min_val = values.iter().fold(f64::INFINITY, |a, &b| a.min(b));
    let max_val = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));

    let probabilities = if min_val == max_val {
        // All values are the same, uniform probability
        vec![1.0 / keys.len() as f64; keys.len()]
    } else {
        // Normalize values
        let normalized: Vec<_> = values
            .iter()
            .map(|&v| {
                // Lower is better, so negate
429
430
                // Note we don't need to do actual min-max norm here, just off by an offset
                let norm = v / (max_val - min_val);
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
                -norm
            })
            .collect();

        // Apply temperature and softmax
        let scaled: Vec<_> = normalized.iter().map(|&v| v / temperature).collect();

        let max_scaled = scaled.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
        let exp_values: Vec<_> = scaled.iter().map(|&v| (v - max_scaled).exp()).collect();

        let sum_exp: f64 = exp_values.iter().sum();
        exp_values.iter().map(|&v| v / sum_exp).collect()
    };

    // Sample from the probability distribution
    let mut rng = rand::rng();
    let sample: f64 = rng.random();

    let mut cumsum = 0.0;
    for (i, &prob) in probabilities.iter().enumerate() {
        cumsum += prob;
        if sample <= cumsum {
453
            return vec![keys[i]];
454
455
456
457
        }
    }

    // Fallback to last key (shouldn't normally reach here)
458
    vec![keys[keys.len() - 1]]
459
460
}

461
// Default implementation matching the Python _cost_function
462
463
464
465
466
467
468
469
470
471
472
473
#[derive(Debug, Clone, Default)]
pub struct DefaultWorkerSelector {
    pub kv_router_config: KvRouterConfig,
}

impl DefaultWorkerSelector {
    pub fn new(kv_router_config: Option<KvRouterConfig>) -> Self {
        Self {
            kv_router_config: kv_router_config.unwrap_or_default(),
        }
    }
}
474
475
476
477

impl WorkerSelector for DefaultWorkerSelector {
    fn select_worker(
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
478
        workers: &HashMap<WorkerId, Option<ModelRuntimeConfig>>,
479
        request: &SchedulingRequest,
480
        block_size: u32,
481
482
483
    ) -> Result<WorkerSelectionResult, KvSchedulerError> {
        assert!(request.isl_tokens > 0);

484
        if workers.is_empty() {
485
486
487
            return Err(KvSchedulerError::NoEndpoints);
        }

488
489
490
491
        let isl = request.isl_tokens;
        let request_blocks = isl.div_ceil(block_size as usize);
        let overlaps = &request.overlaps.scores;

492
493
        let decode_blocks = &request.decode_blocks;
        let prefill_tokens = &request.prefill_tokens;
494

495
        let mut worker_logits = HashMap::new();
496
        let mut max_logit = f64::NEG_INFINITY;
497

Yan Ru Pei's avatar
Yan Ru Pei committed
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
        // Calculate logits for each worker with dp_rank
        // Outer loop: iterate over all workers from runtime config
        // Inner loop: iterate over all dp_ranks for each worker
        for (worker_id, config) in workers.iter() {
            // Get data_parallel_size from runtime config
            // data_parallel_size defaults to 1 in ModelRuntimeConfig
            let data_parallel_size = config.as_ref().map(|c| c.data_parallel_size).unwrap_or(1); // Fallback if config is None

            // Iterate over all dp_ranks for this worker
            for dp_rank in 0..data_parallel_size {
                let worker = WorkerWithDpRank::new(*worker_id, dp_rank);

                // Get overlap for this worker (defaults to 0 if not in overlaps)
                let overlap = *overlaps.get(&worker).unwrap_or(&0);

                // this is the number of prefill tokens the worker would have if the request were scheduled there
                let prefill_token = *prefill_tokens.get(&worker).unwrap_or(&isl);
                let potential_prefill_block = (prefill_token as f64) / (block_size as f64);

                // this is the number of decode blocks the worker would have if the request were scheduled there
                let decode_block = *decode_blocks
                    .get(&worker)
                    .unwrap_or(&(potential_prefill_block.floor() as usize))
                    as f64;

                // Use override if provided, otherwise use default config
                let overlap_weight = request
                    .router_config_override
                    .as_ref()
                    .and_then(|cfg| cfg.overlap_score_weight)
                    .unwrap_or(self.kv_router_config.overlap_score_weight);

                // Calculate logit (lower is better)
                let logit = overlap_weight * potential_prefill_block + decode_block;
                max_logit = max_logit.max(logit);

                worker_logits.insert(worker, logit);

                tracing::info!(
                    "Formula for worker_id={} dp_rank={:?} with {overlap} cached blocks: {logit:.3} \
                     = {overlap_weight:.1} * prefill_blocks + decode_blocks \
                     = {overlap_weight:.1} * {potential_prefill_block:.3} + {decode_block:.3}",
                    worker.worker_id,
                    worker.dp_rank
                );
            }
544
545
        }

546
        // Use softmax sampling to select worker(s)
547
548
549
550
551
552
        // Use override if provided, otherwise use default config
        let temperature = request
            .router_config_override
            .as_ref()
            .and_then(|cfg| cfg.router_temperature)
            .unwrap_or(self.kv_router_config.router_temperature);
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
        let candidates = softmax_sample(&worker_logits, temperature);

        // If multiple candidates (tied), use tree size as tie-breaker
        // If tree sizes are also equal, min_by_key uses HashMap iteration order (pseudo-random)
        let best_worker = if candidates.len() > 1 {
            tracing::info!("Multiple workers tied with same logit, using tree size as tie-breaker");
            *candidates
                .iter()
                .min_by_key(|worker| {
                    request
                        .overlaps
                        .tree_sizes
                        .get(worker)
                        .copied()
                        .unwrap_or(0)
                })
                .expect("candidates should not be empty")
        } else {
            candidates[0]
        };

Yan Ru Pei's avatar
Yan Ru Pei committed
574
575
576
        let best_logit = worker_logits[&best_worker];

        let best_overlap = *overlaps.get(&best_worker).unwrap_or(&0);
577

Yan Ru Pei's avatar
Yan Ru Pei committed
578
        // this is a runtime config set on a per worker basis, not per dp-rank
579
        let total_blocks_info = workers
Yan Ru Pei's avatar
Yan Ru Pei committed
580
            .get(&best_worker.worker_id)
581
582
583
584
585
            .and_then(|cfg| cfg.as_ref())
            .and_then(|cfg| cfg.total_kv_blocks)
            .map(|blocks| format!(", total blocks: {}", blocks))
            .unwrap_or_default();

586
587
588
589
590
591
592
        let tree_size = request
            .overlaps
            .tree_sizes
            .get(&best_worker)
            .copied()
            .unwrap_or(0);

593
        tracing::info!(
594
            "Selected worker: worker_id={} dp_rank={:?}, logit: {:.3}, cached blocks: {}, tree size: {}{}",
Yan Ru Pei's avatar
Yan Ru Pei committed
595
596
            best_worker.worker_id,
            best_worker.dp_rank,
597
598
            best_logit,
            best_overlap,
599
            tree_size,
600
            total_blocks_info
601
        );
602
603

        Ok(WorkerSelectionResult {
Yan Ru Pei's avatar
Yan Ru Pei committed
604
            worker: best_worker,
605
            required_blocks: request_blocks as u64,
Yan Ru Pei's avatar
Yan Ru Pei committed
606
            overlap_blocks: overlaps.get(&best_worker).copied().unwrap_or(0),
607
        })
608
609
    }
}
610
611
612
613
614

#[cfg(test)]
mod tests {
    use super::*;

615
616
617
618
    #[test]
    fn test_softmax_sample_single_key() {
        // Test that with a single key, softmax_sample always returns that key
        let mut logits = HashMap::new();
Yan Ru Pei's avatar
Yan Ru Pei committed
619
620
        let worker = WorkerWithDpRank::from_worker_id(42);
        logits.insert(worker, 0.5); // The value doesn't matter
621
622
623
624

        // Test with different temperatures
        for temperature in &[0.1, 1.0, 10.0] {
            let result = softmax_sample(&logits, *temperature);
625
626
            assert_eq!(result.len(), 1, "Should return exactly one worker");
            assert_eq!(result[0], worker, "Should return the only available worker");
627
628
629
630
        }

        // Test with different logit values
        logits.clear();
Yan Ru Pei's avatar
Yan Ru Pei committed
631
        logits.insert(worker, -100.0); // Very negative value
632
633
634
        let result = softmax_sample(&logits, 1.0);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0], worker);
635
636

        logits.clear();
Yan Ru Pei's avatar
Yan Ru Pei committed
637
        logits.insert(worker, 100.0); // Very positive value
638
639
640
        let result = softmax_sample(&logits, 1.0);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0], worker);
641
642

        logits.clear();
Yan Ru Pei's avatar
Yan Ru Pei committed
643
        logits.insert(worker, 0.0); // Zero value
644
645
646
        let result = softmax_sample(&logits, 1.0);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0], worker);
647
648
    }

649
650
    #[test]
    fn test_softmax_sample_zero_temperature() {
651
        // Test that with temperature 0, softmax_sample returns all keys with smallest logit
652
        let mut logits = HashMap::new();
Yan Ru Pei's avatar
Yan Ru Pei committed
653
654
655
656
657
658
659
660
        let worker1 = WorkerWithDpRank::from_worker_id(1);
        let worker2 = WorkerWithDpRank::from_worker_id(2);
        let worker3 = WorkerWithDpRank::from_worker_id(3);
        let worker4 = WorkerWithDpRank::from_worker_id(4);
        logits.insert(worker1, 5.0);
        logits.insert(worker2, 3.0); // This has the smallest logit
        logits.insert(worker3, 7.0);
        logits.insert(worker4, 3.5);
661

662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
        // With temperature 0, should always return only worker2 (smallest logit)
        let result = softmax_sample(&logits, 0.0);
        assert_eq!(
            result.len(),
            1,
            "Should return one worker when there's no tie"
        );
        assert_eq!(
            result[0], worker2,
            "Should return worker with smallest logit when temperature is 0"
        );

        // Test with tied minimum logits
        logits.clear();
        let worker5 = WorkerWithDpRank::from_worker_id(5);
        let worker6 = WorkerWithDpRank::from_worker_id(6);
        logits.insert(worker1, 5.0);
        logits.insert(worker2, 3.0); // Tied for smallest
        logits.insert(worker5, 3.0); // Tied for smallest
        logits.insert(worker6, 7.0);

        let result = softmax_sample(&logits, 0.0);
        assert_eq!(
            result.len(),
            2,
            "Should return all workers with smallest logit when tied"
        );
        assert!(
            result.contains(&worker2) && result.contains(&worker5),
            "Should contain both tied workers"
        );
693

694
695
        // Test with negative values
        logits.clear();
Yan Ru Pei's avatar
Yan Ru Pei committed
696
697
698
699
700
701
        let worker10 = WorkerWithDpRank::from_worker_id(10);
        let worker20 = WorkerWithDpRank::from_worker_id(20);
        let worker30 = WorkerWithDpRank::from_worker_id(30);
        logits.insert(worker10, -1.0);
        logits.insert(worker20, -5.0); // This has the smallest logit
        logits.insert(worker30, 0.0);
702

703
        let result = softmax_sample(&logits, 0.0);
704
705
706
707
708
        assert_eq!(result.len(), 1);
        assert_eq!(
            result[0], worker20,
            "Should handle negative logits correctly"
        );
709
710
    }
}