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