@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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(®istry, &scene)?;
|
||||
|
||||
let target: HashSet<&str> = scene.models.iter().map(String::as_str).collect();
|
||||
let running = governor::running_ids(®istry);
|
||||
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 = ®istry.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(®istry.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(®istry.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(®istry.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 = ®istry.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(®istry, &scene)?;
|
||||
let plan = plan_transition(®istry, &scene);
|
||||
let mut state = State::load_checked(root)?;
|
||||
state.procs.retain(|_, process| process.owned_alive());
|
||||
preflight(®istry, &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(®istry, &scene)?;
|
||||
let plan = plan_transition(®istry, &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(®istry, &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(®istry.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 = ®istry.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(®istry.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 = ®istry.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
@@ -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(®);
|
||||
let running = governor::running_ids(®);
|
||||
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 = ®.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 ®.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, ®);
|
||||
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
@@ -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, ¶ms.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, ¶ms.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(®istry);
|
||||
let managed = State::load_checked(root)?;
|
||||
let budget = governor::DEFAULT_BUDGET_GB;
|
||||
|
||||
let running = governor::running_ids(®istry)
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
let model = ®istry.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
@@ -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
@@ -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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user