1use 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}