diskann_benchmark_runner/

benchmark.rs

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use serde::{Deserialize, Serialize};

use crate::{internal::visibility::Visibility, Checkpoint, Input, Output};

///////////////
// Benchmark //
///////////////

/// A registered benchmark.
///
/// Benchmarks consist of an [`Input`] and a corresponding serialized `Output`. Inputs will
/// first be validated with the benchmark using [`try_match`](Self::try_match). Only
/// successful matches will be passed to [`run`](Self::run).
pub trait Benchmark: 'static {
    /// The [`Input`] type this benchmark matches against.
    type Input: Input + 'static;

    /// The concrete type of the results generated by this benchmark.
    type Output: Serialize;

    /// Return whether or not this benchmark is compatible with `input`.
    ///
    /// Use [`MatchContext::success`] to create an initial [`Score`], then progressively
    /// refine it with [`Score::penalize`] and [`Score::fail`].
    ///
    /// Among successful matches, the benchmark with the lowest score wins. Ties are broken
    /// by an unspecified procedure.
    ///
    /// When no successful match exists, failure scores rank the "nearest misses".
    /// Implementations should use [`Score::fail`] with descriptive reasons to aid
    /// diagnostics.
    fn try_match(&self, input: &Self::Input, context: &MatchContext) -> Score;

    /// Return descriptive information about the benchmark.
    fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;

    /// Run the benchmark with `input`.
    ///
    /// All prints should be directed to `output`. The `checkpoint` is provided so
    /// long-running benchmarks can periodically save output to prevent data loss due to
    /// an early error.
    ///
    /// Implementors may assume that [`Self::try_match`] returned a successful [`Score`]
    /// for `input`.
    fn run(
        &self,
        input: &Self::Input,
        checkpoint: Checkpoint<'_>,
        output: &mut dyn Output,
    ) -> anyhow::Result<Self::Output>;
}

/// A refinement of [`Benchmark`], that supports before/after comparison of generated results.
///
/// Benchmarks are associated with a "tolerance" input, which may contain runtime values
/// controlling the amount of slack a benchmark is allowed to have between runs before failing.
///
/// The semantics of pass or failure are left solely to the discretion of the [`Regression`]
/// implementation.
///
/// See: [`register_regression`](crate::Registry::register_regression).
pub trait Regression: Benchmark<Output: for<'a> Deserialize<'a>> {
    /// The tolerance [`Input`] associated with this regression check.
    type Tolerances: Input + 'static;

    /// The report summary used to describe a successful regression check.
    type Pass: Serialize + std::fmt::Display + 'static;

    /// The report summary used to describe an unsuccessful regression check.
    type Fail: Serialize + std::fmt::Display + 'static;

    /// Run any regression checks necessary for two benchmark runs `before` and `after`.
    /// Argument `tolerances` contain any tuned runtime tolerances to use when determining
    /// whether or not a regression is detected.
    ///
    /// The `input` is the raw input that would have been provided to [`Benchmark::run`]
    /// when generating the `before` and `after` outputs.
    ///
    /// Implementations of `check` should not attempt to print to `stdout` or any other
    /// stream. Instead, all diagnostics should be encoded in the returned [`PassFail`] type
    /// for reporting upstream.
    fn check(
        &self,
        tolerances: &Self::Tolerances,
        input: &Self::Input,
        before: &Self::Output,
        after: &Self::Output,
    ) -> anyhow::Result<PassFail<Self::Pass, Self::Fail>>;
}

/// Describe whether or not a [`Regression`] passed or failed.
#[derive(Debug, Clone, Copy)]
pub enum PassFail<P, F> {
    Pass(P),
    Fail(F),
}

//////////////
// Matching //
//////////////

/// Context for [`Benchmark::try_match`].
///
/// This is used to create matching [`Score`]s via [`Self::success`] and [`Self::fail`].
/// Users can test their [`Benchmark::try_match`] implementations using [`Self::test`].
///
/// Internally, the [`MatchContext`] decides whether or not failure reasons are evaluated,
/// eliding formatting in situations where the results will not be used.
#[derive(Debug)]
pub struct MatchContext {
    record_failure_reasons: bool,
}

impl MatchContext {
    /// Create a new [`Score`] configured for "success" with the given `score`.
    ///
    /// Lower scores indicate better successful matches.
    pub fn success(&self, score: u32) -> Score {
        Score {
            inner: ScoreInner::Success(SuccessScore(score)),
            context: self.hidden_clone(),
        }
    }

    /// Create a new [`Score`] configured for "failure" with the given `score`.
    ///
    /// Lower scores indicate better "near misses".
    pub fn fail(&self, score: u32, reason: &dyn std::fmt::Display) -> Score {
        let mut s = self.success(u32::MAX);
        s.fail(score, reason);
        s
    }

    /// Test a [`Benchmark::try_match`] implementation, returning a [`TestScore`] with
    /// full failure reasons (if applicable).
    pub fn test<T>(benchmark: &T, input: &T::Input) -> TestScore
    where
        T: Benchmark,
    {
        benchmark
            .try_match(input, &Self::with_reasons())
            .into_test()
    }

    fn hidden_clone(&self) -> Self {
        Self {
            record_failure_reasons: self.record_failure_reasons,
        }
    }

    /// Create a new [`MatchContext`] that does not evaluate failure reasons.
    pub(crate) fn new() -> Self {
        Self {
            record_failure_reasons: false,
        }
    }

    /// Create a new [`MatchContext`] that evaluates failure reasons.
    pub(crate) fn with_reasons() -> Self {
        Self {
            record_failure_reasons: true,
        }
    }
}

/// The result of [`MatchContext::test`], providing low-level access to the final match
/// scoring and failure reasons.
#[derive(Debug)]
pub enum TestScore {
    Success(u32),
    Failure {
        score: u32,
        reasons: Option<Vec<String>>,
    },
}

/// A score for [`Benchmark::try_match`].
///
/// A [`Score`] is in one of two states: success or failure. Lower scores indicate better
/// matches in both states. In the failure case, scores rank "near misses" for diagnostics.
///
/// - [`Score::penalize`] increases the score of a successful match.
/// - [`Score::fail`] transitions to (or worsens) the failure state.
///
/// The current state can be queried with [`Score::is_success`].
#[derive(Debug)]
pub struct Score {
    inner: ScoreInner,
    context: MatchContext,
}

impl Score {
    /// If the score is in the "success" state, penalize it by `by`.
    ///
    /// Has no effect if the score is already in the "failure" state.
    pub fn penalize(&mut self, by: u32) {
        match self.inner {
            ScoreInner::Success(ref mut v) => v.0 = v.0.saturating_add(by),
            ScoreInner::Failure(..) => {}
        }
    }

    /// Transition the score to the "failure" state with penalty `by` and a `reason`.
    ///
    /// If the score is already in the "failure" state, `by` is added to the existing
    /// failure score and `reason` is appended to the list of failure reasons.
    pub fn fail(&mut self, by: u32, reason: &dyn std::fmt::Display) {
        match &mut self.inner {
            ScoreInner::Success(_) => {
                self.inner = ScoreInner::Failure(
                    FailureScore(by),
                    self.context
                        .record_failure_reasons
                        .then(|| vec![reason.to_string()]),
                );
            }
            ScoreInner::Failure(score, reasons) => {
                score.0 = (score.0).saturating_add(by);
                if self.context.record_failure_reasons {
                    reasons
                        .get_or_insert_with(|| Vec::with_capacity(1))
                        .push(reason.to_string())
                }
            }
        }
    }

    /// Return `true` if `self` is in the "success" state. Returning `false` implies the
    /// "failure" state.
    #[must_use = "this function has no side-effects"]
    pub fn is_success(&self) -> bool {
        matches!(self.inner, ScoreInner::Success(_))
    }

    fn into_test(self) -> TestScore {
        match self.inner {
            ScoreInner::Success(score) => TestScore::Success(score.0),
            ScoreInner::Failure(score, reasons) => TestScore::Failure {
                score: score.0,
                reasons,
            },
        }
    }

    pub(crate) fn match_score(&self) -> Option<SuccessScore> {
        match self.inner {
            ScoreInner::Success(score) => Some(score),
            ScoreInner::Failure(..) => None,
        }
    }

    pub(crate) fn reason(&self) -> Reason<'_> {
        match &self.inner {
            ScoreInner::Success(_) => Reason::none(),
            ScoreInner::Failure(_, reasons) => Reason::new(reasons.as_deref()),
        }
    }

    /// Allow `Scores` to be ordered in the following manner:
    ///
    /// * Successes are ordered before failures.
    /// * Successes with lower scores are ordered before those with higher scores.
    /// * Failures with lower scores are ordered before those with higher scores.
    pub(crate) fn order(this: &Self, other: &Self) -> std::cmp::Ordering {
        use ScoreInner::{Failure, Success};

        match (&this.inner, &other.inner) {
            (Success(this), Success(other)) => this.cmp(other),
            (Success(_), Failure(..)) => std::cmp::Ordering::Less,
            (Failure(..), Success(_)) => std::cmp::Ordering::Greater,
            (Failure(this, _), Failure(other, _)) => this.cmp(other),
        }
    }
}

#[derive(Debug)]
enum ScoreInner {
    Success(SuccessScore),
    Failure(FailureScore, Option<Vec<String>>),
}

pub(crate) struct Reason<'a>(Option<&'a [String]>);

impl<'a> Reason<'a> {
    fn new(reasons: Option<&'a [String]>) -> Self {
        Self(reasons)
    }

    fn none() -> Self {
        Self(None)
    }
}

impl std::fmt::Display for Reason<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.0 {
            None => f.write_str("<missing>"),
            Some(reasons) => {
                let mut first = true;
                for reason in reasons.iter() {
                    if !first {
                        writeln!(f)?;
                    }
                    write!(f, "- {}", reason)?;
                    first = false;
                }
                Ok(())
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct SuccessScore(pub(crate) u32);

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct FailureScore(pub(crate) u32);

//////////////
// Internal //
//////////////

pub(crate) mod internal {
    use super::*;

    use crate::{input::internal::Any, Features};

    use anyhow::Context;
    use thiserror::Error;

    #[derive(Debug)]
    pub(crate) enum AnnotatedMatch<'a> {
        /// User matching indicated success.
        User(Score),
        /// We missed because the input type was incorrect.
        WrongTag,
        /// We missed because the input type matches, but the benchmark is feature-gated.
        Gated(&'a Features),
    }

    impl AnnotatedMatch<'_> {
        pub(crate) fn try_into_score(self) -> Option<Score> {
            if let Self::User(score) = self {
                Some(score)
            } else {
                None
            }
        }
    }

    /// Object-safe trait for type-erased benchmarks stored in the registry.
    pub(crate) trait Benchmark {
        fn try_match(&self, input: &Any, context: &MatchContext) -> AnnotatedMatch<'_>;

        fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;

        fn run(
            &self,
            input: &Any,
            checkpoint: Checkpoint<'_>,
            output: &mut dyn Output,
        ) -> anyhow::Result<serde_json::Value>;

        fn visibility(&self) -> Visibility<'_>;

        /// If supported, return an object capable of running regression checks on this benchmark.
        fn as_regression(&self) -> Option<&dyn Regression>;

        fn as_string(&self) -> String {
            Description(self).to_string()
        }
    }

    struct Description<'a, T: ?Sized>(&'a T);

    impl<T> std::fmt::Display for Description<'_, T>
    where
        T: Benchmark + ?Sized,
    {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            self.0.description(f)
        }
    }

    pub(crate) struct Checked {
        pub(crate) json: serde_json::Value,
        pub(crate) display: Box<dyn std::fmt::Display>,
    }

    impl Checked {
        /// Serialize `value` to `serde_json::Value` and box it for future display.
        fn new<T>(value: T) -> Result<Self, serde_json::Error>
        where
            T: Serialize + std::fmt::Display + 'static,
        {
            Ok(Self {
                json: serde_json::to_value(&value)?,
                display: Box::new(value),
            })
        }
    }

    pub(crate) type CheckedPassFail = PassFail<Checked, Checked>;

    pub(crate) trait Regression {
        fn tolerance(&self) -> &dyn crate::input::internal::DynInput;
        fn input_tag(&self) -> &'static str;
        fn check(
            &self,
            tolerances: &Any,
            input: &Any,
            before: &serde_json::Value,
            after: &serde_json::Value,
        ) -> anyhow::Result<CheckedPassFail>;
    }

    impl std::fmt::Debug for dyn Regression + '_ {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("dyn Regression")
                .field("tolerance", &self.tolerance().tag())
                .field("input_tag", &self.input_tag())
                .finish()
        }
    }

    pub(crate) trait AsRegression<T> {
        fn as_regression(benchmark: &T) -> Option<&dyn Regression>;
    }

    #[derive(Debug, Clone, Copy)]
    pub(crate) struct NoRegression;

    impl<T> AsRegression<T> for NoRegression {
        fn as_regression(_benchmark: &T) -> Option<&dyn Regression> {
            None
        }
    }

    #[derive(Debug, Clone, Copy)]
    pub(crate) struct WithRegression;

    impl<T> AsRegression<T> for WithRegression
    where
        T: super::Regression,
    {
        fn as_regression(benchmark: &T) -> Option<&dyn Regression> {
            Some(benchmark)
        }
    }

    impl<T> Regression for T
    where
        T: super::Regression,
    {
        fn tolerance(&self) -> &dyn crate::input::internal::DynInput {
            &crate::input::internal::Wrapper::<T::Tolerances>::INSTANCE
        }

        fn input_tag(&self) -> &'static str {
            T::Input::tag()
        }

        fn check(
            &self,
            tolerance: &Any,
            input: &Any,
            before: &serde_json::Value,
            after: &serde_json::Value,
        ) -> anyhow::Result<CheckedPassFail> {
            let tolerance = tolerance
                .downcast_ref::<T::Tolerances>()
                .ok_or_else(|| BadDownCast::new(T::Tolerances::tag(), tolerance.tag()))
                .context("failed to obtain tolerance")?;

            let input = input
                .downcast_ref::<T::Input>()
                .ok_or_else(|| BadDownCast::new(T::Input::tag(), input.tag()))
                .context("failed to obtain input")?;

            let before = T::Output::deserialize(before)
                .map_err(|err| DeserializationError::new(Kind::Before, err))?;

            let after = T::Output::deserialize(after)
                .map_err(|err| DeserializationError::new(Kind::After, err))?;

            let passfail = match self.check(tolerance, input, &before, &after)? {
                PassFail::Pass(pass) => PassFail::Pass(Checked::new(pass)?),
                PassFail::Fail(fail) => PassFail::Fail(Checked::new(fail)?),
            };

            Ok(passfail)
        }
    }

    #[derive(Debug, Clone, Copy)]
    pub(crate) struct Wrapper<T, R = NoRegression> {
        benchmark: T,
        _regression: R,
    }

    impl<T, R> Wrapper<T, R> {
        pub(crate) const fn new(benchmark: T, regression: R) -> Self {
            Self {
                benchmark,
                _regression: regression,
            }
        }
    }

    impl<T, R> Benchmark for Wrapper<T, R>
    where
        T: super::Benchmark,
        R: AsRegression<T>,
    {
        fn try_match(&self, input: &Any, context: &MatchContext) -> AnnotatedMatch<'_> {
            if let Some(cast) = input.downcast_ref::<T::Input>() {
                AnnotatedMatch::User(self.benchmark.try_match(cast, context))
            } else {
                AnnotatedMatch::WrongTag
            }
        }

        fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            writeln!(f, "tag \"{}\"", <T::Input as Input>::tag())?;
            self.benchmark.description(f)
        }

        fn run(
            &self,
            input: &Any,
            checkpoint: Checkpoint<'_>,
            output: &mut dyn Output,
        ) -> anyhow::Result<serde_json::Value> {
            match input.downcast_ref::<T::Input>() {
                Some(input) => {
                    let result = self.benchmark.run(input, checkpoint, output)?;
                    Ok(serde_json::to_value(result)?)
                }
                None => Err(BadDownCast::new(T::Input::tag(), input.tag()).into()),
            }
        }

        fn visibility(&self) -> Visibility<'_> {
            Visibility::Available
        }

        // Extensions
        fn as_regression(&self) -> Option<&dyn Regression> {
            R::as_regression(&self.benchmark)
        }
    }

    //--------//
    // Errors //
    //--------//

    #[derive(Debug, Clone, Copy, Error)]
    #[error(
        "INTERNAL ERROR: bad downcast - expected \"{}\" but got \"{}\"",
        self.expected,
        self.got
    )]
    struct BadDownCast {
        expected: &'static str,
        got: &'static str,
    }

    impl BadDownCast {
        fn new(expected: &'static str, got: &'static str) -> Self {
            Self { expected, got }
        }
    }

    #[derive(Debug, Error)]
    #[error(
        "the \"{}\" results do not match the output schema expected by this benchmark",
        self.kind
    )]
    struct DeserializationError {
        kind: Kind,
        source: serde_json::Error,
    }

    impl DeserializationError {
        fn new(kind: Kind, source: serde_json::Error) -> Self {
            Self { kind, source }
        }
    }

    #[derive(Debug, Clone, Copy)]
    enum Kind {
        Before,
        After,
    }

    impl std::fmt::Display for Kind {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let as_str = match self {
                Self::Before => "before",
                Self::After => "after",
            };

            write!(f, "{}", as_str)
        }
    }

    // Gated Benchmarks
    //
    // We support two flavors:
    //
    // 1. Benchmarks that have inputs that are valid and registered.
    // 2. Benchmarks that have inputs that are *also* feature gated.

    /// A benchmark that is gated behind `features`, but whose input *is* compiled into the
    /// application. These benchmarks allow their inputs to be inspected and parsed, but
    /// fail matching.
    ///
    /// This custom type allows us to return more precise error messages when benchmark
    /// matching fails.
    pub(crate) struct PartiallyGated<I> {
        features: Features,
        description: String,
        _input: std::marker::PhantomData<I>,
    }

    impl<I> PartiallyGated<I> {
        pub(crate) fn new(features: Features, description: String) -> Self {
            Self {
                features,
                description,
                _input: std::marker::PhantomData,
            }
        }
    }

    impl<I> Benchmark for PartiallyGated<I>
    where
        I: Input,
    {
        fn try_match(&self, input: &Any, _context: &MatchContext) -> AnnotatedMatch<'_> {
            if input.is::<I>() {
                AnnotatedMatch::Gated(&self.features)
            } else {
                AnnotatedMatch::WrongTag
            }
        }

        fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            writeln!(f, "tag \"{}\"", I::tag())?;
            f.write_str(&self.description)
        }

        fn run(
            &self,
            _input: &Any,
            _checkpoint: Checkpoint<'_>,
            _output: &mut dyn Output,
        ) -> anyhow::Result<serde_json::Value> {
            anyhow::bail!("tried to run a gated benchmark - this should be unreachable")
        }

        fn visibility(&self) -> Visibility<'_> {
            Visibility::Gated {
                features: &self.features,
            }
        }

        fn as_regression(&self) -> Option<&dyn Regression> {
            None
        }
    }

    /// A benchmark that is gated behind `features` along with its input.
    ///
    /// These benchmarks are completely unreachable as the front-end has no way to parse
    /// inputs associated with `tag`.
    ///
    /// However, they are included to provide a bread-crumb for users that additional
    /// benchmarks *do* exist behind the associated features.
    pub(crate) struct FullyGated {
        tag: &'static str,
        features: Features,
        description: String,
    }

    impl FullyGated {
        pub(crate) fn new(tag: &'static str, features: Features, description: String) -> Self {
            Self {
                tag,
                features,
                description,
            }
        }
    }

    impl Benchmark for FullyGated {
        fn try_match(&self, _input: &Any, _context: &MatchContext) -> AnnotatedMatch<'_> {
            AnnotatedMatch::WrongTag
        }

        fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            writeln!(f, "tag \"{}\"", self.tag)?;
            f.write_str(&self.description)
        }

        fn run(
            &self,
            _input: &Any,
            _checkpoint: Checkpoint<'_>,
            _output: &mut dyn Output,
        ) -> anyhow::Result<serde_json::Value> {
            anyhow::bail!("tried to run a gated benchmark - this should be unreachable")
        }

        fn visibility(&self) -> Visibility<'_> {
            Visibility::Gated {
                features: &self.features,
            }
        }

        fn as_regression(&self) -> Option<&dyn Regression> {
            None
        }
    }
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug, Clone, Copy, PartialEq)]
    pub(crate) enum RawScore {
        Success(SuccessScore),
        Failure(FailureScore),
    }

    fn as_raw(score: &Score) -> RawScore {
        match &score.inner {
            ScoreInner::Success(s) => RawScore::Success(*s),
            ScoreInner::Failure(s, _) => RawScore::Failure(*s),
        }
    }

    impl RawScore {
        #[cfg(test)]
        fn success(score: u32) -> Self {
            Self::Success(SuccessScore(score))
        }

        #[cfg(test)]
        fn failure(score: u32) -> Self {
            Self::Failure(FailureScore(score))
        }
    }

    #[test]
    fn test_score_no_reasons() {
        let context = MatchContext::new();
        let mut score = context.success(0);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(0));

        score.penalize(10);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(10));

        score.penalize(10);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(20));

        // Ensure that we saturate properly.
        score.penalize(u32::MAX);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(u32::MAX));

        // Switch to failure
        score.fail(5, &"some reason that is not evaluated");
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(5));

        score.fail(10, &"another reason that is not evaluated");
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(15));

        // Calling `penalize` should have no effect.
        score.penalize(5);
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(15));

        // Since we aren't recording reasons - nothing should be returned.
        assert_eq!(score.reason().to_string(), "<missing>");
    }

    #[test]
    fn test_score_with_reasons() {
        let context = MatchContext::with_reasons();
        let mut score = context.success(0);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(0));

        score.penalize(10);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(10));

        score.penalize(10);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(20));
        assert_eq!(score.reason().to_string(), "<missing>");

        // Ensure that we saturate properly.
        score.penalize(u32::MAX);
        assert!(score.is_success());
        assert_eq!(as_raw(&score), RawScore::success(u32::MAX));

        // Switch to failure
        score.fail(5, &"some reason that is evaluated");
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(5));

        score.fail(10, &"another reason that is evaluated");
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(15));

        // Calling `penalize` should have no effect.
        score.penalize(5);
        assert!(!score.is_success());
        assert_eq!(as_raw(&score), RawScore::failure(15));

        // Reasons are recorded, so both failure reasons should be returned.
        let expected = "- some reason that is evaluated\n\
                        - another reason that is evaluated";
        assert_eq!(score.reason().to_string(), expected);
    }
}