Lumbridge Compute — Apache-2.0
ci / rust (push) Failing after 18s

This commit is contained in:
Karti Tripathi
2026-08-03 23:47:51 -07:00
commit a8c8532105
40 changed files with 6101 additions and 0 deletions
+87
View File
@@ -0,0 +1,87 @@
//! Resident Lumbridge Compute supervisor.
use anyhow::{Context, Result};
use std::path::Path;
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use crate::{gateway, lifecycle, mem, proc};
pub fn run(root: &Path, listen: &str, upstream: &str, floor: f64) -> Result<()> {
println!("Lumbridge Compute agent starting");
lifecycle::resume(root).context("resuming persisted desired Scene")?;
let (gateway_done_tx, gateway_done_rx) = mpsc::channel();
let listen_owned = listen.to_string();
let upstream_owned = upstream.to_string();
thread::spawn(move || {
let result = gateway::run(&listen_owned, &upstream_owned);
gateway_done_tx.send(result).ok();
});
println!(
"Lumbridge Compute agent supervising memory floor {:.1} GB",
floor
);
loop {
match gateway_done_rx.try_recv() {
Ok(result) => return result.context("stable gateway stopped"),
Err(mpsc::TryRecvError::Disconnected) => {
anyhow::bail!("stable gateway supervisor disconnected")
}
Err(mpsc::TryRecvError::Empty) => {}
}
enforce_memory_floor(root, floor)?;
thread::sleep(Duration::from_secs(1));
}
}
fn enforce_memory_floor(root: &Path, floor: f64) -> Result<()> {
let memory = mem::read()?;
if memory.available_gb >= floor {
return Ok(());
}
// Take the same lock a Scene transition takes, and skip this tick if a transition
// already holds it. Without this the watchdog races activation: activation is
// deliberately a stop-all-then-start-all sequence, so mid-transition the pool is
// legitimately tight and `state.procs` is being rewritten underneath us. Firing then
// would kill a model the transition just started, and then write a stale `state`
// over the transition's own — losing track of a process that is still alive.
//
// A tick skipped here costs one second. The floor is a backstop, and the transition
// holding the lock is doing its own admission checks.
let _lock = match lifecycle::TransitionLock::acquire(root) {
Ok(lock) => lock,
Err(_) => {
eprintln!(
"agent: MemAvailable {:.1} GB < floor {:.1}, but a Scene transition holds the \
lock; deferring to it for this tick",
memory.available_gb, floor
);
return Ok(());
}
};
let mut state = proc::State::load_checked(root)?;
match state.newest_alive() {
Some((id, process)) => {
eprintln!(
"agent: MemAvailable {:.1} GB < floor {:.1}; stopping newest owned model '{}' (pid {})",
memory.available_gb, floor, id, process.pid
);
proc::stop_owned(&process)?;
state.procs.remove(&id);
state.active_scene = None;
state.last_error = Some(format!(
"watchdog stopped '{id}' after MemAvailable fell to {:.1} GB",
memory.available_gb
));
state.save(root)?;
}
None => eprintln!(
"agent: MemAvailable {:.1} GB < floor {:.1}, but no identity-owned model can be stopped",
memory.available_gb, floor
),
}
Ok(())
}
+151
View File
@@ -0,0 +1,151 @@
//! Config types + loaders for the model registry and scenes.
//!
//! The registry (`registry/models.yaml`) resolves stable model *ids* to weights +
//! launch commands, and evolves as models requantize. Scenes (`scenes/*.scene.yaml`)
//! reference those ids and stay stable — the public, shareable contract.
//!
//! Several fields below are deserialized but never read by the binary
//! (`api_version`, `channel`, `entry`, `app`, `app_dir`, `tags`, `author`).
//! That is deliberate: they are the published `lumbridge/v1` manifest surface,
//! and declaring them is what makes a manifest carrying them parse rather than
//! fail. Deleting them to satisfy the lint would silently narrow the contract.
#![allow(dead_code)]
use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;
#[derive(Debug, Deserialize)]
pub struct Registry {
#[serde(rename = "apiVersion")]
pub api_version: String,
pub models: BTreeMap<String, Model>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct Model {
pub name: String,
/// Worst-case unified memory held once serving (weights + KV/cache + encoder).
pub footprint_gb: f64,
/// `stable` pins the exact weights; `latest` may resolve a newer quant.
#[serde(default)]
pub channel: Option<String>,
#[serde(default)]
pub health: Option<String>,
/// Optional marker that must appear in the health response. This
/// distinguishes mutually-exclusive models that intentionally share a port.
#[serde(default)]
pub health_contains: Option<String>,
pub serve: Serve,
}
/// How to launch a model. Fields are permissive across runtimes (vllm / diffusers /
/// python / uvicorn); only the ones a given `kind` needs are populated.
#[derive(Debug, Clone, Deserialize)]
pub struct Serve {
pub kind: String,
/// Optional runtime executable for kind-based builders (for example a
/// model-specific vLLM virtualenv). Defaults to the kind's command on PATH.
#[serde(default)]
pub executable: Option<String>,
#[serde(default)]
pub port: Option<u16>,
/// Explicit launch argv. If set, it is used verbatim (argv[0] = program) and
/// takes precedence over any `kind`-based builder. This is where launch commands
/// live — never in a scene — so downloaded scenes can't smuggle code.
#[serde(default)]
pub command: Vec<String>,
#[serde(default)]
pub weights: Option<String>,
#[serde(default)]
pub entry: Option<String>,
#[serde(default)]
pub app: Option<String>,
#[serde(default)]
pub app_dir: Option<String>,
#[serde(default)]
pub served_name: Vec<String>,
#[serde(default)]
pub args: BTreeMap<String, serde_yaml::Value>,
#[serde(default)]
pub env: BTreeMap<String, String>,
}
#[derive(Debug, Deserialize)]
pub struct Scene {
#[serde(rename = "apiVersion")]
pub api_version: String,
pub metadata: SceneMeta,
/// Stable model ids resolved via the registry.
pub models: Vec<String>,
/// Optional per-scene budget override (GB); defaults to the Governor's global budget.
#[serde(default)]
pub budget_gb: Option<f64>,
#[serde(default)]
pub activation: Option<Activation>,
}
#[derive(Debug, Deserialize)]
pub struct SceneMeta {
pub name: String,
#[serde(default)]
pub version: u32,
#[serde(default)]
pub description: String,
#[serde(default)]
pub tags: Vec<String>,
#[serde(default)]
pub author: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct Activation {
/// `footprint-asc` (default) | `listed`.
#[serde(default)]
pub order: Option<String>,
#[serde(default)]
pub wait_healthy: Option<bool>,
}
impl Registry {
pub fn load(root: &Path) -> Result<Registry> {
let p = root.join("registry/models.yaml");
let s =
fs::read_to_string(&p).with_context(|| format!("reading registry {}", p.display()))?;
serde_yaml::from_str(&s).with_context(|| format!("parsing registry {}", p.display()))
}
}
/// Load every `*.scene.yaml` under `<root>/scenes`, sorted by name.
pub fn load_scenes(root: &Path) -> Result<Vec<Scene>> {
let dir = root.join("scenes");
let mut out = Vec::new();
if !dir.exists() {
return Ok(out);
}
for entry in fs::read_dir(&dir).with_context(|| format!("reading {}", dir.display()))? {
let path = entry?.path();
let is_scene = path
.file_name()
.and_then(|n| n.to_str())
.map(|n| n.ends_with(".scene.yaml"))
.unwrap_or(false);
if is_scene {
let s = fs::read_to_string(&path)?;
let scene: Scene =
serde_yaml::from_str(&s).with_context(|| format!("parsing {}", path.display()))?;
out.push(scene);
}
}
out.sort_by(|a, b| a.metadata.name.cmp(&b.metadata.name));
Ok(out)
}
pub fn find_scene(root: &Path, name: &str) -> Result<Scene> {
load_scenes(root)?
.into_iter()
.find(|s| s.metadata.name == name)
.with_context(|| format!("no scene named '{name}' in {}/scenes", root.display()))
}
+495
View File
@@ -0,0 +1,495 @@
//! Reproducible evaluations for any OpenAI-compatible model server.
//!
//! Suites are declarative YAML. Results are append-only JSON artifacts suitable
//! for CI, regression comparisons, and future publication to an eval registry.
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
#[derive(Debug, Deserialize)]
struct Suite {
#[serde(rename = "apiVersion")]
api_version: String,
kind: String,
metadata: Metadata,
#[serde(default)]
defaults: Defaults,
cases: Vec<Case>,
}
#[derive(Debug, Deserialize)]
struct Metadata {
name: String,
#[serde(default)]
version: u32,
#[serde(default)]
description: String,
/// Part of the published suite manifest surface; parsed so a suite
/// carrying tags loads, not read by the runner itself.
#[serde(default)]
#[allow(dead_code)]
tags: Vec<String>,
}
#[derive(Debug, Default, Deserialize)]
struct Defaults {
#[serde(default = "default_max_tokens")]
max_tokens: u32,
#[serde(default = "default_temperature")]
temperature: f64,
#[serde(default = "default_repeat")]
repeat: u32,
#[serde(default)]
system: String,
}
fn default_max_tokens() -> u32 {
128
}
fn default_temperature() -> f64 {
0.0
}
fn default_repeat() -> u32 {
1
}
#[derive(Debug, Deserialize)]
struct Case {
id: String,
#[serde(default)]
category: String,
prompt: String,
#[serde(default)]
max_tokens: Option<u32>,
#[serde(default)]
assertions: Vec<Assertion>,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum Assertion {
Exact { value: String },
Contains { value: String },
ContainsAny { values: Vec<String> },
NotContains { value: String },
MaxWords { value: usize },
}
#[derive(Debug, Serialize)]
struct Artifact {
schema_version: u32,
run_id: String,
suite: String,
suite_version: u32,
model: String,
base_url: String,
started_unix_ms: u128,
summary: Summary,
samples: Vec<Sample>,
}
#[derive(Debug, Serialize)]
struct Summary {
passed: usize,
total: usize,
score: f64,
mean_ttft_ms: f64,
p50_ttft_ms: f64,
p95_ttft_ms: f64,
mean_prefill_tps: f64,
mean_decode_tps: f64,
}
#[derive(Debug, Serialize)]
struct Sample {
case_id: String,
category: String,
repetition: u32,
passed: bool,
assertion_results: Vec<bool>,
output: String,
reasoning: String,
prompt_tokens: u64,
completion_tokens: u64,
ttft_ms: f64,
total_ms: f64,
prefill_tps: f64,
decode_tps: f64,
}
struct Completion {
output: String,
reasoning: String,
prompt_tokens: u64,
completion_tokens: u64,
ttft_ms: f64,
total_ms: f64,
}
/// One line of `eval ls`. Suites themselves stay private — a caller has no
/// business reaching into cases and assertions — but the catalogue is the
/// useful part and is shared by the CLI and the MCP server.
#[derive(Debug, Clone)]
pub struct SuiteSummary {
pub name: String,
pub version: u32,
pub cases: usize,
pub description: String,
}
pub fn catalogue(root: &Path) -> Result<Vec<SuiteSummary>> {
Ok(load_all(root)?
.into_iter()
.map(|s| SuiteSummary {
name: s.metadata.name,
version: s.metadata.version,
cases: s.cases.len(),
description: s.metadata.description,
})
.collect())
}
pub fn list(root: &Path) -> Result<()> {
println!("{:<20} {:>4} {:>5} DESCRIPTION", "SUITE", "VER", "CASES");
for s in catalogue(root)? {
println!(
"{:<20} {:>4} {:>5} {}",
s.name, s.version, s.cases, s.description
);
}
Ok(())
}
pub fn run(
root: &Path,
name: &str,
base_url: &str,
model: &str,
repeat: Option<u32>,
) -> Result<()> {
let suite = load(root, name)?;
// `kuda/v1` is the pre-rename contract; still accepted so an older suite keeps running.
let known_contract = matches!(suite.api_version.as_str(), "lumbridge/v1" | "kuda/v1");
if !known_contract || suite.kind != "EvalSuite" {
bail!(
"unsupported eval contract: {}/{}",
suite.api_version,
suite.kind
);
}
let reps = repeat.unwrap_or(suite.defaults.repeat).max(1);
println!(
"eval '{}' v{}{} ({})",
suite.metadata.name, suite.metadata.version, model, base_url
);
let started = now_ms();
let mut samples = Vec::new();
for case in &suite.cases {
for repetition in 1..=reps {
print!(" {:<24} [{}/{}] ", case.id, repetition, reps);
std::io::stdout().flush().ok();
let c = stream_completion(base_url, model, &suite.defaults, case)?;
let assertion_results: Vec<bool> = case
.assertions
.iter()
.map(|a| evaluate(a, &c.output))
.collect();
let passed = assertion_results.iter().all(|v| *v);
let decode_seconds = ((c.total_ms - c.ttft_ms) / 1000.0).max(0.001);
let prefill_seconds = (c.ttft_ms / 1000.0).max(0.001);
let sample = Sample {
case_id: case.id.clone(),
category: case.category.clone(),
repetition,
passed,
assertion_results,
output: c.output,
reasoning: c.reasoning,
prompt_tokens: c.prompt_tokens,
completion_tokens: c.completion_tokens,
ttft_ms: c.ttft_ms,
total_ms: c.total_ms,
prefill_tps: c.prompt_tokens as f64 / prefill_seconds,
decode_tps: c.completion_tokens as f64 / decode_seconds,
};
println!(
"{} ttft={:.0}ms decode={:.1}tok/s",
if passed { "PASS" } else { "FAIL" },
sample.ttft_ms,
sample.decode_tps
);
samples.push(sample);
}
}
let summary = summarize(&samples);
let run_id = format!("{}-{}-{}", started, slug(name), slug(model));
let artifact = Artifact {
schema_version: 1,
run_id: run_id.clone(),
suite: suite.metadata.name,
suite_version: suite.metadata.version,
model: model.to_string(),
base_url: base_url.to_string(),
started_unix_ms: started,
summary,
samples,
};
let out_dir = root.join("eval-results");
fs::create_dir_all(&out_dir)?;
let out = out_dir.join(format!("{run_id}.json"));
fs::write(&out, serde_json::to_string_pretty(&artifact)?)?;
println!(
"\nscore {:.1}% ({}/{}) · mean TTFT {:.0}ms · prefill≈{:.1}tok/s · decode {:.1}tok/s",
artifact.summary.score * 100.0,
artifact.summary.passed,
artifact.summary.total,
artifact.summary.mean_ttft_ms,
artifact.summary.mean_prefill_tps,
artifact.summary.mean_decode_tps
);
println!("result {}", out.display());
Ok(())
}
fn stream_completion(
base_url: &str,
model: &str,
defaults: &Defaults,
case: &Case,
) -> Result<Completion> {
let url = format!("{}/chat/completions", base_url.trim_end_matches('/'));
let mut messages = Vec::new();
if !defaults.system.is_empty() {
messages.push(json!({"role":"system","content":defaults.system}));
}
messages.push(json!({"role":"user","content":case.prompt}));
let request = json!({
"model": model, "messages": messages, "stream": true,
"stream_options": {"include_usage": true},
"max_tokens": case.max_tokens.unwrap_or(defaults.max_tokens),
"temperature": defaults.temperature,
"chat_template_kwargs": {"enable_thinking": false}
});
let mut child = Command::new("curl")
.args([
"-sS",
"-N",
"-X",
"POST",
&url,
"-H",
"Content-Type: application/json",
"--data-binary",
"@-",
])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.context("starting curl (required for eval HTTP streaming)")?;
child
.stdin
.take()
.unwrap()
.write_all(request.to_string().as_bytes())?;
let start = Instant::now();
let mut first = None;
let mut output = String::new();
let mut reasoning = String::new();
let mut prompt_tokens = 0;
let mut completion_tokens = 0;
for line in BufReader::new(child.stdout.take().unwrap()).lines() {
let line = line?;
let Some(data) = line.strip_prefix("data: ") else {
continue;
};
if data == "[DONE]" {
break;
}
let v: Value = serde_json::from_str(data).context("parsing streamed completion")?;
if let Some(usage) = v.get("usage") {
prompt_tokens = usage
.get("prompt_tokens")
.and_then(Value::as_u64)
.unwrap_or(prompt_tokens);
completion_tokens = usage
.get("completion_tokens")
.and_then(Value::as_u64)
.unwrap_or(completion_tokens);
}
let delta = &v["choices"][0]["delta"];
let content = delta.get("content").and_then(Value::as_str).unwrap_or("");
let thought = delta
.get("reasoning")
.or_else(|| delta.get("reasoning_content"))
.and_then(Value::as_str)
.unwrap_or("");
if first.is_none() && (!content.is_empty() || !thought.is_empty()) {
first = Some(start.elapsed());
}
output.push_str(content);
reasoning.push_str(thought);
}
let status = child.wait()?;
if !status.success() {
bail!("completion request failed with {status}");
}
let total = start.elapsed().as_secs_f64() * 1000.0;
Ok(Completion {
output,
reasoning,
prompt_tokens,
completion_tokens,
ttft_ms: first.map(|d| d.as_secs_f64() * 1000.0).unwrap_or(total),
total_ms: total,
})
}
fn evaluate(a: &Assertion, output: &str) -> bool {
let normalized = output.trim();
match a {
Assertion::Exact { value } => normalized.eq_ignore_ascii_case(value.trim()),
Assertion::Contains { value } => normalized.to_lowercase().contains(&value.to_lowercase()),
Assertion::ContainsAny { values } => values
.iter()
.any(|v| normalized.to_lowercase().contains(&v.to_lowercase())),
Assertion::NotContains { value } => {
!normalized.to_lowercase().contains(&value.to_lowercase())
}
Assertion::MaxWords { value } => normalized.split_whitespace().count() <= *value,
}
}
fn summarize(samples: &[Sample]) -> Summary {
let mut ttft: Vec<f64> = samples.iter().map(|s| s.ttft_ms).collect();
ttft.sort_by(f64::total_cmp);
let mean = |f: fn(&Sample) -> f64| {
if samples.is_empty() {
0.0
} else {
samples.iter().map(f).sum::<f64>() / samples.len() as f64
}
};
let percentile = |p: f64| {
if ttft.is_empty() {
0.0
} else {
ttft[((ttft.len() - 1) as f64 * p).round() as usize]
}
};
let passed = samples.iter().filter(|s| s.passed).count();
Summary {
passed,
total: samples.len(),
score: if samples.is_empty() {
0.0
} else {
passed as f64 / samples.len() as f64
},
mean_ttft_ms: mean(|s| s.ttft_ms),
p50_ttft_ms: percentile(0.50),
p95_ttft_ms: percentile(0.95),
mean_prefill_tps: mean(|s| s.prefill_tps),
mean_decode_tps: mean(|s| s.decode_tps),
}
}
fn load(root: &Path, name: &str) -> Result<Suite> {
load_all(root)?
.into_iter()
.find(|s| s.metadata.name == name)
.with_context(|| format!("no eval suite named '{name}' in {}/evals", root.display()))
}
fn load_all(root: &Path) -> Result<Vec<Suite>> {
let dir = root.join("evals");
let mut suites = Vec::new();
if !dir.exists() {
return Ok(suites);
}
for entry in fs::read_dir(&dir)? {
let path: PathBuf = entry?.path();
if path
.file_name()
.and_then(|n| n.to_str())
.map(|n| n.ends_with(".eval.yaml"))
.unwrap_or(false)
{
let suite: Suite = serde_yaml::from_str(&fs::read_to_string(&path)?)
.with_context(|| format!("parsing {}", path.display()))?;
suites.push(suite);
}
}
suites.sort_by(|a, b| a.metadata.name.cmp(&b.metadata.name));
Ok(suites)
}
fn now_ms() -> u128 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
}
fn slug(s: &str) -> String {
s.chars()
.map(|c| {
if c.is_ascii_alphanumeric() {
c.to_ascii_lowercase()
} else {
'-'
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn assertions_are_case_insensitive_and_composable() {
let output = "Market risk remains, while company-specific risk falls.";
assert!(evaluate(
&Assertion::Contains {
value: "MARKET RISK".into()
},
output
));
assert!(evaluate(
&Assertion::ContainsAny {
values: vec!["idiosyncratic".into(), "company-specific".into()],
},
output
));
assert!(evaluate(
&Assertion::NotContains {
value: "markdown".into()
},
output
));
}
#[test]
fn exact_trims_but_does_not_accept_explanation() {
assert!(evaluate(&Assertion::Exact { value: "25".into() }, " 25\n"));
assert!(!evaluate(&Assertion::Exact { value: "25".into() }, "25 GB"));
}
#[test]
fn word_limit_counts_whitespace_tokens() {
assert!(evaluate(
&Assertion::MaxWords { value: 4 },
"one two three four"
));
assert!(!evaluate(
&Assertion::MaxWords { value: 3 },
"one two three four"
));
}
}
+102
View File
@@ -0,0 +1,102 @@
//! Streaming-safe TCP gateway for OpenAI-compatible model servers.
//!
//! The gateway deliberately stays below HTTP: it forwards bytes unchanged, so
//! chunked responses and server-sent-event token streams retain their timing and
//! semantics while clients keep one stable Lumbridge Compute address.
use anyhow::{Context, Result};
use std::io;
use std::net::{Shutdown, TcpListener, TcpStream};
use std::thread;
pub fn run(listen: &str, upstream: &str) -> Result<()> {
let listener = TcpListener::bind(listen)
.with_context(|| format!("binding Lumbridge gateway at {listen}"))?;
println!("Lumbridge gateway {listen} -> {upstream}");
serve_listener(listener, upstream, None)
}
fn serve_listener(
listener: TcpListener,
upstream: &str,
max_connections: Option<usize>,
) -> Result<()> {
let mut accepted = 0usize;
for incoming in listener.incoming() {
let client = match incoming {
Ok(client) => client,
Err(error) => {
eprintln!("gateway accept failed: {error}");
continue;
}
};
let upstream = upstream.to_string();
thread::spawn(move || {
if let Err(error) = proxy(client, &upstream) {
eprintln!("gateway request failed: {error:#}");
}
});
accepted += 1;
if max_connections.is_some_and(|limit| accepted >= limit) {
break;
}
}
Ok(())
}
fn proxy(mut client: TcpStream, upstream_addr: &str) -> Result<()> {
client.set_nodelay(true).ok();
let mut upstream = TcpStream::connect(upstream_addr)
.with_context(|| format!("connecting gateway upstream {upstream_addr}"))?;
upstream.set_nodelay(true).ok();
let mut client_reader = client.try_clone()?;
let mut upstream_writer = upstream.try_clone()?;
let request = thread::spawn(move || -> io::Result<u64> {
let copied = io::copy(&mut client_reader, &mut upstream_writer)?;
upstream_writer.shutdown(Shutdown::Write).ok();
Ok(copied)
});
io::copy(&mut upstream, &mut client)?;
client.shutdown(Shutdown::Write).ok();
request
.join()
.map_err(|_| anyhow::anyhow!("gateway request-copy thread panicked"))??;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Read, Write};
#[test]
fn gateway_forwards_bidirectional_bytes_without_buffering_protocols() {
let upstream = TcpListener::bind("127.0.0.1:0").unwrap();
let upstream_addr = upstream.local_addr().unwrap();
let upstream_thread = thread::spawn(move || {
let (mut socket, _) = upstream.accept().unwrap();
let mut request = [0u8; 4];
socket.read_exact(&mut request).unwrap();
assert_eq!(&request, b"ping");
socket.write_all(b"pong").unwrap();
});
let gateway = TcpListener::bind("127.0.0.1:0").unwrap();
let gateway_addr = gateway.local_addr().unwrap();
let gateway_thread = thread::spawn(move || {
serve_listener(gateway, &upstream_addr.to_string(), Some(1)).unwrap();
});
let mut client = TcpStream::connect(gateway_addr).unwrap();
client.write_all(b"ping").unwrap();
client.shutdown(Shutdown::Write).unwrap();
let mut response = Vec::new();
client.read_to_end(&mut response).unwrap();
assert_eq!(response, b"pong");
upstream_thread.join().unwrap();
gateway_thread.join().unwrap();
}
}
+223
View File
@@ -0,0 +1,223 @@
//! The Governor — the kernel of Lumbridge Compute.
//!
//! On a unified-memory box, over-commit doesn't fail gracefully: the whole machine
//! thrashes and wedges (SSH/ping included) before the OOM killer acts. The Governor
//! makes that impossible via (1) admission control against a hard budget and
//! (2) a watchdog on `MemAvailable` that kills the newest model before thrash.
//!
//! This module currently provides the *sensing* + *admission* half. The watchdog and
//! process spawn/kill land with the `up`/`activate` commands.
use crate::config::{Model, Registry};
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::time::{Duration, Instant};
/// Usable memory for models; the rest is reserved for the OS/desktop.
pub const DEFAULT_BUDGET_GB: f64 = 100.0;
/// Extra headroom required before admitting a new model.
pub const SAFETY_MARGIN_GB: f64 = 8.0;
/// If `MemAvailable` dips below this, the watchdog kills the newest model.
pub const WATCHDOG_FLOOR_GB: f64 = 3.0;
/// A model is considered "running" if its serving port accepts a connection.
pub fn port_open(port: u16) -> bool {
let addr = SocketAddr::from(([127, 0, 0, 1], port));
TcpStream::connect_timeout(&addr, Duration::from_millis(150)).is_ok()
}
pub fn is_running(m: &Model) -> bool {
let Some(port) = m.serve.port else {
return false;
};
if !port_open(port) {
return false;
}
match &m.health_contains {
Some(marker) => health_response(m, port)
.map(|response| response.contains(marker))
.unwrap_or(false),
None => true,
}
}
/// Fetch the local HTTP health endpoint without adding an HTTP client runtime.
/// Registry health URLs are deliberately localhost-only.
fn health_response(m: &Model, port: u16) -> Option<String> {
let path = m
.health
.as_deref()
.and_then(|url| url.split_once("localhost"))
.map(|(_, tail)| tail.trim_start_matches(|c: char| c.is_ascii_digit() || c == ':'))
.filter(|path| path.starts_with('/'))
.unwrap_or("/");
let addr = SocketAddr::from(([127, 0, 0, 1], port));
let mut stream = TcpStream::connect_timeout(&addr, Duration::from_millis(300)).ok()?;
stream
.set_read_timeout(Some(Duration::from_millis(500)))
.ok()?;
stream
.write_all(format!("GET {path} HTTP/1.0\r\nHost: localhost\r\n\r\n").as_bytes())
.ok()?;
let mut response = String::new();
stream.read_to_string(&mut response).ok()?;
Some(response)
}
pub fn running_ids(reg: &Registry) -> Vec<String> {
reg.models
.iter()
.filter(|(_, m)| is_running(m))
.map(|(id, _)| id.clone())
.collect()
}
/// Sum of footprints of all currently-serving registered models.
pub fn committed_gb(reg: &Registry) -> f64 {
reg.models
.values()
.filter(|m| is_running(m))
.map(|m| m.footprint_gb)
.sum()
}
/// Admission control. Two independent ceilings, and both must hold.
///
/// 1. **Declared** — `committed + add + margin <= budget`.
/// 2. **Observed** — `add + margin <= available`.
///
/// The second one is what makes the promise real. `budget_gb`, and every
/// `footprint_gb` feeding `committed_gb`, are numbers a human typed into the
/// registry. If any of them is optimistic — a model that grows past its declared
/// footprint, a KV cache larger than expected, anything started outside Compute —
/// the declared check happily passes while the box is already out of memory, and
/// on unified memory that ends in a wedged machine rather than a failed malloc.
/// `available_gb` comes from `MemAvailable`, which counts reality.
///
/// `available_gb` is `None` when sensing failed. That falls back to the declared
/// ceiling alone: refusing every start because /proc/meminfo was unreadable would
/// turn a sensing failure into a total outage.
pub fn can_admit(
add_gb: f64,
committed_gb: f64,
budget_gb: f64,
available_gb: Option<f64>,
) -> bool {
if committed_gb + add_gb + SAFETY_MARGIN_GB > budget_gb {
return false;
}
match available_gb {
Some(available) => add_gb + SAFETY_MARGIN_GB <= available,
None => true,
}
}
/// What admission will actually enforce right now, for reporting. The lower of the
/// declared headroom and the observed headroom.
pub fn headroom_gb(committed_gb: f64, budget_gb: f64, available_gb: Option<f64>) -> f64 {
let declared = budget_gb - committed_gb - SAFETY_MARGIN_GB;
match available_gb {
Some(available) => declared.min(available - SAFETY_MARGIN_GB),
None => declared,
}
}
/// Block until the exact registered model is healthy, or `timeout` elapses.
/// Models without a port are considered instantly ready.
pub fn wait_healthy(m: &Model, timeout: Duration) -> bool {
let Some(_) = m.serve.port else {
return true;
};
let start = Instant::now();
while start.elapsed() < timeout {
if is_running(m) {
return true;
}
std::thread::sleep(Duration::from_millis(400));
}
false
}
#[cfg(test)]
mod tests {
use super::*;
// A 108 GB Scene budget on a box with plenty free — the ordinary case.
const BUDGET: f64 = 108.0;
#[test]
fn admits_when_both_declared_and_observed_ceilings_allow_it() {
assert!(can_admit(66.0, 9.0, BUDGET, Some(100.0)));
}
#[test]
fn refuses_when_the_declared_budget_is_exceeded() {
// 66 + 40 + 8 margin = 114 > 108, even though the box has memory free.
assert!(!can_admit(66.0, 40.0, BUDGET, Some(100.0)));
}
#[test]
fn refuses_when_the_box_is_out_of_memory_even_though_the_paperwork_agrees() {
// This is the case the declared check alone could never catch: nothing is
// registered as committed, so the budget says there is 100 GB of room, but
// MemAvailable says 20 GB. Something outside Compute is holding the pool.
assert!(can_admit(66.0, 0.0, BUDGET, None), "declared check passes");
assert!(!can_admit(66.0, 0.0, BUDGET, Some(20.0)));
}
#[test]
fn the_safety_margin_is_enforced_against_observed_memory_too() {
// 66 GB model with exactly 66 GB free is a refusal: the margin has to fit.
assert!(!can_admit(66.0, 0.0, BUDGET, Some(66.0)));
assert!(!can_admit(
66.0,
0.0,
BUDGET,
Some(66.0 + SAFETY_MARGIN_GB - 0.1)
));
assert!(can_admit(66.0, 0.0, BUDGET, Some(66.0 + SAFETY_MARGIN_GB)));
}
#[test]
fn failed_sensing_falls_back_to_the_declared_ceiling_rather_than_refusing_everything() {
assert!(can_admit(66.0, 9.0, BUDGET, None));
// ...but it must not become a way to bypass the declared budget.
assert!(!can_admit(66.0, 40.0, BUDGET, None));
}
#[test]
fn a_zero_reading_refuses_everything_rather_than_admitting_everything() {
// parse_meminfo yields 0.0 for an unparseable /proc. That has to read as
// "no memory", not as "no constraint".
assert!(!can_admit(1.0, 0.0, BUDGET, Some(0.0)));
}
#[test]
fn headroom_never_exceeds_what_the_machine_actually_has() {
// The shape of a real regression: a generous declared budget on a smaller box. Reporting
// budget-minus-committed here promises room the next admission will refuse, and both the
// MCP tool and the HTTP API serve this number to callers that cannot check it themselves.
let box_available = 55.8;
let declared_budget = 100.0;
let reported = headroom_gb(0.0, declared_budget, Some(box_available));
assert!(
reported <= box_available,
"reported {reported} GB of headroom on a box with {box_available} GB free",
);
assert!(can_admit(
reported,
0.0,
declared_budget,
Some(box_available)
));
}
#[test]
fn headroom_reports_the_binding_constraint_not_the_generous_one() {
// Declared says 33 GB spare; the box says 12 GB spare minus margin.
assert_eq!(headroom_gb(67.0, BUDGET, None), 33.0);
assert_eq!(headroom_gb(67.0, BUDGET, Some(12.0)), 4.0);
// And the declared ceiling still wins when it is the tighter of the two.
assert_eq!(headroom_gb(100.0, BUDGET, Some(90.0)), 0.0);
}
}
+326
View File
@@ -0,0 +1,326 @@
//! A read-only HTTP control API.
//!
//! This exists so a web UI can read a node's state without shelling out to the CLI. It is a
//! second transport over the operations the MCP server already models — `collect_status`,
//! `collect_models`, `collect_scenes`, `collect_scene`, `collect_evals` — deliberately not a
//! second implementation of them, so the two surfaces cannot drift.
//!
//! Hand-rolled on `std::net`, matching `gateway.rs`. An HTTP framework would pull a dependency
//! tree an order of magnitude larger than the whole rest of this binary, to serve five routes
//! that return pre-serialised JSON.
//!
//! # What this deliberately does not do
//!
//! **It never mutates.** No activate, no stop, no eval run. A read-only surface that a browser
//! can reach is a much smaller thing to get right than one that can move a node's memory around,
//! and the read half is what a dashboard actually needs.
//!
//! # Why the defaults are what they are
//!
//! - **Loopback only.** The listen address defaults to `127.0.0.1`. There is no authentication
//! worth the name here, so a bind to `0.0.0.0` publishes your node's inventory to the network.
//! - **CORS off.** No origin is allowed unless named with `--allow-origin`. Allowing `*` would
//! let *any* page you visit read what models you run, which is a fingerprint of your machine.
//! - **The token is optional but checked in constant time.** Loopback plus an origin allowlist
//! already stops the browser attack; the token is for the case where someone puts this behind
//! a proxy anyway.
use anyhow::{Context, Result};
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::{Path, PathBuf};
use std::thread;
use std::time::Duration;
use crate::mcp::{collect_evals, collect_models, collect_scene, collect_scenes, collect_status};
/// Requests are tiny and come from localhost; anything slower than this is not a browser.
const READ_TIMEOUT: Duration = Duration::from_secs(5);
/// A request line plus headers. Anything larger is not a request we serve.
const MAX_HEAD_BYTES: usize = 8 * 1024;
pub struct Config {
pub root: PathBuf,
/// Origins permitted to read this API from a browser. Empty means none.
pub allowed_origins: Vec<String>,
/// When set, every request must carry `Authorization: Bearer <token>`.
pub token: Option<String>,
}
pub fn run(listen: &str, config: Config) -> Result<()> {
let listener = TcpListener::bind(listen)
.with_context(|| format!("binding the Lumbridge Compute API at {listen}"))?;
println!(
"Lumbridge Compute API on {listen} (read-only) · origins: {} · token: {}",
if config.allowed_origins.is_empty() {
"none".to_string()
} else {
config.allowed_origins.join(", ")
},
if config.token.is_some() {
"required"
} else {
"none"
},
);
if !listen.starts_with("127.") && !listen.starts_with("localhost") {
eprintln!(
"warning: {listen} is not loopback. This API has no authentication by default and \
reveals which models this node runs."
);
}
serve(listener, config)
}
fn serve(listener: TcpListener, config: Config) -> Result<()> {
let config = std::sync::Arc::new(config);
for incoming in listener.incoming() {
let stream = match incoming {
Ok(stream) => stream,
Err(error) => {
eprintln!("api accept failed: {error}");
continue;
}
};
let config = config.clone();
thread::spawn(move || {
if let Err(error) = handle(stream, &config) {
eprintln!("api request failed: {error:#}");
}
});
}
Ok(())
}
struct Request {
method: String,
path: String,
origin: Option<String>,
authorization: Option<String>,
}
/// Read the request line and headers. The body is ignored: every route is a GET.
fn read_request(stream: &TcpStream) -> Result<Option<Request>> {
let mut reader = BufReader::new(stream);
let mut head = String::new();
let mut total = 0usize;
loop {
let mut line = String::new();
let n = reader.read_line(&mut line)?;
if n == 0 {
return Ok(None); // client hung up
}
total += n;
if total > MAX_HEAD_BYTES {
return Ok(None);
}
if line == "\r\n" || line == "\n" {
break;
}
head.push_str(&line);
}
let mut lines = head.lines();
let Some(request_line) = lines.next() else {
return Ok(None);
};
let mut parts = request_line.split_whitespace();
let (Some(method), Some(target)) = (parts.next(), parts.next()) else {
return Ok(None);
};
let mut origin = None;
let mut authorization = None;
for line in lines {
let Some((name, value)) = line.split_once(':') else {
continue;
};
let value = value.trim().to_string();
match name.trim().to_ascii_lowercase().as_str() {
"origin" => origin = Some(value),
"authorization" => authorization = Some(value),
_ => {}
}
}
Ok(Some(Request {
method: method.to_string(),
// Query strings are not used by any route; dropping one keeps routing exact.
path: target.split('?').next().unwrap_or("/").to_string(),
origin,
authorization,
}))
}
/// Constant-time comparison so a token cannot be recovered a byte at a time from response timing.
fn token_ok(expected: &str, supplied: Option<&String>) -> bool {
let Some(supplied) = supplied.and_then(|v| v.strip_prefix("Bearer ")) else {
return false;
};
let a = expected.as_bytes();
let b = supplied.as_bytes();
// Length is compared without branching on it beyond the final AND.
let mut diff = (a.len() ^ b.len()) as u8;
for i in 0..a.len().max(b.len()) {
diff |= a.get(i).copied().unwrap_or(0) ^ b.get(i).copied().unwrap_or(0);
}
diff == 0
}
fn handle(mut stream: TcpStream, config: &Config) -> Result<()> {
stream.set_read_timeout(Some(READ_TIMEOUT))?;
stream.set_write_timeout(Some(READ_TIMEOUT))?;
let Some(request) = read_request(&stream)? else {
return Ok(());
};
// Echo the origin only when it is on the allowlist. Never `*`: this API describes the
// machine it runs on, so any-origin access means any page can fingerprint the node.
let allow_origin = request
.origin
.as_ref()
.filter(|o| config.allowed_origins.iter().any(|a| a == *o))
.cloned();
if request.method == "OPTIONS" {
return write_response(&mut stream, 204, "", allow_origin.as_deref(), true);
}
if let Some(expected) = &config.token {
if !token_ok(expected, request.authorization.as_ref()) {
return write_json(
&mut stream,
401,
r#"{"error":"unauthorized"}"#,
allow_origin.as_deref(),
);
}
}
if request.method != "GET" {
return write_json(
&mut stream,
405,
r#"{"error":"this API is read-only"}"#,
allow_origin.as_deref(),
);
}
let (status, body) = route(&request.path, &config.root);
write_json(&mut stream, status, &body, allow_origin.as_deref())
}
fn route(path: &str, root: &Path) -> (u16, String) {
let rendered = match path {
"/v1/health" => Ok(r#"{"ok":true,"service":"lumbridge-compute","api":"v1"}"#.to_string()),
"/v1/status" => collect_status(root).and_then(|r| Ok(serde_json::to_string(&r)?)),
"/v1/models" => collect_models(root).and_then(|r| Ok(serde_json::to_string(&r)?)),
"/v1/scenes" => collect_scenes(root).and_then(|r| Ok(serde_json::to_string(&r)?)),
"/v1/evals" => collect_evals(root).and_then(|r| Ok(serde_json::to_string(&r)?)),
other => match other.strip_prefix("/v1/scenes/") {
// Exactly one segment: /v1/scenes/a/b is not a route.
Some(name) if !name.is_empty() && !name.contains('/') => {
collect_scene(root, name).and_then(|r| Ok(serde_json::to_string(&r)?))
}
_ => return (404, r#"{"error":"no such route"}"#.to_string()),
},
};
match rendered {
Ok(body) => (200, body),
// A collector fails when the thing does not exist (an unknown Scene) or when the node's
// own config is unreadable. The message is the operator's, and this is a loopback API,
// so passing it through is more useful than flattening it to "error".
Err(error) => (
404,
serde_json::json!({ "error": format!("{error:#}") }).to_string(),
),
}
}
fn write_json(stream: &mut TcpStream, status: u16, body: &str, origin: Option<&str>) -> Result<()> {
write_response(stream, status, body, origin, false)
}
fn write_response(
stream: &mut TcpStream,
status: u16,
body: &str,
origin: Option<&str>,
preflight: bool,
) -> Result<()> {
let reason = match status {
200 => "OK",
204 => "No Content",
401 => "Unauthorized",
404 => "Not Found",
405 => "Method Not Allowed",
_ => "Error",
};
let mut head = format!("HTTP/1.1 {status} {reason}\r\n");
head.push_str("Content-Type: application/json\r\n");
head.push_str(&format!("Content-Length: {}\r\n", body.len()));
// This is live node state; a cached answer is a wrong answer.
head.push_str("Cache-Control: no-store\r\n");
head.push_str("Connection: close\r\n");
if let Some(origin) = origin {
head.push_str(&format!("Access-Control-Allow-Origin: {origin}\r\n"));
// Tell caches the body varies by origin, so an allowed origin's response can never be
// replayed to a disallowed one.
head.push_str("Vary: Origin\r\n");
if preflight {
head.push_str("Access-Control-Allow-Methods: GET, OPTIONS\r\n");
head.push_str("Access-Control-Allow-Headers: Authorization\r\n");
head.push_str("Access-Control-Max-Age: 600\r\n");
}
}
head.push_str("\r\n");
stream.write_all(head.as_bytes())?;
stream.write_all(body.as_bytes())?;
stream.flush()?;
Ok(())
}
/// Drain and discard — kept for symmetry with future routes that accept a body.
#[allow(dead_code)]
fn discard_body(reader: &mut impl Read) {
let mut sink = Vec::new();
let _ = reader.read_to_end(&mut sink);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn token_comparison_rejects_wrong_and_missing_and_prefixless() {
let t = "s3cret-token";
assert!(token_ok(t, Some(&format!("Bearer {t}"))));
assert!(!token_ok(t, None));
assert!(!token_ok(t, Some(&t.to_string()))); // no "Bearer " prefix
assert!(!token_ok(t, Some(&"Bearer wrong".to_string())));
// A correct prefix must not pass — this is the bug constant-time comparison exists for.
assert!(!token_ok(t, Some(&"Bearer s3cret".to_string())));
assert!(!token_ok(t, Some(&"Bearer s3cret-token-plus".to_string())));
}
#[test]
fn unknown_routes_404_and_scene_paths_take_exactly_one_segment() {
let root = Path::new("/nonexistent-root-for-routing-test");
assert_eq!(route("/v1/nope", root).0, 404);
assert_eq!(route("/v1/scenes/a/b", root).0, 404);
assert_eq!(route("/v1/scenes/", root).0, 404);
// Health needs no filesystem, so it answers even on a bogus root.
assert_eq!(route("/v1/health", root).0, 200);
}
#[test]
fn health_body_is_valid_json() {
let (status, body) = route("/v1/health", Path::new("/tmp"));
assert_eq!(status, 200);
let parsed: serde_json::Value = serde_json::from_str(&body).unwrap();
assert_eq!(parsed["ok"], true);
}
}
+689
View File
@@ -0,0 +1,689 @@
//! Transactional Scene lifecycle and persisted desired-state recovery.
use anyhow::{bail, Context, Result};
use std::collections::{BTreeMap, HashSet};
use std::fs::{self, File, OpenOptions};
use std::os::fd::AsRawFd;
use std::path::Path;
use std::time::Duration;
use crate::config::{find_scene, Model, Registry, Scene};
use crate::governor;
use crate::mem;
use crate::proc::{self, Proc, State};
/// What activating a Scene would do, decided before anything is mutated.
///
/// `activate` renders this for the CLI and the MCP server returns it verbatim,
/// so a plan an operator reads and a plan an agent reads can never drift apart.
#[derive(Debug, Clone)]
pub struct Plan {
pub scene: String,
pub budget_gb: f64,
/// Registered models serving outside the target Scene; stopped first.
pub stop: Vec<String>,
/// Scene models not yet serving, in the order they would be admitted.
pub start: Vec<String>,
}
pub struct TransitionLock {
file: File,
}
impl TransitionLock {
pub fn acquire(root: &Path) -> Result<Self> {
let dir = root.join(".kuda");
fs::create_dir_all(&dir)?;
let path = dir.join("transition.lock");
let file = OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(&path)
.with_context(|| format!("opening transition lock {}", path.display()))?;
let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
if result != 0 {
return Err(std::io::Error::last_os_error())
.context("another Scene transition is already running; wait for it to finish");
}
Ok(Self { file })
}
}
impl Drop for TransitionLock {
fn drop(&mut self) {
unsafe {
libc::flock(self.file.as_raw_fd(), libc::LOCK_UN);
}
}
}
pub fn adopt(root: &Path, name: &str) -> Result<()> {
let _lock = TransitionLock::acquire(root)?;
let registry = Registry::load(root)?;
let scene = find_scene(root, name)?;
validate_scene(&registry, &scene)?;
let target: HashSet<&str> = scene.models.iter().map(String::as_str).collect();
let running = governor::running_ids(&registry);
let extras: Vec<String> = running
.iter()
.filter(|id| !target.contains(id.as_str()))
.cloned()
.collect();
if !extras.is_empty() {
bail!(
"cannot adopt '{name}': registered model(s) outside the Scene are serving: {}",
extras.join(", ")
);
}
let previous = State::load_checked(root)?;
let mut adopted = BTreeMap::new();
for id in &scene.models {
let model = &registry.models[id];
if !governor::is_running(model) {
bail!("cannot adopt '{name}': exact model '{id}' is not healthy");
}
let legacy = previous.procs.get(id).with_context(|| {
format!(
"cannot adopt '{name}': no legacy process record for '{id}'; start it through Compute"
)
})?;
let captured = Proc::capture(legacy.pid, legacy.seq, model.serve.port)
.with_context(|| format!("adopting '{id}' pid {}", legacy.pid))?;
adopted.insert(id.clone(), captured);
}
let mut state = previous;
state.procs = adopted;
state.desired_scene = Some(name.to_string());
state.active_scene = Some(name.to_string());
state.last_known_good_scene = Some(name.to_string());
state.transition_scene = None;
state.last_error = None;
state.save(root)?;
println!("adopted exact running Scene '{name}' and captured process identities");
Ok(())
}
pub fn activate(root: &Path, name: &str, dry_run: bool) -> Result<()> {
let _lock = TransitionLock::acquire(root)?;
let prior = State::load_checked(root)?;
let prior_active = prior.active_scene.clone();
if dry_run {
return activate_once(root, name, true);
}
match activate_once(root, name, false) {
Ok(()) => {
let mut state = State::load_checked(root)?;
if prior_active.as_deref() != Some(name) {
if let Some(previous) = prior_active {
state.last_known_good_scene = Some(previous);
}
}
state.desired_scene = Some(name.to_string());
state.active_scene = Some(name.to_string());
if state.last_known_good_scene.is_none() {
state.last_known_good_scene = Some(name.to_string());
}
state.transition_scene = None;
state.last_error = None;
state.save(root)?;
println!("Scene '{name}' is active and persisted as desired");
Ok(())
}
Err(target_error) => {
let target_message = format!("activating '{name}' failed: {target_error:#}");
eprintln!("{target_message}");
let transition_started =
State::load_checked(root)?.transition_scene.as_deref() == Some(name);
let cleanup_error = if transition_started {
stop_all_owned(root).err()
} else {
None
};
let rollback = if transition_started && cleanup_error.is_none() {
prior_active
.as_deref()
.filter(|previous| *previous != name)
.map(|previous| (previous.to_string(), activate_once(root, previous, false)))
} else {
None
};
let mut state = State::load_checked(root)?;
state.transition_scene = None;
state.last_error = Some(target_message.clone());
match rollback {
Some((previous, Ok(()))) => {
state.desired_scene = Some(previous.clone());
state.active_scene = Some(previous.clone());
if state.last_known_good_scene.is_none() {
state.last_known_good_scene = Some(previous.clone());
}
state.save(root)?;
eprintln!("rolled back to Scene '{previous}'");
bail!("{target_message}; rolled back to '{previous}'")
}
Some((previous, Err(rollback_error))) => {
state.active_scene = None;
state.save(root)?;
bail!(
"{target_message}; rollback to '{previous}' also failed: {rollback_error:#}"
)
}
None if !transition_started => {
state.save(root)?;
bail!("{target_message}")
}
None if cleanup_error.is_none() => {
state.active_scene = None;
state.save(root)?;
bail!("{target_message}")
}
None => {
state.active_scene = None;
state.save(root)?;
let cleanup_error = cleanup_error.expect("guarded by match condition");
bail!("{target_message}; cleanup also failed: {cleanup_error:#}")
}
}
}
}
}
fn stop_all_owned(root: &Path) -> Result<()> {
let mut state = State::load_checked(root)?;
let owned: Vec<(String, Proc)> = state
.procs
.iter()
.filter(|(_, process)| process.owned_alive())
.map(|(id, process)| (id.clone(), process.clone()))
.collect();
for (id, process) in owned {
proc::stop_owned(&process).with_context(|| format!("cleaning up '{id}'"))?;
state.procs.remove(&id);
state.save(root)?;
}
Ok(())
}
pub fn resume(root: &Path) -> Result<()> {
let state = State::load_checked(root)?;
let desired = state
.desired_scene
.clone()
.or_else(|| state.last_known_good_scene.clone())
.context("no desired Scene is persisted; activate or adopt one first")?;
let fallback = state.last_known_good_scene.clone();
match activate(root, &desired, false) {
Ok(()) => Ok(()),
Err(desired_error) => {
let Some(fallback) = fallback.filter(|fallback| fallback != &desired) else {
return Err(desired_error).context("resuming desired Scene");
};
eprintln!(
"desired Scene '{desired}' did not resume; trying last-known-good '{fallback}'"
);
activate(root, &fallback, false).with_context(|| {
format!(
"desired Scene '{desired}' failed ({desired_error:#}) and fallback '{fallback}' failed"
)
})
}
}
}
pub fn deactivate(root: &Path) -> Result<()> {
let _lock = TransitionLock::acquire(root)?;
let mut state = State::load_checked(root)?;
let ids: Vec<String> = state.procs.keys().cloned().collect();
for id in ids {
let process = state.procs[&id].clone();
if process.owned_alive() {
println!(" stopping {id} (pid {})", process.pid);
proc::stop_owned(&process)?;
}
state.procs.remove(&id);
state.save(root)?;
}
state.desired_scene = None;
state.active_scene = None;
state.transition_scene = None;
state.last_error = None;
state.save(root)?;
println!("all Compute-managed models stopped; no Scene is desired");
Ok(())
}
/// Diff the target Scene against what is actually serving. Pure: it decides
/// *what* would change, never *whether* the change is allowed.
fn plan_transition(registry: &Registry, scene: &Scene) -> Plan {
let target: HashSet<&str> = scene.models.iter().map(String::as_str).collect();
let running = governor::running_ids(registry);
let stop: Vec<String> = running
.iter()
.filter(|id| !target.contains(id.as_str()))
.cloned()
.collect();
let mut start: Vec<String> = scene
.models
.iter()
.filter(|id| !governor::is_running(&registry.models[*id]))
.cloned()
.collect();
let listed = scene
.activation
.as_ref()
.and_then(|activation| activation.order.as_deref())
.map(|order| order == "listed")
.unwrap_or(false);
if !listed {
start.sort_by(|a, b| {
registry.models[a]
.footprint_gb
.partial_cmp(&registry.models[b].footprint_gb)
.unwrap()
});
}
Plan {
scene: scene.metadata.name.clone(),
budget_gb: scene.budget_gb.unwrap_or(governor::DEFAULT_BUDGET_GB),
stop,
start,
}
}
/// Everything that must hold before the first process is signalled. Runs to
/// completion with nothing mutated, so a rejection here leaves the currently
/// active Scene exactly as it was.
fn preflight(registry: &Registry, scene: &Scene, plan: &Plan, state: &State) -> Result<()> {
for id in &scene.models {
if governor::is_running(&registry.models[id]) {
let process = state.procs.get(id).with_context(|| {
format!(
"'{id}' is already serving but is not identity-owned by Compute; adopt the active Scene first"
)
})?;
if !process.owned_alive() {
bail!(
"ownership record for serving model '{id}' is stale; adopt the active Scene first"
);
}
}
}
for id in &plan.stop {
let process = state.procs.get(id).with_context(|| {
format!(
"'{id}' is serving but is not identity-owned by Compute; adopt the active Scene before switching"
)
})?;
if !process.owned_alive() {
bail!("ownership record for serving model '{id}' is stale; refusing to signal its pid");
}
}
for id in &plan.start {
let model = &registry.models[id];
if let Some(port) = model.serve.port {
// A model in `plan.stop` may currently hold this port. Every stop runs before
// any start, so that is a handoff, not a conflict. Without this exemption every
// same-port swap is rejected — including brain -> brain-laguna/brain-gemma,
// which share :8001 by design because the alias downstream agents call must
// survive a weight swap.
let freed_by_stop = plan.stop.iter().any(|stopping| {
registry
.models
.get(stopping)
.and_then(|stopped| stopped.serve.port)
== Some(port)
});
if !freed_by_stop && governor::port_open(port) && !governor::is_running(model) {
bail!("port {port} is occupied by a different model; refusing to start '{id}'");
}
}
}
Ok(())
}
/// `scene activate --dry-run` as data rather than as printed lines: the same
/// validation, the same plan, nothing written. Callers that need the plan
/// programmatically use this instead of scraping stdout.
pub fn plan(root: &Path, name: &str) -> Result<Plan> {
let _lock = TransitionLock::acquire(root)?;
let registry = Registry::load(root)?;
let scene = find_scene(root, name)?;
validate_scene(&registry, &scene)?;
let plan = plan_transition(&registry, &scene);
let mut state = State::load_checked(root)?;
state.procs.retain(|_, process| process.owned_alive());
preflight(&registry, &scene, &plan, &state)?;
Ok(plan)
}
fn activate_once(root: &Path, name: &str, dry_run: bool) -> Result<()> {
let registry = Registry::load(root)?;
let scene = find_scene(root, name)?;
validate_scene(&registry, &scene)?;
let plan = plan_transition(&registry, &scene);
println!("activate '{name}' (budget {:.0} GB)", plan.budget_gb);
println!(" stop : {}", format_ids(&plan.stop));
println!(" start: {}", format_ids(&plan.start));
let mut state = State::load_checked(root)?;
state.procs.retain(|_, process| process.owned_alive());
preflight(&registry, &scene, &plan, &state)?;
if dry_run {
println!(" dry run: validation passed; nothing changed");
return Ok(());
}
state.transition_scene = Some(name.to_string());
state.save(root)?;
for id in &plan.stop {
let process = state.procs[id].clone();
println!(" stopping {id} (pid {})", process.pid);
proc::stop_owned(&process)?;
state.procs.remove(id);
state.save(root)?;
}
let mut committed: f64 = scene
.models
.iter()
.filter(|id| governor::is_running(&registry.models[*id]))
.map(|id| registry.models[id].footprint_gb)
.sum();
let wait_healthy = scene
.activation
.as_ref()
.and_then(|activation| activation.wait_healthy)
.unwrap_or(true);
for id in &plan.start {
let model: &Model = &registry.models[id];
// Re-read the pool before every start rather than once per activation: each model
// that comes up consumes real memory, and its true appetite is only knowable after
// it has allocated. A footprint that was optimistic shows up here, on the next
// model, instead of taking the box down.
let available_gb = mem::read().ok().map(|m| m.available_gb);
if !governor::can_admit(model.footprint_gb, committed, plan.budget_gb, available_gb) {
match available_gb {
Some(available) => bail!(
"'{id}' ({:.0} GB) refused: {:.0} GB committed against a {:.0} GB budget, \
{:.1} GB actually available, {:.0} GB margin required",
model.footprint_gb,
committed,
plan.budget_gb,
available,
governor::SAFETY_MARGIN_GB
),
None => bail!(
"'{id}' would exceed the {:.0} GB Scene budget with the safety margin",
plan.budget_gb
),
}
}
let pid = proc::spawn(root, id, model)?;
state.seq += 1;
let process = Proc::capture(pid, state.seq, model.serve.port)
.with_context(|| format!("capturing ownership for newly started '{id}'"))?;
state.procs.insert(id.clone(), process.clone());
state.save(root)?;
committed += model.footprint_gb;
println!(" started {id} (pid {pid}); waiting for exact health");
if wait_healthy && !governor::wait_healthy(model, health_timeout()) {
let _ = proc::stop_owned(&process);
state.procs.remove(id);
state.save(root)?;
bail!("'{id}' did not report its exact health marker before timeout");
}
}
let unhealthy: Vec<String> = scene
.models
.iter()
.filter(|id| !governor::is_running(&registry.models[*id]))
.cloned()
.collect();
if !unhealthy.is_empty() {
bail!(
"Scene '{name}' is incomplete; exact health failed for: {}",
unhealthy.join(", ")
);
}
Ok(())
}
fn validate_scene(registry: &Registry, scene: &Scene) -> Result<()> {
let missing: Vec<&String> = scene
.models
.iter()
.filter(|id| !registry.models.contains_key(*id))
.collect();
if !missing.is_empty() {
bail!(
"Scene '{}' references unknown model id(s): {}",
scene.metadata.name,
missing
.iter()
.map(|id| id.as_str())
.collect::<Vec<_>>()
.join(", ")
);
}
let budget = scene.budget_gb.unwrap_or(governor::DEFAULT_BUDGET_GB);
let footprint: f64 = scene
.models
.iter()
.map(|id| registry.models[id].footprint_gb)
.sum();
if footprint + governor::SAFETY_MARGIN_GB > budget {
bail!(
"Scene '{}' needs {:.1} GB including safety margin, above its {:.1} GB budget",
scene.metadata.name,
footprint + governor::SAFETY_MARGIN_GB,
budget
);
}
Ok(())
}
fn health_timeout() -> Duration {
Duration::from_secs(if cfg!(test) { 5 } else { 900 })
}
fn format_ids(ids: &[String]) -> String {
if ids.is_empty() {
"(none)".to_string()
} else {
ids.join(", ")
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::TcpListener;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
fn fixture_root() -> (PathBuf, u16) {
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let root = std::env::temp_dir().join(format!("lumbridge-compute-{unique}"));
fs::create_dir_all(root.join("registry")).unwrap();
fs::create_dir_all(root.join("scenes")).unwrap();
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
drop(listener);
let bad_listener = TcpListener::bind("127.0.0.1:0").unwrap();
let bad_port = bad_listener.local_addr().unwrap().port();
drop(bad_listener);
let registry = format!(
"apiVersion: lumbridge/v1\nmodels:\n fake:\n name: Fake\n footprint_gb: 1\n health: http://localhost:{port}/\n serve:\n kind: exec\n port: {port}\n command: [/usr/bin/python3, -m, http.server, '{port}', --bind, 127.0.0.1]\n bad:\n name: Bad\n footprint_gb: 2\n health: http://localhost:{bad_port}/\n serve:\n kind: exec\n port: {bad_port}\n command: [/usr/bin/false]\n"
);
fs::write(root.join("registry/models.yaml"), registry).unwrap();
fs::write(
root.join("scenes/test.scene.yaml"),
"apiVersion: lumbridge/v1\nmetadata:\n name: test\n version: 1\nmodels: [fake]\nbudget_gb: 100\n",
)
.unwrap();
(root, port)
}
#[test]
fn activation_persists_and_deactivation_stops_owned_process() {
let (root, _port) = fixture_root();
activate(&root, "test", false).unwrap();
let state = State::load_checked(&root).unwrap();
assert_eq!(state.desired_scene.as_deref(), Some("test"));
assert_eq!(state.active_scene.as_deref(), Some("test"));
assert!(state.procs["fake"].owned_alive());
deactivate(&root).unwrap();
assert!(State::load_checked(&root).unwrap().procs.is_empty());
fs::remove_dir_all(root).unwrap();
}
/// Two models sharing one port, distinguishable by `health_contains` — the
/// brain/brain-laguna shape. Each serves its own directory so the health probe
/// can tell which one is actually up.
fn same_port_root() -> PathBuf {
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let root = std::env::temp_dir().join(format!("lumbridge-compute-swap-{unique}"));
fs::create_dir_all(root.join("registry")).unwrap();
fs::create_dir_all(root.join("scenes")).unwrap();
let dir_a = root.join("srv-alpha");
let dir_b = root.join("srv-beta");
fs::create_dir_all(&dir_a).unwrap();
fs::create_dir_all(&dir_b).unwrap();
fs::write(dir_a.join("alpha-marker.txt"), "a").unwrap();
fs::write(dir_b.join("beta-marker.txt"), "b").unwrap();
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
drop(listener);
let (a, b) = (dir_a.display(), dir_b.display());
let registry = format!(
"apiVersion: lumbridge/v1\nmodels:\n alpha:\n name: Alpha\n footprint_gb: 1\n health: http://localhost:{port}/\n health_contains: alpha-marker\n serve:\n kind: exec\n port: {port}\n command: [/usr/bin/python3, -m, http.server, '{port}', --bind, 127.0.0.1, --directory, '{a}']\n beta:\n name: Beta\n footprint_gb: 1\n health: http://localhost:{port}/\n health_contains: beta-marker\n serve:\n kind: exec\n port: {port}\n command: [/usr/bin/python3, -m, http.server, '{port}', --bind, 127.0.0.1, --directory, '{b}']\n"
);
fs::write(root.join("registry/models.yaml"), registry).unwrap();
fs::write(
root.join("scenes/a.scene.yaml"),
"apiVersion: lumbridge/v1\nmetadata:\n name: a\n version: 1\nmodels: [alpha]\nbudget_gb: 100\n",
)
.unwrap();
fs::write(
root.join("scenes/b.scene.yaml"),
"apiVersion: lumbridge/v1\nmetadata:\n name: b\n version: 1\nmodels: [beta]\nbudget_gb: 100\n",
)
.unwrap();
root
}
#[test]
fn same_port_swap_is_allowed_when_the_occupant_is_being_stopped() {
let root = same_port_root();
activate(&root, "a", false).unwrap();
assert!(State::load_checked(&root).unwrap().procs.contains_key("alpha"));
// Regression: this used to fail with "port N is occupied by a different model;
// refusing to start 'beta'". The pre-flight port check ran over `plan.start`
// without exempting ports released by `plan.stop`, so every same-port swap was
// rejected even though stops precede starts.
activate(&root, "b", false).unwrap();
let state = State::load_checked(&root).unwrap();
assert_eq!(state.active_scene.as_deref(), Some("b"));
assert!(state.procs.contains_key("beta"));
assert!(!state.procs.contains_key("alpha"));
deactivate(&root).unwrap();
fs::remove_dir_all(root).unwrap();
}
#[test]
fn already_serving_scene_requires_explicit_adoption() {
let (root, port) = fixture_root();
let registry = Registry::load(&root).unwrap();
let model = &registry.models["fake"];
let pid = proc::spawn(&root, "fake", model).unwrap();
assert!(governor::wait_healthy(model, Duration::from_secs(5)));
let error = activate(&root, "test", true).unwrap_err().to_string();
assert!(error.contains("adopt the active Scene"));
let owned = Proc::capture(pid, 1, Some(port)).unwrap();
proc::stop_owned(&owned).unwrap();
fs::remove_dir_all(root).unwrap();
}
#[test]
fn preflight_rejection_preserves_healthy_active_scene() {
let (root, _port) = fixture_root();
activate(&root, "test", false).unwrap();
let before = State::load_checked(&root).unwrap();
let pid = before.procs["fake"].pid;
fs::write(
root.join("scenes/invalid.scene.yaml"),
"apiVersion: lumbridge/v1\nmetadata:\n name: invalid\n version: 1\nmodels: [missing]\nbudget_gb: 100\n",
)
.unwrap();
assert!(activate(&root, "invalid", false).is_err());
let after = State::load_checked(&root).unwrap();
assert_eq!(after.active_scene.as_deref(), Some("test"));
assert_eq!(after.desired_scene.as_deref(), Some("test"));
assert_eq!(after.procs["fake"].pid, pid);
assert!(after.procs["fake"].owned_alive());
deactivate(&root).unwrap();
fs::remove_dir_all(root).unwrap();
}
#[test]
fn failed_activation_cleans_up_partial_scene() {
let (root, port) = fixture_root();
fs::write(
root.join("scenes/failing.scene.yaml"),
"apiVersion: lumbridge/v1\nmetadata:\n name: failing\n version: 1\nmodels: [fake, bad]\nbudget_gb: 100\n",
)
.unwrap();
assert!(activate(&root, "failing", false).is_err());
let state = State::load_checked(&root).unwrap();
assert!(state.procs.is_empty());
assert!(state.active_scene.is_none());
assert!(!governor::port_open(port));
fs::remove_dir_all(root).unwrap();
}
#[test]
fn checked_state_load_rejects_malformed_yaml() {
let (root, _port) = fixture_root();
fs::create_dir_all(root.join(".kuda")).unwrap();
fs::write(root.join(".kuda/state.yaml"), "desired_scene: [").unwrap();
assert!(State::load_checked(&root).is_err());
fs::remove_dir_all(root).unwrap();
}
#[test]
fn state_write_replaces_complete_yaml_atomically() {
let (root, _port) = fixture_root();
let state = State {
desired_scene: Some("test".to_string()),
..State::default()
};
state.save(&root).unwrap();
let loaded = State::load_checked(&root).unwrap();
assert_eq!(loaded.desired_scene.as_deref(), Some("test"));
fs::remove_dir_all(root).unwrap();
}
}
+431
View File
@@ -0,0 +1,431 @@
//! Lumbridge Compute — safe AI workload orchestration for accelerator nodes.
mod agent;
mod config;
mod eval;
mod gateway;
mod governor;
mod http;
mod lifecycle;
mod mcp;
mod mem;
mod proc;
use anyhow::Result;
use clap::{Parser, Subcommand};
use std::path::{Path, PathBuf};
use std::time::Duration;
use config::{find_scene, load_scenes, Registry, Scene};
use proc::State;
#[derive(Parser)]
#[command(
name = "lumbridge-compute",
version,
about = "Lumbridge Compute — safe AI workload orchestration for accelerator nodes."
)]
struct Cli {
/// Root dir containing registry/ and scenes/ (default: $LUMBRIDGE_COMPUTE_ROOT, $KUDA_ROOT, or current dir)
#[arg(long, global = true)]
root: Option<PathBuf>,
#[command(subcommand)]
cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
/// Governor status: memory, budget, running set, headroom
Status,
/// Manage the model registry
Model {
#[command(subcommand)]
cmd: ModelCmd,
},
/// Manage scenes (bundles of models)
Scene {
#[command(subcommand)]
cmd: SceneCmd,
},
/// Run reproducible capability and serving-performance evaluations
Eval {
#[command(subcommand)]
cmd: EvalCmd,
},
/// Run the memory watchdog in the foreground (kills the newest model before OOM-wedge)
Watchdog {
/// Kill the newest model if MemAvailable dips below this many GB
#[arg(long, default_value_t = governor::WATCHDOG_FLOOR_GB)]
floor: f64,
},
/// Run the stable streaming gateway in the foreground
Gateway {
#[arg(long, default_value = "127.0.0.1:8011")]
listen: String,
#[arg(long, default_value = "127.0.0.1:8001")]
upstream: String,
},
/// Serve the Governor as a read-only JSON API over HTTP
///
/// A second transport over the same operations the MCP server exposes, for a web UI. It
/// never mutates: no activation, no stop, no eval run.
Api {
/// Loopback by default. This API has no authentication unless --token is set, and it
/// reveals which models this node runs, so widening the bind is an explicit act.
#[arg(long, default_value = "127.0.0.1:8012")]
listen: String,
/// Browser origin permitted to read this API. Repeatable. Empty means no browser may
/// read it; `*` is deliberately not supported, because any page you visit would then
/// be able to fingerprint this machine.
#[arg(long = "allow-origin")]
allow_origin: Vec<String>,
/// Require `Authorization: Bearer <token>` on every request.
#[arg(long)]
token: Option<String>,
},
/// Serve the Governor to AI agents as MCP tools over stdio
Mcp {
/// Let an agent perform a real Scene transition, not just plan one.
/// Off by default: the agent writes its own tool arguments, so the only
/// meaningful gate on a production switch is one the operator sets here.
#[arg(long)]
allow_activate: bool,
},
/// Run the resident supervisor: resume desired Scene, gateway, and memory floor
Agent {
#[arg(long, default_value = "127.0.0.1:8011")]
listen: String,
#[arg(long, default_value = "127.0.0.1:8001")]
upstream: String,
#[arg(long, default_value_t = governor::WATCHDOG_FLOOR_GB)]
floor: f64,
},
}
#[derive(Subcommand)]
enum ModelCmd {
/// List registered models and their live state
Ls,
}
#[derive(Subcommand)]
enum SceneCmd {
/// List scenes with total footprint
Ls,
/// Show a scene: models, footprints, and the Governor's admission verdict
Show { name: String },
/// Activate a scene: stop what's not in it, admit + start what is
Activate {
name: String,
/// Print the plan without changing anything
#[arg(long)]
dry_run: bool,
},
/// Adopt an already-running exact Scene and bind legacy PIDs to process identities
Adopt { name: String },
/// Resume the persisted desired Scene, falling back to the previous known-good Scene
Resume,
/// Stop all Lumbridge Compute-managed models
Deactivate,
}
#[derive(Subcommand)]
enum EvalCmd {
/// List available evaluation suites
Ls,
/// Run a suite against an OpenAI-compatible endpoint
Run {
suite: String,
#[arg(long, default_value = "http://127.0.0.1:8001/v1")]
base_url: String,
#[arg(long, default_value = "brain")]
model: String,
/// Override the suite's repetitions per case
#[arg(long)]
repeat: Option<u32>,
},
}
fn root_dir(cli: &Cli) -> PathBuf {
cli.root
.clone()
.or_else(|| {
std::env::var("LUMBRIDGE_COMPUTE_ROOT")
.ok()
.map(PathBuf::from)
})
.or_else(|| std::env::var("KUDA_ROOT").ok().map(PathBuf::from))
.unwrap_or_else(|| PathBuf::from("."))
}
fn main() -> Result<()> {
let cli = Cli::parse();
let root = root_dir(&cli);
match &cli.cmd {
Cmd::Status => cmd_status(&root),
Cmd::Model { cmd } => match cmd {
ModelCmd::Ls => cmd_model_ls(&root),
},
Cmd::Scene { cmd } => match cmd {
SceneCmd::Ls => cmd_scene_ls(&root),
SceneCmd::Show { name } => cmd_scene_show(&root, name),
SceneCmd::Activate { name, dry_run } => lifecycle::activate(&root, name, *dry_run),
SceneCmd::Adopt { name } => lifecycle::adopt(&root, name),
SceneCmd::Resume => lifecycle::resume(&root),
SceneCmd::Deactivate => lifecycle::deactivate(&root),
},
Cmd::Eval { cmd } => match cmd {
EvalCmd::Ls => eval::list(&root),
EvalCmd::Run {
suite,
base_url,
model,
repeat,
} => eval::run(&root, suite, base_url, model, *repeat),
},
Cmd::Watchdog { floor } => cmd_watchdog(&root, *floor),
Cmd::Gateway { listen, upstream } => gateway::run(listen, upstream),
Cmd::Api {
listen,
allow_origin,
token,
} => http::run(
listen,
http::Config {
root: root.clone(),
allowed_origins: allow_origin.clone(),
token: token.clone(),
},
),
Cmd::Mcp { allow_activate } => mcp::run(&root, *allow_activate),
Cmd::Agent {
listen,
upstream,
floor,
} => agent::run(&root, listen, upstream, *floor),
}
}
fn cmd_status(root: &Path) -> Result<()> {
let reg = Registry::load(root)?;
let m = mem::read()?;
let committed = governor::committed_gb(&reg);
let running = governor::running_ids(&reg);
let budget = governor::DEFAULT_BUDGET_GB;
// The binding constraint, not the generous one: admission enforces the declared
// budget AND observed memory, so reporting only the declared headroom would promise
// room the next `scene activate` is going to refuse.
let headroom = governor::headroom_gb(committed, budget, Some(m.available_gb)).max(0.0);
let declared_headroom = (budget - committed - governor::SAFETY_MARGIN_GB).max(0.0);
let managed = State::load_checked(root)?;
println!("Lumbridge Compute · governor");
println!(
" memory {:.1} GB total · {:.1} GB available",
m.total_gb, m.available_gb
);
println!(
" budget {:.0} GB (safety margin {:.0} · watchdog floor {:.0})",
budget,
governor::SAFETY_MARGIN_GB,
governor::WATCHDOG_FLOOR_GB
);
println!(
" committed {committed:.1} GB across {} model(s)",
running.len()
);
if headroom < declared_headroom {
println!(
" headroom {headroom:.1} GB admittable (budget allows {declared_headroom:.1}; \
MemAvailable is the tighter limit)"
);
} else {
println!(" headroom {headroom:.1} GB admittable");
}
println!();
println!(
" scene active={} desired={} fallback={}",
managed.active_scene.as_deref().unwrap_or("-"),
managed.desired_scene.as_deref().unwrap_or("-"),
managed.last_known_good_scene.as_deref().unwrap_or("-")
);
if let Some(error) = &managed.last_error {
println!(" last error {error}");
}
if running.is_empty() {
println!(" (no registered models currently serving)");
} else {
println!(" running:");
for id in &running {
let mdl = &reg.models[id];
let tag = if managed
.procs
.get(id)
.is_some_and(|process| process.owned_alive())
{
"compute"
} else {
"ext "
};
println!(
" ● [{tag}] {:<12} {:>5.0} GB :{:<5} {}",
id,
mdl.footprint_gb,
port_str(mdl.serve.port),
mdl.name
);
}
}
Ok(())
}
fn cmd_model_ls(root: &Path) -> Result<()> {
let reg = Registry::load(root)?;
println!(
"{:<12} {:>6} {:<6} {:<6} MODEL",
"ID", "GB", "STATE", "PORT"
);
for (id, m) in &reg.models {
println!(
"{:<12} {:>6.0} {:<6} {:<6} {}",
id,
m.footprint_gb,
if governor::is_running(m) {
"up"
} else {
"down"
},
port_str(m.serve.port),
m.name
);
}
Ok(())
}
fn cmd_scene_ls(root: &Path) -> Result<()> {
let reg = Registry::load(root)?;
let scenes = load_scenes(root)?;
if scenes.is_empty() {
println!("no scenes in {}/scenes", root.display());
return Ok(());
}
println!("{:<12} {:>6} {:<30} MODELS", "SCENE", "GB", "DESCRIPTION");
for s in &scenes {
let total = scene_footprint(s, &reg);
println!(
"{:<12} {:>6.0} {:<30} {}",
s.metadata.name,
total,
truncate(&s.metadata.description, 30),
s.models.join(", ")
);
}
Ok(())
}
fn cmd_scene_show(root: &Path, name: &str) -> Result<()> {
let reg = Registry::load(root)?;
let scene = find_scene(root, name)?;
let budget = scene.budget_gb.unwrap_or(governor::DEFAULT_BUDGET_GB);
println!(
"scene {} (v{})",
scene.metadata.name, scene.metadata.version
);
if !scene.metadata.description.is_empty() {
println!(" {}", scene.metadata.description);
}
println!(" budget {budget:.0} GB\n models:");
let mut total = 0.0;
let mut missing = Vec::new();
for id in &scene.models {
match reg.models.get(id) {
Some(m) => {
total += m.footprint_gb;
let state = if governor::is_running(m) {
"up"
} else {
"down"
};
println!(
" {:<12} {:>5.0} GB [{:<4}] {}",
id, m.footprint_gb, state, m.name
);
}
None => {
missing.push(id.clone());
println!(" {id:<12} ? [MISSING from registry]");
}
}
}
let needed = total + governor::SAFETY_MARGIN_GB;
println!(
"\n total footprint {total:.1} GB (+{:.0} safety = {needed:.1} GB)",
governor::SAFETY_MARGIN_GB
);
let verdict = if !missing.is_empty() {
format!("{} model(s) missing from registry", missing.len())
} else if needed <= budget {
format!("✓ fits — {:.1} GB to spare", budget - needed)
} else {
format!("✗ exceeds budget by {:.1} GB", needed - budget)
};
println!(" admission {verdict}");
Ok(())
}
fn cmd_watchdog(root: &Path, floor: f64) -> Result<()> {
println!(
"Lumbridge Compute watchdog — killing the newest model if MemAvailable < {floor:.1} GB. Ctrl-C to stop."
);
loop {
let m = mem::read()?;
if m.available_gb < floor {
let mut state = State::load_checked(root)?;
match state.newest_alive() {
Some((id, p)) => {
eprintln!(
"watchdog: MemAvailable {:.1} GB < floor {:.1} — killing newest '{id}' (pid {})",
m.available_gb, floor, p.pid
);
proc::stop_owned(&p)?;
state.procs.remove(&id);
state.save(root)?;
}
None => eprintln!(
"watchdog: MemAvailable {:.1} GB < floor {:.1} but no Lumbridge Compute-managed model to kill!",
m.available_gb, floor
),
}
}
std::thread::sleep(Duration::from_secs(1));
}
}
// ---- helpers ---------------------------------------------------------------
fn scene_footprint(s: &Scene, reg: &Registry) -> f64 {
s.models
.iter()
.filter_map(|id| reg.models.get(id))
.map(|m| m.footprint_gb)
.sum()
}
fn port_str(p: Option<u16>) -> String {
p.map(|p| p.to_string()).unwrap_or_else(|| "-".to_string())
}
fn truncate(s: &str, n: usize) -> String {
if s.chars().count() <= n {
s.to_string()
} else {
format!(
"{}",
s.chars().take(n.saturating_sub(1)).collect::<String>()
)
}
}
+617
View File
@@ -0,0 +1,617 @@
//! MCP server: the Governor's surface, exposed to agents as tools.
//!
//! Agents already drive Compute by shelling out to the CLI and scraping the
//! column-aligned output. That works until a column moves. Speaking MCP means
//! the agent gets the same numbers the Governor reasons about, as data, with a
//! schema attached — and it means the *shape* of what an agent may do becomes
//! something this file decides rather than something `bash` decides.
//!
//! Two deliberate constraints shape everything below.
//!
//! **Read-first.** Every tool here except `scene_activate` is a pure read.
//! Activating a Scene stops every registered model that is not in it, which on
//! this box means taking down whatever is currently serving — a live voice
//! pipeline included. There is no undo an agent can reach for: if the target
//! Scene then fails to come up, recovery depends on rollback that may itself
//! fail. That asymmetry is why `scene adopt`, `scene resume`, and
//! `scene deactivate` are absent entirely; they are operator verbs whose
//! correctness depends on knowing what the box was doing five minutes ago.
//! `scene_activate` is exposed because planning a switch is genuinely the
//! useful thing an agent wants, and it defaults to planning only.
//!
//! **The safety boundary is the operator's, not the agent's.** A `dry_run`
//! argument defaulting to `true` documents intent but guards nothing: the agent
//! writes the arguments, so it can write `false`. The only boundary an agent
//! cannot cross is one set before it connects, so a real transition also
//! requires `lumbridge-compute mcp --allow-activate`, chosen by the human who
//! launched the server. Without that flag `dry_run: false` is refused, and the
//! refusal says so rather than silently planning instead.
use anyhow::{Context, Result};
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{Implementation, ServerCapabilities, ServerInfo};
use rmcp::{tool, tool_handler, tool_router, ErrorData, Json, ServerHandler, ServiceExt};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::os::fd::FromRawFd;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use crate::config::{find_scene, load_scenes, Registry};
use crate::proc::State;
use crate::{eval, governor, lifecycle, mem};
pub fn run(root: &Path, allow_activate: bool) -> Result<()> {
let transport_stdout = hand_over_stdout()?;
// The banner has to go to stderr for the same reason everything else does;
// it doubles as confirmation to the operator that the mutating tool is off.
eprintln!(
"Lumbridge Compute MCP server on stdio · root {} · scene activation {}",
root.display(),
if allow_activate {
"ENABLED (--allow-activate)"
} else {
"disabled; plan only"
}
);
let server = ComputeMcp::new(root.to_path_buf(), allow_activate);
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("building the MCP runtime")?;
runtime.block_on(async move {
let transport = (
tokio::io::stdin(),
tokio::fs::File::from_std(transport_stdout),
);
let service = server
.serve(transport)
.await
.context("negotiating the MCP stdio session")?;
service.waiting().await.context("serving MCP over stdio")?;
Ok(())
})
}
/// Hand the real stdout to the transport and point fd 1 at stderr.
///
/// On stdio transport fd 1 *is* the JSON-RPC framing, and this crate reports
/// progress with `println!` throughout — `lifecycle::activate` narrates every
/// stop and start. One such line interleaved into the framing desynchronises
/// the client mid-transition, which is the worst possible moment to lose it.
/// Rather than audit every print (and every future one), move the file
/// descriptor: library output lands on stderr, where operators already read
/// this server's logs, and the protocol gets a channel nothing else can write.
fn hand_over_stdout() -> Result<std::fs::File> {
let saved = unsafe { libc::dup(libc::STDOUT_FILENO) };
if saved < 0 {
return Err(std::io::Error::last_os_error())
.context("duplicating stdout for the MCP transport");
}
if unsafe { libc::dup2(libc::STDERR_FILENO, libc::STDOUT_FILENO) } < 0 {
let error = std::io::Error::last_os_error();
unsafe { libc::close(saved) };
return Err(error).context("redirecting stdout to stderr");
}
// SAFETY: `saved` is a fresh descriptor from `dup` that nothing else owns.
Ok(unsafe { std::fs::File::from_raw_fd(saved) })
}
#[derive(Clone)]
struct ComputeMcp {
root: Arc<PathBuf>,
allow_activate: bool,
tool_router: ToolRouter<ComputeMcp>,
}
#[tool_router]
impl ComputeMcp {
fn new(root: PathBuf, allow_activate: bool) -> Self {
Self {
root: Arc::new(root),
allow_activate,
tool_router: Self::tool_router(),
}
}
/// Governor status for this node: unified-memory totals, the hard budget
/// and its margins, how much is committed by models that are actually
/// serving, how much is still admittable, the persisted Scene state, and
/// the last transition or watchdog error. Read this before reasoning about
/// whether anything else will fit.
#[tool(
name = "governor_status",
annotations(
title = "Governor status",
read_only_hint = true,
open_world_hint = false
)
)]
async fn governor_status(&self) -> Result<Json<StatusReport>, ErrorData> {
let root = self.root.clone();
offload(move || collect_status(&root)).await.map(Json)
}
/// List every model in the node's registry with its declared worst-case
/// footprint, serving port, and whether it is currently serving. Footprints
/// are what admission control is decided against, not observed usage.
#[tool(
name = "model_list",
annotations(
title = "List registered models",
read_only_hint = true,
open_world_hint = false
)
)]
async fn model_list(&self) -> Result<Json<ModelList>, ErrorData> {
let root = self.root.clone();
offload(move || collect_models(&root)).await.map(Json)
}
/// List the Scenes this node can activate, with each one's total footprint
/// and member models. A Scene is a named set of models brought up as a
/// single unit.
#[tool(
name = "scene_list",
annotations(title = "List scenes", read_only_hint = true, open_world_hint = false)
)]
async fn scene_list(&self) -> Result<Json<SceneList>, ErrorData> {
let root = self.root.clone();
offload(move || collect_scenes(&root)).await.map(Json)
}
/// Show one Scene: its models, their footprints and live state, and the
/// Governor's admission verdict — whether the Scene fits its budget once
/// the safety margin is added, and by how much it fits or misses.
#[tool(
name = "scene_show",
annotations(
title = "Show a scene and its admission verdict",
read_only_hint = true,
open_world_hint = false
)
)]
async fn scene_show(
&self,
Parameters(params): Parameters<SceneNameParams>,
) -> Result<Json<SceneDetail>, ErrorData> {
let root = self.root.clone();
offload(move || collect_scene(&root, &params.name))
.await
.map(Json)
}
/// List the evaluation suites available on this node, with their versions
/// and case counts. Running a suite is deliberately not exposed: it drives
/// real load against a serving model for an unbounded time.
#[tool(
name = "eval_list",
annotations(
title = "List eval suites",
read_only_hint = true,
open_world_hint = false
)
)]
async fn eval_list(&self) -> Result<Json<EvalList>, ErrorData> {
let root = self.root.clone();
offload(move || collect_evals(&root)).await.map(Json)
}
/// Plan — or, if the operator allowed it, perform — a Scene transition.
///
/// With `dry_run` true (the default) this runs every check that guards a
/// real transition and returns the models it would stop and start, changing
/// nothing. With `dry_run` false it actually switches the node: models
/// outside the Scene are stopped, which will interrupt anything they serve,
/// and the call blocks until every Scene model reports healthy or the
/// transition rolls back. A real transition additionally requires the
/// server to have been started with `--allow-activate`; without it,
/// `dry_run: false` is refused. Ask a human before setting it.
#[tool(
name = "scene_activate",
annotations(
title = "Activate a scene (plans by default)",
read_only_hint = false,
destructive_hint = true,
idempotent_hint = false,
open_world_hint = false
)
)]
async fn scene_activate(
&self,
Parameters(params): Parameters<SceneActivateParams>,
) -> Result<Json<ActivationReport>, ErrorData> {
if !params.dry_run && !self.allow_activate {
return Err(ErrorData::invalid_params(
"this MCP server is running plan-only: a real Scene transition needs an \
operator to restart it as `lumbridge-compute mcp --allow-activate`. \
Re-run with dry_run: true to see the plan.",
None,
));
}
let root = self.root.clone();
offload(move || activate_scene(&root, &params.name, params.dry_run))
.await
.map(Json)
}
}
#[tool_handler(router = self.tool_router)]
impl ServerHandler for ComputeMcp {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(Implementation::new(
"lumbridge-compute",
env!("CARGO_PKG_VERSION"),
))
.with_instructions(
"Lumbridge Compute runs many AI models on one unified-memory box. There is a \
single memory pool shared by CPU and GPU, so over-committing does not fail \
gracefully — the machine thrashes and wedges before the OOM killer acts. The \
Governor prevents that by admitting a model only if its declared footprint \
plus a safety margin still fits a hard budget.\n\n\
Start from `governor_status` for what is committed and what is admittable, \
and `scene_show` for whether a named Scene would be admitted. Everything is \
read-only except `scene_activate`, which plans by default. Activating a Scene \
stops every registered model outside it, so treat it as a production change \
and get a human's agreement before asking for a non-dry run.",
)
}
}
// ---- tool parameters -------------------------------------------------------
#[derive(Debug, Deserialize, JsonSchema)]
struct SceneNameParams {
/// Scene name, as reported by `scene_list`.
name: String,
}
#[derive(Debug, Deserialize, JsonSchema)]
struct SceneActivateParams {
/// Scene name, as reported by `scene_list`.
name: String,
/// Validate and report the plan without changing anything. Defaults to
/// true; set it to false only with a human's explicit agreement.
#[serde(default = "yes")]
dry_run: bool,
}
fn yes() -> bool {
true
}
// ---- tool results ----------------------------------------------------------
//
// Lists are wrapped in objects rather than returned bare: MCP structured
// content must be a JSON object, and a named field leaves room to report
// alongside the list later without breaking a client's schema.
#[derive(Debug, Serialize, JsonSchema)]
pub(crate) struct StatusReport {
memory_total_gb: f64,
memory_available_gb: f64,
budget_gb: f64,
safety_margin_gb: f64,
watchdog_floor_gb: f64,
/// Sum of the footprints of every registered model currently serving.
committed_gb: f64,
/// What a new model could still claim without breaching the safety margin.
headroom_gb: f64,
/// Scene whose exact model health was last verified.
active_scene: Option<String>,
/// Scene the node should return to after a restart.
desired_scene: Option<String>,
/// Previous proven Scene, used if the desired one cannot resume.
last_known_good_scene: Option<String>,
last_error: Option<String>,
running: Vec<RunningModel>,
}
#[derive(Debug, Serialize, JsonSchema)]
struct RunningModel {
id: String,
name: String,
footprint_gb: f64,
port: Option<u16>,
/// True when Compute owns this process identity and may signal it. False
/// means the model is serving but was started outside Compute, so a Scene
/// transition will refuse to touch it until it is adopted.
compute_owned: bool,
}
#[derive(Debug, Serialize, JsonSchema)]
pub(crate) struct ModelList {
models: Vec<ModelEntry>,
}
#[derive(Debug, Serialize, JsonSchema)]
struct ModelEntry {
id: String,
name: String,
/// Declared worst-case unified memory once serving, not observed usage.
footprint_gb: f64,
port: Option<u16>,
running: bool,
}
#[derive(Debug, Serialize, JsonSchema)]
pub(crate) struct SceneList {
scenes: Vec<SceneEntry>,
}
#[derive(Debug, Serialize, JsonSchema)]
struct SceneEntry {
name: String,
version: u32,
description: String,
/// Total footprint of the Scene's models that exist in the registry.
footprint_gb: f64,
models: Vec<String>,
}
#[derive(Debug, Serialize, JsonSchema)]
pub(crate) struct SceneDetail {
name: String,
version: u32,
description: String,
budget_gb: f64,
models: Vec<SceneModel>,
/// Model ids the Scene references that the registry does not define.
missing_models: Vec<String>,
footprint_gb: f64,
/// Footprint plus the Governor's safety margin — the figure compared
/// against the budget.
required_gb: f64,
/// The Governor's verdict: would this Scene be admitted?
admits: bool,
verdict: String,
}
#[derive(Debug, Serialize, JsonSchema)]
struct SceneModel {
id: String,
/// Absent when the id is missing from the registry.
name: Option<String>,
footprint_gb: Option<f64>,
running: bool,
}
#[derive(Debug, Serialize, JsonSchema)]
pub(crate) struct EvalList {
suites: Vec<EvalSuite>,
}
#[derive(Debug, Serialize, JsonSchema)]
struct EvalSuite {
name: String,
version: u32,
cases: usize,
description: String,
}
#[derive(Debug, Serialize, JsonSchema)]
struct ActivationReport {
scene: String,
budget_gb: f64,
/// Serving models outside the Scene, which the transition stops first.
stop: Vec<String>,
/// Scene models not yet serving, in the order they would be admitted.
start: Vec<String>,
/// False when this was a plan and the node was left untouched.
applied: bool,
}
// ---- collectors ------------------------------------------------------------
//
// These compose the same primitives the CLI prints from — `governor`, `mem`,
// `config`, `proc::State` — rather than parsing the CLI's output. The numbers
// an agent sees are therefore the numbers admission control uses, by
// construction.
pub(crate) fn collect_status(root: &Path) -> Result<StatusReport> {
let registry = Registry::load(root)?;
let memory = mem::read()?;
let committed = governor::committed_gb(&registry);
let managed = State::load_checked(root)?;
let budget = governor::DEFAULT_BUDGET_GB;
let running = governor::running_ids(&registry)
.into_iter()
.map(|id| {
let model = &registry.models[&id];
RunningModel {
name: model.name.clone(),
footprint_gb: model.footprint_gb,
port: model.serve.port,
compute_owned: managed
.procs
.get(&id)
.is_some_and(|process| process.owned_alive()),
id,
}
})
.collect();
Ok(StatusReport {
memory_total_gb: memory.total_gb,
memory_available_gb: memory.available_gb,
budget_gb: budget,
safety_margin_gb: governor::SAFETY_MARGIN_GB,
watchdog_floor_gb: governor::WATCHDOG_FLOOR_GB,
// `+ 0.0` normalises the negative zero an empty sum can produce, which serialises as
// "-0.0" and reads as a bug to anyone consuming this JSON.
committed_gb: committed + 0.0,
// The BINDING constraint, not the declared one. This used to be
// `budget - committed - margin`, which ignores how much memory the box actually has —
// so on a 62 GB machine with a 100 GB budget it reported 92 GB of headroom that the
// next admission would refuse. cmd_status was fixed when admission started consulting
// MemAvailable; this path was missed, which meant every agent driving the node over
// MCP got the optimistic number.
headroom_gb: governor::headroom_gb(committed, budget, Some(memory.available_gb)).max(0.0),
active_scene: managed.active_scene,
desired_scene: managed.desired_scene,
last_known_good_scene: managed.last_known_good_scene,
last_error: managed.last_error,
running,
})
}
pub(crate) fn collect_models(root: &Path) -> Result<ModelList> {
let registry = Registry::load(root)?;
Ok(ModelList {
models: registry
.models
.iter()
.map(|(id, model)| ModelEntry {
id: id.clone(),
name: model.name.clone(),
footprint_gb: model.footprint_gb,
port: model.serve.port,
running: governor::is_running(model),
})
.collect(),
})
}
pub(crate) fn collect_scenes(root: &Path) -> Result<SceneList> {
let registry = Registry::load(root)?;
Ok(SceneList {
scenes: load_scenes(root)?
.into_iter()
.map(|scene| SceneEntry {
name: scene.metadata.name,
version: scene.metadata.version,
description: scene.metadata.description,
footprint_gb: scene
.models
.iter()
.filter_map(|id| registry.models.get(id))
.map(|model| model.footprint_gb)
.sum(),
models: scene.models,
})
.collect(),
})
}
pub(crate) fn collect_scene(root: &Path, name: &str) -> Result<SceneDetail> {
let registry = Registry::load(root)?;
let scene = find_scene(root, name)?;
let budget = scene.budget_gb.unwrap_or(governor::DEFAULT_BUDGET_GB);
let mut footprint = 0.0;
let mut missing = Vec::new();
let mut models = Vec::new();
for id in &scene.models {
match registry.models.get(id) {
Some(model) => {
footprint += model.footprint_gb;
models.push(SceneModel {
id: id.clone(),
name: Some(model.name.clone()),
footprint_gb: Some(model.footprint_gb),
running: governor::is_running(model),
});
}
None => {
missing.push(id.clone());
models.push(SceneModel {
id: id.clone(),
name: None,
footprint_gb: None,
running: false,
});
}
}
}
// Same three-way verdict the CLI prints: a missing id is fatal regardless
// of arithmetic, because the Scene cannot be resolved at all.
let required = footprint + governor::SAFETY_MARGIN_GB;
let (admits, verdict) = if !missing.is_empty() {
(
false,
format!("{} model(s) missing from the registry", missing.len()),
)
} else if required <= budget {
(
true,
format!("fits with {:.1} GB to spare", budget - required),
)
} else {
(
false,
format!("exceeds the budget by {:.1} GB", required - budget),
)
};
Ok(SceneDetail {
name: scene.metadata.name,
version: scene.metadata.version,
description: scene.metadata.description,
budget_gb: budget,
models,
missing_models: missing,
footprint_gb: footprint,
required_gb: required,
admits,
verdict,
})
}
pub(crate) fn collect_evals(root: &Path) -> Result<EvalList> {
Ok(EvalList {
suites: eval::catalogue(root)?
.into_iter()
.map(|suite| EvalSuite {
name: suite.name,
version: suite.version,
cases: suite.cases,
description: suite.description,
})
.collect(),
})
}
fn activate_scene(root: &Path, name: &str, dry_run: bool) -> Result<ActivationReport> {
// Plan first either way. On a real transition this is the plan we report;
// `activate` re-derives and re-validates its own under the transition lock,
// so the report describes intent and the lock still owns the truth.
let plan = lifecycle::plan(root, name)?;
if !dry_run {
lifecycle::activate(root, name, false)?;
}
Ok(ActivationReport {
scene: plan.scene,
budget_gb: plan.budget_gb,
stop: plan.stop,
start: plan.start,
applied: !dry_run,
})
}
/// Run a synchronous Governor call off the transport's thread.
///
/// Every read here touches `/proc` and probes serving ports with blocking
/// socket timeouts, and a real transition can sit for minutes waiting on model
/// health. Doing that inline would stall the JSON-RPC reader for the duration,
/// so the synchronous core stays on the blocking pool where it belongs.
async fn offload<T, F>(work: F) -> Result<T, ErrorData>
where
F: FnOnce() -> Result<T> + Send + 'static,
T: Send + 'static,
{
match tokio::task::spawn_blocking(work).await {
Ok(Ok(value)) => Ok(value),
// `{:#}` keeps anyhow's context chain, which is where the actionable
// half of a Compute error lives ("adopt the active Scene first").
Ok(Err(error)) => Err(ErrorData::internal_error(format!("{error:#}"), None)),
Err(join) => Err(ErrorData::internal_error(
format!("Compute worker task failed: {join}"),
None,
)),
}
}
+89
View File
@@ -0,0 +1,89 @@
//! Unified-memory sensing. On a shared-memory box there's one pool, so `MemAvailable`
//! from /proc/meminfo is the single source of truth the Governor guards.
use anyhow::{Context, Result};
use std::fs;
pub struct MemInfo {
pub total_gb: f64,
pub available_gb: f64,
}
pub fn read() -> Result<MemInfo> {
let s = fs::read_to_string("/proc/meminfo").context("reading /proc/meminfo")?;
Ok(parse_meminfo(&s))
}
/// Split out from `read` so the parser is testable without a real /proc.
pub fn parse_meminfo(s: &str) -> MemInfo {
let mut total = 0.0;
let mut available = 0.0;
for line in s.lines() {
if let Some(v) = line.strip_prefix("MemTotal:") {
total = kb_to_gb(v);
} else if let Some(v) = line.strip_prefix("MemAvailable:") {
available = kb_to_gb(v);
}
}
MemInfo {
total_gb: total,
available_gb: available,
}
}
/// Parse a " 123456 kB" meminfo value into GB.
fn kb_to_gb(v: &str) -> f64 {
v.split_whitespace()
.next()
.and_then(|n| n.parse::<f64>().ok())
.unwrap_or(0.0)
/ 1024.0
/ 1024.0
}
#[cfg(test)]
mod tests {
use super::*;
const SPARK: &str = "\
MemTotal: 127512345 kB
MemFree: 2048000 kB
MemAvailable: 104857600 kB
Buffers: 123456 kB
";
#[test]
fn reads_total_and_available_and_ignores_other_fields() {
let m = parse_meminfo(SPARK);
assert!((m.total_gb - 121.6).abs() < 0.1, "total was {}", m.total_gb);
assert!(
(m.available_gb - 100.0).abs() < 0.1,
"available was {}",
m.available_gb
);
}
#[test]
fn missing_fields_read_as_zero_rather_than_panicking() {
// A zero here is what makes admission refuse, so an unparseable /proc must not
// look like an empty box with room to spare.
let m = parse_meminfo("SomethingElse: 1 kB\n");
assert_eq!(m.total_gb, 0.0);
assert_eq!(m.available_gb, 0.0);
}
#[test]
fn malformed_values_do_not_panic() {
let m = parse_meminfo("MemTotal: not-a-number kB\nMemAvailable:\n");
assert_eq!(m.total_gb, 0.0);
assert_eq!(m.available_gb, 0.0);
}
#[test]
fn memavailable_is_not_confused_with_memfree() {
// MemFree is much smaller than MemAvailable on a box with page cache; picking the
// wrong one would make the watchdog fire constantly.
let m = parse_meminfo(SPARK);
assert!(m.available_gb > 50.0);
}
}
+470
View File
@@ -0,0 +1,470 @@
//! Process lifecycle: spawn a model in its own process group, track it in a state
//! file, and stop it (SIGTERM → SIGKILL to the whole group). The Governor decides
//! *whether* to start; this module *how*.
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::io::Write;
use std::os::unix::process::CommandExt;
use std::path::{Path, PathBuf};
use std::process::{self, Command, Stdio};
use std::time::Duration;
use crate::config::Model;
/// One Lumbridge Compute-managed process. `seq` is a monotonic launch counter so the watchdog
/// can always find the *newest* model to sacrifice first.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Proc {
pub pid: i32,
pub seq: u64,
pub port: Option<u16>,
/// Linux boot id plus process start ticks bind ownership to one process
/// incarnation, preventing a reused PID from ever being signalled.
#[serde(default)]
pub boot_id: Option<String>,
#[serde(default)]
pub start_time_ticks: Option<u64>,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct State {
#[serde(default)]
pub seq: u64,
#[serde(default)]
pub procs: BTreeMap<String, Proc>,
#[serde(default)]
pub desired_scene: Option<String>,
#[serde(default)]
pub active_scene: Option<String>,
/// Previous proven scene used if the desired scene cannot resume.
#[serde(default)]
pub last_known_good_scene: Option<String>,
#[serde(default)]
pub transition_scene: Option<String>,
#[serde(default)]
pub last_error: Option<String>,
}
fn state_dir(root: &Path) -> PathBuf {
root.join(".kuda")
}
impl State {
pub fn load_checked(root: &Path) -> Result<State> {
let p = state_dir(root).join("state.yaml");
match fs::read_to_string(&p) {
Ok(s) => serde_yaml::from_str(&s)
.with_context(|| format!("parsing persisted state {}", p.display())),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(State::default()),
Err(error) => Err(error).with_context(|| format!("reading {}", p.display())),
}
}
pub fn save(&self, root: &Path) -> Result<()> {
let dir = state_dir(root);
fs::create_dir_all(&dir)?;
let p = dir.join("state.yaml");
let tmp = dir.join(format!(".state.yaml.{}.tmp", process::id()));
let mut file = OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&tmp)
.with_context(|| format!("creating {}", tmp.display()))?;
file.write_all(serde_yaml::to_string(self)?.as_bytes())?;
file.sync_all()?;
fs::rename(&tmp, &p).with_context(|| format!("atomically replacing {}", p.display()))?;
File::open(&dir)
.with_context(|| format!("opening state directory {}", dir.display()))?
.sync_all()
.with_context(|| format!("syncing state directory {}", dir.display()))?;
Ok(())
}
/// Newest still-alive managed model id, if any.
pub fn newest_alive(&self) -> Option<(String, Proc)> {
self.procs
.iter()
.filter(|(_, p)| p.owned_alive())
.max_by_key(|(_, p)| p.seq)
.map(|(id, p)| (id.clone(), p.clone()))
}
}
#[derive(Debug)]
struct ProcessIdentity {
boot_id: String,
start_time_ticks: u64,
process_group_id: i32,
state: char,
}
impl Proc {
/// Capture the identity of a process that Lumbridge Compute just spawned or
/// that an operator explicitly adopted from legacy state.
pub fn capture(pid: i32, seq: u64, port: Option<u16>) -> Result<Proc> {
let identity = process_identity(pid)?;
if identity.process_group_id != pid {
bail!(
"refusing process {pid}: process group {} does not equal pid",
identity.process_group_id
);
}
Ok(Proc {
pid,
seq,
port,
boot_id: Some(identity.boot_id),
start_time_ticks: Some(identity.start_time_ticks),
})
}
pub fn owned_alive(&self) -> bool {
self.validate_owned().is_ok()
}
fn validate_owned(&self) -> Result<()> {
let expected_boot = self
.boot_id
.as_deref()
.context("legacy process record has no boot identity; adopt the active scene first")?;
let expected_start = self
.start_time_ticks
.context("legacy process record has no start identity; adopt the active scene first")?;
let current = process_identity(self.pid)?;
if current.boot_id != expected_boot || current.start_time_ticks != expected_start {
bail!(
"pid {} no longer identifies the process Compute started",
self.pid
);
}
if current.state == 'Z' {
bail!("pid {} is a zombie awaiting reap", self.pid);
}
if current.process_group_id != self.pid {
bail!(
"pid {} now belongs to process group {}; refusing group signal",
self.pid,
current.process_group_id
);
}
Ok(())
}
}
fn process_identity(pid: i32) -> Result<ProcessIdentity> {
let boot_id = fs::read_to_string("/proc/sys/kernel/random/boot_id")
.context("reading Linux boot id")?
.trim()
.to_string();
let stat_path = format!("/proc/{pid}/stat");
let stat = fs::read_to_string(&stat_path)
.with_context(|| format!("reading process identity {stat_path}"))?;
let (state, process_group_id, start_time_ticks) = parse_stat_identity(&stat)?;
Ok(ProcessIdentity {
boot_id,
start_time_ticks,
process_group_id,
state,
})
}
fn parse_stat_identity(stat: &str) -> Result<(char, i32, u64)> {
// `comm` is parenthesized and may contain spaces, so locate its final `)`
// before indexing the fields that follow it. The remainder starts at field 3.
let close = stat
.rfind(')')
.context("malformed /proc stat: missing process name terminator")?;
let fields: Vec<&str> = stat[close + 1..].split_whitespace().collect();
let state = fields
.first()
.and_then(|value| value.chars().next())
.context("malformed /proc stat: missing process state")?;
let process_group_id = fields
.get(2)
.context("malformed /proc stat: missing process group")?
.parse::<i32>()
.context("parsing process group id")?;
let start_time_ticks = fields
.get(19)
.context("malformed /proc stat: missing start time")?
.parse::<u64>()
.context("parsing process start time")?;
Ok((state, process_group_id, start_time_ticks))
}
/// Expand a leading `~/` to `$HOME`.
fn expand(s: &str) -> String {
if let Some(rest) = s.strip_prefix("~/") {
if let Ok(home) = std::env::var("HOME") {
return format!("{home}/{rest}");
}
}
s.to_string()
}
/// The fraction of the pool a vLLM model may reserve, derived from its declared
/// footprint rather than typed separately.
///
/// `footprint_gb` is what the Governor budgets against; `--gpu-memory-utilization` is
/// what actually reserves the memory. Keeping them as two hand-copied numbers means
/// admission control can be arithmetically correct and still wrong about the machine —
/// the registry says a model takes 66 GB while the flag lets it reserve 0.55 of a
/// 121.6 GB pool, which is 67. Deriving one from the other makes the declared footprint
/// the single source of truth.
///
/// An explicit `gpu-memory-utilization` in `serve.args` always wins; this only fills in
/// the gap. Returns `None` when the pool size is unknown, in which case vLLM's own
/// default applies exactly as before.
fn derived_gpu_memory_utilization(m: &Model, mem_total_gb: Option<f64>) -> Option<f64> {
if m.serve.args.contains_key("gpu-memory-utilization") {
return None;
}
let total = mem_total_gb?;
if total <= 0.0 || m.footprint_gb <= 0.0 {
return None;
}
let fraction = m.footprint_gb / total;
// Never hand vLLM a fraction that would reserve the whole box.
(fraction > 0.0 && fraction < 0.95).then_some((fraction * 1000.0).round() / 1000.0)
}
/// `kind`-based builder. argv[0] is the program.
pub fn build_argv(m: &Model, mem_total_gb: Option<f64>) -> Result<Vec<String>> {
if !m.serve.command.is_empty() {
return Ok(m.serve.command.iter().map(|s| expand(s)).collect());
}
match m.serve.kind.as_str() {
"vllm" => {
let executable = m
.serve
.executable
.as_deref()
.map(expand)
.unwrap_or_else(|| "vllm".to_string());
let mut a = vec![executable, "serve".to_string()];
if let Some(w) = &m.serve.weights {
a.push(expand(w));
}
if !m.serve.served_name.is_empty() {
a.push("--served-model-name".into());
a.extend(m.serve.served_name.iter().cloned());
}
a.push("--host".into());
a.push("0.0.0.0".into());
if let Some(p) = m.serve.port {
a.push("--port".into());
a.push(p.to_string());
}
if let Some(fraction) = derived_gpu_memory_utilization(m, mem_total_gb) {
a.push("--gpu-memory-utilization".into());
a.push(format!("{fraction}"));
}
for (k, v) in &m.serve.args {
render_arg(&mut a, k, v);
}
Ok(a)
}
other => bail!(
"model '{}' has no explicit `command` and kind '{}' has no builder yet — \
add a `command: [...]` to the registry entry",
m.name,
other
),
}
}
fn render_arg(out: &mut Vec<String>, k: &str, v: &serde_yaml::Value) {
use serde_yaml::Value;
let flag = format!("--{k}");
match v {
Value::Bool(true) => out.push(flag),
Value::Bool(false) => {}
Value::String(s) => {
out.push(flag);
out.push(s.clone());
}
Value::Number(n) => {
out.push(flag);
out.push(n.to_string());
}
_ => {
out.push(flag);
out.push(
serde_yaml::to_string(v)
.unwrap_or_default()
.trim()
.to_string(),
);
}
}
}
/// Spawn `id`'s model detached in its own process group, logs → compatibility path `.kuda/logs/<id>.log`.
/// Returns the child pid (also its process-group id).
pub fn spawn(root: &Path, id: &str, m: &Model) -> Result<i32> {
let argv = build_argv(m, crate::mem::read().ok().map(|info| info.total_gb))?;
let logdir = state_dir(root).join("logs");
fs::create_dir_all(&logdir)?;
let log = File::create(logdir.join(format!("{id}.log")))?;
let errlog = log.try_clone()?;
let mut cmd = Command::new(&argv[0]);
cmd.args(&argv[1..])
.stdin(Stdio::null())
.stdout(Stdio::from(log))
.stderr(Stdio::from(errlog))
.process_group(0); // own group → clean group-kill, immune to CLI's signals
for (k, v) in &m.serve.env {
cmd.env(k, expand(v));
}
let mut child = cmd
.spawn()
.with_context(|| format!("spawning '{id}': {}", argv.join(" ")))?;
let pid = child.id() as i32;
// A resident agent may outlive many model processes. Reap each child when
// it exits so failed runtimes cannot accumulate as zombies under the agent.
std::thread::spawn(move || {
let _ = child.wait();
});
Ok(pid)
}
/// Stop only the exact process group captured in this ownership record.
pub fn stop_owned(proc: &Proc) -> Result<()> {
if !proc.owned_alive() {
return Ok(());
}
proc.validate_owned()?;
let term = unsafe { libc::kill(-proc.pid, libc::SIGTERM) };
if term != 0 {
return Err(std::io::Error::last_os_error())
.with_context(|| format!("sending SIGTERM to owned process group {}", proc.pid));
}
for _ in 0..60 {
if !proc.owned_alive() {
return Ok(());
}
std::thread::sleep(Duration::from_millis(100));
}
proc.validate_owned()?;
let kill = unsafe { libc::kill(-proc.pid, libc::SIGKILL) };
if kill != 0 {
return Err(std::io::Error::last_os_error())
.with_context(|| format!("sending SIGKILL to owned process group {}", proc.pid));
}
for _ in 0..20 {
if !proc.owned_alive() {
return Ok(());
}
std::thread::sleep(Duration::from_millis(100));
}
bail!("owned process group {} survived SIGKILL", proc.pid)
}
#[cfg(test)]
mod tests {
use super::*;
/// Build a Model the way the registry does — through serde — so these tests exercise
/// the real deserialization path rather than a hand-assembled struct.
fn vllm_model(footprint_gb: f64, extra_args: &str) -> Model {
let yaml = format!(
"name: brain\nfootprint_gb: {footprint_gb}\nserve:\n kind: vllm\n port: 8001\n args:\n max-model-len: 8192\n{extra_args}"
);
serde_yaml::from_str(&yaml).expect("test model yaml")
}
#[test]
fn gpu_memory_utilization_is_derived_from_the_declared_footprint() {
// 66 GB of a 121.6 GB pool is 0.543 — not the 0.55 that used to be typed in
// separately, which is the whole point: one number, not two that can drift.
let m = vllm_model(66.0, "");
assert_eq!(derived_gpu_memory_utilization(&m, Some(121.6)), Some(0.543));
let argv = build_argv(&m, Some(121.6)).unwrap();
let i = argv
.iter()
.position(|a| a == "--gpu-memory-utilization")
.unwrap();
assert_eq!(argv[i + 1], "0.543");
}
#[test]
fn an_explicit_flag_in_the_registry_always_wins() {
let m = vllm_model(66.0, " gpu-memory-utilization: 0.8\n");
assert_eq!(derived_gpu_memory_utilization(&m, Some(121.6)), None);
// ...and it is rendered exactly once, from serve.args.
let argv = build_argv(&m, Some(121.6)).unwrap();
assert_eq!(
argv.iter()
.filter(|a| *a == "--gpu-memory-utilization")
.count(),
1
);
let i = argv
.iter()
.position(|a| a == "--gpu-memory-utilization")
.unwrap();
assert_eq!(argv[i + 1], "0.8");
}
#[test]
fn unknown_pool_size_leaves_vllms_own_default_alone() {
let m = vllm_model(66.0, "");
assert_eq!(derived_gpu_memory_utilization(&m, None), None);
assert!(!build_argv(&m, None)
.unwrap()
.iter()
.any(|a| a == "--gpu-memory-utilization"));
}
#[test]
fn a_footprint_that_would_claim_the_whole_box_is_not_derived() {
// Better to let vLLM apply its own default than to hand it 0.99 and wedge the box.
let m = vllm_model(120.0, "");
assert_eq!(derived_gpu_memory_utilization(&m, Some(121.6)), None);
}
#[test]
fn nonsense_inputs_do_not_produce_a_flag() {
assert_eq!(
derived_gpu_memory_utilization(&vllm_model(0.0, ""), Some(121.6)),
None
);
assert_eq!(
derived_gpu_memory_utilization(&vllm_model(66.0, ""), Some(0.0)),
None
);
}
#[test]
fn an_explicit_command_bypasses_the_builder_entirely() {
let mut m = vllm_model(66.0, "");
m.serve.command = vec!["python".into(), "-m".into(), "server".into()];
assert_eq!(build_argv(&m, Some(121.6)).unwrap(), m.serve.command);
}
#[test]
fn proc_stat_parser_handles_spaces_in_process_name() {
let stat = "123 (worker process) S 1 123 123 0 -1 0 0 0 0 0 0 0 0 0 20 0 1 0 4567";
let (state, pgrp, start) = parse_stat_identity(stat).unwrap();
assert_eq!(state, 'S');
assert_eq!(pgrp, 123);
assert_eq!(start, 4567);
}
#[test]
fn legacy_records_are_never_treated_as_owned() {
let proc = Proc {
pid: std::process::id() as i32,
seq: 1,
port: None,
boot_id: None,
start_time_ticks: None,
};
assert!(!proc.owned_alive());
}
}