use std::collections::VecDeque;
use std::io::Read;
use std::net::{Ipv4Addr, TcpListener};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, ExitStatus, Stdio};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use crate::hardware::RuntimeBackend;
const START_ATTEMPTS: usize = 5;
const DIAGNOSTIC_CAPACITY: usize = 128 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServerLaunchConfig {
pub executable: PathBuf,
pub model: PathBuf,
pub model_identifier: String,
pub backend: RuntimeBackend,
pub context_tokens: u32,
pub logical_cpus: usize,
pub startup_timeout: Duration,
pub health_timeout: Duration,
}
impl ServerLaunchConfig {
pub fn validate(&self) -> Result<(), ServerError> {
if !self.executable.is_file() {
return Err(ServerError::InvalidConfig(format!(
"server executable is missing: {}",
self.executable.display()
)));
}
if !self.model.is_file() {
return Err(ServerError::InvalidConfig(format!(
"model artifact is missing: {}",
self.model.display()
)));
}
if self.model_identifier.is_empty() || self.model_identifier.chars().any(char::is_control) {
return Err(ServerError::InvalidConfig(
"model identifier must be non-empty and contain no control characters".to_owned(),
));
}
if self.context_tokens == 0 {
return Err(ServerError::InvalidConfig(
"context size must be nonzero".to_owned(),
));
}
if self.logical_cpus == 0 {
return Err(ServerError::InvalidConfig(
"logical CPU count must be nonzero".to_owned(),
));
}
if self.startup_timeout.is_zero() || self.health_timeout.is_zero() {
return Err(ServerError::InvalidConfig(
"startup and health timeouts must be nonzero".to_owned(),
));
}
Ok(())
}
#[must_use]
pub fn arguments(&self, port: u16) -> Vec<String> {
let threads = self.logical_cpus.clamp(1, 16).to_string();
let gpu_layers = if self.backend == RuntimeBackend::Cpu {
"0"
} else {
"999"
};
vec![
"--host".to_owned(),
Ipv4Addr::LOCALHOST.to_string(),
"--port".to_owned(),
port.to_string(),
"--model".to_owned(),
self.model.display().to_string(),
"--alias".to_owned(),
self.model_identifier.clone(),
"--ctx-size".to_owned(),
self.context_tokens.to_string(),
"--parallel".to_owned(),
"1".to_owned(),
"--threads".to_owned(),
threads.clone(),
"--threads-batch".to_owned(),
threads,
"--n-gpu-layers".to_owned(),
gpu_layers.to_owned(),
"--flash-attn".to_owned(),
"on".to_owned(),
"--jinja".to_owned(),
"--temp".to_owned(),
"0".to_owned(),
"--spec-type".to_owned(),
"none".to_owned(),
"--no-webui".to_owned(),
]
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ManagedServerEndpoint {
pub endpoint: String,
pub port: u16,
pub pid: u32,
pub reused: bool,
}
#[derive(Debug, Default)]
pub struct ServerManager {
active: Option<ManagedServer>,
}
impl ServerManager {
#[must_use]
pub const fn new() -> Self {
Self { active: None }
}
pub fn ensure_running(
&mut self,
config: ServerLaunchConfig,
) -> Result<ManagedServerEndpoint, ServerError> {
config.validate()?;
let reusable = if let Some(active) = &mut self.active {
active.config == config
&& active.is_live()?
&& matches!(
probe_health(&active.endpoint, config.health_timeout),
Ok(true)
)
&& matches!(
probe_model_identity(
&active.endpoint,
config.health_timeout,
&config.model_identifier
),
Ok(true)
)
} else {
false
};
if reusable {
let active = self.active.as_ref().expect("reusable server is active");
return Ok(ManagedServerEndpoint {
endpoint: active.endpoint.clone(),
port: active.port,
pid: active.child.id(),
reused: true,
});
}
self.shutdown()?;
let mut last_error = None;
for _ in 0..START_ATTEMPTS {
let port = available_loopback_port()?;
match ManagedServer::start(config.clone(), port) {
Ok(server) => {
let endpoint = ManagedServerEndpoint {
endpoint: server.endpoint.clone(),
port,
pid: server.child.id(),
reused: false,
};
self.active = Some(server);
return Ok(endpoint);
}
Err(error @ ServerError::Exited { .. }) => last_error = Some(error),
Err(error) => return Err(error),
}
}
Err(last_error.unwrap_or_else(|| {
ServerError::Spawn("no loopback start attempt completed".to_owned())
}))
}
pub fn shutdown(&mut self) -> Result<(), ServerError> {
if let Some(server) = self.active.take() {
server.shutdown()?;
}
Ok(())
}
#[must_use]
pub fn diagnostics(&self) -> String {
self.active
.as_ref()
.map_or_else(String::new, |server| server.diagnostics.snapshot())
}
}
impl Drop for ServerManager {
fn drop(&mut self) {
let _ = self.shutdown();
}
}
#[derive(Debug, Error)]
pub enum ServerError {
#[error("invalid server configuration: {0}")]
InvalidConfig(String),
#[error("could not reserve a private loopback port: {0}")]
Port(std::io::Error),
#[error("could not start managed llama-server: {0}")]
Spawn(String),
#[error(
"managed llama-server exited during startup with {status}; diagnostics:\n{diagnostics}"
)]
Exited {
status: ExitStatus,
diagnostics: String,
},
#[error(
"managed llama-server was not ready after {seconds} seconds; diagnostics:\n{diagnostics}"
)]
ReadinessTimeout { seconds: u64, diagnostics: String },
#[error("managed llama-server health response was invalid: {0}")]
InvalidHealth(String),
#[error("could not inspect or stop managed llama-server: {0}")]
Process(std::io::Error),
}
#[derive(Debug)]
struct ManagedServer {
config: ServerLaunchConfig,
endpoint: String,
port: u16,
child: Child,
diagnostics: BoundedDiagnostics,
readers: Vec<JoinHandle<()>>,
}
impl ManagedServer {
fn start(config: ServerLaunchConfig, port: u16) -> Result<Self, ServerError> {
let working_directory = config.executable.parent().ok_or_else(|| {
ServerError::InvalidConfig("server executable has no parent directory".to_owned())
})?;
let mut command = Command::new(&config.executable);
command
.args(config.arguments(port))
.current_dir(working_directory)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
configure_library_path(&mut command, working_directory);
let mut child = spawn_with_short_retry(&mut command)?;
let diagnostics = BoundedDiagnostics::new(DIAGNOSTIC_CAPACITY);
let mut readers = Vec::new();
if let Some(stdout) = child.stdout.take() {
readers.push(diagnostics.capture("stdout", stdout));
}
if let Some(stderr) = child.stderr.take() {
readers.push(diagnostics.capture("stderr", stderr));
}
let endpoint = format!("http://{}:{port}", Ipv4Addr::LOCALHOST);
let started = Instant::now();
loop {
if let Some(status) = child.try_wait().map_err(ServerError::Process)? {
join_readers(&mut readers);
return Err(ServerError::Exited {
status,
diagnostics: diagnostics.snapshot(),
});
}
match probe_health(&endpoint, config.health_timeout) {
Ok(true)
if probe_model_identity(
&endpoint,
config.health_timeout,
&config.model_identifier,
)
.unwrap_or(false) =>
{
thread::sleep(Duration::from_millis(50));
if child.try_wait().map_err(ServerError::Process)?.is_none() {
break;
}
}
Ok(false) | Err(ServerError::Process(_)) => {}
Ok(true) => {}
Err(error) => {
let _ = child.kill();
let _ = child.wait();
join_readers(&mut readers);
return Err(error);
}
}
if started.elapsed() >= config.startup_timeout {
let _ = child.kill();
let _ = child.wait();
join_readers(&mut readers);
return Err(ServerError::ReadinessTimeout {
seconds: config.startup_timeout.as_secs(),
diagnostics: diagnostics.snapshot(),
});
}
thread::sleep(Duration::from_millis(100));
}
Ok(Self {
config,
endpoint,
port,
child,
diagnostics,
readers,
})
}
fn is_live(&mut self) -> Result<bool, ServerError> {
self.child
.try_wait()
.map(|status| status.is_none())
.map_err(ServerError::Process)
}
fn shutdown(mut self) -> Result<(), ServerError> {
let result = if self
.child
.try_wait()
.map_err(ServerError::Process)?
.is_none()
{
self.child
.kill()
.and_then(|()| self.child.wait().map(|_| ()))
.map_err(ServerError::Process)
} else {
Ok(())
};
join_readers(&mut self.readers);
result
}
}
impl Drop for ManagedServer {
fn drop(&mut self) {
if self.child.try_wait().ok().flatten().is_none() {
let _ = self.child.kill();
let _ = self.child.wait();
}
join_readers(&mut self.readers);
}
}
fn available_loopback_port() -> Result<u16, ServerError> {
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).map_err(ServerError::Port)?;
listener
.local_addr()
.map(|address| address.port())
.map_err(ServerError::Port)
}
fn spawn_with_short_retry(command: &mut Command) -> Result<Child, ServerError> {
let mut last_error = None;
for _ in 0..3 {
match command.spawn() {
Ok(child) => return Ok(child),
Err(error) if executable_temporarily_busy(&error) => {
last_error = Some(error);
thread::sleep(Duration::from_millis(25));
}
Err(error) => return Err(ServerError::Spawn(error.to_string())),
}
}
Err(ServerError::Spawn(last_error.map_or_else(
|| "spawn retry failed".to_owned(),
|error| error.to_string(),
)))
}
#[cfg(unix)]
fn executable_temporarily_busy(error: &std::io::Error) -> bool {
error.raw_os_error() == Some(26)
}
#[cfg(not(unix))]
fn executable_temporarily_busy(_error: &std::io::Error) -> bool {
false
}
fn probe_health(endpoint: &str, timeout: Duration) -> Result<bool, ServerError> {
let value = probe_json(endpoint, "/health", timeout)?;
Ok(value
.as_ref()
.and_then(|value| value.get("status"))
.and_then(Value::as_str)
== Some("ok"))
}
fn probe_model_identity(
endpoint: &str,
timeout: Duration,
expected_model: &str,
) -> Result<bool, ServerError> {
let value = probe_json(endpoint, "/v1/models", timeout)?;
Ok(value
.as_ref()
.and_then(|value| value.get("data"))
.and_then(Value::as_array)
.is_some_and(|models| {
models
.iter()
.any(|model| model.get("id").and_then(Value::as_str) == Some(expected_model))
}))
}
fn probe_json(endpoint: &str, path: &str, timeout: Duration) -> Result<Option<Value>, ServerError> {
let agent = ureq::AgentBuilder::new().timeout(timeout).build();
let response = match agent.get(&format!("{endpoint}{path}")).call() {
Ok(response) => response,
Err(ureq::Error::Status(503, _)) | Err(ureq::Error::Status(425, _)) => return Ok(None),
Err(ureq::Error::Status(status, _)) => {
return Err(ServerError::InvalidHealth(format!("HTTP {status}")));
}
Err(ureq::Error::Transport(error)) => {
return Err(ServerError::Process(std::io::Error::other(
error.to_string(),
)));
}
};
let body = read_bounded(response.into_reader(), 64 * 1024)
.map_err(|error| ServerError::InvalidHealth(error.to_string()))?;
let value: Value = serde_json::from_str(&body)
.map_err(|error| ServerError::InvalidHealth(error.to_string()))?;
Ok(Some(value))
}
fn read_bounded(mut reader: impl Read, maximum: usize) -> std::io::Result<String> {
let limit = u64::try_from(maximum + 1).expect("health response bound fits u64");
let mut bytes = Vec::new();
reader.by_ref().take(limit).read_to_end(&mut bytes)?;
if bytes.len() > maximum {
return Err(std::io::Error::other("health response exceeded size limit"));
}
String::from_utf8(bytes)
.map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
}
#[cfg(target_os = "linux")]
fn configure_library_path(command: &mut Command, runtime_directory: &Path) {
let value = std::env::var_os("LD_LIBRARY_PATH").map_or_else(
|| runtime_directory.as_os_str().to_owned(),
|existing| {
let mut paths = vec![runtime_directory.to_path_buf()];
paths.extend(std::env::split_paths(&existing));
std::env::join_paths(paths).unwrap_or_else(|_| runtime_directory.as_os_str().to_owned())
},
);
command.env("LD_LIBRARY_PATH", value);
}
#[cfg(target_os = "macos")]
fn configure_library_path(command: &mut Command, runtime_directory: &Path) {
command.env("DYLD_LIBRARY_PATH", runtime_directory);
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn configure_library_path(_command: &mut Command, _runtime_directory: &Path) {}
#[derive(Debug, Clone)]
struct BoundedDiagnostics {
bytes: Arc<Mutex<VecDeque<u8>>>,
capacity: usize,
}
impl BoundedDiagnostics {
fn new(capacity: usize) -> Self {
Self {
bytes: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
capacity,
}
}
fn capture(
&self,
source: &'static str,
mut reader: impl Read + Send + 'static,
) -> JoinHandle<()> {
let diagnostics = self.clone();
thread::spawn(move || {
let mut buffer = [0_u8; 4096];
loop {
match reader.read(&mut buffer) {
Ok(0) | Err(_) => break,
Ok(count) => {
diagnostics.push(format!("[{source}] ").as_bytes());
diagnostics.push(&buffer[..count]);
}
}
}
})
}
fn push(&self, bytes: &[u8]) {
let mut stored = self
.bytes
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
stored.extend(bytes);
while stored.len() > self.capacity {
stored.pop_front();
}
}
fn snapshot(&self) -> String {
let stored = self
.bytes
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let bytes: Vec<u8> = stored.iter().copied().collect();
String::from_utf8_lossy(&bytes).into_owned()
}
}
fn join_readers(readers: &mut Vec<JoinHandle<()>>) {
for reader in readers.drain(..) {
let _ = reader.join();
}
}
#[cfg(test)]
mod tests {
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use std::time::Duration;
use super::{RuntimeBackend, ServerLaunchConfig, ServerManager, probe_health};
#[test]
fn launch_arguments_are_private_bounded_and_non_speculative() {
let config = ServerLaunchConfig {
executable: "/runtime/llama-server".into(),
model: "/models/model.gguf".into(),
model_identifier: "neohorse".to_owned(),
backend: RuntimeBackend::Rocm,
context_tokens: 8192,
logical_cpus: 64,
startup_timeout: Duration::from_secs(1),
health_timeout: Duration::from_secs(1),
};
let arguments = config.arguments(32123);
assert!(
arguments
.windows(2)
.any(|pair| pair == ["--host", "127.0.0.1"])
);
assert!(arguments.windows(2).any(|pair| pair == ["--port", "32123"]));
assert!(arguments.windows(2).any(|pair| pair == ["--threads", "16"]));
assert!(
arguments
.windows(2)
.any(|pair| pair == ["--spec-type", "none"])
);
assert!(arguments.contains(&"--no-webui".to_owned()));
}
#[test]
fn readiness_requires_the_expected_health_document() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept");
let mut request = [0_u8; 4096];
let _ = stream.read(&mut request);
let body = r#"{"status":"ok"}"#;
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
)
.expect("response");
});
assert!(
probe_health(&format!("http://{address}"), Duration::from_secs(1)).expect("health")
);
}
#[test]
fn configuration_rejects_missing_files_before_spawn() {
let config = ServerLaunchConfig {
executable: "/missing/llama-server".into(),
model: "/missing/model.gguf".into(),
model_identifier: "neohorse".to_owned(),
backend: RuntimeBackend::Cpu,
context_tokens: 8192,
logical_cpus: 1,
startup_timeout: Duration::from_secs(1),
health_timeout: Duration::from_secs(1),
};
assert!(config.validate().is_err());
}
#[cfg(unix)]
#[test]
fn owns_reuses_and_shuts_down_only_its_child() {
use std::os::unix::fs::PermissionsExt;
let directory = tempfile::tempdir().expect("tempdir");
let executable = directory.path().join("fake-llama-server");
std::fs::write(
&executable,
r#"#!/usr/bin/env python3
import json
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer
port = int(sys.argv[sys.argv.index("--port") + 1])
class Handler(BaseHTTPRequestHandler):
def do_GET(self):
if self.path == "/v1/models":
value = {"data": [{"id": "fixture"}]}
else:
value = {"status": "ok"}
body = json.dumps(value).encode()
self.send_response(200)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *args):
pass
HTTPServer(("127.0.0.1", port), Handler).serve_forever()
"#,
)
.expect("script");
let mut permissions = std::fs::metadata(&executable)
.expect("metadata")
.permissions();
permissions.set_mode(0o700);
std::fs::set_permissions(&executable, permissions).expect("permissions");
let model = directory.path().join("model.gguf");
std::fs::write(&model, b"fixture").expect("model");
let config = ServerLaunchConfig {
executable,
model,
model_identifier: "fixture".to_owned(),
backend: RuntimeBackend::Cpu,
context_tokens: 128,
logical_cpus: 2,
startup_timeout: Duration::from_secs(3),
health_timeout: Duration::from_millis(250),
};
let mut manager = ServerManager::new();
let first = manager.ensure_running(config.clone()).expect("start");
assert!(!first.reused);
let second = manager.ensure_running(config).expect("reuse");
assert!(second.reused);
assert_eq!(first.pid, second.pid);
assert_eq!(first.endpoint, second.endpoint);
manager.shutdown().expect("shutdown");
assert!(manager.diagnostics().is_empty());
}
}