llama_cpp_2/
llama_backend.rs

1//! Representation of an initialized llama backend
2
3use crate::LLamaCppError;
4use llama_cpp_sys_2::ggml_log_level;
5use std::sync::atomic::AtomicBool;
6use std::sync::atomic::Ordering::SeqCst;
7
8/// Representation of an initialized llama backend
9/// This is required as a parameter for most llama functions as the backend must be initialized
10/// before any llama functions are called. This type is proof of initialization.
11#[derive(Eq, PartialEq, Debug)]
12pub struct LlamaBackend {}
13
14static LLAMA_BACKEND_INITIALIZED: AtomicBool = AtomicBool::new(false);
15
16impl LlamaBackend {
17    /// Mark the llama backend as initialized
18    fn mark_init() -> crate::Result<()> {
19        match LLAMA_BACKEND_INITIALIZED.compare_exchange(false, true, SeqCst, SeqCst) {
20            Ok(_) => Ok(()),
21            Err(_) => Err(LLamaCppError::BackendAlreadyInitialized),
22        }
23    }
24
25    /// Initialize the llama backend (without numa).
26    ///
27    /// # Examples
28    ///
29    /// ```
30    ///# use llama_cpp_2::llama_backend::LlamaBackend;
31    ///# use llama_cpp_2::LLamaCppError;
32    ///# use std::error::Error;
33    ///
34    ///# fn main() -> Result<(), Box<dyn Error>> {
35    ///
36    ///
37    /// let backend = LlamaBackend::init()?;
38    /// // the llama backend can only be initialized once
39    /// assert_eq!(Err(LLamaCppError::BackendAlreadyInitialized), LlamaBackend::init());
40    ///
41    ///# Ok(())
42    ///# }
43    /// ```
44    #[tracing::instrument(skip_all)]
45    pub fn init() -> crate::Result<LlamaBackend> {
46        Self::mark_init()?;
47        unsafe { llama_cpp_sys_2::llama_backend_init() }
48        Ok(LlamaBackend {})
49    }
50
51    /// Initialize the llama backend (with numa).
52    /// ```
53    ///# use llama_cpp_2::llama_backend::LlamaBackend;
54    ///# use std::error::Error;
55    ///# use llama_cpp_2::llama_backend::NumaStrategy;
56    ///
57    ///# fn main() -> Result<(), Box<dyn Error>> {
58    ///
59    /// let llama_backend = LlamaBackend::init_numa(NumaStrategy::MIRROR)?;
60    ///
61    ///# Ok(())
62    ///# }
63    /// ```
64    #[tracing::instrument(skip_all)]
65    pub fn init_numa(strategy: NumaStrategy) -> crate::Result<LlamaBackend> {
66        Self::mark_init()?;
67        unsafe {
68            llama_cpp_sys_2::llama_numa_init(llama_cpp_sys_2::ggml_numa_strategy::from(strategy));
69        }
70        Ok(LlamaBackend {})
71    }
72
73    /// Was the code built for a GPU backend & is a supported one available.
74    pub fn supports_gpu_offload(&self) -> bool {
75        unsafe { llama_cpp_sys_2::llama_supports_gpu_offload() }
76    }
77
78    /// Does this platform support loading the model via mmap.
79    pub fn supports_mmap(&self) -> bool {
80        unsafe { llama_cpp_sys_2::llama_supports_mmap() }
81    }
82
83    /// Does this platform support locking the model in RAM.
84    pub fn supports_mlock(&self) -> bool {
85        unsafe { llama_cpp_sys_2::llama_supports_mlock() }
86    }
87
88    /// Change the output of llama.cpp's logging to be voided instead of pushed to `stderr`.
89    pub fn void_logs(&mut self) {
90        unsafe extern "C" fn void_log(
91            _level: ggml_log_level,
92            _text: *const ::std::os::raw::c_char,
93            _user_data: *mut ::std::os::raw::c_void,
94        ) {
95        }
96
97        unsafe {
98            llama_cpp_sys_2::llama_log_set(Some(void_log), std::ptr::null_mut());
99        }
100    }
101}
102
103/// A rusty wrapper around `numa_strategy`.
104#[derive(Debug, Eq, PartialEq, Copy, Clone)]
105pub enum NumaStrategy {
106    /// The numa strategy is disabled.
107    DISABLED,
108    /// help wanted: what does this do?
109    DISTRIBUTE,
110    /// help wanted: what does this do?
111    ISOLATE,
112    /// help wanted: what does this do?
113    NUMACTL,
114    /// help wanted: what does this do?
115    MIRROR,
116    /// help wanted: what does this do?
117    COUNT,
118}
119
120/// An invalid numa strategy was provided.
121#[derive(Debug, Eq, PartialEq, Copy, Clone)]
122pub struct InvalidNumaStrategy(
123    /// The invalid numa strategy that was provided.
124    pub llama_cpp_sys_2::ggml_numa_strategy,
125);
126
127impl TryFrom<llama_cpp_sys_2::ggml_numa_strategy> for NumaStrategy {
128    type Error = InvalidNumaStrategy;
129
130    fn try_from(value: llama_cpp_sys_2::ggml_numa_strategy) -> Result<Self, Self::Error> {
131        match value {
132            llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISABLED => Ok(Self::DISABLED),
133            llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISTRIBUTE => Ok(Self::DISTRIBUTE),
134            llama_cpp_sys_2::GGML_NUMA_STRATEGY_ISOLATE => Ok(Self::ISOLATE),
135            llama_cpp_sys_2::GGML_NUMA_STRATEGY_NUMACTL => Ok(Self::NUMACTL),
136            llama_cpp_sys_2::GGML_NUMA_STRATEGY_MIRROR => Ok(Self::MIRROR),
137            llama_cpp_sys_2::GGML_NUMA_STRATEGY_COUNT => Ok(Self::COUNT),
138            value => Err(InvalidNumaStrategy(value)),
139        }
140    }
141}
142
143impl From<NumaStrategy> for llama_cpp_sys_2::ggml_numa_strategy {
144    fn from(value: NumaStrategy) -> Self {
145        match value {
146            NumaStrategy::DISABLED => llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISABLED,
147            NumaStrategy::DISTRIBUTE => llama_cpp_sys_2::GGML_NUMA_STRATEGY_DISTRIBUTE,
148            NumaStrategy::ISOLATE => llama_cpp_sys_2::GGML_NUMA_STRATEGY_ISOLATE,
149            NumaStrategy::NUMACTL => llama_cpp_sys_2::GGML_NUMA_STRATEGY_NUMACTL,
150            NumaStrategy::MIRROR => llama_cpp_sys_2::GGML_NUMA_STRATEGY_MIRROR,
151            NumaStrategy::COUNT => llama_cpp_sys_2::GGML_NUMA_STRATEGY_COUNT,
152        }
153    }
154}
155
156/// Drops the llama backend.
157/// ```
158///
159///# use llama_cpp_2::llama_backend::LlamaBackend;
160///# use std::error::Error;
161///
162///# fn main() -> Result<(), Box<dyn Error>> {
163/// let backend = LlamaBackend::init()?;
164/// drop(backend);
165/// // can be initialized again after being dropped
166/// let backend = LlamaBackend::init()?;
167///# Ok(())
168///# }
169///
170/// ```
171impl Drop for LlamaBackend {
172    fn drop(&mut self) {
173        match LLAMA_BACKEND_INITIALIZED.compare_exchange(true, false, SeqCst, SeqCst) {
174            Ok(_) => {}
175            Err(_) => {
176                unreachable!("This should not be reachable as the only ways to obtain a llama backend involve marking the backend as initialized.")
177            }
178        }
179        unsafe { llama_cpp_sys_2::llama_backend_free() }
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn numa_from_and_to() {
189        let numas = [
190            NumaStrategy::DISABLED,
191            NumaStrategy::DISTRIBUTE,
192            NumaStrategy::ISOLATE,
193            NumaStrategy::NUMACTL,
194            NumaStrategy::MIRROR,
195            NumaStrategy::COUNT,
196        ];
197
198        for numa in &numas {
199            let from = llama_cpp_sys_2::ggml_numa_strategy::from(*numa);
200            let to = NumaStrategy::try_from(from).expect("Failed to convert from and to");
201            assert_eq!(*numa, to);
202        }
203    }
204
205    #[test]
206    fn check_invalid_numa() {
207        let invalid = 800;
208        let invalid = NumaStrategy::try_from(invalid);
209        assert_eq!(invalid, Err(InvalidNumaStrategy(invalid.unwrap_err().0)));
210    }
211}
llama_cpp_2/llama_backend.rs

llama_cpp_2/
llama_backend.rs
