metrale_bench/benchmarks/scheduler_equivalence/
driver.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The scheduler-equivalence driver: serve each router under each
4//! lane, issue the draw at each concurrency, and hand the legs to `compare`.
5//!
6//! Owner: bench, scheduler-equivalence gate.
7//! Invariants:
8//! - `plan` serves each (router, lane) variant once, and a lane's sync legs
9//!   (reference, then control) run before its async serve.
10//! - A failed request becomes its reply's outcome (`request`); it does not
11//!   end the pass.
12
13use std::collections::BTreeMap;
14use std::future::Future;
15use std::sync::Arc;
16use std::sync::atomic::{AtomicUsize, Ordering};
17use std::time::{Duration, Instant};
18
19use anyhow::{Context, Result};
20use futures::StreamExt;
21use serde_json::json;
22
23use crate::benchmark::{Benchmark, BenchmarkDescriptor};
24use crate::benchmarks::bfcl::{MAX_NEW_TOKENS, dataset, draw::DrawSpec, provision};
25use crate::benchmarks::transcript::Transcript;
26use crate::hardware::Sensitivity;
27use crate::http;
28use crate::metadata::PluginMetadata;
29use crate::params::{ParamKind, ParamSpec, ParamValue, ParamValues};
30use crate::plugin::{Plugin, PluginHandle, TargetEndpoint};
31use crate::result::{BenchmarkResult, LogLine, RunStatus};
32
33use super::compare::{Leg, Pass, Reply};
34use super::host::{Lane, Router, RouterHost, ServeVariant};
35
36const SUMMARY: &str = "The async router must answer every sample exactly as the sync router does";
37
38/// 2026-09-26: A prefix of the golden draw: `next` truncates the loaded samples
39/// to it and logs a warning that names the cap.
40const DEFAULT_SAMPLE_CAP: usize = 64;
41const DEFAULT_CONCURRENCIES: [i64; 3] = [1, 4, 16];
42
43pub const METADATA: PluginMetadata = PluginMetadata::metrale(SUMMARY);
44
45pub const DESCRIPTOR: BenchmarkDescriptor = BenchmarkDescriptor {
46    id: "scheduler-equivalence",
47    name: "Scheduler Equivalence Gate",
48    summary: SUMMARY,
49    detail: "Serves the checkpoint under --scheduler-config sync and async, once per \
50             speculation lane (speculation off, then --mtp-gate force), issues one BFCL draw \
51             against each at concurrency 1, 4 and 16, and requires every sample_id's reply to \
52             be byte-identical between the routers — text, reasoning, tool calls with raw \
53             arguments, finish reason and completion token count — per concurrency. The sync \
54             serve is measured twice; a control that diverges makes an async difference \
55             unattributable and the run UNPROVEN rather than green. A failed request is \
56             UNMEASURED, never agreement, and an all-empty cell fails as VACUOUS. The async \
57             router's KV-reservation drift is recorded as a diagnostic outside equality.",
58    duration_hint: "~9 draw passes + 4 model loads",
59    expected_secs: 5400,
60    updated: "2026-09-25",
61    needs_confirmation: false,
62    // 2026-09-26: Not pinned to a checkpoint: the host re-serves whichever
63    // checkpoint is being served.
64    intended_for: None,
65    threshold_params: &[],
66    // 2026-09-26: The verdict reads byte equality and token counts, never
67    // timing.
68    sensitivity: Sensitivity::Correctness,
69    ctor: || Box::new(SchedulerEquivalence::default()),
70};
71
72/// 2026-09-26: One unit of the run; `plan` orders them.
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub enum Step {
75    Serve(ServeVariant),
76    /// 2026-09-26: Read the async counters before / after a lane's async legs.
77    Diagnostics {
78        lane: Lane,
79        before: bool,
80    },
81    Generate {
82        lane: Lane,
83        pass: Pass,
84        concurrency: usize,
85    },
86}
87
88/// 2026-09-26: The run, as a sequence: per lane, the sync serve (reference legs,
89/// then the control legs on the same serve), then the async serve with its
90/// legs between two diagnostics reads. Each variant is served once.
91pub fn plan(lanes: &[Lane], concurrencies: &[usize], control: bool) -> Vec<Step> {
92    let mut steps = Vec::new();
93    for &lane in lanes {
94        steps.push(Step::Serve(ServeVariant {
95            router: Router::Sync,
96            lane,
97        }));
98        for &c in concurrencies {
99            steps.push(Step::Generate {
100                lane,
101                pass: Pass::Sync,
102                concurrency: c,
103            });
104        }
105        if control {
106            for &c in concurrencies {
107                steps.push(Step::Generate {
108                    lane,
109                    pass: Pass::Control,
110                    concurrency: c,
111                });
112            }
113        }
114        steps.push(Step::Serve(ServeVariant {
115            router: Router::Async,
116            lane,
117        }));
118        steps.push(Step::Diagnostics { lane, before: true });
119        for &c in concurrencies {
120            steps.push(Step::Generate {
121                lane,
122                pass: Pass::Async,
123                concurrency: c,
124            });
125        }
126        steps.push(Step::Diagnostics {
127            lane,
128            before: false,
129        });
130    }
131    steps
132}
133
134pub fn lanes_for(choice: &str) -> Result<Vec<Lane>> {
135    Ok(match choice {
136        "both" => vec![Lane::SpecOff, Lane::MtpForce],
137        "spec-off" => vec![Lane::SpecOff],
138        "mtp-force" => vec![Lane::MtpForce],
139        other => anyhow::bail!("lanes: {other:?} is not one of both, spec-off, mtp-force"),
140    })
141}
142
143pub struct SchedulerEquivalence {
144    handle: Option<PluginHandle>,
145    host: Option<Arc<dyn RouterHost>>,
146    artifacts: Option<provision::Artifacts>,
147    samples: Vec<dataset::Sample>,
148    steps: Vec<Step>,
149    cursor: usize,
150    target: Option<TargetEndpoint>,
151    legs: Vec<Leg>,
152    diag_before: BTreeMap<String, f64>,
153    diagnostics: BTreeMap<String, f64>,
154    done: bool,
155    sample_cap: usize,
156    max_new_tokens: usize,
157    concurrencies: Vec<usize>,
158    lanes: Vec<Lane>,
159    control: bool,
160    request_timeout: Duration,
161    started: Option<Instant>,
162}
163
164impl Default for SchedulerEquivalence {
165    fn default() -> Self {
166        Self {
167            handle: None,
168            host: None,
169            artifacts: None,
170            samples: Vec::new(),
171            steps: Vec::new(),
172            cursor: 0,
173            target: None,
174            legs: Vec::new(),
175            diag_before: BTreeMap::new(),
176            diagnostics: BTreeMap::new(),
177            done: false,
178            sample_cap: DEFAULT_SAMPLE_CAP,
179            max_new_tokens: MAX_NEW_TOKENS,
180            concurrencies: DEFAULT_CONCURRENCIES.iter().map(|c| *c as usize).collect(),
181            lanes: vec![Lane::SpecOff, Lane::MtpForce],
182            control: true,
183            request_timeout: Duration::from_secs(300),
184            started: None,
185        }
186    }
187}
188
189impl SchedulerEquivalence {
190    fn handle(&self) -> Result<&PluginHandle> {
191        self.handle.as_ref().context("plugin was not loaded")
192    }
193
194    fn host(&self) -> Result<Arc<dyn RouterHost>> {
195        self.host.clone().context(super::host::NO_HOST)
196    }
197
198    fn elapsed(&self) -> Duration {
199        self.started.map(|s| s.elapsed()).unwrap_or_default()
200    }
201
202    fn total_steps(&self) -> u64 {
203        self.steps.len() as u64
204    }
205
206    /// 2026-09-26: The same body the KAT equality gate sends
207    /// (`kat_equality/driver.rs`): streamed with usage, temperature 0, and the
208    /// sample's messages, tools and tool_choice.
209    fn body(&self, sample: &dataset::Sample, model: &str) -> serde_json::Value {
210        json!({
211            "model": model,
212            "stream": true,
213            "stream_options": {"include_usage": true},
214            "temperature": 0.0,
215            "max_tokens": self.max_new_tokens,
216            "messages": sample.messages,
217            "tools": sample.tools,
218            "tool_choice": sample.tool_choice,
219        })
220    }
221
222    /// 2026-09-26: Issue the whole draw with `concurrency` requests in flight.
223    async fn generate(&self, target: &TargetEndpoint, concurrency: usize) -> Result<Vec<Reply>> {
224        let handle = self.handle()?.clone();
225        let bodies: Vec<(String, serde_json::Value)> = self
226            .samples
227            .iter()
228            .map(|s| (s.sample_id.clone(), self.body(s, &target.model)))
229            .collect();
230        let total = bodies.len();
231        let finished = Arc::new(AtomicUsize::new(0));
232        let timeout = self.request_timeout;
233        let mut replies: Vec<Reply> = futures::stream::iter(bodies)
234            .map(|(sample_id, body)| {
235                let handle = handle.clone();
236                let finished = Arc::clone(&finished);
237                async move {
238                    let reply = request(target, sample_id, &body, timeout).await;
239                    let n = finished.fetch_add(1, Ordering::Relaxed) + 1;
240                    handle.status(format!("C={concurrency} · {n}/{total}"));
241                    reply
242                }
243            })
244            .buffer_unordered(concurrency.max(1))
245            .collect()
246            .await;
247        // 2026-09-26: Completion order depends on the concurrency; sorted by
248        // id so a leg's order does not.
249        replies.sort_by(|a, b| a.sample_id.cmp(&b.sample_id));
250        Ok(replies)
251    }
252
253    fn compare(&mut self) -> BenchmarkResult {
254        self.done = true;
255        let total = self.total_steps();
256        super::report::terminal(&self.legs, &self.diagnostics, self.elapsed())
257            .with_progress(total, total)
258    }
259}
260
261/// 2026-09-26: One sample against one serve. Never fails: an error is the
262/// reply's outcome, classified, with the time it took to arrive.
263pub async fn request(
264    target: &TargetEndpoint,
265    sample_id: String,
266    body: &serde_json::Value,
267    timeout: Duration,
268) -> Reply {
269    let started = Instant::now();
270    let outcome = http::chat_stream(target, body, timeout)
271        .await
272        .map(|o| Box::new(Transcript::from(&o)))
273        .map_err(|e| http::classify(&e));
274    Reply {
275        sample_id,
276        outcome,
277        elapsed: started.elapsed(),
278    }
279}
280
281impl Plugin for SchedulerEquivalence {
282    fn metadata(&self) -> &'static PluginMetadata {
283        &METADATA
284    }
285
286    fn load(&mut self, handle: PluginHandle) -> impl Future<Output = Result<()>> + Send {
287        self.started = Some(Instant::now());
288        self.handle = Some(handle.clone());
289        // 2026-09-26: The host check precedes `provision::ensure`, so a process
290        // with no host fails before any download.
291        let host = super::host::installed();
292        async move {
293            let host = host.context(super::host::NO_HOST)?;
294            self.host = Some(host);
295            let artifacts = provision::ensure(handle.artifacts(), &handle).await?;
296            self.artifacts = Some(artifacts);
297            Ok(())
298        }
299    }
300}
301
302impl Benchmark for SchedulerEquivalence {
303    fn descriptor(&self) -> &'static BenchmarkDescriptor {
304        &DESCRIPTOR
305    }
306
307    fn parameters(&self) -> Vec<ParamSpec> {
308        vec![
309            ParamSpec::new(
310                "sample_cap",
311                "Sample cap",
312                "Truncate the draw to this many samples (0 = the whole draw). A PREFIX of the \
313                 sorted-subset concatenation, not a random sample; say the cap when you \
314                 report a green.",
315                ParamKind::Int { min: 0, max: 5000 },
316                ParamValue::Int(DEFAULT_SAMPLE_CAP as i64),
317            ),
318            ParamSpec::new(
319                "max_new_tokens",
320                "Max new tokens",
321                "Generation cap per sample. BFCL's own budget, so the gate answers about the \
322                 regime the accuracy gates measure.",
323                ParamKind::Int { min: 32, max: 4096 },
324                ParamValue::Int(MAX_NEW_TOKENS as i64),
325            ),
326            ParamSpec::new(
327                "concurrencies",
328                "Concurrencies",
329                "Requests in flight per pass. A difference at C >= 4 that is absent at C = 1 \
330                 is a batch-shape difference, not a feed one — the report keeps them apart.",
331                ParamKind::IntList { min: 1, max: 128 },
332                ParamValue::IntList(DEFAULT_CONCURRENCIES.to_vec()),
333            ),
334            ParamSpec::new(
335                "lanes",
336                "Speculation lanes",
337                "Which lane pins to run: both (speculation off, then --mtp-gate force), or one.",
338                ParamKind::Choice(&["both", "spec-off", "mtp-force"]),
339                ParamValue::Text("both".to_string()),
340            ),
341            ParamSpec::new(
342                "control",
343                "Sync-vs-sync control",
344                "Measure the sync serve a second time per lane. Off only for a quick look: \
345                 without it an async difference cannot be attributed.",
346                ParamKind::Bool,
347                ParamValue::Bool(true),
348            ),
349            ParamSpec::new(
350                "request_timeout_s",
351                "Request timeout (s)",
352                "Per-request timeout. A timeout is UNMEASURED, which fails the gate.",
353                ParamKind::Int { min: 10, max: 3600 },
354                ParamValue::Int(300),
355            ),
356        ]
357    }
358
359    fn configure(&mut self, values: &ParamValues) -> Result<()> {
360        let specs = self.parameters();
361        values.validate_against(&specs)?;
362        self.sample_cap = values.usize("sample_cap")?;
363        self.max_new_tokens = values.usize("max_new_tokens")?;
364        self.concurrencies = values
365            .int_list("concurrencies")?
366            .iter()
367            .map(|c| *c as usize)
368            .collect();
369        self.lanes = lanes_for(values.text("lanes")?)?;
370        self.control = values.bool("control")?;
371        self.request_timeout = Duration::from_secs(values.usize("request_timeout_s")? as u64);
372        // 2026-09-26: A re-`configure` starts over: no leg of an earlier run is
373        // compared as this run's.
374        self.steps.clear();
375        self.cursor = 0;
376        self.legs.clear();
377        self.diagnostics.clear();
378        self.target = None;
379        self.done = false;
380        Ok(())
381    }
382
383    async fn next(&mut self) -> Result<BenchmarkResult> {
384        let handle = self.handle()?.clone();
385        handle.check_cancelled()?;
386        if self.done {
387            return Ok(BenchmarkResult {
388                status: RunStatus::Completed,
389                ..BenchmarkResult::running("done", self.elapsed())
390            });
391        }
392        if self.steps.is_empty() {
393            let artifacts = self
394                .artifacts
395                .clone()
396                .context("artifacts were not provisioned")?;
397            self.samples = dataset::load_shard(&artifacts.dataset, &DrawSpec::golden(), None)?;
398            if self.sample_cap > 0 {
399                self.samples.truncate(self.sample_cap);
400            }
401            self.steps = plan(&self.lanes, &self.concurrencies, self.control);
402            let n = self.samples.len();
403            let passes = self
404                .steps
405                .iter()
406                .filter(|s| matches!(s, Step::Generate { .. }))
407                .count();
408            let mut frame = BenchmarkResult::running("draw", self.elapsed())
409                .with_progress(0, self.total_steps())
410                .log_line(LogLine::info(format!(
411                    "{n} samples x {passes} passes ({} lanes x {} concurrencies x {} routers) \
412                     = {} generations, {} model loads",
413                    self.lanes.len(),
414                    self.concurrencies.len(),
415                    if self.control { 3 } else { 2 },
416                    n * passes,
417                    2 * self.lanes.len()
418                )));
419            if self.sample_cap > 0 {
420                frame = frame.log_line(LogLine::warn(format!(
421                    "capped at {n} samples of the golden draw — a prefix, not a sample; \
422                     say the cap when you report a green"
423                )));
424            }
425            return Ok(frame);
426        }
427        let Some(step) = self.steps.get(self.cursor).copied() else {
428            return Ok(self.compare());
429        };
430        let label = match step {
431            Step::Serve(v) => {
432                handle.status(format!("serving {}", v.label()));
433                let target = self.host()?.serve(v).await.with_context(|| {
434                    format!("could not serve the checkpoint under {}", v.label())
435                })?;
436                http::probe(&target, Duration::from_secs(10))
437                    .await
438                    .context("the re-served endpoint does not answer")?;
439                self.target = Some(target);
440                format!("serve {}", v.label())
441            }
442            Step::Diagnostics { lane, before } => {
443                let now = self.host()?.diagnostics()?;
444                if before {
445                    self.diag_before = now;
446                } else {
447                    for (k, v) in now {
448                        let base = self.diag_before.get(&k).copied().unwrap_or(0.0);
449                        self.diagnostics
450                            .insert(format!("{}_{k}", lane.label()), v - base);
451                    }
452                }
453                format!("diagnostics {}", lane.label())
454            }
455            Step::Generate {
456                lane,
457                pass,
458                concurrency,
459            } => {
460                let target = self.target.clone().context("no serve preceded this pass")?;
461                let replies = self.generate(&target, concurrency).await?;
462                self.legs.push(Leg {
463                    lane,
464                    pass,
465                    concurrency,
466                    replies,
467                });
468                format!("{} {} C={concurrency}", lane.label(), pass.label())
469            }
470        };
471        self.cursor += 1;
472        handle.progress(self.cursor as u64, self.total_steps());
473        Ok(BenchmarkResult::running(label.clone(), self.elapsed())
474            .with_progress(self.cursor as u64, self.total_steps())
475            .log_line(LogLine::info(format!("{label} complete"))))
476    }
477
478    async fn cleanup(&mut self) -> Result<()> {
479        match self.host.as_ref() {
480            Some(host) => host.restore().await,
481            None => Ok(()),
482        }
483    }
484}
metrale_bench/benchmarks/scheduler_equivalence/driver.rs

metrale_bench/benchmarks/scheduler_equivalence/
driver.rs
