sequence.rs 56.2 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::indexer::OverlapScores;
26
use crate::tokens::SequenceHash;
27
28
use anyhow::Result;
use dashmap::DashMap;
29
use derive_getters::Getters;
30
31
use dynamo_runtime::component::Component;
use dynamo_runtime::traits::DistributedRuntimeProvider;
32
use dynamo_runtime::traits::events::{EventPublisher, EventSubscriber};
33
use futures::StreamExt;
34
use std::collections::{HashMap, HashSet};
35
use std::rc::{Rc, Weak};
36
use std::sync::Arc;
37
38
use std::time::Duration;
use tokio::time::Instant;
39
40
use uuid::Uuid;

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
409
    senders: Arc<DashMap<WorkerWithDpRank, tokio::sync::mpsc::UnboundedSender<UpdateSequences>>>,
    request_to_worker: Arc<DashMap<RequestId, WorkerWithDpRank>>,
    handles: Arc<DashMap<WorkerWithDpRank, std::thread::JoinHandle<()>>>,
410
    block_size: usize,
411
412
413
    component: Component,
    router_id: Uuid,
    replica_sync: bool,
414
415
416
}

impl ActiveSequencesMultiWorker {
417
418
419
    pub fn new(
        component: Component,
        block_size: usize,
420
        workers_with_configs: HashMap<u64, Option<ModelRuntimeConfig>>,
421
        replica_sync: bool,
422
        router_uuid: String,
423
    ) -> Self {
424
425
        assert!(block_size > 1, "block_size must be greater than 1");

426
427
428
        let senders = Arc::new(DashMap::new());
        let handles = Arc::new(DashMap::new());
        let request_to_worker = Arc::new(DashMap::new());
429
430
431
432
433
434
435
436
        let router_id = Uuid::parse_str(&router_uuid).unwrap_or_else(|e| {
            tracing::warn!(
                "Failed to parse router UUID '{}': {}, using new UUID",
                router_uuid,
                e
            );
            Uuid::new_v4()
        });
437

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

            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);
            }
450
451
        }

452
453
454
        let multi_worker = Self {
            senders: senders.clone(),
            request_to_worker: request_to_worker.clone(),
455
456
            handles,
            block_size,
457
458
459
460
461
462
463
464
465
466
467
            component: component.clone(),
            router_id,
            replica_sync,
        };

        // 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();
            let component_clone = component.clone();
            let router_id_clone = router_id;
468
            let cancel_token = component.drt().runtime().child_token();
469

470
            tokio::spawn(async move {
471
                // NATS subscription loop
472
473
474
475
476
                if let Err(e) = Self::subscribe_to_events(
                    senders_clone,
                    request_to_worker_clone,
                    component_clone,
                    router_id_clone,
477
                    cancel_token,
478
479
480
481
482
483
                )
                .await
                {
                    tracing::error!("Error in active sequences events subscription: {}", e);
                }
            });
484
        }
485
486

        multi_worker
487
488
    }

489
490
491
    /// Helper method to start a worker task
    fn start_worker(
        block_size: usize,
492
        cancel_token: CancellationToken,
493
494
    ) -> (
        tokio::sync::mpsc::UnboundedSender<UpdateSequences>,
495
        std::thread::JoinHandle<()>,
496
    ) {
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
        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,
523
                                    expected_output_tokens,
524
525
                                    resp_tx,
                                } => {
526
                                    let removed = active_sequences.add_request(request_id, token_sequence, isl, overlap, expected_output_tokens);
527
528
529
530
531
532
533
534
                                    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);
                                }
535
536
537
538
539
540
541
542
                                UpdateSequences::AddOutputBlock {
                                    request_id,
                                    decay_fraction,
                                    resp_tx,
                                } => {
                                    let success = active_sequences.add_output_block(&request_id, decay_fraction);
                                    let _ = resp_tx.send(success);
                                }
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
                                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()),
565
566
                                        isl,
                                        overlap,
567
568
569
570
571
572
573
574
575
576
577
578
579
                                    );
                                    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;
580
581
582
                                }
                            }
                        }
583
584
585
586
587
                        // Handle cancellation
                        _ = cancel_token.cancelled() => {
                            tracing::debug!("Worker task cancelled");
                            break;
                        }
588
                    }
589
                }
590
591
592
            });

            tracing::debug!("ActiveSequences worker task completed");
593
594
595
596
597
598
599
        });

        (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
600
601
602
603
        senders: Arc<
            DashMap<WorkerWithDpRank, tokio::sync::mpsc::UnboundedSender<UpdateSequences>>,
        >,
        request_to_worker: Arc<DashMap<RequestId, WorkerWithDpRank>>,
604
605
        component: Component,
        router_id: Uuid,
606
        cancel_token: CancellationToken,
607
608
609
610
611
    ) -> Result<()> {
        let mut subscriber = component
            .subscribe_with_type::<ActiveSequenceEvent>(ACTIVE_SEQUENCES_SUBJECT)
            .await?;

612
613
614
615
616
617
618
619
        loop {
            tokio::select! {
                // Handle incoming events
                result = subscriber.next() => {
                    let Some(result) = result else {
                        // Stream ended
                        break;
                    };
620

621
622
623
624
                    let Ok(event) = result else {
                        tracing::error!(
                            "Error receiving active sequence event: {}",
                            result.unwrap_err()
625
                        );
626
627
628
629
630
631
                        continue;
                    };

                    // Skip events emitted by itself
                    if event.router_id == router_id {
                        continue;
632
                    }
633
634
635
636
637
638

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

Yan Ru Pei's avatar
Yan Ru Pei committed
643
                            if let Some(sender) = senders.get(&event.worker) {
644
645
646
647
648
649
650
                                // 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,
651
                                    expected_output_tokens: *expected_output_tokens,
652
653
654
655
                                    resp_tx,
                                });
                            } else {
                                tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
656
657
                                    "Worker {:?} not found, cannot process AddRequest",
                                    event.worker
658
659
660
661
                                );
                            }
                        }
                        ActiveSequenceEventData::Free => {
Yan Ru Pei's avatar
Yan Ru Pei committed
662
663
                            if let Some((_, worker)) = request_to_worker.remove(&event.request_id)
                                && let Some(sender) = senders.get(&worker)
664
665
666
667
668
669
670
                            {
                                let _ = sender.send(UpdateSequences::Free {
                                    request_id: event.request_id.clone(),
                                });
                            }
                        }
                        ActiveSequenceEventData::MarkPrefillCompleted => {
Yan Ru Pei's avatar
Yan Ru Pei committed
671
672
                            if let Some(worker) = request_to_worker.get(&event.request_id)
                                && let Some(sender) = senders.get(&*worker)
673
674
675
676
677
678
                            {
                                let _ = sender.send(UpdateSequences::MarkPrefillCompleted {
                                    request_id: event.request_id.clone(),
                                });
                            }
                        }
679
                    }
680
                }
681
682
683
684
                // Handle cancellation
                _ = cancel_token.cancelled() => {
                    tracing::debug!("Subscription task cancelled");
                    break;
685
686
                }
            }
687
        }
688

689
        Ok(())
690
691
692
    }

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

Yan Ru Pei's avatar
Yan Ru Pei committed
700
701
702
703
704
705
706
707
708
709
710
        // Expand new workers by their dp_rank
        let mut new_workers: HashSet<WorkerWithDpRank> = HashSet::new();
        for (worker_id, config) in &new_workers_with_configs {
            let dp_size = config.as_ref().map(|c| c.data_parallel_size).unwrap_or(1);

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

        let workers_to_remove: Vec<WorkerWithDpRank> =
711
            current_workers.difference(&new_workers).copied().collect();
Yan Ru Pei's avatar
Yan Ru Pei committed
712
        let workers_to_add: Vec<WorkerWithDpRank> =
713
714
            new_workers.difference(&current_workers).copied().collect();

Yan Ru Pei's avatar
Yan Ru Pei committed
715
716
717
        // Remove workers (this will naturally remove all dp ranks for a worker_id)
        for worker in &workers_to_remove {
            tracing::warn!("Removing worker {:?}", worker);
718
719

            // Send shutdown command to the worker
Yan Ru Pei's avatar
Yan Ru Pei committed
720
            if let Some((_, sender)) = self.senders.remove(worker) {
721
722
                let _ = sender.send(UpdateSequences::Shutdown);
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
723
            self.handles.remove(worker);
724
725
726

            // Clean up request_to_worker mappings for this worker
            self.request_to_worker
Yan Ru Pei's avatar
Yan Ru Pei committed
727
                .retain(|_request_id, mapped_worker| mapped_worker != worker);
728
729
730
        }

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

734
735
736
737
            let (sender, handle) = Self::start_worker(
                self.block_size,
                self.component.drt().runtime().child_token(),
            );
Yan Ru Pei's avatar
Yan Ru Pei committed
738
739
            self.senders.insert(*worker, sender);
            self.handles.insert(*worker, handle);
740
741
742
        }
    }

743
744
    pub async fn add_request(
        &self,
745
        request_id: RequestId,
746
        token_sequence: Option<Vec<SequenceHash>>,
747
        isl: usize,
748
        overlap: u32,
749
        expected_output_tokens: Option<u32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
750
        worker: WorkerWithDpRank,
751
752
    ) -> Result<(), SequenceError> {
        // Check for worker existence
Yan Ru Pei's avatar
Yan Ru Pei committed
753
        if !self.senders.contains_key(&worker) {
754
755
756
757
758
759
760
761
762
            return Err(SequenceError::WorkerNotFound { worker });
        }

        // 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,
            });
763
764
        }

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

768
769
770
771
        // 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
772
                worker,
773
774
775
776
                data: ActiveSequenceEventData::AddRequest {
                    token_sequence: token_sequence.clone(),
                    isl,
                    overlap,
777
                    expected_output_tokens,
778
779
780
781
782
783
                },
                router_id: self.router_id,
            };
            self.component
                .publish(ACTIVE_SEQUENCES_SUBJECT, &event)
                .await?;
784
785
        }

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

789
        self.senders
Yan Ru Pei's avatar
Yan Ru Pei committed
790
            .get(&worker)
791
            .unwrap()
792
793
794
            .send(UpdateSequences::AddRequest {
                request_id,
                token_sequence,
795
                isl,
796
                overlap,
797
                expected_output_tokens,
798
                resp_tx,
799
            })
800
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
801

802
803
804
        // Wait for response and handle removed requests
        let removed_requests = resp_rx
            .await
805
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
806
807
808
809
810
811

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

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

815
        Ok(())
816
817
    }

818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
    /// 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> {
        // Check if request exists - if not, it's already been freed (idempotent)
        let Some(worker) = self.request_to_worker.get(request_id).map(|entry| *entry) else {
            tracing::debug!("Request {request_id} not found, already freed (idempotent)");
            return Ok(());
        };

        // Verify worker still exists
        if !self.senders.contains_key(&worker) {
            return Err(SequenceError::WorkerNotFound { worker });
        }
833

834
835
836
837
        // 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
838
                worker,
839
840
841
842
843
844
845
846
847
848
                data: ActiveSequenceEventData::Free,
                router_id: self.router_id,
            };
            self.component
                .publish(ACTIVE_SEQUENCES_SUBJECT, &event)
                .await?;
        }

        // Update local state
        self.senders
Yan Ru Pei's avatar
Yan Ru Pei committed
849
            .get(&worker)
850
            .unwrap()
851
852
853
            .send(UpdateSequences::Free {
                request_id: request_id.clone(),
            })
854
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
855
856

        self.request_to_worker.remove(request_id);
857

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

861
        Ok(())
862
863
    }

864
    /// Mark prefill as completed for a request
865
866
867
868
869
870
871
    ///
    /// 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> {
Yan Ru Pei's avatar
Yan Ru Pei committed
872
        let worker = self
873
874
            .request_to_worker
            .get(request_id)
875
            .map(|entry| *entry)
876
877
878
879
880
881
882
883
            .ok_or_else(|| SequenceError::RequestNotFound {
                request_id: request_id.clone(),
            })?;

        // Verify worker still exists
        if !self.senders.contains_key(&worker) {
            return Err(SequenceError::WorkerNotFound { worker });
        }
884
885
886
887
888

        // 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
889
                worker,
890
891
892
893
894
895
896
                data: ActiveSequenceEventData::MarkPrefillCompleted,
                router_id: self.router_id,
            };
            self.component
                .publish(ACTIVE_SEQUENCES_SUBJECT, &event)
                .await?;
        }
897

898
899
        // Update local state
        self.senders
Yan Ru Pei's avatar
Yan Ru Pei committed
900
            .get(&worker)
901
            .unwrap()
902
            .send(UpdateSequences::MarkPrefillCompleted {
903
904
                request_id: request_id.clone(),
            })
905
            .map_err(|_| SequenceError::WorkerChannelClosed)?;
906

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

910
        Ok(())
911
912
    }

913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
    /// 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(),
            })?;

        // Verify worker still exists
        if !self.senders.contains_key(&worker) {
            return Err(SequenceError::WorkerNotFound { worker });
        }

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

        // Send command to worker
        self.senders
            .get(&worker)
            .unwrap()
            .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(())
    }

966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
    /// Helper method to query a single worker for active blocks/tokens and publish ActiveLoad
    async fn publish_active_load_for_worker(&self, worker: WorkerWithDpRank) {
        let Some(sender) = self.senders.get(&worker) else {
            tracing::warn!("Worker {worker:?} not found when publishing ActiveLoad");
            return;
        };

        // 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;
            }
        };

        // Publish ActiveLoad
        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),
        };

        if let Err(e) = self
            .component
            .namespace()
            .publish(KV_METRICS_SUBJECT, &active_load)
            .await
        {
            tracing::warn!("Failed to publish ActiveLoad for worker {worker:?}: {e:?}");
        }
    }

1020
1021
1022
1023
1024
1025
    /// Get the number of workers
    pub fn num_workers(&self) -> usize {
        self.senders.len()
    }

    /// Generic method to query all workers with a given command
1026
    async fn query_workers<T: Send + 'static>(
1027
        &self,
1028
        token_sequence: Option<Vec<SequenceHash>>,
1029
1030
1031
1032
        command_fn: impl Fn(
            Option<Arc<Vec<SequenceHash>>>,
            tokio::sync::oneshot::Sender<T>,
        ) -> UpdateSequences,
Yan Ru Pei's avatar
Yan Ru Pei committed
1033
    ) -> HashMap<WorkerWithDpRank, T> {
1034
1035
1036
1037
1038
        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
1039
        for entry in self.senders.iter() {
Yan Ru Pei's avatar
Yan Ru Pei committed
1040
            let worker = *entry.key();
1041
1042
            let sender = entry.value();
            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
Yan Ru Pei's avatar
Yan Ru Pei committed
1043
            receivers.push((worker, resp_rx));
1044
            if let Err(e) = sender.send(command_fn(token_sequence_shared.clone(), resp_tx)) {
Yan Ru Pei's avatar
Yan Ru Pei committed
1045
                tracing::error!("Failed to send command to worker {:?}: {}", worker, e);
1046
            }
1047
1048
1049
        }

        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1050
        for (worker, receiver) in receivers {
1051
1052
            match tokio::time::timeout(tokio::time::Duration::from_secs(1), receiver).await {
                Ok(Ok(result)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1053
                    results.insert(worker, result);
1054
1055
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1056
                    tracing::error!("Worker {:?} dropped response channel", worker);
1057
1058
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1059
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1060
1061
                }
            }
1062
1063
1064
1065
1066
1067
        }

        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
1068
1069
1070
1071
    pub async fn new_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
    ) -> HashMap<WorkerWithDpRank, usize> {
1072
1073
1074
1075
1076
1077
1078
        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"),
        })
1079
        .await
1080
1081
1082
    }

    /// Query all workers for the total number of blocks (new + active) that would be used by a token sequence
1083
1084
1085
    pub async fn potential_blocks(
        &self,
        token_sequence: Vec<SequenceHash>,
Yan Ru Pei's avatar
Yan Ru Pei committed
1086
    ) -> HashMap<WorkerWithDpRank, usize> {
1087
1088
1089
1090
1091
1092
1093
        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"),
        })
1094
        .await
1095
1096
    }

1097
    /// Query all workers for the potential tokens (new + active) that would be used by a token sequence with overlap
1098
    pub async fn potential_blocks_and_tokens(
1099
        &self,
1100
        token_sequence: Option<Vec<SequenceHash>>,
1101
        isl: usize,
1102
        overlaps: OverlapScores,
Yan Ru Pei's avatar
Yan Ru Pei committed
1103
1104
1105
1106
    ) -> (
        HashMap<WorkerWithDpRank, usize>,
        HashMap<WorkerWithDpRank, usize>,
    ) {
1107
1108
        let mut potential_blocks = HashMap::new();
        let mut potential_tokens = HashMap::new();
1109
        let token_sequence_shared = token_sequence.map(Arc::new);
1110
1111
        let mut receivers = Vec::new();

1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
        // 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
                );
1135
            }
1136
1137
1138
        }

        // Collect results from all workers
Yan Ru Pei's avatar
Yan Ru Pei committed
1139
        for (worker, receiver) in receivers {
1140
1141
            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
1142
1143
                    potential_blocks.insert(worker, blocks);
                    potential_tokens.insert(worker, tokens);
1144
1145
                }
                Ok(Err(_)) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1146
                    tracing::error!("Worker {:?} dropped response channel", worker);
1147
1148
                }
                Err(_) => {
Yan Ru Pei's avatar
Yan Ru Pei committed
1149
                    tracing::error!("Timeout waiting for response from worker {:?}", worker);
1150
1151
                }
            }
1152
1153
1154
1155
1156
        }

        (potential_blocks, potential_tokens)
    }

1157
    /// Query all workers for their current number of active blocks
Yan Ru Pei's avatar
Yan Ru Pei committed
1158
    pub async fn active_blocks(&self) -> HashMap<WorkerWithDpRank, usize> {
1159
        self.query_workers(None, |_, resp_tx| UpdateSequences::ActiveBlocks { resp_tx })
1160
            .await
1161
    }
1162
1163

    /// Query all workers for their current number of active tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1164
    pub async fn active_tokens(&self) -> HashMap<WorkerWithDpRank, usize> {
1165
        self.query_workers(None, |_, resp_tx| UpdateSequences::ActiveTokens { resp_tx })
1166
            .await
1167
    }
1168
1169
1170
1171
}

impl Drop for ActiveSequencesMultiWorker {
    fn drop(&mut self) {
1172
1173
1174
        // Send shutdown to all workers
        for entry in self.senders.iter() {
            let _ = entry.value().send(UpdateSequences::Shutdown);
1175
1176
1177
1178
1179
1180
1181
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
1182
    use dynamo_runtime::{DistributedRuntime, Runtime};
1183
    use std::sync::Arc;
1184

1185
1186
1187
1188
1189
    #[test]
    fn test_active_sequences_shared_blocks() {
        let block_size = 4;
        let mut seq_manager = ActiveSequences::new(block_size);

1190
        seq_manager.add_request("request_1".to_string(), Some(vec![1, 2, 3]), 12, 0, None);
1191
1192
1193
        assert_eq!(seq_manager.active_blocks(), 3);
        assert_eq!(seq_manager.active_tokens(), 12);

1194
        seq_manager.add_request("request_2".to_string(), Some(vec![4]), 4, 0, None);
1195
1196
1197
        assert_eq!(seq_manager.active_blocks(), 4);
        assert_eq!(seq_manager.active_tokens(), 16);

1198
        seq_manager.add_request("request_3".to_string(), Some(vec![1, 2, 3, 4]), 16, 4, None);
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
        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);
    }

1215
    #[tokio::test]
1216
    #[ignore]
1217
    async fn test_multi_worker_cross_instance_sync() -> Result<()> {
1218
1219
1220
        // Initialize logging once
        dynamo_runtime::logging::init();

1221
1222
        let block_size = 4; // arbitrary block size

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

1227
1228
        // Create namespace and shared component for both seq_managers
        let namespace = distributed.namespace("test_cross_instance_sync")?;
1229
        let component = namespace.component("sequences")?;
1230

Yan Ru Pei's avatar
Yan Ru Pei committed
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
        // 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;
        workers_with_configs.insert(0, Some(config_worker_0));

        // Create runtime config for worker 1 with dp_size=1 (default)
        let config_worker_1 = crate::local_model::runtime_config::ModelRuntimeConfig::new();
        workers_with_configs.insert(1, Some(config_worker_1));

1247
1248
1249
        let seq_manager_1 = Arc::new(ActiveSequencesMultiWorker::new(
            component.clone(),
            block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
1250
            workers_with_configs.clone(),
1251
1252
1253
1254
1255
1256
            true,
            Uuid::new_v4().to_string(),
        ));
        let seq_manager_2 = Arc::new(ActiveSequencesMultiWorker::new(
            component,
            block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
1257
            workers_with_configs,
1258
1259
1260
1261
1262
1263
1264
1265
1266
            true,
            Uuid::new_v4().to_string(),
        ));

        // 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
1267
        // Add request_0 to worker 0, dp_rank 0: sequence [0, 1, 2]
1268
1269
1270
1271
        seq_manager_1
            .add_request(
                "request_0".to_string(),
                Some(vec![0, 1, 2]),
1272
1273
1274
                12,   // ISL (3 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1275
                WorkerWithDpRank::new(0, 0),
1276
1277
            )
            .await?;
1278

Yan Ru Pei's avatar
Yan Ru Pei committed
1279
        // Add request_1 to worker 0, dp_rank 1: sequence [3, 4]
1280
1281
1282
1283
        seq_manager_1
            .add_request(
                "request_1".to_string(),
                Some(vec![3, 4]),
1284
1285
1286
                8,    // ISL (2 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1287
                WorkerWithDpRank::new(0, 1),
1288
1289
            )
            .await?;
1290

Yan Ru Pei's avatar
Yan Ru Pei committed
1291
        // Add request_2 to worker 1, dp_rank 0: sequence [0, 1, 2, 3] using seq_manager_2
1292
1293
1294
1295
        seq_manager_2
            .add_request(
                "request_2".to_string(),
                Some(vec![0, 1, 2, 3]),
1296
1297
1298
                16,   // ISL (4 blocks * 4 block_size)
                0,    // no overlap
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1299
                WorkerWithDpRank::new(1, 0),
1300
1301
            )
            .await?;
1302

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

1306
1307
1308
        // 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;
1309

Yan Ru Pei's avatar
Yan Ru Pei committed
1310
1311
1312
1313
1314
1315
1316
1317
1318
        // 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);

1319
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1320
1321
            blocks_phase1[&worker_0_dp0], 3,
            "Worker 0 dp_rank 0 should have 3 active blocks (from request_0)"
1322
        );
1323
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1324
1325
            blocks_phase1[&worker_0_dp1], 2,
            "Worker 0 dp_rank 1 should have 2 active blocks (from request_1)"
1326
1327
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1328
1329
            blocks_phase1[&worker_1_dp0], 4,
            "Worker 1 dp_rank 0 should have 4 active blocks (from request_2 added by seq_manager_2)"
1330
1331
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1332
1333
            tokens_phase1[&worker_0_dp0], 12,
            "Worker 0 dp_rank 0 should have 12 active tokens"
1334
        );
1335
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1336
1337
1338
1339
1340
1341
            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)"
1342
        );
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359

        // 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
1360
1361
1362
1363
1364
1365
1366
1367
        // 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 {
1368
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1369
1370
1371
                blocks_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active blocks after all requests freed",
                worker.worker_id, worker.dp_rank
1372
1373
            );
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1374
1375
1376
                tokens_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active tokens after all requests freed",
                worker.worker_id, worker.dp_rank
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
            );
        }

        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")?;
1397
        let component = namespace.component("sequences")?;
1398
1399
1400

        // 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
1401
1402
1403
1404
1405
        let mut workers_with_configs = HashMap::new();
        workers_with_configs.insert(0, None);
        workers_with_configs.insert(1, None);
        workers_with_configs.insert(2, None);

1406
1407
1408
        let seq_manager_1 = Arc::new(ActiveSequencesMultiWorker::new(
            component.clone(),
            block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
1409
            workers_with_configs.clone(),
1410
1411
1412
1413
1414
1415
            true,
            Uuid::new_v4().to_string(),
        ));
        let seq_manager_2 = Arc::new(ActiveSequencesMultiWorker::new(
            component,
            block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
1416
            workers_with_configs,
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
            true,
            Uuid::new_v4().to_string(),
        ));

        // 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
1433
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1434
                WorkerWithDpRank::from_worker_id(0),
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
            )
            .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
1445
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1446
                WorkerWithDpRank::from_worker_id(1),
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
            )
            .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
1457
                None, // expected_output_tokens
Yan Ru Pei's avatar
Yan Ru Pei committed
1458
                WorkerWithDpRank::from_worker_id(2),
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
            )
            .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
1469
1470
1471
1472
        let worker_0 = WorkerWithDpRank::from_worker_id(0);
        let worker_1 = WorkerWithDpRank::from_worker_id(1);
        let worker_2 = WorkerWithDpRank::from_worker_id(2);

1473
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1474
            tokens_phase1[&worker_0], 12,
1475
            "Worker 0 should have 12 active tokens"
1476
1477
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1478
1479
1480
1481
1482
            tokens_phase1[&worker_1], 8,
            "Worker 1 should have 8 active tokens"
        );
        assert_eq!(
            tokens_phase1[&worker_2], 16,
1483
            "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)"
1484
        );
1485

1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
        // 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
1512
            let worker = WorkerWithDpRank::from_worker_id(worker_id);
1513
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
1514
                tokens_phase2[&worker], 0,
1515
1516
1517
1518
1519
                "Worker {} should have 0 active tokens after all requests freed",
                worker_id
            );
        }

1520
        Ok(())
1521
1522
    }
}