metrale_bench/benchmarks/kat_equality/
driver.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The part that talks to a server: issue the draw once per order,
4//! keep every reply, and hand the lot to `compare`.
5//!
6//! Owner: bench, kat_equality.
7//! Invariants: every order issues the same samples, a permutation of the
8//! loaded draw.
9
10use std::future::Future;
11use std::time::{Duration, Instant};
12
13use anyhow::{Context, Result};
14use serde_json::json;
15
16use crate::benchmark::{Benchmark, BenchmarkDescriptor};
17use crate::benchmarks::bfcl::{MAX_NEW_TOKENS, dataset, draw::DrawSpec, provision};
18use crate::benchmarks::transcript::{RequestOutcome, Transcript};
19use crate::hardware::Sensitivity;
20use crate::http;
21use crate::metadata::PluginMetadata;
22use crate::params::{ParamKind, ParamSpec, ParamValue, ParamValues};
23use crate::plugin::{Plugin, PluginHandle};
24use crate::result::{BenchmarkResult, LogLine, RunStatus};
25
26use super::compare::{Observation, OrderRun, permutation, score, verdict};
27
28const SUMMARY: &str = "The same sample must answer the same, whatever ran before it";
29
30/// 2026-09-26: Two orders is the minimum that proves anything, and the second
31/// is the reverse, which changes what preceded every sample.
32const DEFAULT_ORDERS: usize = 2;
33
34/// 2026-09-26: `0` means the whole draw. A cap exists because the whole golden
35/// draw (995 samples) in two orders is 1990 generations; the default is not
36/// smaller because a capped run would report a pass with a fraction of the
37/// samples.
38const DEFAULT_SAMPLE_CAP: usize = 0;
39
40pub const METADATA: PluginMetadata = PluginMetadata::metrale(SUMMARY);
41
42pub const DESCRIPTOR: BenchmarkDescriptor = BenchmarkDescriptor {
43    id: "kat-equality-gate",
44    name: "KAT Equality Gate",
45    summary: SUMMARY,
46    detail: "Issues one BFCL draw against ONE server in two or more request ORDERS \
47             (canonical, then reversed, then rotations) at temperature 0, and requires \
48             every sample_id's reply to be byte-identical across all of them — text, \
49             reasoning, tool calls with RAW arguments, finish reason, and completion \
50             token count. Any difference FAILS: at temperature 0 a reply must be a \
51             function of its own request, and a benchmark whose samples can see each \
52             other cannot be sharded, nor can a score taken under one order describe \
53             another. A failed request is UNMEASURED, never agreement, and a run whose \
54             replies were all empty fails as VACUOUS rather than certifying a dead \
55             endpoint as order-independent. Serve it with --hermetic to measure the \
56             regime a known-answer test actually requires.",
57    duration_hint: "~2x one BFCL leg",
58    expected_secs: 4200,
59    updated: "2026-09-09",
60    needs_confirmation: false,
61    // 2026-09-26: The invariant is a property of the engine, so it is not
62    // pinned to a checkpoint (`ssm_poison` sets `intended_for: None` too).
63    intended_for: None,
64    threshold_params: &[],
65    // 2026-09-26: A busy box can make a run slower; it cannot make one
66    // request's output depend on another's.
67    sensitivity: Sensitivity::Correctness,
68    ctor: || Box::new(KatEquality::default()),
69};
70
71enum Phase {
72    Provision,
73    Generate,
74    Compare,
75    Done,
76}
77
78pub struct KatEquality {
79    handle: Option<PluginHandle>,
80    phase: Phase,
81    artifacts: Option<provision::Artifacts>,
82    /// 2026-09-26: The draw in canonical order. Every order is an index
83    /// permutation of this, so no order can measure a different sample set.
84    samples: Vec<dataset::Sample>,
85    order_index: usize,
86    cursor: usize,
87    runs: Vec<OrderRun>,
88    current: Vec<Observation>,
89    orders: usize,
90    sample_cap: usize,
91    max_new_tokens: usize,
92    request_timeout: Duration,
93    started: Option<Instant>,
94}
95
96impl Default for KatEquality {
97    fn default() -> Self {
98        Self {
99            handle: None,
100            phase: Phase::Provision,
101            artifacts: None,
102            samples: Vec::new(),
103            order_index: 0,
104            cursor: 0,
105            runs: Vec::new(),
106            current: Vec::new(),
107            orders: DEFAULT_ORDERS,
108            sample_cap: DEFAULT_SAMPLE_CAP,
109            max_new_tokens: MAX_NEW_TOKENS,
110            request_timeout: Duration::from_secs(300),
111            started: None,
112        }
113    }
114}
115
116impl KatEquality {
117    fn handle(&self) -> Result<&PluginHandle> {
118        self.handle.as_ref().context("plugin was not loaded")
119    }
120
121    fn elapsed(&self) -> Duration {
122        self.started.map(|s| s.elapsed()).unwrap_or_default()
123    }
124
125    fn order_label(&self, index: usize) -> String {
126        match index {
127            0 => "canonical".to_string(),
128            1 => "reversed".to_string(),
129            k => format!("rotation-{k}"),
130        }
131    }
132
133    fn total_steps(&self) -> u64 {
134        (self.samples.len() * self.orders) as u64
135    }
136
137    fn done_steps(&self) -> u64 {
138        (self.order_index * self.samples.len() + self.cursor) as u64
139    }
140
141    /// 2026-09-26: Issue one sample and keep the whole reply.
142    ///
143    /// The body is BFCL's (`bfcl/exec.rs`: no `seed`) at temperature 0, plus
144    /// `stream_options`: this gate asks whether BFCL's own conditions are
145    /// order-independent. `include_usage` asks for the server's
146    /// `completion_tokens`; without `usage`, `http::chat_stream` counts
147    /// stream deltas instead.
148    async fn issue_one(&mut self, sample_index: usize) -> Result<()> {
149        let handle = self.handle()?.clone();
150        let sample = self.samples[sample_index].clone();
151        let target = handle.target();
152        let body = json!({
153            "model": target.model,
154            "stream": true,
155            "stream_options": {"include_usage": true},
156            "temperature": 0.0,
157            "max_tokens": self.max_new_tokens,
158            "messages": sample.messages,
159            "tools": sample.tools,
160            "tool_choice": sample.tool_choice,
161        });
162        let outcome = match http::chat_stream(target, &body, self.request_timeout).await {
163            // 2026-09-26: A transport failure is its own outcome. BFCL scores
164            // one as "no call", which here would make two failures look like
165            // agreement.
166            Err(e) => RequestOutcome::Error(format!("{e:#}")),
167            Ok(o) => RequestOutcome::Ok(Box::new(Transcript::from(&o))),
168        };
169        self.current.push(Observation {
170            sample_id: sample.sample_id.clone(),
171            outcome,
172        });
173        Ok(())
174    }
175}
176
177impl Plugin for KatEquality {
178    fn metadata(&self) -> &'static PluginMetadata {
179        &METADATA
180    }
181
182    fn load(&mut self, handle: PluginHandle) -> impl Future<Output = Result<()>> + Send {
183        self.started = Some(Instant::now());
184        self.handle = Some(handle.clone());
185        async move {
186            let artifacts = provision::ensure(handle.artifacts(), &handle).await?;
187            self.artifacts = Some(artifacts);
188            Ok(())
189        }
190    }
191}
192
193impl Benchmark for KatEquality {
194    fn descriptor(&self) -> &'static BenchmarkDescriptor {
195        &DESCRIPTOR
196    }
197
198    fn parameters(&self) -> Vec<ParamSpec> {
199        vec![
200            ParamSpec::new(
201                "orders",
202                "Request orders",
203                "How many different orders to issue the draw in. 2 is canonical plus \
204                 reversed, which is the strongest pair; further orders are rotations.",
205                ParamKind::Int { min: 2, max: 6 },
206                ParamValue::Int(DEFAULT_ORDERS as i64),
207            ),
208            ParamSpec::new(
209                "sample_cap",
210                "Sample cap",
211                "Truncate the draw to this many samples (0 = the whole draw). This is \
212                 a PREFIX, not a sample: the draw concatenates subsets in sorted name \
213                 order, so a cap decides which subsets are compared at all. Pick it from \
214                 where the effect lives, and re-run the negative control AT the cap — one \
215                 taken on the whole draw does not describe a capped run, because \
216                 truncating changes what ran before every surviving sample.",
217                ParamKind::Int { min: 0, max: 5000 },
218                ParamValue::Int(DEFAULT_SAMPLE_CAP as i64),
219            ),
220            ParamSpec::new(
221                "max_new_tokens",
222                "Max new tokens",
223                "Generation cap per sample. Defaults to BFCL's own budget: this \
224                 gate asks whether BFCL's conditions are order-independent, so a \
225                 different budget would answer about a regime nobody measures.",
226                ParamKind::Int { min: 32, max: 4096 },
227                ParamValue::Int(MAX_NEW_TOKENS as i64),
228            ),
229            ParamSpec::new(
230                "request_timeout_s",
231                "Request timeout (s)",
232                "Per-request timeout. A timeout is UNMEASURED, which fails the gate — \
233                 it is not evidence that the orders agreed.",
234                ParamKind::Int { min: 10, max: 3600 },
235                ParamValue::Int(300),
236            ),
237        ]
238    }
239
240    fn configure(&mut self, values: &ParamValues) -> Result<()> {
241        let specs = self.parameters();
242        values.validate_against(&specs)?;
243        self.orders = values.usize("orders")?;
244        self.sample_cap = values.usize("sample_cap")?;
245        self.max_new_tokens = values.usize("max_new_tokens")?;
246        self.request_timeout = Duration::from_secs(values.usize("request_timeout_s")? as u64);
247        // 2026-09-26: A re-`configure` must not leave a previous run's replies
248        // behind: they would be compared as if this run had produced them.
249        self.phase = Phase::Provision;
250        self.order_index = 0;
251        self.cursor = 0;
252        self.runs.clear();
253        self.current.clear();
254        Ok(())
255    }
256
257    async fn next(&mut self) -> Result<BenchmarkResult> {
258        let handle = self.handle()?.clone();
259        handle.check_cancelled()?;
260        match self.phase {
261            Phase::Provision => {
262                http::probe(handle.target(), Duration::from_secs(10))
263                    .await
264                    .context("endpoint probe failed — check the target URL and port")?;
265                let artifacts = self
266                    .artifacts
267                    .clone()
268                    .context("artifacts were not provisioned")?;
269                self.samples = dataset::load_shard(&artifacts.dataset, &DrawSpec::golden(), None)?;
270                if self.sample_cap > 0 {
271                    self.samples.truncate(self.sample_cap);
272                }
273                self.phase = Phase::Generate;
274                let n = self.samples.len();
275                let mut frame = BenchmarkResult::running("draw", self.elapsed())
276                    .with_progress(0, self.total_steps())
277                    .log_line(LogLine::info(format!(
278                        "{n} samples x {} orders = {} generations against one server",
279                        self.orders,
280                        n * self.orders
281                    )));
282                if self.sample_cap > 0 {
283                    frame = frame.log_line(LogLine::warn(format!(
284                        "capped at {n} of the golden 995 — the cap TRUNCATES, and the draw \
285                         concatenates in sorted subset order, so this selects which subsets \
286                         are compared rather than a random subsample of them; say the cap \
287                         when you report a green"
288                    )));
289                }
290                Ok(frame)
291            }
292            Phase::Generate => {
293                let n = self.samples.len();
294                if self.cursor >= n {
295                    // 2026-09-26: One order finished; bank it and start the
296                    // next.
297                    let label = self.order_label(self.order_index);
298                    self.runs.push(OrderRun {
299                        label: label.clone(),
300                        observations: std::mem::take(&mut self.current),
301                    });
302                    self.order_index += 1;
303                    self.cursor = 0;
304                    if self.order_index >= self.orders {
305                        self.phase = Phase::Compare;
306                    }
307                    return Ok(BenchmarkResult::running("order done", self.elapsed())
308                        .with_progress(self.done_steps(), self.total_steps())
309                        .log_line(LogLine::info(format!("order `{label}` complete"))));
310                }
311                let index = permutation(self.order_index, n)[self.cursor];
312                self.issue_one(index).await?;
313                self.cursor += 1;
314                let label = self.order_label(self.order_index);
315                handle.progress(self.done_steps(), self.total_steps());
316                handle.status(format!("{label} · {}/{n}", self.cursor));
317                Ok(BenchmarkResult::running(label, self.elapsed())
318                    .with_progress(self.done_steps(), self.total_steps()))
319            }
320            Phase::Compare => {
321                let s = score(&self.runs);
322                let v = verdict(&s);
323                self.phase = Phase::Done;
324                let total = self.total_steps();
325                Ok(BenchmarkResult {
326                    status: RunStatus::Completed,
327                    ..BenchmarkResult::running("done", self.elapsed())
328                }
329                .with_progress(total, total)
330                .with_summary(super::report::summary(&s))
331                .with_table(super::report::table(&s, &self.runs))
332                .with_metrics(super::report::metrics(&s))
333                .with_verdict(v))
334            }
335            Phase::Done => Ok(BenchmarkResult {
336                status: RunStatus::Completed,
337                ..BenchmarkResult::running("done", self.elapsed())
338            }),
339        }
340    }
341}
metrale_bench/benchmarks/kat_equality/driver.rs

metrale_bench/benchmarks/kat_equality/
driver.rs
