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

4
//! Runtime-specific glue for [`ActiveSequencesMultiWorker`].
5
//!
6
7
8
9
10
11
12
13
//! This module provides the concrete [`SequencePublisher`] and [`SequenceSubscriber`]
//! implementations that wire the runtime-agnostic business logic (in `dynamo_kv_router`)
//! to NATS event transport and Prometheus metrics.

pub use dynamo_kv_router::multi_worker_sequence::{
    ActiveSequencesMultiWorker, SequenceError, SequencePublisher, SequenceRequest,
    SequenceSubscriber,
};
14
use dynamo_kv_router::protocols::{ActiveLoad, ActiveSequenceEvent, WorkerWithDpRank};
15
pub use dynamo_kv_router::sequence::{ActiveSequences, RequestId};
16

17
18
19
use anyhow::Result;
use dynamo_runtime::component::Component;
use dynamo_runtime::traits::DistributedRuntimeProvider;
20
use dynamo_runtime::transports::event_plane::{EventPublisher, EventSubscriber};
21
use std::collections::HashMap;
22
23
use std::sync::Arc;

24
use super::metrics::WORKER_LOAD_METRICS;
25
use crate::kv_router::{ACTIVE_SEQUENCES_SUBJECT, KV_METRICS_SUBJECT};
Yan Ru Pei's avatar
Yan Ru Pei committed
26
use crate::local_model::runtime_config::ModelRuntimeConfig;
27

28
29
/// Concrete [`SequencePublisher`] backed by NATS [`EventPublisher`] and Prometheus gauges.
pub struct RuntimeSequencePublisher {
30
    event_publisher: EventPublisher,
31
    metrics_publisher: Arc<EventPublisher>,
32
33
}

34
35
36
impl SequencePublisher for RuntimeSequencePublisher {
    async fn publish_event(&self, event: &ActiveSequenceEvent) -> anyhow::Result<()> {
        self.event_publisher.publish(event).await
37
38
    }

39
    fn publish_load(&self, load: ActiveLoad) {
40
41
        let publisher = self.metrics_publisher.clone();
        tokio::spawn(async move {
42
            if let Err(e) = publisher.publish(&load).await {
43
                tracing::trace!(
44
45
46
                    "Failed to publish ActiveLoad to NATS for worker (id={}, dp_rank={}): {e:?}",
                    load.worker_id,
                    load.dp_rank
47
48
49
                );
            }
        });
50
51
    }

52
    fn observe_load(
53
        &self,
54
55
56
57
58
59
60
61
62
63
64
65
        worker: &WorkerWithDpRank,
        worker_type: &str,
        blocks: usize,
        tokens: usize,
    ) {
        WORKER_LOAD_METRICS.observe(
            worker.worker_id,
            worker.dp_rank,
            worker_type,
            blocks,
            tokens,
        );
66
    }
67
}
68

69
70
71
72
/// Concrete [`SequenceSubscriber`] backed by NATS typed event stream.
pub struct RuntimeSequenceSubscriber {
    inner: dynamo_runtime::transports::event_plane::TypedEventSubscriber<ActiveSequenceEvent>,
}
73

74
75
76
77
78
impl SequenceSubscriber for RuntimeSequenceSubscriber {
    async fn next_event(&mut self) -> Option<anyhow::Result<ActiveSequenceEvent>> {
        match self.inner.next().await? {
            Ok((_envelope, event)) => Some(Ok(event)),
            Err(e) => Some(Err(e)),
79
        }
80
    }
81
}
82

83
84
/// Type alias for the runtime-wired multi-worker sequence tracker.
pub type ActiveSequencesMulti = ActiveSequencesMultiWorker<RuntimeSequencePublisher>;
85

86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
/// Convenience async constructor that creates the NATS publishers/subscribers
/// and returns an `Arc<ActiveSequencesMulti>` with replica sync already running.
pub async fn create_multi_worker_sequences(
    component: Component,
    block_size: usize,
    workers_with_configs: HashMap<u64, ModelRuntimeConfig>,
    replica_sync: bool,
    router_id: u64,
    worker_type: &'static str,
) -> Result<Arc<ActiveSequencesMulti>> {
    let event_publisher =
        EventPublisher::for_component(&component, ACTIVE_SEQUENCES_SUBJECT).await?;
    let metrics_publisher =
        Arc::new(EventPublisher::for_namespace(component.namespace(), KV_METRICS_SUBJECT).await?);

    let publisher = RuntimeSequencePublisher {
        event_publisher,
        metrics_publisher,
    };

106
    let dp_range: HashMap<u64, (u32, u32)> = workers_with_configs
107
        .into_iter()
108
109
110
111
112
113
        .map(|(id, config)| {
            (
                id,
                (config.data_parallel_start_rank, config.data_parallel_size),
            )
        })
114
115
116
117
118
        .collect();

    let multi_worker = ActiveSequencesMultiWorker::new(
        publisher,
        block_size,
119
        dp_range,
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        replica_sync,
        router_id,
        worker_type,
    );

    let arc = Arc::new(multi_worker);

    if replica_sync {
        let subscriber = EventSubscriber::for_component(&component, ACTIVE_SEQUENCES_SUBJECT)
            .await?
            .typed::<ActiveSequenceEvent>();
        let subscriber = RuntimeSequenceSubscriber { inner: subscriber };
        let cancel_token = component.drt().runtime().child_token();
        arc.start_replica_sync(subscriber, cancel_token);
134
    }
135

136
137
138
    let expiry_cancel = component.drt().runtime().child_token();
    arc.start_periodic_force_expiry_across_all_workers(expiry_cancel);

139
    Ok(arc)
140
141
142
143
144
}

#[cfg(test)]
mod tests {
    use super::*;
145
    use dynamo_runtime::{DistributedRuntime, Runtime};
146
    use tokio::time::Instant;
147

148
149
150
151
    #[test]
    fn test_active_sequences_shared_blocks() {
        let block_size = 4;
        let mut seq_manager = ActiveSequences::new(block_size);
152
        let decay_now = Instant::now();
153

154
155
156
157
158
159
160
161
        seq_manager.add_request(
            "request_1".to_string(),
            Some(vec![1, 2, 3]),
            12,
            0,
            None,
            decay_now,
        );
162
        assert_eq!(seq_manager.active_blocks(), 3);
163
        assert_eq!(seq_manager.active_tokens(decay_now), 12);
164

165
166
167
168
169
170
171
172
        seq_manager.add_request(
            "request_2".to_string(),
            Some(vec![4]),
            4,
            0,
            None,
            decay_now,
        );
173
        assert_eq!(seq_manager.active_blocks(), 4);
174
175
176
177
178
179
180
181
182
183
        assert_eq!(seq_manager.active_tokens(decay_now), 16);

        seq_manager.add_request(
            "request_3".to_string(),
            Some(vec![1, 2, 3, 4]),
            16,
            4,
            None,
            decay_now,
        );
184
        assert_eq!(seq_manager.active_blocks(), 4);
185
        assert_eq!(seq_manager.active_tokens(decay_now), 16);
186

187
        seq_manager.free(&"request_2".to_string(), decay_now);
188
        assert_eq!(seq_manager.active_blocks(), 4);
189
        assert_eq!(seq_manager.active_tokens(decay_now), 12);
190

191
        seq_manager.free(&"request_3".to_string(), decay_now);
192
        assert_eq!(seq_manager.active_blocks(), 3);
193
        assert_eq!(seq_manager.active_tokens(decay_now), 12);
194

195
        seq_manager.free(&"request_1".to_string(), decay_now);
196
        assert_eq!(seq_manager.active_blocks(), 0);
197
        assert_eq!(seq_manager.active_tokens(decay_now), 0);
198
199
    }

200
    #[tokio::test]
201
    #[ignore]
202
    async fn test_multi_worker_cross_instance_sync() -> Result<()> {
203
204
        dynamo_runtime::logging::init();

205
        let block_size = 4;
206

207
208
        let runtime = Runtime::from_current()?;
        let distributed = DistributedRuntime::from_settings(runtime.clone()).await?;
209

210
        let namespace = distributed.namespace("test_cross_instance_sync")?;
211
        let component = namespace.component("sequences")?;
212

Yan Ru Pei's avatar
Yan Ru Pei committed
213
214
215
216
        let mut workers_with_configs = HashMap::new();

        let mut config_worker_0 = crate::local_model::runtime_config::ModelRuntimeConfig::new();
        config_worker_0.data_parallel_size = 2;
217
        workers_with_configs.insert(0, config_worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
218
219

        let config_worker_1 = crate::local_model::runtime_config::ModelRuntimeConfig::new();
220
        workers_with_configs.insert(1, config_worker_1);
Yan Ru Pei's avatar
Yan Ru Pei committed
221

222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
        let seq_manager_1 = create_multi_worker_sequences(
            component.clone(),
            block_size,
            workers_with_configs.clone(),
            true,
            1,
            crate::discovery::WORKER_TYPE_DECODE,
        )
        .await?;
        let seq_manager_2 = create_multi_worker_sequences(
            component,
            block_size,
            workers_with_configs,
            true,
            2,
            crate::discovery::WORKER_TYPE_DECODE,
        )
        .await?;
240
241
242

        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        seq_manager_1.add_request(
            SequenceRequest {
                request_id: "request_0".to_string(),
                token_sequence: Some(vec![0, 1, 2]),
                isl: 12,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::new(0, 0),
                lora_name: None,
            },
            Instant::now(),
        )?;

        seq_manager_1.add_request(
            SequenceRequest {
                request_id: "request_1".to_string(),
                token_sequence: Some(vec![3, 4]),
                isl: 8,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::new(0, 1),
                lora_name: None,
            },
            Instant::now(),
        )?;

        seq_manager_2.add_request(
            SequenceRequest {
                request_id: "request_2".to_string(),
                token_sequence: Some(vec![0, 1, 2, 3]),
                isl: 16,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::new(1, 0),
                lora_name: None,
            },
            Instant::now(),
        )?;
287

288
        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;
289

290
        let blocks_phase1 = seq_manager_1.active_blocks();
291
        let tokens_phase1 = seq_manager_1.active_tokens(Instant::now());
292

Yan Ru Pei's avatar
Yan Ru Pei committed
293
294
295
296
        let worker_0_dp0 = WorkerWithDpRank::new(0, 0);
        let worker_0_dp1 = WorkerWithDpRank::new(0, 1);
        let worker_1_dp0 = WorkerWithDpRank::new(1, 0);

297
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
298
299
            blocks_phase1[&worker_0_dp0], 3,
            "Worker 0 dp_rank 0 should have 3 active blocks (from request_0)"
300
        );
301
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
302
303
            blocks_phase1[&worker_0_dp1], 2,
            "Worker 0 dp_rank 1 should have 2 active blocks (from request_1)"
304
305
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
306
307
            blocks_phase1[&worker_1_dp0], 4,
            "Worker 1 dp_rank 0 should have 4 active blocks (from request_2 added by seq_manager_2)"
308
309
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
310
311
            tokens_phase1[&worker_0_dp0], 12,
            "Worker 0 dp_rank 0 should have 12 active tokens"
312
        );
313
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
314
315
316
317
318
319
            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)"
320
        );
321

322
        seq_manager_1.free(&"request_2".to_string(), Instant::now())?;
323

324
325
        seq_manager_2.free(&"request_0".to_string(), Instant::now())?;
        seq_manager_2.free(&"request_1".to_string(), Instant::now())?;
326
327
328

        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

329
        let blocks_phase2 = seq_manager_2.active_blocks();
330
        let tokens_phase2 = seq_manager_2.active_tokens(Instant::now());
331

Yan Ru Pei's avatar
Yan Ru Pei committed
332
333
334
335
336
337
338
        let all_workers = vec![
            WorkerWithDpRank::new(0, 0),
            WorkerWithDpRank::new(0, 1),
            WorkerWithDpRank::new(1, 0),
        ];

        for worker in all_workers {
339
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
340
341
342
                blocks_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active blocks after all requests freed",
                worker.worker_id, worker.dp_rank
343
344
            );
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
345
346
347
                tokens_phase2[&worker], 0,
                "Worker (id={}, dp_rank={}) should have 0 active tokens after all requests freed",
                worker.worker_id, worker.dp_rank
348
349
350
351
352
353
354
355
356
357
358
            );
        }

        Ok(())
    }

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

359
        let block_size = 4;
360
361
362
363
364

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

        let namespace = distributed.namespace("test_no_token_seq_sync")?;
365
        let component = namespace.component("sequences")?;
366

Yan Ru Pei's avatar
Yan Ru Pei committed
367
        let mut workers_with_configs = HashMap::new();
368
369
370
371
372
373
374
375
376
377
378
379
        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
380

381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
        let seq_manager_1 = create_multi_worker_sequences(
            component.clone(),
            block_size,
            workers_with_configs.clone(),
            true,
            1,
            crate::discovery::WORKER_TYPE_DECODE,
        )
        .await?;
        let seq_manager_2 = create_multi_worker_sequences(
            component,
            block_size,
            workers_with_configs,
            true,
            2,
            crate::discovery::WORKER_TYPE_DECODE,
        )
        .await?;
399
400
401

        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
        seq_manager_1.add_request(
            SequenceRequest {
                request_id: "request_0".to_string(),
                token_sequence: None,
                isl: 12,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::from_worker_id(0),
                lora_name: None,
            },
            Instant::now(),
        )?;

        seq_manager_1.add_request(
            SequenceRequest {
                request_id: "request_1".to_string(),
                token_sequence: None,
                isl: 8,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::from_worker_id(1),
                lora_name: None,
            },
            Instant::now(),
        )?;

        seq_manager_2.add_request(
            SequenceRequest {
                request_id: "request_2".to_string(),
                token_sequence: None,
                isl: 16,
                overlap: 0,
                track_prefill_tokens: true,
                expected_output_tokens: None,
                prefill_load_hint: None,
                worker: WorkerWithDpRank::from_worker_id(2),
                lora_name: None,
            },
            Instant::now(),
        )?;
446
447
448

        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

449
        let tokens_phase1 = seq_manager_1.active_tokens(Instant::now());
450

Yan Ru Pei's avatar
Yan Ru Pei committed
451
452
453
454
        let worker_0 = WorkerWithDpRank::from_worker_id(0);
        let worker_1 = WorkerWithDpRank::from_worker_id(1);
        let worker_2 = WorkerWithDpRank::from_worker_id(2);

455
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
456
            tokens_phase1[&worker_0], 12,
457
            "Worker 0 should have 12 active tokens"
458
459
        );
        assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
460
461
462
463
464
            tokens_phase1[&worker_1], 8,
            "Worker 1 should have 8 active tokens"
        );
        assert_eq!(
            tokens_phase1[&worker_2], 16,
465
            "Worker 2 should have 16 active tokens (from request_2 added by seq_manager_2)"
466
        );
467

468
469
        seq_manager_1.mark_prefill_completed(&"request_2".to_string(), Instant::now())?;
        seq_manager_1.free(&"request_2".to_string(), Instant::now())?;
470

471
472
473
474
        seq_manager_2.mark_prefill_completed(&"request_0".to_string(), Instant::now())?;
        seq_manager_2.mark_prefill_completed(&"request_1".to_string(), Instant::now())?;
        seq_manager_2.free(&"request_0".to_string(), Instant::now())?;
        seq_manager_2.free(&"request_1".to_string(), Instant::now())?;
475
476
477

        tokio::time::sleep(tokio::time::Duration::from_millis(300)).await;

478
        let tokens_phase2 = seq_manager_2.active_tokens(Instant::now());
479
480

        for worker_id in 0..=2 {
Yan Ru Pei's avatar
Yan Ru Pei committed
481
            let worker = WorkerWithDpRank::from_worker_id(worker_id);
482
            assert_eq!(
Yan Ru Pei's avatar
Yan Ru Pei committed
483
                tokens_phase2[&worker], 0,
484
485
486
487
488
                "Worker {} should have 0 active tokens after all requests freed",
                worker_id
            );
        }

489
        Ok(())
490
491
    }
}