use anyhow::{Context, Result, bail};
use inferencelayer::EmbedEngine;
use inferencelayer::bench_guard::{
ensure_quiet_machine, foreign_load, load_average_1m, read_pins, self_cpu_seconds, write_pins,
};
use inferencelayer::encoder_weights::EncBatch;
use inferencelayer::gliner::{Gliner, GlinerDevice};
use inferencelayer::pooling::EmbedOut;
use std::time::Instant;
const HOLD_FRACTION: f64 = 0.97;
const CROWN_TARGET: f64 = 0.90;
const ENT_TOKEN_ID: u32 = 128002;
const SENTENCE: &str = "Barack Obama was born in Hawaii and later worked in Chicago.";
const LABELS_2: [&str; 2] = ["person", "location"];
const LABELS_20: [&str; 20] = [
"person",
"location",
"organization",
"date",
"product",
"event",
"disease",
"drug",
"gene",
"protein",
"company",
"city",
"country",
"award",
"job title",
"language",
"nationality",
"book",
"film",
"law",
];
fn main() {
match run(&std::env::args().skip(1).collect::<Vec<_>>()) {
Ok(true) => {}
Ok(false) => std::process::exit(1),
Err(e) => {
eprintln!("encoder-scoreboard: {e:#}");
std::process::exit(2);
}
}
}
struct Opts {
pins_path: String,
write_pins: bool,
ingest_torch: Option<String>,
records: Vec<(String, f64)>,
rows: Vec<String>,
reps: usize,
devices: Vec<Device>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Device {
Cpu,
Wgpu,
}
impl Device {
fn tag(self) -> &'static str {
match self {
Device::Cpu => "cpu",
Device::Wgpu => "wgpu",
}
}
}
fn parse_opts(args: &[String]) -> Result<Opts> {
let mut o = Opts {
pins_path: "bench/encoder_scoreboard_pins.json".into(),
write_pins: false,
ingest_torch: None,
records: vec![],
rows: vec![],
reps: 30,
devices: vec![Device::Cpu, Device::Wgpu],
};
let mut it = args.iter();
while let Some(a) = it.next() {
let mut val = |name: &str| {
it.next()
.cloned()
.with_context(|| format!("{name} needs a value"))
};
match a.as_str() {
"--pins" => o.pins_path = val("--pins")?,
"--pin" => o.write_pins = true,
"--ingest-torch" => o.ingest_torch = Some(val("--ingest-torch")?),
"--rows" => o.rows = val("--rows")?.split(',').map(str::to_string).collect(),
"--reps" => o.reps = val("--reps")?.parse().context("--reps")?,
"--devices" => {
o.devices = val("--devices")?
.split(',')
.map(|d| match d {
"cpu" => Ok(Device::Cpu),
"wgpu" => Ok(Device::Wgpu),
other => bail!("unknown device {other:?} (cpu|wgpu)"),
})
.collect::<Result<_>>()?;
}
"--record" => {
let kv = val("--record")?;
let (k, v) = kv
.split_once('=')
.with_context(|| format!("--record wants row=value_ms, got {kv:?}"))?;
o.records
.push((k.to_string(), v.parse().context("--record value")?));
}
other => bail!("unknown flag {other:?}"),
}
}
Ok(o)
}
fn run(args: &[String]) -> Result<bool> {
let o = parse_opts(args)?;
let quiet = ensure_quiet_machine()?;
if o.write_pins && !quiet {
bail!(
"--pin refused: the quiet-machine guard was overridden, so these numbers carry \
contention and must never become the reference every later run is judged against."
)
}
let mut pins = read_pins(&o.pins_path)?;
if let Some(path) = &o.ingest_torch {
ingest_torch(&mut pins, path, quiet)?;
write_pins(&o.pins_path, &pins)?;
eprintln!("ingested torch baselines from {path} → {}", o.pins_path);
return Ok(true);
}
let load_before = load_average_1m();
let t_start = std::time::Instant::now();
let cpu_start = self_cpu_seconds();
let mut measured: Vec<(String, f64)> = o.records.clone();
for dev in &o.devices {
measured.extend(measure_device(*dev, &o)?);
}
if measured.is_empty() {
bail!(
"nothing measured: set OSFKB_GLINER_DIR / OSFKB_BIOLORD_DIR / OSFKB_CE_DIR, or pass \
--record <row>=<ms>"
);
}
let quiet = quiet && still_quiet(load_before, t_start, cpu_start);
if o.write_pins && !quiet {
bail!(
"--pin refused: the machine did not stay quiet for the whole sweep, so these numbers \
carry contention and must never become the reference every later run is judged against."
)
}
let ok = report_and_gate(&mut pins, &measured, &o, quiet)?;
if !quiet {
eprintln!(
"\nEXPLORATORY: the machine was not quiet for the whole sweep — the rows above are \
indicative, and the gates are NOT enforced. Re-run on a quiet box before believing a \
REGRESSION or a CROWN."
);
return Ok(true);
}
Ok(ok)
}
fn still_quiet(before: Option<f64>, t_start: std::time::Instant, cpu_start: Option<f64>) -> bool {
let cores = std::thread::available_parallelism()
.map(|c| c.get() as f64)
.unwrap_or(8.0);
let limit = cores / 2.0;
let self_cores = match (cpu_start, self_cpu_seconds()) {
(Some(a), Some(b)) => (b - a) / t_start.elapsed().as_secs_f64().max(1e-3),
_ => 0.0,
};
let Some(after) = foreign_load(self_cores) else {
return true; };
if after >= limit {
eprintln!(
"quiet-machine guard: FOREIGN load rose to {after:.2} (≥ {limit:.1}; we were using \
{self_cores:.1} cores) — the later rows were measured on a busier box than the earlier \
ones."
);
return false;
}
if let Some(b) = before
&& after > b * 1.75
&& after > 2.0
{
eprintln!(
"quiet-machine guard: foreign load drifted {b:.2} → {after:.2} during the sweep — the \
rows are not comparable to each other, let alone to a pin."
);
return false;
}
eprintln!(
"quiet-machine guard: foreign load {after:.2} at the end (we used {self_cores:.1} cores) — ok"
);
true
}
fn measure_device(dev: Device, o: &Opts) -> Result<Vec<(String, f64)>> {
let mut out = Vec::new();
let want = |row: &str| o.rows.is_empty() || o.rows.iter().any(|r| r == row);
if let Some(dir) = env_dir("OSFKB_GLINER_DIR") {
let gdev = match dev {
Device::Cpu => GlinerDevice::Cpu,
Device::Wgpu => GlinerDevice::Auto,
};
let mut g = Gliner::load_on(&dir, ENT_TOKEN_ID, gdev)
.with_context(|| format!("load gliner from {}", dir.display()))?;
if dev == Device::Wgpu && g.device() == "cpu" {
eprintln!("SKIP gliner wgpu rows: no usable adapter (auto fell back to cpu)");
} else {
eprintln!("gliner on {} ({})", g.device(), dev.tag());
for (reps, words) in [(1usize, 11usize), (4, 44), (16, 176)] {
let text = vec![SENTENCE; reps].join(" ");
for (nlab, labels) in [
(
2usize,
LABELS_2.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
),
(20, LABELS_20.iter().map(|s| s.to_string()).collect()),
] {
let row = format!("gliner_{words}w_{nlab}lab.{}", dev.tag());
if !want(&row) {
continue;
}
let ms = min_of_n(o.reps, || g.predict_text(&text, &labels, 0.5).map(|_| ()))?;
out.push((row, ms));
}
}
}
} else {
eprintln!("SKIP gliner rows: OSFKB_GLINER_DIR unset");
}
if let Some(dir) = env_dir("OSFKB_BIOLORD_DIR") {
let row = format!("biolord_embed64.{}", dev.tag());
if want(&row) {
match load_engine(dev, &dir)? {
None => eprintln!("SKIP {row}: no usable adapter"),
Some(mut e) => {
eprintln!("biolord on {} ({})", e.device(), dev.tag());
let ids: Vec<u32> = (0..64).map(|i| 5 + (i as u32 * 7) % 2000).collect();
let batch = EncBatch::from_seqs([ids]);
let ms = min_of_n(o.reps, || e.encode(&batch).map(|_| ()))?;
out.push((row, ms));
}
}
}
} else {
eprintln!("SKIP biolord row: OSFKB_BIOLORD_DIR unset");
}
if let Some(dir) = env_dir("OSFKB_CE_DIR") {
let row = format!("ce_score192.{}", dev.tag());
if want(&row) {
if let Some(mut enc) = load_engine(dev, &dir)? {
eprintln!("cross-encoder on {} ({})", enc.device(), dev.tag());
let a: Vec<u32> = (0..96).map(|i| 5 + (i as u32 * 11) % 1500).collect();
let b: Vec<u32> = (0..96).map(|i| 7 + (i as u32 * 13) % 1500).collect();
let batch = EncBatch::from_pairs(vec![(
a.iter().chain(&b).copied().collect::<Vec<u32>>(),
a.iter()
.map(|_| 0u32)
.chain(b.iter().map(|_| 1u32))
.collect::<Vec<u32>>(),
)]);
let ms = min_of_n(o.reps, || {
enc.encode(&batch).map(|o| match o {
EmbedOut::Pooled(_) | EmbedOut::PerToken(_) => (),
})
})?;
out.push((row, ms));
} else {
eprintln!("SKIP {row}: no GPU adapter (would have measured the CPU path)");
}
}
} else {
eprintln!("SKIP ce row: OSFKB_CE_DIR unset");
}
Ok(out)
}
fn load_engine(dev: Device, dir: &std::path::Path) -> Result<Option<EmbedEngine>> {
let e = match dev {
Device::Cpu => EmbedEngine::cpu(dir)?,
Device::Wgpu => {
let e = EmbedEngine::auto(dir, 8192)?;
if e.device() == "cpu" {
return Ok(None);
}
e
}
};
Ok(Some(e))
}
fn env_dir(key: &str) -> Option<std::path::PathBuf> {
let p = std::path::PathBuf::from(std::env::var(key).ok()?);
if p.is_dir() {
Some(p)
} else {
eprintln!("{key} = {} is not a directory", p.display());
None
}
}
fn min_of_n(reps: usize, mut f: impl FnMut() -> Result<()>) -> Result<f64> {
for _ in 0..5 {
f()?;
}
let mut best = f64::MAX;
for _ in 0..reps.max(1) {
let t0 = Instant::now();
f()?;
best = best.min(t0.elapsed().as_secs_f64());
}
Ok(best * 1e3)
}
fn ingest_torch(pins: &mut serde_json::Value, path: &str, quiet: bool) -> Result<()> {
let raw = std::fs::read(path).with_context(|| format!("read {path}"))?;
let doc: serde_json::Value =
serde_json::from_slice(&raw).with_context(|| format!("parse {path}"))?;
let rows = doc
.get("rows")
.and_then(|r| r.as_object())
.context("torch json: expected a top-level `rows` object")?;
let _ = quiet;
if let Some(load) = doc.get("load_1m").and_then(serde_json::Value::as_f64) {
let cores = std::thread::available_parallelism()
.map(std::num::NonZero::get)
.unwrap_or(8) as f64;
let overridden = std::env::var("OSFKB_SCOREBOARD_IGNORE_LOAD").as_deref() == Ok("1");
if load >= cores / 2.0 {
if !overridden {
bail!(
"torch baseline was captured at load {load:.2} on a {cores:.0}-core box — a \
contended torch number would LOWER the bar we claim to clear. Re-run \
bench_encoders.py on a quiet box (or set OSFKB_SCOREBOARD_IGNORE_LOAD=1 to \
ingest it as exploratory)."
)
}
eprintln!(
"WARNING: torch baseline captured at load {load:.2} — exploratory, not a crown."
);
}
}
let machine = doc
.get("machine")
.cloned()
.unwrap_or(serde_json::Value::Null);
let baselines = pins
.as_object_mut()
.context("pins root must be an object")?
.entry("baselines")
.or_insert_with(|| serde_json::json!({}));
for (row, v) in rows {
let value = v
.get("value_ms")
.and_then(serde_json::Value::as_f64)
.with_context(|| format!("torch row {row}: missing value_ms"))?;
baselines[row] = serde_json::json!({
"value_ms": value,
"by_device": v.get("by_device").cloned().unwrap_or(serde_json::Value::Null),
"machine": machine,
"load_1m": doc.get("load_1m").cloned().unwrap_or(serde_json::Value::Null),
"torch": doc.get("torch").cloned().unwrap_or(serde_json::Value::Null),
"device": v.get("device").cloned().unwrap_or(serde_json::Value::Null),
"threads": doc.get("threads").cloned().unwrap_or(serde_json::Value::Null),
"captured_epoch_s": doc.get("captured_epoch_s").cloned().unwrap_or(serde_json::Value::Null),
"method": "min-of-n",
});
eprintln!(" baseline {row:30} {value:8.1} ms");
}
Ok(())
}
fn torch_bar(b: &serde_json::Value, engine_dev: &str) -> Option<f64> {
let by_dev = b.get("by_device").and_then(serde_json::Value::as_object);
if engine_dev == "cpu"
&& let Some(cpu) = by_dev
.and_then(|d| d.get("cpu"))
.and_then(serde_json::Value::as_f64)
{
return Some(cpu);
}
b.get("value_ms").and_then(serde_json::Value::as_f64)
}
fn report_and_gate(
pins: &mut serde_json::Value,
measured: &[(String, f64)],
o: &Opts,
quiet: bool,
) -> Result<bool> {
let load = load_average_1m();
let baselines = pins
.get("baselines")
.cloned()
.unwrap_or(serde_json::json!({}));
let mut ok = true;
let mut aspirational = 0usize;
println!(
"\n{:32} {:>9} {:>9} {:>9} {:>7} gate",
"row", "measured", "pin", "torch", "ratio"
);
for (row, ms) in measured {
let pin = pins
.get("rows")
.and_then(|r| r.get(row))
.and_then(|p| p.get("value_ms"))
.and_then(serde_json::Value::as_f64);
let hold_ok = pin.is_none_or(|p| *ms <= p / HOLD_FRACTION);
let (base_key, dev) = row
.rsplit_once('.')
.map_or((row.as_str(), ""), |(k, d)| (k, d));
let torch = baselines.get(base_key).and_then(|b| torch_bar(b, dev));
let ratio = torch.map(|t| ms / t);
let ratio_max = pins
.get("ratios")
.and_then(|r| r.get(row))
.and_then(serde_json::Value::as_f64);
let crown_ok = match (ratio, ratio_max) {
(Some(r), Some(max)) if max <= 1.0 => r <= max,
(Some(_), _) => {
aspirational += 1;
true
}
(None, _) => true,
};
ok &= hold_ok && crown_ok;
let gate = match (hold_ok, crown_ok, ratio, ratio_max) {
(false, _, _, _) => "REGRESSION".to_string(),
(_, false, Some(r), Some(m)) => format!("CROWN LOST ({r:.2} > {m:.2})"),
(_, _, Some(r), Some(m)) if m <= 1.0 && r <= CROWN_TARGET => "CROWN ✓".into(),
(_, _, Some(_r), Some(m)) if m <= 1.0 => format!("holding (target {CROWN_TARGET:.2})"),
(_, _, Some(r), _) if r > 1.0 => "ASPIRATIONAL (behind torch)".into(),
(_, _, Some(_), _) => "ASPIRATIONAL (ahead — pin it)".into(),
_ => "no torch baseline".into(),
};
println!(
"{row:32} {ms:9.1} {:>9} {:>9} {:>7} {gate}",
pin.map_or("—".into(), |p| format!("{p:.1}")),
torch.map_or("—".into(), |t| format!("{t:.1}")),
ratio.map_or("—".into(), |r| format!("{r:.2}×")),
);
if o.write_pins && (pin.is_none() || *ms < pin.unwrap()) {
pins["rows"][row] = serde_json::json!({
"value_ms": ms,
"load_1m": load,
"quiet": quiet,
"captured_epoch_s": epoch_s(),
});
if let Some(r) = ratio {
let ratcheted = ratio_max.map_or(r, |m| m.min(r)).max(CROWN_TARGET);
pins["ratios"][row] = serde_json::json!(ratcheted);
}
}
}
if o.write_pins {
write_pins(&o.pins_path, pins)?;
eprintln!("\npins written to {}", o.pins_path);
}
if aspirational > 0 {
eprintln!(
"{aspirational} row(s) ASPIRATIONAL — no crown pinned yet (nonfatal). Campaign target: \
every row ≤ {CROWN_TARGET:.2}× torch."
);
}
println!(
"\n{}",
match (quiet, ok) {
(false, _) => "EXPLORATORY — not a verdict (see the quiet-machine guard above)",
(true, true) => "PASS",
(true, false) => "FAIL",
}
);
Ok(ok)
}
fn epoch_s() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}