"examples/backends/trtllm/deploy/agg.yaml" did not exist on "158435cd9c052661d98d507c4272411c4d0f2c2a"
subscriber.rs 19.3 KB
Newer Older
1
2
3
4
5
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Background processes for the KV Router including event consumption and snapshot uploads.

6
use std::{collections::HashSet, time::Duration};
7
8
9
10

use anyhow::Result;
use dynamo_runtime::{
    component::Component,
11
    discovery::DiscoveryQuery,
12
13
14
    prelude::*,
    traits::events::EventPublisher,
    transports::{
15
        etcd::{Client as EtcdClient, WatchEvent},
16
17
18
        nats::{NatsQueue, Slug},
    },
};
19
use futures::StreamExt;
20
use rand::Rng;
21
22
23
24
25
26
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;

use crate::{
    discovery::KV_ROUTERS_ROOT_PATH,
    kv_router::{
27
        KV_EVENT_SUBJECT, RADIX_STATE_BUCKET, RADIX_STATE_FILE,
Yan Ru Pei's avatar
Yan Ru Pei committed
28
29
        indexer::{DumpRequest, GetWorkersRequest, RouterEvent},
        protocols::WorkerId,
30
31
32
    },
};

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
/// Delay between snapshot reads to verify stability
const SNAPSHOT_STABILITY_DELAY: Duration = Duration::from_millis(100);
const MAX_SNAPSHOT_STABILITY_ATTEMPTS: usize = 10;

const CHECK_INTERVAL_BASE: Duration = Duration::from_secs(1);
const CHECK_INTERVAL_JITTER_MS: i64 = 100;

/// Download a stable snapshot from object store and send events to the indexer.
/// Retries until two consecutive reads match or max attempts is reached.
async fn download_stable_snapshot(
    nats_client: &dynamo_runtime::transports::nats::Client,
    bucket_name: &str,
    kv_events_tx: &mpsc::Sender<RouterEvent>,
) -> Result<()> {
    let url = url::Url::parse(&format!(
        "nats://{}/{bucket_name}/{RADIX_STATE_FILE}",
        nats_client.addr()
    ))?;

    // Try to get initial snapshot
    let Ok(mut prev_events) = nats_client
        .object_store_download_data::<Vec<RouterEvent>>(&url)
        .await
    else {
        tracing::debug!(
            "Failed to download snapshots. This is normal for freshly started Router replicas."
        );
        return Ok(());
    };

    // Keep trying until we get two consecutive stable reads
    for attempt in 1..=MAX_SNAPSHOT_STABILITY_ATTEMPTS {
        tokio::time::sleep(SNAPSHOT_STABILITY_DELAY).await;

        let curr_events = match nats_client
            .object_store_download_data::<Vec<RouterEvent>>(&url)
            .await
        {
            Ok(events) => events,
            Err(e) => {
                tracing::warn!(
                    "Snapshot read failed on attempt {attempt}, using previous snapshot with {} events: {e:?}",
                    prev_events.len()
                );
                break;
            }
        };

        // Check if snapshot is stable (two consecutive reads match)
        if prev_events == curr_events {
            tracing::info!(
                "Successfully downloaded stable snapshot with {} events from object store (stable after {attempt} attempts)",
                curr_events.len()
            );
            prev_events = curr_events;
            break;
        }

        tracing::debug!(
            "Snapshot changed between reads on attempt {attempt} ({} -> {} events), retrying",
            prev_events.len(),
            curr_events.len()
        );
        prev_events = curr_events;

        if attempt == MAX_SNAPSHOT_STABILITY_ATTEMPTS {
            tracing::warn!(
                "Max stability attempts reached, using latest snapshot with {} events",
                prev_events.len()
            );
        }
    }

    // Send all events to the indexer
    for event in prev_events {
        if let Err(e) = kv_events_tx.send(event).await {
            tracing::warn!("Failed to send initial event to indexer: {e:?}");
        }
    }
    tracing::info!("Successfully sent all initial events to indexer");

    Ok(())
}

117
118
119
120
121
/// Resources required for snapshot operations
#[derive(Clone)]
struct SnapshotResources {
    nats_client: dynamo_runtime::transports::nats::Client,
    bucket_name: String,
122
123
124
    instances_rx: tokio::sync::watch::Receiver<Vec<dynamo_runtime::component::Instance>>,
    get_workers_tx: mpsc::Sender<GetWorkersRequest>,
    snapshot_tx: mpsc::Sender<DumpRequest>,
125
126
127
}

impl SnapshotResources {
128
    /// Perform snapshot upload and purge operations
129
130
131
132
133
134
135
136
137
138
139
140
141
142
    async fn purge_then_snapshot(
        &self,
        nats_queue: &mut NatsQueue,
        remove_worker_tx: &mpsc::Sender<WorkerId>,
    ) -> anyhow::Result<()> {
        // Purge before snapshot ensures new/warm-restarted routers won't replay already-acknowledged messages.
        // Since KV events are idempotent, this ordering reduces unnecessary reprocessing while maintaining
        // at-least-once delivery guarantees. The snapshot will capture the clean state after purge.
        tracing::info!("Purging acknowledged messages and performing snapshot of radix tree");
        let start_time = std::time::Instant::now();

        // Clean up stale workers before snapshot
        // Get current worker IDs from instances_rx
        let current_instances = self.instances_rx.borrow().clone();
143
        let current_worker_ids: std::collections::HashSet<u64> = current_instances
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
            .iter()
            .map(|instance| instance.instance_id)
            .collect();

        // Get worker IDs from the indexer
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
        let get_workers_req = GetWorkersRequest { resp: resp_tx };

        if let Err(e) = self.get_workers_tx.send(get_workers_req).await {
            tracing::warn!("Failed to send get_workers request during snapshot: {e:?}");
        } else {
            match resp_rx.await {
                Ok(indexer_worker_ids) => {
                    // Find workers in indexer but not in current instances
                    for worker_id in indexer_worker_ids {
                        if !current_worker_ids.contains(&worker_id) {
                            tracing::info!(
161
                                "Removing stale worker {worker_id} from indexer during snapshot"
162
163
164
                            );
                            if let Err(e) = remove_worker_tx.send(worker_id).await {
                                tracing::warn!(
165
                                    "Failed to send remove_worker for stale worker {worker_id}: {e:?}"
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
                                );
                            }
                        }
                    }
                }
                Err(e) => {
                    tracing::warn!("Failed to receive worker IDs from indexer: {e:?}");
                }
            }
        }

        // First, purge acknowledged messages from the stream
        nats_queue.purge_acknowledged().await?;

        // Now request a snapshot from the indexer (which reflects the post-purge state)
        let (resp_tx, resp_rx) = oneshot::channel();
        let dump_req = DumpRequest { resp: resp_tx };

        self.snapshot_tx
            .send(dump_req)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to send dump request: {e:?}"))?;

        // Wait for the dump response
        let events = resp_rx
            .await
            .map_err(|e| anyhow::anyhow!("Failed to receive dump response: {e:?}"))?;

        // Upload the snapshot to NATS object store
        let url = url::Url::parse(&format!(
            "nats://{}/{}/{RADIX_STATE_FILE}",
            self.nats_client.addr(),
            self.bucket_name
        ))?;

        self.nats_client
            .object_store_upload_data(&events, &url)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to upload snapshot: {e:?}"))?;

        tracing::info!(
            "Successfully performed snapshot of radix tree with {} events to bucket {} in {}ms",
            events.len(),
            self.bucket_name,
            start_time.elapsed().as_millis()
        );

        Ok(())
    }
215
216
217
}

/// Start a unified background task for event consumption and optional snapshot management
218
#[allow(clippy::too_many_arguments)]
219
220
221
222
pub async fn start_kv_router_background(
    component: Component,
    consumer_uuid: String,
    kv_events_tx: mpsc::Sender<RouterEvent>,
223
224
225
    remove_worker_tx: mpsc::Sender<WorkerId>,
    maybe_get_workers_tx: Option<mpsc::Sender<GetWorkersRequest>>,
    maybe_snapshot_tx: Option<mpsc::Sender<DumpRequest>>,
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    cancellation_token: CancellationToken,
    router_snapshot_threshold: Option<u32>,
    router_reset_states: bool,
) -> Result<()> {
    // Set up NATS connections
    let stream_name = Slug::slugify(&format!("{}.{}", component.subject(), KV_EVENT_SUBJECT))
        .to_string()
        .replace("_", "-");
    let nats_server =
        std::env::var("NATS_SERVER").unwrap_or_else(|_| "nats://localhost:4222".to_string());

    // Create NatsQueue for event consumption
    let mut nats_queue = NatsQueue::new_with_consumer(
        stream_name.clone(),
        nats_server.clone(),
        std::time::Duration::from_secs(60), // 1 minute timeout
242
        consumer_uuid.clone(),
243
244
245
246
247
248
249
250
251
    );
    nats_queue.connect_with_reset(router_reset_states).await?;

    // Always create NATS client (needed for both reset and snapshots)
    let client_options = dynamo_runtime::transports::nats::Client::builder()
        .server(&nats_server)
        .build()?;
    let nats_client = client_options.connect().await?;

252
253
254
255
256
257
    // Get etcd client (needed for both snapshots and router watching)
    let etcd_client = component
        .drt()
        .etcd_client()
        .ok_or_else(|| anyhow::anyhow!("etcd client not available"))?;

258
259
260
261
262
263
    // Create bucket name for snapshots/state
    let bucket_name = Slug::slugify(&format!("{}-{RADIX_STATE_BUCKET}", component.subject()))
        .to_string()
        .replace("_", "-");

    // Handle initial state based on router_reset_states flag
264
265
266
267
    if !router_reset_states {
        // Try to download initial state from object store with stability check
        download_stable_snapshot(&nats_client, &bucket_name, &kv_events_tx).await?;
    } else {
268
269
270
271
272
273
274
        // Delete the bucket to reset state
        tracing::info!("Resetting router state, deleting bucket: {bucket_name}");
        if let Err(e) = nats_client.object_store_delete_bucket(&bucket_name).await {
            tracing::warn!("Failed to delete bucket (may not exist): {e:?}");
        }
    }

275
276
277
    // Cleanup orphaned consumers on startup
    cleanup_orphaned_consumers(&mut nats_queue, &etcd_client, &component, &consumer_uuid).await;

278
    // Watch for router deletions to clean up orphaned consumers
279
    let (_prefix_str, mut router_replicas_rx) = etcd_client
280
281
282
283
        .kv_get_and_watch_prefix(&format!("{}/", KV_ROUTERS_ROOT_PATH))
        .await?
        .dissolve();

284
285
    // Get the generate endpoint and watch for instance deletions
    let generate_endpoint = component.endpoint("generate");
286
287
288
289
290
291
292
293
294
    let discovery_client = component.drt().discovery();
    let discovery_key = DiscoveryQuery::Endpoint {
        namespace: component.namespace().name().to_string(),
        component: component.name().to_string(),
        endpoint: "generate".to_string(),
    };
    let mut instance_event_stream = discovery_client
        .list_and_watch(discovery_key, Some(cancellation_token.clone()))
        .await?;
295

296
297
    // Get instances_rx for tracking current workers
    let client = generate_endpoint.client().await?;
298
    let instances_rx = client.instance_source.as_ref().clone();
299
300
301
302
303
304
305

    // Only set up snapshot-related resources if snapshot_tx, get_workers_tx, and threshold are provided
    let snapshot_resources = if let (Some(get_workers_tx), Some(snapshot_tx), Some(_)) = (
        maybe_get_workers_tx,
        maybe_snapshot_tx,
        router_snapshot_threshold,
    ) {
306
307
308
        Some(SnapshotResources {
            nats_client,
            bucket_name,
309
310
311
            instances_rx,
            get_workers_tx,
            snapshot_tx,
312
313
314
315
316
        })
    } else {
        None
    };

317
    tokio::spawn(async move {
318
319
320
321
322
323
324
        // Create interval with jitter
        let jitter_ms =
            rand::rng().random_range(-CHECK_INTERVAL_JITTER_MS..=CHECK_INTERVAL_JITTER_MS);
        let interval_duration = Duration::from_millis(
            (CHECK_INTERVAL_BASE.as_millis() as i64 + jitter_ms).max(1) as u64,
        );
        let mut check_interval = tokio::time::interval(interval_duration);
325
326
327
328
329
330
331
332
333
        check_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

        loop {
            tokio::select! {
                biased;

                _ = cancellation_token.cancelled() => {
                    tracing::debug!("KV Router background task received cancellation signal");
                    // Clean up the queue and remove the durable consumer
334
                    // TODO: durable consumer cannot cleanup if ungraceful shutdown (crash)
335
336
337
338
339
340
                    if let Err(e) = nats_queue.shutdown(None).await {
                        tracing::warn!("Failed to shutdown NatsQueue: {e}");
                    }
                    break;
                }

341
                // Handle generate endpoint instance deletion events
342
343
                Some(discovery_event_result) = instance_event_stream.next() => {
                    let Ok(discovery_event) = discovery_event_result else {
344
345
346
                        continue;
                    };

347
                    let dynamo_runtime::discovery::DiscoveryEvent::Removed(worker_id) = discovery_event else {
348
349
350
                        continue;
                    };

351
352
353
354
                    tracing::warn!(
                        worker_id = worker_id,
                        "DISCOVERY: Generate endpoint instance removed, removing worker"
                    );
355
356

                    if let Err(e) = remove_worker_tx.send(worker_id).await {
357
                        tracing::warn!("Failed to send worker removal for worker {worker_id}: {e}");
358
359
360
                    }
                }

361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
                // Handle event consumption
                result = nats_queue.dequeue_task(None) => {
                    match result {
                        Ok(Some(bytes)) => {
                            let event: RouterEvent = match serde_json::from_slice(&bytes) {
                                Ok(event) => event,
                                Err(e) => {
                                    tracing::warn!("Failed to deserialize RouterEvent: {e:?}");
                                    continue;
                                }
                            };

                            // Forward the RouterEvent to the indexer
                            if let Err(e) = kv_events_tx.send(event).await {
                                tracing::warn!(
                                    "failed to send kv event to indexer; shutting down: {e:?}"
                                );
                                break;
                            }
                        },
                        Ok(None) => {
                            tracing::trace!("Dequeue timeout, continuing");
                        },
                        Err(e) => {
                            tracing::error!("Failed to dequeue task: {e:?}");
                            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                        }
                    }
                }

391
                // Handle periodic stream checking and purging (only if snapshot_resources is provided)
392
                _ = check_interval.tick() => {
393
                    let Some(resources) = snapshot_resources.as_ref() else {
394
395
396
397
398
399
400
401
402
403
                        continue;
                    };

                    // Check total messages in the stream
                    let Ok(message_count) = nats_queue.get_stream_messages().await else {
                        tracing::warn!("Failed to get stream message count");
                        continue;
                    };

                    let threshold = router_snapshot_threshold.unwrap_or(u32::MAX) as u64;
404

405
406
407
408
                    if message_count <= threshold {
                        continue;
                    }

409
                    tracing::info!("Stream has {message_count} messages (threshold: {threshold}), performing purge and snapshot");
410

411
                    match resources.purge_then_snapshot(
412
                        &mut nats_queue,
413
                        &remove_worker_tx,
414
415
                    ).await {
                        Ok(_) => tracing::info!("Successfully performed purge and snapshot"),
416
                        Err(e) => tracing::debug!("Could not perform purge and snapshot: {e:?}"),
417
418
419
420
421
422
423
424
425
426
427
                    }
                }

                // Handle router deletion events
                Some(event) = router_replicas_rx.recv() => {
                    let WatchEvent::Delete(kv) = event else {
                        // We only care about deletions for cleaning up consumers
                        continue;
                    };

                    let key = String::from_utf8_lossy(kv.key());
428
                    tracing::info!("Detected router replica deletion: {key}");
429

430
431
432
433
434
435
436
437
438
439
                    // Only process deletions for routers on the same component
                    if !key.contains(component.path().as_str()) {
                        tracing::trace!(
                            "Skipping router deletion from different component (key: {key}, subscriber component: {})",
                            component.path()
                        );
                        continue;
                    }

                    // Extract the router UUID from the key
440
                    let Some(router_uuid) = key.split('/').next_back() else {
441
                        tracing::warn!("Could not extract UUID from router key: {key}");
442
443
444
445
446
447
                        continue;
                    };

                    // The consumer UUID is the router UUID
                    let consumer_to_delete = router_uuid.to_string();

448
                    tracing::info!("Attempting to delete orphaned consumer: {consumer_to_delete}");
449

450
451
452
                    // Delete the consumer (allow race condition if multiple routers try to delete)
                    if let Err(e) = nats_queue.shutdown(Some(consumer_to_delete.clone())).await {
                        tracing::warn!("Failed to delete consumer {consumer_to_delete}: {e}");
453
                    } else {
454
                        tracing::info!("Successfully deleted orphaned consumer: {consumer_to_delete}");
455
456
457
458
459
460
461
462
463
464
465
466
467
468
                    }
                }
            }
        }

        // Clean up the queue and remove the durable consumer
        if let Err(e) = nats_queue.shutdown(None).await {
            tracing::warn!("Failed to shutdown NatsQueue: {e}");
        }
    });

    Ok(())
}

469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
/// Cleanup orphaned NATS consumers that no longer have corresponding etcd router entries
async fn cleanup_orphaned_consumers(
    nats_queue: &mut NatsQueue,
    etcd_client: &EtcdClient,
    component: &Component,
    consumer_uuid: &str,
) {
    let Ok(consumers) = nats_queue.list_consumers().await else {
        return;
    };

    let router_prefix = format!("{}/{}/", KV_ROUTERS_ROOT_PATH, component.path());
    let Ok(router_entries) = etcd_client.kv_get_prefix(&router_prefix).await else {
        return;
    };

    let active_uuids: HashSet<String> = router_entries
        .iter()
        .filter_map(|kv| {
            String::from_utf8_lossy(kv.key())
                .split('/')
                .next_back()
                .map(str::to_string)
        })
        .collect();

    for consumer in consumers {
        if consumer == consumer_uuid {
            // Never delete myself (extra/redundant safeguard)
            continue;
        }
        if !active_uuids.contains(&consumer) {
501
            tracing::info!("Cleaning up orphaned consumer: {consumer}");
502
503
504
505
            let _ = nats_queue.shutdown(Some(consumer)).await;
        }
    }
}