llama_cpp_bindings/
llama_backend.rs

1//! Representation of an initialized llama backend
2
3use crate::LlamaCppError;
4use crate::llama_backend_numa_strategy::NumaStrategy;
5use llama_cpp_bindings_sys::ggml_log_level;
6use std::sync::atomic::AtomicBool;
7use std::sync::atomic::Ordering::SeqCst;
8
9/// Representation of an initialized llama backend.
10///
11/// This is required as a parameter for most llama functions as the backend must be initialized
12/// before any llama functions are called. This type is proof of initialization.
13#[derive(Eq, PartialEq, Debug)]
14pub struct LlamaBackend {}
15
16static LLAMA_BACKEND_INITIALIZED: AtomicBool = AtomicBool::new(false);
17
18impl LlamaBackend {
19    /// Mark the llama backend as initialized
20    fn mark_init() -> crate::Result<()> {
21        match LLAMA_BACKEND_INITIALIZED.compare_exchange(false, true, SeqCst, SeqCst) {
22            Ok(_was_uninitialized) => Ok(()),
23            Err(_was_already_initialized) => Err(LlamaCppError::BackendAlreadyInitialized),
24        }
25    }
26
27    /// Initialize the llama backend (without numa).
28    ///
29    /// # Examples
30    ///
31    /// ```
32    ///# use llama_cpp_bindings::llama_backend::LlamaBackend;
33    ///# use llama_cpp_bindings::LlamaCppError;
34    ///# use std::error::Error;
35    ///
36    ///# fn main() -> Result<(), Box<dyn Error>> {
37    ///
38    ///
39    /// let backend = LlamaBackend::init()?;
40    /// // the llama backend can only be initialized once
41    /// assert!(matches!(LlamaBackend::init(), Err(LlamaCppError::BackendAlreadyInitialized)));
42    ///
43    ///# Ok(())
44    ///# }
45    /// ```
46    /// # Errors
47    /// Returns an error if the backend was already initialized.
48    pub fn init() -> crate::Result<Self> {
49        Self::mark_init()?;
50        unsafe { llama_cpp_bindings_sys::llama_backend_init() }
51        Ok(Self {})
52    }
53
54    /// Initialize the llama backend (with numa).
55    /// ```
56    ///# use llama_cpp_bindings::llama_backend::LlamaBackend;
57    ///# use std::error::Error;
58    ///# use llama_cpp_bindings::llama_backend_numa_strategy::NumaStrategy;
59    ///
60    ///# fn main() -> Result<(), Box<dyn Error>> {
61    ///
62    /// let llama_backend = LlamaBackend::init_numa(NumaStrategy::Mirror)?;
63    ///
64    ///# Ok(())
65    ///# }
66    /// ```
67    /// # Errors
68    /// Returns an error if the backend was already initialized.
69    pub fn init_numa(strategy: NumaStrategy) -> crate::Result<Self> {
70        Self::mark_init()?;
71        unsafe {
72            llama_cpp_bindings_sys::llama_numa_init(
73                llama_cpp_bindings_sys::ggml_numa_strategy::from(strategy),
74            );
75        }
76        Ok(Self {})
77    }
78
79    /// Was the code built for a GPU backend & is a supported one available.
80    #[must_use]
81    pub fn supports_gpu_offload(&self) -> bool {
82        unsafe { llama_cpp_bindings_sys::llama_supports_gpu_offload() }
83    }
84
85    /// Does this platform support loading the model via mmap.
86    #[must_use]
87    pub fn supports_mmap(&self) -> bool {
88        unsafe { llama_cpp_bindings_sys::llama_supports_mmap() }
89    }
90
91    /// Does this platform support locking the model in RAM.
92    #[must_use]
93    pub fn supports_mlock(&self) -> bool {
94        unsafe { llama_cpp_bindings_sys::llama_supports_mlock() }
95    }
96
97    /// Change the output of llama.cpp's logging to be voided instead of pushed to `stderr`.
98    pub fn void_logs(&mut self) {
99        unsafe {
100            llama_cpp_bindings_sys::llama_log_set(Some(void_log), std::ptr::null_mut());
101        }
102    }
103}
104
105const unsafe extern "C" fn void_log(
106    _level: ggml_log_level,
107    _text: *const ::std::os::raw::c_char,
108    _user_data: *mut ::std::os::raw::c_void,
109) {
110}
111
112/// Drops the llama backend.
113/// ```
114///
115///# use llama_cpp_bindings::llama_backend::LlamaBackend;
116///# use std::error::Error;
117///
118///# fn main() -> Result<(), Box<dyn Error>> {
119/// let backend = LlamaBackend::init()?;
120/// drop(backend);
121/// // can be initialized again after being dropped
122/// let backend = LlamaBackend::init()?;
123///# Ok(())
124///# }
125///
126/// ```
127impl Drop for LlamaBackend {
128    fn drop(&mut self) {
129        LLAMA_BACKEND_INITIALIZED.store(false, SeqCst);
130        unsafe { llama_cpp_bindings_sys::llama_backend_free() }
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use serial_test::serial;
137
138    use super::LlamaBackend;
139    use crate::LlamaCppError;
140
141    #[test]
142    fn void_log_callback_does_not_panic() {
143        unsafe {
144            super::void_log(
145                llama_cpp_bindings_sys::GGML_LOG_LEVEL_INFO,
146                c"test".as_ptr(),
147                std::ptr::null_mut(),
148            );
149        }
150    }
151
152    #[test]
153    #[serial]
154    fn init_succeeds() {
155        let backend = LlamaBackend::init();
156        assert!(backend.is_ok());
157    }
158
159    #[test]
160    #[serial]
161    fn double_init_returns_error() {
162        let _backend = LlamaBackend::init().unwrap();
163        let second = LlamaBackend::init();
164        assert!(matches!(
165            second.unwrap_err(),
166            LlamaCppError::BackendAlreadyInitialized
167        ));
168    }
169
170    #[test]
171    #[serial]
172    fn feature_queries_return_bools() {
173        let backend = LlamaBackend::init().unwrap();
174        let _gpu = backend.supports_gpu_offload();
175        let _mmap = backend.supports_mmap();
176        let _mlock = backend.supports_mlock();
177    }
178
179    #[test]
180    #[serial]
181    fn drop_and_reinit_works() {
182        let backend = LlamaBackend::init().unwrap();
183        drop(backend);
184        let backend = LlamaBackend::init();
185        assert!(backend.is_ok());
186    }
187
188    #[test]
189    #[serial]
190    fn init_numa_succeeds() {
191        use crate::llama_backend_numa_strategy::NumaStrategy;
192
193        let backend = LlamaBackend::init_numa(NumaStrategy::Disabled);
194        assert!(backend.is_ok());
195    }
196}
llama_cpp_bindings/llama_backend.rs

llama_cpp_bindings/
llama_backend.rs
