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.
This commit is contained in:
Austin Schaefer
parent
507b25d370
commit
62a92a88bd
5 changed files with 273 additions and 5 deletions
194
Cargo.lock
generated
194
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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<Prompts> = LazyLock::new(|| {
|
||||
toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml")
|
||||
|
|
@ -22,7 +21,9 @@ static PROMPTS: LazyLock<Prompts> = 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<ollama::CompletionModel> {
|
|||
}
|
||||
|
||||
async fn wire_shieldstral() -> anyhow::Result<GenericCompletionModel<LlamafileExt>> {
|
||||
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");
|
||||
|
|
|
|||
68
src/server.rs
Normal file
68
src/server.rs
Normal file
|
|
@ -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<ServerConfig> = LazyLock::new(|| {
|
||||
toml::from_str(include_str!("server.toml")).expect("Could not parse server.toml")
|
||||
});
|
||||
|
||||
static HTTP_CLIENT: LazyLock<reqwest::Client> = 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")
|
||||
}
|
||||
6
src/server.toml
Normal file
6
src/server.toml
Normal file
|
|
@ -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
|
||||
Loading…
Reference in a new issue