Connect evaluations to datasets, experiments, and the CLI

examples/eval_async.rs

@@ -11,8 +11,9 @@ 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.
-    let mut runner = BtEvalRunner::from_env().await?;
+    // Recording mode creates a client; local and list modes need no credentials.
+    let initialization = BtEvalRunner::from_env();
+    let mut runner = initialization.await?;
 
     // Define test cases
     let cases = vec![
@@ -60,17 +61,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 {

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");

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<String>,
 }
 
 /// Inner BTQL query structure.

src/api/experiments.rs

@@ -44,6 +44,12 @@ pub struct ExperimentRegisterRequest {
     pub description: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub base_experiment: Option<String>,
+    /// Dataset associated with this experiment.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dataset_id: Option<String>,
+    /// Explicit dataset snapshot; omitted for a live selection.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub dataset_version: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<Map<String, Value>>,
     #[serde(skip_serializing_if = "Option::is_none")]

src/dataset.rs

@@ -595,6 +595,7 @@ struct DatasetCursor {
     scope: CredentialScope,
     dataset_id: String,
     batch_size: usize,
+    version: Option<String>,
     cursor: Option<String>,
     buffer: VecDeque<DatasetRecord>,
     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<Option<DatasetRecord>> {
@@ -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

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<BraintrustError>);
+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<BraintrustError> for ReadError {
+    fn from(error: BraintrustError) -> Self {
+        Self(Arc::new(error))
+    }
+}
+impl From<crate::ExecutionError> for ReadError {
+    fn from(error: crate::ExecutionError) -> Self {
+        BraintrustError::from(error).into()
+    }
+}
+impl From<crate::EnqueueError> for ReadError {
+    fn from(error: crate::EnqueueError) -> Self {
+        BraintrustError::from(error).into()
+    }
+}
+
+type Reply = oneshot::Sender<std::result::Result<Option<DatasetRecord>, ReadError>>;
+type Driver = futures::future::Shared<BoxFuture<'static, std::result::Result<(), ReadError>>>;
+
+pub(crate) struct ReadSession {
+    requests: Option<mpsc::Sender<Reply>>,
+    ready: Option<oneshot::Receiver<DatasetMetadata>>,
+    metadata: Option<DatasetMetadata>,
+    driver: Driver,
+    done: bool,
+}
+impl ReadSession {
+    pub fn start(
+        dataset: Dataset,
+        batch_size: Option<usize>,
+        version: Option<String>,
+    ) -> std::result::Result<Self, ReadError> {
+        let (requests, mut receiver) = mpsc::channel::<Reply>(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<Option<DatasetRecord>, 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<T>(
+        &self,
+        receiver: oneshot::Receiver<T>,
+    ) -> std::result::Result<T, ReadError> {
+        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()
+}