use crate::rpc::error::RpcError;
use llama_cpp_sys_4 as sys;
use std::ffi::CString;
use std::ptr::NonNull;
pub struct RpcBackend {
backend: NonNull<sys::ggml_backend>,
endpoint: String,
}
impl RpcBackend {
pub fn init(endpoint: &str) -> Result<Self, RpcError> {
let c_endpoint = CString::new(endpoint).map_err(|e| RpcError::StringConversion(e))?;
let backend = unsafe { sys::ggml_backend_rpc_init(c_endpoint.as_ptr()) };
NonNull::new(backend)
.map(|ptr| Self {
backend: ptr,
endpoint: endpoint.to_string(),
})
.ok_or_else(|| RpcError::InitializationFailed {
endpoint: endpoint.to_string(),
})
}
pub fn is_rpc(&self) -> bool {
unsafe { sys::ggml_backend_is_rpc(self.backend.as_ptr()) }
}
pub fn buffer_type(&self) -> Option<NonNull<sys::ggml_backend_buffer_type>> {
let c_endpoint = CString::new(self.endpoint.as_str()).ok()?;
let buffer_type = unsafe { sys::ggml_backend_rpc_buffer_type(c_endpoint.as_ptr()) };
NonNull::new(buffer_type)
}
pub fn get_device_memory(&self) -> Result<(usize, usize), RpcError> {
let c_endpoint =
CString::new(self.endpoint.as_str()).map_err(|e| RpcError::StringConversion(e))?;
let mut free: usize = 0;
let mut total: usize = 0;
unsafe {
sys::ggml_backend_rpc_get_device_memory(c_endpoint.as_ptr(), &mut free, &mut total);
}
if total == 0 {
Err(RpcError::MemoryQueryFailed)
} else {
Ok((free, total))
}
}
pub fn endpoint(&self) -> &str {
&self.endpoint
}
pub(crate) fn as_ptr(&self) -> NonNull<sys::ggml_backend> {
self.backend
}
}
impl Drop for RpcBackend {
fn drop(&mut self) {
unsafe {
sys::ggml_backend_free(self.backend.as_ptr());
}
}
}
unsafe impl Send for RpcBackend {}
unsafe impl Sync for RpcBackend {}
impl std::fmt::Debug for RpcBackend {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RpcBackend")
.field("endpoint", &self.endpoint)
.field("is_rpc", &self.is_rpc())
.finish()
}
}