diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 3d07e85..9b1f4cc 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -1,5 +1,4 @@ use braintrust_sdk_rust::eval::*; -use braintrust_sdk_rust::*; /// Demonstrates using BtEvalRunner, which is compatible with the `bt eval` CLI. /// @@ -10,7 +9,7 @@ use braintrust_sdk_rust::*; /// - BT_EVAL_SSE_SOCK=... — send SSE events over a Unix socket /// - BT_EVAL_SSE_ADDR=... — send SSE events over TCP #[tokio::main] -async fn main() -> Result<()> { +async fn main() -> anyhow::Result<()> { // BtEvalRunner reads BT_EVAL_* env vars automatically. // In non-list mode it creates a BraintrustClient from the environment. let mut runner = BtEvalRunner::from_env().await?; @@ -39,19 +38,35 @@ async fn main() -> Result<()> { }, ]; + // Share callbacks across runs; each run supplies its own source and options. + let definition = EvalDefinition::new( + Box::new(task(|input: &String, _hooks| { + Box::pin(async move { Ok(input.to_uppercase()) }) + })), + vec![ + Box::new(ExactMatch::new()), + Box::new(LevenshteinScorer::new()), + Box::new(scorer( + "nonempty", + |args: ScorerArgs<'_, String, String>| { + Box::pin(async move { + Ok(vec![Score::new( + "nonempty", + if args.output.is_empty() { 0.0 } else { 1.0 }, + )]) + }) + }, + )), + ], + )?; + runner .eval( "my-project", EvalOpts::builder() .name("uppercase-eval".to_string()) - .dataset(cases.into()) - .task(Box::new(SimpleFnTask::new(|input: &String| { - Ok(input.to_uppercase()) - }))) - .scorers(vec![ - Box::new(ExactMatch::new()), - Box::new(LevenshteinScorer::new()), - ]) + .source(cases.into()) + .definition(definition.clone()) .max_concurrency(2) .build(), ) diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index ae065a5..09e9c50 100644 --- a/examples/eval_sync.rs +++ b/examples/eval_sync.rs @@ -1,15 +1,11 @@ use braintrust_sdk_rust::eval::*; use braintrust_sdk_rust::*; +use std::future::IntoFuture; fn main() -> anyhow::Result<()> { - // Note: No #[tokio::main] needed! SyncEvaluator handles async internally - - // Create Braintrust client - let rt = tokio::runtime::Runtime::new().map_err(|e| anyhow::anyhow!(e))?; - let client = rt.block_on(async { BraintrustClient::builder().build().await })?; - - // Create sync evaluator - let evaluator = SyncEvaluator::new(client)?; + // The client's owned runtime drives work; this thread only waits for tokens. + let client = futures::executor::block_on(BraintrustClient::builder().build().into_future())?; + let evaluator = Evaluator::new(client); // Define test cases let cases = vec![ @@ -39,13 +35,17 @@ fn main() -> anyhow::Result<()> { // Run evaluation synchronously println!("Starting synchronous evaluation...\n"); - let result = evaluator.run( - EvalOpts::builder() - .name("sync-uppercase-eval".to_string()) - .dataset(cases.into()) - .task(Box::new(task)) - .scorers(scorers) - .build(), + let definition = EvalDefinition::new(Box::new(task), scorers)?; + let result = futures::executor::block_on( + evaluator + .run( + EvalOpts::builder() + .name("sync-uppercase-eval".to_string()) + .source(cases.into()) + .definition(definition.clone()) + .build(), + ) + .into_future(), )?; // Print results diff --git a/src/eval/blocking.rs b/src/eval/blocking.rs new file mode 100644 index 0000000..e216996 --- /dev/null +++ b/src/eval/blocking.rs @@ -0,0 +1,68 @@ +//! Managed blocking invocation on the run's fixed execution binding. +use crate::internal::execution::ExecutionBinding; + +pub(crate) async fn invoke( + binding: Option<&ExecutionBinding>, + work: impl FnOnce() -> anyhow::Result + Send + 'static, +) -> anyhow::Result { + let binding = binding.ok_or_else(|| { + anyhow::anyhow!("blocking callbacks require a managed evaluation executor") + })?; + let task = binding.spawn_blocking(work); + // Abort queued work on run cancellation. Tokio cannot interrupt a closure + // already running; its owned data remains valid until that closure returns. + struct AbortOnDrop(tokio::task::AbortHandle); + impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } + } + let _abort = AbortOnDrop(task.abort_handle()); + task.await? +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }; + + #[test] + fn dropping_an_invocation_cancels_queued_work_before_it_can_start() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + let binding = ExecutionBinding::external(runtime.handle().clone()); + let (release, waiting) = std::sync::mpsc::channel(); + let (entered, started) = std::sync::mpsc::channel(); + let occupied = binding.spawn_blocking(move || { + entered.send(()).unwrap(); + waiting + .recv_timeout(std::time::Duration::from_secs(3)) + .unwrap() + }); + started + .recv_timeout(std::time::Duration::from_secs(3)) + .unwrap(); + let called = Arc::new(AtomicBool::new(false)); + runtime.block_on(async { + { + let called = called.clone(); + let callback = invoke(Some(&binding), move || { + called.store(true, Ordering::SeqCst); + Ok(()) + }); + futures::pin_mut!(callback); + assert!(futures::poll!(callback).is_pending()); + } + release.send(()).unwrap(); + occupied.await.unwrap(); + binding.spawn_blocking(|| ()).await.unwrap(); + }); + assert!(!called.load(Ordering::SeqCst)); + } +} diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index 1daae18..8662a52 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -10,7 +10,7 @@ use std::io::Write; use std::sync::{Arc, Mutex}; -use serde::{de::DeserializeOwned, Serialize}; +use serde::Serialize; use serde_json::{json, Value}; use crate::{BraintrustClient, BraintrustError, Result}; @@ -186,20 +186,24 @@ fn create_sse_writer() -> Option> { /// ```no_run /// use braintrust_sdk_rust::eval::*; /// -/// # async fn example() -> braintrust_sdk_rust::Result<()> { +/// # async fn example() -> anyhow::Result<()> { /// let mut runner = BtEvalRunner::from_env().await?; /// +/// let definition = EvalDefinition::new( +/// Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), +/// vec![Box::new(ExactMatch::new())], +/// )?; +/// /// runner.eval( /// "my-project", /// EvalOpts::builder() /// .name("my-eval".to_string()) -/// .dataset(vec![Case { +/// .source(vec![Case { /// input: "hello".to_string(), /// expected: Some("HELLO".to_string()), /// ..Default::default() /// }].into()) -/// .task(Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase())))) -/// .scorers(vec![Box::new(ExactMatch::new())]) +/// .definition(definition.clone()) /// .build(), /// ).await?; /// @@ -288,8 +292,8 @@ impl BtEvalRunner { /// If the eval name does not pass the configured filters, it is skipped silently. pub async fn eval(&mut self, project: &str, mut opts: EvalOpts) -> Result<()> where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, { let name = opts.name.clone(); @@ -349,7 +353,7 @@ impl BtEvalRunner { eprintln!("Evaluator '{}' failed: {}", name, e); } if self.config.terminate_on_failure { - return Err(e); + return Err(BraintrustError::Background(e.to_string())); } } } diff --git a/src/eval/case_runner.rs b/src/eval/case_runner.rs new file mode 100644 index 0000000..718368d --- /dev/null +++ b/src/eval/case_runner.rs @@ -0,0 +1,297 @@ +//! One case's computation, independent of recording and reporting. +use super::{Case, CaseError, EvalDefinition, EvalResult, ScorerArgs, TaskHooks}; +use crate::{internal::execution::ExecutionBinding, SpanHandle}; +use futures::{stream::FuturesOrdered, StreamExt}; +use std::sync::Arc; + +pub(crate) struct CaseRunner; +impl CaseRunner { + pub async fn run( + case: Case, + definition: &EvalDefinition, + span: Option<&SpanHandle>, + binding: Option<&ExecutionBinding>, + ) -> EvalResult { + let case = Arc::new(case); + let hooks = TaskHooks::for_case( + span, + case.expected.as_ref(), + case.metadata.clone(), + case.tags.clone(), + ); + let task = definition.task(); + let computed = if task.is_blocking() { + let case = case.clone(); + let state = hooks.rebind(None, None); + let span = span.cloned(); + super::blocking::invoke(binding, move || { + let hooks = state.rebind(span.as_ref(), case.expected.as_ref()); + task.run_blocking(&case.input, &hooks) + }) + .await + } else { + task.run(&case.input, &hooks).await + }; + let mut errors = Vec::new(); + let output = match computed { + Ok(output) => Some(Arc::new(output)), + Err(source) => { + errors.push(CaseError::Task { source }); + None + } + }; + let mut scores = Vec::new(); + if let Some(output) = &output { + let mut pending = FuturesOrdered::new(); + for scorer in definition.scorers() { + let case = case.clone(); + let output = output.clone(); + let scorer = scorer.clone(); + pending.push_back(async move { + let name = scorer.name().to_owned(); + let result = if scorer.is_blocking() { + super::blocking::invoke(binding, move || { + scorer.score_blocking(ScorerArgs { + input: &case.input, + output: &output, + expected: case.expected.as_ref(), + metadata: case.metadata.as_ref(), + }) + }) + .await + } else { + scorer + .score(ScorerArgs { + input: &case.input, + output: &output, + expected: case.expected.as_ref(), + metadata: case.metadata.as_ref(), + }) + .await + }; + result.map_err(|source| CaseError::Scorer { name, source }) + }); + } + while let Some(result) = pending.next().await { + match result { + Ok(values) => scores.extend(values), + Err(error) => errors.push(error), + } + } + } + let metadata = hooks.get_metadata(); + let tags = hooks.get_tags(); + drop(hooks); + // All owned invocations have joined. Borrowed callback arguments cannot + // escape, so no input/output clones are required to return ownership. + let case = + Arc::try_unwrap(case).unwrap_or_else(|_| unreachable!("case invocations joined")); + let output = output.map(|output| { + Arc::try_unwrap(output).unwrap_or_else(|_| unreachable!("scorers joined")) + }); + EvalResult { + input: case.input, + output, + expected: case.expected, + scores, + metadata, + tags, + errors, + case_id: case.id, + xact_id: case.xact_id, + created: case.created, + span_id: None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::{Score, Scorer, Scores, Task}; + use async_trait::async_trait; + use serde_json::json; + use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }; + + #[derive(Debug, thiserror::Error)] + #[error("application failure")] + struct ApplicationFailure; + + // These intentionally implement neither Clone, serialization nor equality. + struct Input(usize); + struct Output(usize); + struct LocalTask; + #[async_trait] + impl Task for LocalTask { + async fn run( + &self, + input: &Input, + hooks: &TaskHooks<'_, Input, Output>, + ) -> anyhow::Result { + assert!(hooks.span.is_none()); + assert_eq!(hooks.expected.unwrap().0, 6); + hooks.add_metadata("task", json!(true)); + hooks.add_tag("task"); + hooks.add_tag("source"); + Ok(Output(input.0 * 2)) + } + } + + #[test] + fn local_computation_owns_nonserializable_values_and_preserves_case_context() { + let case = Case { + input: Input(3), + expected: Some(Output(6)), + metadata: Some(serde_json::from_value(json!({"source":true})).unwrap()), + tags: Some(vec!["source".into()]), + id: Some("case".into()), + xact_id: Some("version".into()), + created: Some("timestamp".into()), + }; + let result = futures::executor::block_on(CaseRunner::run( + case, + &EvalDefinition::new(Box::new(LocalTask), vec![]).unwrap(), + None, + None, + )); + assert_eq!(result.output.unwrap().0, 6); + assert_eq!(result.case_id.as_deref(), Some("case")); + assert_eq!(result.xact_id.as_deref(), Some("version")); + assert_eq!(result.created.as_deref(), Some("timestamp")); + assert_eq!( + result.metadata.unwrap(), + serde_json::from_value::>( + json!({"source":true,"task":true}) + ) + .unwrap() + ); + assert_eq!(result.tags.unwrap(), vec!["source", "task"]); + assert!(result.errors.is_empty()); + assert!(result.span_id.is_none()); + } + + struct FailTask; + #[async_trait] + impl Task for FailTask { + async fn run( + &self, + _: &usize, + hooks: &TaskHooks<'_, usize, usize>, + ) -> anyhow::Result { + hooks.add_metadata("failed", json!(true)); + hooks.add_tag("failed"); + Err(ApplicationFailure.into()) + } + } + struct ProbeScorer(Arc); + #[async_trait] + impl Scorer for ProbeScorer { + fn name(&self) -> &str { + "probe" + } + async fn score(&self, _: ScorerArgs<'_, usize, usize>) -> anyhow::Result { + self.0.store(true, Ordering::SeqCst); + Ok(vec![]) + } + } + #[tokio::test] + async fn task_failure_skips_scorers_and_keeps_cause_and_hook_changes() { + let called = Arc::new(AtomicBool::new(false)); + let scorers: Vec>> = + vec![Box::new(ProbeScorer(called.clone()))]; + let result = CaseRunner::run( + Case { + input: 1, + ..Default::default() + }, + &EvalDefinition::new(Box::new(FailTask), scorers).unwrap(), + None, + None, + ) + .await; + assert!(result.output.is_none()); + assert!(!called.load(Ordering::SeqCst)); + assert_eq!(result.metadata.unwrap()["failed"], true); + assert_eq!(result.tags.unwrap(), ["failed"]); + let [CaseError::Task { source }] = result.errors.as_slice() else { + panic!("task failure must retain its stage") + }; + assert!(source.downcast_ref::().is_some()); + assert_eq!( + serde_json::to_value(&result.errors).unwrap(), + json!([{"stage":"task","message":"application failure"}]) + ); + } + + struct ParallelScorer { + name: &'static str, + gate: Arc, + fail: bool, + } + #[async_trait] + impl Scorer for ParallelScorer { + fn name(&self) -> &str { + self.name + } + async fn score(&self, _: ScorerArgs<'_, usize, usize>) -> anyhow::Result { + self.gate.wait().await; + if self.fail { + return Err(ApplicationFailure.into()); + } + Ok(vec![Score::new(self.name, 0.75)]) + } + } + #[tokio::test] + async fn scorer_failures_are_named_without_discarding_other_scores_or_stopping_parallel_work() { + let gate = Arc::new(tokio::sync::Barrier::new(2)); + let scorers: Vec>> = vec![ + Box::new(ParallelScorer { + name: "bad", + gate: gate.clone(), + fail: true, + }), + Box::new(ParallelScorer { + name: "good", + gate, + fail: false, + }), + ]; + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + CaseRunner::run( + Case { + input: 1, + ..Default::default() + }, + &EvalDefinition::new( + Box::new(crate::eval::task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), + scorers, + ) + .unwrap(), + None, + None, + ), + ) + .await + .unwrap(); + assert_eq!(result.output, Some(1)); + assert_eq!(result.scores[0].name, "good"); + let [CaseError::Scorer { name, source }] = result.errors.as_slice() else { + panic!("expected named scorer failure") + }; + assert_eq!(name, "bad"); + assert!(source.downcast_ref::().is_some()); + assert_eq!( + serde_json::to_value(&result.errors).unwrap(), + json!([{"stage":"scorer","name":"bad","message":"application failure"}]) + ); + let stats = crate::eval::EvalSummary::compute_stats(&[result]); + assert_eq!(stats[0].name, "good"); + assert_eq!(stats[0].mean, 0.75); + } +} diff --git a/src/eval/dataset.rs b/src/eval/dataset.rs deleted file mode 100644 index 8eb9b44..0000000 --- a/src/eval/dataset.rs +++ /dev/null @@ -1,196 +0,0 @@ -use async_trait::async_trait; -use futures::stream::Stream; - -use crate::Result; - -use super::types::Case; - -/// Iterator trait for evaluation test cases -/// -/// This trait allows different dataset sources (in-memory vectors, API-backed datasets, etc.) -/// to provide test cases to the evaluator. Follows the Go SDK pattern where Dataset is -/// an interface/trait rather than a concrete type. -#[async_trait] -pub trait Dataset: Send + Sync -where - I: Clone, - O: Clone, -{ - /// Get the next case from the dataset - async fn next(&mut self) -> Result>>; - - /// Get the dataset ID (if backed by Braintrust API) - fn id(&self) -> Option { - None - } - - /// Get the dataset version (if backed by Braintrust API) - fn version(&self) -> Option { - None - } -} - -/// In-memory dataset backed by a vector -pub struct VecDataset -where - I: Clone, - O: Clone, -{ - cases: Vec>, - index: usize, -} - -impl VecDataset -where - I: Clone, - O: Clone, -{ - pub fn new(cases: Vec>) -> Self { - Self { cases, index: 0 } - } -} - -#[async_trait] -impl Dataset for VecDataset -where - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - if self.index < self.cases.len() { - let case = self.cases[self.index].clone(); - self.index += 1; - Ok(Some(case)) - } else { - Ok(None) - } - } -} - -/// Implement Dataset for any sync iterator over Cases -/// -/// This allows using standard iterators as datasets without manual wrapping. -pub struct IterDataset -where - I: Clone, - O: Clone, - T: Iterator>, -{ - iter: T, -} - -impl IterDataset -where - I: Clone, - O: Clone, - T: Iterator>, -{ - pub fn new(iter: T) -> Self { - Self { iter } - } -} - -#[async_trait] -impl Dataset for IterDataset -where - T: Iterator> + Send + Sync, - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - Ok(self.iter.next()) - } -} - -/// Implement Dataset for async streams -pub struct StreamDataset -where - I: Clone, - O: Clone, - S: Stream> + Send + Unpin, -{ - stream: S, -} - -impl StreamDataset -where - I: Clone, - O: Clone, - S: Stream> + Send + Unpin, -{ - pub fn new(stream: S) -> Self { - Self { stream } - } -} - -#[async_trait] -impl Dataset for StreamDataset -where - S: Stream> + Send + Sync + Unpin, - I: Send + Sync + Clone, - O: Send + Sync + Clone, -{ - async fn next(&mut self) -> Result>> { - use futures::stream::StreamExt; - Ok(self.stream.next().await) - } -} - -/// Convenience conversion from `Vec` to boxed Dataset -impl From>> for Box> -where - I: Send + Sync + Clone + 'static, - O: Send + Sync + Clone + 'static, -{ - fn from(cases: Vec>) -> Self { - Box::new(VecDataset::new(cases)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_vec_dataset() { - let cases = vec![ - Case { - input: "hello", - expected: Some("HELLO"), - ..Default::default() - }, - Case { - input: "world", - expected: Some("WORLD"), - ..Default::default() - }, - ]; - - let mut dataset = VecDataset::new(cases); - - let case1 = dataset.next().await.unwrap(); - assert!(case1.is_some()); - assert_eq!(case1.unwrap().input, "hello"); - - let case2 = dataset.next().await.unwrap(); - assert!(case2.is_some()); - assert_eq!(case2.unwrap().input, "world"); - - let case3 = dataset.next().await.unwrap(); - assert!(case3.is_none()); - } - - #[tokio::test] - async fn test_vec_to_boxed_dataset() { - let cases = vec![Case { - input: "test", - expected: Some("TEST"), - ..Default::default() - }]; - - let mut dataset: Box> = cases.into(); - let case = dataset.next().await.unwrap(); - assert!(case.is_some()); - assert_eq!(case.unwrap().input, "test"); - } -} diff --git a/src/eval/definition.rs b/src/eval/definition.rs new file mode 100644 index 0000000..f4c48c8 --- /dev/null +++ b/src/eval/definition.rs @@ -0,0 +1,207 @@ +//! Reusable computation shared by independent evaluation runs. +use std::sync::Arc; + +use async_trait::async_trait; + +use super::{EvalConfigError, Scorer, ScorerArgs, Scores, Task}; + +/// A validated task and ordered set of scorers, independent of any client or run. +/// +/// Cloning shares the callbacks without cloning inputs or outputs. Callback +/// state is shared across runs; sources, hooks and results belong to each run. +/// Callbacks must therefore support concurrent invocation. +pub struct EvalDefinition { + inner: Arc>, +} + +struct DefinitionInner { + task: Arc>, + scorers: Vec>>, +} + +impl Clone for EvalDefinition { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } +} + +impl EvalDefinition { + /// Validate scorer names and retain the callbacks for reuse. + /// + /// Names are captured once, including for custom scorers whose `name()` + /// could change through interior state. Empty or duplicate names are rejected. + pub fn new( + task: Box>, + scorers: Vec>>, + ) -> Result { + let mut names = std::collections::HashSet::new(); + let mut validated: Vec>> = Vec::with_capacity(scorers.len()); + for (index, scorer) in scorers.into_iter().enumerate() { + let name = scorer.name().to_owned(); + if name.trim().is_empty() { + return Err(EvalConfigError::EmptyScorerName { index }); + } + if !names.insert(name.clone()) { + return Err(EvalConfigError::DuplicateScorerName { name }); + } + validated.push(Arc::new(NamedScorer { name, scorer })); + } + Ok(Self { + inner: Arc::new(DefinitionInner { + task: Arc::from(task), + scorers: validated, + }), + }) + } +} + +impl EvalDefinition { + /// Scorer names in definition order, fixed at construction. + pub fn scorer_names(&self) -> impl ExactSizeIterator { + self.inner.scorers.iter().map(|scorer| scorer.name()) + } + + pub(crate) fn task(&self) -> Arc> { + self.inner.task.clone() + } + + pub(crate) fn scorers(&self) -> &[Arc>] { + &self.inner.scorers + } +} + +struct NamedScorer { + name: String, + scorer: Box>, +} + +#[async_trait] +impl Scorer for NamedScorer { + fn is_blocking(&self) -> bool { + self.scorer.is_blocking() + } + + fn name(&self) -> &str { + &self.name + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result { + self.scorer.score_blocking(args) + } + + async fn score(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result { + self.scorer.score(args).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::{case_runner::CaseRunner, task, Case, CaseError}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + // No Clone, Serialize, equality or runtime is needed for computation. + struct Input(usize); + struct Output(usize); + + #[test] + fn cloned_definitions_share_callbacks_with_independent_case_state() { + let calls = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let definition = EvalDefinition::new( + Box::new(task({ + let calls = calls.clone(); + move |input: &Input, hooks: &crate::eval::TaskHooks<'_, Input, Output>| { + let barrier = barrier.clone(); + let calls = calls.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + hooks.add_tag(format!("case-{}", input.0)); + barrier.wait().await; + Ok(Output(input.0 + hooks.expected.unwrap().0)) + }) + } + })), + vec![], + ) + .unwrap(); + let other = definition.clone(); + drop(definition); + let first = other.clone(); + let case = |input, expected| Case { + input: Input(input), + expected: Some(Output(expected)), + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }; + let (one, two) = futures::executor::block_on(async { + futures::join!( + CaseRunner::run(case(1, 10), &first, None, None), + CaseRunner::run(case(2, 20), &other, None, None), + ) + }); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!(one.output.unwrap().0, 11); + assert_eq!(two.output.unwrap().0, 22); + assert_eq!(one.tags, Some(vec!["case-1".into()])); + assert_eq!(two.tags, Some(vec!["case-2".into()])); + // Reuse the same callbacks after both prior invocations have completed. + let (one, two) = futures::executor::block_on(async { + futures::join!( + CaseRunner::run(case(3, 30), &first, None, None), + CaseRunner::run(case(4, 40), &other, None, None), + ) + }); + assert_eq!(one.output.unwrap().0, 33); + assert_eq!(two.output.unwrap().0, 44); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } + + struct ChangingName(Arc); + #[async_trait] + impl Scorer for ChangingName { + fn name(&self) -> &str { + if self.0.load(Ordering::SeqCst) { + "changed" + } else { + "original" + } + } + async fn score(&self, _: ScorerArgs<'_, usize, usize>) -> anyhow::Result { + anyhow::bail!("scoring failed") + } + } + + #[test] + fn scorer_identity_is_captured_once_for_every_run() { + let changed = Arc::new(AtomicBool::new(false)); + let definition = EvalDefinition::new( + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), + vec![Box::new(ChangingName(changed.clone()))], + ) + .unwrap(); + changed.store(true, Ordering::SeqCst); + assert_eq!(definition.scorer_names().collect::>(), ["original"]); + for definition in [definition.clone(), definition] { + let result = futures::executor::block_on(CaseRunner::run( + Case { + input: 1, + ..Default::default() + }, + &definition, + None, + None, + )); + assert!( + matches!(&result.errors[..], [CaseError::Scorer {name, source}] if name == "original" && source.to_string() == "scoring failed") + ); + } + } +} diff --git a/src/eval/error.rs b/src/eval/error.rs new file mode 100644 index 0000000..37b11dc --- /dev/null +++ b/src/eval/error.rs @@ -0,0 +1,146 @@ +use serde::{ser::SerializeStruct, Serialize, Serializer}; + +/// A computation failure retains its stage and original application error. +/// Other scorers still run after a scorer failure; task failures skip scoring. +#[derive(Debug, thiserror::Error)] +pub enum CaseError { + #[error("task failed: {source}")] + Task { + #[source] + source: anyhow::Error, + }, + #[error("scorer {name} failed: {source}")] + Scorer { + name: String, + #[source] + source: anyhow::Error, + }, +} +impl Serialize for CaseError { + fn serialize(&self, serializer: S) -> Result { + let mut state = serializer.serialize_struct( + "CaseError", + match self { + Self::Task { .. } => 2, + Self::Scorer { .. } => 3, + }, + )?; + let (stage, source) = match self { + Self::Task { source } => ("task", source), + Self::Scorer { name, source } => { + state.serialize_field("name", name)?; + ("scorer", source) + } + }; + state.serialize_field("stage", stage)?; + state.serialize_field("message", &source.to_string())?; + state.end() + } +} + +/// Invalid evaluation definitions are rejected before source access or recording. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum EvalConfigError { + #[error("evaluation concurrency must be greater than zero")] + ZeroConcurrency, + #[error("recorded evaluations inherit their client runtime; omit the run runtime override")] + RecordedRuntime, + #[error("scorer at index {index} must have a nonblank name")] + EmptyScorerName { index: usize }, + #[error("duplicate scorer name: {name}")] + DuplicateScorerName { name: String }, +} + +/// Run-level failures retain their boundary and source location. +/// Case indices are zero-based; None denotes run-level recording work. +#[derive(Debug, thiserror::Error)] +pub enum EvalRunFailure { + #[error("failed to initialize evaluation runtime: {0}")] + Runtime(#[source] crate::BraintrustError), + #[error(transparent)] + Config(#[from] EvalConfigError), + #[error("case source failed at index {case_index}: {source}")] + Source { + case_index: usize, + #[source] + source: super::SourceError, + }, + #[error("recording failed for case {case_index:?}: {source}")] + Recording { + case_index: Option, + #[source] + source: crate::BraintrustError, + }, + #[error(transparent)] + Execution(#[from] crate::ExecutionError), + #[error(transparent)] + Admission(#[from] crate::EnqueueError), +} + +/// An incomplete run preserves finished case results and every run-level failure. +/// Case computation errors remain in those results rather than aborting the run. +/// Recording failures do not discard computed results or change their success counts. +pub struct EvalRunError { + pub failures: Vec, + pub partial: Option>>, +} +impl EvalRunError { + pub(crate) fn new(failure: EvalRunFailure) -> Self { + Self { + failures: vec![failure], + partial: None, + } + } +} +impl std::fmt::Debug for EvalRunError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EvalRunError") + .field("failures", &self.failures) + .field( + "completed_cases", + &self.partial.as_ref().map(|summary| summary.results.len()), + ) + .finish() + } +} +impl std::fmt::Display for EvalRunError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(first) = self.failures.first() { + write!(f, "{first}")?; + } + if self.failures.len() > 1 { + write!(f, " ({} run failures)", self.failures.len())?; + } + Ok(()) + } +} +impl std::error::Error for EvalRunError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.failures + .first() + .map(|failure| failure as &dyn std::error::Error) + } +} +impl From for EvalRunError { + fn from(error: EvalConfigError) -> Self { + Self::new(EvalRunFailure::Config(error)) + } +} +impl From for EvalRunError { + fn from(source: crate::BraintrustError) -> Self { + Self::new(EvalRunFailure::Recording { + case_index: None, + source, + }) + } +} +impl From for EvalRunError { + fn from(error: crate::ExecutionError) -> Self { + Self::new(EvalRunFailure::Execution(error)) + } +} +impl From for EvalRunError { + fn from(error: crate::EnqueueError) -> Self { + Self::new(EvalRunFailure::Admission(error)) + } +} diff --git a/src/eval/executor.rs b/src/eval/executor.rs new file mode 100644 index 0000000..98e4b3a --- /dev/null +++ b/src/eval/executor.rs @@ -0,0 +1,484 @@ +//! Bounded run scheduling and aggregation, independent of recording and output. +use async_trait::async_trait; +use futures::{ + future::{select, Either}, + stream::{FuturesUnordered, StreamExt}, +}; + +use super::{ + case_runner::CaseRunner, Case, EvalDefinition, EvalOpts, EvalResult, EvalRunError, + EvalRunFailure, EvalSummary, +}; +use crate::{BraintrustError, Result, SpanHandle}; + +/// An invocation owns its source and options; only the definition is shared. +pub(crate) struct EvalPlan { + name: String, + source: Box>, + definition: EvalDefinition, + concurrency: usize, +} +impl EvalPlan { + pub fn new(opts: EvalOpts) -> std::result::Result { + opts.validate()?; + Ok(Self { + name: opts.name, + source: opts.source, + definition: opts.definition, + concurrency: opts.max_concurrency, + }) + } +} + +#[async_trait] +pub(crate) trait EvalRecorder: Send + Sync { + async fn start(&self, index: usize, case: &Case) -> Result>>; + async fn finish(&self) -> Vec; +} + +#[async_trait] +pub(crate) trait CaseRecording: Send { + fn span(&self) -> Option<&SpanHandle>; + async fn finish(self: Box, result: &mut EvalResult) -> Vec; +} + +pub(crate) trait EvalReporter { + fn started(&mut self, name: &str); + fn case_finished(&mut self, index: usize, result: &EvalResult); + fn finished(&mut self, summary: &EvalSummary); +} + +struct CaseOutcome { + result: EvalResult, + failures: Vec, +} + +pub(crate) struct EvalExecutor; +impl EvalExecutor { + pub async fn run( + mut plan: EvalPlan, + recorder: Option<&dyn EvalRecorder>, + reporter: &mut impl EvalReporter, + binding: Option<&crate::internal::execution::ExecutionBinding>, + ) -> std::result::Result, EvalRunError> { + reporter.started(&plan.name); + let dataset = plan.source.provenance(); + let mut results = Vec::new(); + let mut failures = Vec::new(); + let mut case_count = 0; + let mut pending = FuturesUnordered::new(); + loop { + // Keep one source read alive while cases run. Sources may await + // case progress and need not support cancellation/restarting next(). + let mut next = Box::pin(plan.source.next()); + let next_case = loop { + if pending.is_empty() { + break next.await; + } + match select(next.as_mut(), pending.next()).await { + Either::Left((case, _)) => break case, + Either::Right((Some((index, outcome)), _)) => { + Self::collect(index, outcome, &mut results, &mut failures, reporter); + if !failures.is_empty() { + break Ok(None); + } + } + Either::Right((None, _)) => unreachable!("nonempty pending cases"), + } + }; + let case = match next_case { + Ok(Some(case)) => case, + Ok(None) => break, + Err(source) => { + failures.push(EvalRunFailure::Source { + case_index: case_count, + source, + }); + break; + } + }; + let index = case_count; + case_count += 1; + let definition = &plan.definition; + pending.push(async move { + ( + index, + Self::run_case(case, definition, recorder, index, binding).await, + ) + }); + if pending.len() >= plan.concurrency { + if let Some((index, outcome)) = pending.next().await { + Self::collect(index, outcome, &mut results, &mut failures, reporter); + } + if !failures.is_empty() { + break; + } + } + } + // Stop reading after failure, but finish all cases already accepted. + while let Some((index, outcome)) = pending.next().await { + Self::collect(index, outcome, &mut results, &mut failures, reporter); + } + if let Some(recorder) = recorder { + failures.extend(recorder.finish().await.into_iter().map(|source| { + EvalRunFailure::Recording { + case_index: None, + source, + } + })); + } + results.sort_unstable_by_key(|(index, _)| *index); + let results: Vec<_> = results.into_iter().map(|(_, result)| result).collect(); + let score_stats = EvalSummary::::compute_stats(&results); + let successful_cases = results + .iter() + .filter(|result| result.errors.is_empty()) + .count(); + let summary = EvalSummary { + name: plan.name, + results, + score_stats, + total_cases: case_count, + successful_cases, + failed_cases: case_count - successful_cases, + dataset, + }; + reporter.finished(&summary); + if failures.is_empty() { + Ok(summary) + } else { + Err(EvalRunError { + failures, + partial: Some(Box::new(summary)), + }) + } + } + + async fn run_case( + case: Case, + definition: &EvalDefinition, + recorder: Option<&dyn EvalRecorder>, + index: usize, + binding: Option<&crate::internal::execution::ExecutionBinding>, + ) -> CaseOutcome { + let mut failures = Vec::new(); + let recording = if let Some(recorder) = recorder { + match recorder.start(index, &case).await { + Ok(recording) => Some(recording), + Err(error) => { + failures.push(error); + None + } + } + } else { + None + }; + let span = recording.as_ref().and_then(|recording| recording.span()); + let mut result = CaseRunner::run(case, definition, span, binding).await; + if let Some(recording) = recording { + failures.extend(recording.finish(&mut result).await); + } + CaseOutcome { result, failures } + } + + fn collect( + index: usize, + outcome: CaseOutcome, + results: &mut Vec<(usize, EvalResult)>, + failures: &mut Vec, + reporter: &mut impl EvalReporter, + ) { + failures.extend( + outcome + .failures + .into_iter() + .map(|source| EvalRunFailure::Recording { + case_index: Some(index), + source, + }), + ); + reporter.case_finished(index, &outcome.result); + results.push((index, outcome.result)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::eval::{task, CaseSource, SourceError}; + use std::{ + cell::Cell, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + }; + + struct Input(usize); + struct Output(usize); + fn case(input: usize) -> Case { + Case { + input: Input(input), + expected: None, + metadata: None, + tags: None, + id: Some(input.to_string()), + xact_id: None, + created: None, + } + } + #[derive(Default)] + struct Observed { + cases: Vec, + finished: bool, + } + impl EvalReporter for Observed { + fn started(&mut self, _: &str) {} + fn case_finished(&mut self, index: usize, _: &EvalResult) { + self.cases.push(index); + } + fn finished(&mut self, _: &EvalSummary) { + self.finished = true; + } + } + struct WaitingSource { + cursor: Cell, + task_progress: Option>, + reads: Arc, + } + #[async_trait] + impl CaseSource for WaitingSource { + async fn next(&mut self) -> std::result::Result>, SourceError> { + self.reads.fetch_add(1, Ordering::SeqCst); + let index = self.cursor.get(); + self.cursor.set(index + 1); + match index { + 0 => Ok(Some(case(1))), + 1 => { + self.task_progress.take().unwrap().await.unwrap(); + Ok(Some(case(2))) + } + _ => anyhow::bail!("source stopped"), + } + } + } + + #[test] + fn local_run_drives_pending_source_and_drains_cases_without_runtime_or_serialization() { + let (progress, receiver) = futures::channel::oneshot::channel(); + let progress = Mutex::new(Some(progress)); + let reads = Arc::new(AtomicUsize::new(0)); + let definition = EvalDefinition::new( + Box::new(task(move |input: &Input, _hooks| { + if let Some(progress) = progress.lock().unwrap().take() { + progress.send(()).unwrap(); + } + Box::pin(async move { Ok(Output(input.0 * 2)) }) + })), + vec![], + ) + .unwrap(); + let plan = EvalPlan::new( + EvalOpts::builder() + .name("local".into()) + .source(Box::new(WaitingSource { + cursor: Cell::new(0), + task_progress: Some(receiver), + reads: reads.clone(), + })) + .definition(definition) + .max_concurrency(2) + .build(), + ) + .unwrap(); + let mut reporter = Observed::default(); + let error = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter, None)) + .err() + .expect("source should fail"); + assert!( + matches!(&error.failures[..], [EvalRunFailure::Source {case_index: 2, source}] if source.to_string() == "source stopped") + ); + let summary = error.partial.unwrap(); + assert_eq!( + reads.load(Ordering::SeqCst), + 3, + "pending next() must not restart" + ); + assert_eq!(summary.total_cases, 2); + assert_eq!(summary.successful_cases, 2); + assert_eq!( + summary + .results + .iter() + .map(|result| result.output.as_ref().unwrap().0) + .collect::>(), + [2, 4] + ); + assert!(summary + .results + .iter() + .all(|result| result.span_id.is_none())); + assert_eq!(reporter.cases, [0, 1]); + assert!(reporter.finished); + } + + #[test] + fn local_success_has_bounded_concurrency_and_source_order_without_data_trait_bounds() { + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let definition = EvalDefinition::new( + Box::new(task({ + let active = active.clone(); + let peak = peak.clone(); + move |input: &Input, _hooks| { + let active = active.clone(); + let peak = peak.clone(); + let barrier = barrier.clone(); + Box::pin(async move { + let count = active.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(count, Ordering::SeqCst); + barrier.wait().await; + active.fetch_sub(1, Ordering::SeqCst); + Ok(Output(input.0 * 2)) + }) + } + })), + vec![], + ) + .unwrap(); + for count in [4, 0] { + let plan = EvalPlan::new( + EvalOpts::builder() + .name("local-success".into()) + .source((0..count).map(case).collect::>().into()) + .definition(definition.clone()) + .max_concurrency(2) + .build(), + ) + .unwrap(); + let mut reporter = Observed::default(); + let summary = + futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter, None)) + .expect("local computation should succeed"); + assert_eq!(summary.total_cases, count); + assert_eq!(summary.successful_cases, count); + assert_eq!( + summary + .results + .iter() + .map(|result| result.output.as_ref().unwrap().0) + .collect::>(), + (0..count).map(|index| index * 2).collect::>() + ); + assert!(reporter.finished); + } + assert_eq!(active.load(Ordering::SeqCst), 0); + assert_eq!(peak.load(Ordering::SeqCst), 2); + } + + struct FailingRecorder { + finished: Arc, + } + struct FailingCase; + #[async_trait] + impl EvalRecorder for FailingRecorder { + async fn start( + &self, + index: usize, + _: &Case, + ) -> Result>> { + if index == 0 { + return Err(BraintrustError::Background("start failed".into())); + } + Ok(Box::new(FailingCase)) + } + async fn finish(&self) -> Vec { + self.finished.fetch_add(1, Ordering::SeqCst); + vec![BraintrustError::Background("run finish failed".into())] + } + } + #[async_trait] + impl CaseRecording for FailingCase { + fn span(&self) -> Option<&SpanHandle> { + None + } + async fn finish(self: Box, _: &mut EvalResult) -> Vec { + vec![BraintrustError::Background("case finish failed".into())] + } + } + #[test] + fn recorder_start_and_finish_errors_preserve_computation_and_finalize_run() { + let finished = Arc::new(AtomicUsize::new(0)); + let recorder = FailingRecorder { + finished: finished.clone(), + }; + let cases = (0..3) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>>(); + let plan = EvalPlan::new( + EvalOpts::builder() + .name("recorded".into()) + .source(cases.into()) + .definition( + EvalDefinition::new( + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(input + 10) }) + })), + vec![], + ) + .unwrap(), + ) + .max_concurrency(2) + .build(), + ) + .unwrap(); + let mut reporter = Observed::default(); + let error = futures::executor::block_on(EvalExecutor::run( + plan, + Some(&recorder), + &mut reporter, + None, + )) + .unwrap_err(); + assert_eq!(error.failures.len(), 3); + assert!(matches!( + &error.failures[0], + EvalRunFailure::Recording { + case_index: Some(0), + .. + } + )); + assert!(matches!( + &error.failures[1], + EvalRunFailure::Recording { + case_index: Some(1), + .. + } + )); + assert!(matches!( + &error.failures[2], + EvalRunFailure::Recording { + case_index: None, + .. + } + )); + let summary = error.partial.unwrap(); + assert_eq!( + summary + .results + .iter() + .map(|result| result.output) + .collect::>(), + [Some(10), Some(11)] + ); + assert_eq!(summary.total_cases, 2, "third case must remain unread"); + assert_eq!(summary.successful_cases, 2); + assert_eq!(summary.failed_cases, 0); + assert_eq!(finished.load(Ordering::SeqCst), 1); + assert!(reporter.finished); + } +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index f8d2d85..2b293c5 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -13,14 +13,12 @@ //! use braintrust_sdk_rust::*; //! use braintrust_sdk_rust::eval::*; //! -//! # async fn example() -> Result<()> { +//! # async fn example() -> anyhow::Result<()> { //! let client = BraintrustClient::builder() //! .api_key("sk-...") //! .build() //! .await?; //! -//! let evaluator = Evaluator::new(client); -//! //! let cases = vec![ //! Case { //! input: "hello".to_string(), @@ -29,35 +27,53 @@ //! }, //! ]; //! -//! let result = evaluator -//! .run( +//! // Reuse this definition with a fresh source and options for every run. +//! let definition = EvalDefinition::new( +//! Box::new(task(|input: &String, _hooks| Box::pin(async move { +//! Ok(input.to_uppercase()) +//! }))), +//! vec![Box::new(ExactMatch::new())], +//! )?; +//! let completion = run(EvalRequest::recorded( //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) -//! .dataset(cases.into()) -//! .task(Box::new(SimpleFnTask::new(|input: &String| { -//! Ok(input.to_uppercase()) -//! }))) -//! .scorers(vec![Box::new(ExactMatch::new())]) +//! .source(cases.into()) +//! .definition(definition.clone()) //! .build(), -//! ) -//! .await?; +//! client, +//! )); +//! let result = completion.await?; //! //! println!("Results: {:?}", result); //! # Ok(()) //! # } //! ``` +mod blocking; mod bt_runner; -mod dataset; +mod case_runner; +mod definition; +mod error; +mod executor; +mod recorder; +mod reporter; +mod run; mod runner; mod scorer; +mod source; mod task; mod types; +pub use definition::EvalDefinition; +pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; + // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; -pub use dataset::{Dataset, IterDataset, StreamDataset, VecDataset}; -pub use runner::{EvalOpts, Evaluator, SyncEvaluator}; -pub use scorer::{ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; -pub use task::{AsyncFnTask, SimpleAsyncFnTask, SimpleFnTask, SyncFnTask, Task}; +pub use run::{run, EvalRequest}; +pub use runner::{EvalOpts, Evaluator}; +pub use scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; +pub use source::{ + CaseSource, DatasetVersionRef, IterCaseSource, SourceError, StreamCaseSource, VecCaseSource, +}; +pub use task::{task, SimpleFnTask, SyncFnTask, Task}; pub use types::{Case, EvalResult, EvalSummary, Score, ScoreStats, ScorerArgs, Scores, TaskHooks}; diff --git a/src/eval/recorder.rs b/src/eval/recorder.rs new file mode 100644 index 0000000..4f8c6e4 --- /dev/null +++ b/src/eval/recorder.rs @@ -0,0 +1,167 @@ +//! Braintrust serialization and write completions, outside the run engine. +use super::{ + executor::{CaseRecording, EvalRecorder}, + Case, EvalResult, +}; +use crate::{ + BraintrustClient, BraintrustError, Completion, DeliveryReceipt, Result, SpanHandle, SpanLog, + WriteError, +}; +use async_trait::async_trait; +use serde::Serialize; +use serde_json::{Map, Value}; +use std::sync::Arc; + +type Write = Completion; + +pub(crate) struct BraintrustRecorder { + client: BraintrustClient, + submitter: Arc, + root: SpanHandle, + root_write: tokio::sync::Mutex>, +} +impl BraintrustRecorder { + pub async fn new( + client: BraintrustClient, + submitter: Arc, + name: String, + metadata: Option>, + tags: Option>, + ) -> Result { + let root = client + .span_builder_for_operation(submitter.clone()) + .await? + .span_type(crate::SpanType::Eval) + .build(); + let mut log = SpanLog::builder().name(format!("eval-{name}")); + if let Some(metadata) = metadata { + log = log.metadata(metadata); + } + if let Some(tags) = tags { + log = log.tags(tags); + } + let root_write = tokio::sync::Mutex::new(Some(root.log(log.build().unwrap()))); + Ok(Self { + client, + submitter, + root, + root_write, + }) + } +} + +struct BraintrustCase { + span: SpanHandle, + initial_write: Write, + failures: Vec, +} +fn serialize( + value: &T, + field: &str, + failures: &mut Vec, +) -> Option { + match serde_json::to_value(value) { + Ok(value) => Some(value), + Err(error) => { + failures.push(BraintrustError::Background(format!( + "failed to serialize evaluation {field}: {error}" + ))); + None + } + } +} +#[async_trait] +impl EvalRecorder + for BraintrustRecorder +{ + async fn start(&self, index: usize, case: &Case) -> Result>> { + let span = self + .client + .span_builder_for_operation(self.submitter.clone()) + .await? + .span_type(crate::SpanType::Eval) + .build(); + let mut failures = Vec::new(); + let mut log = SpanLog::builder().name(format!("case-{}", index + 1)); + if let Some(input) = serialize(&case.input, "input", &mut failures) { + log = log.input(input); + } + if let Some(expected) = &case.expected { + if let Some(expected) = serialize(expected, "expected", &mut failures) { + log = log.expected(expected); + } + } + if let Some(metadata) = &case.metadata { + log = log.metadata(metadata.clone()); + } + if let Some(tags) = &case.tags { + log = log.tags(tags.clone()); + } + let initial_write = span.log(log.build().unwrap()); + Ok(Box::new(BraintrustCase { + span, + initial_write, + failures, + })) + } + async fn finish(&self) -> Vec { + let end = self.root.end(); + let mut failures = Vec::new(); + if let Some(write) = self.root_write.lock().await.take() { + if let Err(error) = write.await { + failures.push(error.into()); + } + } + if let Err(error) = end.await { + failures.push(error.into()); + } + failures + } +} +#[async_trait] +impl CaseRecording + for BraintrustCase +{ + fn span(&self) -> Option<&SpanHandle> { + Some(&self.span) + } + async fn finish(mut self: Box, result: &mut EvalResult) -> Vec { + match self.span.export() { + Ok(export) => result.span_id = export.span_id, + Err(error) => self.failures.push(error), + } + let mut log = SpanLog::builder().scores( + result + .scores + .iter() + .map(|score| (score.name.clone(), score.score)) + .collect(), + ); + if let Some(output) = &result.output { + if let Some(output) = serialize(output, "output", &mut self.failures) { + log = log.output(output); + } + } + if let Some(metadata) = &result.metadata { + log = log.metadata(metadata.clone()); + } + if let Some(tags) = &result.tags { + log = log.tags(tags.clone()); + } + if !result.errors.is_empty() { + if let Some(errors) = serialize(&result.errors, "case errors", &mut self.failures) { + log = log.error(errors); + } + } + let final_write = self.span.log(log.build().unwrap()); + let end = self.span.end(); + // Observe every completion, even if serialization or an earlier write + // failed. The owned computation result is never discarded here. + for outcome in [self.initial_write.await, final_write.await, end.await] { + if let Err(error) = outcome { + self.failures.push(error.into()); + } + } + self.failures + } +} diff --git a/src/eval/reporter.rs b/src/eval/reporter.rs new file mode 100644 index 0000000..149fce5 --- /dev/null +++ b/src/eval/reporter.rs @@ -0,0 +1,44 @@ +//! Console presentation, separate from execution and recording. +use super::{executor::EvalReporter, EvalResult, EvalSummary}; + +pub(crate) struct TerminalReporter { + pub quiet: bool, +} +impl EvalReporter for TerminalReporter { + fn started(&mut self, name: &str) { + if !self.quiet { + println!("Running evaluation: {name}"); + } + } + fn case_finished(&mut self, index: usize, result: &EvalResult) { + if self.quiet { + return; + } + let scores = result + .scores + .iter() + .map(|score| format!("{}={:.3}", score.name, score.score)) + .collect::>() + .join(", "); + println!("Case {}: {}", index + 1, scores); + for error in &result.errors { + eprintln!("Case {}: {}", index + 1, error); + } + } + fn finished(&mut self, summary: &EvalSummary) { + if self.quiet { + return; + } + println!("\nEvaluation complete!"); + println!("Total cases: {}", summary.total_cases); + println!("Successful: {}", summary.successful_cases); + println!("Failed: {}", summary.failed_cases); + println!("\nScore summary:"); + for stat in &summary.score_stats { + println!( + " {}: mean={:.3}, min={:.3}, max={:.3}, count={}", + stat.name, stat.mean, stat.min, stat.max, stat.count + ); + } + } +} diff --git a/src/eval/run.rs b/src/eval/run.rs new file mode 100644 index 0000000..cb784ed --- /dev/null +++ b/src/eval/run.rs @@ -0,0 +1,141 @@ +//! Immediate submission of local and recorded evaluations. +use super::{ + executor::{EvalExecutor, EvalPlan}, + recorder::BraintrustRecorder, + reporter::TerminalReporter, + EvalConfigError, EvalOpts, EvalRunError, EvalRunFailure, EvalSummary, +}; +use crate::{internal::execution::ExecutionBinding, BraintrustClient, Completion}; +use serde::Serialize; + +type RunCompletion = Completion, EvalRunError>; +type SubmitRecorded = fn(BraintrustClient, EvalOpts) -> RunCompletion; + +/// An owned invocation of a reusable definition with its own source and options. +/// +/// Local requests need no client or serialization. Recorded requests require +/// serialization and inherit their client's execution binding and lifecycle. +pub struct EvalRequest { + options: EvalOpts, + target: Target, +} +enum Target { + Local, + Recorded { + client: BraintrustClient, + submit: SubmitRecorded, + }, +} +impl EvalRequest { + /// Run without login, HTTP or recording, on `options.runtime` or an owned runtime. + pub fn local(options: EvalOpts) -> Self { + Self { + options, + target: Target::Local, + } + } +} +impl EvalRequest { + /// Record with this client's runtime. Supplying `options.runtime` is an error. + pub fn recorded(options: EvalOpts, client: BraintrustClient) -> Self { + Self { + options, + target: Target::Recorded { + client, + submit: recorded::, + }, + } + } +} + +/// Submit an evaluation now; awaiting observes the admitted run's result. +/// Dropping the token does not cancel accepted computation or recording. +/// +/// ``` +/// use braintrust_sdk_rust::eval::{self, Case, EvalDefinition, EvalOpts, EvalRequest, SimpleFnTask}; +/// # fn main() -> anyhow::Result<()> { +/// let definition = EvalDefinition::new( +/// Box::new(SimpleFnTask::new(|input: &usize| Ok(input * 2))), vec![], +/// )?; +/// let options = EvalOpts::builder().name("local".into()) +/// .definition(definition).source(vec![Case { input: 3, ..Default::default() }].into()) +/// .quiet(true).build(); +/// let completion = eval::run(EvalRequest::local(options)); +/// assert!(completion.is_ok()); // admission already happened +/// let summary = futures::executor::block_on(async { completion.await })?; +/// assert_eq!(summary.results[0].output, Some(6)); +/// # Ok(()) +/// # } +/// ``` +pub fn run( + request: EvalRequest, +) -> RunCompletion { + if let Err(error) = request.options.validate() { + return Completion::rejected(error.into()); + } + match request.target { + Target::Local => { + let mut options = request.options; + let binding = match ExecutionBinding::new(options.runtime.take()) { + Ok(binding) => binding, + Err(source) => { + return Completion::rejected(EvalRunError::new(EvalRunFailure::Runtime(source))) + } + }; + let callbacks = binding.clone(); + binding.response(async move { + let mut reporter = TerminalReporter { + quiet: options.quiet, + }; + EvalExecutor::run( + EvalPlan::new(options)?, + None, + &mut reporter, + Some(&callbacks), + ) + .await + }) + } + Target::Recorded { client, submit } => { + if request.options.runtime.is_some() { + return Completion::rejected(EvalConfigError::RecordedRuntime.into()); + } + submit(client, request.options) + } + } +} + +fn recorded( + client: BraintrustClient, + options: EvalOpts, +) -> RunCompletion { + let binding = client.execution_binding(); + let callbacks = binding.clone(); + binding.response_try(|| { + let operation = client.span_operation()?; + Ok(async move { + let mut reporter = TerminalReporter { + quiet: options.quiet, + }; + let recorder = BraintrustRecorder::new( + client, + operation.submitter.clone(), + options.name.clone(), + options.metadata.clone(), + options.tags.clone(), + ) + .await?; + let result = EvalExecutor::run( + EvalPlan::new(options)?, + Some(&recorder), + &mut reporter, + Some(&callbacks), + ) + .await; + drop(recorder); + // Revoke write permission before publishing, including escaped spans. + drop(operation); + result + }) + }) +} diff --git a/src/eval/runner.rs b/src/eval/runner.rs index cecb373..87f9bc3 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -1,17 +1,12 @@ -use futures::stream::{FuturesUnordered, StreamExt}; -use serde::{de::DeserializeOwned, Serialize}; -use serde_json::{json, Map, Value}; +use serde::Serialize; +use serde_json::{Map, Value}; use std::marker::PhantomData; use std::sync::Arc; -use tokio::sync::Semaphore; -use crate::{BraintrustClient, Result, SpanHandle}; +use crate::BraintrustClient; use super::{ - dataset::Dataset, - scorer::Scorer, - task::Task, - types::{Case, EvalResult, EvalSummary, TaskHooks}, + source::CaseSource, types::EvalSummary, EvalConfigError, EvalDefinition, EvalRunError, }; /// Options for running an evaluation @@ -19,12 +14,10 @@ use super::{ pub struct EvalOpts { /// Name of the evaluation pub name: String, - /// Dataset providing test cases - pub dataset: Box>, - /// Task to execute on each input - pub task: Box>, - /// Scorers to evaluate outputs - pub scorers: Vec>>, + /// Owned source providing test cases + pub source: Box>, + /// Reusable task and ordered scorers + pub definition: EvalDefinition, /// Braintrust project to log results to pub project_name: Option, /// Optional experiment name for Braintrust logging @@ -39,9 +32,21 @@ pub struct EvalOpts { /// Suppress progress output (default: false) #[builder(default = false)] pub quiet: bool, + /// Explicit runtime for local runs. Recorded runs inherit their client binding. + pub runtime: Option, } -/// Async evaluator for running evaluations +impl EvalOpts { + /// Validate execution options without reading cases or creating recording work. + pub fn validate(&self) -> std::result::Result<(), EvalConfigError> { + if self.max_concurrency == 0 { + return Err(EvalConfigError::ZeroConcurrency); + } + Ok(()) + } +} + +/// Client-bound convenience facade for recorded evaluations. pub struct Evaluator { client: Arc, _phantom: PhantomData<(I, O)>, @@ -49,8 +54,8 @@ pub struct Evaluator { impl Evaluator where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, { /// Create a new evaluator pub fn new(client: BraintrustClient) -> Self { @@ -68,315 +73,25 @@ where } } - /// Run the evaluation - pub async fn run(&self, mut opts: EvalOpts) -> Result> { - if !opts.quiet { - println!("Running evaluation: {}", opts.name); - } - - // Create root span for the evaluation - let root_span = self - .client - .span_builder() - .await? - .span_type(crate::SpanType::Eval) - .build(); - - // Set name and metadata on the span - let mut log_builder = crate::SpanLog::builder().name(format!("eval-{}", opts.name)); - - if let Some(metadata) = &opts.metadata { - log_builder = log_builder.metadata(metadata.clone()); - } - if let Some(tags) = &opts.tags { - log_builder = log_builder.tags(tags.clone()); - } - - root_span.log(log_builder.build().unwrap()); - - // Get dataset ID if available - let dataset_id = opts.dataset.id(); - - // Collect all results - let mut results = Vec::new(); - let mut case_count = 0; - - // Create semaphore for concurrency control - let semaphore = Arc::new(Semaphore::new(opts.max_concurrency)); - - // Process cases with controlled concurrency - let mut futures = FuturesUnordered::new(); - - loop { - // Try to get next case - let case = match opts.dataset.next().await { - Ok(Some(case)) => case, - Ok(None) => break, - Err(e) => { - eprintln!("Error fetching case: {}", e); - break; - } - }; - - case_count += 1; - - // Clone what we need for the async task - let client = Arc::clone(&self.client); - let root_span = root_span.clone(); - let task = &opts.task; - let scorers = &opts.scorers; - let semaphore = Arc::clone(&semaphore); - let quiet = opts.quiet; - - // Acquire permit before spawning - let permit = semaphore.clone().acquire_owned().await.unwrap(); - - // Process this case - let future = async move { - let result = Self::process_case( - case, - task.as_ref(), - scorers, - &root_span, - &client, - case_count, - quiet, - ) - .await; - drop(permit); // Release permit - result - }; - - futures.push(future); - - // Collect completed futures if we have many pending - while futures.len() >= opts.max_concurrency { - if let Some(result) = futures.next().await { - results.push(result?); - } - } - } - - // Collect remaining futures - while let Some(result) = futures.next().await { - results.push(result?); - } - - // End root span - root_span.end().await?; - - // Compute statistics - let score_stats = EvalSummary::::compute_stats(&results); - let successful = results.iter().filter(|r| r.error.is_none()).count(); - let failed = results.len() - successful; - - if !opts.quiet { - println!("\nEvaluation complete!"); - println!("Total cases: {}", results.len()); - println!("Successful: {}", successful); - println!("Failed: {}", failed); - println!("\nScore summary:"); - for stat in &score_stats { - println!( - " {}: mean={:.3}, min={:.3}, max={:.3}, count={}", - stat.name, stat.mean, stat.min, stat.max, stat.count - ); - } - } - - Ok(EvalSummary { - name: opts.name, - results, - score_stats, - total_cases: case_count, - successful_cases: successful, - failed_cases: failed, - dataset_id, - }) - } - - /// Process a single case - async fn process_case( - case: Case, - task: &dyn Task, - scorers: &[Box>], - _root_span: &SpanHandle, - client: &BraintrustClient, - case_num: usize, - quiet: bool, - ) -> Result> { - let case_span = client - .span_builder() - .await? - .span_type(crate::SpanType::Eval) - .build(); - - // Build initial span log with name, input, and expected - let mut log_builder = crate::SpanLog::builder().name(format!("case-{}", case_num)); - - if let Ok(input_json) = serde_json::to_value(&case.input) { - log_builder = log_builder.input(input_json); - } - - if let Some(expected) = &case.expected { - if let Ok(expected_json) = serde_json::to_value(expected) { - log_builder = log_builder.expected(expected_json); - } - } - - // Add case metadata and tags - if let Some(metadata) = &case.metadata { - log_builder = log_builder.metadata(metadata.clone()); - } - if let Some(tags) = &case.tags { - log_builder = log_builder.tags(tags.clone()); - } - - case_span.log(log_builder.build().unwrap()); - - // Create hooks for the task - let hooks = TaskHooks::new(&case_span, case.expected.as_ref()); - - // Run the task - let output = match task.run(&case.input, &hooks).await { - Ok(output) => output, - Err(e) => { - // Log error to span - case_span.log( - crate::SpanLog::builder() - .error(json!(e.to_string())) - .build() - .unwrap(), - ); - case_span.end().await?; - - return Ok(EvalResult { - input: case.input.clone(), - output: None, // Task failed, no output - expected: case.expected.clone(), - scores: vec![], - metadata: case.metadata.clone(), - tags: case.tags.clone(), - error: Some(e.to_string()), - span_id: None, - }); - } - }; - - // Set output in span - if let Ok(output_json) = serde_json::to_value(&output) { - case_span.log( - crate::SpanLog::builder() - .output(output_json) - .build() - .unwrap(), - ); - } - - // Run scorers in parallel - let mut scorer_futures = FuturesUnordered::new(); - for scorer in scorers { - let scorer_args = super::types::ScorerArgs { - input: &case.input, - output: &output, - expected: case.expected.as_ref(), - metadata: case.metadata.as_ref(), - }; - - let future = async move { - match scorer.score(scorer_args).await { - Ok(scores) => scores, - Err(e) => { - eprintln!("Scorer {} failed: {}", scorer.name(), e); - vec![] - } - } - }; - - scorer_futures.push(future); - } - - // Collect all scores - let mut all_scores = vec![]; - while let Some(scores) = scorer_futures.next().await { - all_scores.extend(scores); - } - - // Add scores to span - let mut scores_map = std::collections::HashMap::new(); - for score in &all_scores { - scores_map.insert(score.name.clone(), score.score); - } - - // Get final metadata and tags from hooks - let final_metadata = hooks.get_metadata().or(case.metadata.clone()); - let final_tags = hooks.get_tags().or(case.tags.clone()); - - // Log scores and final metadata/tags - let mut final_log = crate::SpanLog::builder().scores(scores_map); - if let Some(meta) = &final_metadata { - final_log = final_log.metadata(meta.clone()); - } - if let Some(tags) = &final_tags { - final_log = final_log.tags(tags.clone()); - } - case_span.log(final_log.build().unwrap()); - - case_span.end().await?; - - if !quiet { - let score_str = all_scores - .iter() - .map(|s| format!("{}={:.3}", s.name, s.score)) - .collect::>() - .join(", "); - println!("Case {}: {}", case_num, score_str); - } - - Ok(EvalResult { - input: case.input.clone(), - output: Some(output.clone()), - expected: case.expected.clone(), - scores: all_scores, - metadata: final_metadata, - tags: final_tags, - error: None, - span_id: None, - }) - } -} - -/// Sync evaluator (wraps async evaluator with Runtime::block_on) -pub struct SyncEvaluator { - evaluator: Evaluator, - runtime: tokio::runtime::Runtime, -} - -impl SyncEvaluator -where - I: Serialize + DeserializeOwned + Send + Sync + Clone + 'static, - O: Serialize + DeserializeOwned + Send + Sync + PartialEq + Clone + 'static, -{ - /// Create a new sync evaluator - pub fn new(client: BraintrustClient) -> Result { - let runtime = tokio::runtime::Runtime::new() - .map_err(|e| crate::BraintrustError::Background(e.to_string()))?; - Ok(Self { - evaluator: Evaluator::new(client), - runtime, - }) - } - - /// Run the evaluation synchronously - pub fn run(&self, opts: EvalOpts) -> Result> { - self.runtime.block_on(self.evaluator.run(opts)) + /// Submit a recorded run immediately on the client binding. + /// Dropping the completion does not cancel accepted work. + pub fn run( + &self, + opts: EvalOpts, + ) -> crate::Completion, EvalRunError> { + super::run(super::EvalRequest::recorded( + opts, + self.client.as_ref().clone(), + )) } } #[cfg(test)] mod tests { use super::*; - use crate::eval::{dataset::VecDataset, scorer::ExactMatch, task::SimpleFnTask, types::Case}; + use crate::eval::{ + scorer::ExactMatch, source::VecCaseSource, task::SimpleFnTask, types::Case, Scorer, Task, + }; // Note: Full integration tests require a BraintrustClient setup // These tests verify type correctness and compilation @@ -389,16 +104,15 @@ mod tests { ..Default::default() }]; - let dataset: Box> = Box::new(VecDataset::new(cases)); + let dataset: Box> = Box::new(VecCaseSource::new(cases)); let task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let scorers: Vec>> = vec![Box::new(ExactMatch::new())]; let _opts = EvalOpts::builder() .name("test-eval".to_string()) - .dataset(dataset) - .task(task) - .scorers(scorers) + .source(dataset) + .definition(crate::eval::EvalDefinition::new(task, scorers).unwrap()) .build(); } } diff --git a/src/eval/scorer.rs b/src/eval/scorer.rs index 633a917..f051a9e 100644 --- a/src/eval/scorer.rs +++ b/src/eval/scorer.rs @@ -1,26 +1,67 @@ use async_trait::async_trait; +use futures::future::BoxFuture; use std::marker::PhantomData; -use crate::Result; +use anyhow::Result; use super::types::{Score, ScorerArgs, Scores}; /// Scorer trait for evaluating task outputs /// /// Scorers compare the task output against expected values and produce numeric scores. -/// Only async version is needed - SyncEvaluator uses this via Runtime::block_on(). +/// Callback errors use `anyhow::Result` to preserve application error causes. #[async_trait] pub trait Scorer: Send + Sync { /// Name of this scorer (used in results) fn name(&self) -> &str; + /// Whether evaluation runs should invoke this callback on their blocking pool. + /// Override together with `score_blocking` for synchronous work; async scorers use the defaults. + fn is_blocking(&self) -> bool { + false + } + + /// Synchronous callback body used when `is_blocking()` returns true. + /// The evaluation runner calls this on its selected blocking pool. + fn score_blocking(&self, _args: ScorerArgs<'_, I, O>) -> Result { + anyhow::bail!("scorer has no synchronous callback") + } + /// Score the output async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result; } +/// Supply type context for a named async scoring closure. The returned tuple +/// directly implements [`Scorer`]; names may be borrowed or owned strings. +pub fn scorer(name: N, callback: F) -> (N, F) +where + N: AsRef + Send + Sync, + F: for<'a> Fn(ScorerArgs<'a, I, O>) -> BoxFuture<'a, Result> + Send + Sync, +{ + (name, callback) +} + +#[async_trait] +impl Scorer for (N, F) +where + N: AsRef + Send + Sync, + F: for<'a> Fn(ScorerArgs<'a, I, O>) -> BoxFuture<'a, Result> + Send + Sync, + I: Send + Sync, + O: Send + Sync, +{ + fn name(&self) -> &str { + self.0.as_ref() + } + + async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result { + (self.1)(args).await + } +} + /// Wrapper for sync scoring functions /// -/// This allows sync functions to be used as scorers in both async and sync evaluators. +/// Evaluation runs dispatch this callback on their selected runtime's blocking pool. +/// Direct calls to `Scorer::score` execute the callback on the calling thread. pub struct SyncFnScorer where F: Fn(ScorerArgs<'_, I, O>) -> Result + Send + Sync, @@ -50,11 +91,19 @@ where I: Send + Sync, O: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> Result { (self.func)(args) } } @@ -88,11 +137,19 @@ where O: PartialEq + Send + Sync, I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> Result { let score = match args.expected { Some(expected) if expected == args.output => 1.0, Some(_) => 0.0, @@ -177,11 +234,19 @@ impl Scorer for LevenshteinScorer where I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, String>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, String>) -> Result { let expected = match args.expected { Some(e) => e, None => return Ok(vec![Score::new(self.name.clone(), 0.0)]), @@ -246,11 +311,19 @@ impl Scorer for NumericDiffScorer where I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, f64>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, f64>) -> Result { let expected = match args.expected { Some(e) => e, None => return Ok(vec![Score::new(self.name.clone(), 0.0)]), @@ -276,10 +349,11 @@ mod tests { async fn test_exact_match() { let scorer = ExactMatch::new(); + let expected = String::from("HELLO"); let args = ScorerArgs { input: &"hello".to_string(), output: &"HELLO".to_string(), - expected: Some(&"HELLO".to_string()), + expected: Some(&expected), metadata: None, }; @@ -287,10 +361,11 @@ mod tests { assert_eq!(scores.len(), 1); assert_eq!(scores[0].score, 1.0); + let expected = String::from("hello"); let args = ScorerArgs { input: &"hello".to_string(), output: &"HELLO".to_string(), - expected: Some(&"hello".to_string()), + expected: Some(&expected), metadata: None, }; @@ -302,10 +377,11 @@ mod tests { async fn test_levenshtein() { let scorer = LevenshteinScorer::new(); + let expected = String::from("hallo"); let args = ScorerArgs { input: &"test".to_string(), output: &"hello".to_string(), - expected: Some(&"hallo".to_string()), + expected: Some(&expected), metadata: None, }; diff --git a/src/eval/source.rs b/src/eval/source.rs new file mode 100644 index 0000000..025f2ee --- /dev/null +++ b/src/eval/source.rs @@ -0,0 +1,105 @@ +//! Owned, fallible case sources, independent of recording and presentation. +//! +//! ```no_run +//! use braintrust_sdk_rust::eval::{Case, CaseSource, IterCaseSource}; +//! +//! # async fn example() -> anyhow::Result<()> { +//! let cases = (0..3).map(|input| Ok(Case:: { +//! input, ..Default::default() +//! })); +//! let mut source: Box> = +//! Box::new(IterCaseSource::new(cases)); +//! assert_eq!(source.next().await?.unwrap().input, 0); +//! # Ok(()) +//! # } +//! ``` +use super::Case; +use async_trait::async_trait; +use futures::{Stream, StreamExt}; +use serde::Serialize; +use std::pin::Pin; + +/// Application errors from a source retain their original cause. +pub type SourceError = anyhow::Error; + +/// Dataset identity captured by a source and retained with its run result. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct DatasetVersionRef { + pub dataset_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub version: Option, +} + +/// One run owns its source and reads it sequentially. Sources need Send, +/// without requiring Sync, cloning user values, or serializing local cases. +#[async_trait] +pub trait CaseSource: Send { + async fn next(&mut self) -> Result>, SourceError>; + fn provenance(&self) -> Option { + None + } +} + +/// Move cases out of an in-memory vector without retaining consumed values. +pub struct VecCaseSource { + cases: std::vec::IntoIter>, +} +impl VecCaseSource { + pub fn new(cases: Vec>) -> Self { + Self { + cases: cases.into_iter(), + } + } +} +#[async_trait] +impl CaseSource for VecCaseSource { + async fn next(&mut self) -> Result>, SourceError> { + Ok(self.cases.next()) + } +} + +/// Adapt a fallible iterator. For infallible iterators, use `iter.map(Ok)`. +pub struct IterCaseSource { + iter: T, +} +impl IterCaseSource { + pub fn new(iter: T) -> Self { + Self { iter } + } +} +#[async_trait] +impl CaseSource for IterCaseSource +where + T: Iterator, SourceError>> + Send, +{ + async fn next(&mut self) -> Result>, SourceError> { + self.iter.next().transpose() + } +} + +/// Adapt a fallible stream, including streams that are not Unpin or Sync. +/// For infallible streams, use `stream.map(Ok)`. +pub struct StreamCaseSource { + stream: Pin>, +} +impl StreamCaseSource { + pub fn new(stream: S) -> Self { + Self { + stream: Box::pin(stream), + } + } +} +#[async_trait] +impl CaseSource for StreamCaseSource +where + S: Stream, SourceError>> + Send, +{ + async fn next(&mut self) -> Result>, SourceError> { + self.stream.next().await.transpose() + } +} +impl From>> for Box> { + fn from(cases: Vec>) -> Self { + Box::new(VecCaseSource::new(cases)) + } +} diff --git a/src/eval/task.rs b/src/eval/task.rs index d529068..c3ca4ee 100644 --- a/src/eval/task.rs +++ b/src/eval/task.rs @@ -2,16 +2,28 @@ use async_trait::async_trait; use futures::future::BoxFuture; use std::marker::PhantomData; -use crate::Result; +use anyhow::Result; use super::types::TaskHooks; /// Task trait for transforming inputs to outputs /// /// Tasks can be async functions, closures, or custom structs implementing this trait. -/// Only async version is needed - SyncEvaluator uses this via Runtime::block_on(). +/// Callback errors use `anyhow::Result` to preserve application error causes. #[async_trait] pub trait Task: Send + Sync { + /// Whether evaluation runs should invoke this callback on their blocking pool. + /// Override together with `run_blocking` for synchronous work; async tasks use the defaults. + fn is_blocking(&self) -> bool { + false + } + + /// Synchronous callback body used when `is_blocking()` returns true. + /// The evaluation runner calls this on its selected blocking pool. + fn run_blocking(&self, _input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { + anyhow::bail!("task has no synchronous callback") + } + /// Execute the task on the given input async fn run(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result; } @@ -19,7 +31,8 @@ pub trait Task: Send + Sync { /// Wrapper for sync task functions /// /// This allows sync functions to be used as tasks in both async and sync evaluators. -/// The function runs synchronously but is wrapped in an async trait implementation. +/// Evaluation runs dispatch this callback on their selected runtime's blocking pool. +/// Direct calls to `Task::run` execute the callback on the calling thread. pub struct SyncFnTask where F: Fn(&I, &TaskHooks<'_, I, O>) -> Result + Send + Sync, @@ -47,45 +60,42 @@ where I: Send + Sync, O: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + async fn run(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { - (self.func)(input, hooks) + self.run_blocking(input, hooks) } -} -/// Wrapper for async task functions that return BoxFuture -pub struct AsyncFnTask -where - F: for<'a> Fn(&'a I, &'a TaskHooks<'a, I, O>) -> BoxFuture<'a, Result> + Send + Sync, -{ - func: F, - _phantom: PhantomData<(I, O)>, + fn run_blocking(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { + (self.func)(input, hooks) + } } -impl AsyncFnTask +/// Supply type context for a reusable async task closure without wrapping it. +/// The returned callable directly implements [`Task`]. Its future may borrow +/// this invocation's input and hooks across await points. +pub fn task(callback: F) -> F where F: for<'a> Fn(&'a I, &'a TaskHooks<'a, I, O>) -> BoxFuture<'a, Result> + Send + Sync, { - pub fn new(func: F) -> Self { - Self { - func, - _phantom: PhantomData, - } - } + callback } #[async_trait] -impl Task for AsyncFnTask +impl Task for F where F: for<'a> Fn(&'a I, &'a TaskHooks<'a, I, O>) -> BoxFuture<'a, Result> + Send + Sync, I: Send + Sync, O: Send + Sync, { async fn run(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { - (self.func)(input, hooks).await + self(input, hooks).await } } -/// Simple wrapper for tasks that don't need hooks +/// Synchronous task without hooks, dispatched on the run's managed blocking pool. pub struct SimpleFnTask where F: Fn(&I) -> Result + Send + Sync, @@ -113,41 +123,16 @@ where I: Send + Sync, O: Send + Sync, { - async fn run(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { - (self.func)(input) + fn is_blocking(&self) -> bool { + true } -} - -/// Simple async task wrapper -pub struct SimpleAsyncFnTask -where - F: for<'a> Fn(&'a I) -> BoxFuture<'a, Result> + Send + Sync, -{ - func: F, - _phantom: PhantomData<(I, O)>, -} -impl SimpleAsyncFnTask -where - F: for<'a> Fn(&'a I) -> BoxFuture<'a, Result> + Send + Sync, -{ - pub fn new(func: F) -> Self { - Self { - func, - _phantom: PhantomData, - } + async fn run(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { + self.run_blocking(input, _hooks) } -} -#[async_trait] -impl Task for SimpleAsyncFnTask -where - F: for<'a> Fn(&'a I) -> BoxFuture<'a, Result> + Send + Sync, - I: Send + Sync, - O: Send + Sync, -{ - async fn run(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { - (self.func)(input).await + fn run_blocking(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { + (self.func)(input) } } @@ -159,7 +144,9 @@ mod tests { async fn test_simple_fn_task() { let task = SimpleFnTask::new(|input: &String| Ok(input.to_uppercase())); - // We can't test with hooks yet without a full setup, but we can verify the type compiles - let _: Box> = Box::new(task); + let task: Box> = Box::new(task); + let input = String::from("hello"); + let hooks = TaskHooks::local(None); + assert_eq!(task.run(&input, &hooks).await.unwrap(), "HELLO"); } } diff --git a/src/eval/types.rs b/src/eval/types.rs index acbaf0e..49b427c 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -2,15 +2,12 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::sync::{Arc, Mutex}; +use super::{CaseError, DatasetVersionRef}; use crate::SpanHandle; /// A single test case for evaluation #[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Case -where - I: Clone, - O: Clone, -{ +pub struct Case { /// The input to the task pub input: I, /// The expected output (optional) @@ -34,8 +31,7 @@ where impl Default for Case where - I: Default + Clone, - O: Clone, + I: Default, { fn default() -> Self { Self { @@ -118,9 +114,16 @@ pub struct EvalResult { /// Tags from the case and task #[serde(skip_serializing_if = "Option::is_none")] pub tags: Option>, - /// Error message if evaluation failed + /// Computation failures, with task/scorer identity and original causes. + #[serde(skip_serializing_if = "Vec::is_empty")] + pub errors: Vec, + /// Source case identity and dataset transaction information. + #[serde(skip_serializing_if = "Option::is_none")] + pub case_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub xact_id: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, + pub created: Option, /// Span ID for this evaluation #[serde(skip_serializing_if = "Option::is_none")] pub span_id: Option, @@ -152,30 +155,29 @@ pub struct EvalSummary { pub score_stats: Vec, /// Total number of cases evaluated pub total_cases: usize, - /// Number of successful evaluations + /// Cases without task/scorer errors, independently of recording status. pub successful_cases: usize, - /// Number of failed evaluations + /// Cases with task/scorer errors, independently of recording status. pub failed_cases: usize, - /// Dataset ID (if using Braintrust dataset) + /// Full dataset provenance captured before source consumption. #[serde(skip_serializing_if = "Option::is_none")] - pub dataset_id: Option, + pub dataset: Option, } impl EvalSummary { - /// Compute score statistics from results + /// Compute statistics from scores that were actually produced, including + /// successful scorers in cases where another scorer failed. pub fn compute_stats(results: &[EvalResult]) -> Vec { use std::collections::HashMap; let mut score_map: HashMap> = HashMap::new(); for result in results { - if result.error.is_none() { - for score in &result.scores { - score_map - .entry(score.name.clone()) - .or_default() - .push(score.score); - } + for score in &result.scores { + score_map + .entry(score.name.clone()) + .or_default() + .push(score.score); } } @@ -202,8 +204,8 @@ impl EvalSummary { /// Hooks provided to tasks for instrumentation and metadata pub struct TaskHooks<'a, I, O> { - /// Span handle for logging - pub span: &'a SpanHandle, + /// Optional span handle supplied by the recording adapter. + pub span: Option<&'a SpanHandle>, /// Expected output (if available) pub expected: Option<&'a O>, /// Mutable metadata map for the task to add context @@ -216,11 +218,39 @@ pub struct TaskHooks<'a, I, O> { impl<'a, I, O> TaskHooks<'a, I, O> { pub fn new(span: &'a SpanHandle, expected: Option<&'a O>) -> Self { + Self::for_case(Some(span), expected, None, None) + } + + /// Hooks for computation without a client, runtime or recording span. + pub fn local(expected: Option<&'a O>) -> Self { + Self::for_case(None, expected, None, None) + } + + pub(crate) fn for_case( + span: Option<&'a SpanHandle>, + expected: Option<&'a O>, + metadata: Option>, + tags: Option>, + ) -> Self { Self { span, expected, - metadata: Arc::new(Mutex::new(Map::new())), - tags: Arc::new(Mutex::new(Vec::new())), + metadata: Arc::new(Mutex::new(metadata.unwrap_or_default())), + tags: Arc::new(Mutex::new(tags.unwrap_or_default())), + _phantom: std::marker::PhantomData, + } + } + + pub(crate) fn rebind<'b>( + &self, + span: Option<&'b SpanHandle>, + expected: Option<&'b O>, + ) -> TaskHooks<'b, I, O> { + TaskHooks { + span, + expected, + metadata: self.metadata.clone(), + tags: self.tags.clone(), _phantom: std::marker::PhantomData, } } @@ -235,7 +265,10 @@ impl<'a, I, O> TaskHooks<'a, I, O> { /// Add a tag pub fn add_tag(&self, tag: impl Into) { if let Ok(mut tags) = self.tags.lock() { - tags.push(tag.into()); + let tag = tag.into(); + if !tags.contains(&tag) { + tags.push(tag); + } } } diff --git a/src/logger.rs b/src/logger.rs index b106e44..191f73f 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -611,9 +611,31 @@ impl BraintrustClient { .response(async move { client.span_builder_inner().await }) } + pub(crate) fn execution_binding(&self) -> ExecutionBinding { + self.inner.execution.clone() + } + + pub(crate) fn span_operation( + &self, + ) -> std::result::Result { + self.span_submitter().operation() + } + async fn span_builder_inner(&self) -> Result { + self.span_builder_for_operation(self.span_submitter()).await + } + + pub(crate) async fn span_builder_for_operation( + &self, + submitter: Arc, + ) -> Result { + if self.inner.login_skipped { + return Err(BraintrustError::InvalidConfig( + "recording requires login".into(), + )); + } let state = self.wait_for_login_state().await?; - let mut builder = crate::span::SpanBuilder::new(self.span_submitter(), state.scope()?); + let mut builder = crate::span::SpanBuilder::new(submitter, state.scope()?); if let Some(environment) = &self.inner.environment { builder = builder.span_origin_environment(environment.clone()); } diff --git a/tests/blocking_eval.rs b/tests/blocking_eval.rs new file mode 100644 index 0000000..b745a81 --- /dev/null +++ b/tests/blocking_eval.rs @@ -0,0 +1,165 @@ +use braintrust_sdk_rust::eval::*; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +struct Input(usize); +struct Output(usize); +fn options( + task: Box>, + scorers: Vec>>, +) -> EvalOpts { + EvalOpts::builder() + .name("blocking".into()) + .source( + vec![Case { + input: Input(3), + expected: Some(Output(6)), + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }] + .into(), + ) + .definition(EvalDefinition::new(task, scorers).unwrap()) + .quiet(true) + .runtime(tokio::runtime::Handle::current()) + .build() +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_task_uses_selected_blocking_pool_and_preserves_shared_hooks_without_clones() { + let worker = std::thread::current().id(); + let selected = tokio::runtime::Handle::current().id(); + let entered = Arc::new(tokio::sync::Notify::new()); + let started = entered.clone(); + let (release, waiting) = std::sync::mpsc::channel(); + let waiting = Mutex::new(waiting); + let task = SyncFnTask::new(move |input: &Input, hooks: &TaskHooks<'_, Input, Output>| { + assert_ne!(std::thread::current().id(), worker); + assert_eq!(tokio::runtime::Handle::current().id(), selected); + hooks.add_metadata("blocking", serde_json::json!(true)); + hooks.add_tag("from-blocking"); + started.notify_one(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(2))?; + Ok(Output(input.0 + hooks.expected.unwrap().0)) + }); + // An async case on the same single-thread executor must make progress while + // the synchronous task is blocked; it releases the callback through a channel. + let release = tokio::spawn(async move { + entered.notified().await; + release.send(()).unwrap(); + }); + let summary = run(EvalRequest::local(options(Box::new(task), vec![]))) + .await + .unwrap(); + release.await.unwrap(); + assert!(summary.results[0].errors.is_empty()); + assert_eq!(summary.results[0].output.as_ref().unwrap().0, 9); + assert_eq!( + summary.results[0].metadata.as_ref().unwrap()["blocking"], + true + ); + assert_eq!( + summary.results[0].tags.as_ref().unwrap(), + &["from-blocking"] + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_scorers_run_concurrently_and_keep_definition_order_without_clones() { + let worker = std::thread::current().id(); + let selected = tokio::runtime::Handle::current().id(); + let (entered, mut arrivals) = tokio::sync::mpsc::unbounded_channel(); + let mut releases = Vec::new(); + let mut scorers: Vec>> = Vec::new(); + for name in ["first", "second"] { + let entered = entered.clone(); + let (release, waiting) = std::sync::mpsc::channel(); + releases.push(release); + let waiting = Mutex::new(waiting); + scorers.push(Box::new(SyncFnScorer::new( + name, + move |args: ScorerArgs<'_, Input, Output>| { + assert_ne!(std::thread::current().id(), worker); + assert_eq!(tokio::runtime::Handle::current().id(), selected); + assert_eq!(args.input.0, 3); + assert_eq!(args.output.0, 6); + assert_eq!(args.expected.unwrap().0, 6); + entered.send(name).unwrap(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(2))?; + Ok(vec![Score::new(name, 1.0)]) + }, + ))); + } + let task = SimpleFnTask::new(|input: &Input| Ok(Output(input.0 * 2))); + let token = run(EvalRequest::local(options(Box::new(task), scorers))); + let mut names = vec![ + arrivals.recv().await.unwrap(), + arrivals.recv().await.unwrap(), + ]; + names.sort(); + assert_eq!(names, ["first", "second"]); + for release in releases.into_iter().rev() { + release.send(()).unwrap(); + } + let summary = token.await.unwrap(); + assert!(summary.results[0].errors.is_empty()); + assert_eq!( + summary.results[0] + .scores + .iter() + .map(|score| score.name.as_str()) + .collect::>(), + ["first", "second"] + ); +} + +#[tokio::test] +async fn synchronous_panics_become_case_errors_and_later_cases_continue() { + let opts = EvalOpts::builder() + .name("panic".into()) + .source( + (0..2) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>>() + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| { + if *input == 0 { + panic!("synchronous application panic"); + } + Ok(*input) + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let summary = run(EvalRequest::local(opts)).await.unwrap(); + assert_eq!(summary.failed_cases, 1); + assert_eq!(summary.successful_cases, 1); + let [CaseError::Task { source }] = summary.results[0].errors.as_slice() else { + panic!("expected task error") + }; + assert!(source + .downcast_ref::() + .unwrap() + .is_panic()); + assert_eq!(summary.results[1].output, Some(1)); +} diff --git a/tests/eager_eval.rs b/tests/eager_eval.rs new file mode 100644 index 0000000..f7a956b --- /dev/null +++ b/tests/eager_eval.rs @@ -0,0 +1,518 @@ +use braintrust_sdk_rust::{ + eval::*, BraintrustClient, EnqueueError, ExecutionError, SpanLog, WriteError, +}; +use std::{ + future::IntoFuture, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn options( + definition: EvalDefinition, + cases: Vec>, +) -> EvalOpts { + EvalOpts::builder() + .name("eager".into()) + .definition(definition) + .source(cases.into()) + .quiet(true) + .build() +} +fn numeric() -> EvalOpts { + options( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| Ok(input * 2))), + vec![], + ) + .unwrap(), + vec![Case { + input: 3, + ..Default::default() + }], + ) +} + +#[test] +fn local_runs_from_a_plain_thread_without_serialization_clone_or_debug() { + struct Input(usize); + struct Output(usize); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &Input| Ok(Output(input.0 * 2)))), + vec![], + ) + .unwrap(); + let token = run(EvalRequest::local(options( + definition, + vec![Case { + input: Input(3), + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }], + ))); + assert!(token.is_ok()); + let result = futures::executor::block_on(token.into_future()).unwrap(); + assert_eq!(result.results[0].output.as_ref().unwrap().0, 6); + assert!(result.results[0].span_id.is_none()); +} + +#[test] +fn dropping_the_observer_keeps_the_source_and_callbacks_running() { + let (done, observed) = std::sync::mpsc::channel(); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(move |input: &usize| { + done.send(*input).unwrap(); + Ok(*input) + })), + vec![], + ) + .unwrap(); + let token = run(EvalRequest::local(options( + definition, + vec![Case { + input: 7, + ..Default::default() + }], + ))); + assert!(token.is_ok()); + drop(token); + assert_eq!(observed.recv_timeout(Duration::from_secs(3)).unwrap(), 7); +} + +#[tokio::test] +async fn local_runtime_selection_is_explicit_and_never_uses_ambient_by_default() { + let ambient = tokio::runtime::Handle::current().id(); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &usize| { + Ok(tokio::runtime::Handle::current().id()) + })), + vec![], + ) + .unwrap(); + for explicit in [false, true] { + let mut opts = options( + definition.clone(), + vec![Case { + input: 0, + ..Default::default() + }], + ); + opts.runtime = explicit.then(tokio::runtime::Handle::current); + let summary = run(EvalRequest::local(opts)).await.unwrap(); + assert_eq!(summary.results[0].output == Some(ambient), explicit); + } +} + +#[test] +fn external_runtime_loss_before_first_poll_reports_worker_loss_without_fallback() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new({ + let calls = calls.clone(); + move |input: &usize| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(*input) + } + })), + vec![], + ) + .unwrap(); + let mut opts = options( + definition, + vec![Case { + input: 0, + ..Default::default() + }], + ); + opts.runtime = Some(runtime.handle().clone()); + let token = run(EvalRequest::local(opts)); + assert!(token.is_ok()); + assert_eq!(calls.load(Ordering::SeqCst), 0); + drop(runtime); + let error = futures::executor::block_on(token.into_future()).unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Execution(ExecutionError::WorkerLost)] + )); + assert_eq!(calls.load(Ordering::SeqCst), 0); +} + +#[test] +fn invalid_options_return_an_immediate_config_error() { + let mut opts = numeric(); + opts.max_concurrency = 0; + let token = run(EvalRequest::local(opts)); + assert!(token.is_err()); + let error = futures::executor::block_on(token.into_future()).unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Config(EvalConfigError::ZeroConcurrency)] + )); + assert!(error.partial.is_none()); +} + +async fn client(server: &MockServer) -> BraintrustClient { + for (endpoint, body) in [ + ( + "/api/apikey/login", + serde_json::json!({"org_info":[{"id":"org","name":"org"}]}), + ), + ( + "/api/project/register", + serde_json::json!({"project":{"id":"project"}}), + ), + ("/logs3", serde_json::json!({})), + ] { + Mock::given(method("POST")) + .and(path(endpoint)) + .respond_with(ResponseTemplate::new(200).set_body_json(body)) + .mount(server) + .await; + } + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .mount(server) + .await; + BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .app_url(server.uri()) + .api_url(server.uri()) + .api_key("token") + .default_project("project") + .blocking_login(true) + .build() + .await + .unwrap() +} + +#[tokio::test] +async fn recorded_run_admitted_before_shutdown_finishes_writes_and_revokes_escaped_spans() { + let server = MockServer::start().await; + let client = client(&server).await; + let escaped = Arc::new(Mutex::new(None)); + let definition = EvalDefinition::new( + Box::new(task({ + let escaped = escaped.clone(); + move |input: &usize, hooks| { + let escaped = escaped.clone(); + Box::pin(async move { + let child = hooks + .span + .unwrap() + .child_builder() + .unwrap() + .name("callback-child") + .build(); + child + .log( + SpanLog::builder() + .input(serde_json::json!(input)) + .build() + .unwrap(), + ) + .await + .unwrap(); + *escaped.lock().unwrap() = Some(child); + Ok(*input) + }) + } + })), + vec![], + ) + .unwrap(); + let token = run(EvalRequest::recorded( + options( + definition, + vec![Case { + input: 7, + ..Default::default() + }], + ), + client.clone(), + )); + assert!(token.is_ok()); + let shutdown = client.shutdown(Duration::from_secs(5)); + let rejected = run(EvalRequest::recorded(numeric(), client.clone())); + assert!(rejected.is_err()); + assert!(matches!( + rejected.await.unwrap_err().failures.as_slice(), + [EvalRunFailure::Admission(EnqueueError::Closed)] + )); + let (result, shutdown) = tokio::join!(token, shutdown); + let result = result.unwrap(); + shutdown.unwrap(); + assert_eq!(result.results[0].output, Some(7)); + assert!(result.results[0].span_id.is_some()); + let escaped = escaped.lock().unwrap().take().unwrap(); + assert!(matches!( + escaped.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn recorded_runtime_override_is_rejected_and_dropped_observer_still_drains() { + let server = MockServer::start().await; + let client = client(&server).await; + let mut opts = numeric(); + opts.runtime = braintrust_sdk_rust::current_runtime(); + let rejected = run(EvalRequest::recorded(opts, client.clone())); + assert!(rejected.is_err()); + assert!(matches!( + rejected.await.unwrap_err().failures.as_slice(), + [EvalRunFailure::Config(EvalConfigError::RecordedRuntime)] + )); + let calls = Arc::new(AtomicUsize::new(0)); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new({ + let calls = calls.clone(); + move |input: &usize| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(*input) + } + })), + vec![], + ) + .unwrap(); + drop(run(EvalRequest::recorded( + options( + definition, + vec![Case { + input: 1, + ..Default::default() + }], + ), + client.clone(), + ))); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn callback_panic_retires_recorded_run_and_revokes_escaped_permission() { + let server = MockServer::start().await; + let client = client(&server).await; + let escaped = Arc::new(Mutex::new(None)); + let definition: EvalDefinition = EvalDefinition::new( + Box::new(task({ + let escaped = escaped.clone(); + move |_: &usize, hooks| { + *escaped.lock().unwrap() = + Some(hooks.span.unwrap().child_builder().unwrap().build()); + Box::pin(async { panic!("application callback panic") }) + } + })), + vec![], + ) + .unwrap(); + let token = run(EvalRequest::recorded( + options( + definition, + vec![Case { + input: 0, + ..Default::default() + }], + ), + client.clone(), + )); + let error = token.await.unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Execution(ExecutionError::WorkerLost)] + )); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + let escaped = escaped.lock().unwrap().take().unwrap(); + assert!(matches!( + escaped.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn shutdown_deadline_cancels_pending_run_and_revokes_escaped_permission() { + let server = MockServer::start().await; + let client = client(&server).await; + let escaped = Arc::new(Mutex::new(None)); + let entered = Arc::new(tokio::sync::Notify::new()); + let definition: EvalDefinition = EvalDefinition::new( + Box::new(task({ + let escaped = escaped.clone(); + let entered = entered.clone(); + move |_: &usize, hooks| { + *escaped.lock().unwrap() = + Some(hooks.span.unwrap().child_builder().unwrap().build()); + entered.notify_one(); + Box::pin(std::future::pending()) + } + })), + vec![], + ) + .unwrap(); + let token = run(EvalRequest::recorded( + options( + definition, + vec![Case { + input: 0, + ..Default::default() + }], + ), + client.clone(), + )); + entered.notified().await; + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(braintrust_sdk_rust::ShutdownError::Execution( + ExecutionError::TimedOut + )) + )); + let error = token.await.unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Execution(ExecutionError::TimedOut)] + )); + let escaped = escaped.lock().unwrap().take().unwrap(); + assert!(matches!( + escaped.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn running_blocking_callback_keeps_owned_data_after_timeout_but_loses_write_permission() { + #[derive(serde::Serialize)] + struct OwnedInput { + value: usize, + #[serde(skip)] + drops: Arc, + #[serde(skip)] + dropped: Arc, + } + impl Drop for OwnedInput { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + self.dropped.notify_one(); + } + } + let server = MockServer::start().await; + let client = client(&server).await; + let entered = Arc::new(tokio::sync::Notify::new()); + let drops = Arc::new(AtomicUsize::new(0)); + let dropped = Arc::new(tokio::sync::Notify::new()); + let (release, waiting) = std::sync::mpsc::channel(); + let waiting = Mutex::new(waiting); + let (report, observed) = tokio::sync::oneshot::channel(); + let report = Mutex::new(Some(report)); + let callback = SyncFnTask::new({ + let entered = entered.clone(); + move |input: &OwnedInput, hooks: &TaskHooks<'_, OwnedInput, usize>| { + entered.notify_one(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(3))?; + let rejected = hooks.span.unwrap().log(SpanLog::default()).is_err(); + report + .lock() + .unwrap() + .take() + .unwrap() + .send((input.value, rejected)) + .unwrap(); + Ok(input.value) + } + }); + let opts = options( + EvalDefinition::new(Box::new(callback), vec![]).unwrap(), + vec![Case { + input: OwnedInput { + value: 17, + drops: drops.clone(), + dropped: dropped.clone(), + }, + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }], + ); + let token = run(EvalRequest::recorded(opts, client.clone())); + entered.notified().await; + assert!(client.shutdown(Duration::from_millis(20)).await.is_err()); + assert!(matches!( + token + .await + .err() + .expect("run must time out") + .failures + .as_slice(), + [EvalRunFailure::Execution(ExecutionError::TimedOut)] + )); + assert_eq!( + drops.load(Ordering::SeqCst), + 0, + "running callback still owns the input" + ); + release.send(()).unwrap(); + assert_eq!(observed.await.unwrap(), (17, true)); + dropped.notified().await; + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_callback_can_wait_for_recording_without_blocking_delivery() { + let server = MockServer::start().await; + let client = client(&server).await; + let worker = std::thread::current().id(); + let callback = SyncFnTask::new(move |input: &usize, hooks: &TaskHooks<'_, usize, usize>| { + assert_ne!(std::thread::current().id(), worker); + futures::executor::block_on( + hooks + .span + .unwrap() + .log( + SpanLog::builder() + .metadata( + serde_json::json!({"sync_write": true}) + .as_object() + .unwrap() + .clone(), + ) + .build() + .unwrap(), + ) + .into_future(), + )?; + Ok(*input) + }); + let token = run(EvalRequest::recorded( + options( + EvalDefinition::new(Box::new(callback), vec![]).unwrap(), + vec![Case { + input: 9, + ..Default::default() + }], + ), + client.clone(), + )); + let shutdown = client.shutdown(Duration::from_secs(5)); + let (result, shutdown) = tokio::join!(token, shutdown); + assert_eq!(result.unwrap().results[0].output, Some(9)); + shutdown.unwrap(); +} diff --git a/tests/eval_callables.rs b/tests/eval_callables.rs new file mode 100644 index 0000000..5ce314e --- /dev/null +++ b/tests/eval_callables.rs @@ -0,0 +1,235 @@ +use async_trait::async_trait; +use braintrust_sdk_rust::{ + eval::{ + scorer, task, Case, CaseSource, EvalConfigError, EvalDefinition, EvalOpts, Evaluator, + ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, + }, + BraintrustClient, +}; +use futures::future::BoxFuture; +use std::sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, +}; + +// No Clone, serde or equality implementations are needed to adapt callables. +struct Input(usize); +struct Output(usize); + +fn borrowed_task<'a>( + input: &'a Input, + hooks: &'a TaskHooks<'a, Input, Output>, +) -> BoxFuture<'a, anyhow::Result> { + Box::pin(async move { + tokio::task::yield_now().await; + Ok(Output(input.0 + hooks.expected.unwrap().0)) + }) +} +fn borrowed_score<'a>( + args: ScorerArgs<'a, Input, Output>, +) -> BoxFuture<'a, anyhow::Result> { + Box::pin(async move { + tokio::task::yield_now().await; + Ok(vec![Score::new( + "borrowed", + (args.output.0 - args.input.0) as f64, + )]) + }) +} + +#[tokio::test] +async fn functions_and_plain_named_tuples_borrow_values_across_await() { + let task: Box> = Box::new(borrowed_task); + let scorer: Box> = Box::new(("function", borrowed_score)); + let input = Input(4); + let expected = Output(3); + let hooks = TaskHooks::local(Some(&expected)); + let output = task.run(&input, &hooks).await.unwrap(); + assert_eq!(output.0, 7); + assert_eq!(scorer.name(), "function"); + let scores = scorer + .score(ScorerArgs { + input: &input, + output: &output, + expected: Some(&expected), + metadata: None, + }) + .await + .unwrap(); + assert_eq!(scores[0].score, 3.0); +} + +#[tokio::test] +async fn captured_closure_is_reusable_concurrently_without_wrappers_or_value_clones() { + let calls = Arc::new(AtomicUsize::new(0)); + let observed = calls.clone(); + let gate = Arc::new(tokio::sync::Barrier::new(2)); + let task: Arc> = Arc::new(task(move |input: &Input, hooks| { + let calls = calls.clone(); + let gate = gate.clone(); + Box::pin(async move { + gate.wait().await; + calls.fetch_add(1, Ordering::SeqCst); + hooks.add_tag("called"); + Ok(Output(input.0 * 2)) + }) + })); + let one = Input(1); + let two = Input(2); + let first = TaskHooks::local(None); + let second = TaskHooks::local(None); + let (one, two) = tokio::time::timeout(std::time::Duration::from_secs(2), async { + tokio::join!(task.run(&one, &first), task.run(&two, &second)) + }) + .await + .unwrap(); + assert_eq!(one.unwrap().0, 2); + assert_eq!(two.unwrap().0, 4); + assert_eq!(observed.load(Ordering::SeqCst), 2); + assert_eq!(first.get_tags().unwrap(), ["called"]); + assert_eq!(second.get_tags().unwrap(), ["called"]); +} + +#[tokio::test] +async fn scorer_helpers_accept_borrowed_and_owned_names_and_captured_state() { + let local_name = String::from("borrowed-name"); + let borrowed = scorer( + local_name.as_str(), + |args: ScorerArgs<'_, Input, Output>| { + Box::pin(async move { + tokio::task::yield_now().await; + Ok(vec![Score::new("value", args.input.0 as f64)]) + }) + }, + ); + let calls = Arc::new(AtomicUsize::new(0)); + let observed = calls.clone(); + let owned = scorer( + String::from("owned-name"), + move |args: ScorerArgs<'_, Input, Output>| { + let calls = calls.clone(); + Box::pin(async move { + tokio::task::yield_now().await; + calls.fetch_add(1, Ordering::SeqCst); + Ok(vec![Score::new("value", args.output.0 as f64)]) + }) + }, + ); + let input = Input(2); + let output = Output(4); + assert_eq!(borrowed.name(), "borrowed-name"); + assert_eq!(owned.name(), "owned-name"); + for callback in [&borrowed as &dyn Scorer, &owned] { + assert_eq!( + callback + .score(ScorerArgs { + input: &input, + output: &output, + expected: None, + metadata: None + }) + .await + .unwrap() + .len(), + 1 + ); + } + assert_eq!(observed.load(Ordering::SeqCst), 1); +} + +struct UntouchedSource(Arc); +#[async_trait] +impl CaseSource for UntouchedSource { + async fn next(&mut self) -> anyhow::Result>> { + self.0.store(true, Ordering::SeqCst); + panic!("invalid plans must not read their source"); + } +} + +#[tokio::test] +async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_access() { + let server = wiremock::MockServer::start().await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let evaluator = Evaluator::new(client); + let accessed = Arc::new(AtomicBool::new(false)); + for (concurrency, names, expected) in [ + (0, vec!["one"], EvalConfigError::ZeroConcurrency), + ( + 1, + vec![" \t"], + EvalConfigError::EmptyScorerName { index: 0 }, + ), + ( + 1, + vec!["same", "same"], + EvalConfigError::DuplicateScorerName { + name: "same".into(), + }, + ), + ] { + let scorers: Vec>> = names + .into_iter() + .map(|name| { + Box::new(scorer(name, |_args| Box::pin(async { Ok(vec![]) }))) + as Box> + }) + .collect(); + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + scorers, + ); + if concurrency != 0 { + assert!(matches!(definition, Err(actual) if actual == expected)); + continue; + } + let opts = + EvalOpts::builder() + .name("invalid".into()) + .source(Box::new(UntouchedSource(accessed.clone())) + as Box>) + .definition(definition.unwrap()) + .max_concurrency(concurrency) + .quiet(true) + .build(); + assert_eq!(opts.validate(), Err(expected.clone())); + assert!( + matches!(evaluator.run(opts).await, Err(error) if matches!(&error.failures[0], braintrust_sdk_rust::eval::EvalRunFailure::Config(actual) if actual == &expected) && error.partial.is_none()) + ); + } + assert!(!accessed.load(Ordering::SeqCst)); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn duplicate_names_are_rejected_across_custom_tuple_and_builtin_scorers() { + struct Custom; + #[async_trait] + impl Scorer for Custom { + fn name(&self) -> &str { + "exact_match" + } + async fn score(&self, _: ScorerArgs<'_, usize, usize>) -> anyhow::Result { + Ok(vec![]) + } + } + for callback in [ + Box::new(Custom) as Box>, + Box::new(scorer(String::from("exact_match"), |_args| { + Box::pin(async { Ok(vec![]) }) + })), + ] { + let definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + vec![Box::new(ExactMatch::new()), callback], + ); + assert!( + matches!(definition, Err(EvalConfigError::DuplicateScorerName { name }) if name == "exact_match") + ); + } +} diff --git a/tests/eval_sources.rs b/tests/eval_sources.rs new file mode 100644 index 0000000..befff5b --- /dev/null +++ b/tests/eval_sources.rs @@ -0,0 +1,113 @@ +use braintrust_sdk_rust::eval::{ + Case, CaseSource, IterCaseSource, StreamCaseSource, VecCaseSource, +}; +use std::{ + cell::Cell, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; + +fn case(input: I) -> Case { + Case { + input, + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + } +} +struct OwnedInput { + value: Cell, + drops: Arc, +} +impl Drop for OwnedInput { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + } +} + +#[test] +fn vectors_move_nonclone_nonsync_cases_and_release_consumed_values() { + futures::executor::block_on(async { + let drops = Arc::new(AtomicUsize::new(0)); + let cases: Vec>> = (0..2) + .map(|value| { + case(OwnedInput { + value: Cell::new(value), + drops: drops.clone(), + }) + }) + .collect(); + let mut source: Box>> = cases.into(); + let first = source.next().await.unwrap().unwrap(); + assert_eq!(first.input.value.get(), 0); + drop(first); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "the source must release consumed cases" + ); + drop(source); + assert_eq!(drops.load(Ordering::SeqCst), 2); + }); +} + +#[derive(Debug, thiserror::Error)] +#[error("read failed")] +struct ReadFailure; + +#[tokio::test] +async fn fallible_iterators_keep_error_causes_and_do_not_require_sync() { + type LocalCase = Case, Cell>; + let items: Vec> = + vec![Ok(case(Cell::new(7))), Err(ReadFailure.into())]; + let mut source: Box, Cell>> = + Box::new(IterCaseSource::new(items.into_iter())); + assert_eq!(source.next().await.unwrap().unwrap().input.get(), 7); + assert!(source + .next() + .await + .unwrap_err() + .downcast_ref::() + .is_some()); + assert!(source.next().await.unwrap().is_none()); +} + +#[tokio::test] +async fn pinned_streams_allow_nonunpin_nonsync_state_and_owned_error_causes() { + let value = Cell::new(8); + let stream = futures::stream::once(async move { + value.set(value.get() + 1); + tokio::task::yield_now().await; + Ok::, Cell>, anyhow::Error>(case(value)) + }); + let mut source: Box, Cell>> = + Box::new(StreamCaseSource::new(stream)); + assert_eq!(source.next().await.unwrap().unwrap().input.get(), 9); + assert!(source.next().await.unwrap().is_none()); + let stream = + futures::stream::once(async { Err::, anyhow::Error>(ReadFailure.into()) }); + let mut source = StreamCaseSource::new(stream); + assert!(source + .next() + .await + .unwrap_err() + .downcast_ref::() + .is_some()); +} + +#[tokio::test] +async fn empty_and_infallible_sources_end_without_synthetic_cases() { + let mut empty = VecCaseSource::::new(vec![]); + assert!(empty.next().await.unwrap().is_none()); + assert!(empty.next().await.unwrap().is_none()); + let mut source = + IterCaseSource::new((0..2).map(|value| Ok::<_, anyhow::Error>(case::<_, usize>(value)))); + assert_eq!(source.next().await.unwrap().unwrap().input, 0); + assert_eq!(source.next().await.unwrap().unwrap().input, 1); + assert!(source.next().await.unwrap().is_none()); +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index b07eb0d..68e9a42 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -1,20 +1,49 @@ -use braintrust_sdk_rust::eval::Dataset; +use braintrust_sdk_rust::eval::CaseSource; use braintrust_sdk_rust::eval::*; use braintrust_sdk_rust::*; -// Note: Full integration tests require Braintrust credentials. -// These tests verify type correctness and basic functionality. -// To run integration tests, set up credentials and remove #[ignore] - -#[tokio::test] -#[ignore] // Requires Braintrust credentials -async fn test_basic_eval_flow() { +async fn recording_client() -> (BraintrustClient, wiremock::MockServer) { + use serde_json::json; + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + for (endpoint, response) in [ + ( + "/api/apikey/login", + json!({"org_info":[{"id":"org","name":"org"}]}), + ), + ("/api/project/register", json!({"project":{"id":"project"}})), + ("/logs3", json!({})), + ] { + Mock::given(method("POST")) + .and(path(endpoint)) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .mount(&server) + .await; + } + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; let client = BraintrustClient::builder() + .app_url(server.uri()) + .api_url(server.uri()) + .default_project("project") + .blocking_login(true) .api_key("test-key") .build() .await .expect("Failed to create client"); + (client, server) +} + +#[tokio::test] +async fn test_basic_eval_flow() { + let (client, server) = recording_client().await; let evaluator = Evaluator::new(client); let cases = vec![ @@ -30,16 +59,29 @@ async fn test_basic_eval_flow() { }, ]; - let task = SimpleFnTask::new(|input: &String| Ok(input.to_uppercase())); - let scorers: Vec>> = vec![Box::new(ExactMatch::new())]; + let callback = task(|input: &String, _hooks| Box::pin(async move { Ok(input.to_uppercase()) })); + let scorers: Vec>> = vec![Box::new(scorer( + "match", + |args: ScorerArgs<'_, String, String>| { + Box::pin(async move { + Ok(vec![Score::new( + "match", + if args.expected == Some(args.output) { + 1.0 + } else { + 0.0 + }, + )]) + }) + }, + ))]; let result = evaluator .run( EvalOpts::builder() .name("test-basic-eval".to_string()) - .dataset(cases.into()) - .task(Box::new(task)) - .scorers(scorers) + .source(cases.into()) + .definition(EvalDefinition::new(Box::new(callback), scorers).unwrap()) .quiet(true) .build(), ) @@ -48,6 +90,14 @@ async fn test_basic_eval_flow() { assert_eq!(result.total_cases, 2); assert_eq!(result.successful_cases, 2); + assert!(result + .results + .iter() + .all(|case| case.span_id.is_some() && case.errors.is_empty())); + let requests = server.received_requests().await.unwrap(); + assert!(requests + .iter() + .any(|request| request.url.path() == "/logs3")); } #[test] @@ -59,7 +109,7 @@ fn test_eval_types_compile() { ..Default::default() }]; - let _dataset: Box> = cases.into(); + let _dataset: Box> = cases.into(); let _task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let _scorers: Vec>> = vec![Box::new(ExactMatch::new())]; @@ -89,10 +139,11 @@ async fn test_scorers() { // ExactMatch let scorer = ExactMatch::new(); + let expected = String::from("TEST"); let args = ScorerArgs { input: &"test".to_string(), output: &"TEST".to_string(), - expected: Some(&"TEST".to_string()), + expected: Some(&expected), metadata: None, }; let scores = scorer.score(args).await.unwrap(); @@ -100,10 +151,11 @@ async fn test_scorers() { // Levenshtein let scorer = LevenshteinScorer::new(); + let expected = String::from("hallo"); let args = ScorerArgs { input: &"test".to_string(), output: &"hello".to_string(), - expected: Some(&"hallo".to_string()), + expected: Some(&expected), metadata: None, }; let scores = scorer.score(args).await.unwrap(); @@ -130,17 +182,521 @@ fn test_eval_opts_builder() { ..Default::default() }]; - let dataset: Box> = cases.into(); + let dataset: Box> = cases.into(); let task: Box> = Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))); let scorers: Vec>> = vec![Box::new(ExactMatch::new())]; let _opts = EvalOpts::builder() .name("test-eval".to_string()) - .dataset(dataset) - .task(task) - .scorers(scorers) + .source(dataset) + .definition(EvalDefinition::new(task, scorers).unwrap()) .max_concurrency(10) .quiet(true) .build(); } + +#[tokio::test] +async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { + let (client, server) = recording_client().await; + let evaluator = Evaluator::new(client); + let scorers: Vec>> = vec![ + Box::new(ExactMatch::new()), + Box::new(SyncFnScorer::new("broken", |_| { + anyhow::bail!("scoring failed") + })), + ]; + let result = evaluator + .run( + EvalOpts::builder() + .name("failed-scorer".into()) + .source( + vec![Case { + input: "a".to_owned(), + expected: Some("A".to_owned()), + id: Some("source-case".into()), + ..Default::default() + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), + scorers, + ) + .unwrap(), + ) + .quiet(true) + .build(), + ) + .await + .unwrap(); + assert_eq!(result.successful_cases, 0); + assert_eq!(result.failed_cases, 1); + assert_eq!(result.score_stats[0].name, "exact_match"); + assert_eq!(result.score_stats[0].mean, 1.0); + let case = &result.results[0]; + assert_eq!(case.case_id.as_deref(), Some("source-case")); + assert!(matches!(&case.errors[0], CaseError::Scorer {name,..} if name == "broken")); + let rows: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/logs3") + .flat_map(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + assert!(rows + .iter() + .any(|row| row["error"][0]["stage"] == "scorer" && row["error"][0]["name"] == "broken")); + assert!(rows + .iter() + .any(|row| row["span_id"].as_str() == case.span_id.as_deref() + && row["metrics"]["end"].is_number())); +} + +#[tokio::test] +async fn source_failure_drains_active_cases_preserves_provenance_and_finishes_recording() { + use async_trait::async_trait; + use std::{ + cell::Cell, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + }; + use tokio::sync::{Notify, Semaphore}; + #[derive(serde::Serialize)] + struct Input(usize); + #[derive(serde::Serialize)] + struct Output(usize); + #[derive(Debug, thiserror::Error)] + #[error("source read failed")] + struct ReadFailure; + struct Source { + reads: Cell, + started: Arc, + failed: Arc, + } + #[async_trait] + impl CaseSource for Source { + async fn next(&mut self) -> anyhow::Result>> { + let index = self.reads.get(); + self.reads.set(index + 1); + if index == 1 { + self.started.notified().await; + } + if index == 2 { + self.failed.notify_one(); + return Err(ReadFailure.into()); + } + Ok(Some(Case { + input: Input(index), + expected: None, + metadata: None, + tags: None, + id: Some(format!("case-{index}")), + xact_id: Some(format!("row-version-{index}")), + created: None, + })) + } + fn provenance(&self) -> Option { + Some(DatasetVersionRef { + dataset_id: "source-dataset".into(), + version: Some(format!("version-{}", self.reads.get())), + }) + } + } + let (client, server) = recording_client().await; + let started = Arc::new(Notify::new()); + let failed = Arc::new(Notify::new()); + let second_done = Arc::new(Notify::new()); + let gates = Arc::new([Semaphore::new(0), Semaphore::new(0)]); + let calls = Arc::new(AtomicUsize::new(0)); + let callback = { + let started = started.clone(); + let second_done = second_done.clone(); + let gates = gates.clone(); + let calls = calls.clone(); + task(move |input: &Input, _hooks| { + let started = started.clone(); + let second_done = second_done.clone(); + let gates = gates.clone(); + let calls = calls.clone(); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + if input.0 == 0 { + started.notify_one(); + } + gates[input.0].acquire().await.unwrap().forget(); + if input.0 == 1 { + second_done.notify_one(); + } + Ok(Output(input.0 * 2)) + }) + }) + }; + let opts = EvalOpts::builder() + .name("source-failure".into()) + .source(Box::new(Source { + reads: Cell::new(0), + started, + failed: failed.clone(), + }) as Box>) + .definition(EvalDefinition::new(Box::new(callback), vec![]).unwrap()) + .max_concurrency(3) + .quiet(true) + .build(); + let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + tokio::time::timeout(std::time::Duration::from_secs(3), failed.notified()) + .await + .expect("pending source reads must allow active task progress"); + gates[1].add_permits(1); + second_done.notified().await; + assert!( + !run.is_finished(), + "source failure must still wait for the other scheduled case" + ); + gates[0].add_permits(1); + let error = match tokio::time::timeout(std::time::Duration::from_secs(3), run) + .await + .unwrap() + .unwrap() + { + Err(error) => error, + Ok(_) => panic!("source failure must not become a successful summary"), + }; + assert_eq!(calls.load(Ordering::SeqCst), 2); + let [EvalRunFailure::Source { case_index, source }] = error.failures.as_slice() else { + panic!("expected a typed source failure") + }; + assert_eq!(*case_index, 2); + assert!(source.downcast_ref::().is_some()); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 2); + assert_eq!(partial.successful_cases, 2); + assert_eq!(partial.failed_cases, 0); + let dataset = partial.dataset.unwrap(); + assert_eq!(dataset.dataset_id, "source-dataset"); + assert_eq!(dataset.version.as_deref(), Some("version-0")); + assert_eq!( + partial + .results + .iter() + .map(|case| case.input.0) + .collect::>(), + [0, 1] + ); + assert_eq!(partial.results[1].case_id.as_deref(), Some("case-1")); + assert_eq!(partial.results[1].xact_id.as_deref(), Some("row-version-1")); + assert_eq!(partial.results[1].output.as_ref().unwrap().0, 2); + let rows: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/logs3") + .flat_map(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + let ended: std::collections::HashSet<_> = rows + .iter() + .filter(|row| row["metrics"]["end"].is_number()) + .map(|row| row["span_id"].as_str().unwrap()) + .collect(); + assert_eq!( + ended.len(), + 3, + "both cases and the run span must finish before returning partial results" + ); +} + +#[tokio::test] +async fn failure_before_the_first_case_returns_an_empty_partial_run_and_finalizes_root() { + let (client, server) = recording_client().await; + let items: Vec>> = + vec![Err(anyhow::anyhow!("source unavailable"))]; + let opts = EvalOpts::builder() + .name("empty-failure".into()) + .source( + Box::new(IterCaseSource::new(items.into_iter())) as Box> + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &String| -> anyhow::Result { + panic!("no case was read") + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let error = Evaluator::new(client).run(opts).await.unwrap_err(); + assert!( + matches!(&error.failures[0],EvalRunFailure::Source {case_index:0,source} if source.to_string()=="source unavailable") + ); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 0); + assert_eq!(partial.successful_cases, 0); + assert_eq!(partial.failed_cases, 0); + assert!(partial.results.is_empty()); + assert!(partial.dataset.is_none()); + let ended = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .any(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .iter() + .any(|row| row["metrics"]["end"].is_number()) + }); + assert!( + ended, + "the run span must finish even when source failure precedes all cases" + ); +} + +#[tokio::test] +async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + use tokio::sync::{Notify, Semaphore}; + use wiremock::{ + matchers::{method, path}, + Mock, ResponseTemplate, + }; + let (client, server) = recording_client().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(1) + .mount(&server) + .await; + let calls = Arc::new(AtomicUsize::new(0)); + let entered = Arc::new(Notify::new()); + let gate = Arc::new(Semaphore::new(0)); + let callback = { + let calls = calls.clone(); + let entered = entered.clone(); + let gate = gate.clone(); + task(move |input: &usize, _hooks| { + let calls = calls.clone(); + let entered = entered.clone(); + let gate = gate.clone(); + Box::pin(async move { + if calls.fetch_add(1, Ordering::SeqCst) == 1 { + entered.notify_one(); + } + gate.acquire().await.unwrap().forget(); + Ok(*input) + }) + }) + }; + let opts = EvalOpts::builder() + .name("record-failure".into()) + .source( + (0..3) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>>() + .into(), + ) + .definition(EvalDefinition::new(Box::new(callback), vec![]).unwrap()) + .max_concurrency(2) + .quiet(true) + .build(); + let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + tokio::time::timeout(std::time::Duration::from_secs(3), entered.notified()) + .await + .unwrap(); + gate.add_permits(2); + let error = tokio::time::timeout(std::time::Duration::from_secs(3), run) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "the third case must remain unread after recording failure" + ); + assert!(error.failures.iter().any(|failure| matches!( + failure, + EvalRunFailure::Recording { + case_index: Some(0), + .. + } + ))); + assert!(error.failures.iter().any(|failure| matches!( + failure, + EvalRunFailure::Recording { + case_index: Some(1), + .. + } + ))); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 2); + assert_eq!(partial.failed_cases, 0); + assert_eq!(partial.successful_cases, 2); + assert_eq!(partial.results.len(), 2); + assert_eq!( + partial + .results + .iter() + .map(|result| result.output) + .collect::>(), + [Some(0), Some(1)] + ); + assert!(partial + .results + .iter() + .all(|result| result.errors.is_empty() && result.span_id.is_some())); +} + +#[tokio::test] +async fn reusable_definition_keeps_recorded_runs_and_sources_independent() { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + let (client, _server) = recording_client().await; + let evaluator = Evaluator::new(client); + let calls = Arc::new(AtomicUsize::new(0)); + let definition = EvalDefinition::new( + Box::new(task({ + let calls = calls.clone(); + move |input: &usize, hooks| { + calls.fetch_add(1, Ordering::SeqCst); + hooks.add_tag(format!("input-{input}")); + Box::pin(async move { Ok(*input * 2) }) + } + })), + vec![Box::new(ExactMatch::new())], + ) + .unwrap(); + let options = |name: &str, input, expected| { + EvalOpts::builder() + .name(name.into()) + .source( + vec![Case { + input, + expected: Some(expected), + ..Default::default() + }] + .into(), + ) + .definition(definition.clone()) + .quiet(true) + .build() + }; + let (first, second) = tokio::join!( + evaluator.run(options("first", 2, 4)), + evaluator.run(options("second", 3, 100)), + ); + let first = first.unwrap(); + let second = second.unwrap(); + assert_eq!(first.name, "first"); + assert_eq!(second.name, "second"); + assert_eq!(first.results[0].output, Some(4)); + assert_eq!(second.results[0].output, Some(6)); + assert_eq!(first.results[0].scores[0].score, 1.0); + assert_eq!(second.results[0].scores[0].score, 0.0); + assert_eq!(first.results[0].tags, Some(vec!["input-2".into()])); + assert_eq!(second.results[0].tags, Some(vec!["input-3".into()])); + assert_ne!(first.results[0].span_id, second.results[0].span_id); + let later = evaluator.run(options("later", 5, 10)).await.unwrap(); + assert_eq!(later.total_cases, 1); + assert_eq!(later.results[0].output, Some(10)); + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn serialization_failures_keep_outputs_and_finalize_recording() { + #[derive(Debug)] + struct Value(usize); + impl serde::Serialize for Value { + fn serialize(&self, _: S) -> std::result::Result { + Err(serde::ser::Error::custom( + "application serialization failed", + )) + } + } + let (client, server) = recording_client().await; + let options = EvalOpts::builder() + .name("serialization-failure".into()) + .source( + vec![Case { + input: Value(7), + expected: Some(Value(14)), + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &Value| Ok(Value(input.0 * 2)))), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let error = Evaluator::new(client).run(options).await.unwrap_err(); + let fields = error + .failures + .iter() + .map(|failure| match failure { + EvalRunFailure::Recording { + case_index: Some(0), + source, + } => source.to_string(), + other => panic!("unexpected failure: {other}"), + }) + .collect::>(); + assert_eq!(fields.len(), 3); + for (field, message) in ["input", "expected", "output"].into_iter().zip(fields) { + assert!(message.contains(field)); + assert!(message.contains("application serialization failed")); + } + let summary = error.partial.unwrap(); + assert_eq!(summary.successful_cases, 1); + assert_eq!(summary.failed_cases, 0); + assert_eq!(summary.results[0].output.as_ref().unwrap().0, 14); + assert!(summary.results[0].span_id.is_some()); + let requests = server.received_requests().await.unwrap(); + let ended: std::collections::HashSet<_> = requests + .iter() + .filter(|request| request.url.path() == "/logs3") + .flat_map(|request| { + serde_json::from_slice::(&request.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .filter(|row| row["metrics"]["end"].is_number()) + .map(|row| row["id"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(ended.len(), 2, "both case and run recording must finalize"); +}