launcher/src/bench_artifacts.rs

//! Versioned run package for `sm121 bench`.
//!
//! Every bench run lands next to its reports in the recipe's `reports/`
//! dir (same stem as the report triplet):
//!
//! - `<stem>.recipe.toml` — a byte-exact copy of the recipe file the
//!   serve was launched from, so the run's recipe is pinned even after
//!   the recipe (or its image pins) move on.
//! - `<stem>.manifest.json` — provenance: recipe + config SHA-256, the
//!   launcher's git commit, the bench tool's version, the served model's
//!   resolved `refs/main` snapshot (the model reference floats, so the
//!   snapshot pins the run), and the served image's identity per node
//!   (docker id + `sm121.checksum` build label, captured over SSH at bench
//!   time).
//!
//! `sm121 report` reads the manifests back (see [`crate::report`]) and
//! shows each run's provenance — and flags a run whose recipe has since
//! changed (drift).

use std::fs;
use std::path::Path;

use anyhow::Result;
use sha2::{Digest, Sha256};

use crate::config::{Cluster, Config, Recipe};
use crate::ssh::run_ssh;

/// SHA-256 of a byte slice as a 64-char lowercase hex string.
pub fn sha256_hex(bytes: &[u8]) -> String {
    let digest = Sha256::digest(bytes);
    digest.iter().map(|b| format!("{b:02x}")).collect()
}

/// First 12 chars — the display width for hashes in the report UI.
pub fn short12(s: &str) -> String {
    s.chars().take(12).collect()
}

/// UTC "YYYY-MM-DDTHH:MM:SSZ" for the manifest's capture time.
fn utc_now_iso() -> String {
    let secs = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    let (y, m, d) = civil_date((secs / 86400) as i64);
    let rem = (secs % 86400) as u32;
    format!(
        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
        y,
        m,
        d,
        rem / 3600,
        rem % 3600 / 60,
        rem % 60
    )
}

/// Days-since-civil conversion (Hansen/Clayton), same as `bench::civil_date`.
fn civil_date(z: i64) -> (u32, u32, u32) {
    let z = z + 719_468;
    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
    let doe = z - era * 146_097;
    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let day = doy - (153 * mp + 2) / 5 + 1;
    let month = if mp < 10 { mp + 3 } else { mp - 9 };
    let year = era * 400 + yoe + if month <= 2 { 1 } else { 0 };
    (year as u32, month as u32, day as u32)
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestRecipe {
    pub name: String,
    /// Path of the recipe file as resolved by the launcher.
    pub file: String,
    pub sha256: String,
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestCluster {
    pub name: String,
    /// ssh_ip per node (node order from the cluster definition).
    pub nodes: Vec<String>,
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestConfig {
    pub file: String,
    pub sha256: String,
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestLauncher {
    /// Canonical path of the `sm121` binary that ran the bench.
    pub bin: String,
    /// `git rev-parse HEAD` of the repo the binary was built in (cwd),
    /// null when the cwd is not a git checkout.
    pub git_sha: Option<String>,
}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestBench {
    pub tool: String,
    /// `tool_eval_bench_version` from the run's JSON report.
    pub version: Option<String>,
    /// The exact `tool-eval-bench` arguments (base-url, model, seed, …).
    pub args: Vec<String>,
    pub final_score: Option<u32>,
    /// `total/max` points, e.g. "163/176".
    pub points: Option<String>,
}

/// Provenance of the served model. The recipe's `model` reference floats
/// (the engine follows the HF cache's `refs/main` → latest snapshot), so the
/// snapshot it actually resolved to at bench time is what makes a run
/// reproducible. `snapshot` is null when `model` is a local path (no
/// `org/name` id to resolve) or the repo isn't cached on the head node.
#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct ManifestModel {
    /// The recipe's `model` field (an `org/name` repo id or a local path).
    pub id: String,
    /// The `refs/main` snapshot hash the id resolved to on the head node.
    pub snapshot: Option<String>,
}

/// One per bench run; the top-level shape of `<stem>.manifest.json`.
#[derive(serde::Serialize, serde::Deserialize, Default, Debug, Clone)]
pub struct Manifest {
    pub schema_version: u32,
    /// `<YYYY-MM-DD>_<label>` run stem (report triplet base name).
    pub run: String,
    /// UTC ISO-8601 capture time.
    pub created_utc: String,
    pub recipe: ManifestRecipe,
    pub cluster: ManifestCluster,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub config: Option<ManifestConfig>,
    pub launcher: ManifestLauncher,
    pub bench: ManifestBench,
    /// Docker image name → `{ "build_checksum": "…", "nodes": { name →
    /// { "id", "sm121_checksum" } } }`. Nodes without the image are
    /// absent from `nodes`.
    pub images: serde_json::Map<String, serde_json::Value>,
    /// The recipe `model` id and the `refs/main` snapshot it resolved to
    /// at bench time (the model reference floats; the snapshot pins the run).
    /// Absent in pre-model manifests (defaults to empty).
    #[serde(default)]
    pub models: Vec<ManifestModel>,
}

/// `docker image inspect` the served image on every cluster node and
/// collect its id + `sm121.checksum` label. Nodes where the image is
/// missing (not in the launch set) are skipped, not errors.
pub async fn capture_images(config: &Config, cluster: &Cluster, image: &str) -> serde_json::Value {
    let mut nodes = serde_json::Map::new();
    let mut build_checksum = String::new();
    for n in &cluster.nodes {
        let cmd = format!(
            "docker image inspect --format '{{{{.Id}}}}|{{{{index .Config.Labels \"sm121.checksum\"}}}}' {image}"
        );
        let Ok(out) = run_ssh(config, n, cmd, false).await else {
            continue;
        };
        let out = out.trim();
        let (id, sum) = out
            .split_once('|')
            .map(|(a, b)| (a.to_string(), b.to_string()))
            .unwrap_or_else(|| (out.to_string(), String::new()));
        if build_checksum.is_empty() && !sum.is_empty() {
            build_checksum = sum.clone();
        }
        nodes.insert(
            n.name.clone(),
            serde_json::json!({ "id": id, "sm121_checksum": sum }),
        );
    }
    serde_json::json!({ "build_checksum": build_checksum, "nodes": nodes })
}

/// Resolve the recipe `model` id to its `refs/main` snapshot on the head
/// node (the shared-NAS HF cache). Returns the snapshot hash, or `None`
/// when the model is a local path (no `org/name` id to resolve) or the repo
/// isn't cached there. Only `org/name` repo ids carry a `refs/main`.
async fn capture_model_snapshot(
    config: &Config,
    node: &crate::config::Node,
    model: &str,
) -> Option<String> {
    if model.starts_with('/') || model.starts_with('~') || !model.contains('/') {
        return None;
    }
    let model_dir = crate::config::model_cache_dir_name(model);
    let hf_cache = config.resolved_hf_cache_dir();
    let refs_path = format!("{}/hub/{}/refs/main", hf_cache.display(), model_dir);
    let cmd = format!("cat {} 2>/dev/null", crate::docker::shq(&refs_path));
    let out = run_ssh(config, node, cmd, false).await.ok()?;
    let snap = out.trim();
    if snap.is_empty() {
        None
    } else {
        Some(snap.to_string())
    }
}

/// Inputs for [`write_run_package`] (bundled: the signature would
/// otherwise trip `too_many_arguments`).
pub struct RunPackage<'a> {
    pub config: &'a Config,
    pub cluster: &'a Cluster,
    pub recipe: &'a Recipe,
    pub recipe_file: &'a Path,
    /// `<YYYY-MM-DD>_<label>` run stem (report triplet base name).
    pub stem: &'a str,
    pub reports_dir: &'a Path,
    pub bench_args: &'a [String],
    pub bench_json: &'a serde_json::Value,
}

/// Write the run package (`<stem>.recipe.toml` + `<stem>.manifest.json`)
/// into `reports_dir`. The recipe copy is the bytes of `recipe_file` as
/// read at bench time; image identity is captured live per node.
pub async fn write_run_package(p: &RunPackage<'_>) -> Result<()> {
    let RunPackage {
        config,
        cluster,
        recipe,
        recipe_file,
        stem,
        reports_dir,
        bench_args,
        bench_json,
    } = p;
    let bytes = fs::read(recipe_file)?;
    let recipe_sha = sha256_hex(&bytes);
    fs::write(reports_dir.join(format!("{stem}.recipe.toml")), &bytes)?;

    let cfg_path = crate::gconfig::resolve_config_path("config.toml");
    let manifest_config = fs::read_to_string(&cfg_path).ok().map(|s| ManifestConfig {
        file: cfg_path.display().to_string(),
        sha256: sha256_hex(s.as_bytes()),
    });

    let git_sha = std::process::Command::new("git")
        .args(["rev-parse", "HEAD"])
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string());

    let scores = &bench_json["scores"];
    let points = scores["total_points"]
        .as_u64()
        .and_then(|t| scores["max_points"].as_u64().map(|m| format!("{t}/{m}")));

    let mut images = serde_json::Map::new();
    if !recipe.docker_image.is_empty() {
        images.insert(
            recipe.docker_image.clone(),
            capture_images(config, cluster, &recipe.docker_image).await,
        );
    }

    // The model reference floats (engine follows refs/main); capture the
    // snapshot it resolved to on the head node so the run is reproducible.
    let mut models: Vec<ManifestModel> = Vec::new();
    if !recipe.model.is_empty()
        && let Some(head) = cluster.nodes.first()
    {
        models.push(ManifestModel {
            id: recipe.model.clone(),
            snapshot: capture_model_snapshot(config, head, &recipe.model).await,
        });
    }

    let m = Manifest {
        schema_version: 1,
        run: stem.to_string(),
        created_utc: utc_now_iso(),
        recipe: ManifestRecipe {
            name: recipe.name.clone(),
            file: recipe_file.display().to_string(),
            sha256: recipe_sha.clone(),
        },
        cluster: ManifestCluster {
            name: cluster.name.clone(),
            nodes: cluster.nodes.iter().map(|n| n.ssh_ip.clone()).collect(),
        },
        config: manifest_config,
        launcher: ManifestLauncher {
            bin: std::env::current_exe()
                .map(|p| p.display().to_string())
                .unwrap_or_default(),
            git_sha,
        },
        bench: ManifestBench {
            tool: "tool-eval-bench".into(),
            version: bench_json
                .get("tool_eval_bench_version")
                .and_then(|v| v.as_str())
                .map(String::from),
            args: bench_args.to_vec(),
            final_score: bench_json["final_score"].as_u64().map(|v| v as u32),
            points,
        },
        images,
        models,
    };
    let manifest_path = reports_dir.join(format!("{stem}.manifest.json"));
    fs::write(
        &manifest_path,
        format!("{}\n", serde_json::to_string_pretty(&m)?),
    )?;
    eprintln!(
        "bench: run package: {} + {}",
        reports_dir.join(format!("{stem}.recipe.toml")).display(),
        manifest_path.display()
    );
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sha256_hex_known_value() {
        assert_eq!(
            sha256_hex(b"abc"),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
    }

    #[test]
    fn short12_truncates() {
        assert_eq!(short12("deadbeefcafe"), "deadbeefcafe");
        assert_eq!(short12("deadbeefcafe42"), "deadbeefcafe");
    }

    #[test]
    fn utc_now_iso_is_realistic() {
        let s = utc_now_iso();
        // YYYY-MM-DDTHH:MM:SSZ, and the year must be a sane calendar year
        // (regression: civil_date got raw SECONDS and produced year 4900778)
        let year: u32 = s[..4].parse().unwrap();
        assert!(
            (2020..2030).contains(&year),
            "year {year} out of range: {s}"
        );
        assert_eq!(s.len(), 20);
        assert!(s.ends_with('Z'));
    }

    #[test]
    fn old_manifest_without_models_deserializes() {
        // Pre-model manifests lack the `models` field; `#[serde(default)]`
        // keeps `sm121 report` able to read them.
        let json = r#"{"schema_version":1,"run":"r","created_utc":"2026-01-01T00:00:00Z",
            "recipe":{"name":"X","file":"f","sha256":"s"},
            "cluster":{"name":"c","nodes":[]},
            "launcher":{"bin":"b","git_sha":null},
            "bench":{"tool":"tool-eval-bench","version":"1","args":[],"final_score":null,"points":null},
            "images":{}}"#;
        let m: Manifest = serde_json::from_str(json).unwrap();
        assert!(m.models.is_empty());
    }

    #[test]
    fn run_package_writes_copy_and_manifest() {
        let dir = std::env::temp_dir().join(format!("sm121-artifacts-{}-x", std::process::id()));
        fs::create_dir_all(dir.join("reports")).unwrap();
        let recipe_file = dir.join("recipe.toml");
        fs::write(
            &recipe_file,
            "engine = \"vllm\"\nname = \"X\"\nmodel = \"m\"\n",
        )
        .unwrap();

        let config: Config = toml::from_str("ssh_user = \"u\"\n").unwrap();
        let cluster = Cluster {
            name: "c".into(),
            nodes: vec![], // no nodes → no SSH, no image capture
            master_ib_ip: None,
            extra_entities: vec![],
        };
        let recipe: Recipe = toml::from_str(
            "engine = \"vllm\"\nname = \"X\"\nmodel = \"m\"\ndocker_image = \"img:tag\"\n",
        )
        .unwrap();
        let bench_json: serde_json::Value =
            serde_json::from_str(r#"{"tool_eval_bench_version": "1.2.3", "final_score": 93, "scores": {"total_points": 163, "max_points": 176}}"#)
                .unwrap();

        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(write_run_package(&RunPackage {
            config: &config,
            cluster: &cluster,
            recipe: &recipe,
            recipe_file: &recipe_file,
            stem: "2026-01-01_test",
            reports_dir: &dir.join("reports"),
            bench_args: &["--base-url".into(), "x".into()],
            bench_json: &bench_json,
        }))
        .unwrap();

        let copy = fs::read_to_string(dir.join("reports/2026-01-01_test.recipe.toml")).unwrap();
        assert_eq!(copy, "engine = \"vllm\"\nname = \"X\"\nmodel = \"m\"\n");
        let m: Manifest = serde_json::from_str(
            &fs::read_to_string(dir.join("reports/2026-01-01_test.manifest.json")).unwrap(),
        )
        .unwrap();
        assert_eq!(m.schema_version, 1);
        assert_eq!(m.run, "2026-01-01_test");
        assert_eq!(
            m.recipe.sha256,
            sha256_hex(b"engine = \"vllm\"\nname = \"X\"\nmodel = \"m\"\n")
        );
        assert_eq!(m.bench.version.as_deref(), Some("1.2.3"));
        assert_eq!(m.bench.points.as_deref(), Some("163/176"));
        assert_eq!(m.bench.final_score, Some(93));
        // no nodes → the image entry exists but has an empty node map
        let img = &m.images["img:tag"];
        assert_eq!(img["build_checksum"], "");
        assert!(img["nodes"].as_object().unwrap().is_empty());
        fs::remove_dir_all(&dir).ok();
    }
}