dynamo-kv-router 1.3.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

// 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_config = setup_replica_sync(
        config.replica_sync_port,
        &config.replica_sync_peers,
        cancel_token.child_token(),
    )?;
    let (registry, peer_manager) = if let Some(replica_config) = replica_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::clone(&registry);
        let peer_manager = PeerManager::start(
            config.replica_sync_peers,
            cancel_token.child_token(),
            move |event| dispatch_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);
    let listener = TcpListener::bind(("0.0.0.0", config.port)).await?;
    tracing::info!("HTTP server listening on 0.0.0.0:{}", config.port);
    axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            cancel_token.cancelled().await;
            tracing::info!("Received shutdown signal, stopping HTTP server");
        })
        .await?;

    Ok(())
}