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        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
40        // the library owns whatever it returns.
41        let backend = unsafe { sys::ggml_backend_rpc_init(c_endpoint.as_ptr(), device) };
42
43        NonNull::new(backend)
44            .map(|ptr| Self {
45                backend: ptr,
46                endpoint: endpoint.to_string(),
47                device,
48            })
49            .ok_or_else(|| RpcError::InitializationFailed {
50                endpoint: endpoint.to_string(),
51            })
52    }
53
54    /// Check if a backend is an RPC backend
55    #[must_use]
56    pub fn is_rpc(&self) -> bool {
57        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
58        // the library owns whatever it returns.
59        unsafe { sys::ggml_backend_is_rpc(self.backend.as_ptr()) }
60    }
61
62    /// Get the buffer type for this RPC backend
63    #[must_use]
64    pub fn buffer_type(&self) -> Option<NonNull<sys::ggml_backend_buffer_type>> {
65        let c_endpoint = CString::new(self.endpoint.as_str()).ok()?;
66        let buffer_type =
67            // SAFETY: the arguments are either scalars or pointers that outlive the call, and
68            // the library owns whatever it returns.
69            unsafe { sys::ggml_backend_rpc_buffer_type(c_endpoint.as_ptr(), self.device) };
70        NonNull::new(buffer_type)
71    }
72
73    /// Query the available memory on the remote device
74    ///
75    /// Returns (`free_memory`, `total_memory`) in bytes.
76    ///
77    /// # Errors
78    ///
79    /// Returns [`RpcError::MemoryQueryFailed`] when the server reports a total
80    /// of zero, which is how an unreachable endpoint surfaces here.
81    pub fn get_device_memory(&self) -> Result<(usize, usize), RpcError> {
82        let c_endpoint = CString::new(self.endpoint.as_str())?;
83
84        let mut free: usize = 0;
85        let mut total: usize = 0;
86
87        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
88        // the library owns whatever it returns.
89        unsafe {
90            sys::ggml_backend_rpc_get_device_memory(
91                c_endpoint.as_ptr(),
92                self.device,
93                std::ptr::from_mut(&mut free),
94                std::ptr::from_mut(&mut total),
95            );
96        }
97
98        if total == 0 {
99            Err(RpcError::MemoryQueryFailed)
100        } else {
101            Ok((free, total))
102        }
103    }
104
105    /// Get the endpoint this backend is connected to
106    #[must_use]
107    pub fn endpoint(&self) -> &str {
108        &self.endpoint
109    }
110
111    /// Index of the remote device this backend is bound to.
112    #[must_use]
113    pub fn device(&self) -> u32 {
114        self.device
115    }
116
117    /// Get the raw backend pointer for FFI calls.
118    ///
119    /// The pointer is owned by this `RpcBackend` and is freed on drop; do not
120    /// free it, and do not use it after this value goes out of scope.
121    #[must_use]
122    pub fn as_ptr(&self) -> NonNull<sys::ggml_backend> {
123        self.backend
124    }
125}
126
127impl Drop for RpcBackend {
128    fn drop(&mut self) {
129        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
130        // afterwards.
131        unsafe {
132            sys::ggml_backend_free(self.backend.as_ptr());
133        }
134    }
135}
136
137// Safety: RpcBackend can be sent between threads
138unsafe impl Send for RpcBackend {}
139// Safety: RpcBackend can be shared between threads (the C API is thread-safe)
140unsafe impl Sync for RpcBackend {}
141
142impl std::fmt::Debug for RpcBackend {
143    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
144        f.debug_struct("RpcBackend")
145            .field("endpoint", &self.endpoint)
146            .field("device", &self.device)
147            .field("is_rpc", &self.is_rpc())
148            // `backend` is an opaque llama.cpp pointer with nothing useful to show.
149            .finish_non_exhaustive()
150    }
151}
llama_cpp_4/rpc/backend.rs

llama_cpp_4/rpc/
backend.rs
