1use std::{
4 collections::BTreeMap,
5 future::Future,
6 num::NonZeroUsize,
7 pin::Pin,
8 sync::Arc,
9 time::{Duration, Instant},
10};
11
12use futures_util::{StreamExt, stream};
13use runifold_model::{
14 Model, ModelCallContext, ModelErrorKind, ModelRequest, ModelStreamAccumulator, ModelStreamEvent,
15};
16use serde::{Deserialize, Serialize};
17use thiserror::Error;
18
19const REPORT_SCHEMA_VERSION: u32 = 1;
20
21pub type BenchmarkFuture<'a> = Pin<Box<dyn Future<Output = BenchmarkInvocation> + Send + 'a>>;
23
24pub trait BenchmarkTarget: Send + Sync {
29 fn execute(&self) -> BenchmarkFuture<'_>;
31}
32
33#[derive(Clone)]
35pub struct ModelBenchmarkTarget {
36 model: Arc<dyn Model>,
37 request: ModelRequest,
38}
39
40impl ModelBenchmarkTarget {
41 pub fn new(model: Arc<dyn Model>, request: ModelRequest) -> Self {
43 Self { model, request }
44 }
45}
46
47impl std::fmt::Debug for ModelBenchmarkTarget {
48 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49 formatter
50 .debug_struct("ModelBenchmarkTarget")
51 .field("model", &self.request.model)
52 .finish_non_exhaustive()
53 }
54}
55
56impl BenchmarkTarget for ModelBenchmarkTarget {
57 fn execute(&self) -> BenchmarkFuture<'_> {
58 Box::pin(run_model(self.model.as_ref(), self.request.clone()))
59 }
60}
61
62#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
64#[serde(rename_all = "snake_case")]
65#[non_exhaustive]
66pub enum BenchmarkFailureKind {
67 InvalidRequest,
69 UnsupportedFeature,
71 Transport,
73 Protocol,
75 Stream,
77 Provider,
79 Cancelled,
81 DeadlineExceeded,
83 Other,
85}
86
87impl From<&ModelErrorKind> for BenchmarkFailureKind {
88 fn from(kind: &ModelErrorKind) -> Self {
89 match kind {
90 ModelErrorKind::InvalidRequest => Self::InvalidRequest,
91 ModelErrorKind::UnsupportedFeature => Self::UnsupportedFeature,
92 ModelErrorKind::Transport => Self::Transport,
93 ModelErrorKind::Protocol => Self::Protocol,
94 ModelErrorKind::StreamState | ModelErrorKind::MalformedToolArguments => Self::Stream,
95 ModelErrorKind::Provider => Self::Provider,
96 ModelErrorKind::Cancelled => Self::Cancelled,
97 ModelErrorKind::DeadlineExceeded => Self::DeadlineExceeded,
98 _ => Self::Other,
99 }
100 }
101}
102
103#[derive(Clone, Copy, Debug, Eq, PartialEq)]
105pub struct BenchmarkInvocation {
106 ttft: Option<Duration>,
107 total: Duration,
108 failure: Option<BenchmarkFailureKind>,
109}
110
111impl BenchmarkInvocation {
112 pub fn success(
119 ttft: Option<Duration>,
120 total: Duration,
121 ) -> Result<Self, BenchmarkInvocationError> {
122 if ttft.is_some_and(|ttft| ttft > total) {
123 return Err(BenchmarkInvocationError::TtftAfterCompletion);
124 }
125 Ok(Self {
126 ttft,
127 total,
128 failure: None,
129 })
130 }
131
132 pub const fn failure(kind: BenchmarkFailureKind, total: Duration) -> Self {
134 Self {
135 ttft: None,
136 total,
137 failure: Some(kind),
138 }
139 }
140}
141
142#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
144#[non_exhaustive]
145pub enum BenchmarkInvocationError {
146 #[error("time to first output cannot exceed total invocation time")]
148 TtftAfterCompletion,
149}
150
151#[derive(Clone, Debug, Eq, PartialEq)]
153pub struct BenchmarkPlan {
154 measured_runs: NonZeroUsize,
155 warmup_runs: usize,
156 concurrency: NonZeroUsize,
157 environment: BTreeMap<String, String>,
158}
159
160impl BenchmarkPlan {
161 pub const fn new(measured_runs: NonZeroUsize) -> Self {
163 Self {
164 measured_runs,
165 warmup_runs: 0,
166 concurrency: NonZeroUsize::MIN,
167 environment: BTreeMap::new(),
168 }
169 }
170
171 #[must_use]
173 pub const fn with_warmup(mut self, warmup_runs: usize) -> Self {
174 self.warmup_runs = warmup_runs;
175 self
176 }
177
178 #[must_use]
180 pub const fn with_concurrency(mut self, concurrency: NonZeroUsize) -> Self {
181 self.concurrency = concurrency;
182 self
183 }
184
185 #[must_use]
190 pub fn with_environment(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
191 self.environment.insert(key.into(), value.into());
192 self
193 }
194}
195
196#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
198pub struct LatencyDistribution {
199 pub min_us: u64,
201 pub p50_us: u64,
203 pub p95_us: u64,
205 pub p99_us: u64,
207 pub max_us: u64,
209}
210
211#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
213pub struct BenchmarkFailureCount {
214 pub kind: BenchmarkFailureKind,
216 pub count: usize,
218}
219
220#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
222pub struct ProviderBenchmarkReport {
223 pub schema_version: u32,
225 pub label: String,
227 pub measured_runs: usize,
229 pub concurrency: usize,
231 pub environment: BTreeMap<String, String>,
233 pub successes: usize,
235 pub failures: usize,
237 pub success_rate: f64,
239 pub wall_time_us: u64,
241 pub throughput_per_second: f64,
243 pub total_latency: Option<LatencyDistribution>,
245 pub ttft: Option<LatencyDistribution>,
247 pub successes_without_output: usize,
249 pub failure_counts: Vec<BenchmarkFailureCount>,
251}
252
253#[derive(Clone, Debug, Error, PartialEq)]
255#[non_exhaustive]
256pub enum ProviderBenchmarkError {
257 #[error("benchmark label cannot be empty")]
259 EmptyLabel,
260 #[error("benchmark regression ratio `{field}` must be finite and between zero and one")]
262 InvalidRegressionRatio {
263 field: &'static str,
265 },
266}
267
268pub async fn benchmark(
274 label: impl Into<String>,
275 target: Arc<dyn BenchmarkTarget>,
276 plan: BenchmarkPlan,
277) -> Result<ProviderBenchmarkReport, ProviderBenchmarkError> {
278 let label = label.into();
279 if label.trim().is_empty() {
280 return Err(ProviderBenchmarkError::EmptyLabel);
281 }
282 for _ in 0..plan.warmup_runs {
283 let _ = target.execute().await;
284 }
285
286 let started = Instant::now();
287 let outcomes = stream::iter(0..plan.measured_runs.get())
288 .map(|_| {
289 let target = Arc::clone(&target);
290 async move { target.execute().await }
291 })
292 .buffer_unordered(plan.concurrency.get())
293 .collect::<Vec<_>>()
294 .await;
295 let wall_time = started.elapsed();
296 Ok(build_report(label, &plan, wall_time, &outcomes))
297}
298
299#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
301pub struct BenchmarkRegressionPolicy {
302 pub max_success_rate_drop: f64,
304 pub max_throughput_drop: f64,
306 pub max_p95_ttft_increase: f64,
308 pub max_p95_total_latency_increase: f64,
310}
311
312impl BenchmarkRegressionPolicy {
313 pub fn new(
320 max_success_rate_drop: f64,
321 max_throughput_drop: f64,
322 max_p95_ttft_increase: f64,
323 max_p95_total_latency_increase: f64,
324 ) -> Result<Self, ProviderBenchmarkError> {
325 for (field, value) in [
326 ("max_success_rate_drop", max_success_rate_drop),
327 ("max_throughput_drop", max_throughput_drop),
328 ("max_p95_ttft_increase", max_p95_ttft_increase),
329 (
330 "max_p95_total_latency_increase",
331 max_p95_total_latency_increase,
332 ),
333 ] {
334 if !value.is_finite() || !(0.0..=1.0).contains(&value) {
335 return Err(ProviderBenchmarkError::InvalidRegressionRatio { field });
336 }
337 }
338 Ok(Self {
339 max_success_rate_drop,
340 max_throughput_drop,
341 max_p95_ttft_increase,
342 max_p95_total_latency_increase,
343 })
344 }
345}
346
347#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
349pub struct BenchmarkRegressionMetric {
350 pub name: String,
352 pub baseline: Option<f64>,
354 pub candidate: Option<f64>,
356 pub passed: bool,
358}
359
360#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
362pub struct BenchmarkRegressionComparison {
363 pub baseline: String,
365 pub candidate: String,
367 pub metrics: Vec<BenchmarkRegressionMetric>,
369 pub passed: bool,
371}
372
373pub fn compare_benchmarks(
375 baseline: &ProviderBenchmarkReport,
376 candidate: &ProviderBenchmarkReport,
377 policy: BenchmarkRegressionPolicy,
378) -> BenchmarkRegressionComparison {
379 let mut metrics = vec![
380 larger_is_better(
381 "success_rate",
382 Some(baseline.success_rate),
383 Some(candidate.success_rate),
384 baseline.success_rate - policy.max_success_rate_drop,
385 ),
386 larger_is_better(
387 "throughput_per_second",
388 Some(baseline.throughput_per_second),
389 Some(candidate.throughput_per_second),
390 baseline.throughput_per_second * (1.0 - policy.max_throughput_drop),
391 ),
392 smaller_is_better(
393 "p95_ttft_us",
394 baseline.ttft.map(|latency| u64_as_f64(latency.p95_us)),
395 candidate.ttft.map(|latency| u64_as_f64(latency.p95_us)),
396 policy.max_p95_ttft_increase,
397 ),
398 smaller_is_better(
399 "p95_total_latency_us",
400 baseline
401 .total_latency
402 .map(|latency| u64_as_f64(latency.p95_us)),
403 candidate
404 .total_latency
405 .map(|latency| u64_as_f64(latency.p95_us)),
406 policy.max_p95_total_latency_increase,
407 ),
408 ];
409 let passed = metrics.iter().all(|metric| metric.passed);
410 BenchmarkRegressionComparison {
411 baseline: baseline.label.clone(),
412 candidate: candidate.label.clone(),
413 metrics: std::mem::take(&mut metrics),
414 passed,
415 }
416}
417
418async fn run_model(model: &dyn Model, request: ModelRequest) -> BenchmarkInvocation {
419 let started = Instant::now();
420 let mut stream = match model.stream(request, ModelCallContext::new()).await {
421 Ok(stream) => stream,
422 Err(error) => {
423 return BenchmarkInvocation::failure((&error.kind).into(), started.elapsed());
424 }
425 };
426 let mut accumulator = ModelStreamAccumulator::new();
427 let mut ttft = None;
428 while let Some(item) = stream.next().await {
429 let event = match item {
430 Ok(event) => event,
431 Err(error) => {
432 return BenchmarkInvocation::failure((&error.kind).into(), started.elapsed());
433 }
434 };
435 if ttft.is_none() && is_model_output(&event) {
436 ttft = Some(started.elapsed());
437 }
438 match accumulator.push(event) {
439 Ok(Some(_)) => {
440 return BenchmarkInvocation::success(ttft, started.elapsed()).unwrap_or_else(
441 |_| {
442 BenchmarkInvocation::failure(BenchmarkFailureKind::Other, started.elapsed())
443 },
444 );
445 }
446 Ok(None) => {}
447 Err(error) => {
448 return BenchmarkInvocation::failure((&error.kind).into(), started.elapsed());
449 }
450 }
451 }
452 BenchmarkInvocation::failure(BenchmarkFailureKind::Stream, started.elapsed())
453}
454
455fn is_model_output(event: &ModelStreamEvent) -> bool {
456 matches!(
457 event,
458 ModelStreamEvent::TextDelta { .. }
459 | ModelStreamEvent::ReasoningDelta { .. }
460 | ModelStreamEvent::ToolArgumentsDelta { .. }
461 | ModelStreamEvent::RefusalDelta { .. }
462 | ModelStreamEvent::ContentPartCompleted { .. }
463 )
464}
465
466fn build_report(
467 label: String,
468 plan: &BenchmarkPlan,
469 wall_time: Duration,
470 outcomes: &[BenchmarkInvocation],
471) -> ProviderBenchmarkReport {
472 let mut total = Vec::new();
473 let mut ttft = Vec::new();
474 let mut successes_without_output = 0;
475 let mut failures = BTreeMap::new();
476 for outcome in outcomes {
477 if let Some(kind) = outcome.failure {
478 *failures.entry(kind).or_insert(0) += 1;
479 } else {
480 total.push(duration_us(outcome.total));
481 if let Some(value) = outcome.ttft {
482 ttft.push(duration_us(value));
483 } else {
484 successes_without_output += 1;
485 }
486 }
487 }
488 let successes = total.len();
489 let failure_count = outcomes.len() - successes;
490 let wall_seconds = wall_time.as_secs_f64().max(f64::EPSILON);
491 ProviderBenchmarkReport {
492 schema_version: REPORT_SCHEMA_VERSION,
493 label,
494 measured_runs: outcomes.len(),
495 concurrency: plan.concurrency.get(),
496 environment: plan.environment.clone(),
497 successes,
498 failures: failure_count,
499 success_rate: usize_as_f64(successes) / usize_as_f64(outcomes.len()),
500 wall_time_us: duration_us(wall_time),
501 throughput_per_second: usize_as_f64(outcomes.len()) / wall_seconds,
502 total_latency: distribution(&mut total),
503 ttft: distribution(&mut ttft),
504 successes_without_output,
505 failure_counts: failures
506 .into_iter()
507 .map(|(kind, count)| BenchmarkFailureCount { kind, count })
508 .collect(),
509 }
510}
511
512fn distribution(values: &mut [u64]) -> Option<LatencyDistribution> {
513 if values.is_empty() {
514 return None;
515 }
516 values.sort_unstable();
517 Some(LatencyDistribution {
518 min_us: values[0],
519 p50_us: nearest_rank(values, 50),
520 p95_us: nearest_rank(values, 95),
521 p99_us: nearest_rank(values, 99),
522 max_us: values[values.len() - 1],
523 })
524}
525
526fn nearest_rank(values: &[u64], percentile: usize) -> u64 {
527 let rank = values.len().saturating_mul(percentile).saturating_add(99) / 100;
528 values[rank.saturating_sub(1).min(values.len() - 1)]
529}
530
531fn duration_us(duration: Duration) -> u64 {
532 u64::try_from(duration.as_micros()).unwrap_or(u64::MAX)
533}
534
535fn usize_as_f64(value: usize) -> f64 {
536 value.to_string().parse().unwrap_or(f64::MAX)
537}
538
539fn u64_as_f64(value: u64) -> f64 {
540 value.to_string().parse().unwrap_or(f64::MAX)
541}
542
543fn larger_is_better(
544 name: &str,
545 baseline: Option<f64>,
546 candidate: Option<f64>,
547 minimum: f64,
548) -> BenchmarkRegressionMetric {
549 BenchmarkRegressionMetric {
550 name: name.into(),
551 baseline,
552 candidate,
553 passed: candidate.is_some_and(|candidate| candidate >= minimum),
554 }
555}
556
557fn smaller_is_better(
558 name: &str,
559 baseline: Option<f64>,
560 candidate: Option<f64>,
561 allowed_increase: f64,
562) -> BenchmarkRegressionMetric {
563 let passed = match (baseline, candidate) {
564 (Some(baseline), Some(candidate)) => candidate <= baseline * (1.0 + allowed_increase),
565 (None, None) => true,
566 _ => false,
567 };
568 BenchmarkRegressionMetric {
569 name: name.into(),
570 baseline,
571 candidate,
572 passed,
573 }
574}
575
576#[cfg(test)]
577mod tests {
578 use std::sync::atomic::{AtomicUsize, Ordering};
579
580 use futures_executor::block_on;
581
582 use super::*;
583
584 struct SequenceTarget {
585 next: AtomicUsize,
586 outcomes: Vec<BenchmarkInvocation>,
587 }
588
589 impl BenchmarkTarget for SequenceTarget {
590 fn execute(&self) -> BenchmarkFuture<'_> {
591 let index = self.next.fetch_add(1, Ordering::Relaxed) % self.outcomes.len();
592 let outcome = self.outcomes[index];
593 Box::pin(async move { outcome })
594 }
595 }
596
597 fn success(ttft_us: u64, total_us: u64) -> BenchmarkInvocation {
598 BenchmarkInvocation::success(
599 Some(Duration::from_micros(ttft_us)),
600 Duration::from_micros(total_us),
601 )
602 .unwrap()
603 }
604
605 #[test]
606 fn benchmark_report_is_stable_bounded_and_failure_aware() {
607 let target = Arc::new(SequenceTarget {
608 next: AtomicUsize::new(0),
609 outcomes: vec![
610 success(10, 20),
611 success(20, 40),
612 success(30, 60),
613 BenchmarkInvocation::failure(
614 BenchmarkFailureKind::Transport,
615 Duration::from_micros(5),
616 ),
617 ],
618 });
619 let plan = BenchmarkPlan::new(NonZeroUsize::new(4).unwrap())
620 .with_concurrency(NonZeroUsize::new(2).unwrap());
621
622 let report = block_on(benchmark("candidate", target, plan)).unwrap();
623
624 assert_eq!(report.schema_version, 1);
625 assert_eq!(report.successes, 3);
626 assert_eq!(report.failures, 1);
627 assert!((report.success_rate - 0.75).abs() < f64::EPSILON);
628 assert_eq!(report.ttft.unwrap().p95_us, 30);
629 assert_eq!(report.total_latency.unwrap().p50_us, 40);
630 assert_eq!(
631 report.failure_counts,
632 vec![BenchmarkFailureCount {
633 kind: BenchmarkFailureKind::Transport,
634 count: 1,
635 }]
636 );
637 assert!(serde_json::to_value(report).is_ok());
638 }
639
640 #[test]
641 fn regression_gate_rejects_slow_or_unreliable_candidates() {
642 let baseline = report("baseline", 1.0, 100.0, 10, 20);
643 let candidate = report("candidate", 0.8, 70.0, 15, 30);
644 let policy = BenchmarkRegressionPolicy::new(0.05, 0.1, 0.1, 0.1).unwrap();
645
646 let comparison = compare_benchmarks(&baseline, &candidate, policy);
647
648 assert!(!comparison.passed);
649 assert!(comparison.metrics.iter().all(|metric| !metric.passed));
650 }
651
652 fn report(
653 label: &str,
654 success_rate: f64,
655 throughput: f64,
656 p95_ttft: u64,
657 p95_total: u64,
658 ) -> ProviderBenchmarkReport {
659 ProviderBenchmarkReport {
660 schema_version: 1,
661 label: label.into(),
662 measured_runs: 10,
663 concurrency: 1,
664 environment: BTreeMap::new(),
665 successes: 10,
666 failures: 0,
667 success_rate,
668 wall_time_us: 100,
669 throughput_per_second: throughput,
670 total_latency: Some(LatencyDistribution {
671 min_us: p95_total,
672 p50_us: p95_total,
673 p95_us: p95_total,
674 p99_us: p95_total,
675 max_us: p95_total,
676 }),
677 ttft: Some(LatencyDistribution {
678 min_us: p95_ttft,
679 p50_us: p95_ttft,
680 p95_us: p95_ttft,
681 p99_us: p95_ttft,
682 max_us: p95_ttft,
683 }),
684 successes_without_output: 0,
685 failure_counts: Vec::new(),
686 }
687 }
688}