llama_cpp_4/rpc/
backend.rs

1//! RPC backend for distributed inference
2
3use crate::rpc::error::RpcError;
4use llama_cpp_sys_4 as sys;
5use std::ffi::CString;
6use std::ptr::NonNull;
7
8/// RPC backend for distributed inference across multiple machines
9pub struct RpcBackend {
10    backend: NonNull<sys::ggml_backend>,
11    endpoint: String,
12    device: u32,
13}
14
15impl RpcBackend {
16    /// Initialize a new RPC backend for the given endpoint and remote device.
17    ///
18    /// # Arguments
19    /// * `endpoint` - The RPC server endpoint (e.g., "127.0.0.1:50052")
20    /// * `device` - Index of the device to use on the remote server. A single
21    ///   endpoint can expose several devices; pass `0` for the first.
22    ///
23    /// # Errors
24    ///
25    /// Returns [`RpcError::StringConversion`] if `endpoint` contains an interior
26    /// NUL, or [`RpcError::InitializationFailed`] if llama.cpp could not reach
27    /// the endpoint or the device index is out of range.
28    ///
29    /// # Example
30    /// ```no_run
31    /// use llama_cpp_4::rpc::RpcBackend;
32    ///
33    /// let backend = RpcBackend::init("127.0.0.1:50052", 0)?;
34    /// # Ok::<(), llama_cpp_4::rpc::RpcError>(())
35    /// ```
36    pub fn init(endpoint: &str, device: u32) -> Result<Self, RpcError> {
37        let c_endpoint = CString::new(endpoint)?;
38
39        let backend = unsafe { sys::ggml_backend_rpc_init(c_endpoint.as_ptr(), device) };
40
41        NonNull::new(backend)
42            .map(|ptr| Self {
43                backend: ptr,
44                endpoint: endpoint.to_string(),
45                device,
46            })
47            .ok_or_else(|| RpcError::InitializationFailed {
48                endpoint: endpoint.to_string(),
49            })
50    }
51
52    /// Check if a backend is an RPC backend
53    #[must_use]
54    pub fn is_rpc(&self) -> bool {
55        unsafe { sys::ggml_backend_is_rpc(self.backend.as_ptr()) }
56    }
57
58    /// Get the buffer type for this RPC backend
59    #[must_use]
60    pub fn buffer_type(&self) -> Option<NonNull<sys::ggml_backend_buffer_type>> {
61        let c_endpoint = CString::new(self.endpoint.as_str()).ok()?;
62        let buffer_type =
63            unsafe { sys::ggml_backend_rpc_buffer_type(c_endpoint.as_ptr(), self.device) };
64        NonNull::new(buffer_type)
65    }
66
67    /// Query the available memory on the remote device
68    ///
69    /// Returns (`free_memory`, `total_memory`) in bytes.
70    ///
71    /// # Errors
72    ///
73    /// Returns [`RpcError::MemoryQueryFailed`] when the server reports a total
74    /// of zero, which is how an unreachable endpoint surfaces here.
75    pub fn get_device_memory(&self) -> Result<(usize, usize), RpcError> {
76        let c_endpoint = CString::new(self.endpoint.as_str())?;
77
78        let mut free: usize = 0;
79        let mut total: usize = 0;
80
81        unsafe {
82            sys::ggml_backend_rpc_get_device_memory(
83                c_endpoint.as_ptr(),
84                self.device,
85                std::ptr::from_mut(&mut free),
86                std::ptr::from_mut(&mut total),
87            );
88        }
89
90        if total == 0 {
91            Err(RpcError::MemoryQueryFailed)
92        } else {
93            Ok((free, total))
94        }
95    }
96
97    /// Get the endpoint this backend is connected to
98    #[must_use]
99    pub fn endpoint(&self) -> &str {
100        &self.endpoint
101    }
102
103    /// Index of the remote device this backend is bound to.
104    #[must_use]
105    pub fn device(&self) -> u32 {
106        self.device
107    }
108
109    /// Get the raw backend pointer for FFI calls.
110    ///
111    /// The pointer is owned by this `RpcBackend` and is freed on drop; do not
112    /// free it, and do not use it after this value goes out of scope.
113    #[must_use]
114    pub fn as_ptr(&self) -> NonNull<sys::ggml_backend> {
115        self.backend
116    }
117}
118
119impl Drop for RpcBackend {
120    fn drop(&mut self) {
121        unsafe {
122            sys::ggml_backend_free(self.backend.as_ptr());
123        }
124    }
125}
126
127// Safety: RpcBackend can be sent between threads
128unsafe impl Send for RpcBackend {}
129// Safety: RpcBackend can be shared between threads (the C API is thread-safe)
130unsafe impl Sync for RpcBackend {}
131
132impl std::fmt::Debug for RpcBackend {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("RpcBackend")
135            .field("endpoint", &self.endpoint)
136            .field("device", &self.device)
137            .field("is_rpc", &self.is_rpc())
138            // `backend` is an opaque llama.cpp pointer with nothing useful to show.
139            .finish_non_exhaustive()
140    }
141}
llama_cpp_4/rpc/backend.rs

llama_cpp_4/rpc/
backend.rs
