dynamo_memory/nixl/
config.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! NIXL backend configuration with Figment support.
5//!
6//! This module provides configuration extraction for NIXL backends from
7//! environment variables with the pattern: `DYN_KVBM_NIXL_BACKEND_<backend>=<value>`
8
9use anyhow::{Result, bail};
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12
13/// Configuration for NIXL backends.
14///
15/// Supports extracting backend configurations from environment variables:
16/// - `DYN_KVBM_NIXL_BACKEND_UCX=true` - Enable UCX backend with default params
17/// - `DYN_KVBM_NIXL_BACKEND_GDS=false` - Explicitly disable GDS backend
18/// - Valid values: true/false, 1/0, on/off, yes/no (case-insensitive)
19/// - Invalid values (e.g., "maybe", "random") will cause an error
20/// - Custom params (e.g., `DYN_KVBM_NIXL_BACKEND_UCX_PARAM1=value`) will cause an error
21///
22/// # Data Structure
23///
24/// Uses a single HashMap where:
25/// - Key presence = backend is enabled
26/// - Value (inner HashMap) = backend-specific parameters (empty = defaults)
27///
28/// # TOML Example
29///
30/// ```toml
31/// [backends.UCX]
32/// # UCX with default params (empty map)
33///
34/// [backends.GDS]
35/// threads = "4"
36/// buffer_size = "1048576"
37/// ```
38#[derive(Debug, Clone, Default, Serialize, Deserialize)]
39pub struct NixlBackendConfig {
40    /// Map of backend name (uppercase) -> optional parameters.
41    ///
42    /// If a backend is present in the map, it's enabled.
43    /// The inner HashMap contains optional override parameters.
44    /// An empty inner map means use default parameters.
45    #[serde(default)]
46    backends: HashMap<String, HashMap<String, String>>,
47}
48
49impl NixlBackendConfig {
50    /// Creates a new configuration with the given backends.
51    ///
52    /// For an empty configuration with no backends, use [`Default::default()`].
53    pub fn new(backends: HashMap<String, HashMap<String, String>>) -> Self {
54        Self { backends }
55    }
56
57    /// Create configuration from environment variables.
58    ///
59    /// Extracts backends from `DYN_KVBM_NIXL_BACKEND_<backend>=<value>` variables.
60    ///
61    /// # Errors
62    /// Returns an error if:
63    /// - Custom parameters are detected (not yet supported)
64    /// - Invalid boolean values are provided (must be truthy or falsey)
65    pub fn from_env() -> Result<Self> {
66        let mut backends = HashMap::new();
67
68        // Extract all environment variables that match our pattern
69        for (key, value) in std::env::vars() {
70            if let Some(remainder) = key.strip_prefix("DYN_KVBM_NIXL_BACKEND_") {
71                // Check if there's an underscore (indicating custom params)
72                if remainder.contains('_') {
73                    bail!(
74                        "Custom NIXL backend parameters are not yet supported. \
75                         Found: {}. Please use only DYN_KVBM_NIXL_BACKEND_<backend>=true \
76                         to enable backends with default parameters.",
77                        key
78                    );
79                }
80
81                // Simple backend enablement (e.g., DYN_KVBM_NIXL_BACKEND_UCX=true).
82                // Empty or unrecognized values are rejected rather than treated
83                // as false: silently dropping a backend hides misconfiguration.
84                let backend_name = remainder.to_uppercase();
85                match crate::parse_bool_opt(&value) {
86                    Some(true) => {
87                        backends.insert(backend_name, HashMap::new());
88                    }
89                    Some(false) => {
90                        // Explicitly disabled, don't add to backends
91                        continue;
92                    }
93                    None => bail!(
94                        "Invalid value for {}: '{}'. Expected one of: true/false, 1/0, on/off, yes/no",
95                        key,
96                        value
97                    ),
98                }
99            }
100        }
101
102        Ok(Self { backends })
103    }
104
105    /// Add a backend with default parameters.
106    /// Backend name is normalized to uppercase.
107    pub fn with_backend(mut self, backend: impl Into<String>) -> Self {
108        self.backends
109            .insert(backend.into().to_uppercase(), HashMap::new());
110        self
111    }
112
113    /// Add a backend with custom parameters.
114    /// Backend name is normalized to uppercase.
115    pub fn with_backend_params(
116        mut self,
117        backend: impl Into<String>,
118        params: HashMap<String, String>,
119    ) -> Self {
120        self.backends.insert(backend.into().to_uppercase(), params);
121        self
122    }
123
124    /// Get the list of enabled backend names (uppercase).
125    pub fn backends(&self) -> Vec<String> {
126        self.backends.keys().cloned().collect()
127    }
128
129    /// Get parameters for a specific backend.
130    /// Backend name is normalized to uppercase for lookup.
131    ///
132    /// Returns None if the backend is not enabled.
133    pub fn backend_params(&self, backend: &str) -> Option<&HashMap<String, String>> {
134        self.backends.get(&backend.to_uppercase())
135    }
136
137    /// Check if a specific backend is enabled.
138    pub fn has_backend(&self, backend: &str) -> bool {
139        self.backends.contains_key(&backend.to_uppercase())
140    }
141
142    /// Merge another configuration into this one.
143    ///
144    /// Backends from the other configuration will be added to this one.
145    /// If both have the same backend, params from `other` take precedence.
146    pub fn merge(mut self, other: NixlBackendConfig) -> Self {
147        self.backends.extend(other.backends);
148        self
149    }
150
151    /// Iterate over all enabled backends and their parameters.
152    pub fn iter(&self) -> impl Iterator<Item = (&String, &HashMap<String, String>)> {
153        self.backends.iter()
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn test_new_config_is_empty() {
163        let config = NixlBackendConfig::default();
164        assert_eq!(config.backends().len(), 0);
165    }
166
167    #[test]
168    fn test_default_is_empty() {
169        let config = NixlBackendConfig::default();
170        assert!(config.backends().is_empty()); // default() has no backends
171    }
172
173    #[test]
174    fn test_with_backend() {
175        let config = NixlBackendConfig::default()
176            .with_backend("ucx")
177            .with_backend("gds_mt");
178
179        assert!(config.has_backend("ucx"));
180        assert!(config.has_backend("UCX"));
181        assert!(config.has_backend("gds_mt"));
182        assert!(config.has_backend("GDS_MT"));
183        assert!(!config.has_backend("other"));
184    }
185
186    #[test]
187    fn test_with_backend_params() {
188        let mut params = HashMap::new();
189        params.insert("threads".to_string(), "4".to_string());
190        params.insert("buffer_size".to_string(), "1048576".to_string());
191
192        let config = NixlBackendConfig::default()
193            .with_backend("UCX")
194            .with_backend_params("GDS", params);
195
196        // UCX should have empty params
197        let ucx_params = config.backend_params("UCX").unwrap();
198        assert!(ucx_params.is_empty());
199
200        // GDS should have custom params
201        let gds_params = config.backend_params("GDS").unwrap();
202        assert_eq!(gds_params.get("threads"), Some(&"4".to_string()));
203        assert_eq!(gds_params.get("buffer_size"), Some(&"1048576".to_string()));
204    }
205
206    #[test]
207    fn test_merge_configs() {
208        let config1 = NixlBackendConfig::default().with_backend("ucx");
209        let config2 = NixlBackendConfig::default().with_backend("gds");
210
211        let merged = config1.merge(config2);
212
213        assert!(merged.has_backend("ucx"));
214        assert!(merged.has_backend("gds"));
215    }
216
217    #[test]
218    fn test_backend_name_case_insensitive() {
219        let config = NixlBackendConfig::default()
220            .with_backend("ucx")
221            .with_backend("Gds_mt")
222            .with_backend("OTHER");
223
224        assert!(config.has_backend("UCX"));
225        assert!(config.has_backend("ucx"));
226        assert!(config.has_backend("GDS_MT"));
227        assert!(config.has_backend("gds_mt"));
228        assert!(config.has_backend("OTHER"));
229        assert!(config.has_backend("other"));
230    }
231
232    #[test]
233    fn test_iter() {
234        let mut params = HashMap::new();
235        params.insert("key".to_string(), "value".to_string());
236
237        let config = NixlBackendConfig::default()
238            .with_backend("UCX")
239            .with_backend_params("GDS", params);
240
241        let items: Vec<_> = config.iter().collect();
242        assert_eq!(items.len(), 2);
243    }
244
245    // Note: Testing from_env() would require setting environment variables,
246    // which is challenging in unit tests. This is better tested with integration tests.
247}
dynamo_memory/nixl/config.rs

dynamo_memory/nixl/
config.rs
