sequence.rs 61 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// SPDX-License-Identifier: Apache-2.0

//! KV Cache Sequence Management for LLM Inference
//!
//! This module provides efficient management of token sequences and their associated KV cache blocks
//! for distributed LLM inference. It implements a shared block system where multiple requests can
//! reuse the same KV cache blocks for common token prefixes, significantly reducing memory usage.
//!
//! # Key Components
//!
//! - [`ActiveSequences`]: Single-threaded sequence manager that tracks active requests and their
//!   token sequences, managing shared KV cache blocks efficiently.
//!
//! - [`ActiveSequencesMultiWorker`]: Multi-threaded extension that distributes sequence management
//!   across multiple worker threads, enabling parallel processing of requests while maintaining
//!   consistency.
//!
//! # Architecture
//!
//! The system uses a block-based approach where token sequences are divided into fixed-size blocks.
//! Each block is identified by a hash of its contents, allowing for deduplication when multiple
//! requests share common prefixes (e.g., system prompts, few-shot examples).

25
use crate::kv_router::protocols::OverlapScores;
26
27
use anyhow::Result;
use dashmap::DashMap;
28
use derive_getters::Getters;
29
30
use dynamo_runtime::component::Component;
use dynamo_runtime::traits::DistributedRuntimeProvider;
31
use dynamo_runtime::transports::event_plane::{EventPublisher, EventSubscriber};
32
use dynamo_tokens::SequenceHash;
33
use std::collections::{HashMap, HashSet};
34
use std::rc::{Rc, Weak};
35
use std::sync::Arc;
36
37
use std::time::Duration;
use tokio::time::Instant;
38
39
use uuid::Uuid;

40
41
42
use super::protocols::{
    ActiveLoad, ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank,
};
43
use crate::discovery::{WORKER_ACTIVE_DECODE_BLOCKS_GAUGE, WORKER_ACTIVE_PREFILL_TOKENS_GAUGE};
44
use crate::kv_router::{ACTIVE_SEQUENCES_SUBJECT, KV_METRICS_SUBJECT};
Yan Ru Pei's avatar
Yan Ru Pei committed
45
use crate::local_model::runtime_config::ModelRuntimeConfig;
46
use dynamo_runtime::CancellationToken;
47

48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
/// Errors that can occur during sequence management operations
#[derive(Debug, thiserror::Error)]
pub enum SequenceError {
    #[error("Worker {worker:?} not found")]
    WorkerNotFound { worker: WorkerWithDpRank },

    #[error("Request {request_id} already exists (assigned to worker {worker:?})")]
    DuplicateRequest {
        request_id: String,
        worker: WorkerWithDpRank,
    },

    #[error("Request {request_id} not found")]
    RequestNotFound { request_id: String },

    #[error("Failed to publish event: {0}")]
    PublishFailed(#[from] anyhow::Error),

    #[error("Failed to send command to worker: channel closed")]
    WorkerChannelClosed,
}

70
71
72
/// Duration after which stale requests are forcibly expired (5 minutes)
const EXPIRY_DURATION: Duration = Duration::from_secs(300);

73
74
75
76
77
78
// TODO: use the common request_id if it exists in the repo
pub type RequestId = String;

/// A multi-request sequence manager that handles multiple active sequences with shared KV cache
#[derive(Debug, Getters)]
pub struct ActiveSequences {
79
    active_seqs: HashMap<RequestId, Vec<(SequenceHash, Rc<()>)>>,
80

81
82
    prefill_tokens: HashMap<RequestId, usize>,

83
84
85
    /// Expected output tokens per request (used for resource estimation)
    expected_output_tokens: HashMap<RequestId, u32>,

86
    unique_blocks: HashMap<SequenceHash, Weak<()>>,
87

88
89
90
91
92
    /// Fractional block counts for blocks that are partially cached
    /// When a block is in both unique_blocks and fractional_blocks,
    /// it contributes the fractional value instead of 1 to active_blocks()
    fractional_blocks: HashMap<SequenceHash, f64>,

93
94
95
    #[getter(copy)]
    block_size: usize,

96
97
    #[getter(copy)]
    active_tokens: usize,
98
99
100
101
102
103

    /// Timer for when to force expiry of stale requests
    expiry_timer: Instant,

    /// Set of request IDs to check for expiry
    expiry_requests: HashSet<RequestId>,
104
105
106
107
108
109
110
111
112
113
}

impl ActiveSequences {
    /// Create a new SharedSequenceManager instance
    pub fn new(block_size: usize) -> Self {
        // TODO: make this not a hard req
        assert!(block_size > 1, "block_size must be greater than 1");

        Self {
            active_seqs: HashMap::new(),
114
            prefill_tokens: HashMap::new(),
115
            expected_output_tokens: HashMap::new(),
116
            unique_blocks: HashMap::new(),
117
            fractional_blocks: HashMap::new(),
118
            block_size,
119
            active_tokens: 0,
120
121
            expiry_timer: Instant::now() + EXPIRY_DURATION,
            expiry_requests: HashSet::new(),
122
123
124
        }
    }

125
126
127
128
129
    fn touch_block(&mut self, block: &SequenceHash) -> Rc<()> {
        if let Some(weak) = self.unique_blocks.get(block)
            && let Some(rc) = weak.upgrade()
        {
            return rc;
130
131
        }

132
133
134
135
        let rc = Rc::new(());
        self.unique_blocks.insert(*block, Rc::downgrade(&rc));
        rc
    }
136

137
138
139
140
    fn try_remove_block(&mut self, block: &SequenceHash) {
        if let Some(weak) = self.unique_blocks.get(block)
            && weak.strong_count() == 0
        {
141
            self.unique_blocks.remove(block);
142
            self.fractional_blocks.remove(block);
143
144
145
        }
    }

146
    pub fn active_blocks(&self) -> usize {
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
        let mut count = self.unique_blocks.len() as f64;
        for (hash, frac) in &self.fractional_blocks {
            if self.unique_blocks.contains_key(hash) {
                // Subtract 1 (the full block) and add the fractional value
                count = count - 1.0 + frac;
            }
        }
        count.round() as usize
    }

    /// Find all blocks in a request that have only a single strong reference (only used by this request)
    /// and insert them into fractional_blocks with the given fraction value.
    pub fn set_single_ref_blocks_as_fractional(&mut self, request_id: &RequestId, fraction: f64) {
        let Some(blocks) = self.active_seqs.get(request_id) else {
            tracing::warn!(
                "Request {request_id} not found for set_single_ref_blocks_as_fractional"
            );
            return;
        };

        for (hash, rc) in blocks {
            // A block with strong_count == 1 means only this request holds a reference
            if Rc::strong_count(rc) == 1 {
                self.fractional_blocks.insert(*hash, fraction);
            }
        }
173
174
    }

175
    /// Add a new request with its initial tokens
176
    /// Returns the set of expired request IDs that were removed during cleanup
177
178
179
    pub fn add_request(
        &mut self,
        request_id: RequestId,
180
        token_sequence: Option<Vec<SequenceHash>>,
181
        isl: usize,
182
        overlap: u32,
183
        expected_output_tokens: Option<u32>,
184
    ) -> HashSet<RequestId> {
185
        // Check for double-add and log error, returning early
186
        if self.active_seqs.contains_key(&request_id) {
187
188
            tracing::error!("Request {request_id} is already active. Ignoring duplicate add.");
            return HashSet::new();
189
190
        }

191
192
193
        // Lazily check and clean up expired requests, capturing removed IDs
        let removed_requests = self.force_expiry();

194
        let prefill_tokens = self.new_tokens(isl, overlap);
195
196
197
198
        self.prefill_tokens
            .insert(request_id.clone(), prefill_tokens);
        self.active_tokens += prefill_tokens;

199
200
201
202
203
204
        // Store expected output tokens if provided
        if let Some(tokens) = expected_output_tokens {
            self.expected_output_tokens
                .insert(request_id.clone(), tokens);
        }

205
        if let Some(sequence) = token_sequence {
206
207
208
209
210
211
            let sequence_with_refs: Vec<(SequenceHash, Rc<()>)> = sequence
                .iter()
                .map(|block| (*block, self.touch_block(block)))
                .collect();
            self.active_seqs
                .insert(request_id.clone(), sequence_with_refs);
212
213
214
        } else {
            // dummy empty sequence
            self.active_seqs.insert(request_id.clone(), Vec::new());
215
216
        }

217
        removed_requests
218
219
    }

220
221
222
223
224
    /// Mark prefill as completed for a request, removing it from prefill_tokens tracking
    pub fn mark_prefill_completed(&mut self, request_id: &RequestId) {
        if let Some(tokens) = self.prefill_tokens.remove(request_id) {
            self.active_tokens = self
                .active_tokens
225
                .checked_sub(tokens)
226
227
228
229
230
                .expect("active_tokens underflow");
        }
    }

    pub fn new_tokens(&self, isl: usize, overlap: u32) -> usize {
231
232
233
234
235
236
237
238
239
        let cached_tokens = (overlap as usize) * self.block_size;
        isl.checked_sub(cached_tokens)
            .unwrap_or_else(|| {
                tracing::error!(
                    "prefill_tokens < 0 with ISL {isl} < cached_tokens {cached_tokens} (overlap {overlap} * block_size {}), returning 0",
                    self.block_size
                );
                0
            })
240
241
242
243
    }

    pub fn potential_blocks_and_tokens(
        &self,
244
        token_sequence: Option<&[SequenceHash]>,
245
        isl: usize,
246
247
        overlap: u32,
    ) -> (usize, usize) {
248
        let potential_blocks = if let Some(token_seq) = token_sequence {
249
            self.new_blocks(token_seq) + self.active_blocks()
250
        } else {
251
            self.active_blocks()
252
        };
253
        let potential_tokens = self.new_tokens(isl, overlap) + self.active_tokens;
254
255
256
        (potential_blocks, potential_tokens)
    }

257
    /// Match a request against existing blocks and return the number of new blocks that would be added
258
259
    pub fn new_blocks(&self, token_sequence: &[SequenceHash]) -> usize {
        token_sequence
260
261
262
263
264
265
266
            .iter()
            .filter(|block| !self.unique_blocks.contains_key(block))
            .count()
    }

    /// Return the total number of blocks that would be used if the token sequence was added
    /// This is the sum of new blocks that would be added plus the current active blocks
267
    pub fn potential_blocks(&self, token_sequence: &[SequenceHash]) -> usize {
268
        self.new_blocks(token_sequence) + self.active_blocks()
269
270
271
272
    }

    /// Free all blocks associated with a request
    pub fn free(&mut self, request_id: &RequestId) -> usize {
273
        self.mark_prefill_completed(request_id);
274

275
276
        self.expiry_requests.remove(request_id);

277
278
279
        // Remove expected output tokens tracking
        self.expected_output_tokens.remove(request_id);

280
281
282
283
284
        // Remove from active_seqs and get the token sequence
        let token_seq = match self.active_seqs.remove(request_id) {
            Some(seq) => seq,
            None => {
                tracing::warn!("Trying to free non-existent request {request_id}");
285
                return self.active_blocks();
286
            }
287
288
        };

289
290
291
292
        // Drop each Rc reference, then clean up the corresponding weak reference
        for (block_hash, rc) in token_seq {
            drop(rc);
            self.try_remove_block(&block_hash);
293
294
        }

295
        self.active_blocks()
296
    }
297

298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
    /// Add an output block with a random hash and optional fractional decay weight.
    ///
    /// This is used during generation to track output blocks as they are created.
    /// The decay_fraction (if provided) represents how "temporary" the block is:
    /// - 1.0 means fully counted (early in generation)
    /// - 0.0 means not counted (near end of expected output)
    /// - Computed as: 1 - (current_osl / expected_output_tokens)
    ///
    /// Returns true if the block was added, false if the request was not found.
    pub fn add_output_block(
        &mut self,
        request_id: &RequestId,
        decay_fraction: Option<f64>,
    ) -> bool {
        // Check if request exists first (immutable borrow)
        if !self.active_seqs.contains_key(request_id) {
            tracing::warn!("Request {request_id} not found for add_output_block");
            return false;
        }

        // Generate a random block hash using UUID
        let random_hash: SequenceHash = Uuid::new_v4().as_u64_pair().0;

        // Touch the block (adds to unique_blocks)
        let rc = self.touch_block(&random_hash);

        // Now we can safely get_mut and push
        self.active_seqs
            .get_mut(request_id)
            .unwrap()
            .push((random_hash, rc));

        // Apply fractional decay to all single-ref blocks in this request if provided
        if let Some(frac) = decay_fraction {
            self.set_single_ref_blocks_as_fractional(request_id, frac);
        }

        true
    }

338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
    /// Force expiry of stale requests if the timer has elapsed
    /// Returns the set of expired request IDs that were removed
    pub fn force_expiry(&mut self) -> HashSet<RequestId> {
        let now = Instant::now();

        // Early return if timer hasn't expired yet
        if now < self.expiry_timer {
            return HashSet::new();
        }

        // Process expired requests - drain to avoid clone
        let expired_requests: HashSet<RequestId> = self.expiry_requests.drain().collect();
        for request_id in &expired_requests {
            tracing::warn!("Force expiring stale request: {}", request_id);
            self.free(request_id);
        }

        self.expiry_timer = now + EXPIRY_DURATION;
        self.expiry_requests = self.active_seqs.keys().cloned().collect();

        expired_requests
    }
360
361
362
363
364
}

enum UpdateSequences {
    AddRequest {
        request_id: RequestId,
365
        token_sequence: Option<Vec<SequenceHash>>,
366
        isl: usize,
367
        overlap: u32,
368
        expected_output_tokens: Option<u32>,
369
        resp_tx: tokio::sync::oneshot::Sender<HashSet<RequestId>>,
370
371
372
373
    },
    Free {
        request_id: RequestId,
    },
374
    MarkPrefillCompleted {
375
376
        request_id: RequestId,
    },
377
378
379
380
381
    AddOutputBlock {
        request_id: RequestId,
        decay_fraction: Option<f64>,
        resp_tx: tokio::sync::oneshot::Sender<bool>,
    },
382
    NewBlocks {
383
        token_sequence: Arc<Vec<SequenceHash>>,
384
        resp_tx: tokio::sync::oneshot::Sender<usize>,
385
386
    },
    PotentialBlocks {
387
        token_sequence: Arc<Vec<SequenceHash>>,
388
        resp_tx: tokio::sync::oneshot::Sender<usize>,
389
    },
390
    PotentialBlocksAndTokens {
391
        token_sequence: Option<Arc<Vec<SequenceHash>>>,
392
        isl: usize,
393
        overlap: u32,
394
        resp_tx: tokio::sync::oneshot::Sender<(usize, usize)>,
395
    },
396
    ActiveBlocks {
397
        resp_tx: tokio::sync::oneshot::Sender<usize>,
398
    },
399
    ActiveTokens {
400
        resp_tx: tokio::sync::oneshot::Sender<usize>,
401
    },
402
403
404
405
406
    Shutdown,
}

/// Multi-worker extension of ActiveSequences that distributes requests across multiple threads
pub struct ActiveSequencesMultiWorker {
Yan Ru Pei's avatar
Yan Ru Pei committed
407
408
    senders: Arc<DashMap<WorkerWithDpRank, tokio::sync::mpsc::UnboundedSender<UpdateSequences>>>,
    request_to_worker: Arc<DashMap<RequestId, WorkerWithDpRank>>,
409
    request_to_lora: Arc<DashMap<RequestId, String>>,
Yan Ru Pei's avatar
Yan Ru Pei committed
410
    handles: Arc<DashMap<WorkerWithDpRank, std::thread::JoinHandle<()>>>,
411
    block_size: usize,
412
    component: Component,
413
    router_id: u64,
414
415
416
417
    /// Publisher for sequence events
    event_publisher: EventPublisher,
    /// Publisher for metrics (namespace-scoped)
    metrics_publisher: EventPublisher,
418
    replica_sync: bool,
419
420
    /// Worker type for Prometheus metrics labeling ("prefill" or "decode")
    worker_type: &'static str,
421
422
423
}

impl ActiveSequencesMultiWorker {
424
    pub async fn new(
425
426
        component: Component,
        block_size: usize,
427
        workers_with_configs: HashMap<u64, ModelRuntimeConfig>,
428
        replica_sync: bool,
429
        router_id: u64,
430
        worker_type: &'static str,
431
    ) -> Result<Self> {
432
433
        assert!(block_size > 1, "block_size must be greater than 1");

434
435
436
        let senders = Arc::new(DashMap::new());
        let handles = Arc::new(DashMap::new());
        let request_to_worker = Arc::new(DashMap::new());
437
        let request_to_lora = Arc::new(DashMap::new());
438

Yan Ru Pei's avatar
Yan Ru Pei committed
439
440
        // Expand workers by their dp_rank
        for (worker_id, config) in workers_with_configs {
441
            let dp_size = config.data_parallel_size;
Yan Ru Pei's avatar
Yan Ru Pei committed
442
443
444
445
446
447
448
449
450

            for dp_rank in 0..dp_size {
                let worker = WorkerWithDpRank::new(worker_id, dp_rank);
                // Create a child cancellation token from the component's runtime
                let cancel_token = component.drt().runtime().child_token();
                let (sender, handle) = Self::start_worker(block_size, cancel_token);
                senders.insert(worker, sender);
                handles.insert(worker, handle);
            }
451
452
        }

453
454
455
456
457
        let event_publisher =
            EventPublisher::for_component(&component, ACTIVE_SEQUENCES_SUBJECT).await?;
        let metrics_publisher =
            EventPublisher::for_namespace(component.namespace(), KV_METRICS_SUBJECT).await?;

458
459
460
        let multi_worker = Self {
            senders: senders.clone(),
            request_to_worker: request_to_worker.clone(),
461
            request_to_lora: request_to_lora.clone(),
462
463
            handles,
            block_size,
464
            component: component.clone(),
465
466
            event_publisher,
            metrics_publisher,
467
468
            router_id,
            replica_sync,
469
            worker_type,
470
471
472
473
474
475
        };

        // Start the subscription loop only if replica_sync is enabled
        if replica_sync {
            let senders_clone = senders.clone();
            let request_to_worker_clone = request_to_worker.clone();
476
            let request_to_lora_clone = request_to_lora.clone();
477
478
            let component_clone = component.clone();
            let router_id_clone = router_id;
479
            let cancel_token = component.drt().runtime().child_token();
480

481
            tokio::spawn(async move {
482
                // NATS subscription loop
483
484
485
                if let Err(e) = Self::subscribe_to_events(
                    senders_clone,
                    request_to_worker_clone,
486
                    request_to_lora_clone,
487
488
                    component_clone,
                    router_id_clone,
489
                    cancel_token,
490
491
492
493
494
495
                )
                .await
                {
                    tracing::error!("Error in active sequences events subscription: {}", e);
                }
            });
496
        }
497

498
        Ok(multi_worker)
499
500
    }

501
502
503
    /// Helper method to start a worker task
    fn start_worker(
        block_size: usize,
504
        cancel_token: CancellationToken,
505
506
    ) -> (
        tokio::sync::mpsc::UnboundedSender<UpdateSequences>,
507
        std::thread::JoinHandle<()>,
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
        let (request_tx, request_rx) = tokio::sync::mpsc::unbounded_channel();

        let handle = std::thread::spawn(move || {
            // Create a single-threaded tokio runtime
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap();

            runtime.block_on(async move {
                let mut active_sequences = ActiveSequences::new(block_size);
                let mut request_rx = request_rx;

                loop {
                    tokio::select! {
                        command = request_rx.recv() => {
                            let Some(command) = command else {
                                break;
                            };

                            match command {
                                UpdateSequences::AddRequest {
                                    request_id,
                                    token_sequence,
                                    isl,
                                    overlap,
535
                                    expected_output_tokens,
536
537
                                    resp_tx,
                                } => {
538
                                    let removed = active_sequences.add_request(request_id, token_sequence, isl, overlap, expected_output_tokens);
539
540
541
542
543
544
545
546
                                    let _ = resp_tx.send(removed);
                                }
                                UpdateSequences::Free { request_id } => {
                                    active_sequences.free(&request_id);
                                }
                                UpdateSequences::MarkPrefillCompleted { request_id } => {
                                    active_sequences.mark_prefill_completed(&request_id);
                                }
547
548
549
550
551
552
553
554
                                UpdateSequences::AddOutputBlock {
                                    request_id,
                                    decay_fraction,
                                    resp_tx,
                                } => {
                                    let success = active_sequences.add_output_block(&request_id, decay_fraction);
                                    let _ = resp_tx.send(success);
                                }
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
                                UpdateSequences::NewBlocks {
                                    token_sequence,
                                    resp_tx,
                                } => {
                                    let new_blocks = active_sequences.new_blocks(&token_sequence);
                                    let _ = resp_tx.send(new_blocks);
                                }
                                UpdateSequences::PotentialBlocks {
                                    token_sequence,
                                    resp_tx,
                                } => {
                                    let potential_blocks = active_sequences.potential_blocks(&token_sequence);
                                    let _ = resp_tx.send(potential_blocks);
                                }
                                UpdateSequences::PotentialBlocksAndTokens {
                                    token_sequence,
                                    isl,
                                    overlap,
                                    resp_tx,
                                } => {
                                    let potential_tokens = active_sequences.potential_blocks_and_tokens(
                                        token_sequence.as_ref().map(|v| v.as_slice()),
577
578
                                        isl,
                                        overlap,
579
580
581
582
583
584
585
586
587
588
589
590
591
                                    );
                                    let _ = resp_tx.send(potential_tokens);
                                }
                                UpdateSequences::ActiveBlocks { resp_tx } => {
                                    let active_blocks = active_sequences.active_blocks();
                                    let _ = resp_tx.send(active_blocks);
                                }
                                UpdateSequences::ActiveTokens { resp_tx } => {
                                    let active_tokens = active_sequences.active_tokens();
                                    let _ = resp_tx.send(active_tokens);
                                }
                                UpdateSequences::Shutdown => {
                                    break;
592
593
594
                                }
                            }
                        }
595
596
597
598
599
                        // Handle cancellation
                        _ = cancel_token.cancelled() => {
                            tracing::debug!("Worker task cancelled");
                            break;
                        }
600
                    }
601
                }
602
603
604
            });

            tracing::debug!("ActiveSequences worker task completed");
605
606
607
608
609
610
611
        });

        (request_tx, handle)
    }

    /// Background task to subscribe to active sequence events and update all workers
    async fn subscribe_to_events(
Yan Ru Pei's avatar
Yan Ru Pei committed
612
613
614
615
        senders: Arc<
            DashMap<WorkerWithDpRank, tokio::sync::mpsc::UnboundedSender<UpdateSequences>>,
        >,
        request_to_worker: Arc<DashMap<RequestId, WorkerWithDpRank>>,
616
        request_to_lora: Arc<DashMap<RequestId, String>>,
617
        component: Component,
618
        router_id: u64,
619
        cancel_token: CancellationToken,
620
    ) -> Result<()> {
621
622
623
        let mut subscriber = EventSubscriber::for_component(&component, ACTIVE_SEQUENCES_SUBJECT)
            .await?
            .typed::<ActiveSequenceEvent>();
624

625
626
627
628
629
630
631
632
        loop {
            tokio::select! {
                // Handle incoming events
                result = subscriber.next() => {
                    let Some(result) = result else {
                        // Stream ended
                        break;
                    };
633

634
                    let Ok((_envelope, event)) = result else {
635
636
637
                        tracing::error!(
                            "Error receiving active sequence event: {}",
                            result.unwrap_err()
638
                        );
639
640
641
642
643
644
                        continue;
                    };

                    // Skip events emitted by itself
                    if event.router_id == router_id {
                        continue;
645
                    }
646
647
648
649
650
651

                    match &event.data {
                        ActiveSequenceEventData::AddRequest {
                            token_sequence,
                            isl,
                            overlap,
652
                            expected_output_tokens,
653
                        } => {
Yan Ru Pei's avatar
Yan Ru Pei committed
654
                            request_to_worker.insert(event.request_id.clone(), event.worker);
655

656
657
658
659
660
                            // Store lora_name mapping if present
                            if let Some(ref lora_name) = event.lora_name {
                                request_to_lora.insert(event.request_id.clone(), lora_name.clone());
                            }

Yan Ru Pei's avatar
Yan Ru Pei committed
661
                            if let Some(sender) = senders.get(&event.worker) {
662
663
664
665
666
667
668
                                // For replicated events, we create a dummy response channel since we don't need to handle expired requests
                                let (resp_tx, _) = tokio::sync::oneshot::channel();
                                let _ = sender.send(UpdateSequences::AddRequest {
                                    request_id: event.request_id.clone(),
                                    token_sequence: token_sequence.clone(),
                                    isl: *isl,
                                    overlap: *overlap,
669
                                    expected_output_tokens: *expected_output_tokens,
670
671
672
673
                                    resp_tx,
                                });
                            } else {
                                tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
674
675
                                    "Worker {:?} not found, cannot process AddRequest",
                                    event.worker
676
677
678
679
                                );
                            }
                        }
                        ActiveSequenceEventData::Free => {
Yan Ru Pei's avatar
Yan Ru Pei committed
680
681
                            if let Some((_, worker)) = request_to_worker.remove(&event.request_id)
                                && let Some(sender) = senders.get(&worker)
682
683
684
685
686
                            {
                                let _ = sender.send(UpdateSequences::Free {
                                    request_id: event.request_id.clone(),
                                });
                            }
687
688
                            // Clean up lora_name mapping
                            request_to_lora.remove(&event.request_id);
689
690
                        }
                        ActiveSequenceEventData::MarkPrefillCompleted => {
Yan Ru Pei's avatar
Yan Ru Pei committed
691
692
                            if let Some(worker) = request_to_worker.get(&event.request_id)
                                && let Some(sender) = senders.get(&*worker)
693
694
695
696
697
698
                            {
                                let _ = sender.send(UpdateSequences::MarkPrefillCompleted {
                                    request_id: event.request_id.clone(),
                                });
                            }
                        }
699
                    }
700
                }
701
702
703
704
                // Handle cancellation
                _ = cancel_token.cancelled() => {
                    tracing::debug!("Subscription task cancelled");
                    break;
705
706
                }
            }
707
        }
708

709
        Ok(())
710
711
712
    }

    /// Update the set of workers, adding and removing as needed
713
    pub fn update_workers(&self, new_workers_with_configs: HashMap<u64, ModelRuntimeConfig>) {
Yan Ru Pei's avatar
Yan Ru Pei committed
714
        let current_workers: HashSet<WorkerWithDpRank> =
715
            self.senders.iter().map(|entry| *entry.key()).collect();
716

Yan Ru Pei's avatar
Yan Ru Pei committed
717
718
719
        // Expand new workers by their dp_rank
        let mut new_workers: HashSet<WorkerWithDpRank> = HashSet::new();
        for (worker_id, config) in &new_workers_with_configs {
720
            let dp_size = config.data_parallel_size;
Yan Ru Pei's avatar
Yan Ru Pei committed
721
722
723
724
725
726
727

            for dp_rank in 0..dp_size {
                new_workers.insert(WorkerWithDpRank::new(*worker_id, dp_rank));
            }
        }

        let workers_to_remove: Vec<WorkerWithDpRank> =
728
            current_workers.difference(&new_workers).copied().collect();
Yan Ru Pei's avatar
Yan Ru Pei committed
729
        let workers_to_add: Vec<WorkerWithDpRank> =
730
731
            new_workers.difference(&current_workers).copied().collect();

Yan Ru Pei's avatar
Yan Ru Pei committed
732
733
734
        // Remove workers (this will naturally remove all dp ranks for a worker_id)
        for worker in &workers_to_remove {
            tracing::warn!("Removing worker {:?}", worker);
735
736

            // Send shutdown command to the worker
Yan Ru Pei's avatar
Yan Ru Pei committed
737
            if let Some((_, sender)) = self.senders.remove(worker) {
738
739
                let _ = sender.send(UpdateSequences::Shutdown);
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
740
            self.handles.remove(worker);
741

742
743
744
745
746
747
748
749
            // Collect request_ids to remove from request_to_lora
            let requests_to_remove: Vec<RequestId> = self
                .request_to_worker
                .iter()
                .filter(|entry| entry.value() == worker)
                .map(|entry| entry.key().clone())
                .collect();

750
751
            // Clean up request_to_worker mappings for this worker
            self.request_to_worker
Yan Ru Pei's avatar
Yan Ru Pei committed
752
                .retain(|_request_id, mapped_worker| mapped_worker != worker);
753
754
755
756
757

            // Clean up request_to_lora mappings for removed requests
            for request_id in requests_to_remove {
                self.request_to_lora.remove(&request_id);
            }
758
759
760
        }

        // Add new workers
Yan Ru Pei's avatar
Yan Ru Pei committed
761
762
        for worker in &workers_to_add {
            tracing::warn!("Adding worker {:?}", worker);
763

764
765
766
767
            let (sender, handle) = Self::start_worker(
                self.block_size,
                self.component.drt().runtime().child_token(),
            );
Yan Ru Pei's avatar
Yan Ru Pei committed
768
769
            self.senders.insert(*worker, sender);
            self.handles.insert(*worker, handle);
770
771
772
        }
    }

773
    #[allow(clippy::too_many_arguments)]
774
775
    pub async fn add_request(
        &self,
776
        request_id: RequestId,
777
        token_sequence: Option<Vec<SequenceHash>>,
778
        isl: usize,
779
        overlap: u32,
780
        expected_output_tokens: Option<u32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
781
        worker: WorkerWithDpRank,
782
        lora_name: Option<String>,
783
    ) -> Result<(), SequenceError> {
784
785
786
787
788
789
790
791
792
        // Clone the sender upfront so we don't hold the DashMap Ref across
        // the .await points below. Also eliminates the TOCTOU between
        // contains_key and a later get().unwrap().
        let sender = self
            .senders
            .get(&worker)
            .ok_or(SequenceError::WorkerNotFound { worker })?
            .value()
            .clone();
793
794
795
796
797
798
799

        // Check for duplicate request
        if let Some(existing_worker) = self.request_to_worker.get(&request_id) {
            return Err(SequenceError::DuplicateRequest {
                request_id,
                worker: *existing_worker,
            });
800
801
        }

802
803
804
        // Create response channel
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();

805
806
807
808
        // Publish event only if replica_sync is enabled
        if self.replica_sync {
            let event = ActiveSequenceEvent {
                request_id: request_id.clone(),
Yan Ru Pei's avatar
Yan Ru Pei committed
809
                worker,
810
811
812
813
                data: ActiveSequenceEventData::AddRequest {
                    token_sequence: token_sequence.clone(),
                    isl,
                    overlap,
814
                    expected_output_tokens,
815
816
                },
                router_id: self.router_id,
817
                lora_name: lora_name.clone(),
818
            };
819
            self.event_publisher.publish(&event).await?;
820
821
        }

Yan Ru Pei's avatar
Yan Ru Pei committed
822
823
        // Update local state with full WorkerWithDpRank
        self.request_to_worker.insert(request_id.clone(), worker);
824

825
826
827
828
829
        // Store lora_name for later use in Free/MarkPrefillCompleted events
        if let Some(lora) = lora_name {
            self.request_to_lora.insert(request_id.clone(), lora);
        }

830
        sender
831
832
833
            .send(UpdateSequences::AddRequest {
                request_id,
                token_sequence,
834
                isl,
835
                overlap,
836
                expected_output_tokens,
837
                resp_tx,
838
            })
839
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
840

841
842
843
        // Wait for response and handle removed requests
        let removed_requests = resp_rx
            .await
844
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
845
846
847
848

        // Remove expired requests from request_to_worker mapping
        for expired_id in &removed_requests {
            self.request_to_worker.remove(expired_id);
849
            self.request_to_lora.remove(expired_id);
850
851
        }

852
853
854
        // Publish ActiveLoad metrics for this worker
        self.publish_active_load_for_worker(worker).await;

855
        Ok(())
856
857
    }

858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
    /// Send a command to the worker assigned to a request, optionally publishing
    /// a replica-sync event and cleaning up request mappings afterward.
    async fn send_to_request_worker(
        &self,
        request_id: &RequestId,
        event_data: ActiveSequenceEventData,
        command_fn: impl FnOnce(RequestId) -> UpdateSequences,
        remove_mapping: bool,
    ) -> Result<(), SequenceError> {
        let worker = self
            .request_to_worker
            .get(request_id)
            .map(|entry| *entry)
            .ok_or_else(|| SequenceError::RequestNotFound {
                request_id: request_id.clone(),
            })?;
874

875
876
877
878
879
880
        let sender = self
            .senders
            .get(&worker)
            .ok_or(SequenceError::WorkerNotFound { worker })?
            .value()
            .clone();
881

882
        if self.replica_sync {
883
884
885
886
887
            let lora_name = self
                .request_to_lora
                .get(request_id)
                .map(|entry| entry.value().clone());

888
889
            let event = ActiveSequenceEvent {
                request_id: request_id.clone(),
Yan Ru Pei's avatar
Yan Ru Pei committed
890
                worker,
891
                data: event_data,
892
                router_id: self.router_id,
893
                lora_name,
894
            };
895
            self.event_publisher.publish(&event).await?;
896
897
        }

898
899
        sender
            .send(command_fn(request_id.clone()))
900
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
901

902
903
904
905
        if remove_mapping {
            self.request_to_worker.remove(request_id);
            self.request_to_lora.remove(request_id);
        }
906

907
908
        self.publish_active_load_for_worker(worker).await;

909
        Ok(())
910
911
    }

912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
    /// Free all blocks associated with a request
    ///
    /// Note: This operation is idempotent. Calling it multiple times for the same request
    /// will log a warning but not return an error (double free is allowed).
    pub async fn free(&self, request_id: &RequestId) -> Result<(), SequenceError> {
        if !self.request_to_worker.contains_key(request_id) {
            tracing::debug!("Request {request_id} not found, already freed (idempotent)");
            return Ok(());
        }

        self.send_to_request_worker(
            request_id,
            ActiveSequenceEventData::Free,
            |rid| UpdateSequences::Free { request_id: rid },
            true,
        )
        .await
    }

931
    /// Mark prefill as completed for a request
932
933
934
935
936
937
938
    ///
    /// Note: Calling this multiple times for the same request is allowed and will be a no-op
    /// after the first call (idempotent).
    pub async fn mark_prefill_completed(
        &self,
        request_id: &RequestId,
    ) -> Result<(), SequenceError> {
939
940
941
942
943
944
945
        self.send_to_request_worker(
            request_id,
            ActiveSequenceEventData::MarkPrefillCompleted,
            |rid| UpdateSequences::MarkPrefillCompleted { request_id: rid },
            false,
        )
        .await
946
947
    }

948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
    /// Add an output block with optional fractional decay weight
    ///
    /// This is used during generation to track output blocks as they are created.
    /// The decay_fraction represents how "temporary" the block is based on generation progress.
    pub async fn add_output_block(
        &self,
        request_id: &RequestId,
        decay_fraction: Option<f64>,
    ) -> Result<(), SequenceError> {
        let worker = self
            .request_to_worker
            .get(request_id)
            .map(|entry| *entry)
            .ok_or_else(|| SequenceError::RequestNotFound {
                request_id: request_id.clone(),
            })?;

965
966
967
968
969
970
971
        // Clone sender upfront to avoid TOCTOU between contains_key and get().unwrap()
        let sender = self
            .senders
            .get(&worker)
            .ok_or(SequenceError::WorkerNotFound { worker })?
            .value()
            .clone();
972
973
974
975
976

        // Create response channel
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();

        // Send command to worker
977
        sender
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
            .send(UpdateSequences::AddOutputBlock {
                request_id: request_id.clone(),
                decay_fraction,
                resp_tx,
            })
            .map_err(|_| SequenceError::WorkerChannelClosed)?;

        // Wait for response
        let success = resp_rx
            .await
            .map_err(|_| SequenceError::WorkerChannelClosed)?;

        if !success {
            return Err(SequenceError::RequestNotFound {
                request_id: request_id.clone(),
            });
        }

        // Publish ActiveLoad metrics for this worker
        self.publish_active_load_for_worker(worker).await;

        Ok(())
    }

1002
1003
    /// Helper method to query a single worker for active blocks/tokens and publish ActiveLoad
    async fn publish_active_load_for_worker(&self, worker: WorkerWithDpRank) {
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
        // Clone the sender and drop the DashMap Ref immediately.
        // Holding a Ref across .await points can deadlock: if the task yields
        // and update_workers() needs a write lock on the same shard, the
        // runtime thread blocks forever.
        let sender = {
            let Some(entry) = self.senders.get(&worker) else {
                tracing::warn!("Worker {worker:?} not found when publishing ActiveLoad");
                return;
            };
            entry.value().clone()
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
        };

        // Query active blocks
        let (blocks_tx, blocks_rx) = tokio::sync::oneshot::channel();
        if sender
            .send(UpdateSequences::ActiveBlocks { resp_tx: blocks_tx })
            .is_err()
        {
            tracing::warn!("Failed to send ActiveBlocks query to worker {worker:?}");
            return;
        }

        // Query active tokens
        let (tokens_tx, tokens_rx) = tokio::sync::oneshot::channel();
        if sender
            .send(UpdateSequences::ActiveTokens { resp_tx: tokens_tx })
            .is_err()
        {
            tracing::warn!("Failed to send ActiveTokens query to worker {worker:?}");
            return;
        }

        // Await both responses
        let (active_blocks, active_tokens) = match tokio::join!(blocks_rx, tokens_rx) {
            (Ok(blocks), Ok(tokens)) => (blocks, tokens),
            _ => {
                tracing::warn!("Failed to receive active blocks/tokens from worker {worker:?}");
                return;
            }
        };

1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
        // Update Prometheus gauges directly (router's own bookkeeping)
        let worker_id_str = worker.worker_id.to_string();
        let dp_rank_str = worker.dp_rank.to_string();
        WORKER_ACTIVE_DECODE_BLOCKS_GAUGE
            .with_label_values(&[
                worker_id_str.as_str(),
                dp_rank_str.as_str(),
                self.worker_type,
            ])
            .set(active_blocks as i64);
        WORKER_ACTIVE_PREFILL_TOKENS_GAUGE
            .with_label_values(&[
                worker_id_str.as_str(),
                dp_rank_str.as_str(),
                self.worker_type,
            ])
            .set(active_tokens as i64);

        // Also publish ActiveLoad to NATS for other subscribers (if NATS is available)
1064
1065
1066
1067
1068
1069
1070
        let active_load = ActiveLoad {
            worker_id: worker.worker_id,
            dp_rank: worker.dp_rank,
            active_decode_blocks: Some(active_blocks as u64),
            active_prefill_tokens: Some(active_tokens as u64),
        };

1071
        if let Err(e) = self.metrics_publisher.publish(&active_load).await {
1072
1073
            // This is expected if NATS is not available - the local gauge update above already succeeded
            tracing::trace!("Failed to publish ActiveLoad to NATS for worker {worker:?}: {e:?}");
1074
1075
1076
        }
    }

1077
1078
1079
1080
1081
    /// Get the number of workers
    pub fn num_workers(&self) -> usize {
        self.senders.len()
    }

1082
1083
1084
1085
1086
1087
    /// Get the worker type for this router ("prefill" or "decode").
    /// Used for Prometheus metric labeling.
    pub fn worker_type(&self) -> &'static str {
        self.worker_type
    }

1088
    /// Generic method to query all workers with a given command
1089
    async fn query_workers<T: Send + 'static>(
1090
        &self,
1091
        token_sequence: Option<Vec<SequenceHash>>,
1092
1093
1094
1095
        command_fn: impl Fn(
            Option<Arc<Vec<SequenceHash>>>,
            tokio::sync::oneshot::Sender<T>,
        ) -> UpdateSequences,
Yan Ru Pei's avatar
Yan Ru Pei committed
1096
    ) -> HashMap<WorkerWithDpRank, T> {
1097
1098
1099
1100
1101
        let mut results = HashMap::new();
        let token_sequence_shared = token_sequence.map(Arc::new);
        let mut receivers = Vec::new();

        // Send queries to all workers in parallel
1102
        for entry in self.senders.iter() {
Yan Ru Pei's avatar
Yan Ru Pei committed
1103
            let worker = *entry.key();
1104
1105
            let sender = entry.value();
            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
Yan Ru Pei's avatar
Yan Ru Pei committed
1106
            receivers.push((worker, resp_rx));
1107
            if let Err(e) = sender.send(command_fn(token_sequence_shared.clone(), resp_tx)) {
Yan Ru Pei's avatar
Yan Ru Pei committed
1108
                tracing::error!("Failed to send command to worker {:?}: {}", worker, e);
1109
            }
1110
1111
1112
        }

        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1113
        for (worker, receiver) in receivers {
1114
1115
            match tokio::time::timeout(tokio::time::Duration::from_secs(1), receiver).await {
                Ok(Ok(result)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1116
                    results.insert(worker, result);
1117
1118
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1119
                    tracing::error!("Worker {:?} dropped response channel", worker);
1120
1121
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1122
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1123
1124
                }
            }
1125
1126
1127
1128
1129
1130
        }

        results
    }

    /// Query all workers for the number of new blocks that would be added by a token sequence
Yan Ru Pei's avatar
Yan Ru Pei committed
1131
1132
1133
1134
    pub async fn new_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
    ) -> HashMap<WorkerWithDpRank, usize> {
1135
1136
1137
1138
1139
1140
1141
        self.query_workers(Some(token_sequence), |ts, resp_tx| match ts {
            Some(ts) => UpdateSequences::NewBlocks {
                token_sequence: ts,
                resp_tx,
            },
            None => unreachable!("token_sequence should always be Some for new_blocks"),
        })
1142
        .await
1143
1144
1145
    }

    /// Query all workers for the total number of blocks (new + active) that would be used by a token sequence
1146
1147
1148
    pub async fn potential_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
Yan Ru Pei's avatar
Yan Ru Pei committed
1149
    ) -> HashMap<WorkerWithDpRank, usize> {
1150
1151
1152
1153
1154
1155
1156
        self.query_workers(Some(token_sequence), |ts, resp_tx| match ts {
            Some(ts) => UpdateSequences::PotentialBlocks {
                token_sequence: ts,
                resp_tx,
            },
            None => unreachable!("token_sequence should always be Some for potential_blocks"),
        })
1157
        .await
1158
1159
    }

1160
    /// Query all workers for the potential tokens (new + active) that would be used by a token sequence with overlap
1161
    pub async fn potential_blocks_and_tokens(
1162
        &self,
1163
        token_sequence: Option<Vec<SequenceHash>>,
1164
        isl: usize,
1165
        overlaps: OverlapScores,
Yan Ru Pei's avatar
Yan Ru Pei committed
1166
1167
1168
1169
    ) -> (
        HashMap<WorkerWithDpRank, usize>,
        HashMap<WorkerWithDpRank, usize>,
    ) {
1170
1171
1172
1173
1174
        #[cfg(feature = "bench")]
        let start = Instant::now();
        #[cfg(feature = "bench")]
        let num_workers = self.senders.len();

1175
1176
        let mut potential_blocks = HashMap::new();
        let mut potential_tokens = HashMap::new();
1177
        let token_sequence_shared = token_sequence.map(Arc::new);
1178
1179
        let mut receivers = Vec::new();

1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
        // Iterate through all workers, not just those with overlap
        // This ensures we properly account for active tokens/blocks on all workers
        for sender_entry in self.senders.iter() {
            let worker = *sender_entry.key();
            let sender = sender_entry.value();

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

            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
            receivers.push((worker, resp_rx));

            if let Err(e) = sender.send(UpdateSequences::PotentialBlocksAndTokens {
                token_sequence: token_sequence_shared.clone(),
                isl,
                overlap,
                resp_tx,
            }) {
                tracing::error!(
                    "Failed to send potential_tokens command to worker {:?}: {}",
                    worker,
                    e
                );
1203
            }
1204
1205
        }

1206
1207
1208
        #[cfg(feature = "bench")]
        let send_elapsed = start.elapsed();

1209
        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1210
        for (worker, receiver) in receivers {
1211
1212
            match tokio::time::timeout(tokio::time::Duration::from_secs(1), receiver).await {
                Ok(Ok((blocks, tokens))) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1213
1214
                    potential_blocks.insert(worker, blocks);
                    potential_tokens.insert(worker, tokens);
1215
1216
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1217
                    tracing::error!("Worker {:?} dropped response channel", worker);
1218
1219
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1220
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1221
1222
                }
            }
1223
1224
        }

1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
        #[cfg(feature = "bench")]
        {
            let total_elapsed = start.elapsed();
            tracing::info!(
                num_workers,
                send_us = send_elapsed.as_micros() as u64,
                total_us = total_elapsed.as_micros() as u64,
                "potential_blocks_and_tokens completed"
            );
        }

1236
1237
1238
        (potential_blocks, potential_tokens)
    }

1239
    /// Query all workers for their current number of active blocks
Yan Ru Pei's avatar
Yan Ru Pei committed
1240
    pub async fn active_blocks(&self) -> HashMap<WorkerWithDpRank, usize> {
1241
        self.query_workers(None, |_, resp_tx| UpdateSequences::ActiveBlocks { resp_tx })
1242
            .await
1243
    }
1244
1245

    /// Query all workers for their current number of active tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1246
    pub async fn active_tokens(&self) -> HashMap<WorkerWithDpRank, usize> {
1247
        self.query_workers(None, |_, resp_tx| UpdateSequences::ActiveTokens { resp_tx })
1248
            .await
1249
    }
1250
1251
1252
1253
1254
1255
1256
1257
1258

    pub fn get_active_lora_counts(&self) -> HashMap<String, usize> {
        let mut counts: HashMap<String, usize> = HashMap::new();
        for entry in self.request_to_lora.iter() {
            let lora_name = entry.value().clone();
            *counts.entry(lora_name).or_insert(0) += 1;
        }
        counts
    }
1259
1260
1261
1262
}

impl Drop for ActiveSequencesMultiWorker {
    fn drop(&mut self) {
1263
1264
1265
        // Send shutdown to all workers
        for entry in self.senders.iter() {
            let _ = entry.value().send(UpdateSequences::Shutdown);
1266
1267
1268
1269
1270
1271
1272
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
1273
    use dynamo_runtime::{DistributedRuntime, Runtime};
1274
    use std::sync::Arc;
1275

1276
1277
1278
1279
1280
    #[test]
    fn test_active_sequences_shared_blocks() {
        let block_size = 4;
        let mut seq_manager = ActiveSequences::new(block_size);

1281
        seq_manager.add_request("request_1".to_string(), Some(vec![1, 2, 3]), 12, 0, None);
1282
1283
1284
        assert_eq!(seq_manager.active_blocks(), 3);
        assert_eq!(seq_manager.active_tokens(), 12);

1285
        seq_manager.add_request("request_2".to_string(), Some(vec![4]), 4, 0, None);
1286
1287
1288
        assert_eq!(seq_manager.active_blocks(), 4);
        assert_eq!(seq_manager.active_tokens(), 16);

1289
        seq_manager.add_request("request_3".to_string(), Some(vec![1, 2, 3, 4]), 16, 4, None);
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
        assert_eq!(seq_manager.active_blocks(), 4);
        assert_eq!(seq_manager.active_tokens(), 16);

        seq_manager.free(&"request_2".to_string());
        assert_eq!(seq_manager.active_blocks(), 4);
        assert_eq!(seq_manager.active_tokens(), 12);

        seq_manager.free(&"request_3".to_string());
        assert_eq!(seq_manager.active_blocks(), 3);
        assert_eq!(seq_manager.active_tokens(), 12);

        seq_manager.free(&"request_1".to_string());
        assert_eq!(seq_manager.active_blocks(), 0);
        assert_eq!(seq_manager.active_tokens(), 0);
    }

1306
    #[tokio::test]
1307
    #[ignore]
1308
    async fn test_multi_worker_cross_instance_sync() -> Result<()> {
1309
1310
1311
        // Initialize logging once
        dynamo_runtime::logging::init();

1312
1313
        let block_size = 4; // arbitrary block size

1314
1315
1316
        // Create runtime and distributed runtime
        let runtime = Runtime::from_current()?;
        let distributed = DistributedRuntime::from_settings(runtime.clone()).await?;
1317

1318
1319
        // Create namespace and shared component for both seq_managers
        let namespace = distributed.namespace("test_cross_instance_sync")?;
1320
        let component = namespace.component("sequences")?;
1321

Yan Ru Pei's avatar
Yan Ru Pei committed
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
        // Create multi-worker sequence managers with:
        // - Worker 0 with dp_size=2 (dp_ranks 0 and 1)
        // - Worker 1 with dp_size=1 (dp_rank 0)
        // This gives us 3 effective workers total to test dp_rank effect
        // Both seq_managers use the same component to ensure event synchronization works
        let mut workers_with_configs = HashMap::new();

        // Create runtime config for worker 0 with dp_size=2
        let mut config_worker_0 = crate::local_model::runtime_config::ModelRuntimeConfig::new();
        config_worker_0.data_parallel_size = 2;
1332
        workers_with_configs.insert(0, config_worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
1333
1334
1335

        // Create runtime config for worker 1 with dp_size=1 (default)
        let config_worker_1 = crate::local_model::runtime_config::ModelRuntimeConfig::new();
1336
        workers_with_configs.insert(1, config_worker_1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1337

1338
1339
1340
1341
1342
1343
1344
        let seq_manager_1 = Arc::new(
            ActiveSequencesMultiWorker::new(
                component.clone(),
                block_size,
                workers_with_configs.clone(),
                true,
                1,
1345
                crate::discovery::WORKER_TYPE_DECODE,
1346
1347
1348
1349
            )
            .await?,
        );
        let seq_manager_2 = Arc::new(
1350
1351
1352
1353
1354
1355
1356
1357
1358
            ActiveSequencesMultiWorker::new(
                component,
                block_size,
                workers_with_configs,
                true,
                2,
                crate::discovery::WORKER_TYPE_DECODE,
            )
            .await?,
1359
        );
1360
1361
1362
1363
1364
1365

        // Give some time for the subscription loops to start
        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

        // PHASE 1: Add requests using both seq_manager_1 and seq_manager_2

Yan Ru Pei's avatar
Yan Ru Pei committed
1366
        // Add request_0 to worker 0, dp_rank 0: sequence [0, 1, 2]
1367
1368
1369
1370
        seq_manager_1
            .add_request(
                "request_0".to_string(),
                Some(vec![0, 1, 2]),
1371
1372
1373
                12,   // ISL (3 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1374
                WorkerWithDpRank::new(0, 0),
1375
                None, // lora_name
1376
1377
            )
            .await?;
1378

Yan Ru Pei's avatar
Yan Ru Pei committed
1379
        // Add request_1 to worker 0, dp_rank 1: sequence [3, 4]
1380
1381
1382
1383
        seq_manager_1
            .add_request(
                "request_1".to_string(),
                Some(vec![3, 4]),
1384
1385
1386
                8,    // ISL (2 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1387
                WorkerWithDpRank::new(0, 1),
1388
                None, // lora_name
1389
1390
            )
            .await?;
1391

Yan Ru Pei's avatar
Yan Ru Pei committed
1392
        // Add request_2 to worker 1, dp_rank 0: sequence [0, 1, 2, 3] using seq_manager_2
1393
1394
1395
1396
        seq_manager_2
            .add_request(
                "request_2".to_string(),
                Some(vec![0, 1, 2, 3]),
1397
1398
1399
                16,   // ISL (4 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1400
                WorkerWithDpRank::new(1, 0),
1401
                None, // lora_name
1402
1403
            )
            .await?;
1404

1405
1406
        // Give some time for synchronization
        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;
1407

1408
1409
1410
        // Query seq_manager_1 to verify it sees all requests including request_2 from seq_manager_2
        let blocks_phase1 = seq_manager_1.active_blocks().await;
        let tokens_phase1 = seq_manager_1.active_tokens().await;
1411

Yan Ru Pei's avatar
Yan Ru Pei committed
1412
1413
1414
1415
1416
1417
1418
1419
1420
        // Verify that seq_manager_1 sees all requests including request_2 from seq_manager_2
        // We now have:
        // - Worker 0, dp_rank 0: request_0
        // - Worker 0, dp_rank 1: request_1
        // - Worker 1, dp_rank 0: request_2
        let worker_0_dp0 = WorkerWithDpRank::new(0, 0);
        let worker_0_dp1 = WorkerWithDpRank::new(0, 1);
        let worker_1_dp0 = WorkerWithDpRank::new(1, 0);

1421
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1422
1423
            blocks_phase1[&worker_0_dp0], 3,
            "Worker 0 dp_rank 0 should have 3 active blocks (from request_0)"
1424
        );
1425
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1426
1427
            blocks_phase1[&worker_0_dp1], 2,
            "Worker 0 dp_rank 1 should have 2 active blocks (from request_1)"
1428
1429
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1430
1431
            blocks_phase1[&worker_1_dp0], 4,
            "Worker 1 dp_rank 0 should have 4 active blocks (from request_2 added by seq_manager_2)"
1432
1433
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1434
1435
            tokens_phase1[&worker_0_dp0], 12,
            "Worker 0 dp_rank 0 should have 12 active tokens"
1436
        );
1437
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1438
1439
1440
1441
1442
1443
            tokens_phase1[&worker_0_dp1], 8,
            "Worker 0 dp_rank 1 should have 8 active tokens"
        );
        assert_eq!(
            tokens_phase1[&worker_1_dp0], 16,
            "Worker 1 dp_rank 0 should have 16 active tokens (from request_2 added by seq_manager_2)"
1444
        );
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461

        // PHASE 2: Free requests using opposite sequence managers, verify on seq_manager_2

        // Free request_2 (which was added by seq_manager_2) using seq_manager_1
        seq_manager_1.free(&"request_2".to_string()).await?;

        // Free request_0 and request_1 (which were added by seq_manager_1) using seq_manager_2
        seq_manager_2.free(&"request_0".to_string()).await?;
        seq_manager_2.free(&"request_1".to_string()).await?;

        // Give some time for synchronization
        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

        // Query seq_manager_2 to verify everything is empty
        let blocks_phase2 = seq_manager_2.active_blocks().await;
        let tokens_phase2 = seq_manager_2.active_tokens().await;

Yan Ru Pei's avatar
Yan Ru Pei committed
1462
1463
1464
1465
1466
1467
1468
1469
        // Verify phase 2 results - everything should be empty for all 3 workers
        let all_workers = vec![
            WorkerWithDpRank::new(0, 0),
            WorkerWithDpRank::new(0, 1),
            WorkerWithDpRank::new(1, 0),
        ];

        for worker in all_workers {
1470
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1471
1472
1473
                blocks_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active blocks after all requests freed",
                worker.worker_id, worker.dp_rank
1474
1475
            );
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1476
1477
1478
                tokens_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active tokens after all requests freed",
                worker.worker_id, worker.dp_rank
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
            );
        }

        Ok(())
    }

    #[tokio::test]
    #[ignore]
    async fn test_multi_worker_no_token_sequence_sync() -> Result<()> {
        // Initialize logging once
        dynamo_runtime::logging::init();

        let block_size = 4; // arbitrary block size

        // Create runtime and distributed runtime
        let runtime = Runtime::from_current()?;
        let distributed = DistributedRuntime::from_settings(runtime.clone()).await?;

        // Create namespace and shared component for both seq_managers
        let namespace = distributed.namespace("test_no_token_seq_sync")?;
1499
        let component = namespace.component("sequences")?;
1500
1501
1502

        // Create multi-worker sequence managers with ALL workers [0, 1, 2]
        // Both use the same component to ensure event synchronization works
Yan Ru Pei's avatar
Yan Ru Pei committed
1503
        let mut workers_with_configs = HashMap::new();
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
        workers_with_configs.insert(
            0,
            crate::local_model::runtime_config::ModelRuntimeConfig::new(),
        );
        workers_with_configs.insert(
            1,
            crate::local_model::runtime_config::ModelRuntimeConfig::new(),
        );
        workers_with_configs.insert(
            2,
            crate::local_model::runtime_config::ModelRuntimeConfig::new(),
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1516

1517
1518
1519
1520
1521
1522
1523
        let seq_manager_1 = Arc::new(
            ActiveSequencesMultiWorker::new(
                component.clone(),
                block_size,
                workers_with_configs.clone(),
                true,
                1,
1524
                crate::discovery::WORKER_TYPE_DECODE,
1525
1526
1527
1528
            )
            .await?,
        );
        let seq_manager_2 = Arc::new(
1529
1530
1531
1532
1533
1534
1535
1536
1537
            ActiveSequencesMultiWorker::new(
                component,
                block_size,
                workers_with_configs,
                true,
                2,
                crate::discovery::WORKER_TYPE_DECODE,
            )
            .await?,
1538
        );
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551

        // Give some time for the subscription loops to start
        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

        // PHASE 1: Add requests (without token sequences) using both seq_managers

        // Add request_0 to worker 0 with no token sequence
        seq_manager_1
            .add_request(
                "request_0".to_string(),
                None, // No token sequence
                12,   // ISL (12 tokens)
                0,    // no overlap
1552
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1553
                WorkerWithDpRank::from_worker_id(0),
1554
                None, // lora_name
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
            )
            .await?;

        // Add request_1 to worker 1 with no token sequence
        seq_manager_1
            .add_request(
                "request_1".to_string(),
                None, // No token sequence
                8,    // ISL (8 tokens)
                0,    // no overlap
1565
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1566
                WorkerWithDpRank::from_worker_id(1),
1567
                None, // lora_name
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
            )
            .await?;

        // Add request_2 to worker 2 with no token sequence using seq_manager_2
        seq_manager_2
            .add_request(
                "request_2".to_string(),
                None, // No token sequence
                16,   // ISL (16 tokens)
                0,    // no overlap
1578
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1579
                WorkerWithDpRank::from_worker_id(2),
1580
                None, // lora_name
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
            )
            .await?;

        // Give some time for synchronization
        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

        // Query seq_manager_1 to verify it sees all requests including request_2 from seq_manager_2
        let tokens_phase1 = seq_manager_1.active_tokens().await;

        // Verify that seq_manager_1 sees all requests including request_2 from thread 2
Yan Ru Pei's avatar
Yan Ru Pei committed
1591
1592
1593
1594
        let worker_0 = WorkerWithDpRank::from_worker_id(0);
        let worker_1 = WorkerWithDpRank::from_worker_id(1);
        let worker_2 = WorkerWithDpRank::from_worker_id(2);

1595
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1596
            tokens_phase1[&worker_0], 12,
1597
            "Worker 0 should have 12 active tokens"
1598
1599
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1600
1601
1602
1603
1604
            tokens_phase1[&worker_1], 8,
            "Worker 1 should have 8 active tokens"
        );
        assert_eq!(
            tokens_phase1[&worker_2], 16,
1605
            "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)"
1606
        );
1607

1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
        // PHASE 2: Free requests using opposite sequence managers, verify on seq_manager_2

        // Mark prefill completed and free request_2 (which was added by seq_manager_2) using seq_manager_1
        seq_manager_1
            .mark_prefill_completed(&"request_2".to_string())
            .await?;
        seq_manager_1.free(&"request_2".to_string()).await?;

        // Mark prefill completed and free requests 0 and 1 (which were added by seq_manager_1) using seq_manager_2
        seq_manager_2
            .mark_prefill_completed(&"request_0".to_string())
            .await?;
        seq_manager_2
            .mark_prefill_completed(&"request_1".to_string())
            .await?;
        seq_manager_2.free(&"request_0".to_string()).await?;
        seq_manager_2.free(&"request_1".to_string()).await?;

        // Give some time for synchronization
        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

        // Query seq_manager_2 to verify everything is empty
        let tokens_phase2 = seq_manager_2.active_tokens().await;

        // Verify phase 2 results - everything should be empty
        for worker_id in 0..=2 {
Yan Ru Pei's avatar
Yan Ru Pei committed
1634
            let worker = WorkerWithDpRank::from_worker_id(worker_id);
1635
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1636
                tokens_phase2[&worker], 0,
1637
1638
1639
1640
1641
                "Worker {} should have 0 active tokens after all requests freed",
                worker_id
            );
        }

1642
        Ok(())
1643
1644
    }
}