diff --git a/README.md b/README.md index aa927ea6..9cf2b411 100644 --- a/README.md +++ b/README.md @@ -210,6 +210,8 @@ bt eval foo.eval.ts -- --description "Prod" --shard=1/4 - `bt eval --sample 20 --sample-seed 7 qa.eval.ts` — run a deterministic random sample and clearly label the summary as a non-final smoke run. - If you do not pass a sampling flag, `bt eval` runs the full dataset and marks the summary as final. +Use `--max-concurrency ` to limit how many evaluators run at once. This does not change the concurrency configured inside an individual evaluator. + ## `bt datasets` - `bt datasets` works directly against remote Braintrust datasets — no local `bt sync` artifact flow is required. diff --git a/scripts/eval-runner-impl.ts b/scripts/eval-runner-impl.ts index fe58060d..16755c7f 100644 --- a/scripts/eval-runner-impl.ts +++ b/scripts/eval-runner-impl.ts @@ -121,6 +121,7 @@ type RunnerConfig = { jsonl: boolean; list: boolean; terminateOnFailure: boolean; + maxConcurrency: number | null; autoInstrumentation: boolean; filters: EvalFilter[]; first: number | null; @@ -421,6 +422,7 @@ function readRunnerConfig(): RunnerConfig { jsonl: envFlag("BT_EVAL_JSONL"), list: envFlag("BT_EVAL_LIST"), terminateOnFailure: envFlag("BT_EVAL_TERMINATE_ON_FAILURE"), + maxConcurrency: parsePositiveIntegerEnv("BT_EVAL_MAX_CONCURRENCY"), autoInstrumentation: !envFlag("BT_EVAL_NO_AUTO_INSTRUMENTATION"), filters: parseSerializedFilters(process.env.BT_EVAL_FILTER_PARSED), first: parsePositiveIntegerEnv("BT_EVAL_FIRST"), @@ -2269,8 +2271,8 @@ async function createEvalRunner( sse: SseWriter | null, ): Promise { const braintrust = await loadBraintrust(); - const Eval = braintrust.Eval; - if (typeof Eval !== "function") { + const sdkEval = braintrust.Eval; + if (typeof sdkEval !== "function") { throw new Error("Unable to load Eval() from braintrust package."); } const login = braintrust.login; @@ -2281,6 +2283,32 @@ async function createEvalRunner( const parseParent = loadBraintrustUtilParseParent(); const getState = extractGlobalStateGetter(braintrust); + let availableEvalSlots = config.maxConcurrency ?? 0; + const evalSlotWaiters: Array<() => void> = []; + const withEvalSlot = async (run: () => Promise): Promise => { + if (config.maxConcurrency !== null) { + if (availableEvalSlots > 0) { + availableEvalSlots -= 1; + } else { + await new Promise((resolve) => evalSlotWaiters.push(resolve)); + } + } + + try { + return await run(); + } finally { + if (config.maxConcurrency !== null) { + const next = evalSlotWaiters.shift(); + if (next) { + next(); + } else { + availableEvalSlots += 1; + } + } + } + }; + const Eval: EvalFunction = (...args) => withEvalSlot(() => sdkEval(...args)); + const makeEvalOptions = ( evaluatorName: string, overrides?: EvalOptions, @@ -2313,7 +2341,7 @@ async function createEvalRunner( return mergeEvalOptions(base, overrides); }; - const runEval = async ( + const runEvalUnbounded = async ( projectName: string, evaluator: Record, options?: EvalOptions, @@ -2344,7 +2372,7 @@ async function createEvalRunner( ...evaluator, data: sampledData, }); - const result = await Eval(projectName, wrappedEvaluator, opts); + const result = await sdkEval(projectName, wrappedEvaluator, opts); const summary = attachSamplingSummary(result.summary, config); const failingResults = result.results.filter( (r: { error?: unknown }) => r.error !== undefined, @@ -2362,6 +2390,8 @@ async function createEvalRunner( } return result; }; + const runEval: EvalRunner["runEval"] = (...args) => + withEvalSlot(() => runEvalUnbounded(...args)); const runRegisteredEvals = async (evaluators: EvaluatorEntry[]) => { if (sse) { diff --git a/scripts/eval-runner.py b/scripts/eval-runner.py index ab2a785f..699fb5a3 100755 --- a/scripts/eval-runner.py +++ b/scripts/eval-runner.py @@ -82,6 +82,7 @@ class RunnerConfig: jsonl: bool list_only: bool terminate_on_failure: bool + max_concurrency: int | None num_workers: int | None filters: list[EvalFilter] first: int | None @@ -270,6 +271,7 @@ def read_runner_config() -> RunnerConfig: jsonl=env_flag("BT_EVAL_JSONL"), list_only=env_flag("BT_EVAL_LIST"), terminate_on_failure=env_flag("BT_EVAL_TERMINATE_ON_FAILURE"), + max_concurrency=parse_positive_int_env("BT_EVAL_MAX_CONCURRENCY"), num_workers=num_workers, filters=parse_serialized_filters(os.getenv("BT_EVAL_FILTER_PARSED")), first=parse_positive_int_env("BT_EVAL_FIRST"), @@ -1324,6 +1326,11 @@ async def run_once( sse.send("processing", {"evaluators": len(evaluators)}) progress_mode = run_evaluator_progress_mode() + eval_semaphore = ( + asyncio.Semaphore(config.max_concurrency) + if config.max_concurrency is not None + else None + ) async def run_single_evaluator( idx: int, evaluator_instance: EvaluatorInstance @@ -1350,7 +1357,11 @@ async def run_single_evaluator( # command running multiple evaluators doesn't fail on unrelated params. filtered_params = filter_params_for_evaluator(effective_params, evaluator.parameters) evaluator.parameters = validate_parameters(filtered_params, evaluator.parameters) + acquired_eval_slot = False try: + if eval_semaphore is not None: + await eval_semaphore.acquire() + acquired_eval_slot = True result = await run_evaluator_task( evaluator_instance.evaluator, idx, @@ -1363,6 +1374,9 @@ async def run_single_evaluator( except Exception as exc: err = serialize_error(str(exc), traceback.format_exc()) return evaluator_instance, resolved_reporter, None, err + finally: + if acquired_eval_slot and eval_semaphore is not None: + eval_semaphore.release() return evaluator_instance, resolved_reporter, result, None diff --git a/src/eval.rs b/src/eval.rs index 5500de58..214d65a0 100644 --- a/src/eval.rs +++ b/src/eval.rs @@ -312,6 +312,16 @@ pub struct EvalArgs { )] pub terminate_on_failure: bool, + /// Maximum number of evaluators to run concurrently. + #[arg( + long, + env = "BT_EVAL_MAX_CONCURRENCY", + value_name = "COUNT", + value_parser = parse_positive_usize, + conflicts_with = "dev" + )] + pub max_concurrency: Option, + /// Number of worker threads for Python eval execution. #[arg(long, env = "BT_EVAL_NUM_WORKERS", value_name = "COUNT")] pub num_workers: Option, @@ -454,6 +464,7 @@ enum EvalSamplingMode { struct EvalRunOptions { jsonl: bool, terminate_on_failure: bool, + max_concurrency: Option, num_workers: Option, list: bool, filter: Vec, @@ -506,6 +517,7 @@ pub async fn run(base: BaseArgs, args: EvalArgs) -> Result<()> { let options = EvalRunOptions { jsonl: args.jsonl, terminate_on_failure: args.terminate_on_failure, + max_concurrency: args.max_concurrency, num_workers: args.num_workers, list: args.list, filter: args.filter, @@ -867,6 +879,9 @@ async fn spawn_eval_runner( if options.terminate_on_failure { cmd.env("BT_EVAL_TERMINATE_ON_FAILURE", "1"); } + if let Some(max_concurrency) = options.max_concurrency { + cmd.env("BT_EVAL_MAX_CONCURRENCY", max_concurrency.to_string()); + } if options.list { cmd.env("BT_EVAL_LIST", "1"); } @@ -5098,6 +5113,49 @@ mod tests { } } + #[test] + fn eval_args_parse_max_concurrency_flag() { + let _guard = env_test_lock() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let previous = clear_env_var("BT_EVAL_MAX_CONCURRENCY"); + + let parsed = + EvalArgsHarness::try_parse_from(["bt", "--max-concurrency", "2", "sample.eval.ts"]) + .expect("max-concurrency flag should parse"); + assert_eq!(parsed.eval.max_concurrency, Some(2)); + + let err = EvalArgsHarness::try_parse_from(["bt", "--maxConcurrency=3", "sample.eval.ts"]) + .expect_err("camelCase maxConcurrency flag should not parse"); + assert!(err.to_string().contains("unexpected argument")); + + let err = + EvalArgsHarness::try_parse_from(["bt", "--max-concurrency", "0", "sample.eval.ts"]) + .expect_err("zero max concurrency should fail"); + assert!(err.to_string().contains("greater than 0")); + + let err = EvalArgsHarness::try_parse_from([ + "bt", + "--dev", + "--max-concurrency", + "2", + "sample.eval.ts", + ]) + .expect_err("max concurrency with dev mode should fail"); + let message = err.to_string(); + assert!(message.contains("--max-concurrency")); + assert!(message.contains("--dev")); + + set_env_var("BT_EVAL_MAX_CONCURRENCY", "2"); + let err = EvalArgsHarness::try_parse_from(["bt", "--dev", "sample.eval.ts"]) + .expect_err("max concurrency env var with dev mode should fail"); + let message = err.to_string(); + assert!(message.contains("--max-concurrency")); + assert!(message.contains("--dev")); + + restore_env_var("BT_EVAL_MAX_CONCURRENCY", previous); + } + #[test] fn eval_args_from_env_populates_supported_fields() { let _guard = env_test_lock() @@ -5106,6 +5164,7 @@ mod tests { let keys = [ "BT_EVAL_JSONL", "BT_EVAL_TERMINATE_ON_FAILURE", + "BT_EVAL_MAX_CONCURRENCY", "BT_EVAL_NUM_WORKERS", "BT_EVAL_LIST", "BT_EVAL_FILTER", @@ -5123,6 +5182,7 @@ mod tests { keys.iter().map(|key| (*key, clear_env_var(key))).collect(); set_env_var("BT_EVAL_JSONL", "true"); set_env_var("BT_EVAL_TERMINATE_ON_FAILURE", "1"); + set_env_var("BT_EVAL_MAX_CONCURRENCY", "3"); set_env_var("BT_EVAL_NUM_WORKERS", "4"); set_env_var("BT_EVAL_LIST", "yes"); set_env_var("BT_EVAL_FILTER", "metadata.case=smoke.*,metadata.kind=fast"); @@ -5137,6 +5197,7 @@ mod tests { .expect("env vars should parse into eval args"); assert!(parsed.eval.jsonl); assert!(parsed.eval.terminate_on_failure); + assert_eq!(parsed.eval.max_concurrency, Some(3)); assert_eq!(parsed.eval.num_workers, Some(4)); assert!(parsed.eval.list); assert_eq!( diff --git a/tests/eval_fixtures.rs b/tests/eval_fixtures.rs index 62049d28..8e4fe432 100644 --- a/tests/eval_fixtures.rs +++ b/tests/eval_fixtures.rs @@ -926,6 +926,194 @@ fn eval_matrix_param_terminate_on_failure_stops_early() { ); } +fn evaluator_concurrency_stats(path: &Path) -> (usize, usize) { + let contents = fs::read_to_string(path).expect("read evaluator concurrency event log"); + let mut active = BTreeSet::new(); + let mut peak = 0; + let mut starts = 0; + + for line in contents.lines() { + if let Some(name) = line.strip_prefix("start:") { + assert!( + active.insert(name.to_string()), + "evaluator {name} started more than once: {contents}" + ); + starts += 1; + peak = peak.max(active.len()); + } else if let Some(name) = line.strip_prefix("end:") { + assert!( + active.remove(name), + "evaluator {name} ended without being active: {contents}" + ); + } else if !line.trim().is_empty() { + panic!("unexpected evaluator concurrency event {line:?}: {contents}"); + } + } + + assert!( + active.is_empty(), + "evaluators remained active at end of event log: {contents}" + ); + (peak, starts) +} + +#[test] +fn eval_javascript_max_concurrency_limits_evaluators() { + let _guard = test_lock(); + if !command_exists("node") { + if required_runtimes().contains("node") { + panic!("node runtime is required but unavailable for max-concurrency test"); + } + eprintln!( + "Skipping eval_javascript_max_concurrency_limits_evaluators (node not installed)." + ); + return; + } + + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let fixture_dir = root + .join("tests") + .join("evals") + .join("js") + .join("eval-max-concurrency"); + ensure_dependencies(&fixture_dir); + + let out_file = fixture_dir.join(format!( + ".max-concurrency-out-{}.txt", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos() + )); + let output = Command::new(bt_binary_path(&root)) + .args(["eval", "--max-concurrency", "2", "max-concurrency.eval.mjs"]) + .current_dir(&fixture_dir) + .env("BT_EVAL_LOCAL", "1") + .env("BT_MAX_CONCURRENCY_TEST_OUT", &out_file) + .output() + .expect("run JavaScript eval with max concurrency"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "JavaScript max-concurrency eval should succeed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + let stats = evaluator_concurrency_stats(&out_file); + let _ = fs::remove_file(&out_file); + assert_eq!( + stats, + (2, 3), + "expected three evaluators with peak concurrency two.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +#[test] +fn eval_javascript_max_concurrency_limits_sampling() { + let _guard = test_lock(); + if !command_exists("node") { + if required_runtimes().contains("node") { + panic!("node runtime is required but unavailable for max-concurrency sampling test"); + } + eprintln!("Skipping eval_javascript_max_concurrency_limits_sampling (node not installed)."); + return; + } + + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let fixture_dir = root + .join("tests") + .join("evals") + .join("js") + .join("eval-max-concurrency"); + ensure_dependencies(&fixture_dir); + + let out_file = fixture_dir.join(format!( + ".max-concurrency-sampling-out-{}.txt", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos() + )); + let output = Command::new(bt_binary_path(&root)) + .args([ + "eval", + "--max-concurrency", + "2", + "--first", + "1", + "max-concurrency-sampling.eval.mjs", + ]) + .current_dir(&fixture_dir) + .env("BT_EVAL_LOCAL", "1") + .env("BT_MAX_CONCURRENCY_TEST_OUT", &out_file) + .output() + .expect("run JavaScript sampled eval with max concurrency"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "JavaScript sampled max-concurrency eval should succeed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + let stats = evaluator_concurrency_stats(&out_file); + let _ = fs::remove_file(&out_file); + assert_eq!( + stats, + (2, 3), + "expected sampling for three evaluators with peak concurrency two.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +#[test] +fn eval_python_max_concurrency_limits_evaluators() { + let _guard = test_lock(); + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let fixtures_root = root.join("tests").join("evals"); + let fixture_dir = fixtures_root.join("py").join("max_concurrency"); + let python = match ensure_python_env(&fixtures_root.join("py")) { + Some(python) => python, + None => { + if required_runtimes().contains("python") { + panic!("python runtime is required but unavailable for max-concurrency test"); + } + eprintln!( + "Skipping eval_python_max_concurrency_limits_evaluators (python unavailable)." + ); + return; + } + }; + + let out_file = fixture_dir.join(format!( + ".max-concurrency-out-{}.txt", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos() + )); + let output = Command::new(bt_binary_path(&root)) + .args(["eval", "--max-concurrency", "2", "eval_max_concurrency.py"]) + .current_dir(&fixture_dir) + .env("BT_EVAL_LOCAL", "1") + .env("BT_EVAL_PYTHON_RUNNER", &python) + .env("BT_MAX_CONCURRENCY_TEST_OUT", &out_file) + .output() + .expect("run Python eval with max concurrency"); + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success(), + "Python max-concurrency eval should succeed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + let stats = evaluator_concurrency_stats(&out_file); + let _ = fs::remove_file(&out_file); + assert_eq!( + stats, + (2, 3), + "expected three evaluators with peak concurrency two.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + #[test] fn eval_python_callable_list_data_preserves_parallel_scorers() { let _guard = test_lock(); diff --git a/tests/evals/js/eval-max-concurrency/max-concurrency-sampling.eval.mjs b/tests/evals/js/eval-max-concurrency/max-concurrency-sampling.eval.mjs new file mode 100644 index 00000000..28bc4132 --- /dev/null +++ b/tests/evals/js/eval-max-concurrency/max-concurrency-sampling.eval.mjs @@ -0,0 +1,25 @@ +import { appendFileSync } from "node:fs"; +import { Eval } from "braintrust"; + +const outputPath = process.env.BT_MAX_CONCURRENCY_TEST_OUT; + +function registerEvaluator(name) { + Eval(`test-max-concurrency-sampling-${name}`, { + data: async () => { + if (outputPath) { + appendFileSync(outputPath, `start:${name}\n`); + } + await new Promise((resolve) => setTimeout(resolve, 250)); + if (outputPath) { + appendFileSync(outputPath, `end:${name}\n`); + } + return [{ input: name }]; + }, + task: (input) => input, + scores: [], + }); +} + +registerEvaluator("alpha"); +registerEvaluator("beta"); +registerEvaluator("gamma"); diff --git a/tests/evals/js/eval-max-concurrency/max-concurrency.eval.mjs b/tests/evals/js/eval-max-concurrency/max-concurrency.eval.mjs new file mode 100644 index 00000000..3af6d087 --- /dev/null +++ b/tests/evals/js/eval-max-concurrency/max-concurrency.eval.mjs @@ -0,0 +1,25 @@ +import { appendFileSync } from "node:fs"; +import { Eval } from "braintrust"; + +const outputPath = process.env.BT_MAX_CONCURRENCY_TEST_OUT; + +function registerEvaluator(name) { + Eval(`test-max-concurrency-${name}`, { + data: () => [{ input: name }], + task: async (input) => { + if (outputPath) { + appendFileSync(outputPath, `start:${name}\n`); + } + await new Promise((resolve) => setTimeout(resolve, 250)); + if (outputPath) { + appendFileSync(outputPath, `end:${name}\n`); + } + return input; + }, + scores: [], + }); +} + +registerEvaluator("alpha"); +registerEvaluator("beta"); +registerEvaluator("gamma"); diff --git a/tests/evals/js/eval-max-concurrency/package.json b/tests/evals/js/eval-max-concurrency/package.json new file mode 100644 index 00000000..1d183485 --- /dev/null +++ b/tests/evals/js/eval-max-concurrency/package.json @@ -0,0 +1,11 @@ +{ + "name": "bt-eval-max-concurrency", + "private": true, + "type": "module", + "dependencies": { + "braintrust": "^3.25.0" + }, + "devDependencies": { + "tsx": "^4.16.2" + } +} diff --git a/tests/evals/py/max_concurrency/eval_max_concurrency.py b/tests/evals/py/max_concurrency/eval_max_concurrency.py new file mode 100644 index 00000000..fdc49e81 --- /dev/null +++ b/tests/evals/py/max_concurrency/eval_max_concurrency.py @@ -0,0 +1,33 @@ +import os +import time + +from braintrust import Eval + + +OUTPUT_PATH = os.environ.get("BT_MAX_CONCURRENCY_TEST_OUT") + + +def task(name, hooks=None): + if OUTPUT_PATH: + with open(OUTPUT_PATH, "a") as output: + output.write(f"start:{name}\n") + time.sleep(0.25) + if OUTPUT_PATH: + with open(OUTPUT_PATH, "a") as output: + output.write(f"end:{name}\n") + return name + + +def register_evaluator(name): + Eval( + "test-max-concurrency", + data=lambda: [{"input": name}], + task=task, + scores=[], + experiment_name=f"test-{name}", + ) + + +register_evaluator("alpha") +register_evaluator("beta") +register_evaluator("gamma")