kvbm_engine/runtime/
builder.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Builder for KvbmRuntime with optional pre-built components.
5
6use std::sync::Arc;
7
8use anyhow::Result;
9use dynamo_memory::nixl::NixlAgent;
10use kvbm_config::KvbmConfig;
11use tokio::runtime::{Handle, Runtime};
12use velo::Messenger;
13
14/// Runtime handle - either owned or borrowed.
15pub enum RuntimeHandle {
16    /// Owned runtime (created by builder).
17    Owned(Arc<Runtime>),
18    /// Borrowed handle (external runtime).
19    Handle(Handle),
20}
21
22impl RuntimeHandle {
23    /// Get a handle to the runtime.
24    pub fn handle(&self) -> Handle {
25        match self {
26            RuntimeHandle::Owned(rt) => rt.handle().clone(),
27            RuntimeHandle::Handle(h) => h.clone(),
28        }
29    }
30}
31
32/// Builder for KvbmRuntime with optional pre-built components.
33///
34/// The builder allows injecting pre-built components or building them from config:
35/// - If a component is provided, it's used directly
36/// - If not provided, the component is built from the config
37pub struct KvbmRuntimeBuilder {
38    config: KvbmConfig,
39    runtime: Option<RuntimeHandle>,
40    messenger: Option<Arc<Messenger>>,
41    nixl_agent: Option<NixlAgent>,
42}
43
44impl KvbmRuntimeBuilder {
45    /// Create builder from config.
46    pub fn new(config: KvbmConfig) -> Self {
47        Self {
48            config,
49            runtime: None,
50            messenger: None,
51            nixl_agent: None,
52        }
53    }
54
55    /// Create builder from environment.
56    pub fn from_env() -> Result<Self, kvbm_config::ConfigError> {
57        Ok(Self::new(KvbmConfig::from_env()?))
58    }
59
60    /// Create builder from JSON config string (merged with env/files).
61    ///
62    /// JSON has highest priority - overrides env vars, TOML files, and defaults.
63    /// This is the primary entrypoint for vLLM's `kv_connector_extra_config` dict.
64    pub fn from_json(json: &str) -> Result<Self, kvbm_config::ConfigError> {
65        Ok(Self::new(KvbmConfig::from_figment_with_json(json)?))
66    }
67
68    /// Use an existing tokio Runtime (takes ownership via Arc).
69    pub fn with_runtime(mut self, runtime: Arc<Runtime>) -> Self {
70        self.runtime = Some(RuntimeHandle::Owned(runtime));
71        self
72    }
73
74    /// Use an existing tokio Handle (borrowed).
75    pub fn with_runtime_handle(mut self, handle: Handle) -> Self {
76        self.runtime = Some(RuntimeHandle::Handle(handle));
77        self
78    }
79
80    /// Use an existing Messenger instance.
81    pub fn with_messenger(mut self, messenger: Arc<Messenger>) -> Self {
82        self.messenger = Some(messenger);
83        self
84    }
85
86    /// Use an existing NixlAgent instance.
87    pub fn with_nixl_agent(mut self, agent: NixlAgent) -> Self {
88        self.nixl_agent = Some(agent);
89        self
90    }
91
92    /// Build runtime for leader role.
93    pub async fn build_leader(self) -> Result<super::KvbmRuntime> {
94        self.build_internal().await
95    }
96
97    /// Build runtime for worker role.
98    pub async fn build_worker(self) -> Result<super::KvbmRuntime> {
99        self.build_internal().await
100    }
101
102    async fn build_internal(self) -> Result<super::KvbmRuntime> {
103        // 1. Tokio runtime - use provided or build from config
104        let runtime = match self.runtime {
105            Some(rt) => rt,
106            None => RuntimeHandle::Owned(Arc::new(self.config.tokio.build_runtime()?)),
107        };
108
109        // 2. Messenger - use provided or build from config (BEFORE NixL)
110        let messenger = match self.messenger {
111            Some(m) => m,
112            None => self.config.messenger.build_messenger().await?,
113        };
114
115        // 3. NixL - use provided or build from config (AFTER Messenger)
116        //    Only build if config.nixl is Some (NixL enabled)
117        let nixl_agent = match self.nixl_agent {
118            Some(agent) => Some(agent),
119            None => match &self.config.nixl {
120                Some(nixl_config) => {
121                    let agent_name = format!("nixl-{}", messenger.instance_id());
122                    let backend_config = nixl_config.clone().into();
123                    Some(NixlAgent::from_nixl_backend_config(
124                        &agent_name,
125                        backend_config,
126                    )?)
127                }
128                None => None, // NixL disabled
129            },
130        };
131
132        Ok(super::KvbmRuntime {
133            config: self.config,
134            runtime,
135            messenger,
136            nixl_agent,
137        })
138    }
139}
kvbm_engine/runtime/builder.rs

kvbm_engine/runtime/
builder.rs
