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 } } }