sequence.rs 60.6 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
use super::metrics::WORKER_LOAD_METRICS;
41
42
43
44
use super::protocols::{
    ActiveLoad, ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank,
};
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
        // Update Prometheus gauges directly (router's own bookkeeping)
1046
1047
1048
1049
1050
1051
1052
        WORKER_LOAD_METRICS.observe(
            worker.worker_id,
            worker.dp_rank,
            self.worker_type,
            active_blocks,
            active_tokens,
        );
1053
1054

        // Also publish ActiveLoad to NATS for other subscribers (if NATS is available)
1055
1056
1057
1058
1059
1060
1061
        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),
        };

1062
        if let Err(e) = self.metrics_publisher.publish(&active_load).await {
1063
1064
            // 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:?}");
1065
1066
1067
        }
    }

1068
1069
1070
1071
1072
    /// Get the number of workers
    pub fn num_workers(&self) -> usize {
        self.senders.len()
    }

1073
1074
1075
1076
1077
1078
    /// 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
    }

1079
    /// Generic method to query all workers with a given command
1080
    async fn query_workers<T: Send + 'static>(
1081
        &self,
1082
        token_sequence: Option<Vec<SequenceHash>>,
1083
1084
1085
1086
        command_fn: impl Fn(
            Option<Arc<Vec<SequenceHash>>>,
            tokio::sync::oneshot::Sender<T>,
        ) -> UpdateSequences,
Yan Ru Pei's avatar
Yan Ru Pei committed
1087
    ) -> HashMap<WorkerWithDpRank, T> {
1088
1089
1090
1091
1092
        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
1093
        for entry in self.senders.iter() {
Yan Ru Pei's avatar
Yan Ru Pei committed
1094
            let worker = *entry.key();
1095
1096
            let sender = entry.value();
            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
Yan Ru Pei's avatar
Yan Ru Pei committed
1097
            receivers.push((worker, resp_rx));
1098
            if let Err(e) = sender.send(command_fn(token_sequence_shared.clone(), resp_tx)) {
Yan Ru Pei's avatar
Yan Ru Pei committed
1099
                tracing::error!("Failed to send command to worker {:?}: {}", worker, e);
1100
            }
1101
1102
1103
        }

        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1104
        for (worker, receiver) in receivers {
1105
1106
            match tokio::time::timeout(tokio::time::Duration::from_secs(1), receiver).await {
                Ok(Ok(result)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1107
                    results.insert(worker, result);
1108
1109
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1110
                    tracing::error!("Worker {:?} dropped response channel", worker);
1111
1112
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1113
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1114
1115
                }
            }
1116
1117
1118
1119
1120
1121
        }

        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
1122
1123
1124
1125
    pub async fn new_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
    ) -> HashMap<WorkerWithDpRank, usize> {
1126
1127
1128
1129
1130
1131
1132
        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"),
        })
1133
        .await
1134
1135
1136
    }

    /// Query all workers for the total number of blocks (new + active) that would be used by a token sequence
1137
1138
1139
    pub async fn potential_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
Yan Ru Pei's avatar
Yan Ru Pei committed
1140
    ) -> HashMap<WorkerWithDpRank, usize> {
1141
1142
1143
1144
1145
1146
1147
        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"),
        })
1148
        .await
1149
1150
    }

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

1166
1167
        let mut potential_blocks = HashMap::new();
        let mut potential_tokens = HashMap::new();
1168
        let token_sequence_shared = token_sequence.map(Arc::new);
1169
1170
        let mut receivers = Vec::new();

1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
        // 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
                );
1194
            }
1195
1196
        }

1197
1198
1199
        #[cfg(feature = "bench")]
        let send_elapsed = start.elapsed();

1200
        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1201
        for (worker, receiver) in receivers {
1202
1203
            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
1204
1205
                    potential_blocks.insert(worker, blocks);
                    potential_tokens.insert(worker, tokens);
1206
1207
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1208
                    tracing::error!("Worker {:?} dropped response channel", worker);
1209
1210
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1211
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1212
1213
                }
            }
1214
1215
        }

1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
        #[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"
            );
        }

1227
1228
1229
        (potential_blocks, potential_tokens)
    }

1230
    /// Query all workers for their current number of active blocks
Yan Ru Pei's avatar
Yan Ru Pei committed
1231
    pub async fn active_blocks(&self) -> HashMap<WorkerWithDpRank, usize> {
1232
        self.query_workers(None, |_, resp_tx| UpdateSequences::ActiveBlocks { resp_tx })
1233
            .await
1234
    }
1235
1236

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

    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
    }
1250
1251
1252
1253
}

impl Drop for ActiveSequencesMultiWorker {
    fn drop(&mut self) {
1254
1255
1256
        // Send shutdown to all workers
        for entry in self.senders.iter() {
            let _ = entry.value().send(UpdateSequences::Shutdown);
1257
1258
1259
1260
1261
1262
1263
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
1264
    use dynamo_runtime::{DistributedRuntime, Runtime};
1265
    use std::sync::Arc;
1266

1267
1268
1269
1270
1271
    #[test]
    fn test_active_sequences_shared_blocks() {
        let block_size = 4;
        let mut seq_manager = ActiveSequences::new(block_size);

1272
        seq_manager.add_request("request_1".to_string(), Some(vec![1, 2, 3]), 12, 0, None);
1273
1274
1275
        assert_eq!(seq_manager.active_blocks(), 3);
        assert_eq!(seq_manager.active_tokens(), 12);

1276
        seq_manager.add_request("request_2".to_string(), Some(vec![4]), 4, 0, None);
1277
1278
1279
        assert_eq!(seq_manager.active_blocks(), 4);
        assert_eq!(seq_manager.active_tokens(), 16);

1280
        seq_manager.add_request("request_3".to_string(), Some(vec![1, 2, 3, 4]), 16, 4, None);
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
        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);
    }

1297
    #[tokio::test]
1298
    #[ignore]
1299
    async fn test_multi_worker_cross_instance_sync() -> Result<()> {
1300
1301
1302
        // Initialize logging once
        dynamo_runtime::logging::init();

1303
1304
        let block_size = 4; // arbitrary block size

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

1309
1310
        // Create namespace and shared component for both seq_managers
        let namespace = distributed.namespace("test_cross_instance_sync")?;
1311
        let component = namespace.component("sequences")?;
1312

Yan Ru Pei's avatar
Yan Ru Pei committed
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
        // 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;
1323
        workers_with_configs.insert(0, config_worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
1324
1325
1326

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

1329
1330
1331
1332
1333
1334
1335
        let seq_manager_1 = Arc::new(
            ActiveSequencesMultiWorker::new(
                component.clone(),
                block_size,
                workers_with_configs.clone(),
                true,
                1,
1336
                crate::discovery::WORKER_TYPE_DECODE,
1337
1338
1339
1340
            )
            .await?,
        );
        let seq_manager_2 = Arc::new(
1341
1342
1343
1344
1345
1346
1347
1348
1349
            ActiveSequencesMultiWorker::new(
                component,
                block_size,
                workers_with_configs,
                true,
                2,
                crate::discovery::WORKER_TYPE_DECODE,
            )
            .await?,
1350
        );
1351
1352
1353
1354
1355
1356

        // 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
1357
        // Add request_0 to worker 0, dp_rank 0: sequence [0, 1, 2]
1358
1359
1360
1361
        seq_manager_1
            .add_request(
                "request_0".to_string(),
                Some(vec![0, 1, 2]),
1362
1363
1364
                12,   // ISL (3 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1365
                WorkerWithDpRank::new(0, 0),
1366
                None, // lora_name
1367
1368
            )
            .await?;
1369

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

Yan Ru Pei's avatar
Yan Ru Pei committed
1383
        // Add request_2 to worker 1, dp_rank 0: sequence [0, 1, 2, 3] using seq_manager_2
1384
1385
1386
1387
        seq_manager_2
            .add_request(
                "request_2".to_string(),
                Some(vec![0, 1, 2, 3]),
1388
1389
1390
                16,   // ISL (4 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1391
                WorkerWithDpRank::new(1, 0),
1392
                None, // lora_name
1393
1394
            )
            .await?;
1395

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

1399
1400
1401
        // 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;
1402

Yan Ru Pei's avatar
Yan Ru Pei committed
1403
1404
1405
1406
1407
1408
1409
1410
1411
        // 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);

1412
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1413
1414
            blocks_phase1[&worker_0_dp0], 3,
            "Worker 0 dp_rank 0 should have 3 active blocks (from request_0)"
1415
        );
1416
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1417
1418
            blocks_phase1[&worker_0_dp1], 2,
            "Worker 0 dp_rank 1 should have 2 active blocks (from request_1)"
1419
1420
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1421
1422
            blocks_phase1[&worker_1_dp0], 4,
            "Worker 1 dp_rank 0 should have 4 active blocks (from request_2 added by seq_manager_2)"
1423
1424
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1425
1426
            tokens_phase1[&worker_0_dp0], 12,
            "Worker 0 dp_rank 0 should have 12 active tokens"
1427
        );
1428
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1429
1430
1431
1432
1433
1434
            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)"
1435
        );
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452

        // 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
1453
1454
1455
1456
1457
1458
1459
1460
        // 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 {
1461
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1462
1463
1464
                blocks_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active blocks after all requests freed",
                worker.worker_id, worker.dp_rank
1465
1466
            );
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1467
1468
1469
                tokens_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active tokens after all requests freed",
                worker.worker_id, worker.dp_rank
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
            );
        }

        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")?;
1490
        let component = namespace.component("sequences")?;
1491
1492
1493

        // 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
1494
        let mut workers_with_configs = HashMap::new();
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
        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
1507

1508
1509
1510
1511
1512
1513
1514
        let seq_manager_1 = Arc::new(
            ActiveSequencesMultiWorker::new(
                component.clone(),
                block_size,
                workers_with_configs.clone(),
                true,
                1,
1515
                crate::discovery::WORKER_TYPE_DECODE,
1516
1517
1518
1519
            )
            .await?,
        );
        let seq_manager_2 = Arc::new(
1520
1521
1522
1523
1524
1525
1526
1527
1528
            ActiveSequencesMultiWorker::new(
                component,
                block_size,
                workers_with_configs,
                true,
                2,
                crate::discovery::WORKER_TYPE_DECODE,
            )
            .await?,
1529
        );
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542

        // 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
1543
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1544
                WorkerWithDpRank::from_worker_id(0),
1545
                None, // lora_name
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
            )
            .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
1556
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1557
                WorkerWithDpRank::from_worker_id(1),
1558
                None, // lora_name
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
            )
            .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
1569
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1570
                WorkerWithDpRank::from_worker_id(2),
1571
                None, // lora_name
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
            )
            .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
1582
1583
1584
1585
        let worker_0 = WorkerWithDpRank::from_worker_id(0);
        let worker_1 = WorkerWithDpRank::from_worker_id(1);
        let worker_2 = WorkerWithDpRank::from_worker_id(2);

1586
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1587
            tokens_phase1[&worker_0], 12,
1588
            "Worker 0 should have 12 active tokens"
1589
1590
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1591
1592
1593
1594
1595
            tokens_phase1[&worker_1], 8,
            "Worker 1 should have 8 active tokens"
        );
        assert_eq!(
            tokens_phase1[&worker_2], 16,
1596
            "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)"
1597
        );
1598

1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
        // 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
1625
            let worker = WorkerWithDpRank::from_worker_id(worker_id);
1626
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1627
                tokens_phase2[&worker], 0,
1628
1629
1630
1631
1632
                "Worker {} should have 0 active tokens after all requests freed",
                worker_id
            );
        }

1633
        Ok(())
1634
1635
    }
}