From 62a92a88bd4d1259c6438692aca5ebd8634de77e Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 5 Aug 2026 14:45:44 +0200 Subject: [PATCH] feat: auto-start llama-server, extract server module Checks whether llama-server is already healthy on startup and spawns it from configured binary/model paths if not, polling until ready. Server infra config (binary, model path, host, port, context size) split out of prompts.toml into its own server.toml, and all of it lives in a new server module rather than inline in main.rs, alongside a single reused HTTP client and a shared health-check helper. Gemma client setup now runs concurrently with the server health-check/spawn since they're independent. --- Cargo.lock | 194 +++++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 1 + src/main.rs | 9 ++- src/server.rs | 68 +++++++++++++++++ src/server.toml | 6 ++ 5 files changed, 273 insertions(+), 5 deletions(-) create mode 100644 src/server.rs create mode 100644 src/server.toml diff --git a/Cargo.lock b/Cargo.lock index 4cbc4da..849e952 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -253,6 +253,7 @@ name = "doubleo7" version = "0.1.0" dependencies = [ "anyhow", + "reqwest 0.12.28", "rig-core", "serde", "tokio", @@ -325,6 +326,21 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -605,6 +621,22 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-tls" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" +dependencies = [ + "bytes", + "http-body-util", + "hyper", + "hyper-util", + "native-tls", + "tokio", + "tokio-native-tls", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -831,6 +863,12 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.2" @@ -897,6 +935,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "native-tls" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + [[package]] name = "nom" version = "7.1.3" @@ -922,12 +977,49 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "openssl" +version = "0.10.81" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77823a27f0babb03091cb9ed9ef80af3b39dbc82f97e8fa530374b7dafd87a45" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "openssl-probe" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "openssl-sys" +version = "0.9.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b47e7e6bb2c38cd930d25a23b40fa52e068c10e85f3e03a7f5ba5aaca5713695" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "ordered-float" version = "5.3.0" @@ -1196,6 +1288,46 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "encoding_rs", + "futures-core", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-tls", + "hyper-util", + "js-sys", + "log", + "mime", + "native-tls", + "percent-encoding", + "pin-project-lite", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-native-tls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "reqwest" version = "0.13.4" @@ -1260,7 +1392,7 @@ dependencies = [ "mime_guess", "ordered-float", "pin-project-lite", - "reqwest", + "reqwest 0.13.4", "rig-derive", "schemars", "serde", @@ -1314,6 +1446,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.43" @@ -1395,6 +1540,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "same-file" version = "1.0.6" @@ -1542,6 +1693,18 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + [[package]] name = "sha1" version = "0.10.7" @@ -1682,6 +1845,19 @@ dependencies = [ "libc", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.19" @@ -1755,6 +1931,16 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "tokio-native-tls" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +dependencies = [ + "native-tls", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -2013,6 +2199,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "version_check" version = "0.9.5" diff --git a/Cargo.toml b/Cargo.toml index a00633b..44059af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] anyhow = "1.0.104" +reqwest = "0.12" rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" } serde = "1.0.229" tokio = { version = "1.53.1", features = ["full"] } diff --git a/src/main.rs b/src/main.rs index 36e863e..fe5af3e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,8 +13,7 @@ use rig_core::serde_json::json; use crate::models::Prompts; mod models; - -const LLAMA_SERVER_URL: &str = "http://127.0.0.1:8000"; +mod server; static PROMPTS: LazyLock = LazyLock::new(|| { toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml") @@ -22,7 +21,9 @@ static PROMPTS: LazyLock = LazyLock::new(|| { #[tokio::main] async fn main() -> anyhow::Result<()> { - let gemma: ollama::CompletionModel = wire_gemma_client().await?; + // Independent setup steps (talk to unrelated backends, no data dependency) — run concurrently. + let (gemma, ()) = tokio::try_join!(wire_gemma_client(), server::ensure_running())?; + let shieldstral = wire_shieldstral().await?; let prompt_to_test = gemma.completion_request(&PROMPTS.gemma.prompt) @@ -53,7 +54,7 @@ async fn wire_gemma_client() -> anyhow::Result { } async fn wire_shieldstral() -> anyhow::Result> { - let client = llamafile::Client::from_url(LLAMA_SERVER_URL)?; + let client = llamafile::Client::from_url(&server::url())?; // Name doesn't matter here, server just uses whatever is running on it. let shieldstral = client.completion_model("shieldstral"); diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..f3f197d --- /dev/null +++ b/src/server.rs @@ -0,0 +1,68 @@ +use std::sync::LazyLock; +use std::time::Duration; +use serde::Deserialize; +use tokio::process::Command; +use tokio::time::sleep; + +#[derive(Deserialize)] +struct ServerConfig { + llama_server: LlamaServerConfig, +} + +#[derive(Deserialize)] +struct LlamaServerConfig { + binary: String, + model_path: String, + host: String, + port: u16, + context_size: u32, +} + +static SERVER_CONFIG: LazyLock = LazyLock::new(|| { + toml::from_str(include_str!("server.toml")).expect("Could not parse server.toml") +}); + +static HTTP_CLIENT: LazyLock = LazyLock::new(reqwest::Client::new); + +pub(crate) fn url() -> String { + format!("http://{}:{}", SERVER_CONFIG.llama_server.host, SERVER_CONFIG.llama_server.port) +} + +async fn is_healthy(health_url: &str) -> bool { + HTTP_CLIENT.get(health_url).send().await.is_ok_and(|r| r.status().is_success()) +} + +/// Checks whether llama-server is already serving on the configured host/port, +/// and if not, spawns it from the configured binary/model path and waits for +/// it to report healthy before returning. +pub(crate) async fn ensure_running() -> anyhow::Result<()> { + let base_url = url(); + let health_url = format!("{base_url}/health"); + + if is_healthy(&health_url).await { + return Ok(()); + } + + println!("llama-server not running at {base_url}, starting it..."); + + Command::new(&SERVER_CONFIG.llama_server.binary) + .args([ + "-m", &SERVER_CONFIG.llama_server.model_path, + "--jinja", + "-c", &SERVER_CONFIG.llama_server.context_size.to_string(), + "--host", &SERVER_CONFIG.llama_server.host, + "--port", &SERVER_CONFIG.llama_server.port.to_string(), + ]) + .spawn() + .map_err(|e| anyhow::anyhow!("failed to spawn llama-server at {}: {e}", SERVER_CONFIG.llama_server.binary))?; + + for _ in 0..60 { + if is_healthy(&health_url).await { + println!("llama-server is up."); + return Ok(()); + } + sleep(Duration::from_secs(1)).await; + } + + anyhow::bail!("llama-server did not become healthy within 60s") +} diff --git a/src/server.toml b/src/server.toml new file mode 100644 index 0000000..97bf814 --- /dev/null +++ b/src/server.toml @@ -0,0 +1,6 @@ +[llama_server] +binary = "/home/austin/.local/share/llama.cpp/build/bin/llama-server" +model_path = "/home/austin/ai/Shieldstral-1.0-3B-BF16.gguf" +host = "127.0.0.1" +port = 8000 +context_size = 32768