wip

Author: realark

README.md

@@ -82,6 +82,8 @@ var eval = braintrust.<String, String>evalBuilder()
                                 (expected, result) -> expected.equals(result) ? 1.0 : 0.0))
                 .build();
 var result = eval.run();
+// TODO: document the concurrency contract - cases run 10-at-a-time by default;
+// task/scorers must be thread-safe. See Eval.Builder#maxConcurrency and #executor.
 System.out.println("\n\n" + result.createReportString());
 ```
 

braintrust-sdk/src/main/java/dev/braintrust/eval/Classifier.java

@@ -14,6 +14,8 @@
  * @param <INPUT> type of the input data
  * @param <OUTPUT> type of the output data
  */
+// TODO: document the concurrency contract - Eval invokes this from multiple threads
+// (see Eval.Builder#maxConcurrency), so implementations must be thread-safe.
 public interface Classifier<INPUT, OUTPUT> {
     String INVALID_CLASSIFICATION_MESSAGE =
             "When returning structured classifier results, each classification must be a non-empty"

braintrust-sdk/src/main/java/dev/braintrust/eval/Dataset.java

@@ -32,6 +32,8 @@ default void forEach(Consumer<DatasetCase<INPUT, OUTPUT>> consumer) {
         }
     }
 
+    // TODO: document the concurrency contract - Eval drains a cursor from a single coordinator
+    // thread even when evaluating cases concurrently, so implementations need no locking.
     @NotThreadSafe
     interface Cursor<CASE> extends AutoCloseable {
         /**

braintrust-sdk/src/main/java/dev/braintrust/eval/Eval.java

@@ -17,7 +17,13 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
 import java.util.*;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -26,11 +32,20 @@
 /**
  * An evaluation framework for testing AI models.
  *
+ * <p><b>Cases are evaluated concurrently.</b> By default up to {@value #DEFAULT_MAX_CONCURRENCY}
+ * cases run at once, so the {@link Task}, {@link Scorer}s and {@link Classifier}s supplied to an
+ * eval must be safe to invoke from multiple threads. Use {@link Builder#maxConcurrency(int)} to
+ * change the bound, or {@code maxConcurrency(1)} to evaluate cases one at a time. Use {@link
+ * Builder#executor(Executor)} to supply the threads the cases run on.
+ *
  * @param <INPUT> The type of input data for the evaluation
  * @param <OUTPUT> The type of output produced by the task
  */
 @Slf4j
 public final class Eval<INPUT, OUTPUT> {
+    /** Default number of eval cases evaluated concurrently. */
+    public static final int DEFAULT_MAX_CONCURRENCY = 10;
+
     private static final AttributeKey<String> PARENT =
             AttributeKey.stringKey(BraintrustTracing.PARENT_KEY);
     private final @Nonnull String experimentName;
@@ -47,6 +62,8 @@ public final class Eval<INPUT, OUTPUT> {
     private final @Nonnull Map<String, Object> metadata;
     private final @Nonnull Parameters parameters;
     private final boolean ensureNew;
+    private final int maxConcurrency;
+    private final @Nullable Executor executor;
 
     private Eval(Builder<INPUT, OUTPUT> builder) {
         this.experimentName = builder.experimentName;
@@ -65,11 +82,39 @@ private Eval(Builder<INPUT, OUTPUT> builder) {
         this.metadata = Map.copyOf(builder.metadata);
         this.parameters = builder.buildParameters();
         this.ensureNew = builder.ensureNew;
+        this.maxConcurrency = builder.maxConcurrency;
+        this.executor = builder.executor;
     }
 
-    /** Runs the evaluation and returns results. */
+    /**
+     * Runs the evaluation to completion and returns the results.
+     *
+     * <p>Cases are evaluated concurrently; see {@link Builder#maxConcurrency(int)}.
+     */
     public EvalResult run() {
-        try (var cursor = dataset.openCursor()) {
+        var result = start();
+        result.awaitCompletion();
+        return result;
+    }
+
+    /**
+     * Creates the experiment and begins evaluating cases in the background, returning as soon as
+     * the experiment exists.
+     *
+     * <p>The returned {@link EvalResult} carries the experiment id, name and url immediately — so
+     * callers can surface the link right away — while its cases are still being evaluated. Use
+     * {@link EvalResult#isDone()} and {@link EvalResult#awaitCompletion()} to observe the run.
+     *
+     * <p>Errors raised while creating the experiment are thrown from this method. Errors raised
+     * once the run is underway are reported through {@link EvalResult#awaitCompletion()}.
+     */
+    public EvalResult start() {
+        var state = new EvalRunState();
+        var cursor = dataset.openCursor();
+        final EvalResult result;
+        final ExecutorService ownedExecutor;
+        final Executor caseExecutor;
+        try {
             Optional<String> datasetVersion = Optional.empty();
             Optional<String> datasetId = Optional.empty();
             if (dataset instanceof DatasetBrainstoreImpl<INPUT, OUTPUT>) {
@@ -94,8 +139,6 @@ public EvalResult run() {
 
             var experiment = new ExperimentsApi(client).postExperiment(createExperiment);
 
-            cursor.forEach(datasetCase -> evalOne(experiment.getId().toString(), datasetCase));
-
             // Use the experiment's actual name from the response: with ensure_new the backend may
             // dedupe a conflicting name (e.g. "foo" -> "foo-2f8ca776"), and the URL must point at
             // the real, created experiment.
@@ -109,11 +152,143 @@ public EvalResult run() {
                                                     project.getName())
                                             .toASCIIString(),
                                     resolvedName);
-            return new EvalResult(experiment.getId().toString(), resolvedName, experimentUrl);
+            result =
+                    new EvalResult(
+                            experiment.getId().toString(), resolvedName, experimentUrl, state);
+
+            if (executor != null) {
+                ownedExecutor = null;
+                caseExecutor = executor;
+            } else {
+                ownedExecutor = createDefaultExecutor();
+                caseExecutor = ownedExecutor;
+            }
+        } catch (Throwable t) {
+            cursor.close();
+            throw t;
+        }
+
+        // Each case re-establishes the context that was current when the run was started. The eval
+        // span itself is created with setNoParent(), so this carries baggage rather than parentage.
+        var callerContext = Context.current();
+        var experimentId = Objects.requireNonNull(result.getExperimentId());
+        var coordinator =
+                new Thread(
+                        () ->
+                                evalAllCases(
+                                        cursor,
+                                        caseExecutor,
+                                        ownedExecutor,
+                                        experimentId,
+                                        callerContext,
+                                        state),
+                        "braintrust-eval-coordinator");
+        coordinator.setDaemon(true);
+        coordinator.start();
+        return result;
+    }
+
+    /**
+     * Drains the dataset cursor on this (coordinator) thread, submitting each case to {@code
+     * caseExecutor}. A semaphore bounds the number of cases in flight so the whole dataset is never
+     * materialized in memory, and so the executor's queue can't grow without limit.
+     *
+     * <p>Runs on the coordinator thread, never on a worker: waiting for cases to finish from inside
+     * the pool that runs them would deadlock once the pool is saturated.
+     */
+    private void evalAllCases(
+            Dataset.Cursor<DatasetCase<INPUT, OUTPUT>> cursor,
+            Executor caseExecutor,
+            @Nullable ExecutorService ownedExecutor,
+            String experimentId,
+            Context callerContext,
+            EvalRunState state) {
+        var inFlight = new Semaphore(maxConcurrency);
+        Throwable fatal = null;
+        try (cursor) {
+            for (var next = cursor.next(); next.isPresent(); next = cursor.next()) {
+                var datasetCase = next.get();
+                inFlight.acquire();
+                try {
+                    caseExecutor.execute(
+                            () -> {
+                                try (var unused = callerContext.makeCurrent()) {
+                                    if (evalOne(experimentId, datasetCase)) {
+                                        state.caseSucceeded();
+                                    } else {
+                                        state.caseFailed();
+                                    }
+                                } catch (Throwable t) {
+                                    // Contain the failure to this case: one bad case must not
+                                    // abort the rest of the run.
+                                    state.caseFailed();
+                                    log.warn(
+                                            "Eval case failed for input: {}",
+                                            datasetCase.input(),
+                                            t);
+                                } finally {
+                                    inFlight.release();
+                                }
+                            });
+                } catch (RuntimeException e) {
+                    // e.g. RejectedExecutionException from a caller-supplied executor. Release the
+                    // permit we took so the drain below can't hang.
+                    inFlight.release();
+                    throw e;
+                }
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            fatal = e;
+        } catch (Throwable t) {
+            fatal = t;
+        } finally {
+            // Wait for every in-flight case, including when the drain above aborted.
+            try {
+                inFlight.acquire(maxConcurrency);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                if (fatal == null) {
+                    fatal = e;
+                }
+            }
+            if (ownedExecutor != null) {
+                ownedExecutor.shutdown();
+            }
+            // TODO: force-flush spans here so that completion means "results are in Braintrust".
+            // Needs a blocking flush on the eval's OpenTelemetry instance —
+            // BraintrustTracing.attemptForceFlush is package-private and non-blocking.
+            state.complete(fatal);
         }
     }
 
-    private void evalOne(String experimentId, DatasetCase<INPUT, OUTPUT> datasetCase) {
+    /**
+     * The executor used when the caller did not supply one: a fixed pool of daemon threads sized to
+     * {@link #maxConcurrency}, owned by this run and shut down when it finishes.
+     *
+     * <p>On Java 21+, pass {@code Executors.newVirtualThreadPerTaskExecutor()} to {@link
+     * Builder#executor(Executor)} if you want to run many cases concurrently.
+     */
+    private ExecutorService createDefaultExecutor() {
+        var counter = new AtomicInteger();
+        return Executors.newFixedThreadPool(
+                maxConcurrency,
+                r -> {
+                    var thread =
+                            new Thread(r, "braintrust-eval-worker-" + counter.incrementAndGet());
+                    thread.setDaemon(true);
+                    return thread;
+                });
+    }
+
+    /**
+     * Evaluates a single case. Runs entirely on one thread so that the OpenTelemetry scopes it
+     * opens stay thread-confined.
+     *
+     * @return false if the task threw (scorers fell back to {@link Scorer#scoreForTaskException}),
+     *     true otherwise
+     */
+    private boolean evalOne(String experimentId, DatasetCase<INPUT, OUTPUT> datasetCase) {
         var rootSpan =
                 tracer.spanBuilder("eval") // TODO: allow names for eval cases
                         .setNoParent() // each eval case is its own trace
@@ -164,7 +339,7 @@ private void evalOne(String experimentId, DatasetCase<INPUT, OUTPUT> datasetCase
                     for (var scorer : scorers) {
                         runScoreForTaskException(experimentId, rootSpan, scorer, e, datasetCase);
                     }
-                    return;
+                    return false;
                 }
                 taskSpan.end();
             }
@@ -221,6 +396,7 @@ private void evalOne(String experimentId, DatasetCase<INPUT, OUTPUT> datasetCase
         } finally {
             rootSpan.end();
         }
+        return true;
     }
 
     /**
@@ -406,6 +582,8 @@ public static final class Builder<INPUT, OUTPUT> {
         private @Nonnull List<String> tags = List.of();
         private @Nonnull Map<String, Object> metadata = Map.of();
         private boolean ensureNew = false;
+        private int maxConcurrency = DEFAULT_MAX_CONCURRENCY;
+        private @Nullable Executor executor;
 
         public Eval<INPUT, OUTPUT> build() {
             if (config == null) {
@@ -453,6 +631,40 @@ public Builder<INPUT, OUTPUT> apiClient(BraintrustApiClient apiClient) {
             return apiClient(apiClient.openApiClient());
         }
 
+        /**
+         * Sets the maximum number of eval cases evaluated concurrently. Defaults to {@value
+         * Eval#DEFAULT_MAX_CONCURRENCY}.
+         *
+         * <p>Because cases run concurrently, the {@link Task}, {@link Scorer}s and {@link
+         * Classifier}s must be safe to invoke from multiple threads. Pass {@code 1} to evaluate
+         * cases one at a time.
+         *
+         * <p>This bounds how many cases are in flight at once. If you also supply an {@link
+         * #executor(Executor)} with fewer threads than this, that executor is the real limit and
+         * the remaining cases queue.
+         */
+        public Builder<INPUT, OUTPUT> maxConcurrency(int maxConcurrency) {
+            if (maxConcurrency < 1) {
+                throw new IllegalArgumentException(
+                        "maxConcurrency must be at least 1, got " + maxConcurrency);
+            }
+            this.maxConcurrency = maxConcurrency;
+            return this;
+        }
+
+        /**
+         * Sets the executor that eval cases run on. Defaults to a fixed pool of daemon threads
+         * sized to {@link #maxConcurrency(int)}, created and shut down by the eval.
+         *
+         * <p>An executor supplied here is never shut down by the SDK — the caller owns its
+         * lifecycle. On Java 21+, pass {@code Executors.newVirtualThreadPerTaskExecutor()} if you
+         * want to run many cases concurrently.
+         */
+        public Builder<INPUT, OUTPUT> executor(@Nonnull Executor executor) {
+            this.executor = Objects.requireNonNull(executor);
+            return this;
+        }
+
         public Builder<INPUT, OUTPUT> tracer(Tracer tracer) {
             this.tracer = tracer;
             return this;