From dc011de708c210d655e16577ba4abb75672daef9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 00:27:34 +0800 Subject: [PATCH 1/7] Separate evaluation case computation and preserve structured failures --- src/eval/case_runner.rs | 258 ++++++++++++++++++++++++++++++++++++++++ src/eval/error.rs | 39 ++++++ src/eval/mod.rs | 4 + src/eval/runner.rs | 145 +++++++--------------- src/eval/scorer.rs | 4 +- src/eval/task.rs | 4 +- src/eval/types.rs | 63 ++++++---- tests/eval_test.rs | 112 +++++++++++++++-- 8 files changed, 495 insertions(+), 134 deletions(-) create mode 100644 src/eval/case_runner.rs create mode 100644 src/eval/error.rs diff --git a/src/eval/case_runner.rs b/src/eval/case_runner.rs new file mode 100644 index 0000000..ec39207 --- /dev/null +++ b/src/eval/case_runner.rs @@ -0,0 +1,258 @@ +//! One case's computation, independent of scheduling, recording and reporting. +use super::{Case, CaseError, EvalResult, Scorer, ScorerArgs, Task, TaskHooks}; +use crate::SpanHandle; +use futures::{stream::FuturesOrdered, StreamExt}; + +pub(crate) struct CaseRunner; +impl CaseRunner { + pub async fn run( + case: Case, + task: &dyn Task, + scorers: &[Box>], + span: Option<&SpanHandle>, + ) -> EvalResult + where + I: Send + Sync, + O: Send + Sync, + { + let hooks = TaskHooks::for_case( + span, + case.expected.as_ref(), + case.metadata.clone(), + case.tags.clone(), + ); + let mut errors = Vec::new(); + let output = match task.run(&case.input, &hooks).await { + Ok(output) => Some(output), + Err(source) => { + errors.push(CaseError::Task { source }); + None + } + }; + let mut scores = Vec::new(); + if let Some(output) = &output { + // Drive scorers concurrently while keeping definition order in the + // outcome, regardless of their relative completion times. + let mut pending = FuturesOrdered::new(); + for scorer in scorers { + let args = ScorerArgs { + input: &case.input, + output, + expected: case.expected.as_ref(), + metadata: case.metadata.as_ref(), + }; + pending.push_back(async move { + scorer + .score(args) + .await + .map_err(|source| CaseError::Scorer { + name: scorer.name().to_owned(), + 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); + 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, Scores, SimpleFnTask}; + 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, &LocalTask, &[], 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() + }, + &FailTask, + &scorers, + 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() + }, + &SimpleFnTask::new(|input: &usize| Ok(*input)), + &scorers, + 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/error.rs b/src/eval/error.rs new file mode 100644 index 0000000..163fe51 --- /dev/null +++ b/src/eval/error.rs @@ -0,0 +1,39 @@ +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() + } +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index f8d2d85..2798723 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -48,12 +48,16 @@ //! ``` mod bt_runner; +mod case_runner; mod dataset; +mod error; mod runner; mod scorer; mod task; mod types; +pub use error::CaseError; + // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; pub use dataset::{Dataset, IterDataset, StreamDataset, VecDataset}; diff --git a/src/eval/runner.rs b/src/eval/runner.rs index cecb373..9ea4967 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -1,6 +1,6 @@ use futures::stream::{FuturesUnordered, StreamExt}; use serde::{de::DeserializeOwned, Serialize}; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value}; use std::marker::PhantomData; use std::sync::Arc; use tokio::sync::Semaphore; @@ -8,10 +8,11 @@ use tokio::sync::Semaphore; use crate::{BraintrustClient, Result, SpanHandle}; use super::{ + case_runner::CaseRunner, dataset::Dataset, scorer::Scorer, task::Task, - types::{Case, EvalResult, EvalSummary, TaskHooks}, + types::{Case, EvalResult, EvalSummary}, }; /// Options for running an evaluation @@ -167,7 +168,7 @@ where // Compute statistics let score_stats = EvalSummary::::compute_stats(&results); - let successful = results.iter().filter(|r| r.error.is_none()).count(); + let successful = results.iter().filter(|r| r.errors.is_empty()).count(); let failed = results.len() - successful; if !opts.quiet { @@ -232,117 +233,59 @@ where 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); - } + let initial_write = case_span.log(log_builder.build().unwrap()); + let mut result = CaseRunner::run(case, task, scorers, Some(&case_span)).await; + result.span_id = case_span.export()?.span_id; - // 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); + // Recording adapts the completed outcome; computation never waits on + // HTTP or emits console messages. Observe every accepted write before + // reporting recording success, including writes that failed earlier. + let mut final_log = crate::SpanLog::builder().scores( + result + .scores + .iter() + .map(|score| (score.name.clone(), score.score)) + .collect(), + ); + if let Some(output) = &result.output { + if let Ok(output) = serde_json::to_value(output) { + final_log = final_log.output(output); + } } - - // 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(metadata) = &result.metadata { + final_log = final_log.metadata(metadata.clone()); } - if let Some(tags) = &final_tags { + if let Some(tags) = &result.tags { final_log = final_log.tags(tags.clone()); } - case_span.log(final_log.build().unwrap()); - - case_span.end().await?; + if !result.errors.is_empty() { + final_log = final_log.error( + serde_json::to_value(&result.errors) + .map_err(|error| crate::BraintrustError::Background(error.to_string()))?, + ); + } + let final_write = case_span.log(final_log.build().unwrap()); + let end = case_span.end(); + let initial_result = initial_write.await; + let final_result = final_write.await; + let end_result = end.await; + initial_result?; + final_result?; + end_result?; if !quiet { - let score_str = all_scores + let score_str = result + .scores .iter() .map(|s| format!("{}={:.3}", s.name, s.score)) .collect::>() .join(", "); println!("Case {}: {}", case_num, score_str); + for error in &result.errors { + eprintln!("Case {}: {}", case_num, error); + } } - - 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, - }) + Ok(result) } } diff --git a/src/eval/scorer.rs b/src/eval/scorer.rs index 633a917..508e648 100644 --- a/src/eval/scorer.rs +++ b/src/eval/scorer.rs @@ -1,14 +1,14 @@ use async_trait::async_trait; 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) diff --git a/src/eval/task.rs b/src/eval/task.rs index d529068..540e7f3 100644 --- a/src/eval/task.rs +++ b/src/eval/task.rs @@ -2,14 +2,14 @@ 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 { /// Execute the task on the given input diff --git a/src/eval/types.rs b/src/eval/types.rs index acbaf0e..c7356ce 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; 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 error: Option, + pub case_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub xact_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub created: Option, /// Span ID for this evaluation #[serde(skip_serializing_if = "Option::is_none")] pub span_id: Option, @@ -162,20 +165,19 @@ pub struct EvalSummary { } 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,25 @@ 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, } } @@ -235,7 +251,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/tests/eval_test.rs b/tests/eval_test.rs index b07eb0d..26f22fd 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -2,19 +2,48 @@ use braintrust_sdk_rust::eval::Dataset; 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![ @@ -48,6 +77,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] @@ -144,3 +181,64 @@ fn test_eval_opts_builder() { .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()) + .dataset( + vec![Case { + input: "a".to_owned(), + expected: Some("A".to_owned()), + id: Some("source-case".into()), + ..Default::default() + }] + .into(), + ) + .task(Box::new(SimpleFnTask::new(|input: &String| { + Ok(input.to_uppercase()) + }))) + .scorers(scorers) + .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())); +} From 6247c4983eb0ee4e1e0392bb99f43c6f23b861e0 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 00:39:11 +0800 Subject: [PATCH 2/7] Accept reusable evaluation closures and validate run options --- examples/eval_async.rs | 15 ++- src/eval/error.rs | 11 ++ src/eval/mod.rs | 10 +- src/eval/runner.rs | 25 +++++ src/eval/scorer.rs | 37 ++++++- src/eval/task.rs | 64 ++--------- tests/eval_callables.rs | 231 ++++++++++++++++++++++++++++++++++++++++ tests/eval_test.rs | 26 ++++- 8 files changed, 351 insertions(+), 68 deletions(-) create mode 100644 tests/eval_callables.rs diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 3d07e85..130f100 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -45,12 +45,23 @@ async fn main() -> Result<()> { EvalOpts::builder() .name("uppercase-eval".to_string()) .dataset(cases.into()) - .task(Box::new(SimpleFnTask::new(|input: &String| { - Ok(input.to_uppercase()) + .task(Box::new(task(|input: &String, _hooks| { + Box::pin(async move { Ok(input.to_uppercase()) }) }))) .scorers(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 }, + )]) + }) + }, + )), ]) .max_concurrency(2) .build(), diff --git a/src/eval/error.rs b/src/eval/error.rs index 163fe51..fa98412 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -37,3 +37,14 @@ impl Serialize for CaseError { 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("scorer at index {index} must have a nonblank name")] + EmptyScorerName { index: usize }, + #[error("duplicate scorer name: {name}")] + DuplicateScorerName { name: String }, +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 2798723..1754c27 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -34,9 +34,9 @@ //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) //! .dataset(cases.into()) -//! .task(Box::new(SimpleFnTask::new(|input: &String| { +//! .task(Box::new(task(|input: &String, _hooks| Box::pin(async move { //! Ok(input.to_uppercase()) -//! }))) +//! })))) //! .scorers(vec![Box::new(ExactMatch::new())]) //! .build(), //! ) @@ -56,12 +56,12 @@ mod scorer; mod task; mod types; -pub use error::CaseError; +pub use error::{CaseError, EvalConfigError}; // 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 scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; +pub use task::{task, SimpleFnTask, SyncFnTask, Task}; pub use types::{Case, EvalResult, EvalSummary, Score, ScoreStats, ScorerArgs, Scores, TaskHooks}; diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 9ea4967..38ea6e5 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -13,6 +13,7 @@ use super::{ scorer::Scorer, task::Task, types::{Case, EvalResult, EvalSummary}, + EvalConfigError, }; /// Options for running an evaluation @@ -42,6 +43,28 @@ pub struct EvalOpts { pub quiet: bool, } +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); + } + let mut names = std::collections::HashSet::new(); + for (index, scorer) in self.scorers.iter().enumerate() { + let name = scorer.name(); + if name.trim().is_empty() { + return Err(EvalConfigError::EmptyScorerName { index }); + } + if !names.insert(name) { + return Err(EvalConfigError::DuplicateScorerName { + name: name.to_owned(), + }); + } + } + Ok(()) + } +} + /// Async evaluator for running evaluations pub struct Evaluator { client: Arc, @@ -71,6 +94,8 @@ where /// Run the evaluation pub async fn run(&self, mut opts: EvalOpts) -> Result> { + opts.validate() + .map_err(|error| crate::BraintrustError::InvalidConfig(error.to_string()))?; if !opts.quiet { println!("Running evaluation: {}", opts.name); } diff --git a/src/eval/scorer.rs b/src/eval/scorer.rs index 508e648..f64b3e0 100644 --- a/src/eval/scorer.rs +++ b/src/eval/scorer.rs @@ -1,4 +1,5 @@ use async_trait::async_trait; +use futures::future::BoxFuture; use std::marker::PhantomData; use anyhow::Result; @@ -18,6 +19,33 @@ pub trait Scorer: Send + Sync { 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. @@ -276,10 +304,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 +316,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 +332,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/task.rs b/src/eval/task.rs index 540e7f3..18ebb90 100644 --- a/src/eval/task.rs +++ b/src/eval/task.rs @@ -52,36 +52,25 @@ where } } -/// Wrapper for async task functions that return BoxFuture -pub struct 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, { - func: F, - _phantom: PhantomData<(I, O)>, -} - -impl AsyncFnTask -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 } } @@ -118,39 +107,6 @@ where } } -/// 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_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 - } -} - #[cfg(test)] mod tests { use super::*; @@ -159,7 +115,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/tests/eval_callables.rs b/tests/eval_callables.rs new file mode 100644 index 0000000..726b5d3 --- /dev/null +++ b/tests/eval_callables.rs @@ -0,0 +1,231 @@ +use async_trait::async_trait; +use braintrust_sdk_rust::{ + eval::{ + scorer, task, Case, Dataset, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, + Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, + }, + BraintrustClient, BraintrustError, +}; +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 Dataset for UntouchedSource { + async fn next(&mut self) -> braintrust_sdk_rust::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 opts = EvalOpts::builder() + .name("invalid".into()) + .dataset(Box::new(UntouchedSource(accessed.clone())) as Box>) + .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) + .scorers(scorers) + .max_concurrency(concurrency) + .quiet(true) + .build(); + assert_eq!(opts.validate(), Err(expected.clone())); + assert!( + matches!(evaluator.run(opts).await, Err(BraintrustError::InvalidConfig(message)) if message == expected.to_string()) + ); + } + 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 opts = EvalOpts::builder() + .name("duplicate".into()) + .dataset(Vec::>::new().into()) + .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) + .scorers(vec![Box::new(ExactMatch::new()), callback]) + .build(); + assert_eq!( + opts.validate(), + Err(EvalConfigError::DuplicateScorerName { + name: "exact_match".into() + }) + ); + } +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 26f22fd..36c7148 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -59,15 +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)) + .task(Box::new(callback)) .scorers(scorers) .quiet(true) .build(), @@ -126,10 +140,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(); @@ -137,10 +152,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(); From f732ebb54eef0acdad6e497ed66204d14e410d4e Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:05:57 +0800 Subject: [PATCH 3/7] Own evaluation case sources and preserve partial run failures --- examples/eval_async.rs | 2 +- examples/eval_sync.rs | 2 +- src/eval/bt_runner.rs | 10 +- src/eval/dataset.rs | 196 ------------------------- src/eval/error.rs | 91 ++++++++++++ src/eval/mod.rs | 12 +- src/eval/runner.rs | 160 ++++++++++++--------- src/eval/source.rs | 105 ++++++++++++++ src/eval/types.rs | 6 +- tests/eval_callables.rs | 30 ++-- tests/eval_sources.rs | 113 +++++++++++++++ tests/eval_test.rs | 307 +++++++++++++++++++++++++++++++++++++++- 12 files changed, 740 insertions(+), 294 deletions(-) delete mode 100644 src/eval/dataset.rs create mode 100644 src/eval/source.rs create mode 100644 tests/eval_sources.rs diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 130f100..f17894d 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -44,7 +44,7 @@ async fn main() -> Result<()> { "my-project", EvalOpts::builder() .name("uppercase-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(task(|input: &String, _hooks| { Box::pin(async move { Ok(input.to_uppercase()) }) }))) diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index ae065a5..f6fec87 100644 --- a/examples/eval_sync.rs +++ b/examples/eval_sync.rs @@ -42,7 +42,7 @@ fn main() -> anyhow::Result<()> { let result = evaluator.run( EvalOpts::builder() .name("sync-uppercase-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(task)) .scorers(scorers) .build(), diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index 1daae18..fcb2eed 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}; @@ -193,7 +193,7 @@ fn create_sse_writer() -> Option> { /// "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() @@ -288,8 +288,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 +349,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/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/error.rs b/src/eval/error.rs index fa98412..9798a2f 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -48,3 +48,94 @@ pub enum EvalConfigError { #[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(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. +pub struct EvalRunError { + pub failures: Vec, + pub partial: Option>>, +} +impl EvalRunError { + 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/mod.rs b/src/eval/mod.rs index 1754c27..c785cd6 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -13,7 +13,7 @@ //! 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() @@ -33,7 +33,7 @@ //! .run( //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) -//! .dataset(cases.into()) +//! .source(cases.into()) //! .task(Box::new(task(|input: &String, _hooks| Box::pin(async move { //! Ok(input.to_uppercase()) //! })))) @@ -49,19 +49,21 @@ mod bt_runner; mod case_runner; -mod dataset; mod error; mod runner; mod scorer; +mod source; mod task; mod types; -pub use error::{CaseError, EvalConfigError}; +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::{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/runner.rs b/src/eval/runner.rs index 38ea6e5..19453c4 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -1,19 +1,18 @@ use futures::stream::{FuturesUnordered, StreamExt}; -use serde::{de::DeserializeOwned, Serialize}; +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 super::{ case_runner::CaseRunner, - dataset::Dataset, scorer::Scorer, + source::CaseSource, task::Task, types::{Case, EvalResult, EvalSummary}, - EvalConfigError, + EvalConfigError, EvalRunError, EvalRunFailure, }; /// Options for running an evaluation @@ -21,8 +20,8 @@ use super::{ pub struct EvalOpts { /// Name of the evaluation pub name: String, - /// Dataset providing test cases - pub dataset: Box>, + /// Owned source providing test cases + pub source: Box>, /// Task to execute on each input pub task: Box>, /// Scorers to evaluate outputs @@ -73,8 +72,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 { @@ -93,9 +92,11 @@ where } /// Run the evaluation - pub async fn run(&self, mut opts: EvalOpts) -> Result> { - opts.validate() - .map_err(|error| crate::BraintrustError::InvalidConfig(error.to_string()))?; + pub async fn run( + &self, + mut opts: EvalOpts, + ) -> std::result::Result, EvalRunError> { + opts.validate()?; if !opts.quiet { println!("Running evaluation: {}", opts.name); } @@ -118,83 +119,90 @@ where log_builder = log_builder.tags(tags.clone()); } - root_span.log(log_builder.build().unwrap()); + let root_write = root_span.log(log_builder.build().unwrap()); - // Get dataset ID if available - let dataset_id = opts.dataset.id(); - - // Collect all results + let dataset = opts.source.provenance(); let mut results = Vec::new(); + let mut failures = 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 { + // Retain the next-read future while polling active cases. A source + // may wait for prior computation, and need not be cancellation-safe. + let mut next = Box::pin(opts.source.next()); + let next_case = loop { + if futures.is_empty() { + break next.await; + } + tokio::select! { + next_case = next.as_mut() => break next_case, + Some((index, result)) = futures.next() => { + Self::collect_case_result(index, result, &mut results, &mut failures); + if !failures.is_empty() { break Ok(None); } + } + } + }; + let case = match next_case { Ok(Some(case)) => case, Ok(None) => break, - Err(e) => { - eprintln!("Error fetching case: {}", e); + Err(source) => { + failures.push(EvalRunFailure::Source { + case_index: case_count, + source, + }); break; } }; - case_count += 1; - - // Clone what we need for the async task + let index = case_count - 1; 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 { + futures.push(async move { let result = Self::process_case( case, task.as_ref(), scorers, &root_span, &client, - case_count, + index + 1, 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?); + (index, result) + }); + if futures.len() >= opts.max_concurrency { + if let Some((index, result)) = futures.next().await { + Self::collect_case_result(index, result, &mut results, &mut failures); + } + if !failures.is_empty() { + break; } } } - - // Collect remaining futures - while let Some(result) = futures.next().await { - results.push(result?); + // Source/recording failure does not abandon already scheduled cases. + while let Some((index, result)) = futures.next().await { + Self::collect_case_result(index, result, &mut results, &mut failures); } - - // End root span - root_span.end().await?; + let end = root_span.end(); + for result in [root_write.await, end.await] { + if let Err(error) = result { + failures.push(EvalRunFailure::Recording { + case_index: None, + source: error.into(), + }); + } + } + results.sort_unstable_by_key(|(index, _)| *index); + let results: Vec<_> = results.into_iter().map(|(_, result)| result).collect(); // Compute statistics let score_stats = EvalSummary::::compute_stats(&results); let successful = results.iter().filter(|r| r.errors.is_empty()).count(); - let failed = results.len() - successful; + let failed = case_count - successful; if !opts.quiet { println!("\nEvaluation complete!"); @@ -210,15 +218,38 @@ where } } - Ok(EvalSummary { + let summary = EvalSummary { name: opts.name, results, score_stats, total_cases: case_count, successful_cases: successful, failed_cases: failed, - dataset_id, - }) + dataset, + }; + if failures.is_empty() { + Ok(summary) + } else { + Err(EvalRunError { + failures, + partial: Some(Box::new(summary)), + }) + } + } + + fn collect_case_result( + index: usize, + result: Result>, + results: &mut Vec<(usize, EvalResult)>, + failures: &mut Vec, + ) { + match result { + Ok(result) => results.push((index, result)), + Err(source) => failures.push(EvalRunFailure::Recording { + case_index: Some(index), + source, + }), + } } /// Process a single case @@ -322,8 +353,8 @@ pub struct SyncEvaluator { impl SyncEvaluator 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 sync evaluator pub fn new(client: BraintrustClient) -> Result { @@ -336,7 +367,10 @@ where } /// Run the evaluation synchronously - pub fn run(&self, opts: EvalOpts) -> Result> { + pub fn run( + &self, + opts: EvalOpts, + ) -> std::result::Result, EvalRunError> { self.runtime.block_on(self.evaluator.run(opts)) } } @@ -344,7 +378,7 @@ where #[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}; // Note: Full integration tests require a BraintrustClient setup // These tests verify type correctness and compilation @@ -357,14 +391,14 @@ 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) + .source(dataset) .task(task) .scorers(scorers) .build(); 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/types.rs b/src/eval/types.rs index c7356ce..9078148 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::sync::{Arc, Mutex}; -use super::CaseError; +use super::{CaseError, DatasetVersionRef}; use crate::SpanHandle; /// A single test case for evaluation @@ -159,9 +159,9 @@ pub struct EvalSummary { pub successful_cases: usize, /// Number of failed evaluations 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 { diff --git a/tests/eval_callables.rs b/tests/eval_callables.rs index 726b5d3..9b6f47a 100644 --- a/tests/eval_callables.rs +++ b/tests/eval_callables.rs @@ -1,10 +1,10 @@ use async_trait::async_trait; use braintrust_sdk_rust::{ eval::{ - scorer, task, Case, Dataset, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, + scorer, task, Case, CaseSource, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, }, - BraintrustClient, BraintrustError, + BraintrustClient, }; use futures::future::BoxFuture; use std::sync::{ @@ -139,8 +139,8 @@ async fn scorer_helpers_accept_borrowed_and_owned_names_and_captured_state() { struct UntouchedSource(Arc); #[async_trait] -impl Dataset for UntouchedSource { - async fn next(&mut self) -> braintrust_sdk_rust::Result>> { +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"); } @@ -180,17 +180,19 @@ async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_acc as Box> }) .collect(); - let opts = EvalOpts::builder() - .name("invalid".into()) - .dataset(Box::new(UntouchedSource(accessed.clone())) as Box>) - .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) - .scorers(scorers) - .max_concurrency(concurrency) - .quiet(true) - .build(); + let opts = + EvalOpts::builder() + .name("invalid".into()) + .source(Box::new(UntouchedSource(accessed.clone())) + as Box>) + .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) + .scorers(scorers) + .max_concurrency(concurrency) + .quiet(true) + .build(); assert_eq!(opts.validate(), Err(expected.clone())); assert!( - matches!(evaluator.run(opts).await, Err(BraintrustError::InvalidConfig(message)) if message == expected.to_string()) + 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)); @@ -217,7 +219,7 @@ fn duplicate_names_are_rejected_across_custom_tuple_and_builtin_scorers() { ] { let opts = EvalOpts::builder() .name("duplicate".into()) - .dataset(Vec::>::new().into()) + .source(Vec::>::new().into()) .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) .scorers(vec![Box::new(ExactMatch::new()), callback]) .build(); 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 36c7148..e139f1b 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -1,4 +1,4 @@ -use braintrust_sdk_rust::eval::Dataset; +use braintrust_sdk_rust::eval::CaseSource; use braintrust_sdk_rust::eval::*; use braintrust_sdk_rust::*; @@ -80,7 +80,7 @@ async fn test_basic_eval_flow() { .run( EvalOpts::builder() .name("test-basic-eval".to_string()) - .dataset(cases.into()) + .source(cases.into()) .task(Box::new(callback)) .scorers(scorers) .quiet(true) @@ -110,7 +110,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())]; @@ -183,14 +183,14 @@ 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) + .source(dataset) .task(task) .scorers(scorers) .max_concurrency(10) @@ -212,7 +212,7 @@ async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { .run( EvalOpts::builder() .name("failed-scorer".into()) - .dataset( + .source( vec![Case { input: "a".to_owned(), expected: Some("A".to_owned()), @@ -258,3 +258,298 @@ async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { .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>) + .task(Box::new(callback)) + .scorers(vec![]) + .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> + ) + .task(Box::new(SimpleFnTask::new( + |_: &String| -> anyhow::Result { panic!("no case was read") }, + ))) + .scorers(vec![]) + .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(), + ) + .task(Box::new(callback)) + .scorers(vec![]) + .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, 2); + assert_eq!(partial.successful_cases, 0); +} From b3aeef68bb818c091d89badf8e8eafff5b5537e9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:19:05 +0800 Subject: [PATCH 4/7] Separate reusable evaluation definitions from run options --- examples/eval_async.rs | 44 +++++---- examples/eval_sync.rs | 4 +- src/eval/bt_runner.rs | 10 +- src/eval/definition.rs | 197 ++++++++++++++++++++++++++++++++++++++++ src/eval/mod.rs | 14 ++- src/eval/runner.rs | 56 ++++-------- tests/eval_callables.rs | 32 ++++--- tests/eval_test.rs | 93 +++++++++++++++---- 8 files changed, 353 insertions(+), 97 deletions(-) create mode 100644 src/eval/definition.rs diff --git a/examples/eval_async.rs b/examples/eval_async.rs index f17894d..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,30 +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()) .source(cases.into()) - .task(Box::new(task(|input: &String, _hooks| { - Box::pin(async move { Ok(input.to_uppercase()) }) - }))) - .scorers(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 }, - )]) - }) - }, - )), - ]) + .definition(definition.clone()) .max_concurrency(2) .build(), ) diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index f6fec87..ac3c245 100644 --- a/examples/eval_sync.rs +++ b/examples/eval_sync.rs @@ -39,12 +39,12 @@ fn main() -> anyhow::Result<()> { // Run evaluation synchronously println!("Starting synchronous evaluation...\n"); + let definition = EvalDefinition::new(Box::new(task), scorers)?; let result = evaluator.run( EvalOpts::builder() .name("sync-uppercase-eval".to_string()) .source(cases.into()) - .task(Box::new(task)) - .scorers(scorers) + .definition(definition.clone()) .build(), )?; diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index fcb2eed..8662a52 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -186,9 +186,14 @@ 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() @@ -198,8 +203,7 @@ fn create_sse_writer() -> Option> { /// 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?; /// diff --git a/src/eval/definition.rs b/src/eval/definition.rs new file mode 100644 index 0000000..11adddd --- /dev/null +++ b/src/eval/definition.rs @@ -0,0 +1,197 @@ +//! 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: Box>, + 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(Box::new(NamedScorer { name, scorer })); + } + Ok(Self { + inner: Arc::new(DefinitionInner { + 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) -> &dyn Task { + self.inner.task.as_ref() + } + + pub(crate) fn scorers(&self) -> &[Box>] { + &self.inner.scorers + } +} + +struct NamedScorer { + name: String, + scorer: Box>, +} + +#[async_trait] +impl Scorer for NamedScorer { + fn name(&self) -> &str { + &self.name + } + + 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, SimpleFnTask}; + 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.task(), first.scorers(), None), + CaseRunner::run(case(2, 20), other.task(), other.scorers(), 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.task(), first.scorers(), None), + CaseRunner::run(case(4, 40), other.task(), other.scorers(), 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(SimpleFnTask::new(|input: &usize| 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.task(), + definition.scorers(), + None, + )); + assert!( + matches!(&result.errors[..], [CaseError::Scorer {name, source}] if name == "original" && source.to_string() == "scoring failed") + ); + } + } +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index c785cd6..22d02cc 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -29,15 +29,19 @@ //! }, //! ]; //! +//! // 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 result = evaluator //! .run( //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) //! .source(cases.into()) -//! .task(Box::new(task(|input: &String, _hooks| Box::pin(async move { -//! Ok(input.to_uppercase()) -//! })))) -//! .scorers(vec![Box::new(ExactMatch::new())]) +//! .definition(definition.clone()) //! .build(), //! ) //! .await?; @@ -49,6 +53,7 @@ mod bt_runner; mod case_runner; +mod definition; mod error; mod runner; mod scorer; @@ -56,6 +61,7 @@ mod source; mod task; mod types; +pub use definition::EvalDefinition; pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; // Re-export public API diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 19453c4..8263a28 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -8,11 +8,9 @@ use crate::{BraintrustClient, Result, SpanHandle}; use super::{ case_runner::CaseRunner, - scorer::Scorer, source::CaseSource, - task::Task, types::{Case, EvalResult, EvalSummary}, - EvalConfigError, EvalRunError, EvalRunFailure, + EvalConfigError, EvalDefinition, EvalRunError, EvalRunFailure, }; /// Options for running an evaluation @@ -22,10 +20,8 @@ pub struct EvalOpts { pub name: String, /// Owned source providing test cases pub source: Box>, - /// Task to execute on each input - pub task: Box>, - /// Scorers to evaluate outputs - pub scorers: Vec>>, + /// Reusable task and ordered scorers + pub definition: EvalDefinition, /// Braintrust project to log results to pub project_name: Option, /// Optional experiment name for Braintrust logging @@ -48,18 +44,6 @@ impl EvalOpts { if self.max_concurrency == 0 { return Err(EvalConfigError::ZeroConcurrency); } - let mut names = std::collections::HashSet::new(); - for (index, scorer) in self.scorers.iter().enumerate() { - let name = scorer.name(); - if name.trim().is_empty() { - return Err(EvalConfigError::EmptyScorerName { index }); - } - if !names.insert(name) { - return Err(EvalConfigError::DuplicateScorerName { - name: name.to_owned(), - }); - } - } Ok(()) } } @@ -158,20 +142,12 @@ where let index = case_count - 1; let client = Arc::clone(&self.client); let root_span = root_span.clone(); - let task = &opts.task; - let scorers = &opts.scorers; + let definition = &opts.definition; let quiet = opts.quiet; futures.push(async move { - let result = Self::process_case( - case, - task.as_ref(), - scorers, - &root_span, - &client, - index + 1, - quiet, - ) - .await; + let result = + Self::process_case(case, definition, &root_span, &client, index + 1, quiet) + .await; (index, result) }); if futures.len() >= opts.max_concurrency { @@ -255,8 +231,7 @@ where /// Process a single case async fn process_case( case: Case, - task: &dyn Task, - scorers: &[Box>], + definition: &EvalDefinition, _root_span: &SpanHandle, client: &BraintrustClient, case_num: usize, @@ -290,7 +265,13 @@ where } let initial_write = case_span.log(log_builder.build().unwrap()); - let mut result = CaseRunner::run(case, task, scorers, Some(&case_span)).await; + let mut result = CaseRunner::run( + case, + definition.task(), + definition.scorers(), + Some(&case_span), + ) + .await; result.span_id = case_span.export()?.span_id; // Recording adapts the completed outcome; computation never waits on @@ -378,7 +359,9 @@ where #[cfg(test)] mod tests { use super::*; - use crate::eval::{scorer::ExactMatch, source::VecCaseSource, 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 @@ -399,8 +382,7 @@ mod tests { let _opts = EvalOpts::builder() .name("test-eval".to_string()) .source(dataset) - .task(task) - .scorers(scorers) + .definition(crate::eval::EvalDefinition::new(task, scorers).unwrap()) .build(); } } diff --git a/tests/eval_callables.rs b/tests/eval_callables.rs index 9b6f47a..5ce314e 100644 --- a/tests/eval_callables.rs +++ b/tests/eval_callables.rs @@ -1,8 +1,8 @@ use async_trait::async_trait; use braintrust_sdk_rust::{ eval::{ - scorer, task, Case, CaseSource, EvalConfigError, EvalOpts, Evaluator, ExactMatch, Score, - Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, + scorer, task, Case, CaseSource, EvalConfigError, EvalDefinition, EvalOpts, Evaluator, + ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, }, BraintrustClient, }; @@ -180,13 +180,20 @@ async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_acc 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>) - .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) - .scorers(scorers) + .definition(definition.unwrap()) .max_concurrency(concurrency) .quiet(true) .build(); @@ -217,17 +224,12 @@ fn duplicate_names_are_rejected_across_custom_tuple_and_builtin_scorers() { Box::pin(async { Ok(vec![]) }) })), ] { - let opts = EvalOpts::builder() - .name("duplicate".into()) - .source(Vec::>::new().into()) - .task(Box::new(SimpleFnTask::new(|input: &usize| Ok(*input)))) - .scorers(vec![Box::new(ExactMatch::new()), callback]) - .build(); - assert_eq!( - opts.validate(), - Err(EvalConfigError::DuplicateScorerName { - name: "exact_match".into() - }) + 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_test.rs b/tests/eval_test.rs index e139f1b..0c6fdc5 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -81,8 +81,7 @@ async fn test_basic_eval_flow() { EvalOpts::builder() .name("test-basic-eval".to_string()) .source(cases.into()) - .task(Box::new(callback)) - .scorers(scorers) + .definition(EvalDefinition::new(Box::new(callback), scorers).unwrap()) .quiet(true) .build(), ) @@ -191,8 +190,7 @@ fn test_eval_opts_builder() { let _opts = EvalOpts::builder() .name("test-eval".to_string()) .source(dataset) - .task(task) - .scorers(scorers) + .definition(EvalDefinition::new(task, scorers).unwrap()) .max_concurrency(10) .quiet(true) .build(); @@ -221,10 +219,13 @@ async fn recorded_scorer_failure_keeps_valid_scores_and_finishes_the_case() { }] .into(), ) - .task(Box::new(SimpleFnTask::new(|input: &String| { - Ok(input.to_uppercase()) - }))) - .scorers(scorers) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), + scorers, + ) + .unwrap(), + ) .quiet(true) .build(), ) @@ -347,8 +348,7 @@ async fn source_failure_drains_active_cases_preserves_provenance_and_finishes_re started, failed: failed.clone(), }) as Box>) - .task(Box::new(callback)) - .scorers(vec![]) + .definition(EvalDefinition::new(Box::new(callback), vec![]).unwrap()) .max_concurrency(3) .quiet(true) .build(); @@ -430,10 +430,15 @@ async fn failure_before_the_first_case_returns_an_empty_partial_run_and_finalize .source( Box::new(IterCaseSource::new(items.into_iter())) as Box> ) - .task(Box::new(SimpleFnTask::new( - |_: &String| -> anyhow::Result { panic!("no case was read") }, - ))) - .scorers(vec![]) + .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(); @@ -514,8 +519,7 @@ async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() .collect::>>() .into(), ) - .task(Box::new(callback)) - .scorers(vec![]) + .definition(EvalDefinition::new(Box::new(callback), vec![]).unwrap()) .max_concurrency(2) .quiet(true) .build(); @@ -553,3 +557,60 @@ async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() assert_eq!(partial.failed_cases, 2); assert_eq!(partial.successful_cases, 0); } + +#[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); +} From 963fa5da3aee4d8f9a646affa7049009a57fd1c4 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:29:37 +0800 Subject: [PATCH 5/7] Isolate evaluation scheduling from recording and reporting --- src/eval/error.rs | 1 + src/eval/executor.rs | 475 +++++++++++++++++++++++++++++++++++++++++++ src/eval/mod.rs | 3 + src/eval/recorder.rs | 164 +++++++++++++++ src/eval/reporter.rs | 44 ++++ src/eval/runner.rs | 262 ++---------------------- src/eval/types.rs | 4 +- tests/eval_test.rs | 90 +++++++- 8 files changed, 789 insertions(+), 254 deletions(-) create mode 100644 src/eval/executor.rs create mode 100644 src/eval/recorder.rs create mode 100644 src/eval/reporter.rs diff --git a/src/eval/error.rs b/src/eval/error.rs index 9798a2f..6d9dc92 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -75,6 +75,7 @@ pub enum EvalRunFailure { /// 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>>, diff --git a/src/eval/executor.rs b/src/eval/executor.rs new file mode 100644 index 0000000..372939c --- /dev/null +++ b/src/eval/executor.rs @@ -0,0 +1,475 @@ +//! 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, + ) -> 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).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, + ) -> 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.task(), definition.scorers(), span).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, SimpleFnTask, 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)) + .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)) + .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(SimpleFnTask::new(|input: &usize| 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)) + .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 22d02cc..8fe357f 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -55,6 +55,9 @@ mod bt_runner; mod case_runner; mod definition; mod error; +mod executor; +mod recorder; +mod reporter; mod runner; mod scorer; mod source; diff --git a/src/eval/recorder.rs b/src/eval/recorder.rs new file mode 100644 index 0000000..0f8847f --- /dev/null +++ b/src/eval/recorder.rs @@ -0,0 +1,164 @@ +//! 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: Arc, + root: SpanHandle, + root_write: tokio::sync::Mutex>, +} +impl BraintrustRecorder { + pub async fn new( + client: Arc, + name: String, + metadata: Option>, + tags: Option>, + ) -> Result { + let root = client + .span_builder() + .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, + 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() + .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/runner.rs b/src/eval/runner.rs index 8263a28..498b11e 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -1,16 +1,12 @@ -use futures::stream::{FuturesUnordered, StreamExt}; use serde::Serialize; use serde_json::{Map, Value}; use std::marker::PhantomData; use std::sync::Arc; -use crate::{BraintrustClient, Result, SpanHandle}; +use crate::{BraintrustClient, Result}; use super::{ - case_runner::CaseRunner, - source::CaseSource, - types::{Case, EvalResult, EvalSummary}, - EvalConfigError, EvalDefinition, EvalRunError, EvalRunFailure, + source::CaseSource, types::EvalSummary, EvalConfigError, EvalDefinition, EvalRunError, }; /// Options for running an evaluation @@ -78,251 +74,17 @@ where /// Run the evaluation pub async fn run( &self, - mut opts: EvalOpts, + opts: EvalOpts, ) -> std::result::Result, EvalRunError> { - opts.validate()?; - 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()); - } - - let root_write = root_span.log(log_builder.build().unwrap()); - - let dataset = opts.source.provenance(); - let mut results = Vec::new(); - let mut failures = Vec::new(); - let mut case_count = 0; - let mut futures = FuturesUnordered::new(); - - loop { - // Retain the next-read future while polling active cases. A source - // may wait for prior computation, and need not be cancellation-safe. - let mut next = Box::pin(opts.source.next()); - let next_case = loop { - if futures.is_empty() { - break next.await; - } - tokio::select! { - next_case = next.as_mut() => break next_case, - Some((index, result)) = futures.next() => { - Self::collect_case_result(index, result, &mut results, &mut failures); - if !failures.is_empty() { break Ok(None); } - } - } - }; - let case = match next_case { - Ok(Some(case)) => case, - Ok(None) => break, - Err(source) => { - failures.push(EvalRunFailure::Source { - case_index: case_count, - source, - }); - break; - } - }; - case_count += 1; - let index = case_count - 1; - let client = Arc::clone(&self.client); - let root_span = root_span.clone(); - let definition = &opts.definition; - let quiet = opts.quiet; - futures.push(async move { - let result = - Self::process_case(case, definition, &root_span, &client, index + 1, quiet) - .await; - (index, result) - }); - if futures.len() >= opts.max_concurrency { - if let Some((index, result)) = futures.next().await { - Self::collect_case_result(index, result, &mut results, &mut failures); - } - if !failures.is_empty() { - break; - } - } - } - // Source/recording failure does not abandon already scheduled cases. - while let Some((index, result)) = futures.next().await { - Self::collect_case_result(index, result, &mut results, &mut failures); - } - let end = root_span.end(); - for result in [root_write.await, end.await] { - if let Err(error) = result { - failures.push(EvalRunFailure::Recording { - case_index: None, - source: error.into(), - }); - } - } - results.sort_unstable_by_key(|(index, _)| *index); - let results: Vec<_> = results.into_iter().map(|(_, result)| result).collect(); - - // Compute statistics - let score_stats = EvalSummary::::compute_stats(&results); - let successful = results.iter().filter(|r| r.errors.is_empty()).count(); - let failed = case_count - 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 - ); - } - } - - let summary = EvalSummary { - name: opts.name, - results, - score_stats, - total_cases: case_count, - successful_cases: successful, - failed_cases: failed, - dataset, - }; - if failures.is_empty() { - Ok(summary) - } else { - Err(EvalRunError { - failures, - partial: Some(Box::new(summary)), - }) - } - } - - fn collect_case_result( - index: usize, - result: Result>, - results: &mut Vec<(usize, EvalResult)>, - failures: &mut Vec, - ) { - match result { - Ok(result) => results.push((index, result)), - Err(source) => failures.push(EvalRunFailure::Recording { - case_index: Some(index), - source, - }), - } - } - - /// Process a single case - async fn process_case( - case: Case, - definition: &EvalDefinition, - _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()); - } - - let initial_write = case_span.log(log_builder.build().unwrap()); - let mut result = CaseRunner::run( - case, - definition.task(), - definition.scorers(), - Some(&case_span), - ) - .await; - result.span_id = case_span.export()?.span_id; - - // Recording adapts the completed outcome; computation never waits on - // HTTP or emits console messages. Observe every accepted write before - // reporting recording success, including writes that failed earlier. - let mut final_log = crate::SpanLog::builder().scores( - result - .scores - .iter() - .map(|score| (score.name.clone(), score.score)) - .collect(), - ); - if let Some(output) = &result.output { - if let Ok(output) = serde_json::to_value(output) { - final_log = final_log.output(output); - } - } - if let Some(metadata) = &result.metadata { - final_log = final_log.metadata(metadata.clone()); - } - if let Some(tags) = &result.tags { - final_log = final_log.tags(tags.clone()); - } - if !result.errors.is_empty() { - final_log = final_log.error( - serde_json::to_value(&result.errors) - .map_err(|error| crate::BraintrustError::Background(error.to_string()))?, - ); - } - let final_write = case_span.log(final_log.build().unwrap()); - let end = case_span.end(); - let initial_result = initial_write.await; - let final_result = final_write.await; - let end_result = end.await; - initial_result?; - final_result?; - end_result?; - - if !quiet { - let score_str = result - .scores - .iter() - .map(|s| format!("{}={:.3}", s.name, s.score)) - .collect::>() - .join(", "); - println!("Case {}: {}", case_num, score_str); - for error in &result.errors { - eprintln!("Case {}: {}", case_num, error); - } - } - Ok(result) + let name = opts.name.clone(); + let metadata = opts.metadata.clone(); + let tags = opts.tags.clone(); + let mut reporter = super::reporter::TerminalReporter { quiet: opts.quiet }; + let plan = super::executor::EvalPlan::new(opts)?; + let recorder = + super::recorder::BraintrustRecorder::new(self.client.clone(), name, metadata, tags) + .await?; + super::executor::EvalExecutor::run(plan, Some(&recorder), &mut reporter).await } } diff --git a/src/eval/types.rs b/src/eval/types.rs index 9078148..8b34da0 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -155,9 +155,9 @@ 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, /// Full dataset provenance captured before source consumption. #[serde(skip_serializing_if = "Option::is_none")] diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 0c6fdc5..68e9a42 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -554,8 +554,21 @@ async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() ))); let partial = error.partial.unwrap(); assert_eq!(partial.total_cases, 2); - assert_eq!(partial.failed_cases, 2); - assert_eq!(partial.successful_cases, 0); + 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] @@ -614,3 +627,76 @@ async fn reusable_definition_keeps_recorded_runs_and_sources_independent() { 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"); +} From f5cdab804a374ab513a3fdce8a5d9c66215d6cf2 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:40:44 +0800 Subject: [PATCH 6/7] Submit local and recorded evaluation runs eagerly --- examples/eval_sync.rs | 28 +-- src/eval/error.rs | 6 +- src/eval/mod.rs | 14 +- src/eval/recorder.rs | 11 +- src/eval/run.rs | 128 ++++++++++++++ src/eval/runner.rs | 54 ++---- src/logger.rs | 24 ++- tests/eager_eval.rs | 393 ++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 588 insertions(+), 70 deletions(-) create mode 100644 src/eval/run.rs create mode 100644 tests/eager_eval.rs diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index ac3c245..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![ @@ -40,12 +36,16 @@ fn main() -> anyhow::Result<()> { println!("Starting synchronous evaluation...\n"); let definition = EvalDefinition::new(Box::new(task), scorers)?; - let result = evaluator.run( - EvalOpts::builder() - .name("sync-uppercase-eval".to_string()) - .source(cases.into()) - .definition(definition.clone()) - .build(), + 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/error.rs b/src/eval/error.rs index 6d9dc92..37b11dc 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -43,6 +43,8 @@ impl Serialize for CaseError { 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}")] @@ -53,6 +55,8 @@ pub enum EvalConfigError { /// 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}")] @@ -81,7 +85,7 @@ pub struct EvalRunError { pub partial: Option>>, } impl EvalRunError { - fn new(failure: EvalRunFailure) -> Self { + pub(crate) fn new(failure: EvalRunFailure) -> Self { Self { failures: vec![failure], partial: None, diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 8fe357f..1da6849 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -19,8 +19,6 @@ //! .build() //! .await?; //! -//! let evaluator = Evaluator::new(client); -//! //! let cases = vec![ //! Case { //! input: "hello".to_string(), @@ -36,15 +34,15 @@ //! }))), //! vec![Box::new(ExactMatch::new())], //! )?; -//! let result = evaluator -//! .run( +//! let completion = run(EvalRequest::recorded( //! EvalOpts::builder() //! .name("uppercase-eval".to_string()) //! .source(cases.into()) //! .definition(definition.clone()) //! .build(), -//! ) -//! .await?; +//! client, +//! )); +//! let result = completion.await?; //! //! println!("Results: {:?}", result); //! # Ok(()) @@ -58,6 +56,7 @@ mod error; mod executor; mod recorder; mod reporter; +mod run; mod runner; mod scorer; mod source; @@ -69,7 +68,8 @@ pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; -pub use runner::{EvalOpts, Evaluator, SyncEvaluator}; +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, diff --git a/src/eval/recorder.rs b/src/eval/recorder.rs index 0f8847f..4f8c6e4 100644 --- a/src/eval/recorder.rs +++ b/src/eval/recorder.rs @@ -15,19 +15,21 @@ use std::sync::Arc; type Write = Completion; pub(crate) struct BraintrustRecorder { - client: Arc, + client: BraintrustClient, + submitter: Arc, root: SpanHandle, root_write: tokio::sync::Mutex>, } impl BraintrustRecorder { pub async fn new( - client: Arc, + client: BraintrustClient, + submitter: Arc, name: String, metadata: Option>, tags: Option>, ) -> Result { let root = client - .span_builder() + .span_builder_for_operation(submitter.clone()) .await? .span_type(crate::SpanType::Eval) .build(); @@ -41,6 +43,7 @@ impl BraintrustRecorder { let root_write = tokio::sync::Mutex::new(Some(root.log(log.build().unwrap()))); Ok(Self { client, + submitter, root, root_write, }) @@ -74,7 +77,7 @@ impl async fn start(&self, index: usize, case: &Case) -> Result>> { let span = self .client - .span_builder() + .span_builder_for_operation(self.submitter.clone()) .await? .span_type(crate::SpanType::Eval) .build(); diff --git a/src/eval/run.rs b/src/eval/run.rs new file mode 100644 index 0000000..4625872 --- /dev/null +++ b/src/eval/run.rs @@ -0,0 +1,128 @@ +//! 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))) + } + }; + binding.response(async move { + let mut reporter = TerminalReporter { + quiet: options.quiet, + }; + EvalExecutor::run(EvalPlan::new(options)?, None, &mut reporter).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(); + 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).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 498b11e..87f9bc3 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -3,7 +3,7 @@ use serde_json::{Map, Value}; use std::marker::PhantomData; use std::sync::Arc; -use crate::{BraintrustClient, Result}; +use crate::BraintrustClient; use super::{ source::CaseSource, types::EvalSummary, EvalConfigError, EvalDefinition, EvalRunError, @@ -32,6 +32,8 @@ 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, } impl EvalOpts { @@ -44,7 +46,7 @@ impl EvalOpts { } } -/// Async evaluator for running evaluations +/// Client-bound convenience facade for recorded evaluations. pub struct Evaluator { client: Arc, _phantom: PhantomData<(I, O)>, @@ -71,50 +73,16 @@ where } } - /// Run the evaluation - pub async fn run( - &self, - opts: EvalOpts, - ) -> std::result::Result, EvalRunError> { - let name = opts.name.clone(); - let metadata = opts.metadata.clone(); - let tags = opts.tags.clone(); - let mut reporter = super::reporter::TerminalReporter { quiet: opts.quiet }; - let plan = super::executor::EvalPlan::new(opts)?; - let recorder = - super::recorder::BraintrustRecorder::new(self.client.clone(), name, metadata, tags) - .await?; - super::executor::EvalExecutor::run(plan, Some(&recorder), &mut reporter).await - } -} - -/// Sync evaluator (wraps async evaluator with Runtime::block_on) -pub struct SyncEvaluator { - evaluator: Evaluator, - runtime: tokio::runtime::Runtime, -} - -impl SyncEvaluator -where - I: Serialize + Send + Sync + 'static, - O: Serialize + Send + Sync + '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 + /// Submit a recorded run immediately on the client binding. + /// Dropping the completion does not cancel accepted work. pub fn run( &self, opts: EvalOpts, - ) -> std::result::Result, EvalRunError> { - self.runtime.block_on(self.evaluator.run(opts)) + ) -> crate::Completion, EvalRunError> { + super::run(super::EvalRequest::recorded( + opts, + self.client.as_ref().clone(), + )) } } 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/eager_eval.rs b/tests/eager_eval.rs new file mode 100644 index 0000000..64853fd --- /dev/null +++ b/tests/eager_eval.rs @@ -0,0 +1,393 @@ +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)) + )); +} From 0308efbf7a6019797a161e5fd2dfbdd79d016325 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:52:39 +0800 Subject: [PATCH 7/7] Run synchronous evaluation callbacks on the managed blocking pool --- src/eval/blocking.rs | 68 +++++++++++++++++ src/eval/case_runner.rs | 107 +++++++++++++++++--------- src/eval/definition.rs | 42 ++++++---- src/eval/executor.rs | 33 +++++--- src/eval/mod.rs | 1 + src/eval/run.rs | 19 ++++- src/eval/scorer.rs | 47 +++++++++++- src/eval/task.rs | 33 +++++++- src/eval/types.rs | 14 ++++ tests/blocking_eval.rs | 165 ++++++++++++++++++++++++++++++++++++++++ tests/eager_eval.rs | 125 ++++++++++++++++++++++++++++++ 11 files changed, 586 insertions(+), 68 deletions(-) create mode 100644 src/eval/blocking.rs create mode 100644 tests/blocking_eval.rs 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/case_runner.rs b/src/eval/case_runner.rs index ec39207..718368d 100644 --- a/src/eval/case_runner.rs +++ b/src/eval/case_runner.rs @@ -1,29 +1,40 @@ -//! One case's computation, independent of scheduling, recording and reporting. -use super::{Case, CaseError, EvalResult, Scorer, ScorerArgs, Task, TaskHooks}; -use crate::SpanHandle; +//! 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( + pub async fn run( case: Case, - task: &dyn Task, - scorers: &[Box>], + definition: &EvalDefinition, span: Option<&SpanHandle>, - ) -> EvalResult - where - I: Send + Sync, - O: Send + Sync, - { + 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 task.run(&case.input, &hooks).await { - Ok(output) => Some(output), + let output = match computed { + Ok(output) => Some(Arc::new(output)), Err(source) => { errors.push(CaseError::Task { source }); None @@ -31,24 +42,34 @@ impl CaseRunner { }; let mut scores = Vec::new(); if let Some(output) = &output { - // Drive scorers concurrently while keeping definition order in the - // outcome, regardless of their relative completion times. let mut pending = FuturesOrdered::new(); - for scorer in scorers { - let args = ScorerArgs { - input: &case.input, - output, - expected: case.expected.as_ref(), - metadata: case.metadata.as_ref(), - }; + for scorer in definition.scorers() { + let case = case.clone(); + let output = output.clone(); + let scorer = scorer.clone(); pending.push_back(async move { - scorer - .score(args) - .await - .map_err(|source| CaseError::Scorer { - name: scorer.name().to_owned(), - source, + 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 { @@ -61,6 +82,13 @@ impl CaseRunner { 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, @@ -80,7 +108,7 @@ impl CaseRunner { #[cfg(test)] mod tests { use super::*; - use crate::eval::{Score, Scores, SimpleFnTask}; + use crate::eval::{Score, Scorer, Scores, Task}; use async_trait::async_trait; use serde_json::json; use std::sync::{ @@ -123,7 +151,12 @@ mod tests { xact_id: Some("version".into()), created: Some("timestamp".into()), }; - let result = futures::executor::block_on(CaseRunner::run(case, &LocalTask, &[], None)); + 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")); @@ -174,8 +207,8 @@ mod tests { input: 1, ..Default::default() }, - &FailTask, - &scorers, + &EvalDefinition::new(Box::new(FailTask), scorers).unwrap(), + None, None, ) .await; @@ -233,8 +266,14 @@ mod tests { input: 1, ..Default::default() }, - &SimpleFnTask::new(|input: &usize| Ok(*input)), - &scorers, + &EvalDefinition::new( + Box::new(crate::eval::task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), + scorers, + ) + .unwrap(), + None, None, ), ) diff --git a/src/eval/definition.rs b/src/eval/definition.rs index 11adddd..f4c48c8 100644 --- a/src/eval/definition.rs +++ b/src/eval/definition.rs @@ -15,8 +15,8 @@ pub struct EvalDefinition { } struct DefinitionInner { - task: Box>, - scorers: Vec>>, + task: Arc>, + scorers: Vec>>, } impl Clone for EvalDefinition { @@ -37,7 +37,7 @@ impl EvalDefinition { scorers: Vec>>, ) -> Result { let mut names = std::collections::HashSet::new(); - let mut validated: Vec>> = Vec::with_capacity(scorers.len()); + 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() { @@ -46,11 +46,11 @@ impl EvalDefinition { if !names.insert(name.clone()) { return Err(EvalConfigError::DuplicateScorerName { name }); } - validated.push(Box::new(NamedScorer { name, scorer })); + validated.push(Arc::new(NamedScorer { name, scorer })); } Ok(Self { inner: Arc::new(DefinitionInner { - task, + task: Arc::from(task), scorers: validated, }), }) @@ -63,11 +63,11 @@ impl EvalDefinition { self.inner.scorers.iter().map(|scorer| scorer.name()) } - pub(crate) fn task(&self) -> &dyn Task { - self.inner.task.as_ref() + pub(crate) fn task(&self) -> Arc> { + self.inner.task.clone() } - pub(crate) fn scorers(&self) -> &[Box>] { + pub(crate) fn scorers(&self) -> &[Arc>] { &self.inner.scorers } } @@ -79,10 +79,18 @@ struct NamedScorer { #[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 } @@ -91,7 +99,7 @@ impl Scorer for NamedScorer { #[cfg(test)] mod tests { use super::*; - use crate::eval::{case_runner::CaseRunner, task, Case, CaseError, SimpleFnTask}; + 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. @@ -133,8 +141,8 @@ mod tests { }; let (one, two) = futures::executor::block_on(async { futures::join!( - CaseRunner::run(case(1, 10), first.task(), first.scorers(), None), - CaseRunner::run(case(2, 20), other.task(), other.scorers(), None), + CaseRunner::run(case(1, 10), &first, None, None), + CaseRunner::run(case(2, 20), &other, None, None), ) }); assert_eq!(calls.load(Ordering::SeqCst), 2); @@ -145,8 +153,8 @@ mod tests { // 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.task(), first.scorers(), None), - CaseRunner::run(case(4, 40), other.task(), other.scorers(), None), + CaseRunner::run(case(3, 30), &first, None, None), + CaseRunner::run(case(4, 40), &other, None, None), ) }); assert_eq!(one.output.unwrap().0, 33); @@ -173,7 +181,9 @@ mod tests { fn scorer_identity_is_captured_once_for_every_run() { let changed = Arc::new(AtomicBool::new(false)); let definition = EvalDefinition::new( - Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), vec![Box::new(ChangingName(changed.clone()))], ) .unwrap(); @@ -185,8 +195,8 @@ mod tests { input: 1, ..Default::default() }, - definition.task(), - definition.scorers(), + &definition, + None, None, )); assert!( diff --git a/src/eval/executor.rs b/src/eval/executor.rs index 372939c..98e4b3a 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -55,10 +55,11 @@ struct CaseOutcome { pub(crate) struct EvalExecutor; impl EvalExecutor { - pub async fn run( + 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(); @@ -102,7 +103,7 @@ impl EvalExecutor { pending.push(async move { ( index, - Self::run_case(case, definition, recorder, index).await, + Self::run_case(case, definition, recorder, index, binding).await, ) }); if pending.len() >= plan.concurrency { @@ -153,11 +154,12 @@ impl EvalExecutor { } } - async fn run_case( + 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 { @@ -172,7 +174,7 @@ impl EvalExecutor { None }; let span = recording.as_ref().and_then(|recording| recording.span()); - let mut result = CaseRunner::run(case, definition.task(), definition.scorers(), span).await; + let mut result = CaseRunner::run(case, definition, span, binding).await; if let Some(recording) = recording { failures.extend(recording.finish(&mut result).await); } @@ -203,7 +205,7 @@ impl EvalExecutor { #[cfg(test)] mod tests { use super::*; - use crate::eval::{task, CaseSource, SimpleFnTask, SourceError}; + use crate::eval::{task, CaseSource, SourceError}; use std::{ cell::Cell, sync::{ @@ -290,7 +292,7 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let error = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter)) + let error = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter, None)) .err() .expect("source should fail"); assert!( @@ -356,8 +358,9 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let summary = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter)) - .expect("local computation should succeed"); + 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!( @@ -422,7 +425,9 @@ mod tests { .source(cases.into()) .definition( EvalDefinition::new( - Box::new(SimpleFnTask::new(|input: &usize| Ok(input + 10))), + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(input + 10) }) + })), vec![], ) .unwrap(), @@ -432,9 +437,13 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let error = - futures::executor::block_on(EvalExecutor::run(plan, Some(&recorder), &mut reporter)) - .unwrap_err(); + 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], diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 1da6849..2b293c5 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -49,6 +49,7 @@ //! # } //! ``` +mod blocking; mod bt_runner; mod case_runner; mod definition; diff --git a/src/eval/run.rs b/src/eval/run.rs index 4625872..cb784ed 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -82,11 +82,18 @@ pub fn run( 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).await + EvalExecutor::run( + EvalPlan::new(options)?, + None, + &mut reporter, + Some(&callbacks), + ) + .await }) } Target::Recorded { client, submit } => { @@ -103,6 +110,7 @@ fn recorded, ) -> RunCompletion { let binding = client.execution_binding(); + let callbacks = binding.clone(); binding.response_try(|| { let operation = client.span_operation()?; Ok(async move { @@ -117,8 +125,13 @@ fn recorded: 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; } @@ -48,7 +60,8 @@ where /// 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, @@ -78,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) } } @@ -116,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, @@ -205,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)]), @@ -274,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)]), diff --git a/src/eval/task.rs b/src/eval/task.rs index 18ebb90..c3ca4ee 100644 --- a/src/eval/task.rs +++ b/src/eval/task.rs @@ -12,6 +12,18 @@ use super::types::TaskHooks; /// 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,7 +60,15 @@ 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.run_blocking(input, hooks) + } + + fn run_blocking(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { (self.func)(input, hooks) } } @@ -74,7 +95,7 @@ where } } -/// 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, @@ -102,7 +123,15 @@ 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.run_blocking(input, _hooks) + } + + fn run_blocking(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { (self.func)(input) } } diff --git a/src/eval/types.rs b/src/eval/types.rs index 8b34da0..49b427c 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -241,6 +241,20 @@ impl<'a, I, O> TaskHooks<'a, I, O> { } } + 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, + } + } + /// Add metadata key-value pair pub fn add_metadata(&self, key: impl Into, value: Value) { if let Ok(mut metadata) = self.metadata.lock() { 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 index 64853fd..f7a956b 100644 --- a/tests/eager_eval.rs +++ b/tests/eager_eval.rs @@ -391,3 +391,128 @@ async fn shutdown_deadline_cancels_pending_run_and_revokes_escaped_permission() 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(); +}