xtask/
bench.rs

1use crate::{
2    capture, hardware, metrics, report,
3    suite::{self, Case, Configuration, Suite},
4    util,
5};
6use anyhow::{Context, Result, ensure};
7use clap::{Args, ValueEnum};
8use serde_json::{Value, json};
9use std::{
10    fs,
11    path::{Path, PathBuf},
12};
13
14mod model;
15use model::Model;
16
17#[derive(Args)]
18pub struct Options {
19    /// Indexed model alias, source URI, or prepared directory
20    pub model_dir: PathBuf,
21    /// Cherenkov index root, forwarded to inspection and inference
22    #[arg(long)]
23    pub root: Option<PathBuf>,
24    #[arg(long, default_value = "benchmarks/suite.json")]
25    pub suite: PathBuf,
26    #[arg(long)]
27    pub output: Option<PathBuf>,
28    #[arg(long)]
29    pub configs: Option<String>,
30    #[arg(long)]
31    pub cases: Option<String>,
32    #[arg(long)]
33    pub rounds: Option<usize>,
34    #[arg(long)]
35    pub case_cap: Vec<String>,
36    #[arg(long)]
37    pub binary: Option<PathBuf>,
38    #[arg(long)]
39    pub resume: bool,
40    #[arg(long)]
41    pub build_stores: bool,
42    #[arg(long)]
43    pub allow_battery: bool,
44    #[arg(long)]
45    pub dry_run: bool,
46    /// Rebuild the README benchmark section after the suite finishes.
47    #[arg(long)]
48    pub update_readme: bool,
49    /// Free text recorded in the report's provenance (machine, conditions).
50    #[arg(long)]
51    pub note: Option<String>,
52    /// Zip the results directory beside itself when the suite finishes.
53    #[arg(long)]
54    pub archive: bool,
55    /// Preset: `light` runs one round of code, prose and the long prefill on
56    /// the stores already built; `heavy` runs the whole suite. Both archive.
57    #[arg(long, value_enum)]
58    pub mode: Option<Mode>,
59}
60
61#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)]
62pub enum Mode {
63    Light,
64    Heavy,
65}
66
67/// Cases the light preset runs when `--cases` is not given.
68pub const LIGHT_CASES: &str = "code,prose,prefill-long";
69
70/// Hide runner-supplied locations while preserving prompts and configuration arguments.
71pub fn redact_args(args: &mut [Value], config_args: usize) {
72    for (arg, replacement) in args.iter_mut().zip(["<binary>", "<model>"]) {
73        *arg = json!(replacement);
74    }
75
76    let root_flag = 3 + config_args;
77
78    if args.get(root_flag).is_some_and(|arg| arg == "--root")
79        && let Some(value) = args.get_mut(root_flag + 1)
80    {
81        *value = json!("<root>");
82    }
83}
84
85/// Locations the runner may include in diagnostics, ordered before their parent root.
86fn path_replacements(
87    binary: &Path,
88    model: &Model,
89    root: Option<&Path>,
90) -> Vec<(String, &'static str)> {
91    [
92        (Some(binary), "<binary>"),
93        (Some(model.directory.as_path()), "<model>"),
94        (model.source_directory.as_deref(), "<model>"),
95        (root, "<root>"),
96    ]
97    .into_iter()
98    .filter_map(|(path, label)| path.map(|p| (p.to_string_lossy().into_owned(), label)))
99    .collect()
100}
101
102fn redact_paths(text: &str, replacements: &[(String, &str)]) -> String {
103    let mut text = text.to_owned();
104
105    for (path, replacement) in replacements {
106        if !path.is_empty() {
107            text = text.replace(path, replacement);
108        }
109    }
110
111    text
112}
113
114fn redact_saved_metadata(report: &mut Value, binary: &Path, model: &Model, root: Option<&Path>) {
115    let saved_binary = report["provenance"]["binary"]
116        .as_str()
117        .unwrap_or("")
118        .to_owned();
119    let saved_model = report["provenance"]["model"]
120        .as_str()
121        .unwrap_or("")
122        .to_owned();
123    let mut replacements = path_replacements(binary, model, root);
124
125    replacements.insert(0, (saved_binary, "<binary>"));
126    replacements.insert(0, (saved_model, "<model>"));
127
128    let configurations = report["configurations"].clone();
129
130    // Suite contents participate in resume validation. Only redact fields
131    // supplied by the runner, leaving prompts and configuration args intact.
132    for key in ["runs", "previous_attempts"] {
133        if let Some(runs) = report.get_mut(key).and_then(Value::as_array_mut) {
134            for run in runs {
135                let config_args = configurations
136                    .as_array()
137                    .and_then(|configs| {
138                        configs
139                            .iter()
140                            .find(|config| config["id"] == run["configuration"])
141                    })
142                    .and_then(|config| config["args"].as_array())
143                    .map_or(0, Vec::len);
144
145                if let Some(args) = run["args"].as_array_mut() {
146                    redact_args(args, config_args);
147                }
148
149                if let Some(error) = run["error"].as_str() {
150                    run["error"] = json!(redact_paths(error, &replacements));
151                }
152            }
153        }
154    }
155
156    let provenance = &mut report["provenance"];
157    provenance["binary"] = json!("<binary>");
158    provenance["model"] = json!("<model>");
159
160    if let Some(platform) = provenance["platform"].as_str() {
161        let fields: Vec<_> = platform.split_whitespace().collect();
162
163        // Legacy macOS reports used uname -a: sysname, hostname, release,
164        // version (several words), machine. Keep the original machine facts.
165        if fields.first() == Some(&"Darwin") && fields.len() > 3 {
166            provenance["platform"] = json!(format!(
167                "{} {} {}",
168                fields[0],
169                fields[2],
170                fields.last().unwrap()
171            ));
172        }
173    }
174
175    if let Some(settings) = report.get_mut("settings") {
176        settings["suite"] = json!("<suite>");
177    }
178}
179
180fn set_caps(cases: &mut [Case], overrides: &[String]) -> Result<()> {
181    for value in overrides {
182        let (id, limit) = value
183            .split_once('=')
184            .context("--case-cap requires ID=TOKENS")?;
185        let limit = limit.parse()?;
186
187        ensure!(limit > 0, "token cap must be positive");
188
189        cases
190            .iter_mut()
191            .find(|c| c.id == id)
192            .with_context(|| format!("unknown case: {id}"))?
193            .max_tokens = limit;
194    }
195
196    Ok(())
197}
198
199fn check_stores(model: &Model, configs: &[Configuration], allow_build: bool) -> Result<()> {
200    ensure!(
201        model.precisions.contains(&4),
202        "model must already be packed"
203    );
204
205    for config in configs {
206        let Some(bits) = config.store_bits else {
207            continue;
208        };
209        let present = model.precisions.contains(&bits);
210
211        ensure!(
212            present || allow_build,
213            "{bits}-bit store missing; select cached configurations or pass --build-stores"
214        );
215
216        if !present {
217            eprintln!("Allowing first-use {bits}-bit construction, reported in load time.");
218        }
219    }
220
221    Ok(())
222}
223
224pub fn raise_saved_caps(out: &Path, report: &mut Value, signature: Value) -> Result<()> {
225    let revision = json!({"previous_signature":report["signature"],"utc_changed":util::utc()?,"reason":"Raised answer safety caps; prompts, context, binary, and model unchanged."});
226
227    append(report, "suite_revisions", revision);
228
229    let runs = std::mem::take(report["runs"].as_array_mut().context("report runs")?);
230
231    for mut record in runs {
232        let args = record["args"].as_array().context("sample args")?;
233        let pos = args
234            .iter()
235            .position(|a| a == "--max-tokens")
236            .context("sample token cap")?;
237        let old_cap = args[pos + 1]
238            .as_str()
239            .context("sample cap value")?
240            .parse::<u64>()?;
241        let cap = signature["cases"]
242            .as_array()
243            .context("signature cases")?
244            .iter()
245            .find(|c| c["id"] == record["case"])
246            .context("case missing")?["max_tokens"]
247            .as_u64()
248            .context("token cap")?;
249
250        if record["status"] != "incomplete" || cap <= old_cap {
251            append(report, "runs", record);
252
253            continue;
254        }
255
256        let archive = format!(
257            "attempts/{}-cap-{old_cap}.txt",
258            record["id"].as_str().context("sample id")?
259        );
260
261        fs::create_dir_all(out.join("attempts"))?;
262        fs::rename(
263            out.join(record["output"].as_str().context("sample output")?),
264            out.join(&archive),
265        )?;
266
267        record["output"] = json!(archive);
268
269        append(report, "previous_attempts", record);
270    }
271
272    report["cases"] = signature["cases"].clone();
273    report["signature"] = signature;
274
275    Ok(())
276}
277
278pub fn append(report: &mut Value, key: &str, value: Value) {
279    report
280        .as_object_mut()
281        .unwrap()
282        .entry(key.to_owned())
283        .or_insert_with(|| json!([]))
284        .as_array_mut()
285        .unwrap()
286        .push(value);
287}
288
289fn record_result(
290    out: &Path,
291    config: &Configuration,
292    case: &Case,
293    capture: &capture::Captured,
294    record: &mut Value,
295    allow_build: bool,
296    metal_limit_gb: f64,
297) -> Result<()> {
298    if let Some(cycle) = &capture.cycle {
299        record["status"] = json!("cycling");
300
301        anyhow::bail!(
302            "stopped after four repeated word blocks: {}",
303            cycle["example"]
304        );
305    }
306
307    ensure!(
308        capture.code == 0,
309        "engine exited {}: {}",
310        capture.code,
311        capture
312            .stderr
313            .chars()
314            .rev()
315            .take(3000)
316            .collect::<String>()
317            .chars()
318            .rev()
319            .collect::<String>()
320    );
321
322    let m = metrics::parse(&capture.stderr)?;
323    record["metrics"] = m.clone();
324
325    ensure!(
326        m["metal_gb"].as_f64().unwrap() <= metal_limit_gb,
327        "reported Metal allocations exceeded this machine's {metal_limit_gb:.1} GB of memory"
328    );
329    ensure!(
330        capture.stable_power(),
331        "power source changed during this sample"
332    );
333    ensure!(
334        allow_build || m["store_build_seconds"].is_null(),
335        "unexpected store construction; use --build-stores"
336    );
337
338    let reason = if m["output_tokens"].as_u64().unwrap() >= case.max_tokens as u64 {
339        "length"
340    } else {
341        "eos"
342    };
343    record["finish_reason"] = json!(reason);
344
345    if case.stop == "eos" && reason != "eos" {
346        record["status"] = json!("incomplete");
347
348        anyhow::bail!("incomplete answer: reached the token safety cap before EOS");
349    }
350
351    ensure!(
352        case.stop != "length" || reason == "length",
353        "fixed-length decode did not reach its token count"
354    );
355
356    if case.kind == "svg" {
357        record["status"] = json!("invalid_svg");
358        let text = fs::read_to_string(out.join(record["output"].as_str().unwrap()))?;
359        let (svg, elements) = metrics::extract_svg(&text)?;
360        let path = format!("pelicans/{}.svg", config.id);
361
362        fs::write(out.join(&path), format!("{svg}\n"))?;
363
364        record["svg"] = json!(path);
365        record["svg_elements"] = json!(elements);
366        record["within_element_budget"] = json!(elements <= 120);
367    }
368
369    record["status"] = json!("ok");
370
371    Ok(())
372}
373
374struct Run<'a> {
375    binary: &'a Path,
376    model: &'a Model,
377    out: &'a Path,
378    options: &'a Options,
379    max_ctx: usize,
380    /// Physical memory in decimal GB; Metal allocations beyond it are an error.
381    memory_gb: f64,
382}
383
384impl Run<'_> {
385    fn sample(&self, config: &Configuration, case: &Case, round: usize) -> Result<Value> {
386        let id = format!("r{}-{}-{}", round + 1, case.id, config.id);
387        let output = format!("outputs/{id}.txt");
388        let mut args = vec![
389            self.binary.display().to_string(),
390            self.model.reference.clone(),
391            case.prompt(),
392        ];
393
394        args.extend(config.args.clone());
395
396        if let Some(root) = &self.options.root {
397            args.extend(["--root".into(), root.display().to_string()]);
398        }
399
400        args.extend([
401            "--max-tokens".into(),
402            case.max_tokens.to_string(),
403            "--max-ctx".into(),
404            case.max_ctx.unwrap_or(self.max_ctx).to_string(),
405        ]);
406
407        if case.stop == "length" {
408            args.push("--no-eos".into());
409        }
410
411        eprintln!("[{id}] starting");
412
413        // A previous sample may have published additional expert variants.
414        let current = Model::inspect(
415            self.binary,
416            self.model.reference.as_ref(),
417            self.options.root.as_deref(),
418        )?;
419        let present = config
420            .store_bits
421            .is_none_or(|bits| current.precisions.contains(&bits));
422        let c = capture::run(
423            &args,
424            &self.out.join(&output),
425            self.options.allow_battery,
426            None,
427        )?;
428        let mut record = json!({
429            "id": id,
430            "configuration": config.id,
431            "case": case.id,
432            "round": round + 1,
433            "args": args,
434            "output": output,
435            "power_before": c.power_before,
436            "power_after": c.power_after,
437            "power_samples": c.power_samples,
438            "status": "failed",
439            "store_present_before": present,
440            "wall_seconds": c.wall_seconds,
441            "exit_code": c.code,
442        });
443
444        redact_args(record["args"].as_array_mut().unwrap(), config.args.len());
445
446        if let Some(cycle) = &c.cycle {
447            record["cycle"] = cycle.clone();
448        }
449
450        if let Err(error) = record_result(
451            self.out,
452            config,
453            case,
454            &c,
455            &mut record,
456            self.options.build_stores,
457            self.memory_gb,
458        ) {
459            let replacements =
460                path_replacements(self.binary, self.model, self.options.root.as_deref());
461            record["error"] = json!(redact_paths(&error.to_string(), &replacements));
462        }
463
464        eprintln!(
465            "[{id}] {}, tg/s {}",
466            record["status"], record["metrics"]["tg_s"]
467        );
468
469        Ok(record)
470    }
471    fn new_report(
472        &self,
473        signature: Value,
474        configs: &[Configuration],
475        cases: &[Case],
476    ) -> Result<Value> {
477        let Self {
478            binary,
479            model,
480            options,
481            max_ctx,
482            ..
483        } = self;
484
485        let status = util::output(&["git", "status", "--porcelain"])?;
486
487        eprintln!("Describing hardware and sampling store reads.");
488
489        Ok(json!({
490            "version": 1,
491            "signature": signature,
492            "configurations": configs,
493            "cases": cases,
494            "runs": [],
495            "settings": {
496                "suite": "<suite>",
497                "mode": options.mode.map(|m| format!("{m:?}").to_lowercase()),
498                "max_ctx": max_ctx,
499                "case_caps": options.case_cap,
500                "build_stores": options.build_stores,
501                "binary_override": options.binary.is_some(),
502                "pool": "adaptive unless a configuration passes --pool-gb; the server TOML is not read",
503                "environment": "CHERENKOV_* variables are removed from every sample",
504            },
505            "provenance": {
506                "commit": util::output(&["git", "rev-parse", "HEAD"])?,
507                "dirty": !status.is_empty(),
508                "git_status": status,
509                "source_sha256": util::source_digest()?,
510                "binary_sha256": util::digest(binary)?,
511                "binary": "<binary>",
512                "model": "<model>",
513                "model_metadata_sha256": signature["model_metadata_sha256"],
514                "utc_started": util::utc()?,
515                "platform": util::output(&["uname", "-srm"])?,
516                "hardware": util::output(&["sysctl", "-n", "machdep.cpu.brand_string"])?,
517                "memory_bytes": util::output(&["sysctl", "-n", "hw.memsize"])?,
518                "rustc": util::output(&["rustc", "--version"])?,
519                "power": capture::power(),
520                "hardware_detail": hardware::describe(&model.directory),
521                "note": options.note,
522                "method": "fresh processes; rotated interleaved rounds; separate load/prefill/decode; no check; no prefix cache; complete answers through EOS; SVG phase after timings"
523            }
524        }))
525    }
526
527    fn prepare_report(
528        &self,
529        signature: Value,
530        configs: &[Configuration],
531        cases: &[Case],
532    ) -> Result<Value> {
533        let Self {
534            out,
535            options,
536            binary,
537            model,
538            ..
539        } = self;
540
541        if options.resume {
542            let mut report = util::json(&out.join("report.json"))?;
543
544            model.migrate_signature(&mut report["signature"]);
545
546            if let Some(revisions) = report["suite_revisions"].as_array_mut() {
547                for revision in revisions {
548                    model.migrate_signature(&mut revision["previous_signature"]);
549                }
550            }
551
552            ensure!(
553                report["signature"]["model_id"].is_string(),
554                "cannot resume: saved report has no matching model identity; choose a new output directory"
555            );
556
557            if report["signature"] != signature {
558                ensure!(
559                    suite::only_higher_caps(&report["signature"], &signature),
560                    "cannot resume: binary, suite, model, or selections changed"
561                );
562                raise_saved_caps(out, &mut report, signature)?;
563            }
564
565            redact_saved_metadata(&mut report, binary, model, options.root.as_deref());
566
567            return Ok(report);
568        }
569
570        ensure!(
571            !out.exists() || fs::read_dir(out)?.next().is_none(),
572            "output directory is not empty; choose another or --resume"
573        );
574        fs::create_dir_all(out)?;
575
576        self.new_report(signature, configs, cases)
577    }
578}
579
580/// Inspect once before scheduling; each timed sample refreshes variant availability.
581fn resolve_model(options: &Options) -> Result<(PathBuf, Model)> {
582    if options.binary.is_none() {
583        util::build()?;
584    }
585
586    let binary = options
587        .binary
588        .clone()
589        .unwrap_or_else(|| util::root().join("target/release/cherenkov"))
590        .canonicalize()?;
591    let model = Model::inspect(
592        &binary,
593        options.model_dir.as_os_str(),
594        options.root.as_deref(),
595    )?;
596
597    Ok((binary, model))
598}
599
600pub fn run(options: Options) -> Result<()> {
601    let suite: Suite = serde_json::from_value(util::json(&options.suite)?)?;
602    let light = options.mode == Some(Mode::Light);
603    let mut configs = suite::select(&suite.configurations, options.configs.as_deref(), |c| &c.id)?;
604
605    let resolved = if !options.dry_run || (light && options.configs.is_none()) {
606        Some(resolve_model(&options)?)
607    } else {
608        None
609    };
610
611    // The light preset measures what is already on disk instead of building stores.
612    if light && options.configs.is_none() {
613        let (_, model) = resolved
614            .as_ref()
615            .context("light mode requires model inspection")?;
616
617        configs.retain(|c| {
618            c.store_bits
619                .is_none_or(|bits| model.precisions.contains(&bits))
620        });
621    }
622
623    let case_selection = options.cases.as_deref().or(light.then_some(LIGHT_CASES));
624    let mut cases = suite::select(&suite.cases, case_selection, |c| &c.id)?;
625
626    set_caps(&mut cases, &options.case_cap)?;
627
628    let rounds = options
629        .rounds
630        .unwrap_or(if light { 1 } else { suite.rounds });
631
632    ensure!(
633        rounds > 0 && !configs.is_empty() && !cases.is_empty(),
634        "rounds and selections must be nonempty"
635    );
636
637    let jobs = suite::schedule(&configs, &cases, rounds);
638
639    if options.dry_run {
640        let plan: Vec<_> = jobs
641            .iter()
642            .map(|&(round, config, case)| {
643                json!({
644                    "round": round + 1,
645                    "config": configs[config].id,
646                    "case": cases[case].id,
647                    "max_tokens": cases[case].max_tokens,
648                })
649            })
650            .collect();
651
652        println!("{}", serde_json::to_string_pretty(&plan)?);
653
654        return Ok(());
655    }
656
657    let (binary, model) = resolved.context("benchmark model was not resolved")?;
658
659    check_stores(&model, &configs, options.build_stores)?;
660
661    let default_name = match options.mode {
662        Some(Mode::Light) => format!("{}-light", util::utc()?),
663        _ => util::utc()?,
664    };
665    let out = util::absolute(
666        &options
667            .output
668            .clone()
669            .unwrap_or(util::root().join("results").join(default_name)),
670    )?;
671    let metadata = model.metadata()?;
672    let memory_gb = util::output(&["sysctl", "-n", "hw.memsize"])?
673        .parse::<f64>()
674        .context("hw.memsize")?
675        / 1e9;
676
677    let signature = json!({
678        "model_metadata_sha256": metadata,
679        "binary_sha256": util::digest(&binary)?,
680        "suite_sha256": util::digest(&options.suite)?,
681        "model": "<model>",
682        "model_id": model.reference,
683        "configs": configs,
684        "cases": cases,
685        "rounds": rounds,
686        "allow_battery": options.allow_battery,
687    });
688    let runner = Run {
689        binary: &binary,
690        model: &model,
691        out: &out,
692        options: &options,
693        max_ctx: suite.max_ctx,
694        memory_gb,
695    };
696    let mut report = runner.prepare_report(signature, &configs, &cases)?;
697
698    fs::create_dir_all(out.join("outputs"))?;
699    fs::create_dir_all(out.join("pelicans"))?;
700    report::write(&out, &report)?;
701
702    for (r, c, t) in jobs {
703        let id = format!("r{}-{}-{}", r + 1, cases[t].id, configs[c].id);
704
705        if report["runs"].as_array().unwrap().iter().any(|v| {
706            v["id"] == id
707                && matches!(
708                    v["status"].as_str(),
709                    Some("ok" | "incomplete" | "invalid_svg" | "cycling")
710                )
711        }) {
712            continue;
713        }
714
715        let record = runner.sample(&configs[c], &cases[t], r)?;
716        let failed = record["status"] == "failed";
717        let error = record["error"].clone();
718
719        report["runs"]
720            .as_array_mut()
721            .unwrap()
722            .retain(|v| v["id"] != id);
723        append(&mut report, "runs", record);
724        report::write(&out, &report)?;
725        ensure!(
726            !failed,
727            "{error}; outputs saved in {}; fix and --resume",
728            out.display()
729        );
730    }
731
732    report["utc_finished"] = json!(util::utc()?);
733
734    report::write(&out, &report)?;
735
736    if options.update_readme {
737        crate::readme::update(&out)?;
738    }
739
740    println!("{}", out.join("gallery.html").display());
741
742    if options.archive || options.mode.is_some() {
743        let zip = hardware::archive(&out)?;
744
745        println!("{}", zip.display());
746        eprintln!("Attach the archive to a pull request or issue to share this run.");
747    }
748
749    Ok(())
750}
xtask/bench.rs

xtask/
bench.rs
