dynamo_memory/nixl/
agent.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! NIXL agent wrapper and configuration.
5//!
6//! This module provides:
7//! - `NixlAgent`: Wrapper around nixl_sys::Agent that tracks initialized backends
8//! - `NixlBackendConfig`: Configuration for NIXL backends from environment variables
9
10use anyhow::Result;
11use nixl_sys::{Agent, is_stub};
12use std::collections::{HashMap, HashSet};
13
14use crate::nixl::NixlBackendConfig;
15
16/// A NIXL agent wrapper that tracks which backends were successfully initialized.
17///
18/// This wrapper provides:
19/// - Runtime validation of backend availability
20/// - Clear error messages when operations need unavailable backends
21/// - Single source of truth for backend state in tests and production
22///
23/// # Backend Tracking
24///
25/// Since `nixl_sys::Agent` doesn't provide a method to query active backends,
26/// we track them during initialization. The `available_backends` set is populated
27/// based on successful `create_backend()` calls.
28#[derive(Clone, Debug)]
29pub struct NixlAgent {
30    agent: Agent,
31    available_backends: HashSet<String>,
32}
33
34impl NixlAgent {
35    /// Create a NIXL agent without any backends.
36    pub fn new(name: &str) -> Result<Self> {
37        if is_stub() {
38            return Err(anyhow::anyhow!("NIXL is not supported in stub mode"));
39        }
40        let agent = Agent::new(name)?;
41
42        Ok(Self {
43            agent,
44            available_backends: HashSet::new(),
45        })
46    }
47
48    /// Creates a new agent configured with backends from the given config.
49    ///
50    /// This method iterates over all backends in the config and initializes them
51    /// with their associated parameters. If a backend has custom parameters defined
52    /// in the config, those are used; otherwise, default plugin parameters are used.
53    pub fn from_nixl_backend_config(name: &str, config: NixlBackendConfig) -> Result<Self> {
54        let mut agent = Self::new(name)?;
55        for (backend, params) in config.iter() {
56            agent.add_backend_with_params(backend, params)?;
57        }
58        Ok(agent)
59    }
60
61    /// Add a backend to the agent with default parameters.
62    pub fn add_backend(&mut self, backend: &str) -> Result<()> {
63        self.add_backend_with_params(backend, &HashMap::new())
64    }
65
66    /// Add a backend to the agent with optional custom parameters.
67    ///
68    /// If `custom_params` is non-empty, those parameters are used instead of
69    /// the plugin defaults. If empty, default parameters from the plugin are used.
70    ///
71    /// # Errors
72    /// Returns an error if custom parameters are provided (not yet supported until nixl_sys 0.9).
73    pub fn add_backend_with_params(
74        &mut self,
75        backend: &str,
76        custom_params: &HashMap<String, String>,
77    ) -> Result<()> {
78        let backend_upper = backend.to_uppercase();
79        if self.available_backends.contains(&backend_upper) {
80            return Ok(());
81        }
82
83        // TODO(DIS-1310): Custom params require nixl_sys 0.9+ which adds nixl_capi_params_add
84        if !custom_params.is_empty() {
85            anyhow::bail!(
86                "Custom NIXL backend parameters for {} are not yet supported. \
87                 This feature requires nixl_sys 0.9+. Params provided: {:?}",
88                backend_upper,
89                custom_params.keys().collect::<Vec<_>>()
90            );
91        }
92
93        // Get default params from plugin
94        let (_, params) = match self.agent.get_plugin_params(&backend_upper) {
95            Ok(result) => result,
96            Err(_) => anyhow::bail!("No {} plugin found", backend_upper),
97        };
98
99        match self.agent.create_backend(&backend_upper, &params) {
100            Ok(_) => {
101                self.available_backends.insert(backend_upper);
102                Ok(())
103            }
104            Err(e) => anyhow::bail!("Failed to create nixl backend: {}", e),
105        }
106    }
107
108    /// Create a NIXL agent requiring ALL specified backends to be available.
109    ///
110    /// Unlike `new_with_backends()` which continues if some backends fail, this method
111    /// will return an error if ANY backend fails to initialize. Use this in production
112    /// when specific backends are mandatory.
113    ///
114    /// # Arguments
115    /// * `name` - Agent name
116    /// * `backends` - List of backend names that MUST be available
117    ///
118    /// # Returns
119    /// A `NixlAgent` with all requested backends initialized.
120    ///
121    /// # Errors
122    /// Returns an error if:
123    /// - Agent creation fails
124    /// - Any backend fails to initialize
125    pub fn with_backends(name: &str, backends: &[&str]) -> Result<Self> {
126        let mut agent = Self::new(name)?;
127        let mut failed_backends = Vec::new();
128
129        for backend in backends {
130            let backend_upper = backend.to_uppercase();
131            match agent.add_backend(&backend_upper) {
132                Ok(_) => {
133                    tracing::debug!("Initialized NIXL backend: {}", backend_upper);
134                }
135                Err(e) => {
136                    tracing::error!("Failed to initialize {} backend: {}", backend_upper, e);
137                    failed_backends.push((backend_upper, e.to_string()));
138                }
139            }
140        }
141
142        if !failed_backends.is_empty() {
143            let error_details: Vec<String> = failed_backends
144                .iter()
145                .map(|(name, reason)| format!("{}: {}", name, reason))
146                .collect();
147
148            anyhow::bail!(
149                "Failed to initialize required backends: [{}]",
150                error_details.join(", ")
151            );
152        }
153
154        Ok(agent)
155    }
156
157    /// Get a reference to the underlying raw NIXL agent.
158    pub fn raw_agent(&self) -> &Agent {
159        &self.agent
160    }
161
162    /// Consume and return the underlying raw NIXL agent.
163    ///
164    /// **Warning**: Once consumed, backend tracking is lost. Use this only when
165    /// interfacing with code that requires `nixl_sys::Agent` directly.
166    pub fn into_raw_agent(self) -> Agent {
167        self.agent
168    }
169
170    /// Check if a specific backend is available.
171    pub fn has_backend(&self, backend: &str) -> bool {
172        self.available_backends.contains(&backend.to_uppercase())
173    }
174
175    /// Get all available backends.
176    pub fn backends(&self) -> &HashSet<String> {
177        &self.available_backends
178    }
179
180    /// Require a specific backend, returning an error if unavailable.
181    ///
182    /// Use this at the start of operations that need specific backends.
183    ///
184    /// Note: In general, you want to instantiate all your backends before you start registering memory.
185    /// We may change this to a builder pattern in the future to enforce all backends are instantiated
186    /// before you start registering memory.
187    pub fn require_backend(&self, backend: &str) -> Result<()> {
188        let backend_upper = backend.to_uppercase();
189        if self.has_backend(&backend_upper) {
190            Ok(())
191        } else {
192            anyhow::bail!(
193                "Operation requires {} backend, but it was not initialized. Available backends: {:?}",
194                backend_upper,
195                self.available_backends
196            )
197        }
198    }
199}
200
201// Delegate common methods to the underlying agent
202impl std::ops::Deref for NixlAgent {
203    type Target = Agent;
204
205    fn deref(&self) -> &Self::Target {
206        &self.agent
207    }
208}
209
210#[cfg(all(test, feature = "testing-nixl"))]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn test_agent_backend_tracking() {
216        // Try to create agent with UCX
217        let agent = NixlAgent::with_backends("test", &["UCX"]).expect("Need UCX for test");
218
219        // Should succeed if UCX is available
220        assert!(agent.has_backend("UCX"));
221        assert!(agent.has_backend("ucx")); // Case insensitive
222    }
223
224    #[test]
225    fn test_require_backend() {
226        let agent = NixlAgent::with_backends("test", &["UCX"]).expect("Need UCX for test");
227
228        // Should succeed for available backend
229        assert!(agent.require_backend("UCX").is_ok());
230
231        // Should fail for unavailable backend
232        assert!(agent.require_backend("GDS_MT").is_err());
233    }
234
235    #[test]
236    fn test_require_backends_strict() {
237        // Should succeed if UCX is available
238        let agent =
239            NixlAgent::with_backends("test_strict", &["UCX"]).expect("Failed to require backends");
240        assert!(agent.has_backend("UCX"));
241
242        // Should fail if any backend is missing (GDS likely not available)
243        let result = NixlAgent::with_backends("test_strict_fail", &["UCX", "DUDE"]);
244        assert!(result.is_err());
245    }
246
247    #[test]
248    fn test_add_backend_with_empty_params() {
249        let mut agent = NixlAgent::new("test_empty_params").expect("Failed to create agent");
250
251        // Empty params should work (uses plugin defaults)
252        let result = agent.add_backend_with_params("UCX", &HashMap::new());
253        assert!(result.is_ok());
254        assert!(agent.has_backend("UCX"));
255    }
256
257    #[test]
258    fn test_add_backend_with_custom_params_fails() {
259        let mut agent = NixlAgent::new("test_custom_params").expect("Failed to create agent");
260
261        // Custom params should fail until nixl_sys 0.9
262        let mut params = HashMap::new();
263        params.insert("some_key".to_string(), "some_value".to_string());
264
265        let result = agent.add_backend_with_params("UCX", &params);
266        assert!(result.is_err());
267
268        let err_msg = result.unwrap_err().to_string();
269        assert!(err_msg.contains("not yet supported"));
270        assert!(err_msg.contains("nixl_sys 0.9"));
271        assert!(err_msg.contains("some_key"));
272    }
273
274    #[test]
275    fn test_from_nixl_backend_config_with_custom_params_fails() {
276        // Config with custom params should fail
277        let mut params = HashMap::new();
278        params.insert("threads".to_string(), "4".to_string());
279
280        let config = NixlBackendConfig::default().with_backend_params("UCX", params);
281
282        let result = NixlAgent::from_nixl_backend_config("test_config_params", config);
283        assert!(result.is_err());
284
285        let err_msg = result.unwrap_err().to_string();
286        assert!(err_msg.contains("not yet supported"));
287        assert!(err_msg.contains("threads"));
288    }
289
290    #[test]
291    fn test_from_nixl_backend_config_with_empty_params() {
292        // Config with empty params should work
293        let config = NixlBackendConfig::default().with_backend("UCX");
294
295        let result = NixlAgent::from_nixl_backend_config("test_config_empty", config);
296        assert!(result.is_ok());
297
298        let agent = result.unwrap();
299        assert!(agent.has_backend("UCX"));
300    }
301}
dynamo_memory/nixl/agent.rs

dynamo_memory/nixl/
agent.rs
