Skip to main content

runifold_provider_testkit/
benchmark.rs

1//! Framework-neutral latency, throughput, and reliability benchmark contract.
2
3use 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
21/// A boxed benchmark invocation future.
22pub type BenchmarkFuture<'a> = Pin<Box<dyn Future<Output = BenchmarkInvocation> + Send + 'a>>;
23
24/// Framework-neutral operation executed by the benchmark runner.
25///
26/// External comparison adapters can implement this trait for another runtime,
27/// including Rig, while retaining identical scheduling and report semantics.
28pub trait BenchmarkTarget: Send + Sync {
29    /// Executes one isolated invocation.
30    fn execute(&self) -> BenchmarkFuture<'_>;
31}
32
33/// Canonical Runifold model adapter for [`BenchmarkTarget`].
34#[derive(Clone)]
35pub struct ModelBenchmarkTarget {
36    model: Arc<dyn Model>,
37    request: ModelRequest,
38}
39
40impl ModelBenchmarkTarget {
41    /// Creates a target from one model and repeatable canonical request.
42    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/// Stable framework-neutral failure category.
63#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
64#[serde(rename_all = "snake_case")]
65#[non_exhaustive]
66pub enum BenchmarkFailureKind {
67    /// Local request validation failed.
68    InvalidRequest,
69    /// A required feature was unsupported.
70    UnsupportedFeature,
71    /// Network or transport failure.
72    Transport,
73    /// Provider wire-protocol failure.
74    Protocol,
75    /// Stream lifecycle or output assembly failure.
76    Stream,
77    /// Provider-side rejection or failure.
78    Provider,
79    /// Invocation cancellation.
80    Cancelled,
81    /// Invocation deadline exceeded.
82    DeadlineExceeded,
83    /// Runtime-specific failure without a more portable classification.
84    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/// Result and host-observed timing for one invocation.
104#[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    /// Creates one successful invocation.
113    ///
114    /// # Errors
115    ///
116    /// Returns [`BenchmarkInvocationError`] when time to first output exceeds
117    /// total invocation time.
118    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    /// Creates one failed invocation.
133    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/// Invalid timing evidence supplied by a benchmark adapter.
143#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
144#[non_exhaustive]
145pub enum BenchmarkInvocationError {
146    /// The first output was observed after invocation completion.
147    #[error("time to first output cannot exceed total invocation time")]
148    TtftAfterCompletion,
149}
150
151/// Bounded benchmark execution plan.
152#[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    /// Creates a sequential plan with no warmup.
162    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    /// Adds unmeasured warmup invocations.
172    #[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    /// Bounds concurrently executing measured invocations.
179    #[must_use]
180    pub const fn with_concurrency(mut self, concurrency: NonZeroUsize) -> Self {
181        self.concurrency = concurrency;
182        self
183    }
184
185    /// Records one stable environment fact in the generated report.
186    ///
187    /// Use this for framework version, Rust toolchain, target, runtime,
188    /// hardware class, provider endpoint, and model revision.
189    #[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/// Latency distribution expressed in integer microseconds.
197#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
198pub struct LatencyDistribution {
199    /// Minimum observed latency.
200    pub min_us: u64,
201    /// Nearest-rank 50th percentile.
202    pub p50_us: u64,
203    /// Nearest-rank 95th percentile.
204    pub p95_us: u64,
205    /// Nearest-rank 99th percentile.
206    pub p99_us: u64,
207    /// Maximum observed latency.
208    pub max_us: u64,
209}
210
211/// Count of failures in one normalized category.
212#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
213pub struct BenchmarkFailureCount {
214    /// Normalized failure category.
215    pub kind: BenchmarkFailureKind,
216    /// Number of measured invocations in this category.
217    pub count: usize,
218}
219
220/// Stable, serialization-safe benchmark evidence.
221#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
222pub struct ProviderBenchmarkReport {
223    /// Report schema version.
224    pub schema_version: u32,
225    /// User-controlled implementation or framework label.
226    pub label: String,
227    /// Number of measured invocations.
228    pub measured_runs: usize,
229    /// Maximum configured concurrency.
230    pub concurrency: usize,
231    /// Stable environment facts required to reproduce the run.
232    pub environment: BTreeMap<String, String>,
233    /// Successful invocation count.
234    pub successes: usize,
235    /// Failed invocation count.
236    pub failures: usize,
237    /// Successful invocations divided by measured invocations.
238    pub success_rate: f64,
239    /// Complete measured wall-clock interval.
240    pub wall_time_us: u64,
241    /// Measured invocations completed per second.
242    pub throughput_per_second: f64,
243    /// Successful invocation latency distribution.
244    pub total_latency: Option<LatencyDistribution>,
245    /// Successful time-to-first-output distribution.
246    pub ttft: Option<LatencyDistribution>,
247    /// Successful invocations which completed without visible output.
248    pub successes_without_output: usize,
249    /// Deterministically ordered normalized failure counts.
250    pub failure_counts: Vec<BenchmarkFailureCount>,
251}
252
253/// Benchmark construction or comparison failure.
254#[derive(Clone, Debug, Error, PartialEq)]
255#[non_exhaustive]
256pub enum ProviderBenchmarkError {
257    /// Report label was blank.
258    #[error("benchmark label cannot be empty")]
259    EmptyLabel,
260    /// Regression policy included a non-finite or out-of-range ratio.
261    #[error("benchmark regression ratio `{field}` must be finite and between zero and one")]
262    InvalidRegressionRatio {
263        /// Invalid policy field.
264        field: &'static str,
265    },
266}
267
268/// Runs warmup and measured invocations against one framework-neutral target.
269///
270/// # Errors
271///
272/// Returns [`ProviderBenchmarkError`] when `label` is blank.
273pub 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/// Allowed degradation from a baseline report.
300#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
301pub struct BenchmarkRegressionPolicy {
302    /// Maximum absolute success-rate decrease.
303    pub max_success_rate_drop: f64,
304    /// Maximum throughput decrease relative to baseline.
305    pub max_throughput_drop: f64,
306    /// Maximum p95 TTFT increase relative to baseline.
307    pub max_p95_ttft_increase: f64,
308    /// Maximum p95 total-latency increase relative to baseline.
309    pub max_p95_total_latency_increase: f64,
310}
311
312impl BenchmarkRegressionPolicy {
313    /// Creates a validated comparison policy.
314    ///
315    /// # Errors
316    ///
317    /// Returns [`ProviderBenchmarkError`] when any ratio is non-finite or
318    /// outside zero through one.
319    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/// One comparable benchmark metric and its gate decision.
348#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
349pub struct BenchmarkRegressionMetric {
350    /// Stable metric name.
351    pub name: String,
352    /// Baseline value.
353    pub baseline: Option<f64>,
354    /// Candidate value.
355    pub candidate: Option<f64>,
356    /// Whether the candidate satisfies the configured policy.
357    pub passed: bool,
358}
359
360/// Complete baseline-to-candidate benchmark decision.
361#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
362pub struct BenchmarkRegressionComparison {
363    /// Baseline report label.
364    pub baseline: String,
365    /// Candidate report label.
366    pub candidate: String,
367    /// Individual metric decisions.
368    pub metrics: Vec<BenchmarkRegressionMetric>,
369    /// Whether every metric passed.
370    pub passed: bool,
371}
372
373/// Compares two reports using one explicit regression policy.
374pub 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}