overbrainer 0.10.0

//! Prometheus metrics of a project, served at `GET /metrics` while a command holds
//! the project. Counters start from the stage history, then follow the events of
//! every bus the command builds with them as its [`Observer`], so they are
//! cumulative per project. Retries are not in the history: they restart at zero
//! with each process.

use std::collections::BTreeMap;
use std::convert::Infallible;
use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, SystemTime};

use hyper::body::Incoming;
use hyper::header::{CONTENT_TYPE, HeaderValue};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::{TokioIo, TokioTimer};
use prometheus_client::encoding::EncodeLabelSet;
use prometheus_client::metrics::counter::Counter;
use prometheus_client::metrics::family::Family;
use prometheus_client::metrics::gauge::Gauge;
use prometheus_client::registry::{Metric, Registry};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::{JoinHandle, JoinSet};

mod system;

use self::system::{SystemFamilies, SystemSeries};
use crate::events::{Event, Observer, Stage};
use crate::history::Entry;
use crate::llm::Usage;
use crate::runpod::PodRecord;
use crate::runs::Runs;
use crate::train::{Phase, Phases, TrainMetric};

/// Content type of the exposition `prometheus_client` writes.
pub const CONTENT_TYPE_VALUE: &str = "application/openmetrics-text; version=1.0.0; charset=utf-8";

/// Model label of tokens and cost whose stage named no model.
const UNKNOWN_MODEL: &str = "unknown";

/// Connections served at once; a new one beyond is closed at once.
const MAX_CONNECTIONS: usize = 32;

/// How long a connection may take to send a request's headers.
const HEADER_TIMEOUT: Duration = Duration::from_secs(10);

/// Pause after a failed accept (out of file descriptors), so the loop never spins.
const ACCEPT_PAUSE: Duration = Duration::from_millis(100);

type FloatCounter = Counter<f64, AtomicU64>;
type FloatGauge = Gauge<f64, AtomicU64>;

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct ItemLabels {
    stage: &'static str,
    result: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct TokenLabels {
    stage: &'static str,
    model: String,
    direction: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct ModelLabels {
    stage: &'static str,
    model: String,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct StageLabels {
    stage: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct RunLabels {
    run_id: String,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct PhaseLabels {
    run_id: String,
    phase: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, EncodeLabelSet)]
struct VersionLabels {
    version: &'static str,
}

/// What one bus said so far that labels its next events.
#[derive(Debug, Default)]
struct Followed {
    /// The model of each stage, from [`Event::StageModel`].
    models: BTreeMap<Stage, String>,
    /// Stages started and not finished yet.
    running: Vec<Stage>,
    /// The run whose metrics come next, from [`Event::RunWatched`].
    run_id: Option<String>,
    /// The series the last [`Event::System`] set.
    system: Option<SystemSeries>,
    /// What the run's lines said of its phase.
    phases: Phases,
    /// Whether the run's job exited.
    exited: bool,
}

/// The metric families of a project and their registry.
#[derive(Debug)]
pub struct Metrics {
    registry: Registry,
    items: Family<ItemLabels, Counter>,
    tokens: Family<TokenLabels, Counter>,
    cost: Family<ModelLabels, FloatCounter>,
    running: Family<StageLabels, Gauge>,
    retries: Family<StageLabels, Counter>,
    train_step: Family<RunLabels, Gauge>,
    train_max_steps: Family<RunLabels, Gauge>,
    train_phase: Family<PhaseLabels, Gauge>,
    train_loss: Family<RunLabels, FloatGauge>,
    eval_loss: Family<RunLabels, FloatGauge>,
    learning_rate: Family<RunLabels, FloatGauge>,
    spend: Family<RunLabels, FloatGauge>,
    system: SystemFamilies,
    /// Where the Runpod spend is read at each scrape; none leaves it out.
    runs: Option<Runs>,
    /// Per bus ID.
    buses: Mutex<BTreeMap<usize, Followed>>,
    /// Held by a scrape from its spend refill to its encoding, so another never
    /// sees the spend family half refilled.
    scrape: Mutex<()>,
}

impl Metrics {
    /// The metrics, with the items, tokens and cost of `history` already counted.
    #[must_use]
    pub fn new(history: &[Entry]) -> Self {
        let mut registry = Registry::default();
        let items = register(
            &mut registry,
            "overbrainer_stage_items",
            "Items of the pipeline stages, by result",
        );
        let tokens = register(
            &mut registry,
            "overbrainer_tokens",
            "Tokens sent (in) and received (out) by the pipeline stages",
        );
        let cost = register(
            &mut registry,
            "overbrainer_cost_usd",
            "Cost of the pipeline stages in USD, when the model price is known",
        );
        let running = register(
            &mut registry,
            "overbrainer_stage_running",
            "Pipeline stages running now",
        );
        let retries = register(
            &mut registry,
            "overbrainer_item_retries",
            "Failed attempts retried, since this process started",
        );
        let train_step = register(
            &mut registry,
            "overbrainer_train_step",
            "Last optimizer step of a training run",
        );
        let train_max_steps = register(
            &mut registry,
            "overbrainer_train_max_steps",
            "Total optimizer steps of a training run, once known",
        );
        let train_phase = register(
            &mut registry,
            "overbrainer_train_phase",
            "What the job of a followed training run does: 1 for the current phase, 0 for the others",
        );
        let train_loss = register(
            &mut registry,
            "overbrainer_train_loss",
            "Last training loss of a run",
        );
        let eval_loss = register(
            &mut registry,
            "overbrainer_eval_loss",
            "Last evaluation loss of a run",
        );
        let learning_rate = register(
            &mut registry,
            "overbrainer_learning_rate",
            "Last learning rate of a run",
        );
        let spend = register(
            &mut registry,
            "overbrainer_runpod_spend_usd",
            "Estimated Runpod spend of a run in USD: rate times uptime",
        );
        let system = SystemFamilies::new(&mut registry);
        let build: Family<VersionLabels, Gauge> = register(
            &mut registry,
            "overbrainer_build_info",
            "Version of overbrainer",
        );
        build
            .get_or_create(&VersionLabels {
                version: env!("CARGO_PKG_VERSION"),
            })
            .set(1);
        let metrics = Self {
            registry,
            items,
            tokens,
            cost,
            running,
            retries,
            train_step,
            train_max_steps,
            train_phase,
            train_loss,
            eval_loss,
            learning_rate,
            spend,
            system,
            runs: None,
            buses: Mutex::default(),
            scrape: Mutex::default(),
        };
        metrics.seeded(history)
    }

    /// The metrics with every stage shown, running or not, and the items,
    /// tokens and cost of `history` counted.
    fn seeded(self, history: &[Entry]) -> Self {
        for stage in [
            Stage::Subtopics,
            Stage::Questions,
            Stage::Answers,
            Stage::Split,
        ] {
            self.stage_running(stage).set(0);
        }
        for entry in history {
            self.seed(entry);
        }
        self
    }

    /// Also shows the Runpod spend of the pods recorded in `runs`, read at each
    /// scrape.
    #[must_use]
    pub fn with_runs(mut self, runs: Runs) -> Self {
        self.runs = Some(runs);
        self
    }

    fn seed(&self, entry: &Entry) {
        let stage = entry.stage.name();
        for (result, count) in [
            ("done", entry.done),
            ("skipped", entry.skipped),
            ("failed", entry.failed),
            ("excluded", entry.excluded),
        ] {
            if count > 0 {
                self.item(stage, result, count);
            }
        }
        if let Some(model) = &entry.model {
            let usage = Usage {
                input_tokens: entry.input_tokens,
                output_tokens: entry.output_tokens,
            };
            self.spent(stage, model, usage, entry.cost);
        }
    }

    fn item(&self, stage: &'static str, result: &'static str, count: usize) {
        self.items
            .get_or_create(&ItemLabels { stage, result })
            .inc_by(u64::try_from(count).unwrap_or(u64::MAX));
    }

    fn spent(&self, stage: &'static str, model: &str, usage: Usage, cost: Option<f64>) {
        for (direction, count) in [("in", usage.input_tokens), ("out", usage.output_tokens)] {
            self.tokens
                .get_or_create(&TokenLabels {
                    stage,
                    model: model.to_string(),
                    direction,
                })
                .inc_by(count);
        }
        if let Some(cost) = cost {
            self.cost
                .get_or_create(&ModelLabels {
                    stage,
                    model: model.to_string(),
                })
                .inc_by(cost);
        }
    }

    /// Counts the items and tokens of one event, keeping what it must remember in `followed`.
    fn follow(&self, followed: &mut Followed, event: &Event) {
        if let Some((stage, result)) = event.item_result() {
            self.item(stage.name(), result.name(), 1);
        }
        match event {
            Event::StageStarted { stage, .. } => {
                self.stage_running(*stage).inc();
                followed.running.push(*stage);
            },
            Event::StageModel { stage, model } => {
                followed.models.insert(*stage, model.clone());
            },
            Event::ItemDone {
                stage,
                usage: Some(usage),
                cost,
                ..
            } => {
                let model = followed
                    .models
                    .get(stage)
                    .map_or(UNKNOWN_MODEL, String::as_str);
                self.spent(stage.name(), model, *usage, *cost);
            },
            Event::ItemFailed {
                stage,
                retryable: true,
                ..
            } => {
                self.retries
                    .get_or_create(&StageLabels {
                        stage: stage.name(),
                    })
                    .inc();
            },
            Event::StageFinished { stage, stats } => {
                // Split publishes no item events: its exclusions come with its end,
                // as its history line has them. The other stages count them per item.
                if *stage == Stage::Split && stats.excluded > 0 {
                    self.item(stage.name(), "excluded", stats.excluded);
                }
                if let Some(index) = followed.running.iter().position(|s| s == stage) {
                    followed.running.remove(index);
                    self.stage_running(*stage).dec();
                }
            },
            Event::RunWatched { run_id, .. } => {
                if let Some(series) = followed.system.take() {
                    self.system.remove(&series);
                }
                if let Some(previous) = followed.run_id.replace(run_id.clone()) {
                    self.forget_phase(&previous);
                }
                followed.phases = Phases::default();
                followed.exited = false;
            },
            Event::Metric(metric) => {
                followed.phases.metric(metric);
                if let Some(run_id) = &followed.run_id {
                    self.train(run_id, metric);
                    self.set_phase(run_id, followed.phases.phase(followed.exited));
                }
            },
            Event::Mark(mark) => {
                followed.phases.mark(*mark);
                if let Some(run_id) = &followed.run_id {
                    self.set_phase(run_id, followed.phases.phase(followed.exited));
                }
            },
            Event::JobStatus(status) => {
                followed.exited = status.is_finished();
                if let Some(run_id) = &followed.run_id {
                    self.set_phase(run_id, followed.phases.phase(followed.exited));
                }
            },
            Event::System(sample) => {
                if let Some(run_id) = &followed.run_id {
                    let set = self.system.set(run_id, sample, followed.system.as_ref());
                    followed.system = Some(set);
                }
            },
            Event::ItemDone { .. }
            | Event::ItemFailed { .. }
            | Event::PodStatus(_)
            | Event::Push { .. }
            | Event::Judged { .. } => {},
        }
    }

    fn stage_running(&self, stage: Stage) -> impl std::ops::Deref<Target = Gauge> + '_ {
        self.running.get_or_create(&StageLabels {
            stage: stage.name(),
        })
    }

    fn train(&self, run_id: &str, metric: &TrainMetric) {
        let labels = RunLabels {
            run_id: run_id.to_string(),
        };
        self.train_step
            .get_or_create(&labels)
            .set(i64::try_from(metric.step).unwrap_or(i64::MAX));
        if let Some(max_steps) = metric.max_steps {
            self.train_max_steps
                .get_or_create(&labels)
                .set(i64::try_from(max_steps).unwrap_or(i64::MAX));
        }
        for (family, value) in [
            (&self.train_loss, metric.loss),
            (&self.eval_loss, metric.eval_loss),
            (&self.learning_rate, metric.learning_rate),
        ] {
            if let Some(value) = value {
                family.get_or_create(&labels).set(value);
            }
        }
    }

    /// Shows `phase` as the current one of run `run_id`: 1 for it, 0 for every
    /// other phase.
    fn set_phase(&self, run_id: &str, phase: Phase) {
        for name in Phase::NAMES {
            self.train_phase
                .get_or_create(&PhaseLabels {
                    run_id: run_id.to_string(),
                    phase: name,
                })
                .set(i64::from(name == phase.name()));
        }
    }

    /// Removes the phase series of run `run_id`: nothing follows it any more.
    fn forget_phase(&self, run_id: &str) {
        for name in Phase::NAMES {
            self.train_phase.remove(&PhaseLabels {
                run_id: run_id.to_string(),
                phase: name,
            });
        }
    }

    /// Sets the Runpod spend of every recorded pod as of `now`.
    fn read_spend(&self, now: SystemTime) {
        let Some(runs) = &self.runs else {
            return;
        };
        let spends = pod_spends(runs, now);
        // A run removed since the last scrape leaves no series behind.
        self.spend.clear();
        for (run_id, spend) in spends {
            self.spend.get_or_create(&RunLabels { run_id }).set(spend);
        }
    }

    /// The text exposition of every family, the Runpod spend read now.
    ///
    /// # Errors
    ///
    /// Returns an error when a family cannot be written.
    pub fn encode(&self) -> Result<String, fmt::Error> {
        let _scrape = self.scrape.lock().unwrap_or_else(PoisonError::into_inner);
        self.read_spend(SystemTime::now());
        let mut text = String::new();
        prometheus_client::encoding::text::encode(&mut text, &self.registry)?;
        Ok(text)
    }
}

impl Observer for Metrics {
    fn event(&self, bus: usize, event: &Event) {
        let mut buses = self.buses.lock().unwrap_or_else(PoisonError::into_inner);
        self.follow(buses.entry(bus).or_default(), event);
    }

    fn closed(&self, bus: usize) {
        let mut buses = self.buses.lock().unwrap_or_else(PoisonError::into_inner);
        let followed = buses.remove(&bus).unwrap_or_default();
        // A stage that failed or was interrupted never finished.
        for stage in followed.running {
            self.stage_running(stage).dec();
        }
        // The machine of a run nothing watches any more is not shown as it was,
        // nor is its phase.
        if let Some(series) = &followed.system {
            self.system.remove(series);
        }
        if let Some(run_id) = &followed.run_id {
            self.forget_phase(run_id);
        }
    }
}

/// A new metric, registered in `registry` as `name`: the registry keeps a clone,
/// which shares its series.
fn register<M: Metric + Clone + Default>(registry: &mut Registry, name: &str, help: &str) -> M {
    let metric = M::default();
    registry.register(name, help, metric.clone());
    metric
}

/// The Runpod spend at `now` of each run of `runs` with a pod record. A record
/// that cannot be read is skipped.
fn pod_spends(runs: &Runs, now: SystemTime) -> Vec<(String, f64)> {
    let skipped = |id: &str, error: &crate::runs::RunsError| {
        tracing::debug!("metrics: skipping run {id}: {error}");
    };
    let records = runs.list_with(skipped).unwrap_or_else(|error| {
        tracing::debug!("metrics: cannot list the runs: {error}");
        Vec::new()
    });
    records
        .into_iter()
        .filter_map(|record| {
            let pod = PodRecord::load(runs, &record.id)
                .map_err(|error| skipped(&record.id, &error))
                .ok()??;
            Some((record.id, pod.spend_at(now)?))
        })
        .collect()
}

/// The running endpoint. Dropping it stops the server and closes its connections.
#[derive(Debug)]
pub struct MetricsServer {
    address: SocketAddr,
    server: JoinHandle<()>,
}

impl MetricsServer {
    /// The address it listens on, with the port the system chose for port 0.
    #[must_use]
    pub fn address(&self) -> SocketAddr {
        self.address
    }
}

impl Drop for MetricsServer {
    fn drop(&mut self) {
        self.server.abort();
    }
}

/// Serves `metrics` at `GET /metrics` on `address`. The metrics count what the
/// buses observed by them publish.
///
/// # Errors
///
/// Returns an error when `address` cannot be bound.
pub async fn serve(address: SocketAddr, metrics: Arc<Metrics>) -> io::Result<MetricsServer> {
    let listener = TcpListener::bind(address).await?;
    let address = listener.local_addr()?;
    let server = tokio::spawn(accept(listener, metrics));
    Ok(MetricsServer { address, server })
}

/// Serves each connection in its own task, at most [`MAX_CONNECTIONS`] at once,
/// all of them dropped with this one.
async fn accept(listener: TcpListener, metrics: Arc<Metrics>) {
    let mut connections = JoinSet::new();
    loop {
        tokio::select! {
            accepted = listener.accept() => {
                if let Err(error) = admit(&mut connections, accepted, &metrics) {
                    tracing::debug!("cannot accept a metrics connection: {error}");
                    tokio::time::sleep(ACCEPT_PAUSE).await;
                }
            },
            Some(_) = connections.join_next(), if !connections.is_empty() => {},
        }
    }
}

/// Serves the `accepted` connection in `connections`, or closes it when
/// [`MAX_CONNECTIONS`] are served already.
fn admit(
    connections: &mut JoinSet<()>,
    accepted: io::Result<(TcpStream, SocketAddr)>,
    metrics: &Arc<Metrics>,
) -> io::Result<()> {
    let (stream, _) = accepted?;
    if connections.len() >= MAX_CONNECTIONS {
        tracing::debug!("closing a metrics connection: {MAX_CONNECTIONS} are open");
    } else {
        connections.spawn(connection(stream, Arc::clone(metrics)));
    }
    Ok(())
}

async fn connection(stream: TcpStream, metrics: Arc<Metrics>) {
    let service = service_fn(move |request| {
        let response = respond(&request, &metrics);
        async move { Ok::<_, Infallible>(response) }
    });
    if let Err(error) = http1::Builder::new()
        .timer(TokioTimer::new())
        .header_read_timeout(HEADER_TIMEOUT)
        .serve_connection(TokioIo::new(stream), service)
        .await
    {
        tracing::debug!("metrics connection: {error}");
    }
}

fn respond(request: &Request<Incoming>, metrics: &Metrics) -> Response<String> {
    if request.method() != Method::GET || request.uri().path() != "/metrics" {
        return status(StatusCode::NOT_FOUND, "not found\n");
    }
    match metrics.encode() {
        Ok(text) => {
            let mut response = Response::new(text);
            response
                .headers_mut()
                .insert(CONTENT_TYPE, HeaderValue::from_static(CONTENT_TYPE_VALUE));
            response
        },
        Err(error) => {
            tracing::debug!("cannot encode the metrics: {error}");
            status(
                StatusCode::INTERNAL_SERVER_ERROR,
                "cannot encode the metrics\n",
            )
        },
    }
}

fn status(code: StatusCode, text: &str) -> Response<String> {
    let mut response = Response::new(text.to_string());
    *response.status_mut() = code;
    response
}

#[cfg(test)]
mod tests {
    use std::time::UNIX_EPOCH;

    use super::*;
    use crate::events::{EventBus, StageStats};
    use crate::history::Status;
    use crate::runpod::PodState;

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// The value of the sample `series` (name and labels, as written) in `text`.
    fn value(text: &str, series: &str) -> Option<f64> {
        text.lines()
            .find_map(|line| line.strip_prefix(series)?.strip_prefix(' '))
            .and_then(|value| value.parse().ok())
    }

    fn assert_value(text: &str, series: &str, expected: f64) {
        let found = value(text, series);
        assert!(
            found.is_some_and(|found| (found - expected).abs() < 1e-9),
            "{series}: expected {expected}, got {found:?} in\n{text}"
        );
    }

    fn entry(stage: Stage, model: Option<&str>, cost: Option<f64>) -> Entry {
        Entry {
            stage,
            started_at: "2026-09-28T10:00:00Z".into(),
            ended_at: "2026-09-28T10:01:00Z".into(),
            status: Status::Ok,
            provider: model.map(|_| "mock".into()),
            model: model.map(str::to_string),
            done: 3,
            skipped: 1,
            failed: 1,
            excluded: 2,
            input_tokens: 100,
            output_tokens: 40,
            cost,
            split: None,
            backfilled: false,
        }
    }

    fn metric(step: u64, loss: Option<f64>, eval_loss: Option<f64>) -> TrainMetric {
        TrainMetric {
            time: 0.0,
            step,
            epoch: None,
            max_steps: None,
            loss,
            eval_loss,
            learning_rate: Some(2e-5),
            grad_norm: None,
        }
    }

    #[test]
    fn the_history_seeds_items_tokens_and_cost() -> TestResult {
        let metrics = Metrics::new(&[
            entry(Stage::Answers, Some("parent"), Some(0.5)),
            entry(Stage::Answers, Some("parent"), None),
            entry(Stage::Questions, Some("gen"), Some(0.25)),
            entry(Stage::Split, None, None),
        ]);
        let text = metrics.encode()?;
        let answers = r#"stage="answers",model="parent""#;
        assert_value(
            &text,
            &format!("overbrainer_tokens_total{{{answers},direction=\"in\"}}"),
            200.0,
        );
        assert_value(
            &text,
            &format!("overbrainer_tokens_total{{{answers},direction=\"out\"}}"),
            80.0,
        );
        // An unknown cost adds nothing.
        assert_value(
            &text,
            &format!("overbrainer_cost_usd_total{{{answers}}}"),
            0.5,
        );
        assert_value(
            &text,
            r#"overbrainer_cost_usd_total{stage="questions",model="gen"}"#,
            0.25,
        );
        for (result, count) in [
            ("done", 6.0),
            ("skipped", 2.0),
            ("failed", 2.0),
            ("excluded", 4.0),
        ] {
            assert_value(
                &text,
                &format!("overbrainer_stage_items_total{{stage=\"answers\",result=\"{result}\"}}"),
                count,
            );
        }
        // Split has no model: items only.
        assert_value(
            &text,
            r#"overbrainer_stage_items_total{stage="split",result="excluded"}"#,
            2.0,
        );
        assert!(
            !text.contains(r#"overbrainer_tokens_total{stage="split""#),
            "{text}"
        );
        assert_value(
            &text,
            &format!(
                "overbrainer_build_info{{version=\"{}\"}}",
                env!("CARGO_PKG_VERSION")
            ),
            1.0,
        );
        Ok(())
    }

    #[test]
    fn events_move_every_counter_and_gauge() -> TestResult {
        let metrics = Metrics::new(&[]);
        let usage = Usage {
            input_tokens: 10,
            output_tokens: 4,
        };
        for event in [
            Event::StageStarted {
                stage: Stage::Answers,
                total: 3,
            },
            Event::StageModel {
                stage: Stage::Answers,
                model: "parent".into(),
            },
            Event::ItemDone {
                stage: Stage::Answers,
                id: "a".into(),
                usage: Some(usage),
                cost: Some(0.125),
                excluded: false,
            },
            Event::ItemDone {
                stage: Stage::Answers,
                id: "b".into(),
                usage: Some(usage),
                cost: None,
                excluded: false,
            },
            Event::ItemDone {
                stage: Stage::Answers,
                id: "c".into(),
                usage: None,
                cost: None,
                excluded: false,
            },
            Event::ItemFailed {
                stage: Stage::Answers,
                id: "d".into(),
                error: "busy".into(),
                retryable: true,
            },
            Event::ItemFailed {
                stage: Stage::Answers,
                id: "d".into(),
                error: "busy".into(),
                retryable: false,
            },
        ] {
            metrics.event(1, &event);
        }
        let text = metrics.encode()?;
        assert_value(&text, r#"overbrainer_stage_running{stage="answers"}"#, 1.0);
        for (result, count) in [("done", 2.0), ("skipped", 1.0), ("failed", 1.0)] {
            assert_value(
                &text,
                &format!("overbrainer_stage_items_total{{stage=\"answers\",result=\"{result}\"}}"),
                count,
            );
        }
        assert_value(
            &text,
            r#"overbrainer_tokens_total{stage="answers",model="parent",direction="in"}"#,
            20.0,
        );
        assert_value(
            &text,
            r#"overbrainer_cost_usd_total{stage="answers",model="parent"}"#,
            0.125,
        );
        assert_value(
            &text,
            r#"overbrainer_item_retries_total{stage="answers"}"#,
            1.0,
        );
        metrics.event(
            1,
            &Event::StageFinished {
                stage: Stage::Answers,
                stats: StageStats {
                    excluded: 2,
                    ..StageStats::default()
                },
            },
        );
        let text = metrics.encode()?;
        assert_value(&text, r#"overbrainer_stage_running{stage="answers"}"#, 0.0);
        // The end of a stage adds nothing: its items were counted one by one.
        assert_eq!(
            value(
                &text,
                r#"overbrainer_stage_items_total{stage="answers",result="excluded"}"#
            ),
            None,
            "{text}"
        );
        Ok(())
    }

    #[test]
    fn live_counts_mean_what_the_history_seeds() -> TestResult {
        let usage = Usage {
            input_tokens: 100,
            output_tokens: 40,
        };
        let done = |excluded| Event::ItemDone {
            stage: Stage::Answers,
            id: "a".into(),
            usage: Some(usage),
            cost: None,
            excluded,
        };
        let live = Metrics::new(&[]);
        for event in [
            done(false),
            done(false),
            done(false),
            done(true),
            done(true),
        ] {
            live.event(1, &event);
        }
        live.event(
            1,
            &Event::StageFinished {
                stage: Stage::Answers,
                stats: StageStats {
                    done: 3,
                    excluded: 2,
                    ..StageStats::default()
                },
            },
        );
        // Split has no item events: its exclusions come with its end.
        live.event(
            1,
            &Event::StageFinished {
                stage: Stage::Split,
                stats: StageStats {
                    excluded: 4,
                    ..StageStats::default()
                },
            },
        );
        let live = live.encode()?;
        for (stage, result, count) in [
            ("answers", "done", 3.0),
            ("answers", "excluded", 2.0),
            ("split", "excluded", 4.0),
        ] {
            assert_value(
                &live,
                &format!("overbrainer_stage_items_total{{stage=\"{stage}\",result=\"{result}\"}}"),
                count,
            );
        }
        // The same stage recorded in the history: done without the excluded.
        let mut recorded = entry(Stage::Answers, Some("parent"), None);
        (
            recorded.done,
            recorded.skipped,
            recorded.failed,
            recorded.excluded,
        ) = (3, 0, 0, 2);
        let seeded = Metrics::new(&[recorded]).encode()?;
        for result in ["done", "excluded"] {
            let series =
                format!("overbrainer_stage_items_total{{stage=\"answers\",result=\"{result}\"}}");
            assert_eq!(value(&live, &series), value(&seeded, &series), "{series}");
        }
        Ok(())
    }

    /// Training metrics are labelled with the run their bus watches, named by
    /// the last [`Event::RunWatched`].
    #[test]
    fn training_metrics_carry_the_run_their_bus_watches() -> TestResult {
        let metrics = Metrics::new(&[]);
        // Before its bus names the run, a metric has no run to go to.
        metrics.event(2, &Event::Metric(metric(1, Some(9.0), None)));
        metrics.event(
            2,
            &Event::RunWatched {
                run_id: "20260928-100000-a1b2".into(),
                job: None,
            },
        );
        metrics.event(2, &Event::Metric(metric(5, Some(1.5), None)));
        metrics.event(2, &Event::Metric(metric(6, None, Some(1.25))));
        // Another bus watches no run.
        metrics.event(3, &Event::Metric(metric(7, Some(9.0), None)));
        let text = metrics.encode()?;
        let run = r#"{run_id="20260928-100000-a1b2"}"#;
        assert_value(&text, &format!("overbrainer_train_step{run}"), 6.0);
        assert_value(&text, &format!("overbrainer_train_loss{run}"), 1.5);
        assert_value(&text, &format!("overbrainer_eval_loss{run}"), 1.25);
        assert_value(&text, &format!("overbrainer_learning_rate{run}"), 2e-5);
        assert_eq!(text.matches("overbrainer_train_step{").count(), 1, "{text}");
        Ok(())
    }

    /// The phase of a followed run shows as a gauge until its bus closes.
    #[test]
    fn the_phase_of_a_followed_run_shows_until_its_bus_closes() -> TestResult {
        use crate::exec::JobStatus;
        use crate::train::{JobStage, Mark};

        let metrics = Metrics::new(&[]);
        let run = "20260928-100000-a1b2";
        let phase = |text: &str, name: &str| {
            value(
                text,
                &format!("overbrainer_train_phase{{run_id=\"{run}\",phase=\"{name}\"}}"),
            )
        };
        metrics.event(
            4,
            &Event::RunWatched {
                run_id: run.into(),
                job: None,
            },
        );
        metrics.event(
            4,
            &Event::Metric(TrainMetric {
                max_steps: Some(10),
                ..metric(4, Some(1.5), None)
            }),
        );
        let text = metrics.encode()?;
        assert_value(
            &text,
            &format!("overbrainer_train_max_steps{{run_id=\"{run}\"}}"),
            10.0,
        );
        assert_eq!(phase(&text, "training"), Some(1.0));
        assert_eq!(phase(&text, "finalizing"), Some(0.0));
        for (event, current) in [
            (
                Event::Mark(Mark::Eval {
                    step: 3,
                    total: None,
                }),
                "evaluating",
            ),
            (Event::Mark(Mark::End { step: 10 }), "evaluating"),
            (
                Event::Metric(TrainMetric {
                    max_steps: Some(10),
                    ..metric(10, None, Some(1.25))
                }),
                "finalizing",
            ),
            (Event::Mark(Mark::Stage(JobStage::Merge)), "merging"),
            (Event::JobStatus(JobStatus::Exited(0)), "retrieving"),
        ] {
            metrics.event(4, &event);
            let text = metrics.encode()?;
            for name in Phase::NAMES {
                let expected = if name == current { 1.0 } else { 0.0 };
                assert_eq!(
                    phase(&text, name),
                    Some(expected),
                    "{event:?} {name}\n{text}"
                );
            }
        }
        metrics.closed(4);
        let text = metrics.encode()?;
        assert!(!text.contains("overbrainer_train_phase{"), "{text}");
        assert_value(
            &text,
            &format!("overbrainer_train_max_steps{{run_id=\"{run}\"}}"),
            10.0,
        );
        Ok(())
    }

    /// A sample of a pod with `gpus` GPUs, taken at 1000 s past the epoch.
    fn sample(gpus: u32) -> crate::system::SystemSample {
        let mut output = String::from("@gpu\n");
        for index in 0..gpus {
            output.push_str(&index.to_string());
            output.push_str(", NVIDIA H100, 50, 1024, 2048, 60, 300, 700\n");
        }
        output.push_str(
            "@loadavg\n2.5 1 1 1/1 1\n@nproc\n16\n@meminfo\nMemTotal: 1000 kB\n\
             MemAvailable: 250 kB\n@df.run\nh\nd 100 60 40 60% /workspace\n\
             @df.root\nh\nnas:/export 10 5 5 50% /\n",
        );
        crate::system::parse(&output, UNIX_EPOCH + Duration::from_secs(1000), None)
    }

    /// The target machine's figures show in base units until the bus closes.
    #[test]
    fn the_target_machine_shows_in_base_units_until_the_bus_closes() -> TestResult {
        let metrics = Metrics::new(&[]);
        // A sample before any run: nothing to label it with.
        metrics.event(1, &Event::System(sample(1)));
        assert!(!metrics.encode()?.contains("overbrainer_target_"));
        let run_id = "20260928-100000-a1b2";
        metrics.event(
            1,
            &Event::RunWatched {
                run_id: run_id.into(),
                job: None,
            },
        );
        metrics.event(1, &Event::Metric(metric(5, Some(1.5), None)));
        metrics.event(1, &Event::System(sample(2)));
        let text = metrics.encode()?;
        let run = format!("run_id=\"{run_id}\"");
        for (series, expected) in [
            (
                format!("overbrainer_target_disk_used_bytes{{{run},mount=\"/workspace\"}}"),
                61_440.0,
            ),
            (
                format!("overbrainer_target_disk_size_bytes{{{run},mount=\"/workspace\"}}"),
                102_400.0,
            ),
            (format!("overbrainer_target_cpu_load1{{{run}}}"), 2.5),
            (format!("overbrainer_target_cpus{{{run}}}"), 16.0),
            (
                format!("overbrainer_target_memory_used_bytes{{{run}}}"),
                768_000.0,
            ),
            (
                format!("overbrainer_target_memory_limit_bytes{{{run}}}"),
                1_024_000.0,
            ),
            (
                format!("overbrainer_target_sample_timestamp_seconds{{{run}}}"),
                1000.0,
            ),
            (
                format!("overbrainer_gpu_utilization_ratio{{{run},gpu=\"1\"}}"),
                0.5,
            ),
            (
                format!("overbrainer_gpu_memory_used_bytes{{{run},gpu=\"0\"}}"),
                1_073_741_824.0,
            ),
            (
                format!("overbrainer_gpu_memory_total_bytes{{{run},gpu=\"0\"}}"),
                2_147_483_648.0,
            ),
            (
                format!("overbrainer_gpu_temperature_celsius{{{run},gpu=\"0\"}}"),
                60.0,
            ),
            (
                format!("overbrainer_gpu_power_watts{{{run},gpu=\"0\"}}"),
                300.0,
            ),
            (
                format!("overbrainer_gpu_power_limit_watts{{{run},gpu=\"0\"}}"),
                700.0,
            ),
            (
                format!("overbrainer_gpu_info{{{run},gpu=\"1\",name=\"NVIDIA H100\"}}"),
                1.0,
            ),
        ] {
            assert_value(&text, &series, expected);
        }
        // A network file system shows the whole cluster: left out.
        assert!(!text.contains("mount=\"/\""), "{text}");
        // No previous sample: no usage yet.
        assert!(
            !text.contains("overbrainer_target_cpu_usage_ratio{"),
            "{text}"
        );
        // A GPU gone at the next sample takes its series with it, and so does
        // a figure the target no longer reports.
        let mut later = sample(1);
        if let Some(cpu) = later.cpu.as_mut() {
            cpu.load1 = None;
        }
        metrics.event(1, &Event::System(later));
        let text = metrics.encode()?;
        assert!(!text.contains("overbrainer_target_cpu_load1{"), "{text}");
        assert_value(&text, &format!("overbrainer_target_cpus{{{run}}}"), 16.0);
        assert!(!text.contains("gpu=\"1\""), "{text}");
        assert!(text.contains("gpu=\"0\""), "{text}");
        metrics.closed(1);
        let text = metrics.encode()?;
        assert!(
            !text.contains("overbrainer_target_disk_used_bytes{"),
            "{text}"
        );
        assert!(!text.contains("overbrainer_gpu_"), "{text}");
        assert!(
            !text.contains("overbrainer_target_sample_timestamp_seconds{"),
            "{text}"
        );
        // The training series stay: they say how the run ended.
        assert_value(&text, &format!("overbrainer_train_step{{{run}}}"), 5.0);
        Ok(())
    }

    #[test]
    fn a_burst_far_larger_than_any_channel_is_fully_counted() -> TestResult {
        const BURST: usize = 30_000;
        let metrics = Arc::new(Metrics::new(&[]));
        let bus = EventBus::observed(8, Some(Arc::clone(&metrics) as Arc<dyn Observer>));
        let _behind = bus.subscribe();
        for index in 0..BURST {
            bus.publish(Event::ItemDone {
                stage: Stage::Answers,
                id: index.to_string(),
                usage: None,
                cost: None,
                excluded: false,
            });
        }
        let text = metrics.encode()?;
        assert_value(
            &text,
            r#"overbrainer_stage_items_total{stage="answers",result="skipped"}"#,
            30_000.0,
        );
        Ok(())
    }

    #[test]
    fn a_bus_that_ends_mid_stage_stops_it_running() -> TestResult {
        let metrics = Metrics::new(&[]);
        let started = Event::StageStarted {
            stage: Stage::Questions,
            total: 1,
        };
        metrics.event(1, &started);
        metrics.event(2, &started);
        metrics.closed(1);
        assert_value(
            &metrics.encode()?,
            r#"overbrainer_stage_running{stage="questions"}"#,
            1.0,
        );
        metrics.closed(2);
        metrics.closed(2);
        assert_value(
            &metrics.encode()?,
            r#"overbrainer_stage_running{stage="questions"}"#,
            0.0,
        );
        Ok(())
    }

    #[test]
    fn the_runpod_spend_comes_from_the_pod_records() -> TestResult {
        let dir = tempfile::tempdir()?;
        let runs = Runs::new(dir.path());
        let gone = crate::runs::create(&runs, "demo", "/workspace", "cloud")?;
        let mut pod = PodRecord::new(&gone.id, false, 1, "ssh-ed25519 AAAA");
        pod.state = PodState::Deleted;
        pod.estimated_spend = Some(0.64);
        pod.save(&runs)?;
        let live = crate::runs::create(&runs, "demo", "/workspace", "cloud")?;
        let mut pod = PodRecord::new(&live.id, false, 1, "ssh-ed25519 AAAA");
        pod.state = PodState::Running;
        pod.cost_per_hour = Some(0.5);
        pod.created_unix = Some(1000);
        pod.save(&runs)?;
        let local = crate::runs::create(&runs, "demo", "/work", "local")?;
        let metrics = Metrics::new(&[]).with_runs(runs.clone());
        let now = UNIX_EPOCH + Duration::from_secs(1000 + 2 * 3600);
        metrics.read_spend(now);
        let mut text = String::new();
        prometheus_client::encoding::text::encode(&mut text, &metrics.registry)?;
        let series = |id: &str| format!("overbrainer_runpod_spend_usd{{run_id=\"{id}\"}}");
        assert_value(&text, &series(&gone.id), 0.64);
        assert_value(&text, &series(&live.id), 1.0);
        assert_eq!(value(&text, &series(&local.id)), None);
        // A run removed meanwhile leaves no series behind.
        std::fs::remove_dir_all(runs.run_dir(&gone.id)?)?;
        metrics.read_spend(now);
        let mut text = String::new();
        prometheus_client::encoding::text::encode(&mut text, &metrics.registry)?;
        assert_eq!(value(&text, &series(&gone.id)), None, "{text}");
        assert_value(&text, &series(&live.id), 1.0);
        Ok(())
    }

    #[tokio::test]
    async fn the_endpoint_serves_the_families_and_nothing_else() -> TestResult {
        let history = [entry(Stage::Answers, Some("parent"), Some(0.5))];
        let metrics = Arc::new(Metrics::new(&history));
        let server = serve("127.0.0.1:0".parse()?, Arc::clone(&metrics)).await?;
        let base = format!("http://{}", server.address());
        let client = reqwest::Client::new();
        let response = client.get(format!("{base}/metrics")).send().await?;
        assert_eq!(response.status(), reqwest::StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get(reqwest::header::CONTENT_TYPE)
                .map(HeaderValue::as_bytes),
            Some(CONTENT_TYPE_VALUE.as_bytes())
        );
        let text = response.text().await?;
        // A family shows once it has a series: these have one from the start.
        for family in [
            "overbrainer_stage_items",
            "overbrainer_tokens",
            "overbrainer_cost_usd",
            "overbrainer_stage_running",
            "overbrainer_build_info",
        ] {
            assert!(
                text.contains(&format!("# TYPE {family} ")),
                "{family}: {text}"
            );
        }
        assert!(text.ends_with("# EOF\n"), "{text}");
        for response in [
            client.get(format!("{base}/")).send().await?,
            client.get(format!("{base}/metrics/x")).send().await?,
            client.post(format!("{base}/metrics")).send().await?,
        ] {
            assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND);
        }
        // What an observed bus publishes shows at the next scrape.
        let running = r#"overbrainer_stage_running{stage="split"}"#;
        assert_value(&text, running, 0.0);
        let bus = EventBus::observed(1, Some(metrics));
        bus.publish(Event::StageStarted {
            stage: Stage::Split,
            total: 0,
        });
        let text = client
            .get(format!("{base}/metrics"))
            .send()
            .await?
            .text()
            .await?;
        assert_value(&text, running, 1.0);
        Ok(())
    }

    #[tokio::test]
    async fn connections_beyond_the_cap_are_closed() -> TestResult {
        use tokio::io::AsyncReadExt as _;

        let server = serve("127.0.0.1:0".parse()?, Arc::new(Metrics::new(&[]))).await?;
        let mut idle = Vec::new();
        for _ in 0..MAX_CONNECTIONS {
            idle.push(tokio::net::TcpStream::connect(server.address()).await?);
        }
        let mut extra = tokio::net::TcpStream::connect(server.address()).await?;
        let mut byte = [0; 1];
        let read = tokio::time::timeout(Duration::from_secs(5), extra.read(&mut byte)).await;
        assert!(matches!(read, Ok(Ok(0))), "{read:?}");
        drop(idle);
        Ok(())
    }

    #[tokio::test]
    async fn dropping_the_server_frees_its_port() -> TestResult {
        let server = serve("127.0.0.1:0".parse()?, Arc::new(Metrics::new(&[]))).await?;
        let address = server.address();
        assert!(
            TcpListener::bind(address).await.is_err(),
            "the port is free"
        );
        drop(server);
        let mut bound = Err(io::Error::other("never tried"));
        for _ in 0..100 {
            bound = TcpListener::bind(address).await;
            if bound.is_ok() {
                break;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
        bound?;
        Ok(())
    }
}