From 86a3144194db9dd5ab7c194967ef4d1b096203bf Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 02:03:59 +0800 Subject: [PATCH 1/6] Record evaluation cases into resolved experiments --- src/eval/executor.rs | 55 +++++-- src/eval/mod.rs | 4 +- src/eval/recorder.rs | 76 +++++----- src/eval/run.rs | 7 +- src/eval/runner.rs | 9 +- src/eval/types.rs | 11 ++ src/experiments/experiment.rs | 2 +- src/logger.rs | 36 ++++- tests/eager_eval.rs | 48 ++++++ tests/eval_test.rs | 268 +++++++++++++++++++++++++++++++--- 10 files changed, 438 insertions(+), 78 deletions(-) diff --git a/src/eval/executor.rs b/src/eval/executor.rs index 98e4b3a..9161636 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -17,6 +17,8 @@ pub(crate) struct EvalPlan { source: Box>, definition: EvalDefinition, concurrency: usize, + metadata: Option>, + tags: Option>, } impl EvalPlan { pub fn new(opts: EvalOpts) -> std::result::Result { @@ -26,12 +28,35 @@ impl EvalPlan { source: opts.source, definition: opts.definition, concurrency: opts.max_concurrency, + metadata: opts.metadata, + tags: opts.tags, }) } + fn apply_context(&self, case: &mut Case) { + if let Some(defaults) = &self.metadata { + let mut metadata = defaults.clone(); + if let Some(specific) = case.metadata.take() { + metadata.extend(specific); + } + case.metadata = Some(metadata); + } + if let Some(defaults) = &self.tags { + let mut tags = Vec::new(); + for tag in defaults.iter().chain(case.tags.iter().flatten()) { + if !tags.contains(tag) { + tags.push(tag.clone()); + } + } + case.tags = Some(tags); + } + } } #[async_trait] pub(crate) trait EvalRecorder: Send + Sync { + fn target(&self) -> Option { + None + } async fn start(&self, index: usize, case: &Case) -> Result>>; async fn finish(&self) -> Vec; } @@ -70,23 +95,25 @@ impl EvalExecutor { 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); + let next_case = { + let mut next = Box::pin(plan.source.next()); + 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"), } - Either::Right((None, _)) => unreachable!("nonempty pending cases"), } }; - let case = match next_case { + let mut case = match next_case { Ok(Some(case)) => case, Ok(None) => break, Err(source) => { @@ -97,6 +124,7 @@ impl EvalExecutor { break; } }; + plan.apply_context(&mut case); let index = case_count; case_count += 1; let definition = &plan.definition; @@ -142,6 +170,7 @@ impl EvalExecutor { successful_cases, failed_cases: case_count - successful_cases, dataset, + recording: recorder.and_then(|recorder| recorder.target()), }; reporter.finished(&summary); if failures.is_empty() { diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 2b293c5..121f193 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -76,4 +76,6 @@ 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}; +pub use types::{ + Case, EvalRecording, EvalResult, EvalSummary, Score, ScoreStats, ScorerArgs, Scores, TaskHooks, +}; diff --git a/src/eval/recorder.rs b/src/eval/recorder.rs index 4f8c6e4..10bf5bc 100644 --- a/src/eval/recorder.rs +++ b/src/eval/recorder.rs @@ -1,7 +1,7 @@ //! Braintrust serialization and write completions, outside the run engine. use super::{ executor::{CaseRecording, EvalRecorder}, - Case, EvalResult, + Case, EvalRecording, EvalResult, }; use crate::{ BraintrustClient, BraintrustError, Completion, DeliveryReceipt, Result, SpanHandle, SpanLog, @@ -15,38 +15,43 @@ use std::sync::Arc; type Write = Completion; pub(crate) struct BraintrustRecorder { - client: BraintrustClient, - submitter: Arc, - root: SpanHandle, - root_write: tokio::sync::Mutex>, + template: crate::span::SpanBuilder, + target: EvalRecording, } impl BraintrustRecorder { pub async fn new( client: BraintrustClient, submitter: Arc, - name: String, + project_name: Option, + experiment_name: String, metadata: Option>, - tags: Option>, ) -> Result { - let root = client - .span_builder_for_operation(submitter.clone()) + let mut builder = client + .experiment_builder_for_operation(submitter.clone(), project_name) .await? - .span_type(crate::SpanType::Eval) - .build(); - let mut log = SpanLog::builder().name(format!("eval-{name}")); + .experiment_name(experiment_name); if let Some(metadata) = metadata { - log = log.metadata(metadata); - } - if let Some(tags) = tags { - log = log.tags(tags); + builder = builder.metadata(metadata); } - let root_write = tokio::sync::Mutex::new(Some(root.log(log.build().unwrap()))); - Ok(Self { - client, - submitter, - root, - root_write, - }) + let experiment = builder + .build() + .map_err(|error| BraintrustError::InvalidConfig(error.to_string()))?; + let resolved = experiment.register_metadata().await?; + let target = EvalRecording { + project_id: resolved.project_id.clone(), + experiment_id: resolved.experiment_id.clone(), + experiment_name: resolved.experiment_name.clone(), + }; + let template = client + .span_builder_for_operation( + submitter, + Some(crate::types::ParentSpanInfo::Experiment { + object_id: target.experiment_id.clone(), + }), + ) + .await? + .span_type(crate::SpanType::Eval); + Ok(Self { template, target }) } } @@ -76,10 +81,9 @@ impl { async fn start(&self, index: usize, case: &Case) -> Result>> { let span = self - .client - .span_builder_for_operation(self.submitter.clone()) - .await? - .span_type(crate::SpanType::Eval) + .template + .clone() + .name(format!("case-{}", index + 1)) .build(); let mut failures = Vec::new(); let mut log = SpanLog::builder().name(format!("case-{}", index + 1)); @@ -104,18 +108,14 @@ impl failures, })) } + fn target(&self) -> Option { + Some(self.target.clone()) + } + 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 + // Every owned case write and end has been observed before the executor + // finalizes this recorder. The run summary is not a trace span. + Vec::new() } } #[async_trait] diff --git a/src/eval/run.rs b/src/eval/run.rs index cb784ed..694c95b 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -120,9 +120,12 @@ fn recorded { pub source: Box>, /// Reusable task and ordered scorers pub definition: EvalDefinition, - /// Braintrust project to log results to + /// Braintrust project for recording; falls back to the client default project. pub project_name: Option, - /// Optional experiment name for Braintrust logging + /// Experiment name for recording; defaults to the evaluation name. pub experiment_name: Option, - /// Optional metadata for the evaluation + /// Run metadata stored on the experiment and used as case defaults. + /// Case metadata overrides matching keys. pub metadata: Option>, - /// Optional tags for the evaluation + /// Case tag defaults, combined with source case tags without duplicates. pub tags: Option>, /// Maximum concurrent evaluations (default: 5) #[builder(default = 5)] diff --git a/src/eval/types.rs b/src/eval/types.rs index 49b427c..7163980 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -144,6 +144,14 @@ pub struct ScoreStats { pub count: usize, } +/// Experiment identity captured before a recorded run reads cases. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EvalRecording { + pub project_id: String, + pub experiment_id: String, + pub experiment_name: String, +} + /// Summary of evaluation results #[derive(Debug, Serialize)] pub struct EvalSummary { @@ -162,6 +170,9 @@ pub struct EvalSummary { /// Full dataset provenance captured before source consumption. #[serde(skip_serializing_if = "Option::is_none")] pub dataset: Option, + /// Resolved recording target; absent for local runs. + #[serde(skip_serializing_if = "Option::is_none")] + pub recording: Option, } impl EvalSummary { diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 726a32e..4faf5f1 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -545,7 +545,7 @@ impl Experiment { } /// Polled only on the selected context, by preparation or a bound operation. - async fn register_metadata(&self) -> Result<&ExperimentMetadata> { + pub(crate) async fn register_metadata(&self) -> Result<&ExperimentMetadata> { self.state .metadata .get_or_try_init(|| async { diff --git a/src/logger.rs b/src/logger.rs index 191f73f..bfea629 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -622,12 +622,42 @@ impl BraintrustClient { } async fn span_builder_inner(&self) -> Result { - self.span_builder_for_operation(self.span_submitter()).await + self.span_builder_for_operation(self.span_submitter(), None) + .await + } + + pub(crate) async fn experiment_builder_for_operation( + &self, + submitter: Arc, + project: Option, + ) -> Result { + let project = project + .or_else(|| self.inner.default_project.clone()) + .filter(|name| !name.trim().is_empty()) + .ok_or_else(|| { + BraintrustError::InvalidConfig( + "recorded evaluations require project_name or a client default project".into(), + ) + })?; + if self.inner.login_skipped { + return Err(BraintrustError::InvalidConfig( + "recording requires login".into(), + )); + } + let state = self.wait_for_login_state().await?; + Ok(ExperimentBuilder::new( + self.resource_transport(), + self.inner.execution.clone(), + submitter, + state.scope()?, + ) + .project_name(project)) } pub(crate) async fn span_builder_for_operation( &self, submitter: Arc, + parent: Option, ) -> Result { if self.inner.login_skipped { return Err(BraintrustError::InvalidConfig( @@ -642,7 +672,9 @@ impl BraintrustClient { if let Some(span_origin) = &self.inner.span_origin { builder = builder.span_origin(span_origin.clone()); } - if let Some(ref project) = self.inner.default_project { + if let Some(parent) = parent { + builder = builder.parent_info(parent); + } else if let Some(ref project) = self.inner.default_project { builder = builder.project_name(project); } Ok(builder) diff --git a/tests/eager_eval.rs b/tests/eager_eval.rs index f7a956b..ed49156 100644 --- a/tests/eager_eval.rs +++ b/tests/eager_eval.rs @@ -175,6 +175,10 @@ async fn client(server: &MockServer) -> BraintrustClient { "/api/project/register", serde_json::json!({"project":{"id":"project"}}), ), + ( + "/api/experiment/register", + serde_json::json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}}), + ), ("/logs3", serde_json::json!({})), ] { Mock::given(method("POST")) @@ -516,3 +520,47 @@ async fn synchronous_callback_can_wait_for_recording_without_blocking_delivery() assert_eq!(result.unwrap().results[0].output, Some(9)); shutdown.unwrap(); } + +#[test] +fn local_run_context_uses_case_overrides_without_a_recording_target() { + let definition = EvalDefinition::new( + Box::new(SyncFnTask::new( + |input: &usize, hooks: &TaskHooks<'_, usize, usize>| { + assert_eq!(hooks.get_metadata().unwrap()["run"], true); + assert_eq!(hooks.get_metadata().unwrap()["collision"], "case"); + hooks.add_metadata("task", serde_json::json!(true)); + Ok(*input) + }, + )), + vec![], + ) + .unwrap(); + let mut opts = options( + definition, + vec![Case { + input: 1, + metadata: Some( + serde_json::json!({"collision":"case"}) + .as_object() + .unwrap() + .clone(), + ), + tags: Some(vec!["shared".into(), "case".into()]), + ..Default::default() + }], + ); + opts.metadata = Some( + serde_json::json!({"run":true, "collision":"run"}) + .as_object() + .unwrap() + .clone(), + ); + opts.tags = Some(vec!["run".into(), "shared".into()]); + let summary = futures::executor::block_on(run(EvalRequest::local(opts)).into_future()).unwrap(); + assert!(summary.recording.is_none()); + assert_eq!(summary.results[0].metadata.as_ref().unwrap()["task"], true); + assert_eq!( + summary.results[0].tags.as_ref().unwrap(), + &["run", "shared", "case"] + ); +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 68e9a42..7c9bcc8 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -15,6 +15,10 @@ async fn recording_client() -> (BraintrustClient, wiremock::MockServer) { json!({"org_info":[{"id":"org","name":"org"}]}), ), ("/api/project/register", json!({"project":{"id":"project"}})), + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}}), + ), ("/logs3", json!({})), ] { Mock::given(method("POST")) @@ -415,13 +419,13 @@ async fn source_failure_drains_active_cases_preserves_provenance_and_finishes_re .collect(); assert_eq!( ended.len(), - 3, - "both cases and the run span must finish before returning partial results" + 2, + "both case roots must finish before returning partial results" ); } #[tokio::test] -async fn failure_before_the_first_case_returns_an_empty_partial_run_and_finalizes_root() { +async fn failure_before_the_first_case_returns_the_resolved_target_without_trace_rows() { let (client, server) = recording_client().await; let items: Vec>> = vec![Err(anyhow::anyhow!("source unavailable"))]; @@ -451,23 +455,13 @@ async fn failure_before_the_first_case_returns_an_empty_partial_run_and_finalize assert_eq!(partial.failed_cases, 0); assert!(partial.results.is_empty()); assert!(partial.dataset.is_none()); - let ended = server + assert_eq!(partial.recording.unwrap().experiment_id, "experiment"); + assert!(!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" - ); + .any(|request| request.url.path() == "/logs3")); } #[tokio::test] @@ -698,5 +692,245 @@ async fn serialization_failures_keep_outputs_and_finalize_recording() { .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"); + assert_eq!(ended.len(), 1, "case recording must finalize"); +} + +#[tokio::test] +async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independent_traces() { + use serde_json::json; + use std::{ + cell::Cell, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + }; + use wiremock::{ + matchers::{header, method, path}, + Mock, ResponseTemplate, + }; + let (client, server) = recording_client().await; + let registered = Arc::new(AtomicBool::new(false)); + Mock::given(method("POST")).and(path("/api/experiment/register")) + .and(header("authorization", "Bearer test-key")) + .and(header("x-bt-org-name", "org")) + .respond_with({ + let registered = registered.clone(); + move |_: &wiremock::Request| { + registered.store(true, Ordering::SeqCst); + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"chosen-project","name":"chosen"},"experiment":{"id":"chosen-id","name":"resolved-name"}})) + } + }).with_priority(1).expect(1).mount(&server).await; + struct Source { + cursor: Cell, + registered: Arc, + } + #[async_trait::async_trait] + impl CaseSource for Source { + async fn next(&mut self) -> anyhow::Result>> { + assert!(self.registered.load(Ordering::SeqCst)); + let index = self.cursor.get(); + self.cursor.set(index + 1); + Ok((index < 2).then(|| Case { + input: index, + metadata: Some(json!({"collision":"case"}).as_object().unwrap().clone()), + tags: Some(vec!["shared".into(), "case".into()]), + ..Default::default() + })) + } + } + let definition = EvalDefinition::new( + Box::new(task(|input: &usize, hooks| { + Box::pin(async move { + assert_eq!(hooks.get_metadata().unwrap()["shared"], "run"); + assert_eq!(hooks.get_metadata().unwrap()["collision"], "case"); + let child = hooks + .span + .unwrap() + .child_builder() + .unwrap() + .name("nested") + .build(); + child + .log(SpanLog::builder().input(json!(input)).build().unwrap()) + .await?; + child.end().await?; + hooks.add_tag("task"); + Ok(*input * 2) + }) + })), + vec![], + ) + .unwrap(); + let opts = EvalOpts::builder() + .name("named-eval".into()) + .source(Box::new(Source { + cursor: Cell::new(0), + registered, + })) + .definition(definition) + .project_name("chosen".into()) + .experiment_name("requested-name".into()) + .metadata( + json!({"shared":"run", "collision":"run"}) + .as_object() + .unwrap() + .clone(), + ) + .tags(vec!["run".into(), "shared".into()]) + .quiet(true) + .build(); + let summary = run(EvalRequest::recorded(opts, client)).await.unwrap(); + assert_eq!( + summary.recording.as_ref().unwrap(), + &EvalRecording { + project_id: "chosen-project".into(), + experiment_id: "chosen-id".into(), + experiment_name: "resolved-name".into(), + } + ); + assert_eq!( + summary.results[0].tags.as_ref().unwrap(), + &["run", "shared", "case", "task"] + ); + let requests = server.received_requests().await.unwrap(); + assert!(!requests + .iter() + .any(|request| request.url.path() == "/api/project/register")); + let registration: serde_json::Value = serde_json::from_slice( + &requests + .iter() + .find(|request| request.url.path() == "/api/experiment/register") + .unwrap() + .body, + ) + .unwrap(); + assert_eq!(registration["project_name"], "chosen"); + assert_eq!(registration["experiment_name"], "requested-name"); + assert_eq!(registration["metadata"]["collision"], "run"); + let rows: Vec = requests + .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() + .all(|row| row["experiment_id"] == "chosen-id" && row.get("project_id").is_none())); + assert!(rows + .iter() + .all(|row| row["span_attributes"]["name"] != "eval-named-eval")); + let mut roots = std::collections::HashSet::new(); + for case in &summary.results { + let id = case.span_id.as_ref().unwrap(); + let root = rows + .iter() + .find(|row| row["span_id"] == *id && row["metrics"]["end"].is_number()) + .unwrap(); + assert_eq!(root["root_span_id"], *id); + roots.insert(id.clone()); + let child = rows + .iter() + .find(|row| row["span_attributes"]["name"] == "nested" && row["root_span_id"] == *id) + .unwrap(); + assert_eq!(child["span_parents"][0], *id); + } + assert_eq!(roots.len(), 2); + assert_eq!( + summary_to_value(&summary)["recording"]["experiment_id"], + "chosen-id" + ); +} + +#[tokio::test] +async fn registration_failure_reads_no_cases_and_a_later_run_can_record_normally() { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + use wiremock::{ + matchers::{method, path}, + Mock, ResponseTemplate, + }; + let (client, server) = recording_client().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(403).set_body_string("registration denied")) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + let reads = Arc::new(AtomicUsize::new(0)); + 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 opts = EvalOpts::builder() + .name("blocked".into()) + .source(Box::new(IterCaseSource::new({ + let reads = reads.clone(); + (0..1).map(move |input| { + reads.fetch_add(1, Ordering::SeqCst); + Ok(Case { + input, + ..Default::default() + }) + }) + }))) + .definition(definition.clone()) + .quiet(true) + .build(); + let error = run(EvalRequest::recorded(opts, client.clone())) + .await + .unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Recording { + case_index: None, + source: BraintrustError::Api { status: 403, .. } + }] + )); + assert!(error.partial.is_none()); + assert_eq!(reads.load(Ordering::SeqCst), 0); + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(!server + .received_requests() + .await + .unwrap() + .iter() + .any(|request| request.url.path() == "/logs3")); + let opts = EvalOpts::builder() + .name("later".into()) + .source( + vec![Case { + input: 5, + ..Default::default() + }] + .into(), + ) + .definition(definition) + .quiet(true) + .build(); + let summary = run(EvalRequest::recorded(opts, client.clone())) + .await + .unwrap(); + assert_eq!(summary.results[0].output, Some(5)); + assert_eq!(summary.recording.unwrap().experiment_id, "experiment"); + client + .shutdown(std::time::Duration::from_secs(5)) + .await + .unwrap(); } From ba1d0dbce8db0a6a05e1211248d472d821c5a200 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 02:19:59 +0800 Subject: [PATCH 2/6] Adapt versioned dataset reads into evaluation case sources --- src/api/btql.rs | 3 + src/dataset.rs | 5 + src/dataset/read_session.rs | 143 ++++++++++++ src/eval/executor.rs | 9 +- src/eval/mod.rs | 3 +- src/eval/source.rs | 96 ++++++++ src/internal/upload.rs | 6 +- tests/api_client.rs | 4 +- tests/dataset_case_source.rs | 417 +++++++++++++++++++++++++++++++++++ tests/eager_api.rs | 1 + 10 files changed, 683 insertions(+), 4 deletions(-) create mode 100644 src/dataset/read_session.rs create mode 100644 tests/dataset_case_source.rs diff --git a/src/api/btql.rs b/src/api/btql.rs index f4075ab..3fe4d45 100644 --- a/src/api/btql.rs +++ b/src/api/btql.rs @@ -33,6 +33,9 @@ pub struct DatasetRecord { #[derive(Debug, Clone, Serialize)] pub struct BTQLQuery { pub query: BTQLQueryInner, + /// Optional immutable dataset snapshot version. + #[serde(skip_serializing_if = "Option::is_none")] + pub version: Option, } /// Inner BTQL query structure. diff --git a/src/dataset.rs b/src/dataset.rs index dc4a6a4..29ece2f 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -595,6 +595,7 @@ struct DatasetCursor { scope: CredentialScope, dataset_id: String, batch_size: usize, + version: Option, cursor: Option, buffer: VecDeque, done: bool, @@ -634,6 +635,7 @@ impl DatasetIterator { scope, dataset_id, batch_size, + version: None, cursor: None, buffer: VecDeque::new(), done: false, @@ -690,6 +692,8 @@ impl DatasetIterator { } } +pub(crate) mod read_session; + impl DatasetCursor { /// Get the next record, fetching a new batch if needed. async fn next(&mut self) -> Result> { @@ -707,6 +711,7 @@ impl DatasetCursor { limit: Some(self.batch_size), cursor: previous_cursor.clone(), }, + version: self.version.clone(), }; // A failed request leaves the cursor and EOF state untouched. let response = self diff --git a/src/dataset/read_session.rs b/src/dataset/read_session.rs new file mode 100644 index 0000000..b6571d0 --- /dev/null +++ b/src/dataset/read_session.rs @@ -0,0 +1,143 @@ +//! One admitted, demand-driven dataset cursor on its resource's fixed binding. +use super::*; +use std::future::IntoFuture; +use tokio::sync::{mpsc, oneshot}; + +#[derive(Debug, Clone)] +pub(crate) struct ReadError(Arc); +impl std::fmt::Display for ReadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} +impl std::error::Error for ReadError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.0.as_ref()) + } +} +impl From for ReadError { + fn from(error: BraintrustError) -> Self { + Self(Arc::new(error)) + } +} +impl From for ReadError { + fn from(error: crate::ExecutionError) -> Self { + BraintrustError::from(error).into() + } +} +impl From for ReadError { + fn from(error: crate::EnqueueError) -> Self { + BraintrustError::from(error).into() + } +} + +type Reply = oneshot::Sender, ReadError>>; +type Driver = futures::future::Shared>>; + +pub(crate) struct ReadSession { + requests: Option>, + ready: Option>, + metadata: Option, + driver: Driver, + done: bool, +} +impl ReadSession { + pub fn start( + dataset: Dataset, + batch_size: Option, + version: Option, + ) -> std::result::Result { + let (requests, mut receiver) = mpsc::channel::(1); + let (ready, metadata) = oneshot::channel(); + let execution = dataset.execution.clone(); + let completion = execution.response(async move { + let metadata = dataset.ensure_registered().await?.clone(); + let mut cursor = DatasetCursor { + fetcher: dataset.api.clone(), + scope: dataset.scope.clone(), + dataset_id: metadata.dataset_id.clone(), + batch_size: batch_size.unwrap_or(1000).max(1), + version, + cursor: None, + buffer: VecDeque::new(), + done: false, + }; + if ready.send(metadata).is_err() { + return Ok(()); + } + while let Some(reply) = receiver.recv().await { + let result = cursor.next().await.map_err(ReadError::from); + let terminal = !matches!(result, Ok(Some(_))); + let failure = result.as_ref().err().cloned(); + if reply.send(result).is_err() { + break; + } + if let Some(error) = failure { + return Err(error); + } + if terminal { + break; + } + } + Ok(()) + }); + if let Err(error) = &*completion { + return Err(error.clone()); + } + Ok(Self { + requests: Some(requests), + ready: Some(metadata), + metadata: None, + driver: completion.into_future().boxed().shared(), + done: false, + }) + } + + pub async fn prepare(&mut self) -> std::result::Result<&DatasetMetadata, ReadError> { + if let Some(ready) = self.ready.take() { + self.metadata = Some(self.receive(ready).await?); + } + self.metadata.as_ref().ok_or_else(stopped) + } + + pub async fn next(&mut self) -> std::result::Result, ReadError> { + self.prepare().await?; + if self.done { + return Ok(None); + } + let (reply, received) = oneshot::channel(); + if self + .requests + .as_ref() + .ok_or_else(stopped)? + .send(reply) + .await + .is_err() + { + return Err(self.driver.clone().await.err().unwrap_or_else(stopped)); + } + let result = self.receive(received).await?; + let result = result?; + if result.is_none() { + self.done = true; + self.requests = None; + } + Ok(result) + } + + async fn receive( + &self, + receiver: oneshot::Receiver, + ) -> std::result::Result { + use futures::future::{select, Either}; + // Prefer a reply already produced before a terminal driver result. + match select(receiver, self.driver.clone()).await { + Either::Left((Ok(value), _)) => Ok(value), + Either::Left((Err(_), driver)) => Err(driver.await.err().unwrap_or_else(stopped)), + Either::Right((result, _)) => Err(result.err().unwrap_or_else(stopped)), + } + } +} +fn stopped() -> ReadError { + BraintrustError::Background("dataset read session stopped before replying".into()).into() +} diff --git a/src/eval/executor.rs b/src/eval/executor.rs index 9161636..e0dab25 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -87,12 +87,19 @@ impl EvalExecutor { binding: Option<&crate::internal::execution::ExecutionBinding>, ) -> std::result::Result, EvalRunError> { reporter.started(&plan.name); + let prepared = plan.source.prepare().await; let dataset = plan.source.provenance(); let mut results = Vec::new(); let mut failures = Vec::new(); + if let Err(source) = prepared { + failures.push(EvalRunFailure::Source { + case_index: 0, + source, + }); + } let mut case_count = 0; let mut pending = FuturesUnordered::new(); - loop { + while failures.is_empty() { // Keep one source read alive while cases run. Sources may await // case progress and need not support cancellation/restarting next(). let next_case = { diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 121f193..8a556b3 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -73,7 +73,8 @@ 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, + CaseSource, DatasetCaseSource, DatasetVersionRef, IterCaseSource, SourceError, + StreamCaseSource, VecCaseSource, }; pub use task::{task, SimpleFnTask, SyncFnTask, Task}; pub use types::{ diff --git a/src/eval/source.rs b/src/eval/source.rs index 025f2ee..2456b2e 100644 --- a/src/eval/source.rs +++ b/src/eval/source.rs @@ -34,6 +34,11 @@ pub struct DatasetVersionRef { /// without requiring Sync, cloning user values, or serializing local cases. #[async_trait] pub trait CaseSource: Send { + /// Resolve source identity before provenance is captured or cases are read. + async fn prepare(&mut self) -> Result<(), SourceError> { + Ok(()) + } + async fn next(&mut self) -> Result>, SourceError>; fn provenance(&self) -> Option { None @@ -103,3 +108,94 @@ impl From>> for Box anyhow::Result<()> { +/// let mut cases = DatasetCaseSource::::new( +/// dataset, Some("snapshot-version".into()), Some(100), +/// )?; +/// cases.prepare().await?; +/// while let Some(case) = cases.next().await? { println!("{}", case.input); } +/// # Ok(()) +/// # } +/// ``` +pub struct DatasetCaseSource { + session: crate::dataset::read_session::ReadSession, + provenance: Option, + version: Option, + marker: std::marker::PhantomData (I, O)>, +} +impl DatasetCaseSource { + pub fn new( + dataset: crate::Dataset, + version: Option, + batch_size: Option, + ) -> Result { + Ok(Self { + session: crate::dataset::read_session::ReadSession::start( + dataset, + batch_size, + version.clone(), + )?, + provenance: None, + version, + marker: std::marker::PhantomData, + }) + } +} +#[async_trait] +impl CaseSource + for DatasetCaseSource +{ + async fn prepare(&mut self) -> Result<(), SourceError> { + let metadata = self.session.prepare().await?; + if self.provenance.is_none() { + self.provenance = Some(DatasetVersionRef { + dataset_id: metadata.dataset_id.clone(), + version: self.version.clone(), + }); + } + Ok(()) + } + fn provenance(&self) -> Option { + self.provenance.clone() + } + async fn next(&mut self) -> Result>, SourceError> { + self.prepare().await?; + let Some(record) = self.session.next().await? else { + return Ok(None); + }; + let input = record.input.unwrap_or(serde_json::Value::Null); + let input = serde_json::from_value(input).map_err(|error| { + anyhow::Error::new(error).context(format!("dataset record {} input", record.id)) + })?; + let expected = record + .expected + .map(serde_json::from_value) + .transpose() + .map_err(|error| { + anyhow::Error::new(error).context(format!("dataset record {} expected", record.id)) + })?; + Ok(Some(Case { + input, + expected, + metadata: record.metadata, + tags: record.tags, + id: Some(record.id), + xact_id: record.xact_id, + created: record.created, + })) + } +} diff --git a/src/internal/upload.rs b/src/internal/upload.rs index 57a6258..d03b63c 100644 --- a/src/internal/upload.rs +++ b/src/internal/upload.rs @@ -313,7 +313,11 @@ mod tests { .any(|key| key.as_str() == "x-bt-org-name")); assert_eq!(request.body, data); } - let grant: Value = serde_json::from_slice(&requests[0].body).unwrap(); + let grant_request = requests + .iter() + .find(|request| request.url.path() == "/logs3/overflow") + .expect("overflow grant request"); + let grant: Value = serde_json::from_slice(&grant_request.body).unwrap(); assert_eq!(grant["size_bytes"], data.len()); assert_eq!(grant["rows"][0]["object_ids"]["project_id"], "project"); } diff --git a/tests/api_client.rs b/tests/api_client.rs index 450d79a..3a908e2 100644 --- a/tests/api_client.rs +++ b/tests/api_client.rs @@ -147,6 +147,7 @@ async fn btql_readback_preserves_cursor_and_request_errors() { limit: Some(5), cursor: Some("next".into()), }, + version: None, }) .await .unwrap(); @@ -169,7 +170,8 @@ async fn btql_readback_preserves_cursor_and_request_errors() { from: "dataset('id')".into(), limit: None, cursor: None, - } + }, + version: None, }) .await, Err(api::Error::Http { status: 500, .. }) diff --git a/tests/dataset_case_source.rs b/tests/dataset_case_source.rs new file mode 100644 index 0000000..bd45263 --- /dev/null +++ b/tests/dataset_case_source.rs @@ -0,0 +1,417 @@ +use braintrust_sdk_rust::{eval::*, BraintrustClient, BraintrustError, Dataset, EnqueueError}; +use serde::Deserialize; +use serde_json::json; +use std::{ + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + time::Duration, +}; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +async fn setup(recording: bool) -> (BraintrustClient, Dataset, MockServer, MockServer) { + let control = MockServer::start().await; + let data = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"org_info":[{"id":"record-org","name":"record-org","api_url":data.uri()}]}), + )) + .mount(&control) + .await; + Mock::given(method("POST")).and(path("/api/dataset/register")) + .and(header("authorization", "Bearer dataset-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"project"},"dataset":{"id":"dataset-id","name":"cases"}}))) + .mount(&control).await; + Mock::given(method("POST")).and(path("/api/experiment/register")) + .and(header("authorization", "Bearer record-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"record-project","name":"record-project"},"experiment":{"id":"experiment","name":"experiment"}}))) + .mount(&control).await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&data) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .mount(&data) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .app_url(control.uri()) + .api_url(data.uri()) + .api_key("record-token") + .skip_login(!recording) + .blocking_login(recording) + .default_project("record-project") + .build() + .await + .unwrap(); + let dataset = client + .dataset_builder_with_credentials("dataset-token", "dataset-org") + .project_name("project") + .dataset_name("cases") + .build() + .unwrap(); + (client, dataset, control, data) +} +fn page(response: ResponseTemplate) -> Mock { + Mock::given(method("POST")) + .and(path("/btql")) + .and(header("authorization", "Bearer dataset-token")) + .respond_with(response) +} +#[derive(Deserialize)] +struct Input { + value: usize, +} +#[derive(Deserialize)] +struct Output { + value: usize, +} +fn options(source: DatasetCaseSource) -> EvalOpts { + EvalOpts::builder() + .name("dataset".into()) + .source(Box::new(source)) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &Input| { + Ok(Output { + value: input.value * 2, + }) + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build() +} + +#[tokio::test] +async fn typed_pages_pin_version_and_preserve_case_fields_without_clone_or_serialize() { + let (client, dataset, control, data) = setup(false).await; + page(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"one","input":{"value":1},"expected":{"value":2},"metadata":{"source":true},"tags":["source"],"_xact_id":"row-version","created":"created"}],"cursor":"a"}))) + .with_priority(1).up_to_n_times(1).expect(1).mount(&data).await; + page(ResponseTemplate::new(200).set_body_json(json!({"data":[],"cursor":"b"}))) + .with_priority(2) + .up_to_n_times(1) + .expect(1) + .mount(&data) + .await; + page( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"two","input":{"value":2}}],"cursor":null})), + ) + .with_priority(3) + .expect(1) + .mount(&data) + .await; + let source = DatasetCaseSource::new(dataset, Some("snapshot".into()), Some(1)).unwrap(); + let summary = run(EvalRequest::local(options(source))).await.unwrap(); + assert_eq!( + summary.dataset.unwrap(), + DatasetVersionRef { + dataset_id: "dataset-id".into(), + version: Some("snapshot".into()) + } + ); + assert!(summary.recording.is_none()); + assert_eq!( + summary + .results + .iter() + .map(|case| case.output.as_ref().unwrap().value) + .collect::>(), + [2, 4] + ); + let case = &summary.results[0]; + assert_eq!(case.expected.as_ref().unwrap().value, 2); + assert_eq!(case.case_id.as_deref(), Some("one")); + assert_eq!(case.xact_id.as_deref(), Some("row-version")); + assert_eq!(case.created.as_deref(), Some("created")); + assert_eq!(case.metadata.as_ref().unwrap()["source"], true); + assert_eq!(case.tags.as_ref().unwrap(), &["source"]); + let queries: Vec = data + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/btql") + .map(|request| serde_json::from_slice(&request.body).unwrap()) + .collect(); + assert_eq!(queries.len(), 3); + for query in &queries { + assert_eq!(query["version"], "snapshot"); + assert_eq!(query["query"]["from"], "dataset('dataset-id')"); + assert_eq!(query["query"]["limit"], 1); + } + assert!(queries[0]["query"].get("cursor").is_none()); + assert_eq!(queries[1]["query"]["cursor"], "a"); + assert_eq!(queries[2]["query"]["cursor"], "b"); + assert!(!control + .received_requests() + .await + .unwrap() + .iter() + .any(|request| request.url.path() == "/btql")); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn recorded_source_admitted_before_shutdown_can_fetch_all_pages_and_finish_recording() { + let (client, dataset, _, data) = setup(true).await; + page( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"one","input":1}],"cursor":"next"})), + ) + .with_priority(1) + .up_to_n_times(1) + .mount(&data) + .await; + page( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"two","input":2}],"cursor":null})), + ) + .with_priority(2) + .mount(&data) + .await; + let source = DatasetCaseSource::::new(dataset, None, Some(1)).unwrap(); + let opts = EvalOpts::builder() + .name("recorded-dataset".into()) + .source(Box::new(source)) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let token = run(EvalRequest::recorded(opts, client.clone())); + assert!(token.is_ok()); + let shutdown = client.shutdown(Duration::from_secs(5)); + let (result, shutdown) = tokio::join!(token, shutdown); + let summary = result.unwrap(); + shutdown.unwrap(); + assert_eq!(summary.total_cases, 2); + assert_eq!(summary.dataset.unwrap().version, None); + assert_eq!(summary.recording.unwrap().experiment_id, "experiment"); + assert!(summary.results.iter().all(|case| case.span_id.is_some())); +} + +#[tokio::test] +async fn registration_failure_is_a_typed_source_error_before_cases_are_consumed() { + let (client, dataset, control, data) = setup(false).await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(403).set_body_string("denied")) + .with_priority(1) + .mount(&control) + .await; + let calls = Arc::new(AtomicUsize::new(0)); + let source = DatasetCaseSource::::new(dataset, None, None).unwrap(); + let opts = EvalOpts::builder() + .name("failure".into()) + .source(Box::new(source)) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new({ + let calls = calls.clone(); + move |input: &usize| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(*input) + } + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let error = run(EvalRequest::local(opts)).await.unwrap_err(); + let [EvalRunFailure::Source { + case_index: 0, + source, + }] = error.failures.as_slice() + else { + panic!("expected source failure") + }; + assert!(source.chain().any(|error| matches!( + error.downcast_ref::(), + Some(BraintrustError::Api { status: 403, .. }) + ))); + assert_eq!(calls.load(Ordering::SeqCst), 0); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 0); + assert!(partial.dataset.is_none()); + assert!(data.received_requests().await.unwrap().is_empty()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn later_page_failure_retains_completed_cases_and_snapshot_identity() { + let (client, dataset, _, data) = setup(false).await; + page( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"one","input":{"value":4}}],"cursor":"next"})), + ) + .with_priority(1) + .up_to_n_times(1) + .mount(&data) + .await; + page(ResponseTemplate::new(502).set_body_string("page failed")) + .with_priority(2) + .mount(&data) + .await; + let error = run(EvalRequest::local(options( + DatasetCaseSource::new(dataset, Some("pinned".into()), Some(1)).unwrap(), + ))) + .await + .err() + .expect("page should fail"); + assert!(matches!( + &error.failures[..], + [EvalRunFailure::Source { case_index: 1, .. }] + )); + let summary = error.partial.unwrap(); + assert_eq!(summary.dataset.unwrap().version.as_deref(), Some("pinned")); + assert_eq!(summary.results[0].output.as_ref().unwrap().value, 8); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn malformed_row_keeps_decode_cause_and_preceding_results() { + let (client, dataset, _, data) = setup(false).await; + page(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"good","input":{"value":2}},{"id":"bad","input":"invalid input"}],"cursor":null}))).mount(&data).await; + let error = run(EvalRequest::local(options( + DatasetCaseSource::new(dataset, None, None).unwrap(), + ))) + .await + .err() + .expect("decode should fail"); + let [EvalRunFailure::Source { + case_index: 1, + source, + }] = error.failures.as_slice() + else { + panic!("expected row failure") + }; + assert!(source + .chain() + .any(|error| error.downcast_ref::().is_some())); + assert!(source.to_string().contains("bad input")); + assert_eq!( + error.partial.unwrap().results[0] + .output + .as_ref() + .unwrap() + .value, + 4 + ); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn dropping_an_unread_source_releases_accepted_work_and_closed_admission_is_immediate() { + let (client, dataset, _, data) = setup(false).await; + let source = DatasetCaseSource::::new(dataset.clone(), None, None).unwrap(); + drop(source); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert!(data.received_requests().await.unwrap().is_empty()); + let error = DatasetCaseSource::::new(dataset, None, None) + .err() + .expect("source admission must be closed"); + assert!(error.chain().any(|error| matches!( + error.downcast_ref::(), + Some(BraintrustError::Admission(EnqueueError::Closed)) + ))); +} + +#[tokio::test] +async fn empty_live_source_reports_identity_without_claiming_a_snapshot() { + let (client, dataset, _, data) = setup(false).await; + page(ResponseTemplate::new(200).set_body_json(json!({"data":[],"cursor":null}))) + .expect(1) + .mount(&data) + .await; + let summary = run(EvalRequest::local(options( + DatasetCaseSource::new(dataset, None, None).unwrap(), + ))) + .await + .unwrap(); + assert_eq!(summary.total_cases, 0); + assert_eq!( + summary.dataset.unwrap(), + DatasetVersionRef { + dataset_id: "dataset-id".into(), + version: None + } + ); + let requests = data.received_requests().await.unwrap(); + let query: serde_json::Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert!(query.get("version").is_none()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn source_runtime_loss_reports_original_worker_loss_without_running_reads_elsewhere() { + let control = MockServer::start().await; + let data = MockServer::start().await; + let app_url = control.uri(); + let api_url = data.uri(); + let source = std::thread::spawn(move || { + use std::future::IntoFuture; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let client = runtime + .block_on( + BraintrustClient::builder() + .runtime(Some(runtime.handle().clone())) + .app_url(app_url) + .api_url(api_url) + .skip_login(true) + .build() + .into_future(), + ) + .unwrap(); + let dataset = client + .dataset_builder_with_credentials("token", "org") + .project_name("project") + .dataset_name("dataset") + .build() + .unwrap(); + let source = DatasetCaseSource::new(dataset, None, None).unwrap(); + drop(runtime); + source + }) + .join() + .unwrap(); + let error = run(EvalRequest::local(options(source))) + .await + .err() + .expect("source runtime was lost"); + let [EvalRunFailure::Source { + case_index: 0, + source, + }] = error.failures.as_slice() + else { + panic!("expected source failure") + }; + assert!(source.chain().any(|error| matches!( + error.downcast_ref::(), + Some(BraintrustError::Execution( + braintrust_sdk_rust::ExecutionError::WorkerLost + )) + ))); + assert!(control.received_requests().await.unwrap().is_empty()); + assert!(data.received_requests().await.unwrap().is_empty()); +} diff --git a/tests/eager_api.rs b/tests/eager_api.rs index 2e15533..b36f52c 100644 --- a/tests/eager_api.rs +++ b/tests/eager_api.rs @@ -107,6 +107,7 @@ async fn all_endpoint_families_run_unobserved_from_a_plain_thread_and_drain_afte limit: None, cursor: None, }, + version: None, })); abandon(client.experiments().register(ExperimentRegisterRequest { project_name: "p".into(), From 4318106bfc1881c94bc9f18a41301133f5dd02e4 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 02:31:11 +0800 Subject: [PATCH 3/6] Associate recorded evaluations with dataset selections and row origins --- src/api/experiments.rs | 6 + src/eval/executor.rs | 21 +++- src/eval/recorder.rs | 82 +++++++++++--- src/eval/run.rs | 7 +- src/eval/source.rs | 7 ++ src/experiments/builder.rs | 14 +++ src/experiments/experiment.rs | 6 + src/span.rs | 19 ++++ src/types.rs | 21 ++++ tests/api_client.rs | 2 + tests/dataset_case_source.rs | 202 ++++++++++++++++++++++++++++++++++ tests/eager_api.rs | 2 + tests/eval_test.rs | 20 +++- 13 files changed, 384 insertions(+), 25 deletions(-) diff --git a/src/api/experiments.rs b/src/api/experiments.rs index 65cf443..148abb5 100644 --- a/src/api/experiments.rs +++ b/src/api/experiments.rs @@ -44,6 +44,12 @@ pub struct ExperimentRegisterRequest { pub description: Option, #[serde(skip_serializing_if = "Option::is_none")] pub base_experiment: Option, + /// Dataset associated with this experiment. + #[serde(skip_serializing_if = "Option::is_none")] + pub dataset_id: Option, + /// Explicit dataset snapshot; omitted for a live selection. + #[serde(skip_serializing_if = "Option::is_none")] + pub dataset_version: Option, #[serde(skip_serializing_if = "Option::is_none")] pub metadata: Option>, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/src/eval/executor.rs b/src/eval/executor.rs index e0dab25..af24c3f 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -54,6 +54,10 @@ impl EvalPlan { #[async_trait] pub(crate) trait EvalRecorder: Send + Sync { + /// Initialize recording only after source identity is frozen, before reading cases. + async fn prepare(&mut self, _dataset: Option<&super::DatasetVersionRef>) -> Result<()> { + Ok(()) + } fn target(&self) -> Option { None } @@ -82,7 +86,7 @@ pub(crate) struct EvalExecutor; impl EvalExecutor { pub async fn run( mut plan: EvalPlan, - recorder: Option<&dyn EvalRecorder>, + mut recorder: Option<&mut dyn EvalRecorder>, reporter: &mut impl EvalReporter, binding: Option<&crate::internal::execution::ExecutionBinding>, ) -> std::result::Result, EvalRunError> { @@ -97,6 +101,17 @@ impl EvalExecutor { source, }); } + if failures.is_empty() { + if let Some(recorder) = recorder.as_deref_mut() { + if let Err(source) = recorder.prepare(dataset.as_ref()).await { + failures.push(EvalRunFailure::Recording { + case_index: None, + source, + }); + } + } + } + let recorder = recorder.as_deref(); let mut case_count = 0; let mut pending = FuturesUnordered::new(); while failures.is_empty() { @@ -446,7 +461,7 @@ mod tests { #[test] fn recorder_start_and_finish_errors_preserve_computation_and_finalize_run() { let finished = Arc::new(AtomicUsize::new(0)); - let recorder = FailingRecorder { + let mut recorder = FailingRecorder { finished: finished.clone(), }; let cases = (0..3) @@ -475,7 +490,7 @@ mod tests { let mut reporter = Observed::default(); let error = futures::executor::block_on(EvalExecutor::run( plan, - Some(&recorder), + Some(&mut recorder), &mut reporter, None, )) diff --git a/src/eval/recorder.rs b/src/eval/recorder.rs index 10bf5bc..d756310 100644 --- a/src/eval/recorder.rs +++ b/src/eval/recorder.rs @@ -1,7 +1,7 @@ //! Braintrust serialization and write completions, outside the run engine. use super::{ executor::{CaseRecording, EvalRecorder}, - Case, EvalRecording, EvalResult, + Case, DatasetVersionRef, EvalRecording, EvalResult, }; use crate::{ BraintrustClient, BraintrustError, Completion, DeliveryReceipt, Result, SpanHandle, SpanLog, @@ -15,23 +15,47 @@ use std::sync::Arc; type Write = Completion; pub(crate) struct BraintrustRecorder { + client: BraintrustClient, + submitter: Arc, + project_name: Option, + experiment_name: String, + metadata: Option>, + ready: Option, +} +struct RecordingState { template: crate::span::SpanBuilder, target: EvalRecording, + dataset: Option, } impl BraintrustRecorder { - pub async fn new( + pub fn new( client: BraintrustClient, submitter: Arc, project_name: Option, experiment_name: String, metadata: Option>, - ) -> Result { - let mut builder = client - .experiment_builder_for_operation(submitter.clone(), project_name) + ) -> Self { + Self { + client, + submitter, + project_name, + experiment_name, + metadata, + ready: None, + } + } + async fn initialize(&mut self, dataset: Option<&DatasetVersionRef>) -> Result<()> { + let mut builder = self + .client + .experiment_builder_for_operation(self.submitter.clone(), self.project_name.clone()) .await? - .experiment_name(experiment_name); - if let Some(metadata) = metadata { - builder = builder.metadata(metadata); + .experiment_name(self.experiment_name.clone()); + if let Some(metadata) = &self.metadata { + builder = builder.metadata(metadata.clone()); + } + if let Some(dataset) = dataset { + builder = + builder.dataset_reference(dataset.dataset_id.clone(), dataset.version.clone()); } let experiment = builder .build() @@ -42,16 +66,22 @@ impl BraintrustRecorder { experiment_id: resolved.experiment_id.clone(), experiment_name: resolved.experiment_name.clone(), }; - let template = client + let template = self + .client .span_builder_for_operation( - submitter, + self.submitter.clone(), Some(crate::types::ParentSpanInfo::Experiment { object_id: target.experiment_id.clone(), }), ) .await? .span_type(crate::SpanType::Eval); - Ok(Self { template, target }) + self.ready = Some(RecordingState { + template, + target, + dataset: dataset.cloned(), + }); + Ok(()) } } @@ -79,12 +109,30 @@ fn serialize( impl EvalRecorder for BraintrustRecorder { + async fn prepare(&mut self, dataset: Option<&DatasetVersionRef>) -> Result<()> { + self.initialize(dataset).await + } async fn start(&self, index: usize, case: &Case) -> Result>> { - let span = self - .template - .clone() - .name(format!("case-{}", index + 1)) - .build(); + let ready = self.ready.as_ref().ok_or_else(|| { + BraintrustError::Background("evaluation recorder is not initialized".into()) + })?; + let mut builder = ready.template.clone().name(format!("case-{}", index + 1)); + // Match the dataset object-reference contract: both row identity and row + // transaction are required. The selected dataset version is run metadata, + // not a substitute for a particular record's transaction. + if let (Some(dataset), Some(id), Some(xact_id)) = ( + &ready.dataset, + case.id.as_ref().filter(|id| !id.is_empty()), + case.xact_id.as_ref().filter(|id| !id.is_empty()), + ) { + builder = builder.dataset_record_origin(crate::types::DatasetRecordOrigin { + object_id: dataset.dataset_id.clone(), + id: id.clone(), + xact_id: xact_id.clone(), + created: case.created.clone(), + }); + } + let span = builder.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) { @@ -109,7 +157,7 @@ impl })) } fn target(&self) -> Option { - Some(self.target.clone()) + self.ready.as_ref().map(|ready| ready.target.clone()) } async fn finish(&self) -> Vec { diff --git a/src/eval/run.rs b/src/eval/run.rs index 694c95b..7d1eb4e 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -117,7 +117,7 @@ fn recorded: Send { /// Resolve source identity before provenance is captured or cases are read. + /// The executor calls this once, before initializing recorded experiments. async fn prepare(&mut self) -> Result<(), SourceError> { Ok(()) } async fn next(&mut self) -> Result>, SourceError>; + /// The executor captures this once after preparation. A recorded run uses + /// it for experiment association and for case origins with nonempty `id` + /// and `xact_id`; ordinary case IDs alone never imply a dataset association. fn provenance(&self) -> Option { None } diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index c4d0d49..baf138e 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -44,6 +44,8 @@ pub struct ExperimentBuilder { experiment_name: Option, description: Option, base_experiment: Option, + dataset_id: Option, + dataset_version: Option, metadata: Option>, public: Option, update: Option, @@ -61,6 +63,8 @@ impl Clone for ExperimentBuilder { experiment_name: self.experiment_name.clone(), description: self.description.clone(), base_experiment: self.base_experiment.clone(), + dataset_id: self.dataset_id.clone(), + dataset_version: self.dataset_version.clone(), metadata: self.metadata.clone(), public: self.public, update: self.update, @@ -86,6 +90,8 @@ impl ExperimentBuilder { experiment_name: None, description: None, base_experiment: None, + dataset_id: None, + dataset_version: None, metadata: None, public: None, update: None, @@ -123,6 +129,12 @@ impl ExperimentBuilder { self } + pub(crate) fn dataset_reference(mut self, id: String, version: Option) -> Self { + self.dataset_id = Some(id); + self.dataset_version = version; + self + } + /// Set the experiment metadata. pub fn metadata(mut self, metadata: Map) -> Self { self.metadata = Some(metadata); @@ -176,6 +188,8 @@ impl ExperimentBuilder { experiment_name: self.experiment_name, description: self.description, base_experiment: self.base_experiment, + dataset_id: self.dataset_id, + dataset_version: self.dataset_version, metadata: self.metadata, public: self.public, update: self.update, diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 4faf5f1..85a8821 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -45,6 +45,8 @@ pub struct Experiment { pub(crate) experiment_name: Option, pub(crate) description: Option, pub(crate) base_experiment: Option, + pub(crate) dataset_id: Option, + pub(crate) dataset_version: Option, pub(crate) metadata: Option>, pub(crate) public: Option, pub(crate) update: Option, @@ -96,6 +98,8 @@ impl Clone for Experiment { experiment_name: self.experiment_name.clone(), description: self.description.clone(), base_experiment: self.base_experiment.clone(), + dataset_id: self.dataset_id.clone(), + dataset_version: self.dataset_version.clone(), update: self.update, repo_info: self.repo_info.clone(), metadata: self.metadata.clone(), @@ -555,6 +559,8 @@ impl Experiment { experiment_name: self.experiment_name.clone(), description: self.description.clone(), base_experiment: self.base_experiment.clone(), + dataset_id: self.dataset_id.clone(), + dataset_version: self.dataset_version.clone(), metadata: self.metadata.clone(), public: self.public, update: self.update, diff --git a/src/span.rs b/src/span.rs index 5ef0b19..4b517bb 100644 --- a/src/span.rs +++ b/src/span.rs @@ -250,6 +250,7 @@ impl SpanLogBuilder { } pub struct SpanBuilder { + dataset_record_origin: Option, submitter: Arc, scope: CredentialScope, project_name: Option, @@ -290,6 +291,7 @@ impl Clone for SpanBuilder { span_id_override: self.span_id_override.clone(), row_id_override: self.row_id_override.clone(), span_origin: self.span_origin.clone(), + dataset_record_origin: self.dataset_record_origin.clone(), span_attributes_extra: self.span_attributes_extra.clone(), } } @@ -378,10 +380,19 @@ impl SpanBuilder { span_id_override: None, row_id_override: None, span_origin: None, + dataset_record_origin: None, span_attributes_extra: HashMap::new(), } } + pub(crate) fn dataset_record_origin( + mut self, + origin: crate::types::DatasetRecordOrigin, + ) -> Self { + self.dataset_record_origin = Some(origin); + self + } + /// Set the span name before its first snapshot is submitted. pub fn name(mut self, name: impl Into) -> Self { self.name = Some(name.into()); @@ -539,6 +550,7 @@ impl SpanBuilder { name: self.name, environment: self.environment, span_origin: self.span_origin, + dataset_record_origin: self.dataset_record_origin, span_attributes_extra: self.span_attributes_extra, propagated_event, ..Default::default() @@ -807,6 +819,7 @@ fn current_span_payload(inner: &mut SpanData) -> SpanPayload { #[derive(Clone, Default)] struct SpanData { + dataset_record_origin: Option, row_id: String, span_id: String, has_flushed: bool, @@ -873,6 +886,12 @@ impl From for SpanPayload { event_data.apply_propagated_event(&propagated_event); } + if let Some(origin) = data.dataset_record_origin { + event_data + .extra + .insert("origin".into(), origin.into_value()); + } + Self { row_id: data.row_id, span_id: data.span_id, diff --git a/src/types.rs b/src/types.rs index d7d3b23..df133d1 100644 --- a/src/types.rs +++ b/src/types.rs @@ -97,6 +97,27 @@ pub(crate) struct SpanAttributes { pub extra: HashMap, } +/// Dataset row provenance, independent of instrumentation span_origin context. +#[derive(Debug, Clone)] +pub(crate) struct DatasetRecordOrigin { + pub object_id: String, + pub id: String, + pub xact_id: String, + pub created: Option, +} +impl DatasetRecordOrigin { + pub fn into_value(self) -> Value { + let mut value = serde_json::json!({ + "object_type": "dataset", "object_id": self.object_id, + "id": self.id, "_xact_id": self.xact_id, + }); + if let Some(created) = self.created { + value["created"] = Value::String(created); + } + value + } +} + #[derive(Debug, Clone, Default)] pub(crate) struct SpanEventData { pub input: Option, diff --git a/tests/api_client.rs b/tests/api_client.rs index 3a908e2..6fcc8a7 100644 --- a/tests/api_client.rs +++ b/tests/api_client.rs @@ -267,6 +267,8 @@ async fn selected_resource_endpoints_keep_their_wire_contracts_and_planes() { experiment_name: Some("e".into()), description: None, base_experiment: None, + dataset_id: None, + dataset_version: None, metadata: None, public: None, update: None, diff --git a/tests/dataset_case_source.rs b/tests/dataset_case_source.rs index bd45263..883fb35 100644 --- a/tests/dataset_case_source.rs +++ b/tests/dataset_case_source.rs @@ -415,3 +415,205 @@ async fn source_runtime_loss_reports_original_worker_loss_without_running_reads_ assert!(control.received_requests().await.unwrap().is_empty()); assert!(data.received_requests().await.unwrap().is_empty()); } + +#[tokio::test] +async fn recorded_dataset_selection_and_row_origins_preserve_distinct_versions() { + for version in [Some("snapshot".to_owned()), None] { + let (client, dataset, control, data) = setup(true).await; + page(ResponseTemplate::new(200).set_body_json(json!({"data":[ + {"id":"one","input":1,"_xact_id":"row-one","created":"timestamp"}, + {"id":"two","input":2,"_xact_id":"row-two"}, + {"id":"three","input":3}, + {"id":"four","input":4,"_xact_id":""}, + {"id":"","input":5,"_xact_id":"row-five"} + ],"cursor":null}))) + .expect(1) + .mount(&data) + .await; + let source = + DatasetCaseSource::::new(dataset, version.clone(), None).unwrap(); + let definition = EvalDefinition::new( + Box::new(task(|input: &usize, hooks| { + Box::pin(async move { + let child = hooks.span.unwrap().child_builder()?.name("child").build(); + child.end().await?; + Ok(*input) + }) + })), + vec![], + ) + .unwrap(); + let opts = EvalOpts::builder() + .name("dataset-origins".into()) + .source(Box::new(source)) + .definition(definition) + .quiet(true) + .build(); + let summary = run(EvalRequest::recorded(opts, client.clone())) + .await + .unwrap(); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert_eq!(summary.dataset.as_ref().unwrap().version, version); + let requests = control.received_requests().await.unwrap(); + let dataset_index = requests + .iter() + .position(|r| r.url.path() == "/api/dataset/register") + .unwrap(); + let experiment_index = requests + .iter() + .position(|r| r.url.path() == "/api/experiment/register") + .unwrap(); + assert!(dataset_index < experiment_index); + let registration: serde_json::Value = + serde_json::from_slice(&requests[experiment_index].body).unwrap(); + assert_eq!(registration["dataset_id"], "dataset-id"); + assert_eq!( + registration.get("dataset_version"), + version.as_ref().map(|v| json!(v)).as_ref() + ); + let requests = data.received_requests().await.unwrap(); + let rows: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .flat_map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + assert!(!rows.is_empty()); + for (index, result) in summary.results.iter().enumerate() { + let case_rows: Vec<_> = rows + .iter() + .filter(|r| r["span_id"] == *result.span_id.as_ref().unwrap()) + .collect(); + assert!(!case_rows.is_empty()); + for row in case_rows { + assert_eq!(row["experiment_id"], "experiment"); + // Instrumentation context and dataset provenance are separate fields. + assert_eq!(row["context"]["span_origin"]["name"], "braintrust.sdk.rust"); + let expected = match index { + 0 => Some( + json!({"object_type":"dataset","object_id":"dataset-id","id":"one","_xact_id":"row-one","created":"timestamp"}), + ), + 1 => Some( + json!({"object_type":"dataset","object_id":"dataset-id","id":"two","_xact_id":"row-two"}), + ), + _ => None, + }; + assert_eq!(row.get("origin"), expected.as_ref()); + } + } + let children: Vec<_> = rows + .iter() + .filter(|r| r["span_attributes"]["name"] == "child") + .collect(); + assert_eq!(children.len(), 5); + assert!(children.iter().all(|r| r.get("origin").is_none())); + } +} + +#[tokio::test] +async fn failed_dataset_preparation_creates_no_experiment_and_returns_empty_partial() { + let (client, dataset, control, data) = setup(true).await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(403).set_body_string("no dataset access")) + .with_priority(1) + .mount(&control) + .await; + let source = DatasetCaseSource::::new(dataset, None, None).unwrap(); + let opts = EvalOpts::builder() + .name("failed-source".into()) + .source(Box::new(source)) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &usize| -> anyhow::Result { + panic!("must not execute") + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let error = run(EvalRequest::recorded(opts, client.clone())) + .await + .unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Source { case_index: 0, .. }] + )); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 0); + assert!(partial.recording.is_none()); + assert!(partial.dataset.is_none()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert!(!control + .received_requests() + .await + .unwrap() + .iter() + .any(|r| r.url.path() == "/api/experiment/register")); + assert!(!data + .received_requests() + .await + .unwrap() + .iter() + .any(|r| matches!(r.url.path(), "/btql" | "/logs3"))); +} + +#[tokio::test] +async fn failed_experiment_registration_retains_prepared_dataset_without_reading_pages() { + let (client, dataset, control, data) = setup(true).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(403).set_body_string("no experiment access")) + .with_priority(1) + .mount(&control) + .await; + let source = + DatasetCaseSource::::new(dataset, Some("snapshot".into()), None).unwrap(); + let opts = EvalOpts::builder() + .name("failed-target".into()) + .source(Box::new(source)) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &usize| -> anyhow::Result { + panic!("must not execute") + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let error = run(EvalRequest::recorded(opts, client.clone())) + .await + .unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Recording { + case_index: None, + .. + }] + )); + let partial = error.partial.unwrap(); + assert_eq!(partial.total_cases, 0); + assert!(partial.recording.is_none()); + assert_eq!( + partial.dataset, + Some(DatasetVersionRef { + dataset_id: "dataset-id".into(), + version: Some("snapshot".into()) + }) + ); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert!(!data + .received_requests() + .await + .unwrap() + .iter() + .any(|r| matches!(r.url.path(), "/btql" | "/logs3"))); +} diff --git a/tests/eager_api.rs b/tests/eager_api.rs index b36f52c..9b66015 100644 --- a/tests/eager_api.rs +++ b/tests/eager_api.rs @@ -115,6 +115,8 @@ async fn all_endpoint_families_run_unobserved_from_a_plain_thread_and_drain_afte experiment_name: None, description: None, base_experiment: None, + dataset_id: None, + dataset_version: None, metadata: None, public: None, update: None, diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 7c9bcc8..29ebdb2 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -711,12 +711,15 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ }; let (client, server) = recording_client().await; let registered = Arc::new(AtomicBool::new(false)); + let prepared = Arc::new(AtomicBool::new(false)); Mock::given(method("POST")).and(path("/api/experiment/register")) .and(header("authorization", "Bearer test-key")) .and(header("x-bt-org-name", "org")) .respond_with({ let registered = registered.clone(); + let prepared = prepared.clone(); move |_: &wiremock::Request| { + assert!(prepared.load(Ordering::SeqCst)); registered.store(true, Ordering::SeqCst); ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"chosen-project","name":"chosen"},"experiment":{"id":"chosen-id","name":"resolved-name"}})) } @@ -724,15 +727,26 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ struct Source { cursor: Cell, registered: Arc, + prepared: Arc, } #[async_trait::async_trait] impl CaseSource for Source { + async fn prepare(&mut self) -> anyhow::Result<()> { + assert!( + !self.prepared.swap(true, Ordering::SeqCst), + "prepare exactly once" + ); + assert!(!self.registered.load(Ordering::SeqCst)); + Ok(()) + } async fn next(&mut self) -> anyhow::Result>> { assert!(self.registered.load(Ordering::SeqCst)); let index = self.cursor.get(); self.cursor.set(index + 1); Ok((index < 2).then(|| Case { input: index, + id: Some(format!("application-case-{index}")), + xact_id: Some("application-version".into()), metadata: Some(json!({"collision":"case"}).as_object().unwrap().clone()), tags: Some(vec!["shared".into(), "case".into()]), ..Default::default() @@ -767,6 +781,7 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ .source(Box::new(Source { cursor: Cell::new(0), registered, + prepared, })) .definition(definition) .project_name("chosen".into()) @@ -808,6 +823,8 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ assert_eq!(registration["project_name"], "chosen"); assert_eq!(registration["experiment_name"], "requested-name"); assert_eq!(registration["metadata"]["collision"], "run"); + assert!(registration.get("dataset_id").is_none()); + assert!(registration.get("dataset_version").is_none()); let rows: Vec = requests .iter() .filter(|request| request.url.path() == "/logs3") @@ -824,6 +841,7 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ assert!(rows .iter() .all(|row| row["span_attributes"]["name"] != "eval-named-eval")); + assert!(rows.iter().all(|row| row.get("origin").is_none())); let mut roots = std::collections::HashSet::new(); for case in &summary.results { let id = case.span_id.as_ref().unwrap(); @@ -903,7 +921,7 @@ async fn registration_failure_reads_no_cases_and_a_later_run_can_record_normally source: BraintrustError::Api { status: 403, .. } }] )); - assert!(error.partial.is_none()); + assert_eq!(error.partial.as_ref().unwrap().total_cases, 0); assert_eq!(reads.load(Ordering::SeqCst), 0); assert_eq!(calls.load(Ordering::SeqCst), 0); assert!(!server From 2943bf2dd98c03d483fc87b1912b895cad4ecb50 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 07:22:53 +0800 Subject: [PATCH 4/6] Separate evaluation reporting and route CLI output through run adapters --- src/eval/bt_runner.rs | 214 ++++++++++++++--- src/eval/error.rs | 9 +- src/eval/executor.rs | 42 ++-- src/eval/mod.rs | 3 + src/eval/reporter.rs | 203 ++++++++++++++-- src/eval/run.rs | 37 ++- src/eval/runner.rs | 2 +- tests/eval_reporting.rs | 516 ++++++++++++++++++++++++++++++++++++++++ tests/eval_test.rs | 2 +- 9 files changed, 936 insertions(+), 92 deletions(-) create mode 100644 tests/eval_reporting.rs diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index 8662a52..b775461 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -16,8 +16,7 @@ use serde_json::{json, Value}; use crate::{BraintrustClient, BraintrustError, Result}; use super::{ - runner::{EvalOpts, Evaluator}, - types::EvalSummary, + runner::EvalOpts, types::EvalSummary, EvalReporter, EvalRequest, JsonlReporter, ReportError, }; /// A filter for selecting which evaluators to run, read from `BT_EVAL_FILTER_PARSED`. @@ -115,12 +114,27 @@ struct SseWriter { } impl SseWriter { - fn send(&self, event: &str, data: &Value) { - let data_str = serde_json::to_string(data).unwrap_or_default(); - if let Ok(mut w) = self.writer.lock() { - let _ = write!(w, "event: {}\ndata: {}\n\n", event, data_str); - let _ = w.flush(); - } + fn send(&self, event: &str, data: &Value) -> std::result::Result<(), ReportError> { + let data_str = serde_json::to_string(data)?; + let mut writer = self + .writer + .lock() + .map_err(|_| std::io::Error::other("SSE writer lock poisoned"))?; + write!(writer, "event: {}\ndata: {}\n\n", event, data_str)?; + writer.flush()?; + Ok(()) + } +} + +struct SseReporter { + sse: Arc, +} +impl EvalReporter for SseReporter { + fn started(&mut self, name: &str) -> std::result::Result<(), ReportError> { + self.sse.send("start", &json!({"name": name})) + } + fn finished(&mut self, summary: &EvalSummary) -> std::result::Result<(), ReportError> { + self.sse.send("summary", &summary_to_value(summary)?) } } @@ -224,14 +238,17 @@ pub struct BtEvalRunner { impl BtEvalRunner { /// Create a new runner, reading configuration from environment variables. /// - /// In non-list mode, this creates a `BraintrustClient` from environment variables - /// (e.g. `BRAINTRUST_API_KEY`). In list mode, no client is created. + /// Recording mode creates a client from environment variables (e.g. + /// `BRAINTRUST_API_KEY`). List and local modes create no client. pub async fn from_env() -> Result { let config = RunnerConfig::from_env(); let sse = create_sse_writer(); - // Only create a client when we'll actually run evals - let client = if !config.list { + Self::from_config(config, sse).await + } + + async fn from_config(config: RunnerConfig, sse: Option>) -> Result { + let client = if !config.list && !config.no_send_logs { Some(Arc::new(BraintrustClient::builder().build().await?)) } else { None @@ -311,30 +328,27 @@ impl BtEvalRunner { opts.project_name = Some(project.to_string()); } - let client = self - .client - .clone() - .ok_or_else(|| BraintrustError::Background("BtEvalRunner has no client".to_string()))?; - - let evaluator = Evaluator::::from_arc(client); - - if let Some(sse) = &self.sse { - sse.send("start", &json!({ "name": name })); - } + let request = if self.config.no_send_logs { + EvalRequest::local(opts) + } else { + let client = self.client.as_ref().ok_or_else(|| { + BraintrustError::Background("BtEvalRunner has no client".to_string()) + })?; + EvalRequest::recorded(opts, client.as_ref().clone()) + }; + // Protocol output replaces terminal progress even when quiet is false. + let request = if let Some(sse) = &self.sse { + request.with_reporter(SseReporter { sse: sse.clone() }) + } else if self.config.jsonl { + request.with_reporter(JsonlReporter::new(std::io::stdout())) + } else { + request + }; - match evaluator.run(opts).await { + match super::run(request).await { Ok(summary) => { let failed = summary.failed_cases; - if let Some(sse) = &self.sse { - let summary_value = serde_json::to_value(&summary).unwrap_or(Value::Null); - sse.send("summary", &summary_value); - } else if self.config.jsonl { - if let Ok(json) = serde_json::to_string(&summary) { - println!("{}", json); - } - } - if failed > 0 { self.all_passed = false; if self.config.terminate_on_failure { @@ -347,8 +361,17 @@ impl BtEvalRunner { } Err(e) => { self.all_passed = false; - if let Some(sse) = &self.sse { - sse.send("error", &json!({ "message": e.to_string() })); + let output_failed = e + .failures + .iter() + .any(|failure| matches!(failure, super::EvalRunFailure::Reporting { .. })); + if let Some(sse) = self.sse.as_ref().filter(|_| !output_failed) { + if let Err(output) = sse.send("error", &json!({ "message": e.to_string() })) { + eprintln!( + "Evaluator '{}' failed: {}; error output also failed: {}", + name, e, output + ); + } } else { eprintln!("Evaluator '{}' failed: {}", name, e); } @@ -394,13 +417,13 @@ impl BtEvalRunner { /// Serialize an [`EvalSummary`] in a way compatible with the SSE summary event. /// -/// This is used internally by [`BtEvalRunner`] and is public only for advanced use cases. -pub fn summary_to_value(summary: &EvalSummary) -> Value +/// Serialization failures are returned with their cause; no null placeholder is emitted. +pub fn summary_to_value(summary: &EvalSummary) -> serde_json::Result where I: Serialize, O: Serialize, { - serde_json::to_value(summary).unwrap_or(Value::Null) + serde_json::to_value(summary) } #[cfg(test)] @@ -469,4 +492,121 @@ mod tests { assert!(filter.matches("foo-eval")); assert!(!filter.matches("bar-eval")); } + #[derive(Clone, Default)] + struct Buffer(Arc>>); + impl Write for Buffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + fn local_config(terminate_on_failure: bool) -> RunnerConfig { + RunnerConfig { + jsonl: true, + list: false, + terminate_on_failure, + filters: vec![], + no_send_logs: true, + } + } + fn local_options(fail: bool) -> EvalOpts { + super::super::EvalOpts::builder() + .name("cli".into()) + .quiet(false) + .source( + vec![super::super::Case { + input: 1, + ..Default::default() + }] + .into(), + ) + .definition( + super::super::EvalDefinition::new( + Box::new(super::super::SimpleFnTask::new(move |input: &usize| { + if fail { + anyhow::bail!("case failed") + } else { + Ok(*input) + } + })), + vec![], + ) + .unwrap(), + ) + .build() + } + #[tokio::test] + async fn sse_protocol_is_the_only_reporter_and_case_failures_update_cli_status() { + let output = Buffer::default(); + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(output.clone())), + }); + let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) + .await + .unwrap(); + assert!(runner.client.is_none()); + runner.eval("project", local_options(false)).await.unwrap(); + runner.eval("project", local_options(true)).await.unwrap(); + assert!(!runner.all_passed()); + assert!(!runner.finish().await.unwrap()); + let bytes = output.0.lock().unwrap(); + let text = std::str::from_utf8(&bytes).unwrap(); + let events: Vec<_> = text.split("\n\n").filter(|s| !s.is_empty()).collect(); + assert_eq!(events.len(), 4); + assert!(events[0].starts_with("event: start\ndata: ")); + assert!(events[1].starts_with("event: summary\ndata: ")); + assert!(events[2].starts_with("event: start\ndata: ")); + let summary: Value = + serde_json::from_str(events[3].strip_prefix("event: summary\ndata: ").unwrap()) + .unwrap(); + assert_eq!(summary["failed_cases"], 1); + } + #[tokio::test] + async fn terminate_on_failure_returns_an_error_and_retains_failed_session_status() { + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(Buffer::default())), + }); + let mut runner = BtEvalRunner::from_config(local_config(true), Some(sse)) + .await + .unwrap(); + assert!(runner.eval("project", local_options(true)).await.is_err()); + assert!(!runner.all_passed()); + assert!(!runner.finish().await.unwrap()); + } + #[tokio::test] + async fn sse_write_failure_is_a_report_failure_without_losing_computation() { + struct Broken; + impl Write for Broken { + fn write(&mut self, _: &[u8]) -> std::io::Result { + Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "closed socket", + )) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(Broken)), + }); + let result = super::super::run( + EvalRequest::local(local_options(false)) + .with_reporter(SseReporter { sse: sse.clone() }), + ) + .await; + let error = result.unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [super::super::EvalRunFailure::Reporting { + stage: super::super::ReportStage::Started, + source: ReportError::Io(_), + .. + }] + )); + assert_eq!(error.partial.unwrap().successful_cases, 1); + } } diff --git a/src/eval/error.rs b/src/eval/error.rs index 37b11dc..694d9df 100644 --- a/src/eval/error.rs +++ b/src/eval/error.rs @@ -71,6 +71,13 @@ pub enum EvalRunFailure { #[source] source: crate::BraintrustError, }, + #[error("reporting failed at {stage:?} for case {case_index:?}: {source}")] + Reporting { + stage: super::ReportStage, + case_index: Option, + #[source] + source: super::ReportError, + }, #[error(transparent)] Execution(#[from] crate::ExecutionError), #[error(transparent)] @@ -79,7 +86,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. +/// Recording and reporting failures preserve computed results and their success counts. pub struct EvalRunError { pub failures: Vec, pub partial: Option>>, diff --git a/src/eval/executor.rs b/src/eval/executor.rs index af24c3f..40686e7 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -9,6 +9,7 @@ use super::{ case_runner::CaseRunner, Case, EvalDefinition, EvalOpts, EvalResult, EvalRunError, EvalRunFailure, EvalSummary, }; +use super::{reporter::ReporterSession, EvalReporter}; use crate::{BraintrustError, Result, SpanHandle}; /// An invocation owns its source and options; only the definition is shared. @@ -71,12 +72,6 @@ pub(crate) trait CaseRecording: Send { 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, @@ -87,9 +82,10 @@ impl EvalExecutor { pub async fn run( mut plan: EvalPlan, mut recorder: Option<&mut dyn EvalRecorder>, - reporter: &mut impl EvalReporter, + reporter: &mut dyn EvalReporter, binding: Option<&crate::internal::execution::ExecutionBinding>, ) -> std::result::Result, EvalRunError> { + let mut reporter = ReporterSession::new(reporter); reporter.started(&plan.name); let prepared = plan.source.prepare().await; let dataset = plan.source.provenance(); @@ -126,7 +122,13 @@ impl EvalExecutor { 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); + Self::collect( + index, + outcome, + &mut results, + &mut failures, + &mut reporter, + ); if !failures.is_empty() { break Ok(None); } @@ -158,7 +160,7 @@ impl EvalExecutor { }); if pending.len() >= plan.concurrency { if let Some((index, outcome)) = pending.next().await { - Self::collect(index, outcome, &mut results, &mut failures, reporter); + Self::collect(index, outcome, &mut results, &mut failures, &mut reporter); } if !failures.is_empty() { break; @@ -167,7 +169,7 @@ impl EvalExecutor { } // 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); + Self::collect(index, outcome, &mut results, &mut failures, &mut reporter); } if let Some(recorder) = recorder { failures.extend(recorder.finish().await.into_iter().map(|source| { @@ -194,7 +196,9 @@ impl EvalExecutor { dataset, recording: recorder.and_then(|recorder| recorder.target()), }; - reporter.finished(&summary); + if let Some(failure) = reporter.finished(&summary) { + failures.push(failure); + } if failures.is_empty() { Ok(summary) } else { @@ -237,7 +241,7 @@ impl EvalExecutor { outcome: CaseOutcome, results: &mut Vec<(usize, EvalResult)>, failures: &mut Vec, - reporter: &mut impl EvalReporter, + reporter: &mut ReporterSession<'_, I, O>, ) { failures.extend( outcome @@ -284,12 +288,20 @@ mod tests { finished: bool, } impl EvalReporter for Observed { - fn started(&mut self, _: &str) {} - fn case_finished(&mut self, index: usize, _: &EvalResult) { + fn case_finished( + &mut self, + index: usize, + _: &EvalResult, + ) -> std::result::Result<(), super::super::ReportError> { self.cases.push(index); + Ok(()) } - fn finished(&mut self, _: &EvalSummary) { + fn finished( + &mut self, + _: &EvalSummary, + ) -> std::result::Result<(), super::super::ReportError> { self.finished = true; + Ok(()) } } struct WaitingSource { diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 8a556b3..2306503 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -69,6 +69,9 @@ pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; +pub use reporter::{ + EvalReporter, JsonlReporter, ReportError, ReportStage, SilentReporter, TerminalReporter, +}; pub use run::{run, EvalRequest}; pub use runner::{EvalOpts, Evaluator}; pub use scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; diff --git a/src/eval/reporter.rs b/src/eval/reporter.rs index 149fce5..aadeb4f 100644 --- a/src/eval/reporter.rs +++ b/src/eval/reporter.rs @@ -1,44 +1,197 @@ -//! Console presentation, separate from execution and recording. -use super::{executor::EvalReporter, EvalResult, EvalSummary}; +//! Run-owned output adapters. Reporting never owns computation or recording. +use std::io::{self, Write}; -pub(crate) struct TerminalReporter { - pub quiet: bool, +use serde::Serialize; + +use super::{EvalResult, EvalRunFailure, EvalSummary}; + +/// A failure to present evaluation progress or results. +#[derive(Debug, thiserror::Error)] +pub enum ReportError { + #[error("output failed: {0}")] + Io(#[from] io::Error), + #[error("result serialization failed: {0}")] + Serialization(#[from] serde_json::Error), + #[error(transparent)] + Other(#[from] anyhow::Error), } -impl EvalReporter for TerminalReporter { - fn started(&mut self, name: &str) { - if !self.quiet { - println!("Running evaluation: {name}"); - } + +/// Where a reporter failed. Case indices in run failures are zero-based. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReportStage { + Started, + CaseFinished, + Finished, +} + +/// Per-run presentation, without serialization bounds on local computation. +/// +/// Callbacks run sequentially on the run's execution binding and must return +/// promptly. Case notifications follow completion order and carry source indices; +/// the final summary is in source order. `finished` also receives partial summaries +/// after source or recording failures. +/// +/// The first returned error disables this reporter for the rest of the run. +/// Computation and recording continue; the completion reports a `Reporting` +/// failure with the computed summary in `EvalRunError::partial`. Reporter panics +/// follow the run's ordinary worker-loss behavior. +pub trait EvalReporter: Send { + fn started(&mut self, _name: &str) -> Result<(), ReportError> { + Ok(()) + } + fn case_finished( + &mut self, + _index: usize, + _result: &EvalResult, + ) -> Result<(), ReportError> { + Ok(()) + } + fn finished(&mut self, _summary: &EvalSummary) -> Result<(), ReportError> { + Ok(()) + } +} + +/// Suppress all progress and result output. +pub struct SilentReporter; +impl EvalReporter for SilentReporter {} + +/// Human-readable progress on stdout and case errors on stderr. +pub struct TerminalReporter { + output: Box, + errors: Box, +} +impl Default for TerminalReporter { + fn default() -> Self { + Self::with_writers(io::stdout(), io::stderr()) } - fn case_finished(&mut self, index: usize, result: &EvalResult) { - if self.quiet { - return; +} +impl TerminalReporter { + pub fn with_writers( + output: impl Write + Send + 'static, + errors: impl Write + Send + 'static, + ) -> Self { + Self { + output: Box::new(output), + errors: Box::new(errors), } + } +} +impl EvalReporter for TerminalReporter { + fn started(&mut self, name: &str) -> Result<(), ReportError> { + writeln!(self.output, "Running evaluation: {name}")?; + Ok(()) + } + fn case_finished( + &mut self, + index: usize, + result: &EvalResult, + ) -> Result<(), ReportError> { let scores = result .scores .iter() .map(|score| format!("{}={:.3}", score.name, score.score)) .collect::>() .join(", "); - println!("Case {}: {}", index + 1, scores); + writeln!(self.output, "Case {}: {}", index + 1, scores)?; for error in &result.errors { - eprintln!("Case {}: {}", index + 1, error); + writeln!(self.errors, "Case {}: {}", index + 1, error)?; } + Ok(()) } - 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:"); + fn finished(&mut self, summary: &EvalSummary) -> Result<(), ReportError> { + writeln!(self.output, "\nEvaluation results:")?; + writeln!(self.output, "Total cases: {}", summary.total_cases)?; + writeln!(self.output, "Successful: {}", summary.successful_cases)?; + writeln!(self.output, "Failed: {}", summary.failed_cases)?; + writeln!(self.output, "\nScore summary:")?; for stat in &summary.score_stats { - println!( + writeln!( + self.output, " {}: mean={:.3}, min={:.3}, max={:.3}, count={}", stat.name, stat.mean, stat.min, stat.max, stat.count - ); + )?; + } + self.output.flush()?; + self.errors.flush()?; + Ok(()) + } +} + +/// Emit one serialized summary followed by a newline, with no progress text. +/// Serialization is required only when this adapter is selected. Serialization +/// errors are returned before writing a line; I/O failures may leave a partial line. +/// +/// ```no_run +/// use braintrust_sdk_rust::eval::{EvalOpts, EvalRequest, JsonlReporter}; +/// # fn example(options: EvalOpts) { +/// let completion = braintrust_sdk_rust::eval::run( +/// EvalRequest::local(options).with_reporter(JsonlReporter::new(std::io::stdout())), +/// ); +/// # } +/// ``` +pub struct JsonlReporter { + writer: W, +} +impl JsonlReporter { + pub fn new(writer: W) -> Self { + Self { writer } + } +} +impl EvalReporter for JsonlReporter { + fn finished(&mut self, summary: &EvalSummary) -> Result<(), ReportError> { + let line = serde_json::to_vec(summary)?; + self.writer.write_all(&line)?; + self.writer.write_all(b"\n")?; + self.writer.flush()?; + Ok(()) + } +} + +/// Keep reporting failures out of the scheduler's stop-reading condition. +/// Disable after the first error so broken output cannot accumulate errors per case. +pub(crate) struct ReporterSession<'a, I, O> { + reporter: Option<&'a mut dyn EvalReporter>, + failure: Option, +} +impl<'a, I, O> ReporterSession<'a, I, O> { + pub fn new(reporter: &'a mut dyn EvalReporter) -> Self { + Self { + reporter: Some(reporter), + failure: None, + } + } + fn observe( + &mut self, + outcome: Result<(), ReportError>, + stage: ReportStage, + case_index: Option, + ) { + if let Err(source) = outcome { + self.reporter = None; + self.failure = Some(EvalRunFailure::Reporting { + stage, + case_index, + source, + }); + } + } + pub fn started(&mut self, name: &str) { + if let Some(reporter) = self.reporter.as_mut() { + let outcome = reporter.started(name); + self.observe(outcome, ReportStage::Started, None); + } + } + pub fn case_finished(&mut self, index: usize, result: &EvalResult) { + if let Some(reporter) = self.reporter.as_mut() { + let outcome = reporter.case_finished(index, result); + self.observe(outcome, ReportStage::CaseFinished, Some(index)); + } + } + pub fn finished(&mut self, summary: &EvalSummary) -> Option { + if let Some(reporter) = self.reporter.as_mut() { + let outcome = reporter.finished(summary); + self.observe(outcome, ReportStage::Finished, None); } + self.failure.take() } } diff --git a/src/eval/run.rs b/src/eval/run.rs index 7d1eb4e..04ab6b6 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -2,14 +2,16 @@ use super::{ executor::{EvalExecutor, EvalPlan}, recorder::BraintrustRecorder, - reporter::TerminalReporter, - EvalConfigError, EvalOpts, EvalRunError, EvalRunFailure, EvalSummary, + reporter::{SilentReporter, TerminalReporter}, + EvalConfigError, EvalOpts, EvalReporter, EvalRunError, EvalRunFailure, EvalSummary, }; use crate::{internal::execution::ExecutionBinding, BraintrustClient, Completion}; use serde::Serialize; type RunCompletion = Completion, EvalRunError>; -type SubmitRecorded = fn(BraintrustClient, EvalOpts) -> RunCompletion; +type RunReporter = Box>; +type SubmitRecorded = + fn(BraintrustClient, EvalOpts, RunReporter) -> RunCompletion; /// An owned invocation of a reusable definition with its own source and options. /// @@ -17,6 +19,7 @@ type SubmitRecorded = fn(BraintrustClient, EvalOpts) -> RunCompletio /// serialization and inherit their client's execution binding and lifecycle. pub struct EvalRequest { options: EvalOpts, + reporter: Option>, target: Target, } enum Target { @@ -27,10 +30,17 @@ enum Target { }, } impl EvalRequest { + /// Select run-owned output. This overrides `EvalOpts::quiet`, which only + /// controls the default terminal reporter. Reporter failures preserve results. + pub fn with_reporter(mut self, reporter: impl EvalReporter + 'static) -> Self { + self.reporter = Some(Box::new(reporter)); + self + } /// Run without login, HTTP or recording, on `options.runtime` or an owned runtime. pub fn local(options: EvalOpts) -> Self { Self { options, + reporter: None, target: Target::Local, } } @@ -40,6 +50,7 @@ impl pub fn recorded(options: EvalOpts, client: BraintrustClient) -> Self { Self { options, + reporter: None, target: Target::Recorded { client, submit: recorded::, @@ -73,6 +84,13 @@ pub fn run( if let Err(error) = request.options.validate() { return Completion::rejected(error.into()); } + let mut reporter = request.reporter.unwrap_or_else(|| { + if request.options.quiet { + Box::new(SilentReporter) + } else { + Box::new(TerminalReporter::default()) + } + }); match request.target { Target::Local => { let mut options = request.options; @@ -84,13 +102,10 @@ pub fn run( }; let callbacks = binding.clone(); binding.response(async move { - let mut reporter = TerminalReporter { - quiet: options.quiet, - }; EvalExecutor::run( EvalPlan::new(options)?, None, - &mut reporter, + reporter.as_mut(), Some(&callbacks), ) .await @@ -100,7 +115,7 @@ pub fn run( if request.options.runtime.is_some() { return Completion::rejected(EvalConfigError::RecordedRuntime.into()); } - submit(client, request.options) + submit(client, request.options, reporter) } } } @@ -108,15 +123,13 @@ pub fn run( fn recorded( client: BraintrustClient, options: EvalOpts, + mut reporter: RunReporter, ) -> RunCompletion { let binding = client.execution_binding(); let callbacks = binding.clone(); binding.response_try(|| { let operation = client.span_operation()?; Ok(async move { - let mut reporter = TerminalReporter { - quiet: options.quiet, - }; let mut recorder = BraintrustRecorder::new( client, operation.submitter.clone(), @@ -130,7 +143,7 @@ fn recorded { /// Maximum concurrent evaluations (default: 5) #[builder(default = 5)] pub max_concurrency: usize, - /// Suppress progress output (default: false) + /// Suppress default terminal output (default: false). An explicit request reporter takes precedence. #[builder(default = false)] pub quiet: bool, /// Explicit runtime for local runs. Recorded runs inherit their client binding. diff --git a/tests/eval_reporting.rs b/tests/eval_reporting.rs new file mode 100644 index 0000000..bd47379 --- /dev/null +++ b/tests/eval_reporting.rs @@ -0,0 +1,516 @@ +use braintrust_sdk_rust::{eval::*, BraintrustClient}; +use serde::{Serialize, Serializer}; +use serde_json::json; +use std::{ + cell::Cell, + io::{self, Write}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; + +#[derive(Clone, Default)] +struct Buffer(Arc>>); +impl Write for Buffer { + fn write(&mut self, bytes: &[u8]) -> io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} +fn options(calls: Arc) -> EvalOpts { + EvalOpts::builder() + .name("reporting".into()) + .quiet(true) + .max_concurrency(1) + .source( + (0..4) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>() + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(move |input: &usize| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(input * 2) + })), + vec![], + ) + .unwrap(), + ) + .build() +} +struct FailedReporter { + stage: ReportStage, + attempts: Arc, +} +impl FailedReporter { + fn report(&self, stage: ReportStage) -> Result<(), ReportError> { + self.attempts.fetch_add(1, Ordering::SeqCst); + if self.stage == stage { + Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed output").into()) + } else { + Ok(()) + } + } +} +impl EvalReporter for FailedReporter { + fn started(&mut self, _: &str) -> Result<(), ReportError> { + self.report(ReportStage::Started) + } + fn case_finished(&mut self, _: usize, _: &EvalResult) -> Result<(), ReportError> { + self.report(ReportStage::CaseFinished) + } + fn finished(&mut self, _: &EvalSummary) -> Result<(), ReportError> { + self.report(ReportStage::Finished) + } +} +#[test] +fn reporting_errors_disable_output_without_stopping_cases_or_losing_results() { + for stage in [ + ReportStage::Started, + ReportStage::CaseFinished, + ReportStage::Finished, + ] { + let calls = Arc::new(AtomicUsize::new(0)); + let attempts = Arc::new(AtomicUsize::new(0)); + let completion = run(EvalRequest::local(options(calls.clone())).with_reporter( + FailedReporter { + stage, + attempts: attempts.clone(), + }, + )); + assert!(completion.is_ok()); + let Err(error) = futures::executor::block_on(async { completion.await }) else { + panic!("reporting must fail"); + }; + let [EvalRunFailure::Reporting { + stage: actual, + case_index, + source: ReportError::Io(source), + }] = error.failures.as_slice() + else { + panic!("unexpected failures") + }; + assert_eq!(*actual, stage); + assert_eq!( + *case_index, + (stage == ReportStage::CaseFinished).then_some(0) + ); + assert_eq!(source.kind(), io::ErrorKind::BrokenPipe); + assert_eq!( + attempts.load(Ordering::SeqCst), + match stage { + ReportStage::Started => 1, + ReportStage::CaseFinished => 2, + ReportStage::Finished => 6, + } + ); + assert_eq!(calls.load(Ordering::SeqCst), 4); + let summary = error.partial.unwrap(); + assert_eq!(summary.successful_cases, 4); + assert_eq!( + summary + .results + .iter() + .map(|r| r.output.unwrap()) + .collect::>(), + [0, 2, 4, 6] + ); + } +} +#[test] +fn custom_local_reporter_needs_no_serialization_clone_debug_or_sync() { + struct Input(usize); + struct Output(usize); + struct Reporter { + count: Cell, + observed: Arc, + } + impl EvalReporter for Reporter { + fn case_finished( + &mut self, + index: usize, + result: &EvalResult, + ) -> Result<(), ReportError> { + assert_eq!(index, 0); + assert_eq!(result.output.as_ref().unwrap().0, 9); + self.count.set(self.count.get() + 1); + Ok(()) + } + fn finished(&mut self, _: &EvalSummary) -> Result<(), ReportError> { + self.observed.store(self.count.get(), Ordering::SeqCst); + Ok(()) + } + } + let observed = Arc::new(AtomicUsize::new(0)); + let opts = EvalOpts::builder() + .name("opaque".into()) + .source( + vec![Case { + input: Input(3), + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &Input| Ok(Output(input.0 * 3)))), + vec![], + ) + .unwrap(), + ) + .build(); + let completion = run(EvalRequest::local(opts).with_reporter(Reporter { + count: Cell::new(0), + observed: observed.clone(), + })); + futures::executor::block_on(async { completion.await }).unwrap(); + assert_eq!(observed.load(Ordering::SeqCst), 1); +} +#[test] +fn jsonl_outputs_one_complete_summary_and_terminal_uses_separate_writers() { + let output = Buffer::default(); + let completion = run(EvalRequest::local(options(Arc::new(AtomicUsize::new(0)))) + .with_reporter(JsonlReporter::new(output.clone()))); + futures::executor::block_on(async { completion.await }).unwrap(); + let bytes = output.0.lock().unwrap(); + assert_eq!(bytes.iter().filter(|&&b| b == b'\n').count(), 1); + let summary: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(summary["total_cases"], 4); + drop(bytes); + let terminal = Buffer::default(); + let errors = Buffer::default(); + let mut opts = options(Arc::new(AtomicUsize::new(0))); + opts.definition = EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &usize| -> anyhow::Result { + anyhow::bail!("task-error") + })), + vec![], + ) + .unwrap(); + let completion = run( + EvalRequest::local(opts).with_reporter(TerminalReporter::with_writers( + terminal.clone(), + errors.clone(), + )), + ); + assert_eq!( + futures::executor::block_on(async { completion.await }) + .unwrap() + .failed_cases, + 4 + ); + assert!(String::from_utf8(terminal.0.lock().unwrap().clone()) + .unwrap() + .contains("Failed: 4")); + assert!(String::from_utf8(errors.0.lock().unwrap().clone()) + .unwrap() + .contains("task-error")); +} +struct Unserializable; +impl Serialize for Unserializable { + fn serialize(&self, _: S) -> Result { + Err(serde::ser::Error::custom("cannot serialize")) + } +} +#[test] +fn serialization_errors_preserve_computed_results_without_writing_fake_json() { + let output = Buffer::default(); + let opts = EvalOpts::builder() + .name("serialize".into()) + .source( + vec![Case { + input: Unserializable, + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|_: &Unserializable| Ok(7usize))), + vec![], + ) + .unwrap(), + ) + .build(); + let completion = + run(EvalRequest::local(opts).with_reporter(JsonlReporter::new(output.clone()))); + let Err(error) = futures::executor::block_on(async { completion.await }) else { + panic!("reporting must fail"); + }; + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Reporting { + stage: ReportStage::Finished, + source: ReportError::Serialization(_), + .. + }] + )); + let partial = error.partial.unwrap(); + assert!(summary_to_value(&partial).is_err()); + assert_eq!(partial.results[0].output, Some(7)); + assert!(output.0.lock().unwrap().is_empty()); +} +#[tokio::test] +async fn reporter_failure_preserves_recorded_results_and_does_not_poison_shutdown() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + for (endpoint, body) in [ + ( + "/api/apikey/login", + json!({"org_info":[{"id":"org","name":"org"}]}), + ), + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}}), + ), + ("/logs3", 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(json!({}))) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .api_key("key") + .app_url(server.uri()) + .api_url(server.uri()) + .default_project("project") + .blocking_login(true) + .build() + .await + .unwrap(); + let calls = Arc::new(AtomicUsize::new(0)); + let completion = run( + EvalRequest::recorded(options(calls.clone()), client.clone()).with_reporter( + FailedReporter { + stage: ReportStage::Started, + attempts: Arc::new(AtomicUsize::new(0)), + }, + ), + ); + let shutdown = client.shutdown(Duration::from_secs(5)); + let (result, shutdown) = tokio::join!(completion, shutdown); + shutdown.unwrap(); + let error = result.unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Reporting { .. }] + )); + let partial = error.partial.unwrap(); + assert_eq!(partial.recording.unwrap().experiment_id, "experiment"); + assert_eq!(partial.successful_cases, 4); + let requests = server.received_requests().await.unwrap(); + let rows: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .flat_map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + for case in partial.results { + assert!(rows + .iter() + .any(|r| r["span_id"] == *case.span_id.as_ref().unwrap() + && r["metrics"]["end"].is_number())); + } + assert_eq!(calls.load(Ordering::SeqCst), 4); +} + +// A subprocess exercises actual environment parsing and stdout without mutating +// the test runner's process-wide environment or racing other tests. +#[test] +fn cli_output_child() { + let Ok(mode) = std::env::var("BT_REPORTING_CHILD") else { + return; + }; + futures::executor::block_on(async { + let mut runner = BtEvalRunner::from_env().await.unwrap(); + assert!(runner.no_send_logs()); + let calls = Arc::new(AtomicUsize::new(0)); + for name in ["excluded", "selected"] { + let mut opts = options(calls.clone()); + opts.name = name.into(); + opts.quiet = false; + // No recording span is provided in local CLI mode. + opts.definition = EvalDefinition::new( + Box::new(task({ + let calls = calls.clone(); + move |input: &usize, hooks| { + assert!(hooks.span.is_none()); + calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { Ok(*input) }) + } + })), + vec![], + ) + .unwrap(); + runner.eval("project", opts).await.unwrap(); + } + assert!(runner.finish().await.unwrap()); + assert_eq!( + calls.load(Ordering::SeqCst), + if mode == "list" { 0 } else { 4 } + ); + }); +} +#[test] +fn cli_local_jsonl_and_list_modes_need_no_credentials_and_keep_stdout_clean() { + for mode in ["jsonl", "no-send", "list"] { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "cli_output_child", "--nocapture"]) + .env("BT_REPORTING_CHILD", mode) + .env("BT_EVAL_LOCAL", if mode == "no-send" { "0" } else { "1" }) + .env( + "BT_EVAL_NO_SEND_LOGS", + if mode == "no-send" { "1" } else { "0" }, + ) + .env("BT_EVAL_LIST", if mode == "list" { "1" } else { "0" }) + .env("BT_EVAL_JSONL", "1") + .env("BT_EVAL_TERMINATE_ON_FAILURE", "0") + .env( + "BT_EVAL_FILTER_PARSED", + r#"[{"path":["name"],"pattern":"^selected$"}]"#, + ) + .env_remove("BT_EVAL_SSE_SOCK") + .env_remove("BT_EVAL_SSE_ADDR") + .env_remove("BRAINTRUST_API_KEY") + .env("BRAINTRUST_APP_URL", "http://127.0.0.1:9") + .output() + .unwrap(); + let stdout = String::from_utf8(output.stdout).unwrap(); + assert!( + output.status.success(), + "{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(!stdout.contains("Running evaluation:")); + assert!(!stdout.contains("Case 1:")); + if mode != "list" { + let lines: Vec<_> = stdout + .lines() + .filter(|line| line.starts_with('{')) + .collect(); + assert_eq!(lines.len(), 1, "{stdout}"); + let summary: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(summary["name"], "selected"); + assert_eq!(summary["total_cases"], 4); + assert!(summary.get("recording").is_none()); + } else { + assert!(stdout.lines().any(|line| line == "selected"), "{stdout}"); + assert!(!stdout.contains("excluded")); + assert!(!stdout.contains('{')); + } + } +} + +#[test] +fn source_and_reporting_failures_both_survive_in_partial_result() { + let mut opts = options(Arc::new(AtomicUsize::new(0))); + opts.source = Box::new(IterCaseSource::new( + vec![ + Ok(Case { + input: 3, + ..Default::default() + }), + Err(anyhow::anyhow!("source failed")), + ] + .into_iter(), + )); + let completion = run(EvalRequest::local(opts).with_reporter(FailedReporter { + stage: ReportStage::Finished, + attempts: Arc::new(AtomicUsize::new(0)), + })); + let error = futures::executor::block_on(async { completion.await }).unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [ + EvalRunFailure::Source { case_index: 1, .. }, + EvalRunFailure::Reporting { + stage: ReportStage::Finished, + .. + } + ] + )); + assert_eq!(error.partial.unwrap().results[0].output, Some(6)); +} +#[test] +fn dropping_the_observer_does_not_cancel_run_owned_reporting() { + struct Notice(Option>); + impl EvalReporter for Notice { + fn finished(&mut self, summary: &EvalSummary) -> Result<(), ReportError> { + self.0.take().unwrap().send(summary.total_cases).unwrap(); + Ok(()) + } + } + let (sender, receiver) = std::sync::mpsc::channel(); + let token = run(EvalRequest::local(options(Arc::new(AtomicUsize::new(0)))) + .with_reporter(Notice(Some(sender)))); + assert!(token.is_ok()); + drop(token); + assert_eq!(receiver.recv_timeout(Duration::from_secs(5)).unwrap(), 4); +} +#[test] +fn jsonl_write_and_flush_failures_remain_io_errors_with_computed_results() { + struct Broken { + fail_flush: bool, + } + impl Write for Broken { + fn write(&mut self, bytes: &[u8]) -> io::Result { + if self.fail_flush { + Ok(bytes.len()) + } else { + Err(io::ErrorKind::BrokenPipe.into()) + } + } + fn flush(&mut self) -> io::Result<()> { + Err(io::ErrorKind::BrokenPipe.into()) + } + } + for fail_flush in [false, true] { + let token = run(EvalRequest::local(options(Arc::new(AtomicUsize::new(0)))) + .with_reporter(JsonlReporter::new(Broken { fail_flush }))); + let error = futures::executor::block_on(async { token.await }).unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Reporting { + stage: ReportStage::Finished, + source: ReportError::Io(_), + .. + }] + )); + assert_eq!(error.partial.unwrap().successful_cases, 4); + } +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 29ebdb2..765e937 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -859,7 +859,7 @@ async fn experiment_is_resolved_before_source_reads_and_case_roots_keep_independ } assert_eq!(roots.len(), 2); assert_eq!( - summary_to_value(&summary)["recording"]["experiment_id"], + summary_to_value(&summary).unwrap()["recording"]["experiment_id"], "chosen-id" ); } From 481cd22e0f55cb243ed8cfee8bb9cea608eef852 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 15:32:56 +0800 Subject: [PATCH 5/6] Make CLI evaluations eager and join accepted session work --- examples/eval_async.rs | 24 +- src/eval/bt_runner.rs | 448 +++++++++-------------- src/eval/bt_runner/tests.rs | 643 +++++++++++++++++++++++++++++++++ src/eval/cli_session.rs | 195 ++++++++++ src/eval/mod.rs | 3 + src/eval/run.rs | 48 +-- src/eval/submission.rs | 129 +++++++ src/internal/execution.rs | 44 ++- src/internal/response_tasks.rs | 4 +- 9 files changed, 1201 insertions(+), 337 deletions(-) create mode 100644 src/eval/bt_runner/tests.rs create mode 100644 src/eval/cli_session.rs create mode 100644 src/eval/submission.rs diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 9b1f4cc..5e0c833 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -11,7 +11,7 @@ use braintrust_sdk_rust::eval::*; #[tokio::main] async fn main() -> anyhow::Result<()> { // BtEvalRunner reads BT_EVAL_* env vars automatically. - // In non-list mode it creates a BraintrustClient from the environment. + // Recording mode creates a client; local and list modes need no credentials. let mut runner = BtEvalRunner::from_env().await?; // Define test cases @@ -60,17 +60,17 @@ async fn main() -> anyhow::Result<()> { ], )?; - runner - .eval( - "my-project", - EvalOpts::builder() - .name("uppercase-eval".to_string()) - .source(cases.into()) - .definition(definition.clone()) - .max_concurrency(2) - .build(), - ) - .await?; + let completion = runner.eval( + "my-project", + EvalOpts::builder() + .name("uppercase-eval".to_string()) + .source(cases.into()) + .definition(definition.clone()) + .max_concurrency(2) + .build(), + ); + assert!(completion.is_ok()); // Submission already happened. + completion.await?; let passed = runner.finish().await?; if !passed { diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index b775461..3972051 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -13,10 +13,16 @@ use std::sync::{Arc, Mutex}; use serde::Serialize; use serde_json::{json, Value}; -use crate::{BraintrustClient, BraintrustError, Result}; +use crate::{ + internal::execution::ExecutionBinding, BraintrustClient, BraintrustError, Completion, Result, +}; use super::{ - runner::EvalOpts, types::EvalSummary, EvalReporter, EvalRequest, JsonlReporter, ReportError, + cli_session::{ready, Session}, + runner::EvalOpts, + submission::{self, RunReporter}, + types::EvalSummary, + BtEvalError, BtFinishError, EvalReporter, JsonlReporter, ReportError, }; /// A filter for selecting which evaluators to run, read from `BT_EVAL_FILTER_PARSED`. @@ -208,7 +214,7 @@ fn create_sse_writer() -> Option> { /// vec![Box::new(ExactMatch::new())], /// )?; /// -/// runner.eval( +/// let completion = runner.eval( /// "my-project", /// EvalOpts::builder() /// .name("my-eval".to_string()) @@ -219,7 +225,9 @@ fn create_sse_writer() -> Option> { /// }].into()) /// .definition(definition.clone()) /// .build(), -/// ).await?; +/// ); +/// assert!(completion.is_ok()); // accepted before awaiting +/// completion.await?; /// /// let _passed = runner.finish().await?; /// Ok(()) @@ -231,8 +239,8 @@ pub struct BtEvalRunner { client: Option>, /// Names registered in list mode, in order of registration. registered_names: Vec, - /// Whether all evals that have run so far passed (no failures). - all_passed: bool, + binding: Option, + session: Session, } impl BtEvalRunner { @@ -254,28 +262,41 @@ impl BtEvalRunner { None }; + let binding = if let Some(client) = &client { + Some(client.execution_binding()) + } else if !config.list { + Some(ExecutionBinding::new(None)?) + } else { + None + }; + let session = Session::new(config.terminate_on_failure); Ok(Self { config, sse, client, + binding, + session, registered_names: Vec::new(), - all_passed: true, }) } /// Create a new runner with an explicit client. /// /// Environment variables are still read for mode configuration (list, JSONL, SSE, etc.), - /// but the provided client is used instead of creating a new one. + /// but the provided client is used instead of creating a new one. All runs, + /// including local mode, inherit its fixed execution binding. pub fn with_client(client: Arc) -> Self { let config = RunnerConfig::from_env(); let sse = create_sse_writer(); + let binding = Some(client.execution_binding()); + let session = Session::new(config.terminate_on_failure); Self { config, sse, client: Some(client), + binding, + session, registered_names: Vec::new(), - all_passed: true, } } @@ -289,9 +310,10 @@ impl BtEvalRunner { self.config.no_send_logs } - /// Returns `true` if all evals that have run so far passed. + /// Returns true if no completed evaluation or admission has failed. + /// Pending evaluations are included once `finish()` joins the session. pub fn all_passed(&self) -> bool { - self.all_passed + self.session.all_passed() } fn passes_filter(&self, name: &str) -> bool { @@ -301,117 +323,163 @@ impl BtEvalRunner { self.config.filters.iter().any(|f| f.matches(name)) } - /// Register and/or run an evaluator. - /// - /// - In **list mode**: records the eval name and returns without running. - /// - In **normal mode**: runs the evaluator and sends results via SSE if configured. + /// Submit a CLI evaluation now, in call order on the session/client binding. + /// Dropping its completion does not cancel it. Filtered/list entries complete + /// immediately without evaluation admission. Local runs use an owned session + /// runtime; an explicit client supplies its binding. Omit `opts.runtime`. /// - /// 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<()> + /// With terminate-on-failure enabled, failures are returned with their typed + /// results, and queued/later evaluations are skipped. Otherwise failures mark + /// `all_passed()` false and the session continues, as with the CLI protocol. + /// Admission/configuration errors are still reported synchronously on the token. + pub fn eval( + &mut self, + project: &str, + mut opts: EvalOpts, + ) -> Completion<(), BtEvalError> where I: Serialize + Send + Sync + 'static, O: Serialize + Send + Sync + 'static, { - let name = opts.name.clone(); - - if !self.passes_filter(&name) { - return Ok(()); + if !self.passes_filter(&opts.name) { + return ready(()); } - if self.config.list { - self.registered_names.push(name); - return Ok(()); + self.registered_names.push(opts.name); + return ready(()); + } + if self.session.stopped() { + return Completion::rejected(BtEvalError::Stopped); + } + if opts.runtime.is_some() { + self.session.fail(); + return Completion::rejected(BtEvalError::RuntimeOverride); } - - // Set the project name from the argument if not already set in opts if opts.project_name.is_none() { opts.project_name = Some(project.to_string()); } - - let request = if self.config.no_send_logs { - EvalRequest::local(opts) - } else { - let client = self.client.as_ref().ok_or_else(|| { - BraintrustError::Background("BtEvalRunner has no client".to_string()) - })?; - EvalRequest::recorded(opts, client.as_ref().clone()) - }; - // Protocol output replaces terminal progress even when quiet is false. - let request = if let Some(sse) = &self.sse { - request.with_reporter(SseReporter { sse: sse.clone() }) + let name = opts.name.clone(); + let reporter: RunReporter = if let Some(sse) = &self.sse { + Box::new(SseReporter { sse: sse.clone() }) } else if self.config.jsonl { - request.with_reporter(JsonlReporter::new(std::io::stdout())) + Box::new(JsonlReporter::new(std::io::stdout())) + } else if opts.quiet { + Box::new(super::SilentReporter) } else { - request + Box::new(super::TerminalReporter::default()) }; - - match super::run(request).await { - Ok(summary) => { - let failed = summary.failed_cases; - - if failed > 0 { - self.all_passed = false; - if self.config.terminate_on_failure { - return Err(BraintrustError::Background(format!( - "Evaluator '{}' failed with {} error(s)", - name, failed - ))); - } - } - } - Err(e) => { - self.all_passed = false; - let output_failed = e - .failures - .iter() - .any(|failure| matches!(failure, super::EvalRunFailure::Reporting { .. })); - if let Some(sse) = self.sse.as_ref().filter(|_| !output_failed) { - if let Err(output) = sse.send("error", &json!({ "message": e.to_string() })) { - eprintln!( - "Evaluator '{}' failed: {}; error output also failed: {}", - name, e, output - ); - } + let session = self.session.clone(); + let sse = self.sse.clone(); + let terminate = self.config.terminate_on_failure; + let make = || { + let (previous, ticket) = session.reserve::()?; + let sequence = ticket.sequence(); + let before = async move { + previous.await; + if session.should_skip(sequence) { + Err(BtEvalError::Stopped) } else { - eprintln!("Evaluator '{}' failed: {}", name, e); + Ok(()) } - if self.config.terminate_on_failure { - return Err(BraintrustError::Background(e.to_string())); + }; + let outcome_session = self.session.clone(); + let after = move |outcome: std::result::Result< + EvalSummary, + BtEvalError, + >| { + let failed = !matches!(&outcome, Ok(summary) if summary.failed_cases == 0); + let result = match outcome { + Ok(summary) if summary.failed_cases > 0 && terminate => { + Err(BtEvalError::CaseFailures { + summary: Box::new(summary), + }) + } + Ok(_) => Ok(()), + Err(BtEvalError::Stopped) => Err(BtEvalError::Stopped), + Err(error) => { + let output_failed = matches!(&error, BtEvalError::Run(run) if run.failures.iter().any(|failure| matches!(failure, super::EvalRunFailure::Reporting { .. }))); + if let Some(sse) = sse.as_ref().filter(|_| !output_failed) { + if let Err(output) = + sse.send("error", &json!({"message":error.to_string()})) + { + eprintln!( + "Evaluator '{}' failed: {}; error output also failed: {}", + name, error, output + ); + } + } else { + eprintln!("Evaluator '{}' failed: {}", name, error); + } + if terminate { + Err(error) + } else { + Ok(()) + } + } + }; + if failed { + outcome_session.fail_at(sequence); } - } + result + }; + Ok((before, after, move |published: bool| { + ticket.complete(!published) + })) + }; + let token = if self.config.no_send_logs { + submission::local( + opts, + self.binding + .as_ref() + .expect("non-list session has a binding") + .clone(), + reporter, + make, + ) + } else if let Some(client) = &self.client { + submission::recorded(client.as_ref().clone(), opts, reporter, make) + } else { + Completion::rejected(BtEvalError::Run( + BraintrustError::Background("BtEvalRunner has no client".into()).into(), + )) + }; + if token.is_err() { + self.session.fail(); } - - Ok(()) + token } - /// Finish the runner session. - /// - /// In list mode, prints all registered eval names to stdout (one per line). - /// The SSE connection (if any) is closed when this method returns. - /// - /// Returns `true` if all evals passed, `false` if any failed. - pub async fn finish(self) -> Result { + /// Close the session and join every accepted evaluation, including discarded + /// observers. Output and recording finish before this completion is published. + /// The supplied client remains usable; this does not shut down shared clients. + /// Returns false if any evaluation/admission failed or lost its worker. + pub fn finish(self) -> Completion { if self.config.list { + let mut output = std::io::stdout().lock(); for name in &self.registered_names { - println!("{}", name); + if let Err(error) = writeln!(output, "{}", name) { + return Completion::rejected(BtFinishError::Output(error.into())); + } } + if let Err(error) = output.flush() { + return Completion::rejected(BtFinishError::Output(error.into())); + } + return ready(self.session.all_passed()); } - // SSE connection is closed when `self.sse` is dropped - Ok(self.all_passed) - } - - /// Convenience method: run multiple evaluations whose results can be summarized. - /// - /// Returns `true` if all evals passed. - pub async fn run_all(&mut self, evals: Vec) -> Result - where - F: FnOnce(&mut Self) -> Fut, - Fut: std::future::Future>, - { - for eval_fn in evals { - eval_fn(self).await?; - } - Ok(self.all_passed) + let binding = self + .binding + .as_ref() + .expect("non-list session has a binding") + .clone(); + let (sender, completion) = Completion::pair(); + // One consuming control operation, independent of full run capacity. + binding.spawn(async move { + self.session.tail().await; + let passed = self.session.all_passed(); + drop(self); // Close the CLI socket and release ownership before publication. + let _ = sender.send(Ok(passed)); + }); + completion } } @@ -427,186 +495,4 @@ where } #[cfg(test)] -mod tests { - use std::sync::Mutex; - - use super::*; - - // Serialize tests that mutate BT_EVAL_FILTER_PARSED to avoid race conditions. - static FILTER_ENV_MUTEX: Mutex<()> = Mutex::new(()); - - #[test] - fn test_env_flag() { - // Not set → false - assert!(!env_flag("BT_TEST_NONEXISTENT_FLAG_XYZ")); - } - - #[test] - fn test_parse_filters_empty() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - // When BT_EVAL_FILTER_PARSED is not set, filters should be empty - std::env::remove_var("BT_EVAL_FILTER_PARSED"); - let filters = parse_filters(); - assert!(filters.is_empty()); - } - - #[test] - fn test_parse_filters_valid() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - std::env::set_var( - "BT_EVAL_FILTER_PARSED", - r#"[{"path":["name"],"pattern":"foo.*"}]"#, - ); - let filters = parse_filters(); - assert_eq!(filters.len(), 1); - assert!(filters[0].matches("foobar")); - assert!(!filters[0].matches("bar")); - std::env::remove_var("BT_EVAL_FILTER_PARSED"); - } - - #[test] - fn test_parse_filters_invalid_json() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - std::env::set_var("BT_EVAL_FILTER_PARSED", "not-json"); - let filters = parse_filters(); - assert!(filters.is_empty()); - std::env::remove_var("BT_EVAL_FILTER_PARSED"); - } - - #[test] - fn test_eval_filter_matches_by_name() { - let filter = EvalFilter { - path: vec!["name".to_string()], - pattern: regex::Regex::new("^my-.*").unwrap(), - }; - assert!(filter.matches("my-eval")); - assert!(!filter.matches("your-eval")); - } - - #[test] - fn test_eval_filter_empty_path_matches_name() { - let filter = EvalFilter { - path: vec![], - pattern: regex::Regex::new("foo").unwrap(), - }; - assert!(filter.matches("foo-eval")); - assert!(!filter.matches("bar-eval")); - } - #[derive(Clone, Default)] - struct Buffer(Arc>>); - impl Write for Buffer { - fn write(&mut self, bytes: &[u8]) -> std::io::Result { - self.0.lock().unwrap().extend_from_slice(bytes); - Ok(bytes.len()) - } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - } - fn local_config(terminate_on_failure: bool) -> RunnerConfig { - RunnerConfig { - jsonl: true, - list: false, - terminate_on_failure, - filters: vec![], - no_send_logs: true, - } - } - fn local_options(fail: bool) -> EvalOpts { - super::super::EvalOpts::builder() - .name("cli".into()) - .quiet(false) - .source( - vec![super::super::Case { - input: 1, - ..Default::default() - }] - .into(), - ) - .definition( - super::super::EvalDefinition::new( - Box::new(super::super::SimpleFnTask::new(move |input: &usize| { - if fail { - anyhow::bail!("case failed") - } else { - Ok(*input) - } - })), - vec![], - ) - .unwrap(), - ) - .build() - } - #[tokio::test] - async fn sse_protocol_is_the_only_reporter_and_case_failures_update_cli_status() { - let output = Buffer::default(); - let sse = Arc::new(SseWriter { - writer: Mutex::new(Box::new(output.clone())), - }); - let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) - .await - .unwrap(); - assert!(runner.client.is_none()); - runner.eval("project", local_options(false)).await.unwrap(); - runner.eval("project", local_options(true)).await.unwrap(); - assert!(!runner.all_passed()); - assert!(!runner.finish().await.unwrap()); - let bytes = output.0.lock().unwrap(); - let text = std::str::from_utf8(&bytes).unwrap(); - let events: Vec<_> = text.split("\n\n").filter(|s| !s.is_empty()).collect(); - assert_eq!(events.len(), 4); - assert!(events[0].starts_with("event: start\ndata: ")); - assert!(events[1].starts_with("event: summary\ndata: ")); - assert!(events[2].starts_with("event: start\ndata: ")); - let summary: Value = - serde_json::from_str(events[3].strip_prefix("event: summary\ndata: ").unwrap()) - .unwrap(); - assert_eq!(summary["failed_cases"], 1); - } - #[tokio::test] - async fn terminate_on_failure_returns_an_error_and_retains_failed_session_status() { - let sse = Arc::new(SseWriter { - writer: Mutex::new(Box::new(Buffer::default())), - }); - let mut runner = BtEvalRunner::from_config(local_config(true), Some(sse)) - .await - .unwrap(); - assert!(runner.eval("project", local_options(true)).await.is_err()); - assert!(!runner.all_passed()); - assert!(!runner.finish().await.unwrap()); - } - #[tokio::test] - async fn sse_write_failure_is_a_report_failure_without_losing_computation() { - struct Broken; - impl Write for Broken { - fn write(&mut self, _: &[u8]) -> std::io::Result { - Err(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "closed socket", - )) - } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - } - let sse = Arc::new(SseWriter { - writer: Mutex::new(Box::new(Broken)), - }); - let result = super::super::run( - EvalRequest::local(local_options(false)) - .with_reporter(SseReporter { sse: sse.clone() }), - ) - .await; - let error = result.unwrap_err(); - assert!(matches!( - error.failures.as_slice(), - [super::super::EvalRunFailure::Reporting { - stage: super::super::ReportStage::Started, - source: ReportError::Io(_), - .. - }] - )); - assert_eq!(error.partial.unwrap().successful_cases, 1); - } -} +mod tests; diff --git a/src/eval/bt_runner/tests.rs b/src/eval/bt_runner/tests.rs new file mode 100644 index 0000000..57ebb36 --- /dev/null +++ b/src/eval/bt_runner/tests.rs @@ -0,0 +1,643 @@ +use std::sync::Mutex; + +use super::*; +use crate::eval::EvalRequest; + +// Serialize tests that mutate BT_EVAL_FILTER_PARSED to avoid race conditions. +static FILTER_ENV_MUTEX: Mutex<()> = Mutex::new(()); + +#[test] +fn test_env_flag() { + // Not set → false + assert!(!env_flag("BT_TEST_NONEXISTENT_FLAG_XYZ")); +} + +#[test] +fn test_parse_filters_empty() { + let _guard = FILTER_ENV_MUTEX.lock().unwrap(); + // When BT_EVAL_FILTER_PARSED is not set, filters should be empty + std::env::remove_var("BT_EVAL_FILTER_PARSED"); + let filters = parse_filters(); + assert!(filters.is_empty()); +} + +#[test] +fn test_parse_filters_valid() { + let _guard = FILTER_ENV_MUTEX.lock().unwrap(); + std::env::set_var( + "BT_EVAL_FILTER_PARSED", + r#"[{"path":["name"],"pattern":"foo.*"}]"#, + ); + let filters = parse_filters(); + assert_eq!(filters.len(), 1); + assert!(filters[0].matches("foobar")); + assert!(!filters[0].matches("bar")); + std::env::remove_var("BT_EVAL_FILTER_PARSED"); +} + +#[test] +fn test_parse_filters_invalid_json() { + let _guard = FILTER_ENV_MUTEX.lock().unwrap(); + std::env::set_var("BT_EVAL_FILTER_PARSED", "not-json"); + let filters = parse_filters(); + assert!(filters.is_empty()); + std::env::remove_var("BT_EVAL_FILTER_PARSED"); +} + +#[test] +fn test_eval_filter_matches_by_name() { + let filter = EvalFilter { + path: vec!["name".to_string()], + pattern: regex::Regex::new("^my-.*").unwrap(), + }; + assert!(filter.matches("my-eval")); + assert!(!filter.matches("your-eval")); +} + +#[test] +fn test_eval_filter_empty_path_matches_name() { + let filter = EvalFilter { + path: vec![], + pattern: regex::Regex::new("foo").unwrap(), + }; + assert!(filter.matches("foo-eval")); + assert!(!filter.matches("bar-eval")); +} +#[derive(Clone, Default)] +struct Buffer(Arc>>); +impl Write for Buffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} +fn local_config(terminate_on_failure: bool) -> RunnerConfig { + RunnerConfig { + jsonl: true, + list: false, + terminate_on_failure, + filters: vec![], + no_send_logs: true, + } +} +fn local_options(fail: bool) -> EvalOpts { + super::super::EvalOpts::builder() + .name("cli".into()) + .quiet(false) + .source( + vec![super::super::Case { + input: 1, + ..Default::default() + }] + .into(), + ) + .definition( + super::super::EvalDefinition::new( + Box::new(super::super::SimpleFnTask::new(move |input: &usize| { + if fail { + anyhow::bail!("case failed") + } else { + Ok(*input) + } + })), + vec![], + ) + .unwrap(), + ) + .build() +} +#[tokio::test] +async fn sse_protocol_is_the_only_reporter_and_case_failures_update_cli_status() { + let output = Buffer::default(); + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(output.clone())), + }); + let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) + .await + .unwrap(); + assert!(runner.client.is_none()); + runner.eval("project", local_options(false)).await.unwrap(); + runner.eval("project", local_options(true)).await.unwrap(); + assert!(!runner.all_passed()); + assert!(!runner.finish().await.unwrap()); + let bytes = output.0.lock().unwrap(); + let text = std::str::from_utf8(&bytes).unwrap(); + let events: Vec<_> = text.split("\n\n").filter(|s| !s.is_empty()).collect(); + assert_eq!(events.len(), 4); + assert!(events[0].starts_with("event: start\ndata: ")); + assert!(events[1].starts_with("event: summary\ndata: ")); + assert!(events[2].starts_with("event: start\ndata: ")); + let summary: Value = + serde_json::from_str(events[3].strip_prefix("event: summary\ndata: ").unwrap()).unwrap(); + assert_eq!(summary["failed_cases"], 1); +} +#[tokio::test] +async fn terminate_on_failure_returns_an_error_and_retains_failed_session_status() { + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(Buffer::default())), + }); + let mut runner = BtEvalRunner::from_config(local_config(true), Some(sse)) + .await + .unwrap(); + assert!(runner.eval("project", local_options(true)).await.is_err()); + assert!(!runner.all_passed()); + assert!(!runner.finish().await.unwrap()); +} +#[tokio::test] +async fn sse_write_failure_is_a_report_failure_without_losing_computation() { + struct Broken; + impl Write for Broken { + fn write(&mut self, _: &[u8]) -> std::io::Result { + Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "closed socket", + )) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(Broken)), + }); + let result = super::super::run( + EvalRequest::local(local_options(false)).with_reporter(SseReporter { sse: sse.clone() }), + ) + .await; + let error = result.unwrap_err(); + assert!(matches!( + error.failures.as_slice(), + [super::super::EvalRunFailure::Reporting { + stage: super::super::ReportStage::Started, + source: ReportError::Io(_), + .. + }] + )); + assert_eq!(error.partial.unwrap().successful_cases, 1); +} +fn quiet_options(input: usize) -> EvalOpts { + let mut opts = local_options(false); + opts.quiet = true; + opts.source = vec![super::super::Case { + input, + ..Default::default() + }] + .into(); + opts +} +fn runner_on(binding: ExecutionBinding, terminate: bool) -> BtEvalRunner { + let mut config = local_config(terminate); + config.jsonl = false; + BtEvalRunner { + config, + sse: None, + client: None, + binding: Some(binding), + session: Session::new(terminate), + registered_names: Vec::new(), + } +} +#[tokio::test] +async fn eager_runs_keep_call_order_and_finish_joins_discarded_observers() { + use crate::eval::{task, EvalDefinition}; + let mut runner = runner_on(ExecutionBinding::new(None).unwrap(), false); + let order = Arc::new(Mutex::new(Vec::new())); + let entered = Arc::new(tokio::sync::Notify::new()); + let (release, wait) = futures::channel::oneshot::channel(); + let wait = Mutex::new(Some(wait)); + let mut first = quiet_options(1); + first.definition = EvalDefinition::new( + Box::new(task({ + let entered = entered.clone(); + let order = order.clone(); + move |input: &usize, _| { + let wait = wait.lock().unwrap().take().unwrap(); + let entered = entered.clone(); + let order = order.clone(); + Box::pin(async move { + entered.notify_one(); + wait.await?; + order.lock().unwrap().push(*input); + Ok(*input) + }) + } + })), + vec![], + ) + .unwrap(); + let mut second = quiet_options(2); + second.definition = EvalDefinition::new( + Box::new(super::super::SimpleFnTask::new({ + let order = order.clone(); + move |input: &usize| { + order.lock().unwrap().push(*input); + Ok(*input) + } + })), + vec![], + ) + .unwrap(); + let first = runner.eval("project", first); + let second = runner.eval("project", second); + assert!(first.is_ok()); + assert!(second.is_ok()); + drop(first); + drop(second); + let finished = runner.finish(); + assert!(finished.is_ok()); + entered.notified().await; + assert!(order.lock().unwrap().is_empty()); + release.send(()).unwrap(); + assert!(finished.await.unwrap()); + assert_eq!(*order.lock().unwrap(), [1, 2]); +} +#[tokio::test] +async fn terminate_on_failure_skips_already_admitted_successors_without_reading_sources() { + use crate::eval::{task, EvalDefinition, IterCaseSource}; + use std::sync::atomic::{AtomicUsize, Ordering}; + let mut runner = runner_on( + ExecutionBinding::external(tokio::runtime::Handle::current()), + true, + ); + let (release, wait) = futures::channel::oneshot::channel(); + let wait = Mutex::new(Some(wait)); + let mut first = quiet_options(1); + first.definition = EvalDefinition::new( + Box::new(task(move |_: &usize, _| { + let wait = wait.lock().unwrap().take().unwrap(); + Box::pin(async move { + wait.await?; + anyhow::bail!("case failed") + }) + })), + vec![], + ) + .unwrap(); + let reads = Arc::new(AtomicUsize::new(0)); + let mut second = quiet_options(2); + second.source = Box::new(IterCaseSource::new({ + let reads = reads.clone(); + (0..1).map(move |input| { + reads.fetch_add(1, Ordering::SeqCst); + Ok(super::super::Case { + input, + ..Default::default() + }) + }) + })); + let first = runner.eval("project", first); + let second = runner.eval("project", second); + assert!(first.is_ok()); + assert!(second.is_ok()); + release.send(()).unwrap(); + let BtEvalError::CaseFailures { summary } = first.await.unwrap_err() else { + panic!("case failure must retain summary") + }; + assert_eq!(summary.failed_cases, 1); + assert!(matches!(second.await, Err(BtEvalError::Stopped))); + assert_eq!(reads.load(Ordering::SeqCst), 0); + let later = runner.eval("project", quiet_options(3)); + assert!(later.is_err()); + assert!(matches!(later.await, Err(BtEvalError::Stopped))); + assert!(!runner.finish().await.unwrap()); +} +#[tokio::test] +async fn full_cli_admission_uses_one_slot_per_run_and_finish_has_control_capacity() { + use crate::{ + eval::{task, EvalDefinition}, + internal::response_tasks::MAX_RESPONSE_TASKS, + }; + use std::sync::atomic::{AtomicUsize, Ordering}; + let binding = ExecutionBinding::external(tokio::runtime::Handle::current()); + let mut runner = runner_on(binding.clone(), false); + let count = Arc::new(AtomicUsize::new(0)); + let (release, wait) = futures::channel::oneshot::channel(); + let wait = Mutex::new(Some(wait)); + let mut first = quiet_options(0); + first.definition = EvalDefinition::new( + Box::new(task({ + let count = count.clone(); + move |input: &usize, _| { + let wait = wait.lock().unwrap().take().unwrap(); + let count = count.clone(); + Box::pin(async move { + wait.await?; + count.fetch_add(1, Ordering::SeqCst); + Ok(*input) + }) + } + })), + vec![], + ) + .unwrap(); + let first = runner.eval("project", first); + assert!(first.is_ok()); + drop(first); + for index in 1..MAX_RESPONSE_TASKS { + let mut opts = quiet_options(index); + opts.definition = EvalDefinition::new( + Box::new(super::super::SimpleFnTask::new({ + let count = count.clone(); + move |input: &usize| { + count.fetch_add(1, Ordering::SeqCst); + Ok(*input) + } + })), + vec![], + ) + .unwrap(); + let token = runner.eval("project", opts); + assert!(token.is_ok()); + drop(token); + } + assert_eq!(binding.outstanding_responses(), MAX_RESPONSE_TASKS); + let rejected = runner.eval("project", quiet_options(99)); + assert!(rejected.is_err()); + assert!( + matches!(rejected.await, Err(BtEvalError::Run(error)) if matches!(error.failures.as_slice(), [super::super::EvalRunFailure::Admission(crate::EnqueueError::Full)])) + ); + let finish = runner.finish(); + assert!(finish.is_ok()); + release.send(()).unwrap(); + assert!(!finish.await.unwrap()); + assert_eq!(count.load(Ordering::SeqCst), MAX_RESPONSE_TASKS); + assert_eq!( + binding.outstanding_responses(), + 0, + "finish includes response publication" + ); +} +#[test] +fn external_runtime_loss_before_first_poll_retires_cli_work_without_fallback() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let mut runner = runner_on(ExecutionBinding::external(runtime.handle().clone()), false); + let status = runner.session.clone(); + let count = Arc::new(AtomicUsize::new(0)); + let mut opts = quiet_options(1); + opts.definition = super::super::EvalDefinition::new( + Box::new(super::super::SimpleFnTask::new({ + let count = count.clone(); + move |input: &usize| { + count.fetch_add(1, Ordering::SeqCst); + Ok(*input) + } + })), + vec![], + ) + .unwrap(); + let run = runner.eval("project", opts); + assert!(run.is_ok()); + let finish = runner.finish(); + assert!(finish.is_ok()); + drop(runtime); + assert!( + matches!(futures::executor::block_on(async { run.await }), Err(BtEvalError::Run(error)) if matches!(error.failures.as_slice(), [super::super::EvalRunFailure::Execution(crate::ExecutionError::WorkerLost)])) + ); + assert!(matches!( + futures::executor::block_on(async { finish.await }), + Err(BtFinishError::Execution(crate::ExecutionError::WorkerLost)) + )); + assert!(!status.all_passed()); + assert_eq!(count.load(Ordering::SeqCst), 0); +} +#[tokio::test] +async fn panic_releases_the_cli_order_guard_and_obeys_termination_policy() { + use crate::eval::{task, EvalDefinition}; + for terminate in [false, true] { + let mut runner = runner_on( + ExecutionBinding::external(tokio::runtime::Handle::current()), + terminate, + ); + let mut first = quiet_options(1); + first.definition = EvalDefinition::new( + Box::new(task(|_: &usize, _| { + Box::pin(async { panic!("async callback panic") }) + })), + vec![], + ) + .unwrap(); + let first = runner.eval("project", first); + let second = runner.eval("project", quiet_options(2)); + assert!( + matches!(first.await, Err(BtEvalError::Run(error)) if matches!(error.failures.as_slice(), [super::super::EvalRunFailure::Execution(crate::ExecutionError::WorkerLost)])) + ); + let second = second.await; + if terminate { + assert!(matches!(second, Err(BtEvalError::Stopped))); + } else { + second.unwrap(); + } + assert!(!runner.finish().await.unwrap()); + } +} +#[tokio::test] +async fn immediate_rejections_do_not_advance_the_cli_order_tail() { + let binding = ExecutionBinding::external(tokio::runtime::Handle::current()); + let mut runner = runner_on(binding.clone(), false); + let tail = runner.session.tail(); + let mut opts = quiet_options(1); + opts.max_concurrency = 0; + let rejected = runner.eval("project", opts); + assert!(rejected.is_err()); + assert!( + matches!(rejected.await, Err(BtEvalError::Run(error)) if matches!(error.failures.as_slice(), [super::super::EvalRunFailure::Config(super::super::EvalConfigError::ZeroConcurrency)])) + ); + let mut opts = quiet_options(1); + opts.runtime = crate::current_runtime(); + let rejected = runner.eval("project", opts); + assert!(rejected.is_err()); + assert!(matches!(rejected.await, Err(BtEvalError::RuntimeOverride))); + assert!(tail.ptr_eq(&runner.session.tail())); + runner.eval("project", quiet_options(2)).await.unwrap(); + assert!(!runner.finish().await.unwrap()); + assert_eq!(binding.outstanding_responses(), 0); +} +#[tokio::test] +async fn recorded_cli_runs_use_existing_admission_and_finish_does_not_close_shared_clients() { + use std::time::Duration; + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + for shutdown_first in [true, false] { + let server = MockServer::start().await; + for (endpoint, body) in [ + ( + "/api/apikey/login", + json!({"org_info":[{"id":"org","name":"org"}]}), + ), + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}}), + ), + ("/logs3", 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(json!({}))) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(crate::current_runtime()) + .api_key("key") + .app_url(server.uri()) + .api_url(server.uri()) + .default_project("project") + .blocking_login(true) + .build() + .await + .unwrap(); + let mut runner = runner_on(client.execution_binding(), false); + runner.config.no_send_logs = false; + runner.client = Some(Arc::new(client.clone())); + let first = runner.eval("project", quiet_options(1)); + let second = runner.eval("project", quiet_options(2)); + assert!(first.is_ok()); + assert!(second.is_ok()); + assert_eq!(client.execution_binding().outstanding_responses(), 2); + let shutdown = shutdown_first.then(|| client.shutdown(Duration::from_secs(5))); + let finished = runner.finish(); + assert!(finished.is_ok()); + assert!(finished.await.unwrap()); + // Joining the session includes publication of every earlier token. + use futures::FutureExt; + use std::future::IntoFuture; + first.into_future().now_or_never().unwrap().unwrap(); + second.into_future().now_or_never().unwrap().unwrap(); + if let Some(shutdown) = shutdown { + shutdown.await.unwrap(); + } else { + client.wait_for_login().await.unwrap(); + client.flush().await.unwrap(); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + } + let requests = server.received_requests().await.unwrap(); + let rows: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .flat_map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + assert_eq!( + rows.iter() + .filter(|r| r["metrics"]["end"].is_number()) + .count(), + 2 + ); + } +} +#[tokio::test] +async fn source_failure_keeps_its_typed_partial_summary_at_cli_boundary() { + use crate::eval::{Case, EvalRunFailure, IterCaseSource}; + let mut runner = runner_on( + ExecutionBinding::external(tokio::runtime::Handle::current()), + true, + ); + let mut opts = quiet_options(1); + opts.max_concurrency = 1; + opts.source = Box::new(IterCaseSource::new( + vec![ + Ok(Case { + input: 1, + ..Default::default() + }), + Err(anyhow::anyhow!("source failed")), + ] + .into_iter(), + )); + let token = runner.eval("project", opts); + assert!(token.is_ok()); + let BtEvalError::Run(error) = token.await.unwrap_err() else { + panic!("expected typed run failure") + }; + assert!(matches!( + error.failures.as_slice(), + [EvalRunFailure::Source { case_index: 1, .. }] + )); + assert_eq!(error.partial.unwrap().successful_cases, 1); + assert!(!runner.finish().await.unwrap()); +} +#[tokio::test] +async fn finish_closes_sse_before_publication_even_when_its_observer_is_dropped() { + use std::sync::atomic::{AtomicBool, Ordering}; + struct Close { + output: Buffer, + closed: Arc, + notice: Option>, + } + impl Write for Close { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.output.write(bytes) + } + fn flush(&mut self) -> std::io::Result<()> { + self.output.flush() + } + } + impl Drop for Close { + fn drop(&mut self) { + self.closed.store(true, Ordering::SeqCst); + let _ = self.notice.take().unwrap().send(()); + } + } + for discard in [true, false] { + let closed = Arc::new(AtomicBool::new(false)); + let (notice, closing) = tokio::sync::oneshot::channel(); + let sse = Arc::new(SseWriter { + writer: Mutex::new(Box::new(Close { + output: Buffer::default(), + closed: closed.clone(), + notice: Some(notice), + })), + }); + let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) + .await + .unwrap(); + drop(runner.eval("project", quiet_options(1))); + let finish = runner.finish(); + assert!(finish.is_ok()); + if discard { + drop(finish); + closing.await.unwrap(); + } else { + assert!(finish.await.unwrap()); + } + assert!(closed.load(Ordering::SeqCst)); + } +} + +#[tokio::test] +async fn later_issuance_failure_does_not_skip_earlier_accepted_work() { + let mut runner = runner_on( + ExecutionBinding::external(tokio::runtime::Handle::current()), + true, + ); + let earlier = runner.eval("project", quiet_options(1)); + let mut invalid = quiet_options(2); + invalid.runtime = crate::current_runtime(); + let rejected = runner.eval("project", invalid); + assert!(rejected.is_err()); + assert!(matches!(rejected.await, Err(BtEvalError::RuntimeOverride))); + earlier.await.unwrap(); + assert!(matches!( + runner.eval("project", quiet_options(3)).await, + Err(BtEvalError::Stopped) + )); + assert!(!runner.finish().await.unwrap()); +} diff --git a/src/eval/cli_session.rs b/src/eval/cli_session.rs new file mode 100644 index 0000000..7aa3f18 --- /dev/null +++ b/src/eval/cli_session.rs @@ -0,0 +1,195 @@ +//! Bounded live CLI ordering; no retained per-run outcomes. +use super::{EvalRunError, EvalSummary, ReportError}; +use crate::{Completion, EnqueueError, ExecutionError}; +use futures::{ + future::{BoxFuture, Shared}, + FutureExt, +}; +use std::{ + fmt, + sync::{Arc, Mutex}, +}; +use tokio::sync::oneshot; + +/// CLI run failure. Detailed computation/recording results remain available. +pub enum BtEvalError { + Run(EvalRunError), + CaseFailures { + summary: Box>, + }, + /// A previous evaluation failed with terminate-on-failure enabled. + Stopped, + /// CLI runs inherit their session/client binding, rather than a run override. + RuntimeOverride, +} +impl fmt::Debug for BtEvalError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Run(error) => f.debug_tuple("Run").field(error).finish(), + Self::CaseFailures { summary } => f + .debug_struct("CaseFailures") + .field("failed_cases", &summary.failed_cases) + .finish(), + Self::Stopped => f.write_str("Stopped"), + Self::RuntimeOverride => f.write_str("RuntimeOverride"), + } + } +} +impl fmt::Display for BtEvalError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Run(error) => error.fmt(f), + Self::CaseFailures { summary } => write!(f, "Evaluator '{}' failed with {} error(s)", summary.name, summary.failed_cases), + Self::Stopped => f.write_str("CLI evaluation session stopped after a previous failure"), + Self::RuntimeOverride => f.write_str("CLI evaluations inherit their session/client runtime; omit the run runtime override"), + } + } +} +impl std::error::Error for BtEvalError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Run(error) => Some(error), + _ => None, + } + } +} +impl From> for BtEvalError { + fn from(error: EvalRunError) -> Self { + Self::Run(error) + } +} +impl From for BtEvalError { + fn from(error: ExecutionError) -> Self { + Self::Run(error.into()) + } +} +impl From for BtEvalError { + fn from(error: EnqueueError) -> Self { + Self::Run(error.into()) + } +} + +/// Failure to join a CLI session or write its list output. +#[derive(Debug, thiserror::Error)] +pub enum BtFinishError { + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Output(#[from] ReportError), +} + +type Barrier = Shared>; +#[derive(Clone)] +pub(crate) struct Session(Arc>); +struct State { + all_passed: bool, + stopped_after: Option, + next: u64, + terminate: bool, + tail: Barrier, +} +impl Session { + pub fn new(terminate: bool) -> Self { + Self(Arc::new(Mutex::new(State { + all_passed: true, + stopped_after: None, + next: 0, + terminate, + tail: futures::future::ready(()).boxed().shared(), + }))) + } + fn state(&self) -> std::sync::MutexGuard<'_, State> { + self.0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + pub fn all_passed(&self) -> bool { + self.state().all_passed + } + pub fn stopped(&self) -> bool { + self.state().stopped_after.is_some() + } + pub fn should_skip(&self, sequence: u64) -> bool { + self.state() + .stopped_after + .is_some_and(|failed| sequence > failed) + } + pub fn fail(&self) { + let sequence = self.state().next; + self.fail_at(sequence); + } + pub fn fail_at(&self, sequence: u64) { + let mut state = self.state(); + state.all_passed = false; + if state.terminate { + state.stopped_after = Some( + state + .stopped_after + .map_or(sequence, |previous| previous.min(sequence)), + ); + } + } + /// Called inside response admission, after acquiring any write operation. + pub fn reserve(&self) -> Result<(Barrier, Ticket), BtEvalError> { + let mut state = self.state(); + if state.stopped_after.is_some() { + return Err(BtEvalError::Stopped); + } + let sequence = state.next; + state.next = sequence + .checked_add(1) + .ok_or_else(|| BtEvalError::from(EnqueueError::Full))?; + let previous = state.tail.clone(); + let (sender, receiver) = oneshot::channel(); + state.tail = async move { + let _ = receiver.await; + } + .boxed() + .shared(); + Ok(( + previous, + Ticket { + session: self.clone(), + sender: Some(sender), + completed: false, + sequence, + }, + )) + } + pub fn tail(&self) -> Barrier { + self.state().tail.clone() + } +} +pub(crate) struct Ticket { + session: Session, + sender: Option>, + completed: bool, + sequence: u64, +} +impl Ticket { + pub fn sequence(&self) -> u64 { + self.sequence + } + pub fn complete(mut self, failed: bool) { + if failed { + self.session.fail_at(self.sequence); + } + self.completed = true; + } +} +impl Drop for Ticket { + fn drop(&mut self) { + if !self.completed { + self.session.fail_at(self.sequence); + } + // Unblock successors even on panic, timeout, or loss before first poll. + if let Some(sender) = self.sender.take() { + let _ = sender.send(()); + } + } +} +pub(crate) fn ready(value: T) -> Completion { + let (sender, completion) = Completion::pair(); + let _ = sender.send(Ok(value)); + completion +} diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 2306503..dd4a750 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -52,6 +52,7 @@ mod blocking; mod bt_runner; mod case_runner; +mod cli_session; mod definition; mod error; mod executor; @@ -61,9 +62,11 @@ mod run; mod runner; mod scorer; mod source; +mod submission; mod task; mod types; +pub use cli_session::{BtEvalError, BtFinishError}; pub use definition::EvalDefinition; pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; diff --git a/src/eval/run.rs b/src/eval/run.rs index 04ab6b6..63187b0 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -1,15 +1,13 @@ //! Immediate submission of local and recorded evaluations. use super::{ - executor::{EvalExecutor, EvalPlan}, - recorder::BraintrustRecorder, reporter::{SilentReporter, TerminalReporter}, + submission::{self, RunReporter}, EvalConfigError, EvalOpts, EvalReporter, EvalRunError, EvalRunFailure, EvalSummary, }; use crate::{internal::execution::ExecutionBinding, BraintrustClient, Completion}; use serde::Serialize; type RunCompletion = Completion, EvalRunError>; -type RunReporter = Box>; type SubmitRecorded = fn(BraintrustClient, EvalOpts, RunReporter) -> RunCompletion; @@ -84,7 +82,7 @@ pub fn run( if let Err(error) = request.options.validate() { return Completion::rejected(error.into()); } - let mut reporter = request.reporter.unwrap_or_else(|| { + let reporter = request.reporter.unwrap_or_else(|| { if request.options.quiet { Box::new(SilentReporter) } else { @@ -100,15 +98,8 @@ pub fn run( return Completion::rejected(EvalRunError::new(EvalRunFailure::Runtime(source))) } }; - let callbacks = binding.clone(); - binding.response(async move { - EvalExecutor::run( - EvalPlan::new(options)?, - None, - reporter.as_mut(), - Some(&callbacks), - ) - .await + submission::local(options, binding, reporter, || { + Ok((std::future::ready(Ok(())), std::convert::identity, |_| {})) }) } Target::Recorded { client, submit } => { @@ -123,34 +114,9 @@ pub fn run( fn recorded( client: BraintrustClient, options: EvalOpts, - mut reporter: RunReporter, + reporter: RunReporter, ) -> RunCompletion { - let binding = client.execution_binding(); - let callbacks = binding.clone(); - binding.response_try(|| { - let operation = client.span_operation()?; - Ok(async move { - let mut recorder = BraintrustRecorder::new( - client, - operation.submitter.clone(), - options.project_name.clone(), - options - .experiment_name - .clone() - .unwrap_or_else(|| options.name.clone()), - options.metadata.clone(), - ); - let result = EvalExecutor::run( - EvalPlan::new(options)?, - Some(&mut recorder), - reporter.as_mut(), - Some(&callbacks), - ) - .await; - drop(recorder); - // Revoke write permission before publishing, including escaped spans. - drop(operation); - result - }) + submission::recorded(client, options, reporter, || { + Ok((std::future::ready(Ok(())), std::convert::identity, |_| {})) }) } diff --git a/src/eval/submission.rs b/src/eval/submission.rs new file mode 100644 index 0000000..1d9a4fd --- /dev/null +++ b/src/eval/submission.rs @@ -0,0 +1,129 @@ +//! One response reservation per run, shared by public and CLI entry points. +use super::{ + executor::{EvalExecutor, EvalPlan}, + recorder::BraintrustRecorder, + EvalConfigError, EvalOpts, EvalReporter, EvalRunError, EvalSummary, +}; +use crate::{ + internal::execution::ExecutionBinding, BraintrustClient, Completion, EnqueueError, + ExecutionError, +}; +use serde::Serialize; +use std::future::Future; + +pub(crate) type RunReporter = Box>; + +// Hooks are made only after response/write admission. They provide CLI ordering +// and result translation without a second response task or runtime. +pub(crate) fn local( + options: EvalOpts, + binding: ExecutionBinding, + mut reporter: RunReporter, + make: M, +) -> Completion +where + I: Send + Sync + 'static, + O: Send + Sync + 'static, + T: Send + 'static, + E: From> + From + From + Send + 'static, + M: FnOnce() -> Result<(B, A, C), E>, + B: Future> + Send + 'static, + A: FnOnce(Result, E>) -> Result + Send + 'static, + C: FnOnce(bool) + Send + 'static, +{ + if let Err(error) = options.validate() { + return Completion::rejected(E::from(EvalRunError::from(error))); + } + let callbacks = binding.clone(); + binding.response_try_finalized(|| { + let (before, after, finalized) = make()?; + Ok(( + async move { + let result = match before.await { + Ok(()) => { + let plan = EvalPlan::new(options) + .map_err(EvalRunError::from) + .map_err(E::from)?; + EvalExecutor::run(plan, None, reporter.as_mut(), Some(&callbacks)) + .await + .map_err(E::from) + } + Err(error) => Err(error), + }; + drop(reporter); + after(result) + }, + finalized, + )) + }) +} + +pub(crate) fn recorded( + client: BraintrustClient, + options: EvalOpts, + mut reporter: RunReporter, + make: M, +) -> Completion +where + I: Serialize + Send + Sync + 'static, + O: Serialize + Send + Sync + 'static, + T: Send + 'static, + E: From> + From + From + Send + 'static, + M: FnOnce() -> Result<(B, A, C), E>, + B: Future> + Send + 'static, + A: FnOnce(Result, E>) -> Result + Send + 'static, + C: FnOnce(bool) + Send + 'static, +{ + if let Err(error) = options.validate() { + return Completion::rejected(E::from(EvalRunError::from(error))); + } + if options.runtime.is_some() { + return Completion::rejected(E::from(EvalRunError::from( + EvalConfigError::RecordedRuntime, + ))); + } + let binding = client.execution_binding(); + let callbacks = binding.clone(); + binding.response_try_finalized(|| { + let operation = client + .span_operation() + .map_err(EvalRunError::from) + .map_err(E::from)?; + let (before, after, finalized) = make()?; + Ok(( + async move { + let result = match before.await { + Ok(()) => { + let mut recorder = BraintrustRecorder::new( + client, + operation.submitter.clone(), + options.project_name.clone(), + options + .experiment_name + .clone() + .unwrap_or_else(|| options.name.clone()), + options.metadata.clone(), + ); + let plan = EvalPlan::new(options) + .map_err(EvalRunError::from) + .map_err(E::from)?; + EvalExecutor::run( + plan, + Some(&mut recorder), + reporter.as_mut(), + Some(&callbacks), + ) + .await + .map_err(E::from) + } + Err(error) => Err(error), + }; + drop(reporter); + // Revoke even escaped spans before status, ordering or token publication. + drop(operation); + after(result) + }, + finalized, + )) + }) +} diff --git a/src/internal/execution.rs b/src/internal/execution.rs index eebfde4..9f6c419 100644 --- a/src/internal/execution.rs +++ b/src/internal/execution.rs @@ -159,6 +159,21 @@ impl ExecutionBinding { E: From + From + Send + 'static, F: Future> + Send + 'static, M: FnOnce() -> Result, + { + self.response_try_finalized(|| make().map(|work| (work, |_| {}))) + } + + /// A bounded operation may release its private join/order guard after token + /// publication. The finalizer runs outside registry locks; false means a + /// cancellation already published instead. Its captures are dropped on task + /// loss, including before first poll, just like the operation's captures. + pub fn response_try_finalized(&self, make: M) -> Completion + where + T: Send + 'static, + E: From + From + Send + 'static, + F: Future> + Send + 'static, + M: FnOnce() -> Result<(F, C), E>, + C: FnOnce(bool) + Send + 'static, { let (sender, mut completion) = Completion::pair(); let (guard, response) = match self.responses.admit(sender) { @@ -168,7 +183,7 @@ impl ExecutionBinding { return completion; } }; - let work = match make() { + let (work, finalized) = match make() { Ok(work) => work, Err(error) => { completion.record_issuance(Err(error)); @@ -177,7 +192,8 @@ impl ExecutionBinding { }; let id = guard.id; let task = self.spawn(async move { - guard.complete(&response, work.await); + let published = guard.complete(&response, work.await); + finalized(published); }); self.responses.install_abort(id, task.abort_handle()); completion @@ -430,4 +446,28 @@ mod tests { .unwrap(); }); } + #[tokio::test] + async fn operation_finalizer_observes_cancellation_that_won_publication() { + let binding = ExecutionBinding::external(Handle::current()); + let cancel = binding.clone(); + let (sender, finalized) = oneshot::channel(); + let token = binding.response_try_finalized::(|| { + Ok(( + async move { + cancel.cancel_responses(ExecutionError::TimedOut); + Ok(7) + }, + move |published| { + let _ = sender.send(published); + }, + )) + }); + assert!(token.is_ok()); + assert!(matches!( + token.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(!finalized.await.unwrap()); + assert_eq!(binding.outstanding_responses(), 0); + } } diff --git a/src/internal/response_tasks.rs b/src/internal/response_tasks.rs index 0d45a31..9ce54da 100644 --- a/src/internal/response_tasks.rs +++ b/src/internal/response_tasks.rs @@ -167,7 +167,7 @@ pub(crate) struct TaskGuard { finished: bool, } impl TaskGuard { - pub fn complete(mut self, response: &Response, result: Result) { + pub fn complete(mut self, response: &Response, result: Result) -> bool { self.finished = true; if let Some((entry, publication)) = self.tasks.completing(self.id) { let unused = response.send(result); @@ -176,8 +176,10 @@ impl TaskGuard { // under the registry lock; both can reenter SDK scheduling. drop(unused); drop(publication); + true } else { drop(result); // Cancellation already published this task's outcome. + false } } } From d2bee75f96fbdb8db6c31047dc147a1c93fc16e5 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 17:06:50 +0800 Subject: [PATCH 6/6] Make CLI construction eager and consolidate the evaluation API --- examples/eval_async.rs | 3 +- examples/eval_sync.rs | 22 +- src/eval/bt_runner.rs | 272 +++++++++------------- src/eval/bt_runner/config.rs | 89 +++++++ src/eval/bt_runner/tests.rs | 112 ++++++--- src/eval/cli_session.rs | 25 ++ src/eval/mod.rs | 6 +- src/eval/{runner.rs => options.rs} | 49 +--- src/logger.rs | 51 ++-- tests/cli_initialization.rs | 360 +++++++++++++++++++++++++++++ tests/eval_callables.rs | 8 +- tests/eval_test.rs | 90 ++++---- 12 files changed, 765 insertions(+), 322 deletions(-) create mode 100644 src/eval/bt_runner/config.rs rename src/eval/{runner.rs => options.rs} (67%) create mode 100644 tests/cli_initialization.rs diff --git a/examples/eval_async.rs b/examples/eval_async.rs index 5e0c833..914526a 100644 --- a/examples/eval_async.rs +++ b/examples/eval_async.rs @@ -12,7 +12,8 @@ use braintrust_sdk_rust::eval::*; async fn main() -> anyhow::Result<()> { // BtEvalRunner reads BT_EVAL_* env vars automatically. // Recording mode creates a client; local and list modes need no credentials. - let mut runner = BtEvalRunner::from_env().await?; + let initialization = BtEvalRunner::from_env(); + let mut runner = initialization.await?; // Define test cases let cases = vec![ diff --git a/examples/eval_sync.rs b/examples/eval_sync.rs index 09e9c50..1931d54 100644 --- a/examples/eval_sync.rs +++ b/examples/eval_sync.rs @@ -5,7 +5,6 @@ use std::future::IntoFuture; fn main() -> anyhow::Result<()> { // 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![ @@ -36,17 +35,16 @@ fn main() -> anyhow::Result<()> { println!("Starting synchronous evaluation...\n"); let definition = EvalDefinition::new(Box::new(task), scorers)?; - let result = futures::executor::block_on( - evaluator - .run( - EvalOpts::builder() - .name("sync-uppercase-eval".to_string()) - .source(cases.into()) - .definition(definition.clone()) - .build(), - ) - .into_future(), - )?; + let completion = run(EvalRequest::recorded( + EvalOpts::builder() + .name("sync-uppercase-eval".to_string()) + .source(cases.into()) + .definition(definition.clone()) + .build(), + client, + )); + assert!(completion.is_ok()); + let result = futures::executor::block_on(completion.into_future())?; // Print results println!("\n=== Results ===\n"); diff --git a/src/eval/bt_runner.rs b/src/eval/bt_runner.rs index 3972051..0783972 100644 --- a/src/eval/bt_runner.rs +++ b/src/eval/bt_runner.rs @@ -13,99 +13,20 @@ use std::sync::{Arc, Mutex}; use serde::Serialize; use serde_json::{json, Value}; -use crate::{ - internal::execution::ExecutionBinding, BraintrustClient, BraintrustError, Completion, Result, -}; +use crate::{internal::execution::ExecutionBinding, BraintrustClient, BraintrustError, Completion}; use super::{ cli_session::{ready, Session}, - runner::EvalOpts, + options::EvalOpts, submission::{self, RunReporter}, types::EvalSummary, - BtEvalError, BtFinishError, EvalReporter, JsonlReporter, ReportError, + BtEvalError, BtFinishError, BtInitError, EvalReporter, JsonlReporter, ReportError, }; -/// A filter for selecting which evaluators to run, read from `BT_EVAL_FILTER_PARSED`. -struct EvalFilter { - path: Vec, - pattern: regex::Regex, -} - -impl EvalFilter { - /// Returns true if the given eval name matches this filter. - fn matches(&self, name: &str) -> bool { - // Match against the eval name when path is ["name"] or empty - if self.path.is_empty() || self.path == ["name"] { - return self.pattern.is_match(name); - } - // Unknown path — include by default - true - } -} - -/// Configuration read from `BT_EVAL_*` environment variables. -struct RunnerConfig { - /// Output summaries as JSONL to stdout (`BT_EVAL_JSONL=1`) - jsonl: bool, - /// List evaluator names without running them (`BT_EVAL_LIST=1`) - list: bool, - /// Stop running after the first failure (`BT_EVAL_TERMINATE_ON_FAILURE=1`) - terminate_on_failure: bool, - /// Filters to select which evaluators to run (`BT_EVAL_FILTER_PARSED`) - filters: Vec, - /// Skip sending logs to Braintrust (`BT_EVAL_NO_SEND_LOGS=1` or `BT_EVAL_LOCAL=1`) - no_send_logs: bool, -} - -impl RunnerConfig { - fn from_env() -> Self { - Self { - jsonl: env_flag("BT_EVAL_JSONL"), - list: env_flag("BT_EVAL_LIST"), - terminate_on_failure: env_flag("BT_EVAL_TERMINATE_ON_FAILURE"), - filters: parse_filters(), - no_send_logs: env_flag("BT_EVAL_NO_SEND_LOGS") || env_flag("BT_EVAL_LOCAL"), - } - } -} - -fn env_flag(name: &str) -> bool { - !matches!( - std::env::var(name).ok().as_deref(), - None | Some("") | Some("0") | Some("false") | Some("no") | Some("off") - ) -} - -fn parse_filters() -> Vec { - let serialized = match std::env::var("BT_EVAL_FILTER_PARSED").ok() { - Some(s) if !s.is_empty() => s, - _ => return vec![], - }; - - let parsed: Vec = match serde_json::from_str(&serialized) { - Ok(v) => v, - Err(e) => { - eprintln!("Invalid BT_EVAL_FILTER_PARSED: {}", e); - return vec![]; - } - }; - - parsed - .into_iter() - .filter_map(|v| { - let path = v["path"] - .as_array()? - .iter() - .filter_map(|s| s.as_str().map(String::from)) - .collect(); - let pattern_str = v["pattern"].as_str()?; - let pattern = regex::Regex::new(pattern_str) - .map_err(|e| eprintln!("Invalid filter pattern '{}': {}", pattern_str, e)) - .ok()?; - Some(EvalFilter { path, pattern }) - }) - .collect() -} +mod config; +#[cfg(test)] +use config::{env_flag, parse_filters, EvalFilter}; +use config::{RunnerConfig, SseEndpoint}; /// SSE event writer that sends events over a socket to the `bt eval` CLI. /// @@ -144,43 +65,47 @@ impl EvalReporter for SseReporter { } } -fn create_sse_writer() -> Option> { - // Unix socket SSE is only available on Unix platforms; Windows continues - // to the TCP fallback below. - #[cfg(unix)] - if let Ok(sock_path) = std::env::var("BT_EVAL_SSE_SOCK") { - if !sock_path.is_empty() { - match std::os::unix::net::UnixStream::connect(&sock_path) { - Ok(stream) => { - return Some(Arc::new(SseWriter { - writer: Mutex::new(Box::new(stream)), - })); - } - Err(e) => { - eprintln!("Failed to connect to SSE socket {}: {}", sock_path, e); - } - } +fn connect_sse(endpoint: SseEndpoint) -> std::result::Result, BtInitError> { + let writer: Box = match endpoint { + #[cfg(unix)] + SseEndpoint::Unix(path) => Box::new( + std::os::unix::net::UnixStream::connect(path).map_err(|source| BtInitError::Sse { + transport: "Unix socket", + source, + })?, + ), + SseEndpoint::Tcp(address) => { + let stream = + std::net::TcpStream::connect(address).map_err(|source| BtInitError::Sse { + transport: "TCP", + source, + })?; + let _ = stream.set_nodelay(true); + Box::new(stream) } - } - - // Try TCP address (BT_EVAL_SSE_ADDR) - if let Ok(addr) = std::env::var("BT_EVAL_SSE_ADDR") { - if !addr.is_empty() { - match std::net::TcpStream::connect(&addr) { - Ok(stream) => { - let _ = stream.set_nodelay(true); - return Some(Arc::new(SseWriter { - writer: Mutex::new(Box::new(stream)), - })); - } - Err(e) => { - eprintln!("Failed to connect to SSE address {}: {}", addr, e); - } - } + }; + Ok(Arc::new(SseWriter { + writer: Mutex::new(writer), + })) +} +async fn prepare_sse( + binding: &ExecutionBinding, + endpoint: Option, +) -> std::result::Result>, BtInitError> { + let Some(endpoint) = endpoint else { + return Ok(None); + }; + let work = binding.spawn_blocking(move || connect_sse(endpoint)); + struct Abort(tokio::task::AbortHandle); + impl Drop for Abort { + fn drop(&mut self) { + self.0.abort(); } } - - None + let _abort = Abort(work.abort_handle()); + work.await + .map_err(|_| BtInitError::Execution(crate::ExecutionError::WorkerLost))? + .map(Some) } /// A runner designed to work with the `bt eval` CLI. @@ -207,7 +132,8 @@ fn create_sse_writer() -> Option> { /// use braintrust_sdk_rust::eval::*; /// /// # async fn example() -> anyhow::Result<()> { -/// let mut runner = BtEvalRunner::from_env().await?; +/// let initialization = BtEvalRunner::from_env(); // initialization is admitted now +/// let mut runner = initialization.await?; /// /// let definition = EvalDefinition::new( /// Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), @@ -247,56 +173,82 @@ impl BtEvalRunner { /// Create a new runner, reading configuration from environment variables. /// /// Recording mode creates a client from environment variables (e.g. - /// `BRAINTRUST_API_KEY`). List and local modes create no client. - pub async fn from_env() -> Result { - let config = RunnerConfig::from_env(); - let sse = create_sse_writer(); - - Self::from_config(config, sse).await - } - - async fn from_config(config: RunnerConfig, sse: Option>) -> Result { - let client = if !config.list && !config.no_send_logs { - Some(Arc::new(BraintrustClient::builder().build().await?)) - } else { - None + /// `BRAINTRUST_API_KEY`). List and local modes create no client. Configuration + /// is captured and initialization admitted at the call; awaiting only observes + /// the result. Invalid filters reject immediately; SSE failures are returned + /// without switching output transports. Socket setup precedes client login. + pub fn from_env() -> Completion { + let config = match RunnerConfig::from_env() { + Ok(config) => config, + Err(error) => return Completion::rejected(error), }; + Self::initialize(config, None) + } - let binding = if let Some(client) = &client { - Some(client.execution_binding()) - } else if !config.list { - Some(ExecutionBinding::new(None)?) - } else { - None - }; + fn assemble( + config: RunnerConfig, + sse: Option>, + client: Option>, + binding: Option, + ) -> Self { let session = Session::new(config.terminate_on_failure); - Ok(Self { + Self { config, sse, client, binding, session, registered_names: Vec::new(), - }) + } } - /// Create a new runner with an explicit client. - /// - /// Environment variables are still read for mode configuration (list, JSONL, SSE, etc.), - /// but the provided client is used instead of creating a new one. All runs, - /// including local mode, inherit its fixed execution binding. - pub fn with_client(client: Arc) -> Self { - let config = RunnerConfig::from_env(); - let sse = create_sse_writer(); - let binding = Some(client.execution_binding()); - let session = Session::new(config.terminate_on_failure); - Self { - config, - sse, - client: Some(client), - binding, - session, - registered_names: Vec::new(), + /// Initialize now on this client's fixed binding. Dropping the completion + /// does not cancel socket setup. List mode uses neither its runtime nor SSE. + /// Configuration is captured from the environment at the call. + pub fn with_client(client: Arc) -> Completion { + let config = match RunnerConfig::from_env() { + Ok(config) => config, + Err(error) => return Completion::rejected(error), + }; + Self::initialize(config, Some(client)) + } + + fn initialize( + config: RunnerConfig, + supplied: Option>, + ) -> Completion { + if config.list { + return ready(Self::assemble(config, None, supplied, None)); + } + if let Some(client) = supplied { + let binding = client.execution_binding(); + let context = binding.clone(); + return binding.response(async move { + let sse = prepare_sse(&context, config.endpoint.clone()).await?; + Ok(Self::assemble(config, sse, Some(client), Some(context))) + }); + } + if config.no_send_logs { + let binding = match ExecutionBinding::new(None) { + Ok(binding) => binding, + Err(error) => return Completion::rejected(error.into()), + }; + let context = binding.clone(); + binding.response(async move { + let sse = prepare_sse(&context, config.endpoint.clone()).await?; + Ok(Self::assemble(config, sse, None, Some(context))) + }) + } else { + BraintrustClient::builder().build_with(|builder, binding| async move { + let sse = prepare_sse(&binding, config.endpoint.clone()).await?; + let client = builder.initialize(binding.clone()).await?; + Ok(Self::assemble( + config, + sse, + Some(Arc::new(client)), + Some(binding), + )) + }) } } diff --git a/src/eval/bt_runner/config.rs b/src/eval/bt_runner/config.rs new file mode 100644 index 0000000..bea9a8b --- /dev/null +++ b/src/eval/bt_runner/config.rs @@ -0,0 +1,89 @@ +//! Capture and validate CLI settings before admitting initialization. +use crate::eval::BtInitError; +use serde::Deserialize; + +pub(super) struct EvalFilter { + pub path: Vec, + pub pattern: regex::Regex, +} +impl EvalFilter { + pub fn matches(&self, name: &str) -> bool { + if self.path.is_empty() || self.path == ["name"] { + self.pattern.is_match(name) + } else { + true + } // Preserve existing handling of unsupported CLI filter paths. + } +} +#[derive(Clone)] +pub(super) enum SseEndpoint { + #[cfg(unix)] + Unix(std::path::PathBuf), + Tcp(String), +} +pub(super) struct RunnerConfig { + pub jsonl: bool, + pub list: bool, + pub terminate_on_failure: bool, + pub filters: Vec, + pub no_send_logs: bool, + pub endpoint: Option, +} +impl RunnerConfig { + pub fn from_env() -> Result { + let list = env_flag("BT_EVAL_LIST"); + // List output is stdout-only and requires no socket or runtime setup. + let mut endpoint = None; + if !list { + #[cfg(unix)] + if let Some(path) = std::env::var_os("BT_EVAL_SSE_SOCK").filter(|s| !s.is_empty()) { + endpoint = Some(SseEndpoint::Unix(path.into())); + } + if endpoint.is_none() { + if let Some(address) = std::env::var("BT_EVAL_SSE_ADDR") + .ok() + .filter(|s| !s.is_empty()) + { + endpoint = Some(SseEndpoint::Tcp(address)); + } + } + } + Ok(Self { + jsonl: env_flag("BT_EVAL_JSONL"), + list, + terminate_on_failure: env_flag("BT_EVAL_TERMINATE_ON_FAILURE"), + filters: parse_filters(std::env::var("BT_EVAL_FILTER_PARSED").ok().as_deref())?, + no_send_logs: env_flag("BT_EVAL_NO_SEND_LOGS") || env_flag("BT_EVAL_LOCAL"), + endpoint, + }) + } +} +pub(super) fn env_flag(name: &str) -> bool { + !matches!( + std::env::var(name).ok().as_deref(), + None | Some("") | Some("0") | Some("false") | Some("no") | Some("off") + ) +} +pub(super) fn parse_filters(serialized: Option<&str>) -> Result, BtInitError> { + let Some(serialized) = serialized.filter(|s| !s.is_empty()) else { + return Ok(vec![]); + }; + #[derive(Deserialize)] + struct Filter { + path: Vec, + pattern: String, + } + let filters: Vec = serde_json::from_str(serialized).map_err(BtInitError::Filters)?; + filters + .into_iter() + .enumerate() + .map(|(index, filter)| { + let pattern = regex::Regex::new(&filter.pattern) + .map_err(|source| BtInitError::FilterPattern { index, source })?; + Ok(EvalFilter { + path: filter.path, + pattern, + }) + }) + .collect() +} diff --git a/src/eval/bt_runner/tests.rs b/src/eval/bt_runner/tests.rs index 57ebb36..ccf8086 100644 --- a/src/eval/bt_runner/tests.rs +++ b/src/eval/bt_runner/tests.rs @@ -3,45 +3,28 @@ use std::sync::Mutex; use super::*; use crate::eval::EvalRequest; -// Serialize tests that mutate BT_EVAL_FILTER_PARSED to avoid race conditions. -static FILTER_ENV_MUTEX: Mutex<()> = Mutex::new(()); - #[test] fn test_env_flag() { - // Not set → false assert!(!env_flag("BT_TEST_NONEXISTENT_FLAG_XYZ")); } - #[test] fn test_parse_filters_empty() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - // When BT_EVAL_FILTER_PARSED is not set, filters should be empty - std::env::remove_var("BT_EVAL_FILTER_PARSED"); - let filters = parse_filters(); - assert!(filters.is_empty()); + assert!(parse_filters(None).unwrap().is_empty()); + assert!(parse_filters(Some("")).unwrap().is_empty()); } - #[test] fn test_parse_filters_valid() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - std::env::set_var( - "BT_EVAL_FILTER_PARSED", - r#"[{"path":["name"],"pattern":"foo.*"}]"#, - ); - let filters = parse_filters(); + let filters = parse_filters(Some(r#"[{"path":["name"],"pattern":"foo.*"}]"#)).unwrap(); assert_eq!(filters.len(), 1); assert!(filters[0].matches("foobar")); assert!(!filters[0].matches("bar")); - std::env::remove_var("BT_EVAL_FILTER_PARSED"); } - #[test] fn test_parse_filters_invalid_json() { - let _guard = FILTER_ENV_MUTEX.lock().unwrap(); - std::env::set_var("BT_EVAL_FILTER_PARSED", "not-json"); - let filters = parse_filters(); - assert!(filters.is_empty()); - std::env::remove_var("BT_EVAL_FILTER_PARSED"); + assert!(matches!( + parse_filters(Some("not-json")), + Err(BtInitError::Filters(_)) + )); } #[test] @@ -81,8 +64,17 @@ fn local_config(terminate_on_failure: bool) -> RunnerConfig { terminate_on_failure, filters: vec![], no_send_logs: true, + endpoint: None, } } +fn local_reported_runner(config: RunnerConfig, sse: Arc) -> BtEvalRunner { + BtEvalRunner::assemble( + config, + Some(sse), + None, + Some(ExecutionBinding::new(None).unwrap()), + ) +} fn local_options(fail: bool) -> EvalOpts { super::super::EvalOpts::builder() .name("cli".into()) @@ -115,9 +107,7 @@ async fn sse_protocol_is_the_only_reporter_and_case_failures_update_cli_status() let sse = Arc::new(SseWriter { writer: Mutex::new(Box::new(output.clone())), }); - let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) - .await - .unwrap(); + let mut runner = local_reported_runner(local_config(false), sse); assert!(runner.client.is_none()); runner.eval("project", local_options(false)).await.unwrap(); runner.eval("project", local_options(true)).await.unwrap(); @@ -139,9 +129,7 @@ async fn terminate_on_failure_returns_an_error_and_retains_failed_session_status let sse = Arc::new(SseWriter { writer: Mutex::new(Box::new(Buffer::default())), }); - let mut runner = BtEvalRunner::from_config(local_config(true), Some(sse)) - .await - .unwrap(); + let mut runner = local_reported_runner(local_config(true), sse); assert!(runner.eval("project", local_options(true)).await.is_err()); assert!(!runner.all_passed()); assert!(!runner.finish().await.unwrap()); @@ -606,9 +594,7 @@ async fn finish_closes_sse_before_publication_even_when_its_observer_is_dropped( notice: Some(notice), })), }); - let mut runner = BtEvalRunner::from_config(local_config(false), Some(sse)) - .await - .unwrap(); + let mut runner = local_reported_runner(local_config(false), sse); drop(runner.eval("project", quiet_options(1))); let finish = runner.finish(); assert!(finish.is_ok()); @@ -641,3 +627,63 @@ async fn later_issuance_failure_does_not_skip_earlier_accepted_work() { )); assert!(!runner.finish().await.unwrap()); } + +#[tokio::test] +async fn constructor_uses_one_client_response_slot_and_rejects_before_socket_work() { + use crate::internal::response_tasks::MAX_RESPONSE_TASKS; + let client = BraintrustClient::builder() + .runtime(crate::current_runtime()) + .skip_login(true) + .build() + .await + .unwrap(); + let binding = client.execution_binding(); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + for _ in 0..MAX_RESPONSE_TASKS { + let gate = gate.clone(); + let token = binding.response::<(), BraintrustError, _>(async move { + gate.acquire().await.unwrap().forget(); + Ok(()) + }); + assert!(token.is_ok()); + drop(token); + } + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let mut config = local_config(false); + config.endpoint = Some(SseEndpoint::Tcp(listener.local_addr().unwrap().to_string())); + let rejected = BtEvalRunner::initialize(config, Some(Arc::new(client.clone()))); + assert!(rejected.is_err()); + assert!(matches!( + rejected.await, + Err(BtInitError::Admission(crate::EnqueueError::Full)) + )); + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); + gate.add_permits(MAX_RESPONSE_TASKS); + binding.responses_drained().await.unwrap(); + let token = BtEvalRunner::initialize(local_config(false), Some(Arc::new(client.clone()))); + assert!(token.is_ok()); + assert_eq!(binding.outstanding_responses(), 1); + let runner = token.await.unwrap(); + assert!(runner.finish().await.unwrap()); + client + .shutdown(std::time::Duration::from_secs(5)) + .await + .unwrap(); + let rejected = BtEvalRunner::initialize(local_config(false), Some(Arc::new(client.clone()))); + assert!(rejected.is_err()); + assert!(matches!( + rejected.await, + Err(BtInitError::Admission(crate::EnqueueError::Closed)) + )); + let mut list = local_config(false); + list.list = true; + let token = BtEvalRunner::initialize(list, Some(Arc::new(client))); + assert!(token.is_ok()); + let list = token.await.unwrap(); + assert!(list.binding.is_none()); + assert!(list.sse.is_none()); + assert!(list.finish().await.unwrap()); +} diff --git a/src/eval/cli_session.rs b/src/eval/cli_session.rs index 7aa3f18..feeb0d7 100644 --- a/src/eval/cli_session.rs +++ b/src/eval/cli_session.rs @@ -11,6 +11,31 @@ use std::{ }; use tokio::sync::oneshot; +/// Failure to configure or initialize a CLI session. +#[derive(Debug, thiserror::Error)] +pub enum BtInitError { + #[error(transparent)] + Client(#[from] crate::BraintrustError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error("invalid BT_EVAL_FILTER_PARSED: {0}")] + Filters(#[source] serde_json::Error), + #[error("invalid pattern in CLI filter {index}: {source}")] + FilterPattern { + index: usize, + #[source] + source: regex::Error, + }, + #[error("failed to connect to CLI SSE {transport}: {source}")] + Sse { + transport: &'static str, + #[source] + source: std::io::Error, + }, +} + /// CLI run failure. Detailed computation/recording results remain available. pub enum BtEvalError { Run(EvalRunError), diff --git a/src/eval/mod.rs b/src/eval/mod.rs index dd4a750..fe890f6 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -56,27 +56,27 @@ mod cli_session; mod definition; mod error; mod executor; +mod options; mod recorder; mod reporter; mod run; -mod runner; mod scorer; mod source; mod submission; mod task; mod types; -pub use cli_session::{BtEvalError, BtFinishError}; +pub use cli_session::{BtEvalError, BtFinishError, BtInitError}; pub use definition::EvalDefinition; pub use error::{CaseError, EvalConfigError, EvalRunError, EvalRunFailure}; // Re-export public API pub use bt_runner::{summary_to_value, BtEvalRunner}; +pub use options::EvalOpts; pub use reporter::{ EvalReporter, JsonlReporter, ReportError, ReportStage, SilentReporter, TerminalReporter, }; pub use run::{run, EvalRequest}; -pub use runner::{EvalOpts, Evaluator}; pub use scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer}; pub use source::{ CaseSource, DatasetCaseSource, DatasetVersionRef, IterCaseSource, SourceError, diff --git a/src/eval/runner.rs b/src/eval/options.rs similarity index 67% rename from src/eval/runner.rs rename to src/eval/options.rs index a1afb66..c00b4b7 100644 --- a/src/eval/runner.rs +++ b/src/eval/options.rs @@ -1,13 +1,6 @@ -use serde::Serialize; use serde_json::{Map, Value}; -use std::marker::PhantomData; -use std::sync::Arc; -use crate::BraintrustClient; - -use super::{ - source::CaseSource, types::EvalSummary, EvalConfigError, EvalDefinition, EvalRunError, -}; +use super::{source::CaseSource, EvalConfigError, EvalDefinition}; /// Options for running an evaluation #[derive(bon::Builder)] @@ -47,46 +40,6 @@ impl EvalOpts { } } -/// Client-bound convenience facade for recorded evaluations. -pub struct Evaluator { - client: Arc, - _phantom: PhantomData<(I, O)>, -} - -impl Evaluator -where - I: Serialize + Send + Sync + 'static, - O: Serialize + Send + Sync + 'static, -{ - /// Create a new evaluator - pub fn new(client: BraintrustClient) -> Self { - Self { - client: Arc::new(client), - _phantom: PhantomData, - } - } - - /// Create a new evaluator from an Arc client (avoids cloning) - pub fn from_arc(client: Arc) -> Self { - Self { - client, - _phantom: PhantomData, - } - } - - /// Submit a recorded run immediately on the client binding. - /// Dropping the completion does not cancel accepted work. - pub fn run( - &self, - opts: EvalOpts, - ) -> crate::Completion, EvalRunError> { - super::run(super::EvalRequest::recorded( - opts, - self.client.as_ref().clone(), - )) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/logger.rs b/src/logger.rs index bfea629..e67fdd2 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -340,29 +340,48 @@ impl BraintrustClientBuilder { /// Schedule client initialization immediately and observe its result. /// Dropping the completion does not cancel initialization or login. pub fn build(self) -> crate::Completion { - let validate = || -> Result<()> { - Url::parse(self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL)) - .map_err(|_| BraintrustError::InvalidConfig("invalid app_url".into()))?; - Url::parse(self.api_url.as_deref().unwrap_or(DEFAULT_API_URL)) - .map_err(|_| BraintrustError::InvalidConfig("invalid api_url".into()))?; - if !self.skip_login && self.api_key.is_none() { - return Err(BraintrustError::InvalidConfig( - "API key required: set BRAINTRUST_API_KEY or call .api_key()".into(), - )); - } - Ok(()) - }; - let binding = validate().and_then(|()| ExecutionBinding::new(self.runtime.clone())); + self.build_with(|builder, binding| async move { builder.initialize(binding).await }) + } + + /// Compose initialization with an internal adapter in one response task on + /// the same validated binding. No nested build token or second runtime. + pub(crate) fn build_with(self, work: M) -> crate::Completion + where + T: Send + 'static, + E: From + + From + + From + + Send + + 'static, + M: FnOnce(Self, ExecutionBinding) -> F + Send + 'static, + F: std::future::Future> + Send + 'static, + { + let binding = self + .validate() + .and_then(|()| ExecutionBinding::new(self.runtime.clone())); match binding { Ok(binding) => { let context = binding.clone(); - binding.response(async move { self.initialize(context).await }) + binding.response(async move { work(self, context).await }) } - Err(error) => crate::Completion::rejected(error), + Err(error) => crate::Completion::rejected(E::from(error)), } } - async fn initialize(self, execution: ExecutionBinding) -> Result { + fn validate(&self) -> Result<()> { + Url::parse(self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL)) + .map_err(|_| BraintrustError::InvalidConfig("invalid app_url".into()))?; + Url::parse(self.api_url.as_deref().unwrap_or(DEFAULT_API_URL)) + .map_err(|_| BraintrustError::InvalidConfig("invalid api_url".into()))?; + if !self.skip_login && self.api_key.is_none() { + return Err(BraintrustError::InvalidConfig( + "API key required: set BRAINTRUST_API_KEY or call .api_key()".into(), + )); + } + Ok(()) + } + + pub(crate) async fn initialize(self, execution: ExecutionBinding) -> Result { let app_url_str = self.app_url.unwrap_or_else(|| DEFAULT_APP_URL.into()); let api_url_str = self.api_url.unwrap_or_else(|| DEFAULT_API_URL.into()); diff --git a/tests/cli_initialization.rs b/tests/cli_initialization.rs new file mode 100644 index 0000000..1e70a16 --- /dev/null +++ b/tests/cli_initialization.rs @@ -0,0 +1,360 @@ +use braintrust_sdk_rust::{eval::*, BraintrustClient, ExecutionError}; +use std::{future::IntoFuture, io::Read, net::TcpListener, sync::Arc, time::Duration}; + +fn command(mode: &str) -> std::process::Command { + let mut command = std::process::Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", "initialization_child", "--nocapture"]) + .env("BT_INIT_CHILD", mode) + .env("BT_EVAL_LOCAL", "1") + .env("BT_EVAL_NO_SEND_LOGS", "0") + .env("BT_EVAL_LIST", "0") + .env("BT_EVAL_JSONL", "1") + .env("BT_EVAL_TERMINATE_ON_FAILURE", "0") + .env_remove("BT_EVAL_FILTER_PARSED") + .env_remove("BT_EVAL_SSE_SOCK") + .env_remove("BT_EVAL_SSE_ADDR") + .env_remove("BRAINTRUST_API_KEY") + .env_remove("BRAINTRUST_ORG_NAME"); + command +} +fn check(command: &mut std::process::Command) { + let output = command.output().unwrap(); + assert!( + output.status.success(), + "stdout: {}\nstderr: {}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} +fn options() -> EvalOpts { + EvalOpts::builder() + .name("init".into()) + .source( + vec![Case { + input: 1, + ..Default::default() + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + vec![], + ) + .unwrap(), + ) + .quiet(false) + .build() +} +#[test] +fn initialization_child() { + let Ok(mode) = std::env::var("BT_INIT_CHILD") else { + return; + }; + match mode.as_str() { + "filters-json" | "filters-pattern" => { + let token = BtEvalRunner::from_env(); + assert!(token.is_err()); + let Err(error) = futures::executor::block_on(token.into_future()) else { + panic!("invalid filters must reject") + }; + if mode == "filters-json" { + assert!(matches!(error, BtInitError::Filters(_))); + } else { + assert!(matches!(error, BtInitError::FilterPattern { index: 0, .. })); + } + assert!(std::error::Error::source(&error).is_some()); + } + "list" => { + let token = BtEvalRunner::from_env(); + assert!(token.is_ok()); + let mut runner = futures::executor::block_on(token.into_future()).unwrap(); + assert!(runner.is_list_mode()); + drop(runner.eval("project", options())); + assert!(futures::executor::block_on(runner.finish().into_future()).unwrap()); + } + "sse-error" | "unix-error" => { + let token = BtEvalRunner::from_env(); + assert!(token.is_ok()); + let Err(error) = futures::executor::block_on(token.into_future()) else { + panic!("failed SSE must not fall back") + }; + match error { + BtInitError::Sse { transport, source } => { + assert_eq!( + transport, + if mode == "unix-error" { + "Unix socket" + } else { + "TCP" + } + ); + assert!(matches!( + source.kind(), + std::io::ErrorKind::NotFound | std::io::ErrorKind::ConnectionRefused + )); + } + other => panic!("unexpected error: {other:?}"), + } + } + "snapshot" | "lost" => { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let client = runtime + .block_on( + BraintrustClient::builder() + .runtime(Some(runtime.handle().clone())) + .skip_login(true) + .build() + .into_future(), + ) + .unwrap(); + let token = BtEvalRunner::with_client(Arc::new(client)); + assert!(token.is_ok()); + if mode == "lost" { + drop(runtime); + let Err(error) = futures::executor::block_on(token.into_future()) else { + panic!("runtime loss must fail") + }; + assert!(matches!( + error, + BtInitError::Execution(ExecutionError::WorkerLost) + )); + } else { + // No initialization task can run until this runtime is driven. + std::env::set_var("BT_EVAL_LIST", "1"); + std::env::set_var("BT_EVAL_FILTER_PARSED", "invalid"); + std::env::set_var("BT_EVAL_SSE_ADDR", "127.0.0.1:9"); + std::env::set_var("BT_EVAL_LOCAL", "0"); + runtime.block_on(async { + let mut runner = token.await.unwrap(); + assert!(!runner.is_list_mode()); + assert!(runner.no_send_logs()); + runner.eval("project", options()).await.unwrap(); + assert!(runner.finish().await.unwrap()); + }); + } + } + "record" => { + let token = BtEvalRunner::from_env(); + assert!(token.is_ok()); + futures::executor::block_on(async { + let mut runner = token.await.unwrap(); + assert!(!runner.no_send_logs()); + runner.eval("project", options()).await.unwrap(); + assert!(runner.finish().await.unwrap()); + }); + } + "drop" => { + let client = futures::executor::block_on( + BraintrustClient::builder() + .skip_login(true) + .build() + .into_future(), + ) + .unwrap(); + let token = BtEvalRunner::with_client(Arc::new(client.clone())); + assert!(token.is_ok()); + drop(token); + futures::executor::block_on(client.shutdown(Duration::from_secs(5)).into_future()) + .unwrap(); + } + other => panic!("unknown mode {other}"), + } +} +#[test] +fn invalid_filters_fail_before_client_or_socket_setup() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + for (mode, filters) in [ + ("filters-json", "{"), + ("filters-json", r#"[{"path":[1],"pattern":"a"}]"#), + ("filters-pattern", r#"[{"path":["name"],"pattern":"("}]"#), + ] { + check( + command(mode) + .env("BT_EVAL_LOCAL", "0") + .env("BT_EVAL_FILTER_PARSED", filters) + .env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + ), + ); + } + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); +} +#[test] +fn list_initialization_ignores_sse_and_needs_no_credentials() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + check( + command("list") + .env("BT_EVAL_LIST", "1") + .env("BT_EVAL_LOCAL", "0") + .env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + ), + ); + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); +} +#[test] +fn constructor_configuration_is_captured_before_first_poll() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + check(command("snapshot").env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + )); + let (mut stream, _) = listener + .accept() + .expect("captured SSE endpoint must be used"); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + let mut text = String::new(); + stream.read_to_string(&mut text).unwrap(); + assert!(text.contains("event: start\n")); + assert!(text.contains("event: summary\n")); + assert!(!text.contains("event: error\n")); +} +#[test] +fn dropped_constructor_observer_still_connects_and_shutdown_joins_cleanup() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + check(command("drop").env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + )); + let (mut stream, _) = listener + .accept() + .expect("accepted initialization must survive observer drop"); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + let mut bytes = Vec::new(); + stream.read_to_end(&mut bytes).unwrap(); + assert!( + bytes.is_empty(), + "discarded runner closes its socket without running evals" + ); +} +#[test] +fn external_runtime_loss_before_first_poll_does_not_connect_elsewhere() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + check(command("lost").env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + )); + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); +} +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sse_failure_precedes_recording_client_login() { + let server = wiremock::MockServer::start().await; + let closed = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = closed.local_addr().unwrap(); + drop(closed); + check( + command("sse-error") + .env("BT_EVAL_LOCAL", "0") + .env("BRAINTRUST_API_KEY", "test-key") + .env("BRAINTRUST_APP_URL", server.uri()) + .env("BRAINTRUST_API_URL", server.uri()) + .env("BT_EVAL_SSE_ADDR", address.to_string()), + ); + assert!(server.received_requests().await.unwrap().is_empty()); +} +#[cfg(unix)] +#[test] +fn explicit_unix_failure_does_not_fall_back_to_tcp() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let missing = + std::path::PathBuf::from("/tmp").join(format!("bt-init-{}.sock", uuid::Uuid::new_v4())); + check(command("unix-error").env("BT_EVAL_SSE_SOCK", missing).env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + )); + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn recording_constructor_composes_login_and_evaluation_with_the_same_client() { + use serde_json::json; + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + for (endpoint, body) in [ + ( + "/api/apikey/login", + json!({"org_info":[{"id":"org","name":"org"}]}), + ), + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}}), + ), + ("/logs3", 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(json!({}))) + .mount(&server) + .await; + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + check( + command("record") + .env("BT_EVAL_LOCAL", "0") + .env("BRAINTRUST_API_KEY", "test-key") + .env("BRAINTRUST_APP_URL", server.uri()) + .env("BRAINTRUST_API_URL", server.uri()) + .env( + "BT_EVAL_SSE_ADDR", + listener.local_addr().unwrap().to_string(), + ), + ); + let (mut stream, _) = listener.accept().unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + let mut text = String::new(); + stream.read_to_string(&mut text).unwrap(); + assert!(text.contains("event: start\n")); + assert!(text.contains("\"experiment_id\":\"experiment\"")); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|r| r.url.path() == "/api/apikey/login") + .count(), + 1 + ); + assert_eq!( + requests + .iter() + .filter(|r| r.url.path() == "/api/experiment/register") + .count(), + 1 + ); + assert!(requests.iter().any(|r| r.url.path() == "/logs3")); +} diff --git a/tests/eval_callables.rs b/tests/eval_callables.rs index 5ce314e..1fe35ac 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, EvalDefinition, EvalOpts, Evaluator, - ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, + run, scorer, task, Case, CaseSource, EvalConfigError, EvalDefinition, EvalOpts, + EvalRequest, ExactMatch, Score, Scorer, ScorerArgs, Scores, SimpleFnTask, Task, TaskHooks, }, BraintrustClient, }; @@ -156,7 +156,7 @@ async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_acc .build() .await .unwrap(); - let evaluator = Evaluator::new(client); + let execute = |options| run(EvalRequest::recorded(options, client.clone())); let accessed = Arc::new(AtomicBool::new(false)); for (concurrency, names, expected) in [ (0, vec!["one"], EvalConfigError::ZeroConcurrency), @@ -199,7 +199,7 @@ async fn invalid_concurrency_and_names_fail_before_login_recording_or_source_acc .build(); assert_eq!(opts.validate(), Err(expected.clone())); assert!( - matches!(evaluator.run(opts).await, Err(error) if matches!(&error.failures[0], braintrust_sdk_rust::eval::EvalRunFailure::Config(actual) if actual == &expected) && error.partial.is_none()) + matches!(execute(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)); diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 765e937..87ac862 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -48,7 +48,7 @@ async fn recording_client() -> (BraintrustClient, wiremock::MockServer) { #[tokio::test] async fn test_basic_eval_flow() { let (client, server) = recording_client().await; - let evaluator = Evaluator::new(client); + let execute = |options| run(EvalRequest::recorded(options, client.clone())); let cases = vec![ Case { @@ -80,17 +80,16 @@ async fn test_basic_eval_flow() { }, ))]; - let result = evaluator - .run( - EvalOpts::builder() - .name("test-basic-eval".to_string()) - .source(cases.into()) - .definition(EvalDefinition::new(Box::new(callback), scorers).unwrap()) - .quiet(true) - .build(), - ) - .await - .expect("Evaluation failed"); + let result = execute( + EvalOpts::builder() + .name("test-basic-eval".to_string()) + .source(cases.into()) + .definition(EvalDefinition::new(Box::new(callback), scorers).unwrap()) + .quiet(true) + .build(), + ) + .await + .expect("Evaluation failed"); assert_eq!(result.total_cases, 2); assert_eq!(result.successful_cases, 2); @@ -203,38 +202,37 @@ fn test_eval_opts_builder() { #[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 execute = |options| run(EvalRequest::recorded(options, client.clone())); let scorers: Vec>> = vec![ Box::new(ExactMatch::new()), Box::new(SyncFnScorer::new("broken", |_| { anyhow::bail!("scoring failed") })), ]; - let result = evaluator - .run( - EvalOpts::builder() - .name("failed-scorer".into()) - .source( - vec![Case { - input: "a".to_owned(), - expected: Some("A".to_owned()), - id: Some("source-case".into()), - ..Default::default() - }] - .into(), - ) - .definition( - EvalDefinition::new( - Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), - scorers, - ) - .unwrap(), + let result = execute( + EvalOpts::builder() + .name("failed-scorer".into()) + .source( + vec![Case { + input: "a".to_owned(), + expected: Some("A".to_owned()), + id: Some("source-case".into()), + ..Default::default() + }] + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &String| Ok(input.to_uppercase()))), + scorers, ) - .quiet(true) - .build(), - ) - .await - .unwrap(); + .unwrap(), + ) + .quiet(true) + .build(), + ) + .await + .unwrap(); assert_eq!(result.successful_cases, 0); assert_eq!(result.failed_cases, 1); assert_eq!(result.score_stats[0].name, "exact_match"); @@ -356,7 +354,7 @@ async fn source_failure_drains_active_cases_preserves_provenance_and_finishes_re .max_concurrency(3) .quiet(true) .build(); - let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + let run = tokio::spawn(async move { run(EvalRequest::recorded(opts, client)).await }); tokio::time::timeout(std::time::Duration::from_secs(3), failed.notified()) .await .expect("pending source reads must allow active task progress"); @@ -445,7 +443,7 @@ async fn failure_before_the_first_case_returns_the_resolved_target_without_trace ) .quiet(true) .build(); - let error = Evaluator::new(client).run(opts).await.unwrap_err(); + let error = run(EvalRequest::recorded(opts, client)).await.unwrap_err(); assert!( matches!(&error.failures[0],EvalRunFailure::Source {case_index:0,source} if source.to_string()=="source unavailable") ); @@ -517,7 +515,7 @@ async fn recording_failure_stops_new_cases_but_drains_already_scheduled_cases() .max_concurrency(2) .quiet(true) .build(); - let run = tokio::spawn(async move { Evaluator::new(client).run(opts).await }); + let run = tokio::spawn(async move { run(EvalRequest::recorded(opts, client)).await }); tokio::time::timeout(std::time::Duration::from_secs(3), entered.notified()) .await .unwrap(); @@ -572,7 +570,7 @@ async fn reusable_definition_keeps_recorded_runs_and_sources_independent() { Arc, }; let (client, _server) = recording_client().await; - let evaluator = Evaluator::new(client); + let execute = |options| run(EvalRequest::recorded(options, client.clone())); let calls = Arc::new(AtomicUsize::new(0)); let definition = EvalDefinition::new( Box::new(task({ @@ -602,8 +600,8 @@ async fn reusable_definition_keeps_recorded_runs_and_sources_independent() { .build() }; let (first, second) = tokio::join!( - evaluator.run(options("first", 2, 4)), - evaluator.run(options("second", 3, 100)), + execute(options("first", 2, 4)), + execute(options("second", 3, 100)), ); let first = first.unwrap(); let second = second.unwrap(); @@ -616,7 +614,7 @@ async fn reusable_definition_keeps_recorded_runs_and_sources_independent() { 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(); + let later = execute(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); @@ -657,7 +655,9 @@ async fn serialization_failures_keep_outputs_and_finalize_recording() { ) .quiet(true) .build(); - let error = Evaluator::new(client).run(options).await.unwrap_err(); + let error = run(EvalRequest::recorded(options, client)) + .await + .unwrap_err(); let fields = error .failures .iter()