From 0e9c2c0adb7d02b0878f08073fd5b366f2fcebe6 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 8 Aug 2026 14:45:33 -0700 Subject: [PATCH 01/10] Making runs started from the remote benchmark registry resumable --- cli/README.md | 2 + .../benchmark/v1/benchmark_registry.proto | 13 +- cli/src/benchmark_registry/benchmark.rs | 1 + cli/src/benchmark_registry/client.rs | 5 +- cli/src/benchmark_registry/manifest.rs | 7 + cli/src/benchmark_registry/mod.rs | 2 +- cli/src/benchmark_registry/resolver.rs | 80 +++++- cli/src/commands/resume.rs | 236 ++++++++++++++++-- cli/src/commands/run.rs | 208 ++++++++------- cli/src/db/entities/mod.rs | 1 + cli/src/db/entities/remote_benchmark_run.rs | 30 +++ cli/src/db/entities/workflow_run.rs | 8 + cli/src/db/mod.rs | 7 +- cli/src/db/runs.rs | 91 ++++++- cli/src/db/schema.rs | 3 +- cli/src/db/summaries.rs | 9 + 16 files changed, 581 insertions(+), 122 deletions(-) create mode 100644 cli/src/db/entities/remote_benchmark_run.rs diff --git a/cli/README.md b/cli/README.md index 4183922..446dc80 100644 --- a/cli/README.md +++ b/cli/README.md @@ -79,6 +79,8 @@ When you run `qt run `, the CLI first looks in the local configuratio When `qt` uses the remote benchmark service, downloaded remote definitions and prompt templates are verified and kept in memory for the run. They will not be cached on disk. +If you run an eval from the remote benchmark registry, the CLI will persist the registry endpoint, immutable benchmark version, and manifest hash. If that run needs to be resumed later with `qt resume`, the CLI will re-download that exact benchmark version again and reject it if the manifest hash changed (the registry guarantees that versions are immutable once published). This behavior means that resuming runs that were started from the benchmark registry requires internet access. + ## Architecture The Quantiles CLI, `qt`, keeps execution simple: your code runs locally, while `qt` handles durability and observability. diff --git a/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto b/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto index 038a74f..42ea58e 100644 --- a/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto +++ b/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto @@ -11,15 +11,24 @@ message ResolveBenchmarkRequest { // Stable benchmark name, such as "simpleqa-verified". string benchmark_name = 1; - // Immutable version to resolve. An empty value requests the latest published version. + // Immutable version to resolve. + // + // Versions that were previously published are immutable, so + // any valid version request should return a valid response. + // + // If you pass an empty value for a valid benchmark, you'll + // get the latest published version. string version = 2; } message ResolveBenchmarkResponse { string benchmark_name = 1; + + // Immutable published version resolved by the registry. string version = 2; - // SHA-256 digest of the canonical resource manifest. + // SHA-256 digest of the canonical resource manifest. After a + // version is published, this value must be stable. string manifest_sha256 = 3; repeated BenchmarkResource resources = 4; diff --git a/cli/src/benchmark_registry/benchmark.rs b/cli/src/benchmark_registry/benchmark.rs index cb501ce..d503d20 100644 --- a/cli/src/benchmark_registry/benchmark.rs +++ b/cli/src/benchmark_registry/benchmark.rs @@ -8,6 +8,7 @@ use super::proto::v1::{ResolveBenchmarkResponse, ResourceKind}; use crate::config::{BenchmarkConfig, CustomNoCodeBenchmarkConfig, WorkspaceConfig}; /// A downloaded benchmark ready to execute without materializing its resources on disk. +#[derive(Debug)] pub struct RemoteBenchmark { pub config: CustomNoCodeBenchmarkConfig, pub prompt_template: String, diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs index a9623b9..78a32a5 100644 --- a/cli/src/benchmark_registry/client.rs +++ b/cli/src/benchmark_registry/client.rs @@ -61,6 +61,7 @@ pub(super) fn validate_remote_url(remote_url: &str) -> Result { /// Resolve benchmark metadata from the remote `ConnectRPC` service. pub(super) async fn resolve_manifest( benchmark_name: &str, + version: &str, endpoint: &Url, ) -> Result> { let uri = endpoint @@ -80,7 +81,7 @@ pub(super) async fn resolve_manifest( let client = BenchmarkRegistryServiceClient::new(transport, config); let request = ResolveBenchmarkRequest { benchmark_name: benchmark_name.to_owned(), - version: String::new(), + version: version.to_owned(), ..Default::default() }; @@ -135,7 +136,7 @@ mod tests { let endpoint = validate_remote_url(&server.uri()).unwrap(); assert!( - resolve_manifest("missing", &endpoint) + resolve_manifest("missing", "", &endpoint) .await .unwrap() .is_none() diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs index e88daec..3cc20d4 100644 --- a/cli/src/benchmark_registry/manifest.rs +++ b/cli/src/benchmark_registry/manifest.rs @@ -16,6 +16,7 @@ const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; /// Validate that a response identifies the requested immutable benchmark manifest. pub(super) fn validate_response_identity( benchmark_name: &str, + requested_version: &str, response: &ResolveBenchmarkResponse, ) -> Result<()> { if response.benchmark_name != benchmark_name { @@ -27,6 +28,12 @@ pub(super) fn validate_response_identity( if response.version.is_empty() { bail!("remote benchmark response is missing an immutable version"); } + if !requested_version.is_empty() && response.version != requested_version { + bail!( + "remote benchmark response version `{}` does not match requested version `{requested_version}`", + response.version + ); + } validate_sha256("manifest", &response.manifest_sha256)?; Ok(()) } diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index 6a373f0..b32c557 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -2,7 +2,7 @@ pub use self::benchmark::RemoteBenchmark; pub use self::client::select_remote_url; -pub use self::resolver::resolve_and_download; +pub use self::resolver::{resolve_and_download, resolve_and_download_version}; mod benchmark; mod client; diff --git a/cli/src/benchmark_registry/resolver.rs b/cli/src/benchmark_registry/resolver.rs index 68b6645..d2e0d63 100644 --- a/cli/src/benchmark_registry/resolver.rs +++ b/cli/src/benchmark_registry/resolver.rs @@ -17,13 +17,40 @@ use super::manifest::{validate_resources, validate_response_identity}; pub async fn resolve_and_download( benchmark_name: &str, remote_url: &str, +) -> Result> { + resolve_and_download_inner(benchmark_name, "", remote_url).await +} + +/// Resolve and download one exact immutable benchmark version. +/// +/// `Ok(None)` means the registry no longer exposes the requested version. +/// +/// # Errors +/// +/// Returns an error for an empty version, invalid endpoints, RPC failures, malformed manifests, +/// failed downloads, digest mismatches, invalid UTF-8, or invalid no-code definitions. +pub async fn resolve_and_download_version( + benchmark_name: &str, + version: &str, + remote_url: &str, +) -> Result> { + if version.is_empty() { + anyhow::bail!("remote benchmark version must not be empty"); + } + resolve_and_download_inner(benchmark_name, version, remote_url).await +} + +async fn resolve_and_download_inner( + benchmark_name: &str, + version: &str, + remote_url: &str, ) -> Result> { let endpoint = validate_remote_url(remote_url)?; - let Some(response) = resolve_manifest(benchmark_name, &endpoint).await? else { + let Some(response) = resolve_manifest(benchmark_name, version, &endpoint).await? else { return Ok(None); }; - validate_response_identity(benchmark_name, &response)?; + validate_response_identity(benchmark_name, version, &response)?; let resources = validate_resources(&response.resources, endpoint.scheme() == "http")?; let downloaded = download_resources(&resources).await?; let remote = RemoteBenchmark::new(benchmark_name, response, downloaded)?; @@ -35,9 +62,11 @@ mod tests { use buffa::Message as _; use sha2::{Digest as _, Sha256}; use wiremock::matchers::{method, path}; - use wiremock::{Mock, MockServer, ResponseTemplate}; + use wiremock::{Match, Mock, MockServer, Request, ResponseTemplate}; - use super::super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; + use super::super::proto::v1::{ + BenchmarkResource, ResolveBenchmarkRequest, ResolveBenchmarkResponse, ResourceKind, + }; use super::*; #[tokio::test] @@ -80,6 +109,7 @@ mod tests { .and(path( "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", )) + .and(RequestedVersion("")) .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "application/proto") @@ -114,6 +144,48 @@ mod tests { )); } + #[tokio::test] + async fn rejects_a_response_for_a_different_version() { + let server = MockServer::start().await; + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-test".to_owned(), + version: "v2".to_owned(), + manifest_sha256: "a".repeat(64), + ..Default::default() + }; + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .and(RequestedVersion("v1")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .mount(&server) + .await; + + let error = resolve_and_download_version("remote-test", "v1", &server.uri()) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("does not match requested version") + ); + } + + struct RequestedVersion(&'static str); + + impl Match for RequestedVersion { + fn matches(&self, request: &Request) -> bool { + ResolveBenchmarkRequest::decode_from_slice(&request.body) + .is_ok_and(|request| request.version == self.0) + } + } + fn resource( id: &str, logical_path: &str, diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index cb5b2c0..a16bfbc 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -12,6 +12,12 @@ use qt::metrics_store::MetricsStore; pub(crate) enum ResumePlan { Builtin, CustomCode(Vec), + RemoteBenchmark, +} + +struct RemoteResume { + builtin: Box, + manifest_sha256: String, } /// Plan how to resume a run without doing any IO. @@ -24,6 +30,7 @@ pub(crate) fn plan_resume( workflow_name: &str, run_status: &RunStatus, bench_config: Option<&qt::config::BenchmarkConfig>, + remote_provenance: Option<&qt::db::RemoteBenchmarkProvenance>, ) -> Result { if *run_status == RunStatus::Completed { bail!( @@ -32,6 +39,10 @@ pub(crate) fn plan_resume( ); } + if remote_provenance.is_some() { + return Ok(ResumePlan::RemoteBenchmark); + } + match bench_config { Some(bench) => { bench.validate()?; @@ -60,8 +71,8 @@ pub(crate) fn plan_resume( /// /// # Errors /// -/// Returns an error when the run does not exist, is already completed, the -/// config file is missing or invalid, or execution fails. +/// Returns an error when the run does not exist, is already completed, its local +/// configuration or immutable remote benchmark cannot be restored, or execution fails. pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<()> { let cwd = std::env::current_dir()?; let root = db::resolve_workspace_root(&cwd, true).await?; @@ -79,21 +90,119 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( let workflow_name = run.workflow_name.as_str(); let stored_input = run.input.as_deref(); - - let config = qt::config::load()?; - let bench_config = config.benchmarks.get(workflow_name); - - let plan = plan_resume(workflow_name, &run.status, bench_config)?; + let remote_provenance = db::get_remote_benchmark_provenance(&db, run_id).await?; + + let config = if remote_provenance.is_some() { + None + } else { + Some(qt::config::load()?) + }; + let bench_config = config + .as_ref() + .and_then(|config| config.benchmarks.get(workflow_name)); + + let plan = plan_resume( + workflow_name, + &run.status, + bench_config, + remote_provenance.as_ref(), + )?; + let remote_resume = if matches!(plan, ResumePlan::RemoteBenchmark) { + Some(prepare_remote_resume(workflow_name, stored_input, remote_provenance.as_ref()).await?) + } else { + None + }; db::resume_run(&db, run_id).await?; if !json { println!("Resuming eval run {run_id} ({workflow_name})"); } + execute_resume_plan(ExecuteResumeArgs { + plan, + bench_config, + remote_resume, + db: &db, + metrics_store: &metrics_store, + run_id, + workflow_name, + stored_input, + json, + process_start, + }) + .await +} + +async fn prepare_remote_resume( + workflow_name: &str, + stored_input: Option<&str>, + provenance: Option<&qt::db::RemoteBenchmarkProvenance>, +) -> Result { + let provenance = provenance.context("remote benchmark run is missing registry provenance")?; + if provenance.benchmark_name != workflow_name { + bail!( + "stored remote benchmark name `{}` does not match run workflow `{workflow_name}`", + provenance.benchmark_name + ); + } + let remote = qt::benchmark_registry::resolve_and_download_version( + workflow_name, + &provenance.version, + &provenance.registry_url, + ) + .await? + .with_context(|| { + format!( + "remote benchmark `{workflow_name}` version `{}` is no longer available", + provenance.version + ) + })?; + if remote.manifest_sha256 != provenance.manifest_sha256 { + bail!( + "remote benchmark `{workflow_name}` version `{}` manifest changed: expected `{}`, got `{}`", + provenance.version, + provenance.manifest_sha256, + remote.manifest_sha256 + ); + } + let input = stored_input.context("remote benchmark run is missing stored input")?; + let builtin = super::run::remote_benchmark_builtin(workflow_name, input, remote)?; + Ok(RemoteResume { + builtin, + manifest_sha256: provenance.manifest_sha256.clone(), + }) +} + +struct ExecuteResumeArgs<'a> { + plan: ResumePlan, + bench_config: Option<&'a qt::config::BenchmarkConfig>, + remote_resume: Option, + db: &'a sea_orm::DatabaseConnection, + metrics_store: &'a MetricsStore, + run_id: i64, + workflow_name: &'a str, + stored_input: Option<&'a str>, + json: bool, + process_start: Instant, +} + +async fn execute_resume_plan(args: ExecuteResumeArgs<'_>) -> Result<()> { // TODO: we always re-read the command from the config file on resume. // This means that if the config file is edited between `qt run` and // `qt resume`, the resumed run will use the updated command. It may be // wise to revisit this policy. + let ExecuteResumeArgs { + plan, + bench_config, + remote_resume, + db, + metrics_store, + run_id, + workflow_name, + stored_input, + json, + process_start, + } = args; match plan { ResumePlan::Builtin => { let builtin: Box = match bench_config { @@ -110,8 +219,8 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( _ => None, }; super::run::execute_builtin(super::run::ExecuteBuiltinArgs { - db: &db, - metrics_store: &metrics_store, + db, + metrics_store, run_id, workflow_name, builtin, @@ -122,6 +231,22 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( }) .await } + ResumePlan::RemoteBenchmark => { + let remote_resume = + remote_resume.context("remote benchmark resume was not prepared")?; + super::run::execute_builtin(super::run::ExecuteBuiltinArgs { + db, + metrics_store, + run_id, + workflow_name, + builtin: remote_resume.builtin, + input: stored_input, + json, + process_start, + remote_hash: Some(&remote_resume.manifest_sha256), + }) + .await + } ResumePlan::CustomCode(command) => { super::run::execute_custom( run_id, @@ -141,6 +266,8 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( mod tests { use super::*; + static CWD_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + /// Resuming a run whose status is `completed` must be rejected before any execution /// begins, because a completed run cannot be meaningfully resumed. #[test] @@ -152,7 +279,7 @@ mod tests { model: None, max_workers: None, }); - let err = plan_resume("demo", &RunStatus::Completed, Some(&bench)).unwrap_err(); + let err = plan_resume("demo", &RunStatus::Completed, Some(&bench), None).unwrap_err(); assert!(err.to_string().contains("already completed")); } @@ -167,7 +294,7 @@ mod tests { model: None, max_workers: None, }); - let plan = plan_resume("demo", &RunStatus::Failed, Some(&bench)).unwrap(); + let plan = plan_resume("demo", &RunStatus::Failed, Some(&bench), None).unwrap(); assert!(matches!(plan, ResumePlan::Builtin)); } @@ -175,10 +302,30 @@ mod tests { /// falling back to the hardcoded builtin registry. #[test] fn plan_resume_builtin_without_config() { - let plan = plan_resume("pubmedqa", &RunStatus::Failed, None).unwrap(); + let plan = plan_resume("pubmedqa", &RunStatus::Failed, None, None).unwrap(); assert!(matches!(plan, ResumePlan::Builtin)); } + #[test] + fn plan_resume_prefers_persisted_remote_provenance_over_builtin_name() { + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "simpleqa-verified".to_owned(), + registry_url: "https://api.quantiles.io".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + }; + + let plan = plan_resume( + "simpleqa-verified", + &RunStatus::Failed, + None, + Some(&provenance), + ) + .unwrap(); + + assert!(matches!(plan, ResumePlan::RemoteBenchmark)); + } + /// A `custom_code` benchmark with a config section should plan to resume by re-running /// the command array from the config file with the stored DB input. #[test] @@ -189,7 +336,7 @@ mod tests { command: vec!["python".to_owned(), "eval.py".to_owned()], input: None, }); - let plan = plan_resume("my-eval", &RunStatus::Failed, Some(&bench)).unwrap(); + let plan = plan_resume("my-eval", &RunStatus::Failed, Some(&bench), None).unwrap(); assert!(matches!(&plan, ResumePlan::CustomCode(cmd) if cmd == &["python", "eval.py"])); } @@ -197,7 +344,7 @@ mod tests { /// CLI has no source of truth for what command to execute. #[test] fn plan_resume_custom_code_without_config_errors() { - let err = plan_resume("my-eval", &RunStatus::Failed, None).unwrap_err(); + let err = plan_resume("my-eval", &RunStatus::Failed, None, None).unwrap_err(); assert!(err.to_string().contains("no config section found")); } @@ -205,7 +352,7 @@ mod tests { /// fail immediately with a clear "no config section found" message. #[test] fn plan_resume_unknown_without_config_errors() { - let err = plan_resume("unknown-eval", &RunStatus::Failed, None).unwrap_err(); + let err = plan_resume("unknown-eval", &RunStatus::Failed, None, None).unwrap_err(); assert!(err.to_string().contains("no config section found")); } @@ -219,7 +366,7 @@ mod tests { command: vec![], input: None, }); - let err = plan_resume("my-eval", &RunStatus::Failed, Some(&bench)).unwrap_err(); + let err = plan_resume("my-eval", &RunStatus::Failed, Some(&bench), None).unwrap_err(); assert!(err.to_string().contains("non-empty `command`")); } @@ -249,10 +396,64 @@ mod tests { }, }, )); - let plan = plan_resume("nocode_custom", &RunStatus::Failed, Some(&bench)).unwrap(); + let plan = plan_resume("nocode_custom", &RunStatus::Failed, Some(&bench), None).unwrap(); assert!(matches!(plan, ResumePlan::Builtin)); } + #[tokio::test] + async fn unavailable_remote_version_does_not_reset_run_status() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(404) + .insert_header("content-type", "application/json") + .set_body_raw( + r#"{"code":"not_found","message":"version does not exist"}"#, + "application/json", + ), + ) + .mount(&server) + .await; + + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = + qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "remote-resume-test".to_owned(), + registry_url: server.uri(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + }; + let run_id = + qt::db::create_remote_benchmark_run(&db, "remote-resume-test", Some("{}"), &provenance) + .await + .unwrap(); + qt::db::fail_run(&db, &metrics_store, run_id, "simulated failure") + .await + .unwrap(); + + let _cwd_guard = CWD_LOCK.lock().await; + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + + let error = result.unwrap_err(); + assert!(error.to_string().contains("is no longer available")); + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Failed); + assert_eq!(run.error.as_deref(), Some("simulated failure")); + } + /// A failed `custom_nocode` run can be resumed and re-execute successfully /// through the `CustomNoCodeBuiltin`, verifying that the resume path wires /// the correct builtin and `ExecuteBuiltinArgs`. @@ -261,6 +462,7 @@ mod tests { use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; + let _cwd_guard = CWD_LOCK.lock().await; let server = MockServer::start().await; let tmpdir = tempfile::tempdir().unwrap(); let root = tmpdir.path(); diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index 195cdc9..08fe465 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -35,87 +35,21 @@ pub async fn run( match bench_config { Some(bench) => { bench.validate()?; - match bench { - qt::config::BenchmarkConfig::Builtin(b) => { - let (effective_input, _) = assemble_builtin_input(Some(b), cli_input); - run_builtin_workflow( - workflow_name, - effective_input.as_deref(), - json, - process_start, - ) - .await - } - qt::config::BenchmarkConfig::CustomCode(c) => { - let (merged_input, overridden_keys) = - merge_inputs(c.input.as_ref(), cli_input)?; - let warning = if overridden_keys.is_empty() { - None - } else { - Some(format!( - "--input overrides config input for keys: {}", - overridden_keys.join(", ") - )) - }; - let command = &c.command; - - let cwd = std::env::current_dir()?; - let root = db::resolve_workspace_root(&cwd, true).await?; - let db = db::open_workspace(&root).await?; - let run_id = - db::create_run(&db, workflow_name, merged_input.as_deref()).await?; - - if !json { - println!("Created run {run_id}"); - } - - execute_custom( - run_id, - workflow_name, - merged_input.as_deref(), - command, - json, - process_start, - warning.as_deref(), - ) - .await - } - qt::config::BenchmarkConfig::CustomNoCode(c) => { - let input = assemble_custom_nocode_input(c, cli_input)?; - - let cwd = std::env::current_dir()?; - let root = db::resolve_workspace_root(&cwd, true).await?; - let db = db::open_workspace(&root).await?; - let metrics_store = MetricsStore::new(db::metrics_dir(&root))?; - let run_id = db::create_run(&db, workflow_name, Some(&input)).await?; - - if !json { - println!("Created run {run_id}"); - } - - let builtin = Box::new(qt::builtins::CustomNoCodeBuiltin::new( - workflow_name.to_owned(), - )); - execute_builtin(ExecuteBuiltinArgs { - db: &db, - metrics_store: &metrics_store, - run_id, - workflow_name, - builtin, - input: Some(&input), - json, - process_start, - remote_hash: None, - }) - .await - } - } + run_configured_benchmark(workflow_name, cli_input, json, process_start, bench).await } None => { if let Some(remote) = qt::benchmark_registry::resolve_and_download(workflow_name, &remote_url).await? { - run_remote_benchmark(workflow_name, cli_input, json, process_start, remote).await + run_remote_benchmark( + workflow_name, + cli_input, + json, + process_start, + &remote_url, + remote, + ) + .await } else if builtins::resolve(workflow_name).is_some() { let (effective_input, _) = assemble_builtin_input(None, cli_input); run_builtin_workflow( @@ -132,24 +66,104 @@ pub async fn run( } } +async fn run_configured_benchmark( + workflow_name: &str, + cli_input: Option<&str>, + json: bool, + process_start: Instant, + bench: &qt::config::BenchmarkConfig, +) -> Result<()> { + match bench { + qt::config::BenchmarkConfig::Builtin(config) => { + let (effective_input, _) = assemble_builtin_input(Some(config), cli_input); + run_builtin_workflow( + workflow_name, + effective_input.as_deref(), + json, + process_start, + ) + .await + } + qt::config::BenchmarkConfig::CustomCode(config) => { + let (merged_input, overridden_keys) = merge_inputs(config.input.as_ref(), cli_input)?; + let warning = (!overridden_keys.is_empty()).then(|| { + format!( + "--input overrides config input for keys: {}", + overridden_keys.join(", ") + ) + }); + + let cwd = std::env::current_dir()?; + let root = db::resolve_workspace_root(&cwd, true).await?; + let db = db::open_workspace(&root).await?; + let run_id = db::create_run(&db, workflow_name, merged_input.as_deref()).await?; + if !json { + println!("Created run {run_id}"); + } + + execute_custom( + run_id, + workflow_name, + merged_input.as_deref(), + &config.command, + json, + process_start, + warning.as_deref(), + ) + .await + } + qt::config::BenchmarkConfig::CustomNoCode(config) => { + let input = assemble_custom_nocode_input(config, cli_input)?; + let cwd = std::env::current_dir()?; + let root = db::resolve_workspace_root(&cwd, true).await?; + let db = db::open_workspace(&root).await?; + let metrics_store = MetricsStore::new(db::metrics_dir(&root))?; + let run_id = db::create_run(&db, workflow_name, Some(&input)).await?; + if !json { + println!("Created run {run_id}"); + } + + execute_builtin(ExecuteBuiltinArgs { + db: &db, + metrics_store: &metrics_store, + run_id, + workflow_name, + builtin: Box::new(qt::builtins::CustomNoCodeBuiltin::new( + workflow_name.to_owned(), + )), + input: Some(&input), + json, + process_start, + remote_hash: None, + }) + .await + } + } +} + async fn run_remote_benchmark( workflow_name: &str, cli_input: Option<&str>, json: bool, process_start: Instant, + registry_url: &str, remote: qt::benchmark_registry::RemoteBenchmark, ) -> Result<()> { - let configured_template_path = remote.config.params.prompt_template_file.clone(); let remote_hash = remote.manifest_sha256.clone(); let input = assemble_custom_nocode_input(&remote.config, cli_input)?; - let effective_params: qt::config::CustomNoCodeParams = serde_json::from_str(&input) - .context("failed to parse assembled remote custom_nocode input")?; + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: workflow_name.to_owned(), + registry_url: registry_url.to_owned(), + version: remote.version.clone(), + manifest_sha256: remote_hash.clone(), + }; let cwd = std::env::current_dir()?; let root = db::resolve_workspace_root(&cwd, true).await?; let db = db::open_workspace(&root).await?; let metrics_store = MetricsStore::new(db::metrics_dir(&root))?; - let run_id = db::create_run(&db, workflow_name, Some(&input)).await?; + let run_id = + db::create_remote_benchmark_run(&db, workflow_name, Some(&input), &provenance).await?; if !json { println!( @@ -159,16 +173,7 @@ async fn run_remote_benchmark( println!("Created run {run_id}"); } - let builtin = if effective_params.prompt_template_file == configured_template_path { - Box::new(qt::builtins::CustomNoCodeBuiltin::with_prompt_template( - workflow_name.to_owned(), - remote.prompt_template, - )) - } else { - Box::new(qt::builtins::CustomNoCodeBuiltin::new( - workflow_name.to_owned(), - )) - }; + let builtin = remote_benchmark_builtin(workflow_name, &input, remote)?; execute_builtin(ExecuteBuiltinArgs { db: &db, metrics_store: &metrics_store, @@ -183,6 +188,29 @@ async fn run_remote_benchmark( .await } +pub(super) fn remote_benchmark_builtin( + workflow_name: &str, + input: &str, + remote: qt::benchmark_registry::RemoteBenchmark, +) -> Result> { + let configured_template_path = remote.config.params.prompt_template_file.clone(); + let effective_params: qt::config::CustomNoCodeParams = serde_json::from_str(input) + .context("failed to parse assembled remote custom_nocode input")?; + + if effective_params.prompt_template_file == configured_template_path { + Ok(Box::new( + qt::builtins::CustomNoCodeBuiltin::with_prompt_template( + workflow_name.to_owned(), + remote.prompt_template, + ), + )) + } else { + Ok(Box::new(qt::builtins::CustomNoCodeBuiltin::new( + workflow_name.to_owned(), + ))) + } +} + fn assemble_builtin_input( bench: Option<&qt::config::BuiltinBenchmarkConfig>, cli_input: Option<&str>, diff --git a/cli/src/db/entities/mod.rs b/cli/src/db/entities/mod.rs index 8c6033c..c9433d1 100644 --- a/cli/src/db/entities/mod.rs +++ b/cli/src/db/entities/mod.rs @@ -1,4 +1,5 @@ pub mod event; +pub mod remote_benchmark_run; pub mod step; pub mod workflow; pub mod workflow_run; diff --git a/cli/src/db/entities/remote_benchmark_run.rs b/cli/src/db/entities/remote_benchmark_run.rs new file mode 100644 index 0000000..250c196 --- /dev/null +++ b/cli/src/db/entities/remote_benchmark_run.rs @@ -0,0 +1,30 @@ +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "remote_benchmark_runs")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub run_id: i64, + pub benchmark_name: String, + pub registry_url: String, + pub version: String, + pub manifest_sha256: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::workflow_run::Entity", + from = "Column::RunId", + to = "super::workflow_run::Column::Id" + )] + WorkflowRun, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::WorkflowRun.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/cli/src/db/entities/workflow_run.rs b/cli/src/db/entities/workflow_run.rs index 0334d84..ec318e8 100644 --- a/cli/src/db/entities/workflow_run.rs +++ b/cli/src/db/entities/workflow_run.rs @@ -28,6 +28,8 @@ pub enum Relation { Steps, #[sea_orm(has_many = "super::event::Entity")] Events, + #[sea_orm(has_one = "super::remote_benchmark_run::Entity")] + RemoteBenchmarkRun, } impl Related for Entity { @@ -48,4 +50,10 @@ impl Related for Entity { } } +impl Related for Entity { + fn to() -> RelationDef { + Relation::RemoteBenchmarkRun.def() + } +} + impl ActiveModelBehavior for ActiveModel {} diff --git a/cli/src/db/mod.rs b/cli/src/db/mod.rs index a43b35d..d7609c2 100644 --- a/cli/src/db/mod.rs +++ b/cli/src/db/mod.rs @@ -11,12 +11,13 @@ pub use db_url::{DBUrl, SQLitePathURL}; pub use observability::list_events_for_run; pub use runs::{ - complete_run, create_run, fail_run, get_run, list_runs, resume_run, set_run_input, - set_run_output, + complete_run, create_remote_benchmark_run, create_run, fail_run, + get_remote_benchmark_provenance, get_run, list_runs, resume_run, set_run_input, set_run_output, }; pub use steps::list_steps_for_run; pub use summaries::{ - EventSummary, MetricPointSummary, RunStatus, StepStatus, StepSummary, WorkflowRun, + EventSummary, MetricPointSummary, RemoteBenchmarkProvenance, RunStatus, StepStatus, + StepSummary, WorkflowRun, }; pub use workspace::{ init_workspace, metrics_dir, open_database, open_workspace, resolve_workspace_root, diff --git a/cli/src/db/runs.rs b/cli/src/db/runs.rs index 640d6de..88a495b 100644 --- a/cli/src/db/runs.rs +++ b/cli/src/db/runs.rs @@ -6,8 +6,8 @@ use sea_orm::{ }; use serde::Deserialize; -use crate::db::entities::{event, step, workflow, workflow_run}; -use crate::db::summaries::{RunStatus, WorkflowRun, WorkflowRunSummary}; +use crate::db::entities::{event, remote_benchmark_run, step, workflow, workflow_run}; +use crate::db::summaries::{RemoteBenchmarkProvenance, RunStatus, WorkflowRun, WorkflowRunSummary}; use crate::metrics_store::MetricsStore; use crate::time::now_utc; @@ -44,6 +44,30 @@ pub async fn create_run( db: &DatabaseConnection, workflow_name: &str, input: Option<&str>, +) -> Result { + create_run_with_remote_provenance(db, workflow_name, input, None).await +} + +/// Create a running eval run with immutable remote benchmark provenance. +/// +/// # Errors +/// +/// Returns an error if the run, provenance, or initial event cannot be persisted. +pub async fn create_remote_benchmark_run( + db: &DatabaseConnection, + workflow_name: &str, + input: Option<&str>, + provenance: &RemoteBenchmarkProvenance, +) -> Result { + create_run_with_remote_provenance(db, workflow_name, input, Some(provenance)).await +} + +/// Create a running eval run with optional remote provenance. Used by `create_run` and `create_remote_benchmark_run`. +async fn create_run_with_remote_provenance( + db: &DatabaseConnection, + workflow_name: &str, + input: Option<&str>, + provenance: Option<&RemoteBenchmarkProvenance>, ) -> Result { let tx = db.begin().await?; @@ -79,6 +103,18 @@ pub async fn create_run( let run_id = result.last_insert_id; + if let Some(provenance) = provenance { + remote_benchmark_run::Entity::insert(remote_benchmark_run::ActiveModel { + run_id: Set(run_id), + benchmark_name: Set(provenance.benchmark_name.clone()), + registry_url: Set(provenance.registry_url.clone()), + version: Set(provenance.version.clone()), + manifest_sha256: Set(provenance.manifest_sha256.clone()), + }) + .exec(&tx) + .await?; + } + event::Entity::insert(event::ActiveModel { run_id: Set(run_id), event_type: Set("run.started".to_owned()), @@ -94,6 +130,26 @@ pub async fn create_run( Ok(run_id) } +/// Fetch immutable registry provenance for a run, if it was resolved remotely. +/// +/// # Errors +/// +/// Returns an error when the provenance cannot be read. +pub async fn get_remote_benchmark_provenance( + db: &DatabaseConnection, + run_id: i64, +) -> Result> { + Ok(remote_benchmark_run::Entity::find_by_id(run_id) + .one(db) + .await? + .map(|provenance| RemoteBenchmarkProvenance { + benchmark_name: provenance.benchmark_name, + registry_url: provenance.registry_url, + version: provenance.version, + manifest_sha256: provenance.manifest_sha256, + })) +} + /// Update an eval run's output without changing status or emitting an event. /// /// # Errors @@ -381,6 +437,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn create_remote_run_persists_registry_provenance() -> Result<()> { + let (db, _store, _tmpdir) = test_db().await?; + let provenance = RemoteBenchmarkProvenance { + benchmark_name: "simpleqa-verified".to_owned(), + registry_url: "https://api.quantiles.io".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + }; + + let run_id = create_remote_benchmark_run( + &db, + "simpleqa-verified", + Some(r#"{"limit":10}"#), + &provenance, + ) + .await?; + + assert_eq!( + get_remote_benchmark_provenance(&db, run_id).await?, + Some(provenance) + ); + assert!( + get_remote_benchmark_provenance(&db, run_id + 1) + .await? + .is_none() + ); + + Ok(()) + } + #[tokio::test] async fn set_run_output_leaves_other_fields_untouched() -> Result<()> { let (db, _store, _tmpdir) = test_db().await?; diff --git a/cli/src/db/schema.rs b/cli/src/db/schema.rs index e5d2f13..87b1f81 100644 --- a/cli/src/db/schema.rs +++ b/cli/src/db/schema.rs @@ -2,7 +2,7 @@ use anyhow::{Context, Result}; use sea_orm::schema::Schema; use sea_orm::{ConnectionTrait, DatabaseConnection, EntityName, Statement}; -use crate::db::entities::{event, step, workflow, workflow_run}; +use crate::db::entities::{event, remote_benchmark_run, step, workflow, workflow_run}; pub(super) async fn apply_schema(db: &DatabaseConnection) -> Result<()> { db.execute(Statement::from_string( @@ -36,6 +36,7 @@ pub(super) async fn apply_schema(db: &DatabaseConnection) -> Result<()> { create_table!(workflow); create_table!(workflow_run); + create_table!(remote_benchmark_run); create_table!(step); create_table!(event); diff --git a/cli/src/db/summaries.rs b/cli/src/db/summaries.rs index 36e3c65..46956e2 100644 --- a/cli/src/db/summaries.rs +++ b/cli/src/db/summaries.rs @@ -24,6 +24,15 @@ pub struct WorkflowRun { pub error: Option, } +/// Immutable registry provenance needed to resume a remotely resolved benchmark. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteBenchmarkProvenance { + pub benchmark_name: String, + pub registry_url: String, + pub version: String, + pub manifest_sha256: String, +} + #[derive(Debug, Clone, serde::Serialize)] pub struct StepSummary { pub id: i64, From 855297cdee17c9e3e3ac32e73953e4708c2d6749 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 8 Aug 2026 14:53:58 -0700 Subject: [PATCH 02/10] progress --- cli/src/benchmark_registry/manifest.rs | 8 ++-- cli/src/benchmark_registry/mod.rs | 2 +- cli/src/benchmark_registry/resolver.rs | 51 +++++++++++--------------- cli/src/commands/resume.rs | 8 +++- cli/src/commands/run.rs | 6 ++- 5 files changed, 38 insertions(+), 37 deletions(-) diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs index 3cc20d4..3808266 100644 --- a/cli/src/benchmark_registry/manifest.rs +++ b/cli/src/benchmark_registry/manifest.rs @@ -16,7 +16,7 @@ const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; /// Validate that a response identifies the requested immutable benchmark manifest. pub(super) fn validate_response_identity( benchmark_name: &str, - requested_version: &str, + requested_version: Option<&str>, response: &ResolveBenchmarkResponse, ) -> Result<()> { if response.benchmark_name != benchmark_name { @@ -28,9 +28,11 @@ pub(super) fn validate_response_identity( if response.version.is_empty() { bail!("remote benchmark response is missing an immutable version"); } - if !requested_version.is_empty() && response.version != requested_version { + if let Some(version) = requested_version + && response.version != version + { bail!( - "remote benchmark response version `{}` does not match requested version `{requested_version}`", + "remote benchmark response version `{}` does not match requested version `{version}`", response.version ); } diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index b32c557..6a373f0 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -2,7 +2,7 @@ pub use self::benchmark::RemoteBenchmark; pub use self::client::select_remote_url; -pub use self::resolver::{resolve_and_download, resolve_and_download_version}; +pub use self::resolver::resolve_and_download; mod benchmark; mod client; diff --git a/cli/src/benchmark_registry/resolver.rs b/cli/src/benchmark_registry/resolver.rs index d2e0d63..3343c32 100644 --- a/cli/src/benchmark_registry/resolver.rs +++ b/cli/src/benchmark_registry/resolver.rs @@ -5,10 +5,14 @@ use super::client::{resolve_manifest, validate_remote_url}; use super::download::download_resources; use super::manifest::{validate_resources, validate_response_identity}; -/// Resolve a benchmark and download all of its resources into memory. +/// Resolve a benchmark, optionally with a version, and download all of its +/// resources into memory. /// -/// `Ok(None)` means the registry returned Connect's `not_found` status. Other transport and -/// service failures are returned to the caller rather than treated as absence. +/// If you pass `None` for `version`, this function returns the latest latest +/// published version for that benchmark. +/// +/// A return value of `Ok(None)` means the registry did not find that benchmark +/// name and/or version. /// /// # Errors /// @@ -16,35 +20,13 @@ use super::manifest::{validate_resources, validate_response_identity}; /// digest mismatches, invalid UTF-8, or invalid no-code benchmark definitions. pub async fn resolve_and_download( benchmark_name: &str, + version: Option<&str>, remote_url: &str, ) -> Result> { - resolve_and_download_inner(benchmark_name, "", remote_url).await -} - -/// Resolve and download one exact immutable benchmark version. -/// -/// `Ok(None)` means the registry no longer exposes the requested version. -/// -/// # Errors -/// -/// Returns an error for an empty version, invalid endpoints, RPC failures, malformed manifests, -/// failed downloads, digest mismatches, invalid UTF-8, or invalid no-code definitions. -pub async fn resolve_and_download_version( - benchmark_name: &str, - version: &str, - remote_url: &str, -) -> Result> { - if version.is_empty() { - anyhow::bail!("remote benchmark version must not be empty"); + if version.is_some_and(str::is_empty) { + anyhow::bail!("remote benchmark version must not be passed as the empty string"); } - resolve_and_download_inner(benchmark_name, version, remote_url).await -} -async fn resolve_and_download_inner( - benchmark_name: &str, - version: &str, - remote_url: &str, -) -> Result> { let endpoint = validate_remote_url(remote_url)?; let Some(response) = resolve_manifest(benchmark_name, version, &endpoint).await? else { return Ok(None); @@ -128,7 +110,7 @@ mod tests { .mount(&server) .await; - let benchmark = resolve_and_download("remote-test", &server.uri()) + let benchmark = resolve_and_download("remote-test", None, &server.uri()) .await .unwrap() .unwrap(); @@ -166,7 +148,7 @@ mod tests { .mount(&server) .await; - let error = resolve_and_download_version("remote-test", "v1", &server.uri()) + let error = resolve_and_download("remote-test", Some("v1"), &server.uri()) .await .unwrap_err(); @@ -177,6 +159,15 @@ mod tests { ); } + #[tokio::test] + async fn rejects_an_empty_explicit_version() { + let error = resolve_and_download("remote-test", Some(""), "https://api.quantiles.io") + .await + .unwrap_err(); + + assert!(error.to_string().contains("must not be empty")); + } + struct RequestedVersion(&'static str); impl Match for RequestedVersion { diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index a16bfbc..6ebafa0 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -15,6 +15,7 @@ pub(crate) enum ResumePlan { RemoteBenchmark, } +/// Executable state reconstructed from immutable remote benchmark provenance. struct RemoteResume { builtin: Box, manifest_sha256: String, @@ -133,6 +134,7 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( .await } +/// Resolves, verifies, and reconstructs a remotely sourced benchmark for resume. async fn prepare_remote_resume( workflow_name: &str, stored_input: Option<&str>, @@ -145,9 +147,9 @@ async fn prepare_remote_resume( provenance.benchmark_name ); } - let remote = qt::benchmark_registry::resolve_and_download_version( + let remote = qt::benchmark_registry::resolve_and_download( workflow_name, - &provenance.version, + Some(&provenance.version), &provenance.registry_url, ) .await? @@ -173,6 +175,7 @@ async fn prepare_remote_resume( }) } +/// Inputs needed to execute a prepared resume plan. struct ExecuteResumeArgs<'a> { plan: ResumePlan, bench_config: Option<&'a qt::config::BenchmarkConfig>, @@ -186,6 +189,7 @@ struct ExecuteResumeArgs<'a> { process_start: Instant, } +/// Executes a prepared resume plan using the existing run and stored input. async fn execute_resume_plan(args: ExecuteResumeArgs<'_>) -> Result<()> { // TODO: we always re-read the command from the config file on resume. // This means that if the config file is edited between `qt run` and diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index 08fe465..a031565 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -39,7 +39,8 @@ pub async fn run( } None => { if let Some(remote) = - qt::benchmark_registry::resolve_and_download(workflow_name, &remote_url).await? + qt::benchmark_registry::resolve_and_download(workflow_name, None, &remote_url) + .await? { run_remote_benchmark( workflow_name, @@ -66,6 +67,7 @@ pub async fn run( } } +/// Runs a benchmark defined in the local configuration. async fn run_configured_benchmark( workflow_name: &str, cli_input: Option<&str>, @@ -141,6 +143,7 @@ async fn run_configured_benchmark( } } +/// Persists and runs a benchmark resolved from the remote registry. async fn run_remote_benchmark( workflow_name: &str, cli_input: Option<&str>, @@ -188,6 +191,7 @@ async fn run_remote_benchmark( .await } +/// Builds an executable no-code workflow from a downloaded remote benchmark. pub(super) fn remote_benchmark_builtin( workflow_name: &str, input: &str, From 3527035470cd3045ce2798c5fbdc52ed8088c489 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 8 Aug 2026 14:54:05 -0700 Subject: [PATCH 03/10] more fixups --- cli/src/benchmark_registry/client.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs index 78a32a5..2d716aa 100644 --- a/cli/src/benchmark_registry/client.rs +++ b/cli/src/benchmark_registry/client.rs @@ -61,7 +61,7 @@ pub(super) fn validate_remote_url(remote_url: &str) -> Result { /// Resolve benchmark metadata from the remote `ConnectRPC` service. pub(super) async fn resolve_manifest( benchmark_name: &str, - version: &str, + version: Option<&str>, endpoint: &Url, ) -> Result> { let uri = endpoint @@ -81,7 +81,9 @@ pub(super) async fn resolve_manifest( let client = BenchmarkRegistryServiceClient::new(transport, config); let request = ResolveBenchmarkRequest { benchmark_name: benchmark_name.to_owned(), - version: version.to_owned(), + // If the version is passed as `None` to this function, send the empty string + // over RPC + version: version.unwrap_or_default().to_owned(), ..Default::default() }; @@ -136,7 +138,7 @@ mod tests { let endpoint = validate_remote_url(&server.uri()).unwrap(); assert!( - resolve_manifest("missing", "", &endpoint) + resolve_manifest("missing", None, &endpoint) .await .unwrap() .is_none() From 7853fab2e515b3f2469af02e08e1522f04d6b0ae Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:10:33 -0700 Subject: [PATCH 04/10] progress --- cli/src/benchmark_registry/resolver.rs | 6 +- cli/src/commands/resume.rs | 223 +++++++++++++++++++++++++ 2 files changed, 228 insertions(+), 1 deletion(-) diff --git a/cli/src/benchmark_registry/resolver.rs b/cli/src/benchmark_registry/resolver.rs index 3343c32..fd0b123 100644 --- a/cli/src/benchmark_registry/resolver.rs +++ b/cli/src/benchmark_registry/resolver.rs @@ -165,7 +165,11 @@ mod tests { .await .unwrap_err(); - assert!(error.to_string().contains("must not be empty")); + assert!( + error + .to_string() + .contains("must not be passed as the empty string") + ); } struct RequestedVersion(&'static str); diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index 6ebafa0..e446500 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -270,6 +270,20 @@ async fn execute_resume_plan(args: ExecuteResumeArgs<'_>) -> Result<()> { mod tests { use super::*; + /// We're re-including generated proto stubs here, rather than increasing the visibility of + /// generated code in proto.rs, since this is just test code and expanding visibility of + /// the proto stubs just so test code can use it isn't a great idea. + #[expect( + clippy::allow_attributes, + clippy::pedantic, + reason = "ConnectRPC and Buffa generated code uses allow attributes" + )] + mod registry_proto { + connectrpc::include_generated!(); + } + + // TODO: Remove this lock by extracting a `resume_in_workspace` helper and a + // path-based config loader so tests can pass a workspace root without changing process CWD. static CWD_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); /// Resuming a run whose status is `completed` must be rejected before any execution @@ -458,6 +472,215 @@ mod tests { assert_eq!(run.error.as_deref(), Some("simulated failure")); } + #[tokio::test] + async fn remote_benchmark_resumes_from_exact_published_version() { + use wiremock::MockServer; + + let _cwd_guard = CWD_LOCK.lock().await; + let server = MockServer::start().await; + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + let cache_dir = root.join("cache"); + let manifest_sha256 = mock_remote_registry(&server).await; + + let original_hf = std::env::var("HF_DATASETS_SERVER").ok(); + let original_cache = std::env::var("QUANTILES_DATASET_CACHE_DIR").ok(); + unsafe { + std::env::set_var("HF_DATASETS_SERVER", server.uri()); + std::env::set_var("QUANTILES_DATASET_CACHE_DIR", &cache_dir); + } + cache_fixture_rows(&cache_dir).await; + + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = + qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let stored_input = serde_json::json!({ + "dataset": { "name": "fixture/qa" }, + "model": "random", + "prompt_template_file": "prompts/qa.txt", + "limit": 2, + "style": { "type": "exact_match", "golden_column": "answer" } + }) + .to_string(); + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "remote-resume-test".to_owned(), + registry_url: server.uri(), + version: "v1".to_owned(), + manifest_sha256, + }; + let run_id = qt::db::create_remote_benchmark_run( + &db, + "remote-resume-test", + Some(&stored_input), + &provenance, + ) + .await + .unwrap(); + qt::db::fail_run(&db, &metrics_store, run_id, "simulated failure") + .await + .unwrap(); + + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + restore_env("HF_DATASETS_SERVER", original_hf); + restore_env("QUANTILES_DATASET_CACHE_DIR", original_cache); + result.unwrap(); + + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Completed); + assert_eq!( + qt::db::list_steps_for_run(&db, run_id).await.unwrap().len(), + 2 + ); + let metrics = metrics_store.list_aggregate_for_run(run_id).await.unwrap(); + assert!( + metrics + .iter() + .any(|metric| metric.metric_name == "accuracy") + ); + } + + async fn mock_remote_registry(server: &wiremock::MockServer) -> String { + use buffa::Message as _; + use sha2::{Digest as _, Sha256}; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, ResponseTemplate}; + + use registry_proto::quantiles::benchmark::v1::{ + BenchmarkResource, ResolveBenchmarkResponse, ResourceKind, + }; + + let definition = br#" +[benchmarks.remote-resume-test] +type = "custom_nocode" +dataset = { name = "fixture/qa" } +model = "random" +prompt_template_file = "prompts/qa.txt" +limit = 2 +style = { type = "exact_match", golden_column = "answer" } +"#; + let prompt = b"{{ row.question }}\nAnswer:"; + let resource = + |id: &str, logical_path: &str, kind: ResourceKind, route: &str, bytes: &[u8]| { + BenchmarkResource { + resource_id: id.to_owned(), + logical_path: logical_path.to_owned(), + kind: kind.into(), + download_url: format!("{}{route}", server.uri()), + sha256: format!("{:x}", Sha256::digest(bytes)), + size_bytes: u64::try_from(bytes.len()).unwrap(), + content_type: "application/octet-stream".to_owned(), + ..Default::default() + } + }; + let manifest_sha256 = "a".repeat(64); + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-resume-test".to_owned(), + version: "v1".to_owned(), + manifest_sha256: manifest_sha256.clone(), + resources: vec![ + resource( + "definition", + "bundle/quantiles.toml", + ResourceKind::Definition, + "/definition", + definition, + ), + resource( + "prompt", + "bundle/prompts/qa.txt", + ResourceKind::PromptTemplate, + "/prompt", + prompt, + ), + ], + ..Default::default() + }; + + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .and(RequestedVersion("v1")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .expect(1) + .mount(server) + .await; + for (route, body) in [ + ("/definition", definition.as_slice()), + ("/prompt", prompt.as_slice()), + ] { + Mock::given(method("GET")) + .and(path(route)) + .respond_with(ResponseTemplate::new(200).set_body_bytes(body)) + .mount(server) + .await; + } + mock_dataset_metadata(server).await; + manifest_sha256 + } + + async fn mock_dataset_metadata(server: &wiremock::MockServer) { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, ResponseTemplate}; + + Mock::given(method("GET")) + .and(path("/splits")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "splits": [{"config": "default", "split": "train"}] + }))) + .mount(server) + .await; + Mock::given(method("GET")) + .and(path("/size")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "size": {"splits": [{"num_rows": 2}]} + }))) + .mount(server) + .await; + } + + async fn cache_fixture_rows(cache_dir: &std::path::Path) { + let cache = qt::dataset::cache::DatasetCache::new(cache_dir.to_owned()); + let rows = vec![ + serde_json::json!({"question": "what is 2+2", "answer": "4"}), + serde_json::json!({"question": "what is 3+3", "answer": "6"}), + ]; + let key = qt::dataset::cache::cache_key("fixture/qa", "default", "train", None); + cache + .write_batch(&cache.batch_path(&key, 0, 2), &rows) + .await + .unwrap(); + } + + fn restore_env(name: &str, value: Option) { + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + + struct RequestedVersion(&'static str); + + impl wiremock::Match for RequestedVersion { + fn matches(&self, request: &wiremock::Request) -> bool { + use buffa::Message as _; + use registry_proto::quantiles::benchmark::v1::ResolveBenchmarkRequest; + + ResolveBenchmarkRequest::decode_from_slice(&request.body) + .is_ok_and(|request| request.version == self.0) + } + } + /// A failed `custom_nocode` run can be resumed and re-execute successfully /// through the `CustomNoCodeBuiltin`, verifying that the resume path wires /// the correct builtin and `ExecuteBuiltinArgs`. From f008ecb65601d21dc7c3f594b3e8d35365f76056 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:19:03 -0700 Subject: [PATCH 05/10] making qt resume --json error in JSON format --- cli/src/main.rs | 26 +++++++++++++++++++++----- cli/tests/resume_json.rs | 30 ++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 5 deletions(-) create mode 100644 cli/tests/resume_json.rs diff --git a/cli/src/main.rs b/cli/src/main.rs index 62deb0d..c06251b 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,12 +1,30 @@ mod cli; mod commands; +use std::process::ExitCode; use std::time::Instant; use anyhow::Result; use clap::Parser; -fn main() -> Result<()> { +fn main() -> ExitCode { + let cli = cli::Cli::parse(); + let json_errors = matches!(&cli.command, Some(cli::Command::Resume { json: true, .. })); + + match try_main(cli) { + Ok(()) => ExitCode::SUCCESS, + Err(error) if json_errors => { + println!("{}", serde_json::json!({ "error": format!("{error:#}") })); + ExitCode::FAILURE + } + Err(error) => { + eprintln!("Error: {error:#}"); + ExitCode::FAILURE + } + } +} + +fn try_main(cli: cli::Cli) -> Result<()> { // `fastembed`'s dependency tree enables another Rustls crypto provider alongside AWS-LC. // Install AWS-LC explicitly, to avoid Rustls-related panics when multiple providers are // enabled. @@ -22,12 +40,10 @@ fn main() -> Result<()> { tokio::runtime::Builder::new_multi_thread() .enable_all() .build()? - .block_on(async_main(process_start)) + .block_on(async_main(cli, process_start)) } -async fn async_main(process_start: Instant) -> Result<()> { - let cli = cli::Cli::parse(); - +async fn async_main(cli: cli::Cli, process_start: Instant) -> Result<()> { if cli.version { println!("{}", cli::VERSION); return Ok(()); diff --git a/cli/tests/resume_json.rs b/cli/tests/resume_json.rs new file mode 100644 index 0000000..18016a0 --- /dev/null +++ b/cli/tests/resume_json.rs @@ -0,0 +1,30 @@ +use assert_cmd::Command; +use predicates::prelude::*; + +#[tokio::test] +async fn completed_run_error_is_json_when_requested() { + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let run_id = qt::db::create_run(&db, "completed-test", None) + .await + .unwrap(); + qt::db::complete_run(&db, &metrics_store, run_id) + .await + .unwrap(); + + let message = format!( + "run {run_id} is already completed; create a new run or resume a running/failed one" + ); + let expected = format!("{}\n", serde_json::json!({ "error": message })); + + Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(root) + .args(["resume", &run_id.to_string(), "--json"]) + .assert() + .failure() + .stdout(predicate::eq(expected)) + .stderr(predicate::str::is_empty()); +} From cc1568db74f247fb42ff27d04fa4fbf72c84a648 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:02:38 -0700 Subject: [PATCH 06/10] adding provenance for prompt template file --- cli/README.md | 2 + cli/src/commands/resume.rs | 78 +++++++++++++++++++- cli/src/commands/run.rs | 80 +++++++++++++++++---- cli/src/db/entities/remote_benchmark_run.rs | 1 + cli/src/db/runs.rs | 3 + cli/src/db/summaries.rs | 2 + 6 files changed, 149 insertions(+), 17 deletions(-) diff --git a/cli/README.md b/cli/README.md index 446dc80..dcfbe36 100644 --- a/cli/README.md +++ b/cli/README.md @@ -81,6 +81,8 @@ When `qt` uses the remote benchmark service, downloaded remote definitions and p If you run an eval from the remote benchmark registry, the CLI will persist the registry endpoint, immutable benchmark version, and manifest hash. If that run needs to be resumed later with `qt resume`, the CLI will re-download that exact benchmark version again and reject it if the manifest hash changed (the registry guarantees that versions are immutable once published). This behavior means that resuming runs that were started from the benchmark registry requires internet access. +When `--input` overrides a remote benchmark's `prompt_template_file` with a local file, `qt` also persists the file's SHA-256 hash. On resume, `qt` reads the local prompt into memory and rejects the resume before changing the run status if its contents no longer match the stored hash. + ## Architecture The Quantiles CLI, `qt`, keeps execution simple: your code runs locally, while `qt` handles durability and observability. diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index e446500..4e3fef1 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -168,9 +168,19 @@ async fn prepare_remote_resume( ); } let input = stored_input.context("remote benchmark run is missing stored input")?; - let builtin = super::run::remote_benchmark_builtin(workflow_name, input, remote)?; + let prepared = super::run::remote_benchmark_builtin(workflow_name, input, remote)?; + if prepared.prompt_template_sha256 != provenance.prompt_template_sha256 { + bail!( + "remote benchmark `{workflow_name}` local prompt template changed: expected SHA-256 `{}`, got `{}`", + provenance + .prompt_template_sha256 + .as_deref() + .unwrap_or("none"), + prepared.prompt_template_sha256.as_deref().unwrap_or("none") + ); + } Ok(RemoteResume { - builtin, + builtin: prepared.builtin, manifest_sha256: provenance.manifest_sha256.clone(), }) } @@ -331,6 +341,7 @@ mod tests { registry_url: "https://api.quantiles.io".to_owned(), version: "v1".to_owned(), manifest_sha256: "a".repeat(64), + prompt_template_sha256: None, }; let plan = plan_resume( @@ -450,6 +461,7 @@ mod tests { registry_url: server.uri(), version: "v1".to_owned(), manifest_sha256: "a".repeat(64), + prompt_template_sha256: None, }; let run_id = qt::db::create_remote_benchmark_run(&db, "remote-resume-test", Some("{}"), &provenance) @@ -472,6 +484,67 @@ mod tests { assert_eq!(run.error.as_deref(), Some("simulated failure")); } + #[tokio::test] + async fn changed_local_prompt_override_does_not_reset_run_status() { + use sha2::{Digest as _, Sha256}; + use wiremock::MockServer; + + let _cwd_guard = CWD_LOCK.lock().await; + let server = MockServer::start().await; + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + let manifest_sha256 = mock_remote_registry(&server).await; + let prompt_path = root.join("local-prompt.txt"); + let original_prompt = "{{ row.question }}\nOriginal answer:"; + std::fs::write(&prompt_path, original_prompt).unwrap(); + + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = + qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let stored_input = serde_json::json!({ + "dataset": { "name": "fixture/qa" }, + "model": "random", + "prompt_template_file": prompt_path, + "limit": 2, + "style": { "type": "exact_match", "golden_column": "answer" } + }) + .to_string(); + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "remote-resume-test".to_owned(), + registry_url: server.uri(), + version: "v1".to_owned(), + manifest_sha256, + prompt_template_sha256: Some(format!( + "{:x}", + Sha256::digest(original_prompt.as_bytes()) + )), + }; + let run_id = qt::db::create_remote_benchmark_run( + &db, + "remote-resume-test", + Some(&stored_input), + &provenance, + ) + .await + .unwrap(); + qt::db::fail_run(&db, &metrics_store, run_id, "simulated failure") + .await + .unwrap(); + std::fs::write(&prompt_path, "{{ row.question }}\nChanged answer:").unwrap(); + + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + + let error = result.unwrap_err(); + assert!(error.to_string().contains("local prompt template changed")); + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Failed); + assert_eq!(run.error.as_deref(), Some("simulated failure")); + } + #[tokio::test] async fn remote_benchmark_resumes_from_exact_published_version() { use wiremock::MockServer; @@ -508,6 +581,7 @@ mod tests { registry_url: server.uri(), version: "v1".to_owned(), manifest_sha256, + prompt_template_sha256: None, }; let run_id = qt::db::create_remote_benchmark_run( &db, diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index a031565..323afb5 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -9,6 +9,7 @@ use anyhow::{Context, Result, bail}; use comfy_table::{Cell, ContentArrangement, Table, presets::NOTHING}; use sea_orm::DatabaseConnection; use serde::Serialize; +use sha2::{Digest as _, Sha256}; use qt::builtins; use qt::client::QuantilesClient; @@ -154,11 +155,13 @@ async fn run_remote_benchmark( ) -> Result<()> { let remote_hash = remote.manifest_sha256.clone(); let input = assemble_custom_nocode_input(&remote.config, cli_input)?; + let prepared = remote_benchmark_builtin(workflow_name, &input, remote)?; let provenance = qt::db::RemoteBenchmarkProvenance { benchmark_name: workflow_name.to_owned(), registry_url: registry_url.to_owned(), - version: remote.version.clone(), + version: prepared.version.clone(), manifest_sha256: remote_hash.clone(), + prompt_template_sha256: prepared.prompt_template_sha256.clone(), }; let cwd = std::env::current_dir()?; @@ -171,18 +174,17 @@ async fn run_remote_benchmark( if !json { println!( "Resolved remote benchmark {workflow_name} version {} ({})", - remote.version, remote.manifest_sha256 + prepared.version, remote_hash ); println!("Created run {run_id}"); } - let builtin = remote_benchmark_builtin(workflow_name, &input, remote)?; execute_builtin(ExecuteBuiltinArgs { db: &db, metrics_store: &metrics_store, run_id, workflow_name, - builtin, + builtin: prepared.builtin, input: Some(&input), json, process_start, @@ -191,28 +193,43 @@ async fn run_remote_benchmark( .await } +/// Executable remote benchmark and any provenance not covered by its registry manifest. +pub(super) struct RemoteBenchmarkBuiltin { + pub builtin: Box, + pub version: String, + pub prompt_template_sha256: Option, +} + /// Builds an executable no-code workflow from a downloaded remote benchmark. pub(super) fn remote_benchmark_builtin( workflow_name: &str, input: &str, remote: qt::benchmark_registry::RemoteBenchmark, -) -> Result> { +) -> Result { let configured_template_path = remote.config.params.prompt_template_file.clone(); let effective_params: qt::config::CustomNoCodeParams = serde_json::from_str(input) .context("failed to parse assembled remote custom_nocode input")?; - if effective_params.prompt_template_file == configured_template_path { - Ok(Box::new( - qt::builtins::CustomNoCodeBuiltin::with_prompt_template( - workflow_name.to_owned(), - remote.prompt_template, - ), - )) + let (prompt_template, prompt_template_sha256) = if effective_params.prompt_template_file + == configured_template_path + { + (remote.prompt_template, None) } else { - Ok(Box::new(qt::builtins::CustomNoCodeBuiltin::new( + let prompt_path = &effective_params.prompt_template_file; + let prompt_template = std::fs::read_to_string(prompt_path) + .with_context(|| format!("failed to read prompt template override `{prompt_path}`"))?; + let sha256 = format!("{:x}", Sha256::digest(prompt_template.as_bytes())); + (prompt_template, Some(sha256)) + }; + + Ok(RemoteBenchmarkBuiltin { + builtin: Box::new(qt::builtins::CustomNoCodeBuiltin::with_prompt_template( workflow_name.to_owned(), - ))) - } + prompt_template, + )), + version: remote.version, + prompt_template_sha256, + }) } fn assemble_builtin_input( @@ -1005,6 +1022,39 @@ mod tests { assert_eq!(parsed["prompt_template_file"], "prompts/other.txt"); } + #[test] + fn remote_benchmark_builtin_hashes_local_prompt_override() { + use sha2::{Digest as _, Sha256}; + + let tmpdir = tempfile::tempdir().unwrap(); + let prompt_path = tmpdir.path().join("override.txt"); + let prompt = "{{ row.question }}\nLocal answer:"; + std::fs::write(&prompt_path, prompt).unwrap(); + let benchmark = custom_nocode_benchmark_for_override_tests(); + let input = super::assemble_custom_nocode_input( + &benchmark, + Some(&format!( + r#"{{"prompt_template_file":{}}}"#, + serde_json::to_string(&prompt_path).unwrap() + )), + ) + .unwrap(); + let remote = qt::benchmark_registry::RemoteBenchmark { + config: benchmark, + prompt_template: "{{ row.question }}\nRegistry answer:".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + }; + + let prepared = super::remote_benchmark_builtin("remote-test", &input, remote).unwrap(); + + assert_eq!(prepared.version, "v1"); + assert_eq!( + prepared.prompt_template_sha256, + Some(format!("{:x}", Sha256::digest(prompt.as_bytes()))) + ); + } + /// Fields outside the intentionally narrow custom no-code override surface should /// fail before a run is created. #[test] diff --git a/cli/src/db/entities/remote_benchmark_run.rs b/cli/src/db/entities/remote_benchmark_run.rs index 250c196..8799d0f 100644 --- a/cli/src/db/entities/remote_benchmark_run.rs +++ b/cli/src/db/entities/remote_benchmark_run.rs @@ -9,6 +9,7 @@ pub struct Model { pub registry_url: String, pub version: String, pub manifest_sha256: String, + pub prompt_template_sha256: Option, } #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] diff --git a/cli/src/db/runs.rs b/cli/src/db/runs.rs index 88a495b..18a6e44 100644 --- a/cli/src/db/runs.rs +++ b/cli/src/db/runs.rs @@ -110,6 +110,7 @@ async fn create_run_with_remote_provenance( registry_url: Set(provenance.registry_url.clone()), version: Set(provenance.version.clone()), manifest_sha256: Set(provenance.manifest_sha256.clone()), + prompt_template_sha256: Set(provenance.prompt_template_sha256.clone()), }) .exec(&tx) .await?; @@ -147,6 +148,7 @@ pub async fn get_remote_benchmark_provenance( registry_url: provenance.registry_url, version: provenance.version, manifest_sha256: provenance.manifest_sha256, + prompt_template_sha256: provenance.prompt_template_sha256, })) } @@ -445,6 +447,7 @@ mod tests { registry_url: "https://api.quantiles.io".to_owned(), version: "v1".to_owned(), manifest_sha256: "a".repeat(64), + prompt_template_sha256: Some("b".repeat(64)), }; let run_id = create_remote_benchmark_run( diff --git a/cli/src/db/summaries.rs b/cli/src/db/summaries.rs index 46956e2..567509b 100644 --- a/cli/src/db/summaries.rs +++ b/cli/src/db/summaries.rs @@ -31,6 +31,8 @@ pub struct RemoteBenchmarkProvenance { pub registry_url: String, pub version: String, pub manifest_sha256: String, + /// SHA-256 of a local `prompt_template_file` override, when one was used. + pub prompt_template_sha256: Option, } #[derive(Debug, Clone, serde::Serialize)] From 1a7fdc904ffb0e72b7161411cde876252db62063 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:18:33 -0700 Subject: [PATCH 07/10] adding newtype for version --- cli/src/benchmark_registry/client.rs | 10 +++- cli/src/benchmark_registry/manifest.rs | 6 +- cli/src/benchmark_registry/mod.rs | 2 + cli/src/benchmark_registry/resolver.rs | 33 ++++------- cli/src/benchmark_registry/version.rs | 23 ++++++++ cli/src/commands/resume.rs | 79 +++++++++++++++++++++++++- 6 files changed, 126 insertions(+), 27 deletions(-) create mode 100644 cli/src/benchmark_registry/version.rs diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs index 2d716aa..03daf17 100644 --- a/cli/src/benchmark_registry/client.rs +++ b/cli/src/benchmark_registry/client.rs @@ -8,6 +8,8 @@ use connectrpc::client::{ClientConfig, HttpClient}; use reqwest::Url; use rustls_platform_verifier::ConfigVerifierExt as _; +use crate::benchmark_registry::version::Version; + use super::proto::v1::{ BenchmarkRegistryServiceClient, ResolveBenchmarkRequest, ResolveBenchmarkResponse, }; @@ -61,7 +63,7 @@ pub(super) fn validate_remote_url(remote_url: &str) -> Result { /// Resolve benchmark metadata from the remote `ConnectRPC` service. pub(super) async fn resolve_manifest( benchmark_name: &str, - version: Option<&str>, + version: Option, endpoint: &Url, ) -> Result> { let uri = endpoint @@ -83,7 +85,11 @@ pub(super) async fn resolve_manifest( benchmark_name: benchmark_name.to_owned(), // If the version is passed as `None` to this function, send the empty string // over RPC - version: version.unwrap_or_default().to_owned(), + version: if let Some(ver) = version { + ver.to_string() + } else { + String::new() + }, ..Default::default() }; diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs index 3808266..d4bb01a 100644 --- a/cli/src/benchmark_registry/manifest.rs +++ b/cli/src/benchmark_registry/manifest.rs @@ -4,6 +4,8 @@ use std::path::{Component, Path, PathBuf}; use anyhow::{Context, Result, bail}; use reqwest::Url; +use crate::benchmark_registry::version::Version; + use super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; /// Maximum number of resources allowed in a remote benchmark manifest. @@ -16,7 +18,7 @@ const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; /// Validate that a response identifies the requested immutable benchmark manifest. pub(super) fn validate_response_identity( benchmark_name: &str, - requested_version: Option<&str>, + requested_version: Option, response: &ResolveBenchmarkResponse, ) -> Result<()> { if response.benchmark_name != benchmark_name { @@ -29,7 +31,7 @@ pub(super) fn validate_response_identity( bail!("remote benchmark response is missing an immutable version"); } if let Some(version) = requested_version - && response.version != version + && response.version != version.clone().to_string() { bail!( "remote benchmark response version `{}` does not match requested version `{version}`", diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index 6a373f0..32ef784 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -3,6 +3,7 @@ pub use self::benchmark::RemoteBenchmark; pub use self::client::select_remote_url; pub use self::resolver::resolve_and_download; +pub use self::version::Version; mod benchmark; mod client; @@ -10,3 +11,4 @@ mod download; mod manifest; mod proto; mod resolver; +mod version; diff --git a/cli/src/benchmark_registry/resolver.rs b/cli/src/benchmark_registry/resolver.rs index fd0b123..fb93d1d 100644 --- a/cli/src/benchmark_registry/resolver.rs +++ b/cli/src/benchmark_registry/resolver.rs @@ -1,5 +1,7 @@ use anyhow::Result; +use crate::benchmark_registry::version::Version; + use super::RemoteBenchmark; use super::client::{resolve_manifest, validate_remote_url}; use super::download::download_resources; @@ -20,15 +22,11 @@ use super::manifest::{validate_resources, validate_response_identity}; /// digest mismatches, invalid UTF-8, or invalid no-code benchmark definitions. pub async fn resolve_and_download( benchmark_name: &str, - version: Option<&str>, + version: Option, remote_url: &str, ) -> Result> { - if version.is_some_and(str::is_empty) { - anyhow::bail!("remote benchmark version must not be passed as the empty string"); - } - let endpoint = validate_remote_url(remote_url)?; - let Some(response) = resolve_manifest(benchmark_name, version, &endpoint).await? else { + let Some(response) = resolve_manifest(benchmark_name, version.clone(), &endpoint).await? else { return Ok(None); }; @@ -148,9 +146,13 @@ mod tests { .mount(&server) .await; - let error = resolve_and_download("remote-test", Some("v1"), &server.uri()) - .await - .unwrap_err(); + let error = resolve_and_download( + "remote-test", + Some(Version::new("v1").unwrap()), + &server.uri(), + ) + .await + .unwrap_err(); assert!( error @@ -159,19 +161,6 @@ mod tests { ); } - #[tokio::test] - async fn rejects_an_empty_explicit_version() { - let error = resolve_and_download("remote-test", Some(""), "https://api.quantiles.io") - .await - .unwrap_err(); - - assert!( - error - .to_string() - .contains("must not be passed as the empty string") - ); - } - struct RequestedVersion(&'static str); impl Match for RequestedVersion { diff --git a/cli/src/benchmark_registry/version.rs b/cli/src/benchmark_registry/version.rs new file mode 100644 index 0000000..00b125e --- /dev/null +++ b/cli/src/benchmark_registry/version.rs @@ -0,0 +1,23 @@ +use anyhow::{Result, bail}; + +#[derive(Clone, Debug)] +pub struct Version(String); + +impl Version { + pub fn new(ver: &str) -> Result { + if ver.is_empty() { + bail!("version must not be empty"); + } + Ok(Self(ver.to_string())); + } + + pub(crate) fn to_string(self) -> String { + self.0.clone() + } +} + +impl std::fmt::Display for Version { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index 4e3fef1..b3c4eb7 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -2,6 +2,7 @@ use std::time::Instant; use anyhow::{Context, Result, bail}; +use qt::benchmark_registry::Version; use qt::builtins; use qt::db; use qt::db::RunStatus; @@ -149,7 +150,7 @@ async fn prepare_remote_resume( } let remote = qt::benchmark_registry::resolve_and_download( workflow_name, - Some(&provenance.version), + Some(Version::new(&provenance.version)?), &provenance.registry_url, ) .await? @@ -545,6 +546,82 @@ mod tests { assert_eq!(run.error.as_deref(), Some("simulated failure")); } + #[tokio::test] + async fn unchanged_local_prompt_override_resumes_successfully() { + use sha2::{Digest as _, Sha256}; + use wiremock::MockServer; + + let _cwd_guard = CWD_LOCK.lock().await; + let server = MockServer::start().await; + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + let cache_dir = root.join("cache"); + let manifest_sha256 = mock_remote_registry(&server).await; + let prompt_path = root.join("local-prompt.txt"); + let prompt = "{{ row.question }}\nLocal answer:"; + std::fs::write(&prompt_path, prompt).unwrap(); + + let original_hf = std::env::var("HF_DATASETS_SERVER").ok(); + let original_cache = std::env::var("QUANTILES_DATASET_CACHE_DIR").ok(); + unsafe { + std::env::set_var("HF_DATASETS_SERVER", server.uri()); + std::env::set_var("QUANTILES_DATASET_CACHE_DIR", &cache_dir); + } + cache_fixture_rows(&cache_dir).await; + + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = + qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let stored_input = serde_json::json!({ + "dataset": { "name": "fixture/qa" }, + "model": "random", + "prompt_template_file": prompt_path, + "limit": 2, + "style": { "type": "exact_match", "golden_column": "answer" } + }) + .to_string(); + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "remote-resume-test".to_owned(), + registry_url: server.uri(), + version: "v1".to_owned(), + manifest_sha256, + prompt_template_sha256: Some(format!("{:x}", Sha256::digest(prompt.as_bytes()))), + }; + let run_id = qt::db::create_remote_benchmark_run( + &db, + "remote-resume-test", + Some(&stored_input), + &provenance, + ) + .await + .unwrap(); + qt::db::fail_run(&db, &metrics_store, run_id, "simulated failure") + .await + .unwrap(); + + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + restore_env("HF_DATASETS_SERVER", original_hf); + restore_env("QUANTILES_DATASET_CACHE_DIR", original_cache); + result.unwrap(); + + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Completed); + assert_eq!( + qt::db::list_steps_for_run(&db, run_id).await.unwrap().len(), + 2 + ); + let metrics = metrics_store.list_aggregate_for_run(run_id).await.unwrap(); + assert!( + metrics + .iter() + .any(|metric| metric.metric_name == "accuracy") + ); + } + #[tokio::test] async fn remote_benchmark_resumes_from_exact_published_version() { use wiremock::MockServer; From d10f7462f2c8cd2d59a975cf15c0e2e8b4f4deef Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:18:40 -0700 Subject: [PATCH 08/10] adding unit tests for version --- cli/src/benchmark_registry/version.rs | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/cli/src/benchmark_registry/version.rs b/cli/src/benchmark_registry/version.rs index 00b125e..3637878 100644 --- a/cli/src/benchmark_registry/version.rs +++ b/cli/src/benchmark_registry/version.rs @@ -21,3 +21,30 @@ impl std::fmt::Display for Version { formatter.write_str(&self.0) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_an_empty_version() { + let error = Version::new("").unwrap_err(); + + assert!(error.to_string().contains("version must not be empty")); + } + + #[test] + fn preserves_and_displays_a_nonempty_version() { + let version = Version::new("v1.2.3").unwrap(); + + assert_eq!(format!("{version}"), "v1.2.3"); + assert_eq!(version.to_string(), "v1.2.3"); + } + + #[test] + fn clone_preserves_the_version() { + let version = Version::new("release-42").unwrap(); + + assert_eq!(version.clone().to_string(), version.to_string()); + } +} From 870f44e3b78d4efd301fb2938b83d9512a2d3515 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:21:37 -0700 Subject: [PATCH 09/10] small fixups and docs --- cli/src/benchmark_registry/manifest.rs | 2 +- cli/src/benchmark_registry/version.rs | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs index d4bb01a..4a76f62 100644 --- a/cli/src/benchmark_registry/manifest.rs +++ b/cli/src/benchmark_registry/manifest.rs @@ -31,7 +31,7 @@ pub(super) fn validate_response_identity( bail!("remote benchmark response is missing an immutable version"); } if let Some(version) = requested_version - && response.version != version.clone().to_string() + && response.version != version.as_str() { bail!( "remote benchmark response version `{}` does not match requested version `{version}`", diff --git a/cli/src/benchmark_registry/version.rs b/cli/src/benchmark_registry/version.rs index 3637878..04099e5 100644 --- a/cli/src/benchmark_registry/version.rs +++ b/cli/src/benchmark_registry/version.rs @@ -4,15 +4,21 @@ use anyhow::{Result, bail}; pub struct Version(String); impl Version { + /// Create a new Version from a version string. + /// + /// # Errors + /// + /// Returns an error if the given `ver` string is empty. pub fn new(ver: &str) -> Result { if ver.is_empty() { bail!("version must not be empty"); } - Ok(Self(ver.to_string())); + Ok(Self(ver.to_string())) } - pub(crate) fn to_string(self) -> String { - self.0.clone() + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 } } From d1968a6bfaa7dae9efd367643c906bd8ec9494eb Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 9 Aug 2026 11:27:26 -0700 Subject: [PATCH 10/10] adding more tests --- cli/src/commands/resume.rs | 169 +++++++++++++++++++++++++++++++++---- 1 file changed, 153 insertions(+), 16 deletions(-) diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index b3c4eb7..1b2b3f0 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -485,6 +485,67 @@ mod tests { assert_eq!(run.error.as_deref(), Some("simulated failure")); } + #[tokio::test] + async fn changed_remote_manifest_does_not_reset_run_status() { + use wiremock::MockServer; + + let _cwd_guard = CWD_LOCK.lock().await; + let server = MockServer::start().await; + let returned_manifest_sha256 = "b".repeat(64); + mock_remote_registry_versions(&server, &["v1"], &returned_manifest_sha256).await; + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + qt::db::init_workspace(root).await.unwrap(); + let db = qt::db::open_workspace(root).await.unwrap(); + let metrics_store = + qt::metrics_store::MetricsStore::new(qt::db::metrics_dir(root)).unwrap(); + let stored_input = serde_json::json!({ + "dataset": { "name": "fixture/qa" }, + "model": "random", + "prompt_template_file": "prompts/qa.txt", + "limit": 2, + "style": { "type": "exact_match", "golden_column": "answer" } + }) + .to_string(); + let provenance = qt::db::RemoteBenchmarkProvenance { + benchmark_name: "remote-resume-test".to_owned(), + registry_url: server.uri(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + prompt_template_sha256: None, + }; + let run_id = qt::db::create_remote_benchmark_run( + &db, + "remote-resume-test", + Some(&stored_input), + &provenance, + ) + .await + .unwrap(); + qt::db::fail_run(&db, &metrics_store, run_id, "simulated failure") + .await + .unwrap(); + + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + + let error = result.unwrap_err(); + assert!(error.to_string().contains("manifest changed")); + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Failed); + assert_eq!(run.error.as_deref(), Some("simulated failure")); + assert!( + qt::db::list_steps_for_run(&db, run_id) + .await + .unwrap() + .is_empty() + ); + let events = qt::db::list_events_for_run(&db, run_id).await.unwrap(); + assert!(!events.iter().any(|event| event.event_type == "run.resumed")); + } + #[tokio::test] async fn changed_local_prompt_override_does_not_reset_run_status() { use sha2::{Digest as _, Sha256}; @@ -694,7 +755,82 @@ mod tests { ); } + #[tokio::test] + async fn remote_run_failure_resumes_from_its_persisted_provenance() { + use wiremock::MockServer; + + let _cwd_guard = CWD_LOCK.lock().await; + let server = MockServer::start().await; + let expected_manifest_sha256 = "a".repeat(64); + mock_remote_registry_versions(&server, &["", "v1"], &expected_manifest_sha256).await; + let tmpdir = tempfile::tempdir().unwrap(); + let root = tmpdir.path(); + let cache_dir = root.join("cache"); + + let original_hf = std::env::var("HF_DATASETS_SERVER").ok(); + let original_cache = std::env::var("QUANTILES_DATASET_CACHE_DIR").ok(); + let original_remote = std::env::var("QUANTILES_REMOTE_URL").ok(); + unsafe { + std::env::set_var("HF_DATASETS_SERVER", server.uri()); + std::env::set_var("QUANTILES_DATASET_CACHE_DIR", &cache_dir); + std::env::remove_var("QUANTILES_REMOTE_URL"); + } + + let original_cwd = std::env::current_dir().unwrap(); + std::env::set_current_dir(root).unwrap(); + let initial_result = crate::commands::run::run( + "remote-resume-test", + None, + Some(&server.uri()), + true, + std::time::Instant::now(), + ) + .await; + assert!(initial_result.is_err()); + + let db = qt::db::open_workspace(root).await.unwrap(); + let runs = qt::db::list_runs(&db).await.unwrap(); + assert_eq!(runs.len(), 1); + let run_id = runs[0].id; + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Failed); + let provenance = qt::db::get_remote_benchmark_provenance(&db, run_id) + .await + .unwrap() + .unwrap(); + assert_eq!(provenance.benchmark_name, "remote-resume-test"); + assert_eq!(provenance.registry_url, server.uri()); + assert_eq!(provenance.version, "v1"); + assert_eq!(provenance.manifest_sha256, expected_manifest_sha256); + assert_eq!(provenance.prompt_template_sha256, None); + + cache_fixture_rows(&cache_dir).await; + let resume_result = resume(run_id, true, std::time::Instant::now()).await; + std::env::set_current_dir(original_cwd).unwrap(); + restore_env("HF_DATASETS_SERVER", original_hf); + restore_env("QUANTILES_DATASET_CACHE_DIR", original_cache); + restore_env("QUANTILES_REMOTE_URL", original_remote); + resume_result.unwrap(); + + let run = qt::db::get_run(&db, run_id).await.unwrap(); + assert_eq!(run.status, qt::db::RunStatus::Completed); + assert_eq!( + qt::db::list_steps_for_run(&db, run_id).await.unwrap().len(), + 2 + ); + } + async fn mock_remote_registry(server: &wiremock::MockServer) -> String { + let manifest_sha256 = "a".repeat(64); + mock_remote_registry_versions(server, &["v1"], &manifest_sha256).await; + manifest_sha256 + } + + async fn mock_remote_registry_versions( + server: &wiremock::MockServer, + requested_versions: &[&'static str], + manifest_sha256: &str, + ) { use buffa::Message as _; use sha2::{Digest as _, Sha256}; use wiremock::matchers::{method, path}; @@ -727,11 +863,10 @@ style = { type = "exact_match", golden_column = "answer" } ..Default::default() } }; - let manifest_sha256 = "a".repeat(64); let response = ResolveBenchmarkResponse { benchmark_name: "remote-resume-test".to_owned(), version: "v1".to_owned(), - manifest_sha256: manifest_sha256.clone(), + manifest_sha256: manifest_sha256.to_owned(), resources: vec![ resource( "definition", @@ -751,19 +886,22 @@ style = { type = "exact_match", golden_column = "answer" } ..Default::default() }; - Mock::given(method("POST")) - .and(path( - "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", - )) - .and(RequestedVersion("v1")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "application/proto") - .set_body_bytes(response.encode_to_vec()), - ) - .expect(1) - .mount(server) - .await; + let response_body = response.encode_to_vec(); + for requested_version in requested_versions { + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .and(RequestedVersion(requested_version)) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response_body.clone()), + ) + .expect(1) + .mount(server) + .await; + } for (route, body) in [ ("/definition", definition.as_slice()), ("/prompt", prompt.as_slice()), @@ -775,7 +913,6 @@ style = { type = "exact_match", golden_column = "answer" } .await; } mock_dataset_metadata(server).await; - manifest_sha256 } async fn mock_dataset_metadata(server: &wiremock::MockServer) {