llama_cpp_4/rpc/server.rs
1//! RPC server for hosting backends
2
3use crate::rpc::error::RpcError;
4use llama_cpp_sys_4 as sys;
5use std::ffi::CString;
6use std::path::Path;
7use std::ptr::NonNull;
8
9/// Serve one or more local devices over RPC.
10///
11/// **This call blocks.** `ggml_backend_rpc_start_server` runs the accept loop on
12/// the calling thread and does not return while the server is live, so run it on
13/// a dedicated thread if the caller needs to stay responsive.
14///
15/// # Arguments
16/// * `endpoint` - Address to listen on (e.g. `"0.0.0.0:50052"`)
17/// * `cache_dir` - Directory for the server-side tensor cache, or `None` to
18/// disable caching
19/// * `n_threads` - Worker threads used to service requests
20/// * `devices` - Local devices to expose; clients address them by index, in the
21/// order given here
22///
23/// # Errors
24///
25/// Returns [`RpcError::StringConversion`] if `endpoint` or `cache_dir` contains
26/// an interior NUL, [`RpcError::InvalidEndpoint`] if `cache_dir` is not valid
27/// UTF-8, and [`RpcError::ServerError`] if `devices` is empty or exceeds
28/// llama.cpp's server limit.
29///
30/// # Example
31/// ```no_run
32/// use llama_cpp_4::rpc::serve;
33///
34/// // `devices` comes from the ggml backend registry.
35/// # let devices: Vec<std::ptr::NonNull<llama_cpp_sys_4::ggml_backend_device>> = vec![];
36/// serve("0.0.0.0:50052", None, 4, &devices)?;
37/// # Ok::<(), llama_cpp_4::rpc::RpcError>(())
38/// ```
39pub fn serve(
40 endpoint: &str,
41 cache_dir: Option<&Path>,
42 n_threads: usize,
43 devices: &[NonNull<sys::ggml_backend_device>],
44) -> Result<(), RpcError> {
45 if devices.is_empty() {
46 return Err(RpcError::ServerError {
47 message: "at least one device must be exposed".to_owned(),
48 });
49 }
50 if devices.len() > sys::GGML_RPC_MAX_SERVERS as usize {
51 return Err(RpcError::ServerError {
52 message: format!(
53 "{} devices requested but llama.cpp serves at most {}",
54 devices.len(),
55 sys::GGML_RPC_MAX_SERVERS
56 ),
57 });
58 }
59
60 let c_endpoint = CString::new(endpoint)?;
61 let c_cache_dir = cache_dir
62 .map(|dir| {
63 let dir = dir.to_str().ok_or_else(|| RpcError::InvalidEndpoint {
64 endpoint: dir.display().to_string(),
65 })?;
66 CString::new(dir).map_err(RpcError::from)
67 })
68 .transpose()?;
69
70 // `ggml_backend_dev_t` is a raw pointer, so a `NonNull` slice has the same
71 // layout; copy into an owned Vec rather than casting to keep that implicit.
72 let mut device_ptrs: Vec<sys::ggml_backend_dev_t> =
73 devices.iter().map(|d| d.as_ptr()).collect();
74
75 unsafe {
76 sys::ggml_backend_rpc_start_server(
77 c_endpoint.as_ptr(),
78 c_cache_dir
79 .as_ref()
80 .map_or(std::ptr::null(), |d| d.as_ptr()),
81 n_threads,
82 device_ptrs.len(),
83 device_ptrs.as_mut_ptr(),
84 );
85 }
86
87 Ok(())
88}
89
90/// Register a remote RPC server with the ggml backend registry.
91///
92/// Returns the backend *registration* covering every device the endpoint
93/// exposes; use [`RpcBackend::init`](crate::rpc::RpcBackend::init) to bind one
94/// of those devices.
95///
96/// # Errors
97///
98/// Returns [`RpcError::StringConversion`] if `endpoint` contains an interior
99/// NUL, or [`RpcError::InitializationFailed`] if the endpoint could not be
100/// registered.
101pub fn add_rpc_server(endpoint: &str) -> Result<NonNull<sys::ggml_backend_reg>, RpcError> {
102 let c_endpoint = CString::new(endpoint)?;
103
104 let reg = unsafe { sys::ggml_backend_rpc_add_server(c_endpoint.as_ptr()) };
105
106 NonNull::new(reg).ok_or_else(|| RpcError::InitializationFailed {
107 endpoint: endpoint.to_string(),
108 })
109}