Skip to main content

xtask/
readme.rs

1//! Static README tables and images, generated from a saved benchmark report.
2use crate::report::markdown::{cell, drawings, label, link};
3use crate::{report, util};
4use anyhow::{Context, Result, ensure};
5use serde_json::Value;
6use std::{fmt::Write as _, fs, path::Path};
7
8const START: &str = "<!-- benchmarks:start -->";
9const END: &str = "<!-- benchmarks:end -->";
10
11fn table(data: &Value, rows: &[Value], kind: &str, metric: &str) -> Result<String> {
12    let cases: Vec<_> = data["cases"]
13        .as_array()
14        .context("report cases")?
15        .iter()
16        .filter(|case| case["kind"].as_str().unwrap_or("decode") == kind)
17        .collect();
18
19    if cases.is_empty() {
20        return Ok(String::new());
21    }
22
23    let mut text = String::from("| Experts |");
24
25    for case in &cases {
26        write!(text, " {} {metric} |", cell(report::text(&case["id"])))?;
27    }
28
29    text.push_str("\n| --- |");
30    text.push_str(&" ---: |".repeat(cases.len()));
31    text.push('\n');
32
33    for config in data["configurations"]
34        .as_array()
35        .context("configurations")?
36    {
37        write!(text, "| {} |", cell(label(config)))?;
38
39        for case in &cases {
40            let row = rows
41                .iter()
42                .find(|r| r["case"] == case["id"] && r["configuration"] == config["id"]);
43            let key = if kind == "prefill" { "pp_s" } else { "tg_s" };
44
45            match row.and_then(|r| r[key].as_f64()) {
46                Some(rate) => write!(text, " {rate:.2} |")?,
47                None => text.push_str(" n/a |"),
48            }
49        }
50
51        text.push('\n');
52    }
53
54    text.push('\n');
55
56    Ok(text)
57}
58
59pub fn render(data: &Value, directory: &Path) -> Result<String> {
60    ensure!(
61        data["utc_finished"].is_string(),
62        "finish the benchmark before updating the README"
63    );
64
65    let rows = report::rows(data)?;
66
67    ensure!(!rows.is_empty(), "no valid benchmark samples");
68
69    let runs = data["runs"].as_array().context("runs")?;
70    let valid = runs.iter().filter(|r| r["status"] == "ok").count();
71    let provenance = &data["provenance"];
72    let commit = report::text(&provenance["commit"]);
73    let memory = provenance["memory_bytes"]
74        .as_str()
75        .and_then(|v| v.parse::<f64>().ok())
76        .context("machine memory")?
77        / 1073741824.0;
78    let metal = rows
79        .iter()
80        .filter_map(|r| r["metal_gb"].as_f64())
81        .fold(0.0, f64::max);
82    let report_link = link(&directory.join("report.json"));
83    let summary_link = link(&directory.join("summary.md"));
84    let mut text = format!(
85        "The benchmark ran on {} hardware with {memory:.0} GiB of memory.\nThe engine reported {metal:.2} GB of Metal allocations.\nThe report contains {valid} valid samples from revision `{}`.\n\nThe inference rates exclude loading and store construction. Answer lengths vary,\nso compare completion times in the full report.\n\n[Full report]({report_link}).\n\n",
86        cell(report::text(&provenance["hardware"])),
87        cell(commit.get(..7).unwrap_or(commit)),
88    );
89
90    text.push_str(&table(data, &rows, "decode", "tg/s")?);
91    text.push_str(&table(data, &rows, "prefill", "pp/s")?);
92    append_cut_warning(&mut text, data)?;
93    writeln!(text, "[All timings and outputs]({summary_link}).\n")?;
94
95    let pelicans = drawings(data, directory)?;
96
97    if !pelicans.is_empty() {
98        text.push_str("### Pelicans\n\nThese are unedited model outputs from the benchmark.\n\n");
99        text.push_str(&pelicans);
100    }
101
102    Ok(text)
103}
104
105fn append_cut_warning(text: &mut String, data: &Value) -> Result<()> {
106    let configs = data["configurations"]
107        .as_array()
108        .context("configurations")?;
109    let cut = configs
110        .iter()
111        .filter_map(|c| c["args"].as_array())
112        .any(|args| {
113            args.windows(2).any(|pair| {
114                pair[0] == "--cut-weak"
115                    && pair[1]
116                        .as_str()
117                        .and_then(|v| v.parse::<f32>().ok())
118                        .is_some_and(|v| v > 0.0)
119            })
120        });
121
122    if cut {
123        text.push_str("Settings with a deadline cut are not reproducible.\n\n");
124    }
125
126    Ok(())
127}
128
129pub fn replace_section(readme: &str, generated: &str) -> Result<String> {
130    ensure!(
131        readme.matches(START).count() == 1 && readme.matches(END).count() == 1,
132        "README needs exactly one benchmark marker pair"
133    );
134
135    let start = readme.find(START).unwrap() + START.len();
136    let end = readme.find(END).unwrap();
137
138    ensure!(start < end, "README benchmark markers are reversed");
139
140    Ok(format!(
141        "{}\n\n{}\n{}",
142        &readme[..start],
143        generated.trim_end(),
144        &readme[end..]
145    ))
146}
147
148pub fn update(directory: &Path) -> Result<()> {
149    let root = util::root().canonicalize()?;
150    let directory = directory.canonicalize()?;
151    let relative = directory
152        .strip_prefix(&root)
153        .context("README results must be inside the repository")?;
154    let data = util::json(&directory.join("report.json"))?;
155
156    for run in data["runs"].as_array().context("runs")? {
157        if let Some(svg) = run["svg"].as_str() {
158            ensure!(
159                directory.join(svg).is_file(),
160                "missing pelican image: {svg}"
161            );
162        }
163    }
164
165    let generated = render(&data, relative)?;
166    let path = root.join("README.md");
167    let text = replace_section(&fs::read_to_string(&path)?, &generated)?;
168
169    fs::write(&path, text)?;
170    println!("Updated {} from {}", path.display(), directory.display());
171
172    Ok(())
173}