xtask/
metrics.rs

1use anyhow::{Context, Result, ensure};
2use regex::Regex;
3use serde_json::{Value, json};
4
5const BYTES_PER_GB: f64 = 1_000_000_000.0;
6
7fn fields(text: &str, label: &str, pattern: &str) -> Result<Vec<f64>> {
8    let regex = Regex::new(pattern)?;
9    let captures = regex
10        .captures(text)
11        .with_context(|| format!("incomplete engine telemetry: {label}"))?;
12
13    captures
14        .iter()
15        .skip(1)
16        .map(|m| Ok(m.unwrap().as_str().parse()?))
17        .collect()
18}
19
20pub fn parse(text: &str) -> Result<Value> {
21    let load = fields(text, "load", r"load ([\d.]+)s, clock probe ([\d.]+) ms")?;
22    let prefill = fields(
23        text,
24        "prefill",
25        r"prefill (\d+) tokens in ([\d.]+)s \(([\d.]+) tok/s\)",
26    )?;
27    let decode = fields(
28        text,
29        "decode",
30        r"decode (\d+) tokens in ([\d.]+)s \(([\d.]+) tok/s\).*?\| (\d+) steps, mean step ([\d.]+) ms.*?gpu-(?:active|span) ([\d.]+) ms.*?io wait ([\d.]+) ms",
31    )?;
32    let memory = memory_stats(text)?;
33    let metal_gb = match &memory {
34        Some(stats) => {
35            stats["metal_allocated_bytes_observed"]
36                .as_u64()
37                .context("memory_stats: expected integer Metal bytes")? as f64
38                / BYTES_PER_GB
39        }
40        None => fields(text, "memory", r"\| ([\d.]+) GB Metal")?[0],
41    };
42    let clock = fields(text, "clock_end", r"clock probe at end ([\d.]+) ms")?;
43    let drafts = fields(text, "drafts", r"drafts (\d+), accepted ([\d.]+)/step")?;
44    let build = fields(
45        text,
46        "store_build",
47        r"built the [23]-bit store in ([\d.]+)s",
48    )
49    .ok();
50
51    Ok(json!({
52        "store_build_seconds":build.map(|v|v[0]), "load_seconds":load[0], "clock_start_ms":load[1],
53        "prompt_tokens":prefill[0] as u64, "prefill_seconds":prefill[1], "pp_s":prefill[2],
54        "output_tokens":decode[0] as u64, "decode_seconds":decode[1], "tg_s":decode[2],
55        "steps":decode[3] as u64, "mean_step_ms":decode[4], "gpu_span_ms":decode[5], "io_wait_ms":decode[6],
56        "metal_gb":metal_gb, "memory":memory,
57        "clock_end_ms":clock[0], "drafts":drafts[0] as u64, "accepted_drafts_per_step":drafts[1]
58    }))
59}
60
61fn memory_stats(text: &str) -> Result<Option<Value>> {
62    let mut snapshots = text
63        .lines()
64        .filter_map(|line| line.strip_prefix("memory_stats "));
65    let Some(line) = snapshots.next() else {
66        return Ok(None);
67    };
68
69    ensure!(
70        snapshots.next().is_none(),
71        "multiple memory_stats snapshots in one run"
72    );
73
74    let value: Value = serde_json::from_str(line).context("invalid memory_stats JSON")?;
75
76    ensure!(value.is_object(), "memory_stats must be an object");
77
78    Ok(Some(value))
79}
80
81pub fn extract_svg(text: &str) -> Result<(&str, usize)> {
82    let regex = Regex::new(r"(?is)<svg\b.*?</svg\s*>")?;
83    let svg = regex
84        .find(text)
85        .context("no complete SVG: output may have hit its token limit")?
86        .as_str();
87    let document = roxmltree::Document::parse(svg)?;
88
89    anyhow::ensure!(
90        document.root_element().tag_name().name() == "svg",
91        "root element is not svg"
92    );
93
94    Ok((
95        svg,
96        document.descendants().filter(|n| n.is_element()).count(),
97    ))
98}
99
100pub fn repeated_tail(text: &str) -> Option<Value> {
101    let mut words: Vec<&str> = text.split_whitespace().rev().skip(1).take(4096).collect();
102
103    words.reverse();
104
105    for period in 32..=512.min(words.len() / 4) {
106        let tail = &words[words.len() - 4 * period..];
107        let block = &words[words.len() - period..];
108
109        if tail.chunks_exact(period).all(|chunk| chunk == block) {
110            return Some(
111                json!({"words_per_cycle":period,"repetitions":4,"example":block.join(" ").chars().take(240).collect::<String>()}),
112            );
113        }
114    }
115
116    None
117}
xtask/metrics.rs

xtask/
metrics.rs
