dynamo-kv-router 1.4.0

KV Router - Radix tree for LLM KV cache routing
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Standalone HTTP slot tracker.
//!
//! Hosts an Axum HTTP server that exposes manual worker registration,
//! request-lifecycle updates, and advisory load reads. The service intentionally
//! stays independent of Dynamo runtime and LLM-layer dependencies.

pub mod registry;
pub mod server;

use std::sync::Arc;

use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;

use crate::services::common::replica_sync::{PeerManager, setup_replica_sync};
use registry::SlotTrackerRegistry;
use server::{AppState, create_router};

pub struct SlotTrackerConfig {
    pub port: u16,
    pub replica_sync_port: Option<u16>,
    pub replica_sync_peers: Vec<String>,
}

pub async fn run_server(config: SlotTrackerConfig) -> anyhow::Result<()> {
    let cancel_token = CancellationToken::new();
    let shutdown_token = cancel_token.clone();
    tokio::spawn(async move {
        tokio::signal::ctrl_c().await.ok();
        tracing::info!("Received shutdown signal");
        shutdown_token.cancel();
    });

    let replica_runtime = setup_replica_sync(
        config.replica_sync_port,
        &config.replica_sync_peers,
        cancel_token.child_token(),
    )?;
    let (registry, peer_manager) = if let Some(replica_runtime) = &replica_runtime {
        let replica_config = replica_runtime.config();
        let process_id = replica_config.process_id();
        let registry = Arc::new(SlotTrackerRegistry::new_with_replica_sync(
            cancel_token.clone(),
            replica_config,
        ));
        let dispatch_registry = Arc::downgrade(&registry);
        let peer_manager = Arc::new(PeerManager::start(
            config.replica_sync_peers,
            cancel_token.child_token(),
            move |event| {
                if let Some(registry) = dispatch_registry.upgrade() {
                    registry.dispatch_replica_event(event);
                }
            },
        )?);
        tracing::info!(
            port = config.port,
            replica_sync_port = config.replica_sync_port,
            process_id,
            "Starting standalone slot tracker with replica sync"
        );
        (registry, Some(peer_manager))
    } else {
        tracing::info!(
            port = config.port,
            "Starting standalone slot tracker (HTTP-only mode)"
        );
        (
            Arc::new(SlotTrackerRegistry::new(cancel_token.clone())),
            None,
        )
    };

    let app = create_router(
        Arc::new(AppState { registry }),
        peer_manager.as_ref().map(Arc::clone),
    );
    let listener = TcpListener::bind(("0.0.0.0", config.port)).await?;
    tracing::info!("HTTP server listening on 0.0.0.0:{}", config.port);
    let result = axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            cancel_token.cancelled().await;
            tracing::info!("Received shutdown signal, stopping HTTP server");
        })
        .await;
    if let (Some(peer_manager), Some(replica_runtime)) = (&peer_manager, &replica_runtime) {
        tokio::join!(peer_manager.shutdown(), replica_runtime.shutdown());
    }
    result?;

    Ok(())
}