Skip to main content

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}