From b0999ef2288da39fa90b4ab4b70d91446d5c8f65 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 11 Jul 2026 12:39:43 -0700 Subject: [PATCH 01/18] Adding hugging-face dataset prefixes --- CONFIG.md | 18 ++++++++ README.md | 1 + cli/README.md | 3 +- cli/examples/configs/anthropic/quantiles.toml | 2 + .../configs/both_error/.quantiles.toml | 3 ++ .../configs/cloudflare/quantiles.toml | 2 + cli/examples/configs/gemini/quantiles.toml | 2 + cli/examples/configs/openai/quantiles.toml | 2 + cli/examples/configs/simple/quantiles.toml | 2 + cli/src/builtins/common.rs | 3 ++ cli/src/builtins/input.rs | 3 ++ cli/src/builtins/pubmedqa/mod.rs | 12 ++++- cli/src/builtins/similarity.rs | 19 +++++--- cli/src/commands/resume.rs | 2 + cli/src/commands/run.rs | 13 +++++- cli/src/config.rs | 7 +++ cli/src/dataset/mod.rs | 45 ++++++++++++++++++- cli/src/server/mod.rs | 11 +---- 18 files changed, 130 insertions(+), 20 deletions(-) diff --git a/CONFIG.md b/CONFIG.md index 6ed2a10..0e28e75 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -26,6 +26,7 @@ For example, if you want to override default parameters for the built-in PubMedQ # Configure the model and sample limit for the built-in PubMedQA benchmark. [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" ``` @@ -41,6 +42,7 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below | Field | Type | Required | Description | | ------------- | --------------- | -------- | ---------------------------------------------------------------- | | `type` | string | no | Defaults to `"builtin"`. May be omitted for built-in benchmarks. | +| `dataset` | string | no | Dataset source. Hugging Face datasets must use `hf://...`. | | `samples` | integer | no | Number of dataset rows to evaluate. | | `model` | string or table | no | Model sampler. See [model format](#model-format). | | `max_workers` | integer | no | Maximum concurrent workers. | @@ -48,10 +50,24 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below If none of these fields are customized, the built-in benchmark uses the following defaults: - `type`: `builtin` +- `dataset`: The benchmark's default Hugging Face dataset. - `samples`: All samples available in the benchmark's dataset, in order - `model`: The "demo" model, which outputs random values - `max_workers`: The default parallelism provided by the Rust [Tokio runtime](https://tokio.rs/) +#### `dataset` naming + +The `dataset` field accepts a Hugging Face dataset source prefixed with `hf://`. +For example: + +```toml +dataset = "hf://quantiles/PubMedQA" +``` + +The CLI resolves `hf://quantiles/PubMedQA` to the Hugging Face dataset ID +`quantiles/PubMedQA` before using the normal Hugging Face download path. +Other URI-style prefixes are rejected. + #### `model` naming The `model` field described above accepts a provider-prefixed string, for example: @@ -156,6 +172,7 @@ When you run `qt resume `, the CLI looks up the stored eval name and inp ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" model = "openai:gpt-5.6" ``` @@ -163,6 +180,7 @@ model = "openai:gpt-5.6" ```toml [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 10 ``` diff --git a/README.md b/README.md index 553dd56..5ed1603 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ For **built-in benchmarks**, configure settings like `samples`, `model`, and `ma ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" max_workers = 100 diff --git a/cli/README.md b/cli/README.md index 42747ce..7521479 100644 --- a/cli/README.md +++ b/cli/README.md @@ -60,10 +60,11 @@ You can customize how the CLI executes built-in benchmarks and custom evaluation ### Built-in benchmarks -For built-in benchmarks, configure settings like `samples`, `model`, and `max_workers`: +For built-in benchmarks, configure settings like `dataset`, `samples`, `model`, and `max_workers`: ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" max_workers = 100 diff --git a/cli/examples/configs/anthropic/quantiles.toml b/cli/examples/configs/anthropic/quantiles.toml index 7405860..05f6e8a 100644 --- a/cli/examples/configs/anthropic/quantiles.toml +++ b/cli/examples/configs/anthropic/quantiles.toml @@ -5,6 +5,8 @@ # variable. To get an API key, go to https://platform.claude.com/settings/workspaces/default/keys [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since Anthropic charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/both_error/.quantiles.toml b/cli/examples/configs/both_error/.quantiles.toml index e2a96a1..040121c 100644 --- a/cli/examples/configs/both_error/.quantiles.toml +++ b/cli/examples/configs/both_error/.quantiles.toml @@ -1,8 +1,11 @@ [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 10 [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 100 [benchmarks.financebench] +dataset = "hf://quantiles/financebench" samples = 50 diff --git a/cli/examples/configs/cloudflare/quantiles.toml b/cli/examples/configs/cloudflare/quantiles.toml index 29bf1cd..ae4f45c 100644 --- a/cli/examples/configs/cloudflare/quantiles.toml +++ b/cli/examples/configs/cloudflare/quantiles.toml @@ -13,6 +13,8 @@ # https://dash.cloudflare.com//ai/ai-gateway [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since Cloudflare charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/gemini/quantiles.toml b/cli/examples/configs/gemini/quantiles.toml index cfc348c..6c5c893 100644 --- a/cli/examples/configs/gemini/quantiles.toml +++ b/cli/examples/configs/gemini/quantiles.toml @@ -7,6 +7,8 @@ # https://aistudio.google.com/prompts/new_chat [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since Google charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/openai/quantiles.toml b/cli/examples/configs/openai/quantiles.toml index db46a7a..4e31d8d 100644 --- a/cli/examples/configs/openai/quantiles.toml +++ b/cli/examples/configs/openai/quantiles.toml @@ -8,6 +8,8 @@ # https://platform.openai.com/home [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1000 samples. Since OpenAI charges for # usage, we are limiting the number of samples to 50 in this demo, to keep costs lower. samples = 50 diff --git a/cli/examples/configs/simple/quantiles.toml b/cli/examples/configs/simple/quantiles.toml index 06b287e..9f2998f 100644 --- a/cli/examples/configs/simple/quantiles.toml +++ b/cli/examples/configs/simple/quantiles.toml @@ -8,7 +8,9 @@ # number of dataset rows processed. [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 10 [benchmarks.simpleqa-verified] +dataset = "hf://quantiles/simpleqa-verified" samples = 100 diff --git a/cli/src/builtins/common.rs b/cli/src/builtins/common.rs index 53ea370..6ed37fa 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -18,6 +18,9 @@ pub(crate) struct BuiltinConfig { /// Number of dataset rows to evaluate. If omitted, the entire dataset is used. #[serde(default)] pub(crate) limit: Option, + /// Dataset source to evaluate. Builtins currently support Hugging Face via `hf://...`. + #[serde(default)] + pub(crate) dataset: Option, /// Which model sampler to use. If omitted, the builtin chooses a sensible default. #[serde(default)] pub(crate) model: Option, diff --git a/cli/src/builtins/input.rs b/cli/src/builtins/input.rs index 7d443b5..437a124 100644 --- a/cli/src/builtins/input.rs +++ b/cli/src/builtins/input.rs @@ -5,6 +5,7 @@ use crate::llm::Sampler; /// Normalized run input schema for all builtins. #[derive(Serialize)] pub(crate) struct BuiltinRunInput { + pub(crate) dataset: String, pub(crate) model: String, pub(crate) num_samples: usize, #[serde(skip_serializing_if = "Option::is_none")] @@ -15,11 +16,13 @@ pub(crate) struct BuiltinRunInput { pub(crate) async fn set_builtin_run_input( db: &sea_orm::DatabaseConnection, run_id: i64, + dataset: &str, model: Option<&Sampler>, num_samples: usize, max_workers: Option, ) -> anyhow::Result<()> { let input = serde_json::to_string(&BuiltinRunInput { + dataset: dataset.to_owned(), model: builtin_model_name(model), num_samples, max_workers, diff --git a/cli/src/builtins/pubmedqa/mod.rs b/cli/src/builtins/pubmedqa/mod.rs index 7ea162c..d6e6cc5 100644 --- a/cli/src/builtins/pubmedqa/mod.rs +++ b/cli/src/builtins/pubmedqa/mod.rs @@ -7,7 +7,7 @@ use crate::builtins::dataset_runner::DatasetRunner; use crate::builtins::input::set_builtin_run_input; use crate::builtins::output::set_builtin_run_output; use crate::builtins::{BuiltinContext, BuiltinWorkflow}; -use crate::dataset::DatasetManager; +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::llm::LLMSampler; use crate::llm::random_label::RandomLabelSampler; @@ -22,6 +22,8 @@ mod eval; /// `PubMedQA` builtin using the quantiles/PubMedQA dataset. pub struct PubmedqaBuiltin; +const DEFAULT_DATASET_SOURCE: &str = "hf://quantiles/PubMedQA"; + #[expect(clippy::too_many_lines)] #[async_trait::async_trait] impl BuiltinWorkflow for PubmedqaBuiltin { @@ -49,7 +51,12 @@ impl BuiltinWorkflow for PubmedqaBuiltin { }; let manager = DatasetManager::new()?; - let dataset_id = "quantiles/PubMedQA"; + let dataset_source = config + .base + .dataset + .as_deref() + .unwrap_or(DEFAULT_DATASET_SOURCE); + let dataset_id = resolve_hf_dataset_source(dataset_source)?; let info = manager .init(dataset_id, Some("pqa_labeled"), Some("train"), None) .await?; @@ -62,6 +69,7 @@ impl BuiltinWorkflow for PubmedqaBuiltin { set_builtin_run_input( ctx.db, ctx.run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, diff --git a/cli/src/builtins/similarity.rs b/cli/src/builtins/similarity.rs index 051a260..0fd9083 100644 --- a/cli/src/builtins/similarity.rs +++ b/cli/src/builtins/similarity.rs @@ -9,7 +9,7 @@ use crate::builtins::dataset_runner::DatasetRunner; use crate::builtins::input::set_builtin_run_input; use crate::builtins::output::set_builtin_run_output; use crate::builtins::{BuiltinContext, BuiltinWorkflow}; -use crate::dataset::DatasetManager; +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::llm::LLMSampler; use crate::llm::random::RandomSampler; use crate::similarity::{ @@ -45,7 +45,7 @@ struct RowOutput { #[derive(Clone, Copy)] pub struct SimilarityBenchmark { name: &'static str, - dataset_id: &'static str, + dataset_source: &'static str, input_field: &'static str, target_field: &'static str, } @@ -53,7 +53,7 @@ pub struct SimilarityBenchmark { /// `simpleqa-verified` builtin. pub const SIMPLEQA: SimilarityBenchmark = SimilarityBenchmark { name: "simpleqa-verified", - dataset_id: "quantiles/simpleqa-verified", + dataset_source: "hf://quantiles/simpleqa-verified", input_field: "problem", target_field: "answer", }; @@ -61,7 +61,7 @@ pub const SIMPLEQA: SimilarityBenchmark = SimilarityBenchmark { /// `financebench` builtin. pub const FINANCEBENCH: SimilarityBenchmark = SimilarityBenchmark { name: "financebench", - dataset_id: "quantiles/financebench", + dataset_source: "hf://quantiles/financebench", input_field: "question", target_field: "answer", }; @@ -96,7 +96,13 @@ impl BuiltinWorkflow for SimilarityBenchmark { }; let manager = DatasetManager::new()?; - let info = manager.init(self.dataset_id, None, None, None).await?; + let dataset_source = config + .base + .dataset + .as_deref() + .unwrap_or(self.dataset_source); + let dataset_id = resolve_hf_dataset_source(dataset_source)?; + let info = manager.init(dataset_id, None, None, None).await?; let total = info .total_rows @@ -109,6 +115,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { set_builtin_run_input( db, run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, @@ -120,7 +127,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { let metric = &metric; let max_workers = config.base.max_workers.unwrap_or_else(get_max_workers); - let scores = DatasetRunner::new(&manager, self.dataset_id, &info, limit) + let scores = DatasetRunner::new(&manager, dataset_id, &info, limit) .desc(self.name) .set_quiet(ctx.quiet) .for_each_concurrent(max_workers, |i, row| { diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index 91b61af..37eb86c 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -132,6 +132,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, + dataset: None, model: None, max_workers: None, }); @@ -146,6 +147,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + dataset: None, model: None, max_workers: None, }); diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index 18fcb06..e33796e 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -103,12 +103,17 @@ fn assemble_builtin_input( } if let Some(bench) = bench { - if bench.samples.is_none() && bench.model.is_none() && bench.max_workers.is_none() { + if bench.samples.is_none() + && bench.dataset.is_none() + && bench.model.is_none() + && bench.max_workers.is_none() + { return (None, Vec::new()); } let input = BuiltinConfigInput { limit: bench.samples, + dataset: bench.dataset.clone(), model: bench.model.clone(), max_workers: bench.max_workers, }; @@ -383,6 +388,8 @@ struct BuiltinConfigInput { #[serde(skip_serializing_if = "Option::is_none")] limit: Option, #[serde(skip_serializing_if = "Option::is_none")] + dataset: Option, + #[serde(skip_serializing_if = "Option::is_none")] model: Option, #[serde(skip_serializing_if = "Option::is_none")] max_workers: Option, @@ -623,6 +630,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + dataset: None, model: None, max_workers: None, }; @@ -637,12 +645,14 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(5), + dataset: Some("hf://quantiles/PubMedQA".to_owned()), model: Some(qt::llm::Sampler::Random {}), max_workers: Some(8), }; let (input, _) = super::assemble_builtin_input(Some(&bench), None); let parsed: serde_json::Value = serde_json::from_str(&input.unwrap()).unwrap(); assert_eq!(parsed["limit"], 5); + assert_eq!(parsed["dataset"], "hf://quantiles/PubMedQA"); assert_eq!(parsed["model"], "random"); assert_eq!(parsed["max_workers"], 8); } @@ -654,6 +664,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, + dataset: None, model: None, max_workers: None, }; diff --git a/cli/src/config.rs b/cli/src/config.rs index 0da4a18..da17ebe 100644 --- a/cli/src/config.rs +++ b/cli/src/config.rs @@ -76,6 +76,8 @@ pub struct BuiltinBenchmarkConfig { pub type_: String, /// Number of samples (rows) to evaluate. pub samples: Option, + /// Dataset source for this benchmark. + pub dataset: Option, /// Which model sampler to use for this benchmark. pub model: Option, /// Maximum concurrent workers for this benchmark. @@ -167,6 +169,7 @@ mod tests { if let BenchmarkConfig::Builtin(b) = bench { assert_eq!(b.type_, "builtin"); assert_eq!(b.samples, Some(10)); + assert!(b.dataset.is_none()); assert!(b.model.is_none()); } } @@ -177,11 +180,15 @@ mod tests { [benchmarks.demo] type = "builtin" samples = 5 + dataset = "hf://quantiles/demo" model = "openai:gpt-4" "#; let config: WorkspaceConfig = toml::from_str(toml).unwrap(); let bench = config.benchmarks.get("demo").unwrap(); assert!(matches!(bench, BenchmarkConfig::Builtin(_))); + if let BenchmarkConfig::Builtin(b) = bench { + assert_eq!(b.dataset.as_deref(), Some("hf://quantiles/demo")); + } } #[test] diff --git a/cli/src/dataset/mod.rs b/cli/src/dataset/mod.rs index b740fd8..d08dfd9 100644 --- a/cli/src/dataset/mod.rs +++ b/cli/src/dataset/mod.rs @@ -3,7 +3,7 @@ pub mod hf_client; use crate::dataset::cache::DatasetCache; use crate::dataset::hf_client::HuggingFaceClient; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, bail}; use serde::Serialize; use serde_json::Value; @@ -145,3 +145,46 @@ impl DatasetManager { splits.first().cloned().context("dataset has no splits") } } + +/// Resolve a configured Hugging Face dataset source to the dataset ID expected +/// by the existing Hugging Face download client. +pub fn resolve_hf_dataset_source(source: &str) -> Result<&str> { + if let Some(dataset_id) = source.strip_prefix("hf://") { + if dataset_id.is_empty() { + bail!("dataset source `hf://` is missing a Hugging Face dataset id"); + } + Ok(dataset_id) + } else if source.contains("://") { + bail!("unsupported dataset source `{source}`; expected `hf://...`"); + } else { + bail!("dataset source `{source}` is missing required `hf://` prefix"); + } +} + +#[cfg(test)] +mod tests { + use super::resolve_hf_dataset_source; + + #[test] + fn resolve_hf_dataset_source_strips_hf_prefix() { + assert_eq!( + resolve_hf_dataset_source("hf://quantiles/PubMedQA").unwrap(), + "quantiles/PubMedQA" + ); + } + + #[test] + fn resolve_hf_dataset_source_rejects_other_prefixes() { + let err = resolve_hf_dataset_source("s3://bucket/dataset").unwrap_err(); + assert!( + err.to_string() + .contains("unsupported dataset source `s3://bucket/dataset`") + ); + } + + #[test] + fn resolve_hf_dataset_source_requires_prefix() { + let err = resolve_hf_dataset_source("quantiles/PubMedQA").unwrap_err(); + assert!(err.to_string().contains("missing required `hf://` prefix")); + } +} diff --git a/cli/src/server/mod.rs b/cli/src/server/mod.rs index 40b9631..94d9fbd 100644 --- a/cli/src/server/mod.rs +++ b/cli/src/server/mod.rs @@ -15,7 +15,7 @@ use sea_orm::DatabaseConnection; use serde_json::json; use tokio::net::TcpListener; -use crate::dataset::DatasetManager; +use crate::dataset::{DatasetManager, resolve_hf_dataset_source}; use crate::db::steps::{self, StepDecision}; use crate::db::{self, DBUrl, SQLitePathURL}; use crate::metrics_store::MetricsStore; @@ -321,14 +321,7 @@ async fn dataset_batch( } fn parse_hf_source(source: &str) -> Result<&str, ApiError> { - source - .strip_prefix("huggingface://") - .or_else(|| source.strip_prefix("hf://")) - .ok_or_else(|| { - ApiError(anyhow::anyhow!( - "unsupported dataset source `{source}`; expected `huggingface://...` or `hf://...`" - )) - }) + resolve_hf_dataset_source(source).map_err(ApiError) } struct ApiError(anyhow::Error); From 13e2835517103fe87f3e1f14252e4511e9560d8a Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 11 Jul 2026 12:43:39 -0700 Subject: [PATCH 02/18] progress --- CONFIG.md | 8 ++++---- cli/src/commands/resume.rs | 4 ++-- cli/src/commands/run.rs | 26 +++++++++----------------- cli/src/config.rs | 17 ++++++++++++++--- cli/src/dataset/mod.rs | 24 +++++++++++++++++++----- 5 files changed, 48 insertions(+), 31 deletions(-) diff --git a/CONFIG.md b/CONFIG.md index 0e28e75..45382e1 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -42,7 +42,7 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below | Field | Type | Required | Description | | ------------- | --------------- | -------- | ---------------------------------------------------------------- | | `type` | string | no | Defaults to `"builtin"`. May be omitted for built-in benchmarks. | -| `dataset` | string | no | Dataset source. Hugging Face datasets must use `hf://...`. | +| `dataset` | string | yes | Dataset source. Hugging Face datasets must use `hf://...` or `huggingface://...`. | | `samples` | integer | no | Number of dataset rows to evaluate. | | `model` | string or table | no | Model sampler. See [model format](#model-format). | | `max_workers` | integer | no | Maximum concurrent workers. | @@ -50,15 +50,15 @@ Built-in benchmarks run natively inside the CLI, without any custom code. Below If none of these fields are customized, the built-in benchmark uses the following defaults: - `type`: `builtin` -- `dataset`: The benchmark's default Hugging Face dataset. +- `dataset`: The benchmark's default Hugging Face dataset, used when no config section is provided. - `samples`: All samples available in the benchmark's dataset, in order - `model`: The "demo" model, which outputs random values - `max_workers`: The default parallelism provided by the Rust [Tokio runtime](https://tokio.rs/) #### `dataset` naming -The `dataset` field accepts a Hugging Face dataset source prefixed with `hf://`. -For example: +The `dataset` field accepts a Hugging Face dataset source prefixed with `hf://` +or `huggingface://`. For example: ```toml dataset = "hf://quantiles/PubMedQA" diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index 37eb86c..86e68ed 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -132,7 +132,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, - dataset: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }); @@ -147,7 +147,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), - dataset: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }); diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index e33796e..ab536d8 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -103,14 +103,6 @@ fn assemble_builtin_input( } if let Some(bench) = bench { - if bench.samples.is_none() - && bench.dataset.is_none() - && bench.model.is_none() - && bench.max_workers.is_none() - { - return (None, Vec::new()); - } - let input = BuiltinConfigInput { limit: bench.samples, dataset: bench.dataset.clone(), @@ -387,8 +379,7 @@ struct BuiltinRunJsonOutput { struct BuiltinConfigInput { #[serde(skip_serializing_if = "Option::is_none")] limit: Option, - #[serde(skip_serializing_if = "Option::is_none")] - dataset: Option, + dataset: String, #[serde(skip_serializing_if = "Option::is_none")] model: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -630,7 +621,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), - dataset: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }; @@ -645,7 +636,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(5), - dataset: Some("hf://quantiles/PubMedQA".to_owned()), + dataset: "hf://quantiles/PubMedQA".to_owned(), model: Some(qt::llm::Sampler::Random {}), max_workers: Some(8), }; @@ -657,19 +648,20 @@ mod tests { assert_eq!(parsed["max_workers"], 8); } - /// When the builtin config section exists but has no runtime-relevant fields, the - /// input should be `None` rather than an empty JSON object. + /// When the builtin config section only has the required dataset, the input should + /// still carry that dataset source into builtin execution. #[test] - fn assemble_builtin_input_none_when_no_config_fields() { + fn assemble_builtin_input_with_dataset_only_config() { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, - dataset: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }; let (input, _) = super::assemble_builtin_input(Some(&bench), None); - assert!(input.is_none()); + let parsed: serde_json::Value = serde_json::from_str(&input.unwrap()).unwrap(); + assert_eq!(parsed["dataset"], "hf://quantiles/PubMedQA"); } /// When there is no config section at all and no CLI `--input`, builtin runs should diff --git a/cli/src/config.rs b/cli/src/config.rs index da17ebe..2afa4b8 100644 --- a/cli/src/config.rs +++ b/cli/src/config.rs @@ -77,7 +77,7 @@ pub struct BuiltinBenchmarkConfig { /// Number of samples (rows) to evaluate. pub samples: Option, /// Dataset source for this benchmark. - pub dataset: Option, + pub dataset: String, /// Which model sampler to use for this benchmark. pub model: Option, /// Maximum concurrent workers for this benchmark. @@ -161,6 +161,7 @@ mod tests { fn deserialize_builtin_without_type() { let toml = r#" [benchmarks.demo] + dataset = "hf://quantiles/demo" samples = 10 "#; let config: WorkspaceConfig = toml::from_str(toml).unwrap(); @@ -169,7 +170,7 @@ mod tests { if let BenchmarkConfig::Builtin(b) = bench { assert_eq!(b.type_, "builtin"); assert_eq!(b.samples, Some(10)); - assert!(b.dataset.is_none()); + assert_eq!(b.dataset, "hf://quantiles/demo"); assert!(b.model.is_none()); } } @@ -187,10 +188,20 @@ mod tests { let bench = config.benchmarks.get("demo").unwrap(); assert!(matches!(bench, BenchmarkConfig::Builtin(_))); if let BenchmarkConfig::Builtin(b) = bench { - assert_eq!(b.dataset.as_deref(), Some("hf://quantiles/demo")); + assert_eq!(b.dataset, "hf://quantiles/demo"); } } + #[test] + fn builtin_requires_dataset_field() { + let toml = r#" + [benchmarks.demo] + samples = 5 + "#; + let result: Result = toml::from_str(toml); + assert!(result.is_err(), "builtin should require dataset field"); + } + #[test] fn deserialize_custom_code() { let toml = r#" diff --git a/cli/src/dataset/mod.rs b/cli/src/dataset/mod.rs index d08dfd9..5242ff2 100644 --- a/cli/src/dataset/mod.rs +++ b/cli/src/dataset/mod.rs @@ -149,15 +149,18 @@ impl DatasetManager { /// Resolve a configured Hugging Face dataset source to the dataset ID expected /// by the existing Hugging Face download client. pub fn resolve_hf_dataset_source(source: &str) -> Result<&str> { - if let Some(dataset_id) = source.strip_prefix("hf://") { + if let Some(dataset_id) = source + .strip_prefix("hf://") + .or_else(|| source.strip_prefix("huggingface://")) + { if dataset_id.is_empty() { - bail!("dataset source `hf://` is missing a Hugging Face dataset id"); + bail!("dataset source `{source}` is missing a Hugging Face dataset id"); } Ok(dataset_id) } else if source.contains("://") { - bail!("unsupported dataset source `{source}`; expected `hf://...`"); + bail!("unsupported dataset source `{source}`; expected `hf://...` or `huggingface://...`"); } else { - bail!("dataset source `{source}` is missing required `hf://` prefix"); + bail!("dataset source `{source}` is missing required `hf://` or `huggingface://` prefix"); } } @@ -173,6 +176,14 @@ mod tests { ); } + #[test] + fn resolve_hf_dataset_source_strips_huggingface_prefix() { + assert_eq!( + resolve_hf_dataset_source("huggingface://quantiles/PubMedQA").unwrap(), + "quantiles/PubMedQA" + ); + } + #[test] fn resolve_hf_dataset_source_rejects_other_prefixes() { let err = resolve_hf_dataset_source("s3://bucket/dataset").unwrap_err(); @@ -185,6 +196,9 @@ mod tests { #[test] fn resolve_hf_dataset_source_requires_prefix() { let err = resolve_hf_dataset_source("quantiles/PubMedQA").unwrap_err(); - assert!(err.to_string().contains("missing required `hf://` prefix")); + assert!( + err.to_string() + .contains("missing required `hf://` or `huggingface://` prefix") + ); } } From 48f64c4c2d44fde803e1b840bde10b48385cbb83 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 1 Aug 2026 12:12:55 -0700 Subject: [PATCH 03/18] Adding client for benchmark registry --- cli/Cargo.lock | 85 +++++++++++++++++++ cli/Cargo.toml | 3 + cli/build.rs | 11 +++ .../benchmark/v1/benchmark_registry.proto | 52 ++++++++++++ 4 files changed, 151 insertions(+) create mode 100644 cli/build.rs create mode 100644 cli/proto/quantiles/benchmark/v1/benchmark_registry.proto diff --git a/cli/Cargo.lock b/cli/Cargo.lock index 3c83a5f..b80d782 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -795,6 +795,51 @@ dependencies = [ "serde", ] +[[package]] +name = "buffa" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33f29a40702df4b86ccd84211bfde8cee0bce6d0811450ade4a86a7d0958a23a" +dependencies = [ + "base64 0.22.1", + "bytes", + "foldhash 0.1.5", + "hashbrown 0.15.5", + "once_cell", + "rustversion", + "serde", + "serde_json", + "smoothutf8", + "thiserror", +] + +[[package]] +name = "buffa-codegen" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6681b562b18ea719622d0d12684e88ecbdca600d517557dc7bcef7704631e28" +dependencies = [ + "buffa", + "buffa-descriptor", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.117", + "thiserror", +] + +[[package]] +name = "buffa-descriptor" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57ed423c4ecec86d1500879ce42e7e5f6def01bc632985ac756c1fd723fa21fc" +dependencies = [ + "buffa", + "rustversion", + "serde", + "serde_json", +] + [[package]] name = "built" version = "0.8.1" @@ -1033,6 +1078,35 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "connectrpc-build" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8cdec2dacacd00437898dad38f98105c82d51086a78baa62cee6dc66c031070" +dependencies = [ + "anyhow", + "buffa", + "buffa-codegen", + "connectrpc-codegen", + "tempfile", +] + +[[package]] +name = "connectrpc-codegen" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c6e3bb62a4a77f41d5cd3d290a488bf529dc7e26b444cda67a3e5544c97542" +dependencies = [ + "anyhow", + "buffa", + "buffa-codegen", + "heck 0.5.0", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "console" version = "0.16.3" @@ -2558,6 +2632,7 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.1.5", + "serde", ] [[package]] @@ -4339,6 +4414,7 @@ dependencies = [ "axum", "clap", "comfy-table", + "connectrpc-build", "datafusion", "dirs", "fastembed", @@ -5454,6 +5530,15 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8e2fb0f499abb4d162f2bedad68f5ef91a1682b5a03596ddb67efd37768d100" +[[package]] +name = "smoothutf8" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36358427d32ecdb1624616deed99eccfef0a167fe5bf40ddb51efe6980bc1ec8" +dependencies = [ + "simdutf8", +] + [[package]] name = "snafu" version = "0.8.9" diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 631f074..bae7a06 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -38,6 +38,9 @@ pedantic = "deny" [lints.rust] unfulfilled_lint_expectations = "deny" +[build-dependencies] +connectrpc-build = "0.8" + [dev-dependencies] assert_cmd = "2.1.1" predicates = "3.1.3" diff --git a/cli/build.rs b/cli/build.rs new file mode 100644 index 0000000..f76245d --- /dev/null +++ b/cli/build.rs @@ -0,0 +1,11 @@ +fn main() -> Result<(), Box> { + println!("cargo:rerun-if-changed=proto/quantiles/benchmark/v1/benchmark_registry.proto"); + + connectrpc_build::Config::new() + .files(&["proto/quantiles/benchmark/v1/benchmark_registry.proto"]) + .use_buf() + .include_file("_connectrpc.rs") + .compile()?; + + Ok(()) +} diff --git a/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto b/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto new file mode 100644 index 0000000..038a74f --- /dev/null +++ b/cli/proto/quantiles/benchmark/v1/benchmark_registry.proto @@ -0,0 +1,52 @@ +syntax = "proto3"; + +package quantiles.benchmark.v1; + +// Resolves immutable benchmark metadata and the resources needed to run it. +service BenchmarkRegistryService { + rpc ResolveBenchmark(ResolveBenchmarkRequest) returns (ResolveBenchmarkResponse); +} + +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. + string version = 2; +} + +message ResolveBenchmarkResponse { + string benchmark_name = 1; + string version = 2; + + // SHA-256 digest of the canonical resource manifest. + string manifest_sha256 = 3; + + repeated BenchmarkResource resources = 4; +} + +message BenchmarkResource { + // Opaque registry identifier used to identify this resource. + string resource_id = 1; + + // Relative path at which the client materializes the resource. + string logical_path = 2; + + ResourceKind kind = 3; + + // HTTP(S) URL from which the client downloads the resource bytes. + string download_url = 4; + + // Lowercase hexadecimal SHA-256 digest of the resource bytes. + string sha256 = 5; + + uint64 size_bytes = 6; + string content_type = 7; +} + +enum ResourceKind { + RESOURCE_KIND_UNSPECIFIED = 0; + RESOURCE_KIND_DEFINITION = 1; + RESOURCE_KIND_PROMPT_TEMPLATE = 2; + RESOURCE_KIND_AUXILIARY = 3; +} From 31df5e5ad2813076a5380098ad290a27ca692516 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 1 Aug 2026 12:46:41 -0700 Subject: [PATCH 04/18] Implementing the client --- README.md | 2 + cli/Cargo.lock | 69 +++ cli/Cargo.toml | 3 + cli/README.md | 8 + cli/src/benchmark_registry/mod.rs | 539 ++++++++++++++++++++++ cli/src/builtins/custom_nocode/mod.rs | 28 +- cli/src/builtins/custom_nocode/runtime.rs | 31 +- cli/src/cli.rs | 3 + cli/src/commands/run.rs | 60 ++- cli/src/lib.rs | 1 + cli/src/main.rs | 12 +- 11 files changed, 747 insertions(+), 9 deletions(-) create mode 100644 cli/src/benchmark_registry/mod.rs diff --git a/README.md b/README.md index 2313164..5d505dd 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,8 @@ qt show qt compare ``` +When `qt run ` does not find the requested name in the local configuration, it queries the Quantiles remote benchmark service for a benchmark that has the given name. See the [CLI's README file](./cli/README.md#remote-benchmark-fallback) for details. + > Note: Pass `--json` to any of these commands to output machine- and agent-friendly JSON instead of human-formatted output. See the [CLI reference](https://quantiles.io/documentation/reference/cli) for available commands, options, and usage details. diff --git a/cli/Cargo.lock b/cli/Cargo.lock index b80d782..762b6a8 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -897,6 +897,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "bzip2" @@ -1078,6 +1081,41 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "connectrpc" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4d1bbfc95ed1f4f31a027f5b97ff981be2617bc8a5c800fd6fcc701867bb97f" +dependencies = [ + "async-compression", + "async-trait", + "base64 0.22.1", + "buffa", + "bytes", + "flate2", + "futures", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "percent-encoding", + "pin-project", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tracing", + "wasm-bindgen-futures", + "zstd", +] + [[package]] name = "connectrpc-build" version = "0.8.0" @@ -2300,6 +2338,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -4155,6 +4194,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -4412,8 +4471,10 @@ dependencies = [ "assert_cmd", "async-trait", "axum", + "buffa", "clap", "comfy-table", + "connectrpc", "connectrpc-build", "datafusion", "dirs", @@ -4426,6 +4487,7 @@ dependencies = [ "rand 0.8.6", "reqwest 0.13.4", "rstest", + "rustls-platform-verifier", "sea-orm", "seq-macro", "serde", @@ -6315,6 +6377,7 @@ dependencies = [ "pin-project-lite", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -7447,6 +7510,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "zlib-rs" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" + [[package]] name = "zmij" version = "1.0.21" diff --git a/cli/Cargo.toml b/cli/Cargo.toml index bae7a06..9476ee4 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -10,6 +10,8 @@ async-trait = "0.1" axum = "0.8.9" clap = { version = "4.6.1", features = ["derive"] } comfy-table = "7.2.1" +buffa = "0.8" +connectrpc = { version = "0.8", features = ["client-tls"] } datafusion = "46.0.0" dirs = "6" fastembed = "5" @@ -19,6 +21,7 @@ jinja = "0.1" parquet = { version = "54", features = ["arrow"] } rand = "0.8" reqwest = { version = "0.13.3", features = ["json", "rustls"] } +rustls-platform-verifier = "0.7" sea-orm = { version = "1.1", features = ["sqlx-sqlite", "runtime-tokio-native-tls", "time"] } seq-macro = "0.3" serde = { version = "1.0.228", features = ["derive"] } diff --git a/cli/README.md b/cli/README.md index 0343732..eec9b6e 100644 --- a/cli/README.md +++ b/cli/README.md @@ -51,6 +51,14 @@ max_workers = 100 See the [configuration guide](https://quantiles.io/documentation/configuration) for file location, supported fields, validation behavior, and examples. Additional runnable configurations are available in [CLI configuration examples](./examples/configs) and [custom no-code examples](../custom-nocode-examples/quantiles.toml). +### Remote benchmark fallback + +When you run `qt run `, the CLI first looks in the local configuration file for an evaluation called `eval_name`. If one is found, the CLI runs it immediately. If none is found, `qt` looks in 1the Quantiles remote benchmark service for an evaluation called `eval_name` and, if a match is found, downloads the benchmark definition and runs it. + +>If you want to override the location of the remote benchmark service, use the `--remote-url` flag or the `QUANTILES_REMOTE_URL` environment variable. + +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. + ## Architecture The Quantiles CLI, `qt`, keeps execution simple: your code runs locally, while `qt` handles durability and observability. diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs new file mode 100644 index 0000000..58519cb --- /dev/null +++ b/cli/src/benchmark_registry/mod.rs @@ -0,0 +1,539 @@ +//! Client support for resolving remotely hosted benchmark definitions. + +use std::collections::{HashMap, HashSet}; +use std::ffi::OsString; +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use connectrpc::ErrorCode; +use connectrpc::client::{ClientConfig, HttpClient}; +use futures::StreamExt as _; +use reqwest::Url; +use rustls_platform_verifier::ConfigVerifierExt as _; +use sha2::{Digest as _, Sha256}; + +use crate::config::{BenchmarkConfig, CustomNoCodeBenchmarkConfig, WorkspaceConfig}; + +#[expect( + clippy::allow_attributes, + clippy::pedantic, + reason = "ConnectRPC and Buffa generated code uses allow attributes" +)] +pub mod proto { + connectrpc::include_generated!(); +} + +use proto::quantiles::benchmark::v1::{ + BenchmarkRegistryServiceClient, BenchmarkResource, ResolveBenchmarkRequest, + ResolveBenchmarkResponse, ResourceKind, +}; + +pub const DEFAULT_REMOTE_URL: &str = "https://api.quantiles.io"; +const MAX_MANIFEST_BYTES: usize = 1024 * 1024; +const MAX_RESOURCE_COUNT: usize = 32; +const MAX_RESOURCE_BYTES: u64 = 10 * 1024 * 1024; +const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; + +/// A downloaded benchmark ready to execute without materializing its resources on disk. +pub struct RemoteBenchmark { + pub config: CustomNoCodeBenchmarkConfig, + pub prompt_template: String, + pub version: String, + pub manifest_sha256: String, +} + +/// Select the remote service URL and reject ambiguous CLI/environment configuration. +/// +/// # Errors +/// +/// Returns an error when both sources are set or the environment value is not valid UTF-8. +pub fn select_remote_url(cli_url: Option<&str>, env_url: Option) -> Result { + if cli_url.is_some() && env_url.is_some() { + bail!("cannot use both `--remote-url` and `QUANTILES_REMOTE_URL`"); + } + + if let Some(url) = cli_url { + return Ok(url.to_owned()); + } + + env_url.map_or_else( + || Ok(DEFAULT_REMOTE_URL.to_owned()), + |url| { + url.into_string() + .map_err(|_| anyhow::anyhow!("`QUANTILES_REMOTE_URL` must be valid UTF-8")) + }, + ) +} + +/// Resolve a benchmark 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. +/// +/// # Errors +/// +/// Returns an error for invalid endpoints, RPC failures, malformed manifests, failed downloads, +/// digest mismatches, invalid UTF-8, or invalid no-code benchmark definitions. +pub async fn resolve_and_download( + benchmark_name: &str, + remote_url: &str, +) -> Result> { + let endpoint = validate_remote_url(remote_url)?; + let Some(response) = resolve_manifest(benchmark_name, &endpoint).await? else { + return Ok(None); + }; + + validate_response_identity(benchmark_name, &response)?; + let resources = validate_resources(&response.resources, endpoint.scheme() == "http")?; + let downloaded = download_resources(&resources).await?; + let remote = build_remote_benchmark(benchmark_name, response, downloaded)?; + Ok(Some(remote)) +} + +fn validate_remote_url(remote_url: &str) -> Result { + let mut url = Url::parse(remote_url) + .with_context(|| format!("invalid remote benchmark service URL `{remote_url}`"))?; + if !matches!(url.scheme(), "http" | "https") { + bail!("remote benchmark service URL must use http or https"); + } + if url.host_str().is_none() { + bail!("remote benchmark service URL must include a host"); + } + if url.query().is_some() || url.fragment().is_some() { + bail!("remote benchmark service URL must not include a query or fragment"); + } + while url.path().ends_with('/') && url.path() != "/" { + let trimmed = url.path().trim_end_matches('/').to_owned(); + url.set_path(&trimmed); + } + Ok(url) +} + +async fn resolve_manifest( + benchmark_name: &str, + endpoint: &Url, +) -> Result> { + let uri = endpoint + .as_str() + .parse() + .with_context(|| format!("invalid remote benchmark service URL `{endpoint}`"))?; + let transport = if endpoint.scheme() == "https" { + let tls = connectrpc::rustls::ClientConfig::with_platform_verifier() + .context("failed to configure TLS for remote benchmark service")?; + HttpClient::with_tls(Arc::new(tls)) + } else { + HttpClient::plaintext() + }; + let config = ClientConfig::new(uri) + .with_default_timeout(Duration::from_secs(15)) + .with_default_max_message_size(MAX_MANIFEST_BYTES); + let client = BenchmarkRegistryServiceClient::new(transport, config); + let request = ResolveBenchmarkRequest { + benchmark_name: benchmark_name.to_owned(), + version: String::new(), + ..Default::default() + }; + + match client.resolve_benchmark(request).await { + Ok(response) => Ok(Some(response.into_owned())), + Err(error) if error.code == ErrorCode::NotFound => Ok(None), + Err(error) => Err(error).with_context(|| { + format!("failed to resolve benchmark `{benchmark_name}` from `{endpoint}`") + }), + } +} + +fn validate_response_identity( + benchmark_name: &str, + response: &ResolveBenchmarkResponse, +) -> Result<()> { + if response.benchmark_name != benchmark_name { + bail!( + "remote benchmark response name `{}` does not match requested name `{benchmark_name}`", + response.benchmark_name + ); + } + if response.version.is_empty() { + bail!("remote benchmark response is missing an immutable version"); + } + validate_sha256("manifest", &response.manifest_sha256)?; + Ok(()) +} + +fn validate_resources( + resources: &[BenchmarkResource], + allow_plaintext_downloads: bool, +) -> Result> { + if resources.is_empty() { + bail!("remote benchmark manifest contains no resources"); + } + if resources.len() > MAX_RESOURCE_COUNT { + bail!( + "remote benchmark manifest contains {} resources; maximum is {MAX_RESOURCE_COUNT}", + resources.len() + ); + } + + let mut resource_ids = HashSet::new(); + let mut logical_paths = HashSet::new(); + let mut total_bytes = 0_u64; + let mut definition_count = 0; + + for resource in resources { + if resource.resource_id.is_empty() || !resource_ids.insert(resource.resource_id.clone()) { + bail!("remote benchmark resource IDs must be non-empty and unique"); + } + let logical_path = validate_logical_path(&resource.logical_path)?; + if !logical_paths.insert(logical_path) { + bail!("remote benchmark resource logical paths must be unique"); + } + let kind = resource + .kind + .as_known() + .with_context(|| format!("resource `{}` has an unknown kind", resource.logical_path))?; + if kind == ResourceKind::Unspecified { + bail!( + "resource `{}` has an unspecified kind", + resource.logical_path + ); + } + if kind == ResourceKind::Definition { + definition_count += 1; + } + validate_download_url(&resource.download_url, allow_plaintext_downloads)?; + validate_sha256( + &format!("resource `{}`", resource.logical_path), + &resource.sha256, + )?; + if resource.size_bytes > MAX_RESOURCE_BYTES { + bail!( + "resource `{}` declares {} bytes; maximum is {MAX_RESOURCE_BYTES}", + resource.logical_path, + resource.size_bytes + ); + } + total_bytes = total_bytes + .checked_add(resource.size_bytes) + .context("remote benchmark resource sizes overflowed")?; + if total_bytes > MAX_BUNDLE_BYTES { + bail!("remote benchmark declares more than {MAX_BUNDLE_BYTES} bytes"); + } + } + + if definition_count != 1 { + bail!("remote benchmark manifest must contain exactly one definition resource"); + } + Ok(resources.to_vec()) +} + +fn validate_logical_path(path: &str) -> Result { + if path.is_empty() || path.contains('\\') { + bail!("resource logical path `{path}` is invalid"); + } + let path = Path::new(path); + if path.is_absolute() + || path + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + bail!( + "resource logical path `{}` must be a safe relative path", + path.display() + ); + } + Ok(path.to_path_buf()) +} + +fn validate_download_url(download_url: &str, allow_plaintext: bool) -> Result<()> { + let url = Url::parse(download_url) + .with_context(|| format!("invalid benchmark resource URL `{download_url}`"))?; + if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() { + bail!("benchmark resource URL must be an absolute http or https URL"); + } + if url.scheme() == "http" && !allow_plaintext { + bail!("an HTTPS benchmark service must return HTTPS resource URLs"); + } + Ok(()) +} + +fn validate_sha256(label: &str, digest: &str) -> Result<()> { + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + bail!("{label} SHA-256 must be 64 lowercase hexadecimal characters"); + } + Ok(()) +} + +async fn download_resources(resources: &[BenchmarkResource]) -> Result>> { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .context("failed to create benchmark resource HTTP client")?; + let mut downloaded = HashMap::with_capacity(resources.len()); + + for resource in resources { + let response = client + .get(&resource.download_url) + .send() + .await + .with_context(|| format!("failed to download resource `{}`", resource.logical_path))? + .error_for_status() + .with_context(|| format!("failed to download resource `{}`", resource.logical_path))?; + let mut bytes = Vec::with_capacity(usize::try_from(resource.size_bytes).unwrap_or(0)); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| { + format!( + "failed while downloading resource `{}`", + resource.logical_path + ) + })?; + let next_len = bytes + .len() + .checked_add(chunk.len()) + .context("downloaded benchmark resource size overflowed")?; + if next_len > usize::try_from(MAX_RESOURCE_BYTES).unwrap_or(usize::MAX) { + bail!( + "resource `{}` exceeded the download size limit", + resource.logical_path + ); + } + bytes.extend_from_slice(&chunk); + } + + if u64::try_from(bytes.len()).ok() != Some(resource.size_bytes) { + bail!( + "resource `{}` size mismatch: expected {}, downloaded {}", + resource.logical_path, + resource.size_bytes, + bytes.len() + ); + } + let actual_sha256 = format!("{:x}", Sha256::digest(&bytes)); + if actual_sha256 != resource.sha256 { + bail!( + "resource `{}` failed SHA-256 verification", + resource.logical_path + ); + } + downloaded.insert(validate_logical_path(&resource.logical_path)?, bytes); + } + + Ok(downloaded) +} + +fn build_remote_benchmark( + benchmark_name: &str, + response: ResolveBenchmarkResponse, + mut downloaded: HashMap>, +) -> Result { + let definition = response + .resources + .iter() + .find(|resource| resource.kind.as_known() == Some(ResourceKind::Definition)) + .context("remote benchmark definition resource is missing")?; + let definition_path = validate_logical_path(&definition.logical_path)?; + let definition_bytes = downloaded + .remove(&definition_path) + .context("downloaded benchmark definition is missing")?; + let definition_toml = std::str::from_utf8(&definition_bytes) + .context("remote benchmark definition is not valid UTF-8")?; + let mut workspace: WorkspaceConfig = toml::from_str(definition_toml) + .context("failed to parse remote benchmark definition as Quantiles TOML")?; + let benchmark = workspace + .benchmarks + .remove(benchmark_name) + .with_context(|| { + format!("remote definition does not contain benchmark `{benchmark_name}`") + })?; + let BenchmarkConfig::CustomNoCode(config) = benchmark else { + bail!("remote benchmark `{benchmark_name}` must have type `custom_nocode`"); + }; + config.params.style.validate()?; + config.params.validate_metrics()?; + + let prompt_path = definition_path + .parent() + .unwrap_or_else(|| Path::new("")) + .join(validate_logical_path(&config.params.prompt_template_file)?); + let prompt_resource = response.resources.iter().find(|resource| { + resource.kind.as_known() == Some(ResourceKind::PromptTemplate) + && validate_logical_path(&resource.logical_path).ok().as_ref() == Some(&prompt_path) + }); + if prompt_resource.is_none() { + bail!( + "remote benchmark prompt template `{}` is not declared as a prompt-template resource", + prompt_path.display() + ); + } + let prompt_bytes = downloaded.remove(&prompt_path).with_context(|| { + format!( + "downloaded prompt template `{}` is missing", + prompt_path.display() + ) + })?; + let prompt_template = String::from_utf8(prompt_bytes).with_context(|| { + format!( + "prompt template `{}` is not valid UTF-8", + prompt_path.display() + ) + })?; + + Ok(RemoteBenchmark { + config: *config, + prompt_template, + version: response.version, + manifest_sha256: response.manifest_sha256, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use buffa::Message as _; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + #[test] + fn remote_url_sources_are_mutually_exclusive() { + let error = select_remote_url( + Some("http://127.0.0.1:8787"), + Some(OsString::from("https://api.quantiles.io")), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot use both")); + } + + #[test] + fn remote_url_defaults_to_production() { + assert_eq!(select_remote_url(None, None).unwrap(), DEFAULT_REMOTE_URL); + } + + #[test] + fn logical_paths_reject_traversal_and_absolute_paths() { + for path in ["../prompt.txt", "/prompt.txt", "prompts/../prompt.txt", ""] { + assert!(validate_logical_path(path).is_err(), "accepted `{path}`"); + } + assert_eq!( + validate_logical_path("prompts/qa.txt").unwrap(), + PathBuf::from("prompts/qa.txt") + ); + } + + #[tokio::test] + async fn resolves_and_downloads_a_custom_nocode_benchmark_in_memory() { + let server = MockServer::start().await; + let definition = br#" + [benchmarks.remote-test] + type = "custom_nocode" + dataset = { name = "quantiles/example" } + prompt_template_file = "prompts/qa.txt" + style = { type = "exact_match", golden_column = "answer" } + "#; + let prompt = b"{{ row.question }}"; + let definition_url = format!("{}/resources/definition", server.uri()); + let prompt_url = format!("{}/resources/prompt", server.uri()); + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-test".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + resources: vec![ + resource( + "definition", + "bundle/quantiles.toml", + ResourceKind::Definition, + &definition_url, + definition, + ), + resource( + "prompt", + "bundle/prompts/qa.txt", + ResourceKind::PromptTemplate, + &prompt_url, + prompt, + ), + ], + ..Default::default() + }; + + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/definition")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(definition)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/prompt")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(prompt)) + .mount(&server) + .await; + + let benchmark = resolve_and_download("remote-test", &server.uri()) + .await + .unwrap() + .unwrap(); + assert_eq!(benchmark.version, "v1"); + assert_eq!(benchmark.prompt_template, "{{ row.question }}"); + assert_eq!(benchmark.config.params.dataset.name, "quantiles/example"); + } + + #[tokio::test] + async fn connect_not_found_is_reported_as_benchmark_absence() { + 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":"benchmark does not exist"}"#, + "application/json", + ), + ) + .mount(&server) + .await; + + assert!( + resolve_and_download("missing", &server.uri()) + .await + .unwrap() + .is_none() + ); + } + + fn resource( + id: &str, + logical_path: &str, + kind: ResourceKind, + download_url: &str, + contents: &[u8], + ) -> BenchmarkResource { + BenchmarkResource { + resource_id: id.to_owned(), + logical_path: logical_path.to_owned(), + kind: kind.into(), + download_url: download_url.to_owned(), + sha256: format!("{:x}", Sha256::digest(contents)), + size_bytes: u64::try_from(contents.len()).unwrap(), + content_type: "application/octet-stream".to_owned(), + ..Default::default() + } + } +} diff --git a/cli/src/builtins/custom_nocode/mod.rs b/cli/src/builtins/custom_nocode/mod.rs index a18152a..b6bb007 100644 --- a/cli/src/builtins/custom_nocode/mod.rs +++ b/cli/src/builtins/custom_nocode/mod.rs @@ -5,7 +5,10 @@ use anyhow::Result; use self::data::{DatasetRow, prepare_row}; use self::evaluation::{EvaluateRowArgs, evaluate_row}; use self::metrics::emit_default_aggregate_metrics; -use self::runtime::{load_template, parse_input, resolve_dataset_limit, resolve_sampler_for_style}; +use self::runtime::{ + LoadedTemplate, load_template, load_template_string, parse_input, resolve_dataset_limit, + resolve_sampler_for_style, +}; use crate::builtins::common::get_max_workers; use crate::builtins::dataset_runner::DatasetRunner; use crate::builtins::output::set_builtin_run_output; @@ -19,13 +22,26 @@ mod runtime; /// No-code custom benchmark builtin. pub struct CustomNoCodeBuiltin { name: String, + prompt_template: Option, } impl CustomNoCodeBuiltin { /// Create a new builtin with the workflow name from the config file. #[must_use] pub fn new(name: String) -> Self { - Self { name } + Self { + name, + prompt_template: None, + } + } + + /// Create a no-code benchmark backed by an in-memory prompt template. + #[must_use] + pub fn with_prompt_template(name: String, prompt_template: String) -> Self { + Self { + name, + prompt_template: Some(prompt_template), + } } } @@ -37,7 +53,13 @@ impl BuiltinWorkflow for CustomNoCodeBuiltin { async fn execute(&self, ctx: BuiltinContext<'_>) -> Result<()> { let mut config = parse_input(ctx.input)?; - let (template_str, env) = load_template(&config.prompt_template_file)?; + let LoadedTemplate { + template: template_str, + environment: env, + } = match &self.prompt_template { + Some(template) => load_template_string(template.clone())?, + None => load_template(&config.prompt_template_file)?, + }; let max_workers = config.max_workers.unwrap_or_else(get_max_workers); let llm = resolve_sampler_for_style(config.model.as_ref(), &config.style)?; diff --git a/cli/src/builtins/custom_nocode/runtime.rs b/cli/src/builtins/custom_nocode/runtime.rs index c538ba4..0e95320 100644 --- a/cli/src/builtins/custom_nocode/runtime.rs +++ b/cli/src/builtins/custom_nocode/runtime.rs @@ -8,6 +8,12 @@ use crate::dataset::DatasetManager; use crate::llm::random::RandomSampler; use crate::llm::random_label::RandomLabelSampler; +/// A validated prompt template and the environment used to render it. +pub(super) struct LoadedTemplate { + pub(super) template: String, + pub(super) environment: jinja::Environment<'static>, +} + /// Resolve the configured sampler, using configured choice labels for random multiple-choice runs. pub(super) fn resolve_sampler_for_style( model: Option<&crate::llm::Sampler>, @@ -38,16 +44,33 @@ pub(super) fn parse_input(input: Option<&str>) -> Result Result<(String, jinja::Environment<'_>)> { +pub(super) fn load_template(path: &str) -> Result { let template_str = std::fs::read_to_string(path) .with_context(|| format!("failed to read prompt template file `{path}`"))?; + let env = validate_template(&template_str, &format!("prompt template file `{path}`"))?; + Ok(LoadedTemplate { + template: template_str, + environment: env, + }) +} + +/// Validate an in-memory prompt template and return its rendering environment. +pub(super) fn load_template_string(template_str: String) -> Result { + let env = validate_template(&template_str, "remote prompt template")?; + Ok(LoadedTemplate { + template: template_str, + environment: env, + }) +} + +fn validate_template(template_str: &str, source: &str) -> Result> { let env = jinja::Environment::new(); env.render_str( - &template_str, + template_str, jinja::context!(row => serde_json::json!({}), choices => Vec::::new()), ) - .with_context(|| format!("invalid jinja syntax in prompt template file `{path}`"))?; - Ok((template_str, env)) + .with_context(|| format!("invalid jinja syntax in {source}"))?; + Ok(env) } /// Initialize the configured dataset and clamp the requested limit to its available row count. diff --git a/cli/src/cli.rs b/cli/src/cli.rs index 4777bfa..7cc8e03 100644 --- a/cli/src/cli.rs +++ b/cli/src/cli.rs @@ -49,6 +49,9 @@ pub enum Command { workflow_name: String, #[arg(long)] input: Option, + /// Override the remote benchmark service URL. + #[arg(long)] + remote_url: Option, /// Emit machine-readable JSON. #[arg(long)] json: bool, diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index c64acaa..b6e0ab1 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -20,9 +20,14 @@ use qt::server::{self, ServerConfig}; pub async fn run( workflow_name: &str, cli_input: Option<&str>, + cli_remote_url: Option<&str>, json: bool, process_start: Instant, ) -> Result<()> { + let remote_url = qt::benchmark_registry::select_remote_url( + cli_remote_url, + std::env::var_os("QUANTILES_REMOTE_URL"), + )?; let config = qt::config::load()?; let bench_config = config.benchmarks.get(workflow_name); @@ -106,7 +111,11 @@ pub async fn run( } } None => { - if builtins::resolve(workflow_name).is_some() { + 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 + } else if builtins::resolve(workflow_name).is_some() { let (effective_input, _) = assemble_builtin_input(None, cli_input); run_builtin_workflow( workflow_name, @@ -122,6 +131,55 @@ pub async fn run( } } +async fn run_remote_benchmark( + workflow_name: &str, + cli_input: Option<&str>, + json: bool, + process_start: Instant, + remote: qt::benchmark_registry::RemoteBenchmark, +) -> Result<()> { + let configured_template_path = remote.config.params.prompt_template_file.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 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!( + "Resolved remote benchmark {workflow_name} version {} ({})", + remote.version, remote.manifest_sha256 + ); + 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(), + )) + }; + execute_builtin(ExecuteBuiltinArgs { + db: &db, + metrics_store: &metrics_store, + run_id, + workflow_name, + builtin, + input: Some(&input), + json, + process_start, + }) + .await +} + fn assemble_builtin_input( bench: Option<&qt::config::BuiltinBenchmarkConfig>, cli_input: Option<&str>, diff --git a/cli/src/lib.rs b/cli/src/lib.rs index b89db48..c4ed883 100644 --- a/cli/src/lib.rs +++ b/cli/src/lib.rs @@ -1,3 +1,4 @@ +pub mod benchmark_registry; pub mod builtins; pub mod client; pub mod config; diff --git a/cli/src/main.rs b/cli/src/main.rs index e542ba9..8e3dd87 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -33,8 +33,18 @@ async fn async_main(process_start: Instant) -> Result<()> { cli::Command::Run { workflow_name, input, + remote_url, json, - } => commands::run(&workflow_name, input.as_deref(), json, process_start).await, + } => { + commands::run( + &workflow_name, + input.as_deref(), + remote_url.as_deref(), + json, + process_start, + ) + .await + } cli::Command::Resume { run_id, json } => { commands::resume(run_id, json, process_start).await } From d1e1796cebfbe3fed07ee70916293326ea75342d Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:24:29 -0700 Subject: [PATCH 05/18] fixup rustls installation --- cli/src/main.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cli/src/main.rs b/cli/src/main.rs index 8e3dd87..22b8d82 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -7,6 +7,12 @@ use anyhow::Result; use clap::Parser; fn main() -> Result<()> { + // `fastembed`'s dependency tree enables another Rustls crypto provider alongside AWS-LC. + // Install AWS-LC explicitly because Rustls may panic when multiple providers are enabled. + connectrpc::rustls::crypto::aws_lc_rs::default_provider() + .install_default() + .map_err(|_| anyhow::anyhow!("failed to install the AWS-LC Rustls crypto provider"))?; + let process_start = Instant::now(); // TODO: allow number of total threads to be configurable, and possibly default From 456ce6612db362f27b53329adc816d31fe9010ad Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:09:44 -0700 Subject: [PATCH 06/18] adding remote hash --- cli/src/commands/resume.rs | 1 + cli/src/commands/run.rs | 41 +++++++++++++++++++++++++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/cli/src/commands/resume.rs b/cli/src/commands/resume.rs index 710e996..cb5b2c0 100644 --- a/cli/src/commands/resume.rs +++ b/cli/src/commands/resume.rs @@ -118,6 +118,7 @@ pub async fn resume(run_id: i64, json: bool, process_start: Instant) -> Result<( input: custom_nocode_input.as_deref().or(stored_input), json, process_start, + remote_hash: None, }) .await } diff --git a/cli/src/commands/run.rs b/cli/src/commands/run.rs index b6e0ab1..195cdc9 100644 --- a/cli/src/commands/run.rs +++ b/cli/src/commands/run.rs @@ -105,6 +105,7 @@ pub async fn run( input: Some(&input), json, process_start, + remote_hash: None, }) .await } @@ -139,6 +140,7 @@ async fn run_remote_benchmark( 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")?; @@ -176,6 +178,7 @@ async fn run_remote_benchmark( input: Some(&input), json, process_start, + remote_hash: Some(&remote_hash), }) .await } @@ -314,6 +317,7 @@ async fn run_builtin_workflow( input, json, process_start, + remote_hash: None, }) .await } @@ -328,6 +332,10 @@ pub struct ExecuteBuiltinArgs<'a> { pub input: Option<&'a str>, pub json: bool, pub process_start: Instant, + /// If the benchmark was loaded from local configuration, this + /// field will be `None`. If it was loaded from remote configuration, + /// it will contain a hash of that configuration. + pub remote_hash: Option<&'a str>, } pub async fn execute_builtin(args: ExecuteBuiltinArgs<'_>) -> Result<()> { @@ -389,6 +397,7 @@ pub async fn execute_builtin(args: ExecuteBuiltinArgs<'_>) -> Result<()> { aggregate_metrics: metrics_map, run_id: args.run_id, warning: None, + remote_hash: args.remote_hash, }; println!("{}", serde_json::to_string(&output)?); } else { @@ -523,11 +532,13 @@ fn print_aggregate_metrics_table(metrics: &[MetricPointSummary]) { /// JSON payload emitted by `qt run --json`. #[derive(Serialize)] -struct BuiltinRunJsonOutput { +struct BuiltinRunJsonOutput<'a> { aggregate_metrics: HashMap, run_id: i64, #[serde(skip_serializing_if = "Option::is_none")] warning: Option, + #[serde(skip_serializing_if = "Option::is_none")] + remote_hash: Option<&'a str>, } /// Config input shape auto-generated from `quantiles.toml` `[benchmarks.*]`. @@ -717,6 +728,34 @@ mod tests { use serde_json::json; + #[test] + fn builtin_run_json_includes_remote_hash_when_resolved_remotely() { + let output = super::BuiltinRunJsonOutput { + aggregate_metrics: HashMap::new(), + run_id: 1, + warning: None, + remote_hash: Some("0123456789abcdef"), + }; + + let serialized = serde_json::to_value(output).unwrap(); + + assert_eq!(serialized["remote_hash"], "0123456789abcdef"); + } + + #[test] + fn builtin_run_json_omits_remote_hash_for_local_runs() { + let output = super::BuiltinRunJsonOutput { + aggregate_metrics: HashMap::new(), + run_id: 1, + warning: None, + remote_hash: None, + }; + + let serialized = serde_json::to_value(output).unwrap(); + + assert!(serialized.get("remote_hash").is_none()); + } + /// When only config input is present and no `--input` CLI flag is given, the result /// should be the exact config object with no overridden keys. #[test] From 730bb323691b49220d52b356c6b8dc01f5c5e80d Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:05:56 -0700 Subject: [PATCH 07/18] comments --- cli/src/builtins/custom_nocode/mod.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cli/src/builtins/custom_nocode/mod.rs b/cli/src/builtins/custom_nocode/mod.rs index b6bb007..191e412 100644 --- a/cli/src/builtins/custom_nocode/mod.rs +++ b/cli/src/builtins/custom_nocode/mod.rs @@ -21,7 +21,10 @@ mod runtime; /// No-code custom benchmark builtin. pub struct CustomNoCodeBuiltin { + /// The name of the custom no-code evaluation name: String, + /// The template used to construct the prompt for a dataset + /// row, if applicable. prompt_template: Option, } From df30704480382f8fc17e5acc7c073bf9a7aba46b Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:07:37 -0700 Subject: [PATCH 08/18] comment --- cli/src/builtins/common.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cli/src/builtins/common.rs b/cli/src/builtins/common.rs index 9e2e22e..b4298df 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -20,7 +20,9 @@ pub(crate) struct BuiltinConfig { /// Number of dataset rows to evaluate. If omitted, the entire dataset is used. #[serde(default)] pub(crate) limit: Option, - /// Dataset source to evaluate. Builtins currently support Hugging Face via `hf://...`. + /// The dataset to use for the evaluation. + /// Currently HuggingFace is the only supported source, and all sources + /// must start with `hf://...` or `huggingface://...` #[serde(default)] pub(crate) dataset: Option, /// Which model sampler to use. If omitted, the builtin chooses a sensible default. From b944fc435589447f48fd22958fab4ec8c62e037c Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:10:31 -0700 Subject: [PATCH 09/18] comments --- cli/src/main.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cli/src/main.rs b/cli/src/main.rs index 22b8d82..62deb0d 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -8,7 +8,8 @@ use clap::Parser; fn main() -> Result<()> { // `fastembed`'s dependency tree enables another Rustls crypto provider alongside AWS-LC. - // Install AWS-LC explicitly because Rustls may panic when multiple providers are enabled. + // Install AWS-LC explicitly, to avoid Rustls-related panics when multiple providers are + // enabled. connectrpc::rustls::crypto::aws_lc_rs::default_provider() .install_default() .map_err(|_| anyhow::anyhow!("failed to install the AWS-LC Rustls crypto provider"))?; From 6824e8f1ac5ccbaaf607904f047da2ccd268aaf3 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:12:28 -0700 Subject: [PATCH 10/18] fixup --- cli/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/README.md b/cli/README.md index 07430a2..c5859c3 100644 --- a/cli/README.md +++ b/cli/README.md @@ -53,7 +53,7 @@ See the [configuration guide](https://quantiles.io/documentation/configuration) ### Remote benchmark fallback -When you run `qt run `, the CLI first looks in the local configuration file for an evaluation called `eval_name`. If one is found, the CLI runs it immediately. If none is found, `qt` looks in 1the Quantiles remote benchmark service for an evaluation called `eval_name` and, if a match is found, downloads the benchmark definition and runs it. +When you run `qt run `, the CLI first looks in the local configuration file for an evaluation called `eval_name`. If one is found, the CLI runs it immediately. If none is found, `qt` looks in the Quantiles remote benchmark service for an evaluation of the same name. If a match is found, the CLI downloads the benchmark definition and runs it. >If you want to override the location of the remote benchmark service, use the `--remote-url` flag or the `QUANTILES_REMOTE_URL` environment variable. From bfa264013ec0c62532b9b06ffdab7fd9de7de6b0 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:13:30 -0700 Subject: [PATCH 11/18] progress --- cli/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/README.md b/cli/README.md index c5859c3..7e8d299 100644 --- a/cli/README.md +++ b/cli/README.md @@ -51,7 +51,7 @@ max_workers = 100 See the [configuration guide](https://quantiles.io/documentation/configuration) for file location, supported fields, validation behavior, and examples. See the [model configuration guide](https://quantiles.io/documentation/model-configuration) for guidance on setting up provider models, managing credentials, and troubleshooting configuration issues. Additional runnable configurations are available in [CLI configuration examples](./examples/configs) and [custom no-code examples](../custom-nocode-examples/quantiles.toml). -### Remote benchmark fallback +### Custom evaluations and the remote benchmark service When you run `qt run `, the CLI first looks in the local configuration file for an evaluation called `eval_name`. If one is found, the CLI runs it immediately. If none is found, `qt` looks in the Quantiles remote benchmark service for an evaluation of the same name. If a match is found, the CLI downloads the benchmark definition and runs it. From 5ea3d0f9c691328f75200f4766d94e67668fe8e0 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:15:29 -0700 Subject: [PATCH 12/18] fixup --- cli/src/benchmark_registry/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index 58519cb..22caac3 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -21,7 +21,7 @@ use crate::config::{BenchmarkConfig, CustomNoCodeBenchmarkConfig, WorkspaceConfi clippy::pedantic, reason = "ConnectRPC and Buffa generated code uses allow attributes" )] -pub mod proto { +mod proto { connectrpc::include_generated!(); } @@ -30,7 +30,7 @@ use proto::quantiles::benchmark::v1::{ ResolveBenchmarkResponse, ResourceKind, }; -pub const DEFAULT_REMOTE_URL: &str = "https://api.quantiles.io"; +pub(crate) const DEFAULT_REMOTE_URL: &str = "https://api.quantiles.io"; const MAX_MANIFEST_BYTES: usize = 1024 * 1024; const MAX_RESOURCE_COUNT: usize = 32; const MAX_RESOURCE_BYTES: u64 = 10 * 1024 * 1024; From d591815cf40d0b01ce8b56d849c6907f9fa83002 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:23:55 -0700 Subject: [PATCH 13/18] splitting up benchmark registry files --- cli/src/benchmark_registry/benchmark.rs | 81 ++++ cli/src/benchmark_registry/client.rs | 142 +++++++ cli/src/benchmark_registry/download.rs | 70 ++++ cli/src/benchmark_registry/manifest.rs | 151 +++++++ cli/src/benchmark_registry/mod.rs | 521 +----------------------- cli/src/benchmark_registry/proto.rs | 10 + cli/src/benchmark_registry/tests.rs | 93 +++++ 7 files changed, 560 insertions(+), 508 deletions(-) create mode 100644 cli/src/benchmark_registry/benchmark.rs create mode 100644 cli/src/benchmark_registry/client.rs create mode 100644 cli/src/benchmark_registry/download.rs create mode 100644 cli/src/benchmark_registry/manifest.rs create mode 100644 cli/src/benchmark_registry/proto.rs create mode 100644 cli/src/benchmark_registry/tests.rs diff --git a/cli/src/benchmark_registry/benchmark.rs b/cli/src/benchmark_registry/benchmark.rs new file mode 100644 index 0000000..c783620 --- /dev/null +++ b/cli/src/benchmark_registry/benchmark.rs @@ -0,0 +1,81 @@ +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, bail}; + +use super::manifest::validate_logical_path; +use super::proto::v1::{ResolveBenchmarkResponse, ResourceKind}; +use crate::config::{BenchmarkConfig, CustomNoCodeBenchmarkConfig, WorkspaceConfig}; + +/// A downloaded benchmark ready to execute without materializing its resources on disk. +pub struct RemoteBenchmark { + pub config: CustomNoCodeBenchmarkConfig, + pub prompt_template: String, + pub version: String, + pub manifest_sha256: String, +} + +pub(super) fn build_remote_benchmark( + benchmark_name: &str, + response: ResolveBenchmarkResponse, + mut downloaded: HashMap>, +) -> Result { + let definition = response + .resources + .iter() + .find(|resource| resource.kind.as_known() == Some(ResourceKind::Definition)) + .context("remote benchmark definition resource is missing")?; + let definition_path = validate_logical_path(&definition.logical_path)?; + let definition_bytes = downloaded + .remove(&definition_path) + .context("downloaded benchmark definition is missing")?; + let definition_toml = std::str::from_utf8(&definition_bytes) + .context("remote benchmark definition is not valid UTF-8")?; + let mut workspace: WorkspaceConfig = toml::from_str(definition_toml) + .context("failed to parse remote benchmark definition as Quantiles TOML")?; + let benchmark = workspace + .benchmarks + .remove(benchmark_name) + .with_context(|| { + format!("remote definition does not contain benchmark `{benchmark_name}`") + })?; + let BenchmarkConfig::CustomNoCode(config) = benchmark else { + bail!("remote benchmark `{benchmark_name}` must have type `custom_nocode`"); + }; + config.params.style.validate()?; + config.params.validate_metrics()?; + + let prompt_path = definition_path + .parent() + .unwrap_or_else(|| Path::new("")) + .join(validate_logical_path(&config.params.prompt_template_file)?); + let prompt_resource = response.resources.iter().find(|resource| { + resource.kind.as_known() == Some(ResourceKind::PromptTemplate) + && validate_logical_path(&resource.logical_path).ok().as_ref() == Some(&prompt_path) + }); + if prompt_resource.is_none() { + bail!( + "remote benchmark prompt template `{}` is not declared as a prompt-template resource", + prompt_path.display() + ); + } + let prompt_bytes = downloaded.remove(&prompt_path).with_context(|| { + format!( + "downloaded prompt template `{}` is missing", + prompt_path.display() + ) + })?; + let prompt_template = String::from_utf8(prompt_bytes).with_context(|| { + format!( + "prompt template `{}` is not valid UTF-8", + prompt_path.display() + ) + })?; + + Ok(RemoteBenchmark { + config: *config, + prompt_template, + version: response.version, + manifest_sha256: response.manifest_sha256, + }) +} diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs new file mode 100644 index 0000000..cc3fc2c --- /dev/null +++ b/cli/src/benchmark_registry/client.rs @@ -0,0 +1,142 @@ +use std::ffi::OsString; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use connectrpc::ErrorCode; +use connectrpc::client::{ClientConfig, HttpClient}; +use reqwest::Url; +use rustls_platform_verifier::ConfigVerifierExt as _; + +use super::proto::v1::{ + BenchmarkRegistryServiceClient, ResolveBenchmarkRequest, ResolveBenchmarkResponse, +}; + +const DEFAULT_REMOTE_URL: &str = "https://api.quantiles.io"; +const MAX_MANIFEST_BYTES: usize = 1024 * 1024; + +/// Select the remote service URL and reject ambiguous CLI/environment configuration. +/// +/// # Errors +/// +/// Returns an error when both sources are set or the environment value is not valid UTF-8. +pub fn select_remote_url(cli_url: Option<&str>, env_url: Option) -> Result { + if cli_url.is_some() && env_url.is_some() { + bail!("cannot use both `--remote-url` and `QUANTILES_REMOTE_URL`"); + } + + if let Some(url) = cli_url { + return Ok(url.to_owned()); + } + + env_url.map_or_else( + || Ok(DEFAULT_REMOTE_URL.to_owned()), + |url| { + url.into_string() + .map_err(|_| anyhow::anyhow!("`QUANTILES_REMOTE_URL` must be valid UTF-8")) + }, + ) +} + +pub(super) fn validate_remote_url(remote_url: &str) -> Result { + let mut url = Url::parse(remote_url) + .with_context(|| format!("invalid remote benchmark service URL `{remote_url}`"))?; + if !matches!(url.scheme(), "http" | "https") { + bail!("remote benchmark service URL must use http or https"); + } + if url.host_str().is_none() { + bail!("remote benchmark service URL must include a host"); + } + if url.query().is_some() || url.fragment().is_some() { + bail!("remote benchmark service URL must not include a query or fragment"); + } + while url.path().ends_with('/') && url.path() != "/" { + let trimmed = url.path().trim_end_matches('/').to_owned(); + url.set_path(&trimmed); + } + Ok(url) +} + +pub(super) async fn resolve_manifest( + benchmark_name: &str, + endpoint: &Url, +) -> Result> { + let uri = endpoint + .as_str() + .parse() + .with_context(|| format!("invalid remote benchmark service URL `{endpoint}`"))?; + let transport = if endpoint.scheme() == "https" { + let tls = connectrpc::rustls::ClientConfig::with_platform_verifier() + .context("failed to configure TLS for remote benchmark service")?; + HttpClient::with_tls(Arc::new(tls)) + } else { + HttpClient::plaintext() + }; + let config = ClientConfig::new(uri) + .with_default_timeout(Duration::from_secs(15)) + .with_default_max_message_size(MAX_MANIFEST_BYTES); + let client = BenchmarkRegistryServiceClient::new(transport, config); + let request = ResolveBenchmarkRequest { + benchmark_name: benchmark_name.to_owned(), + version: String::new(), + ..Default::default() + }; + + match client.resolve_benchmark(request).await { + Ok(response) => Ok(Some(response.into_owned())), + Err(error) if error.code == ErrorCode::NotFound => Ok(None), + Err(error) => Err(error).with_context(|| { + format!("failed to resolve benchmark `{benchmark_name}` from `{endpoint}`") + }), + } +} + +#[cfg(test)] +mod tests { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + #[test] + fn remote_url_sources_are_mutually_exclusive() { + let error = select_remote_url( + Some("http://127.0.0.1:8787"), + Some(OsString::from("https://api.quantiles.io")), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot use both")); + } + + #[test] + fn remote_url_defaults_to_production() { + assert_eq!(select_remote_url(None, None).unwrap(), DEFAULT_REMOTE_URL); + } + + #[tokio::test] + async fn connect_not_found_is_reported_as_benchmark_absence() { + 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":"benchmark does not exist"}"#, + "application/json", + ), + ) + .mount(&server) + .await; + + let endpoint = validate_remote_url(&server.uri()).unwrap(); + assert!( + resolve_manifest("missing", &endpoint) + .await + .unwrap() + .is_none() + ); + } +} diff --git a/cli/src/benchmark_registry/download.rs b/cli/src/benchmark_registry/download.rs new file mode 100644 index 0000000..9c70774 --- /dev/null +++ b/cli/src/benchmark_registry/download.rs @@ -0,0 +1,70 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use futures::StreamExt as _; +use sha2::{Digest as _, Sha256}; + +use super::manifest::{MAX_RESOURCE_BYTES, validate_logical_path}; +use super::proto::v1::BenchmarkResource; + +pub(super) async fn download_resources( + resources: &[BenchmarkResource], +) -> Result>> { + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .context("failed to create benchmark resource HTTP client")?; + let mut downloaded = HashMap::with_capacity(resources.len()); + + for resource in resources { + let response = client + .get(&resource.download_url) + .send() + .await + .with_context(|| format!("failed to download resource `{}`", resource.logical_path))? + .error_for_status() + .with_context(|| format!("failed to download resource `{}`", resource.logical_path))?; + let mut bytes = Vec::with_capacity(usize::try_from(resource.size_bytes).unwrap_or(0)); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| { + format!( + "failed while downloading resource `{}`", + resource.logical_path + ) + })?; + let next_len = bytes + .len() + .checked_add(chunk.len()) + .context("downloaded benchmark resource size overflowed")?; + if next_len > usize::try_from(MAX_RESOURCE_BYTES).unwrap_or(usize::MAX) { + bail!( + "resource `{}` exceeded the download size limit", + resource.logical_path + ); + } + bytes.extend_from_slice(&chunk); + } + + if u64::try_from(bytes.len()).ok() != Some(resource.size_bytes) { + bail!( + "resource `{}` size mismatch: expected {}, downloaded {}", + resource.logical_path, + resource.size_bytes, + bytes.len() + ); + } + let actual_sha256 = format!("{:x}", Sha256::digest(&bytes)); + if actual_sha256 != resource.sha256 { + bail!( + "resource `{}` failed SHA-256 verification", + resource.logical_path + ); + } + downloaded.insert(validate_logical_path(&resource.logical_path)?, bytes); + } + + Ok(downloaded) +} diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs new file mode 100644 index 0000000..f6602b5 --- /dev/null +++ b/cli/src/benchmark_registry/manifest.rs @@ -0,0 +1,151 @@ +use std::collections::HashSet; +use std::path::{Component, Path, PathBuf}; + +use anyhow::{Context, Result, bail}; +use reqwest::Url; + +use super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; + +const MAX_RESOURCE_COUNT: usize = 32; +pub(super) const MAX_RESOURCE_BYTES: u64 = 10 * 1024 * 1024; +const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; + +pub(super) fn validate_response_identity( + benchmark_name: &str, + response: &ResolveBenchmarkResponse, +) -> Result<()> { + if response.benchmark_name != benchmark_name { + bail!( + "remote benchmark response name `{}` does not match requested name `{benchmark_name}`", + response.benchmark_name + ); + } + if response.version.is_empty() { + bail!("remote benchmark response is missing an immutable version"); + } + validate_sha256("manifest", &response.manifest_sha256)?; + Ok(()) +} + +pub(super) fn validate_resources( + resources: &[BenchmarkResource], + allow_plaintext_downloads: bool, +) -> Result> { + if resources.is_empty() { + bail!("remote benchmark manifest contains no resources"); + } + if resources.len() > MAX_RESOURCE_COUNT { + bail!( + "remote benchmark manifest contains {} resources; maximum is {MAX_RESOURCE_COUNT}", + resources.len() + ); + } + + let mut resource_ids = HashSet::new(); + let mut logical_paths = HashSet::new(); + let mut total_bytes = 0_u64; + let mut definition_count = 0; + + for resource in resources { + if resource.resource_id.is_empty() || !resource_ids.insert(resource.resource_id.clone()) { + bail!("remote benchmark resource IDs must be non-empty and unique"); + } + let logical_path = validate_logical_path(&resource.logical_path)?; + if !logical_paths.insert(logical_path) { + bail!("remote benchmark resource logical paths must be unique"); + } + let kind = resource + .kind + .as_known() + .with_context(|| format!("resource `{}` has an unknown kind", resource.logical_path))?; + if kind == ResourceKind::Unspecified { + bail!( + "resource `{}` has an unspecified kind", + resource.logical_path + ); + } + if kind == ResourceKind::Definition { + definition_count += 1; + } + validate_download_url(&resource.download_url, allow_plaintext_downloads)?; + validate_sha256( + &format!("resource `{}`", resource.logical_path), + &resource.sha256, + )?; + if resource.size_bytes > MAX_RESOURCE_BYTES { + bail!( + "resource `{}` declares {} bytes; maximum is {MAX_RESOURCE_BYTES}", + resource.logical_path, + resource.size_bytes + ); + } + total_bytes = total_bytes + .checked_add(resource.size_bytes) + .context("remote benchmark resource sizes overflowed")?; + if total_bytes > MAX_BUNDLE_BYTES { + bail!("remote benchmark declares more than {MAX_BUNDLE_BYTES} bytes"); + } + } + + if definition_count != 1 { + bail!("remote benchmark manifest must contain exactly one definition resource"); + } + Ok(resources.to_vec()) +} + +pub(super) fn validate_logical_path(path: &str) -> Result { + if path.is_empty() || path.contains('\\') { + bail!("resource logical path `{path}` is invalid"); + } + let path = Path::new(path); + if path.is_absolute() + || path + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + bail!( + "resource logical path `{}` must be a safe relative path", + path.display() + ); + } + Ok(path.to_path_buf()) +} + +fn validate_download_url(download_url: &str, allow_plaintext: bool) -> Result<()> { + let url = Url::parse(download_url) + .with_context(|| format!("invalid benchmark resource URL `{download_url}`"))?; + if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() { + bail!("benchmark resource URL must be an absolute http or https URL"); + } + if url.scheme() == "http" && !allow_plaintext { + bail!("an HTTPS benchmark service must return HTTPS resource URLs"); + } + Ok(()) +} + +fn validate_sha256(label: &str, digest: &str) -> Result<()> { + if digest.len() != 64 + || !digest + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + { + bail!("{label} SHA-256 must be 64 lowercase hexadecimal characters"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn logical_paths_reject_traversal_and_absolute_paths() { + for path in ["../prompt.txt", "/prompt.txt", "prompts/../prompt.txt", ""] { + assert!(validate_logical_path(path).is_err(), "accepted `{path}`"); + } + assert_eq!( + validate_logical_path("prompts/qa.txt").unwrap(), + PathBuf::from("prompts/qa.txt") + ); + } +} diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index 22caac3..e8d14c8 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -1,71 +1,19 @@ //! Client support for resolving remotely hosted benchmark definitions. -use std::collections::{HashMap, HashSet}; -use std::ffi::OsString; -use std::path::{Component, Path, PathBuf}; -use std::sync::Arc; -use std::time::Duration; +use anyhow::Result; -use anyhow::{Context, Result, bail}; -use connectrpc::ErrorCode; -use connectrpc::client::{ClientConfig, HttpClient}; -use futures::StreamExt as _; -use reqwest::Url; -use rustls_platform_verifier::ConfigVerifierExt as _; -use sha2::{Digest as _, Sha256}; +pub use self::benchmark::RemoteBenchmark; +use self::benchmark::build_remote_benchmark; +pub use self::client::select_remote_url; +use self::client::{resolve_manifest, validate_remote_url}; +use self::download::download_resources; +use self::manifest::{validate_resources, validate_response_identity}; -use crate::config::{BenchmarkConfig, CustomNoCodeBenchmarkConfig, WorkspaceConfig}; - -#[expect( - clippy::allow_attributes, - clippy::pedantic, - reason = "ConnectRPC and Buffa generated code uses allow attributes" -)] -mod proto { - connectrpc::include_generated!(); -} - -use proto::quantiles::benchmark::v1::{ - BenchmarkRegistryServiceClient, BenchmarkResource, ResolveBenchmarkRequest, - ResolveBenchmarkResponse, ResourceKind, -}; - -pub(crate) const DEFAULT_REMOTE_URL: &str = "https://api.quantiles.io"; -const MAX_MANIFEST_BYTES: usize = 1024 * 1024; -const MAX_RESOURCE_COUNT: usize = 32; -const MAX_RESOURCE_BYTES: u64 = 10 * 1024 * 1024; -const MAX_BUNDLE_BYTES: u64 = 50 * 1024 * 1024; - -/// A downloaded benchmark ready to execute without materializing its resources on disk. -pub struct RemoteBenchmark { - pub config: CustomNoCodeBenchmarkConfig, - pub prompt_template: String, - pub version: String, - pub manifest_sha256: String, -} - -/// Select the remote service URL and reject ambiguous CLI/environment configuration. -/// -/// # Errors -/// -/// Returns an error when both sources are set or the environment value is not valid UTF-8. -pub fn select_remote_url(cli_url: Option<&str>, env_url: Option) -> Result { - if cli_url.is_some() && env_url.is_some() { - bail!("cannot use both `--remote-url` and `QUANTILES_REMOTE_URL`"); - } - - if let Some(url) = cli_url { - return Ok(url.to_owned()); - } - - env_url.map_or_else( - || Ok(DEFAULT_REMOTE_URL.to_owned()), - |url| { - url.into_string() - .map_err(|_| anyhow::anyhow!("`QUANTILES_REMOTE_URL` must be valid UTF-8")) - }, - ) -} +mod benchmark; +mod client; +mod download; +mod manifest; +mod proto; /// Resolve a benchmark and download all of its resources into memory. /// @@ -92,448 +40,5 @@ pub async fn resolve_and_download( Ok(Some(remote)) } -fn validate_remote_url(remote_url: &str) -> Result { - let mut url = Url::parse(remote_url) - .with_context(|| format!("invalid remote benchmark service URL `{remote_url}`"))?; - if !matches!(url.scheme(), "http" | "https") { - bail!("remote benchmark service URL must use http or https"); - } - if url.host_str().is_none() { - bail!("remote benchmark service URL must include a host"); - } - if url.query().is_some() || url.fragment().is_some() { - bail!("remote benchmark service URL must not include a query or fragment"); - } - while url.path().ends_with('/') && url.path() != "/" { - let trimmed = url.path().trim_end_matches('/').to_owned(); - url.set_path(&trimmed); - } - Ok(url) -} - -async fn resolve_manifest( - benchmark_name: &str, - endpoint: &Url, -) -> Result> { - let uri = endpoint - .as_str() - .parse() - .with_context(|| format!("invalid remote benchmark service URL `{endpoint}`"))?; - let transport = if endpoint.scheme() == "https" { - let tls = connectrpc::rustls::ClientConfig::with_platform_verifier() - .context("failed to configure TLS for remote benchmark service")?; - HttpClient::with_tls(Arc::new(tls)) - } else { - HttpClient::plaintext() - }; - let config = ClientConfig::new(uri) - .with_default_timeout(Duration::from_secs(15)) - .with_default_max_message_size(MAX_MANIFEST_BYTES); - let client = BenchmarkRegistryServiceClient::new(transport, config); - let request = ResolveBenchmarkRequest { - benchmark_name: benchmark_name.to_owned(), - version: String::new(), - ..Default::default() - }; - - match client.resolve_benchmark(request).await { - Ok(response) => Ok(Some(response.into_owned())), - Err(error) if error.code == ErrorCode::NotFound => Ok(None), - Err(error) => Err(error).with_context(|| { - format!("failed to resolve benchmark `{benchmark_name}` from `{endpoint}`") - }), - } -} - -fn validate_response_identity( - benchmark_name: &str, - response: &ResolveBenchmarkResponse, -) -> Result<()> { - if response.benchmark_name != benchmark_name { - bail!( - "remote benchmark response name `{}` does not match requested name `{benchmark_name}`", - response.benchmark_name - ); - } - if response.version.is_empty() { - bail!("remote benchmark response is missing an immutable version"); - } - validate_sha256("manifest", &response.manifest_sha256)?; - Ok(()) -} - -fn validate_resources( - resources: &[BenchmarkResource], - allow_plaintext_downloads: bool, -) -> Result> { - if resources.is_empty() { - bail!("remote benchmark manifest contains no resources"); - } - if resources.len() > MAX_RESOURCE_COUNT { - bail!( - "remote benchmark manifest contains {} resources; maximum is {MAX_RESOURCE_COUNT}", - resources.len() - ); - } - - let mut resource_ids = HashSet::new(); - let mut logical_paths = HashSet::new(); - let mut total_bytes = 0_u64; - let mut definition_count = 0; - - for resource in resources { - if resource.resource_id.is_empty() || !resource_ids.insert(resource.resource_id.clone()) { - bail!("remote benchmark resource IDs must be non-empty and unique"); - } - let logical_path = validate_logical_path(&resource.logical_path)?; - if !logical_paths.insert(logical_path) { - bail!("remote benchmark resource logical paths must be unique"); - } - let kind = resource - .kind - .as_known() - .with_context(|| format!("resource `{}` has an unknown kind", resource.logical_path))?; - if kind == ResourceKind::Unspecified { - bail!( - "resource `{}` has an unspecified kind", - resource.logical_path - ); - } - if kind == ResourceKind::Definition { - definition_count += 1; - } - validate_download_url(&resource.download_url, allow_plaintext_downloads)?; - validate_sha256( - &format!("resource `{}`", resource.logical_path), - &resource.sha256, - )?; - if resource.size_bytes > MAX_RESOURCE_BYTES { - bail!( - "resource `{}` declares {} bytes; maximum is {MAX_RESOURCE_BYTES}", - resource.logical_path, - resource.size_bytes - ); - } - total_bytes = total_bytes - .checked_add(resource.size_bytes) - .context("remote benchmark resource sizes overflowed")?; - if total_bytes > MAX_BUNDLE_BYTES { - bail!("remote benchmark declares more than {MAX_BUNDLE_BYTES} bytes"); - } - } - - if definition_count != 1 { - bail!("remote benchmark manifest must contain exactly one definition resource"); - } - Ok(resources.to_vec()) -} - -fn validate_logical_path(path: &str) -> Result { - if path.is_empty() || path.contains('\\') { - bail!("resource logical path `{path}` is invalid"); - } - let path = Path::new(path); - if path.is_absolute() - || path - .components() - .any(|component| !matches!(component, Component::Normal(_))) - { - bail!( - "resource logical path `{}` must be a safe relative path", - path.display() - ); - } - Ok(path.to_path_buf()) -} - -fn validate_download_url(download_url: &str, allow_plaintext: bool) -> Result<()> { - let url = Url::parse(download_url) - .with_context(|| format!("invalid benchmark resource URL `{download_url}`"))?; - if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() { - bail!("benchmark resource URL must be an absolute http or https URL"); - } - if url.scheme() == "http" && !allow_plaintext { - bail!("an HTTPS benchmark service must return HTTPS resource URLs"); - } - Ok(()) -} - -fn validate_sha256(label: &str, digest: &str) -> Result<()> { - if digest.len() != 64 - || !digest - .bytes() - .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) - { - bail!("{label} SHA-256 must be 64 lowercase hexadecimal characters"); - } - Ok(()) -} - -async fn download_resources(resources: &[BenchmarkResource]) -> Result>> { - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .build() - .context("failed to create benchmark resource HTTP client")?; - let mut downloaded = HashMap::with_capacity(resources.len()); - - for resource in resources { - let response = client - .get(&resource.download_url) - .send() - .await - .with_context(|| format!("failed to download resource `{}`", resource.logical_path))? - .error_for_status() - .with_context(|| format!("failed to download resource `{}`", resource.logical_path))?; - let mut bytes = Vec::with_capacity(usize::try_from(resource.size_bytes).unwrap_or(0)); - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.with_context(|| { - format!( - "failed while downloading resource `{}`", - resource.logical_path - ) - })?; - let next_len = bytes - .len() - .checked_add(chunk.len()) - .context("downloaded benchmark resource size overflowed")?; - if next_len > usize::try_from(MAX_RESOURCE_BYTES).unwrap_or(usize::MAX) { - bail!( - "resource `{}` exceeded the download size limit", - resource.logical_path - ); - } - bytes.extend_from_slice(&chunk); - } - - if u64::try_from(bytes.len()).ok() != Some(resource.size_bytes) { - bail!( - "resource `{}` size mismatch: expected {}, downloaded {}", - resource.logical_path, - resource.size_bytes, - bytes.len() - ); - } - let actual_sha256 = format!("{:x}", Sha256::digest(&bytes)); - if actual_sha256 != resource.sha256 { - bail!( - "resource `{}` failed SHA-256 verification", - resource.logical_path - ); - } - downloaded.insert(validate_logical_path(&resource.logical_path)?, bytes); - } - - Ok(downloaded) -} - -fn build_remote_benchmark( - benchmark_name: &str, - response: ResolveBenchmarkResponse, - mut downloaded: HashMap>, -) -> Result { - let definition = response - .resources - .iter() - .find(|resource| resource.kind.as_known() == Some(ResourceKind::Definition)) - .context("remote benchmark definition resource is missing")?; - let definition_path = validate_logical_path(&definition.logical_path)?; - let definition_bytes = downloaded - .remove(&definition_path) - .context("downloaded benchmark definition is missing")?; - let definition_toml = std::str::from_utf8(&definition_bytes) - .context("remote benchmark definition is not valid UTF-8")?; - let mut workspace: WorkspaceConfig = toml::from_str(definition_toml) - .context("failed to parse remote benchmark definition as Quantiles TOML")?; - let benchmark = workspace - .benchmarks - .remove(benchmark_name) - .with_context(|| { - format!("remote definition does not contain benchmark `{benchmark_name}`") - })?; - let BenchmarkConfig::CustomNoCode(config) = benchmark else { - bail!("remote benchmark `{benchmark_name}` must have type `custom_nocode`"); - }; - config.params.style.validate()?; - config.params.validate_metrics()?; - - let prompt_path = definition_path - .parent() - .unwrap_or_else(|| Path::new("")) - .join(validate_logical_path(&config.params.prompt_template_file)?); - let prompt_resource = response.resources.iter().find(|resource| { - resource.kind.as_known() == Some(ResourceKind::PromptTemplate) - && validate_logical_path(&resource.logical_path).ok().as_ref() == Some(&prompt_path) - }); - if prompt_resource.is_none() { - bail!( - "remote benchmark prompt template `{}` is not declared as a prompt-template resource", - prompt_path.display() - ); - } - let prompt_bytes = downloaded.remove(&prompt_path).with_context(|| { - format!( - "downloaded prompt template `{}` is missing", - prompt_path.display() - ) - })?; - let prompt_template = String::from_utf8(prompt_bytes).with_context(|| { - format!( - "prompt template `{}` is not valid UTF-8", - prompt_path.display() - ) - })?; - - Ok(RemoteBenchmark { - config: *config, - prompt_template, - version: response.version, - manifest_sha256: response.manifest_sha256, - }) -} - #[cfg(test)] -mod tests { - use super::*; - use buffa::Message as _; - use wiremock::matchers::{method, path}; - use wiremock::{Mock, MockServer, ResponseTemplate}; - - #[test] - fn remote_url_sources_are_mutually_exclusive() { - let error = select_remote_url( - Some("http://127.0.0.1:8787"), - Some(OsString::from("https://api.quantiles.io")), - ) - .unwrap_err(); - assert!(error.to_string().contains("cannot use both")); - } - - #[test] - fn remote_url_defaults_to_production() { - assert_eq!(select_remote_url(None, None).unwrap(), DEFAULT_REMOTE_URL); - } - - #[test] - fn logical_paths_reject_traversal_and_absolute_paths() { - for path in ["../prompt.txt", "/prompt.txt", "prompts/../prompt.txt", ""] { - assert!(validate_logical_path(path).is_err(), "accepted `{path}`"); - } - assert_eq!( - validate_logical_path("prompts/qa.txt").unwrap(), - PathBuf::from("prompts/qa.txt") - ); - } - - #[tokio::test] - async fn resolves_and_downloads_a_custom_nocode_benchmark_in_memory() { - let server = MockServer::start().await; - let definition = br#" - [benchmarks.remote-test] - type = "custom_nocode" - dataset = { name = "quantiles/example" } - prompt_template_file = "prompts/qa.txt" - style = { type = "exact_match", golden_column = "answer" } - "#; - let prompt = b"{{ row.question }}"; - let definition_url = format!("{}/resources/definition", server.uri()); - let prompt_url = format!("{}/resources/prompt", server.uri()); - let response = ResolveBenchmarkResponse { - benchmark_name: "remote-test".to_owned(), - version: "v1".to_owned(), - manifest_sha256: "a".repeat(64), - resources: vec![ - resource( - "definition", - "bundle/quantiles.toml", - ResourceKind::Definition, - &definition_url, - definition, - ), - resource( - "prompt", - "bundle/prompts/qa.txt", - ResourceKind::PromptTemplate, - &prompt_url, - prompt, - ), - ], - ..Default::default() - }; - - Mock::given(method("POST")) - .and(path( - "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", - )) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "application/proto") - .set_body_bytes(response.encode_to_vec()), - ) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/resources/definition")) - .respond_with(ResponseTemplate::new(200).set_body_bytes(definition)) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/resources/prompt")) - .respond_with(ResponseTemplate::new(200).set_body_bytes(prompt)) - .mount(&server) - .await; - - let benchmark = resolve_and_download("remote-test", &server.uri()) - .await - .unwrap() - .unwrap(); - assert_eq!(benchmark.version, "v1"); - assert_eq!(benchmark.prompt_template, "{{ row.question }}"); - assert_eq!(benchmark.config.params.dataset.name, "quantiles/example"); - } - - #[tokio::test] - async fn connect_not_found_is_reported_as_benchmark_absence() { - 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":"benchmark does not exist"}"#, - "application/json", - ), - ) - .mount(&server) - .await; - - assert!( - resolve_and_download("missing", &server.uri()) - .await - .unwrap() - .is_none() - ); - } - - fn resource( - id: &str, - logical_path: &str, - kind: ResourceKind, - download_url: &str, - contents: &[u8], - ) -> BenchmarkResource { - BenchmarkResource { - resource_id: id.to_owned(), - logical_path: logical_path.to_owned(), - kind: kind.into(), - download_url: download_url.to_owned(), - sha256: format!("{:x}", Sha256::digest(contents)), - size_bytes: u64::try_from(contents.len()).unwrap(), - content_type: "application/octet-stream".to_owned(), - ..Default::default() - } - } -} +mod tests; diff --git a/cli/src/benchmark_registry/proto.rs b/cli/src/benchmark_registry/proto.rs new file mode 100644 index 0000000..f12c02b --- /dev/null +++ b/cli/src/benchmark_registry/proto.rs @@ -0,0 +1,10 @@ +#[expect( + clippy::allow_attributes, + clippy::pedantic, + reason = "ConnectRPC and Buffa generated code uses allow attributes" +)] +pub(super) mod generated { + connectrpc::include_generated!(); +} + +pub(super) use generated::quantiles::benchmark::v1; diff --git a/cli/src/benchmark_registry/tests.rs b/cli/src/benchmark_registry/tests.rs new file mode 100644 index 0000000..5b20acb --- /dev/null +++ b/cli/src/benchmark_registry/tests.rs @@ -0,0 +1,93 @@ +use buffa::Message as _; +use sha2::{Digest as _, Sha256}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; +use super::resolve_and_download; + +#[tokio::test] +async fn resolves_and_downloads_a_custom_nocode_benchmark_in_memory() { + let server = MockServer::start().await; + let definition = br#" + [benchmarks.remote-test] + type = "custom_nocode" + dataset = { name = "quantiles/example" } + prompt_template_file = "prompts/qa.txt" + style = { type = "exact_match", golden_column = "answer" } + "#; + let prompt = b"{{ row.question }}"; + let definition_url = format!("{}/resources/definition", server.uri()); + let prompt_url = format!("{}/resources/prompt", server.uri()); + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-test".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + resources: vec![ + resource( + "definition", + "bundle/quantiles.toml", + ResourceKind::Definition, + &definition_url, + definition, + ), + resource( + "prompt", + "bundle/prompts/qa.txt", + ResourceKind::PromptTemplate, + &prompt_url, + prompt, + ), + ], + ..Default::default() + }; + + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/definition")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(definition)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/prompt")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(prompt)) + .mount(&server) + .await; + + let benchmark = resolve_and_download("remote-test", &server.uri()) + .await + .unwrap() + .unwrap(); + assert_eq!(benchmark.version, "v1"); + assert_eq!(benchmark.prompt_template, "{{ row.question }}"); + assert_eq!(benchmark.config.params.dataset.name, "quantiles/example"); +} + +fn resource( + id: &str, + logical_path: &str, + kind: ResourceKind, + download_url: &str, + contents: &[u8], +) -> BenchmarkResource { + BenchmarkResource { + resource_id: id.to_owned(), + logical_path: logical_path.to_owned(), + kind: kind.into(), + download_url: download_url.to_owned(), + sha256: format!("{:x}", Sha256::digest(contents)), + size_bytes: u64::try_from(contents.len()).unwrap(), + content_type: "application/octet-stream".to_owned(), + ..Default::default() + } +} From 0214bd9dfdf5f5a97954700d62def05f3e16402f Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:24:16 -0700 Subject: [PATCH 14/18] removing unnecessary code --- cli/src/benchmark_registry/mod.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index e8d14c8..a5e136b 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -39,6 +39,3 @@ pub async fn resolve_and_download( let remote = build_remote_benchmark(benchmark_name, response, downloaded)?; Ok(Some(remote)) } - -#[cfg(test)] -mod tests; From 09885c7ea6dfa7732bdb8d9b721981c2afbeddf6 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:24:55 -0700 Subject: [PATCH 15/18] adding some docs etc --- cli/src/benchmark_registry/benchmark.rs | 1 + cli/src/benchmark_registry/client.rs | 2 ++ cli/src/benchmark_registry/download.rs | 1 + cli/src/benchmark_registry/manifest.rs | 7 +++++++ 4 files changed, 11 insertions(+) diff --git a/cli/src/benchmark_registry/benchmark.rs b/cli/src/benchmark_registry/benchmark.rs index c783620..e4b7017 100644 --- a/cli/src/benchmark_registry/benchmark.rs +++ b/cli/src/benchmark_registry/benchmark.rs @@ -15,6 +15,7 @@ pub struct RemoteBenchmark { pub manifest_sha256: String, } +/// Assemble a runnable no-code benchmark from its manifest and downloaded resources. pub(super) fn build_remote_benchmark( benchmark_name: &str, response: ResolveBenchmarkResponse, diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs index cc3fc2c..05d8770 100644 --- a/cli/src/benchmark_registry/client.rs +++ b/cli/src/benchmark_registry/client.rs @@ -38,6 +38,7 @@ pub fn select_remote_url(cli_url: Option<&str>, env_url: Option) -> Re ) } +/// Parse and validate a remote benchmark service URL. pub(super) fn validate_remote_url(remote_url: &str) -> Result { let mut url = Url::parse(remote_url) .with_context(|| format!("invalid remote benchmark service URL `{remote_url}`"))?; @@ -57,6 +58,7 @@ pub(super) fn validate_remote_url(remote_url: &str) -> Result { Ok(url) } +/// Resolve benchmark metadata from the remote ConnectRPC service. pub(super) async fn resolve_manifest( benchmark_name: &str, endpoint: &Url, diff --git a/cli/src/benchmark_registry/download.rs b/cli/src/benchmark_registry/download.rs index 9c70774..303c434 100644 --- a/cli/src/benchmark_registry/download.rs +++ b/cli/src/benchmark_registry/download.rs @@ -9,6 +9,7 @@ use sha2::{Digest as _, Sha256}; use super::manifest::{MAX_RESOURCE_BYTES, validate_logical_path}; use super::proto::v1::BenchmarkResource; +/// Download validated resources into memory and verify their sizes and digests. pub(super) async fn download_resources( resources: &[BenchmarkResource], ) -> Result>> { diff --git a/cli/src/benchmark_registry/manifest.rs b/cli/src/benchmark_registry/manifest.rs index f6602b5..e88daec 100644 --- a/cli/src/benchmark_registry/manifest.rs +++ b/cli/src/benchmark_registry/manifest.rs @@ -6,10 +6,14 @@ use reqwest::Url; use super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; +/// Maximum number of resources allowed in a remote benchmark manifest. const MAX_RESOURCE_COUNT: usize = 32; +/// Maximum declared and downloaded size of one benchmark resource. pub(super) const MAX_RESOURCE_BYTES: u64 = 10 * 1024 * 1024; +/// Maximum aggregate declared size of all resources in a benchmark manifest. 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, response: &ResolveBenchmarkResponse, @@ -27,6 +31,7 @@ pub(super) fn validate_response_identity( Ok(()) } +/// Validate a manifest's resource identities, locations, kinds, digests, and sizes. pub(super) fn validate_resources( resources: &[BenchmarkResource], allow_plaintext_downloads: bool, @@ -93,6 +98,7 @@ pub(super) fn validate_resources( Ok(resources.to_vec()) } +/// Validate and normalize a resource path relative to its benchmark bundle. pub(super) fn validate_logical_path(path: &str) -> Result { if path.is_empty() || path.contains('\\') { bail!("resource logical path `{path}` is invalid"); @@ -111,6 +117,7 @@ pub(super) fn validate_logical_path(path: &str) -> Result { Ok(path.to_path_buf()) } +/// Validate a resource URL and enforce the benchmark service's transport security. fn validate_download_url(download_url: &str, allow_plaintext: bool) -> Result<()> { let url = Url::parse(download_url) .with_context(|| format!("invalid benchmark resource URL `{download_url}`"))?; From 787c5ad9ab225f90739b1cb49f0d295ef383f11c Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:28:42 -0700 Subject: [PATCH 16/18] progress --- cli/src/benchmark_registry/benchmark.rs | 124 ++++++++++++------------ cli/src/benchmark_registry/client.rs | 2 +- cli/src/benchmark_registry/mod.rs | 3 +- cli/src/builtins/common.rs | 2 +- 4 files changed, 66 insertions(+), 65 deletions(-) diff --git a/cli/src/benchmark_registry/benchmark.rs b/cli/src/benchmark_registry/benchmark.rs index e4b7017..cb501ce 100644 --- a/cli/src/benchmark_registry/benchmark.rs +++ b/cli/src/benchmark_registry/benchmark.rs @@ -15,68 +15,70 @@ pub struct RemoteBenchmark { pub manifest_sha256: String, } -/// Assemble a runnable no-code benchmark from its manifest and downloaded resources. -pub(super) fn build_remote_benchmark( - benchmark_name: &str, - response: ResolveBenchmarkResponse, - mut downloaded: HashMap>, -) -> Result { - let definition = response - .resources - .iter() - .find(|resource| resource.kind.as_known() == Some(ResourceKind::Definition)) - .context("remote benchmark definition resource is missing")?; - let definition_path = validate_logical_path(&definition.logical_path)?; - let definition_bytes = downloaded - .remove(&definition_path) - .context("downloaded benchmark definition is missing")?; - let definition_toml = std::str::from_utf8(&definition_bytes) - .context("remote benchmark definition is not valid UTF-8")?; - let mut workspace: WorkspaceConfig = toml::from_str(definition_toml) - .context("failed to parse remote benchmark definition as Quantiles TOML")?; - let benchmark = workspace - .benchmarks - .remove(benchmark_name) - .with_context(|| { - format!("remote definition does not contain benchmark `{benchmark_name}`") +impl RemoteBenchmark { + /// Assemble a runnable no-code benchmark from its manifest and downloaded resources. + pub(super) fn new( + benchmark_name: &str, + response: ResolveBenchmarkResponse, + mut downloaded: HashMap>, + ) -> Result { + let definition = response + .resources + .iter() + .find(|resource| resource.kind.as_known() == Some(ResourceKind::Definition)) + .context("remote benchmark definition resource is missing")?; + let definition_path = validate_logical_path(&definition.logical_path)?; + let definition_bytes = downloaded + .remove(&definition_path) + .context("downloaded benchmark definition is missing")?; + let definition_toml = std::str::from_utf8(&definition_bytes) + .context("remote benchmark definition is not valid UTF-8")?; + let mut workspace: WorkspaceConfig = toml::from_str(definition_toml) + .context("failed to parse remote benchmark definition as Quantiles TOML")?; + let benchmark = workspace + .benchmarks + .remove(benchmark_name) + .with_context(|| { + format!("remote definition does not contain benchmark `{benchmark_name}`") + })?; + let BenchmarkConfig::CustomNoCode(config) = benchmark else { + bail!("remote benchmark `{benchmark_name}` must have type `custom_nocode`"); + }; + config.params.style.validate()?; + config.params.validate_metrics()?; + + let prompt_path = definition_path + .parent() + .unwrap_or_else(|| Path::new("")) + .join(validate_logical_path(&config.params.prompt_template_file)?); + let prompt_resource = response.resources.iter().find(|resource| { + resource.kind.as_known() == Some(ResourceKind::PromptTemplate) + && validate_logical_path(&resource.logical_path).ok().as_ref() == Some(&prompt_path) + }); + if prompt_resource.is_none() { + bail!( + "remote benchmark prompt template `{}` is not declared as a prompt-template resource", + prompt_path.display() + ); + } + let prompt_bytes = downloaded.remove(&prompt_path).with_context(|| { + format!( + "downloaded prompt template `{}` is missing", + prompt_path.display() + ) + })?; + let prompt_template = String::from_utf8(prompt_bytes).with_context(|| { + format!( + "prompt template `{}` is not valid UTF-8", + prompt_path.display() + ) })?; - let BenchmarkConfig::CustomNoCode(config) = benchmark else { - bail!("remote benchmark `{benchmark_name}` must have type `custom_nocode`"); - }; - config.params.style.validate()?; - config.params.validate_metrics()?; - let prompt_path = definition_path - .parent() - .unwrap_or_else(|| Path::new("")) - .join(validate_logical_path(&config.params.prompt_template_file)?); - let prompt_resource = response.resources.iter().find(|resource| { - resource.kind.as_known() == Some(ResourceKind::PromptTemplate) - && validate_logical_path(&resource.logical_path).ok().as_ref() == Some(&prompt_path) - }); - if prompt_resource.is_none() { - bail!( - "remote benchmark prompt template `{}` is not declared as a prompt-template resource", - prompt_path.display() - ); + Ok(RemoteBenchmark { + config: *config, + prompt_template, + version: response.version, + manifest_sha256: response.manifest_sha256, + }) } - let prompt_bytes = downloaded.remove(&prompt_path).with_context(|| { - format!( - "downloaded prompt template `{}` is missing", - prompt_path.display() - ) - })?; - let prompt_template = String::from_utf8(prompt_bytes).with_context(|| { - format!( - "prompt template `{}` is not valid UTF-8", - prompt_path.display() - ) - })?; - - Ok(RemoteBenchmark { - config: *config, - prompt_template, - version: response.version, - manifest_sha256: response.manifest_sha256, - }) } diff --git a/cli/src/benchmark_registry/client.rs b/cli/src/benchmark_registry/client.rs index 05d8770..a9623b9 100644 --- a/cli/src/benchmark_registry/client.rs +++ b/cli/src/benchmark_registry/client.rs @@ -58,7 +58,7 @@ pub(super) fn validate_remote_url(remote_url: &str) -> Result { Ok(url) } -/// Resolve benchmark metadata from the remote ConnectRPC service. +/// Resolve benchmark metadata from the remote `ConnectRPC` service. pub(super) async fn resolve_manifest( benchmark_name: &str, endpoint: &Url, diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index a5e136b..16d73fe 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -3,7 +3,6 @@ use anyhow::Result; pub use self::benchmark::RemoteBenchmark; -use self::benchmark::build_remote_benchmark; pub use self::client::select_remote_url; use self::client::{resolve_manifest, validate_remote_url}; use self::download::download_resources; @@ -36,6 +35,6 @@ pub async fn resolve_and_download( validate_response_identity(benchmark_name, &response)?; let resources = validate_resources(&response.resources, endpoint.scheme() == "http")?; let downloaded = download_resources(&resources).await?; - let remote = build_remote_benchmark(benchmark_name, response, downloaded)?; + let remote = RemoteBenchmark::new(benchmark_name, response, downloaded)?; Ok(Some(remote)) } diff --git a/cli/src/builtins/common.rs b/cli/src/builtins/common.rs index b4298df..75fc999 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -21,7 +21,7 @@ pub(crate) struct BuiltinConfig { #[serde(default)] pub(crate) limit: Option, /// The dataset to use for the evaluation. - /// Currently HuggingFace is the only supported source, and all sources + /// Currently `HuggingFace` is the only supported source, and all sources /// must start with `hf://...` or `huggingface://...` #[serde(default)] pub(crate) dataset: Option, From cfbb7e9ab681846cb5e7cef419aefea7295cb105 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:32:16 -0700 Subject: [PATCH 17/18] test fixup --- cli/src/benchmark_registry/mod.rs | 32 +------ cli/src/benchmark_registry/resolver.rs | 128 +++++++++++++++++++++++++ cli/src/benchmark_registry/tests.rs | 93 ------------------ 3 files changed, 130 insertions(+), 123 deletions(-) create mode 100644 cli/src/benchmark_registry/resolver.rs delete mode 100644 cli/src/benchmark_registry/tests.rs diff --git a/cli/src/benchmark_registry/mod.rs b/cli/src/benchmark_registry/mod.rs index 16d73fe..6a373f0 100644 --- a/cli/src/benchmark_registry/mod.rs +++ b/cli/src/benchmark_registry/mod.rs @@ -1,40 +1,12 @@ //! Client support for resolving remotely hosted benchmark definitions. -use anyhow::Result; - pub use self::benchmark::RemoteBenchmark; pub use self::client::select_remote_url; -use self::client::{resolve_manifest, validate_remote_url}; -use self::download::download_resources; -use self::manifest::{validate_resources, validate_response_identity}; +pub use self::resolver::resolve_and_download; mod benchmark; mod client; mod download; mod manifest; mod proto; - -/// Resolve a benchmark 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. -/// -/// # Errors -/// -/// Returns an error for invalid endpoints, RPC failures, malformed manifests, failed downloads, -/// digest mismatches, invalid UTF-8, or invalid no-code benchmark definitions. -pub async fn resolve_and_download( - benchmark_name: &str, - remote_url: &str, -) -> Result> { - let endpoint = validate_remote_url(remote_url)?; - let Some(response) = resolve_manifest(benchmark_name, &endpoint).await? else { - return Ok(None); - }; - - validate_response_identity(benchmark_name, &response)?; - let resources = validate_resources(&response.resources, endpoint.scheme() == "http")?; - let downloaded = download_resources(&resources).await?; - let remote = RemoteBenchmark::new(benchmark_name, response, downloaded)?; - Ok(Some(remote)) -} +mod resolver; diff --git a/cli/src/benchmark_registry/resolver.rs b/cli/src/benchmark_registry/resolver.rs new file mode 100644 index 0000000..86a1408 --- /dev/null +++ b/cli/src/benchmark_registry/resolver.rs @@ -0,0 +1,128 @@ +use anyhow::Result; + +use super::RemoteBenchmark; +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. +/// +/// `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. +/// +/// # Errors +/// +/// Returns an error for invalid endpoints, RPC failures, malformed manifests, failed downloads, +/// digest mismatches, invalid UTF-8, or invalid no-code benchmark definitions. +pub async fn resolve_and_download( + benchmark_name: &str, + remote_url: &str, +) -> Result> { + let endpoint = validate_remote_url(remote_url)?; + let Some(response) = resolve_manifest(benchmark_name, &endpoint).await? else { + return Ok(None); + }; + + validate_response_identity(benchmark_name, &response)?; + let resources = validate_resources(&response.resources, endpoint.scheme() == "http")?; + let downloaded = download_resources(&resources).await?; + let remote = RemoteBenchmark::new(benchmark_name, response, downloaded)?; + Ok(Some(remote)) +} + +#[cfg(test)] +mod tests { + use buffa::Message as _; + use sha2::{Digest as _, Sha256}; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; + use super::*; + + #[tokio::test] + async fn resolves_and_downloads_a_custom_nocode_benchmark_in_memory() { + let server = MockServer::start().await; + let definition = br#" + [benchmarks.remote-test] + type = "custom_nocode" + dataset = { name = "quantiles/example" } + prompt_template_file = "prompts/qa.txt" + style = { type = "exact_match", golden_column = "answer" } + "#; + let prompt = b"{{ row.question }}"; + let definition_url = format!("{}/resources/definition", server.uri()); + let prompt_url = format!("{}/resources/prompt", server.uri()); + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-test".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + resources: vec![ + resource( + "definition", + "bundle/quantiles.toml", + ResourceKind::Definition, + &definition_url, + definition, + ), + resource( + "prompt", + "bundle/prompts/qa.txt", + ResourceKind::PromptTemplate, + &prompt_url, + prompt, + ), + ], + ..Default::default() + }; + + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/definition")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(definition)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/resources/prompt")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(prompt)) + .mount(&server) + .await; + + let benchmark = resolve_and_download("remote-test", &server.uri()) + .await + .unwrap() + .unwrap(); + assert_eq!(benchmark.version, "v1"); + assert_eq!(benchmark.prompt_template, "{{ row.question }}"); + assert_eq!(benchmark.config.params.dataset.name, "quantiles/example"); + } + + fn resource( + id: &str, + logical_path: &str, + kind: ResourceKind, + download_url: &str, + contents: &[u8], + ) -> BenchmarkResource { + BenchmarkResource { + resource_id: id.to_owned(), + logical_path: logical_path.to_owned(), + kind: kind.into(), + download_url: download_url.to_owned(), + sha256: format!("{:x}", Sha256::digest(contents)), + size_bytes: u64::try_from(contents.len()).unwrap(), + content_type: "application/octet-stream".to_owned(), + ..Default::default() + } + } +} diff --git a/cli/src/benchmark_registry/tests.rs b/cli/src/benchmark_registry/tests.rs deleted file mode 100644 index 5b20acb..0000000 --- a/cli/src/benchmark_registry/tests.rs +++ /dev/null @@ -1,93 +0,0 @@ -use buffa::Message as _; -use sha2::{Digest as _, Sha256}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use super::proto::v1::{BenchmarkResource, ResolveBenchmarkResponse, ResourceKind}; -use super::resolve_and_download; - -#[tokio::test] -async fn resolves_and_downloads_a_custom_nocode_benchmark_in_memory() { - let server = MockServer::start().await; - let definition = br#" - [benchmarks.remote-test] - type = "custom_nocode" - dataset = { name = "quantiles/example" } - prompt_template_file = "prompts/qa.txt" - style = { type = "exact_match", golden_column = "answer" } - "#; - let prompt = b"{{ row.question }}"; - let definition_url = format!("{}/resources/definition", server.uri()); - let prompt_url = format!("{}/resources/prompt", server.uri()); - let response = ResolveBenchmarkResponse { - benchmark_name: "remote-test".to_owned(), - version: "v1".to_owned(), - manifest_sha256: "a".repeat(64), - resources: vec![ - resource( - "definition", - "bundle/quantiles.toml", - ResourceKind::Definition, - &definition_url, - definition, - ), - resource( - "prompt", - "bundle/prompts/qa.txt", - ResourceKind::PromptTemplate, - &prompt_url, - prompt, - ), - ], - ..Default::default() - }; - - Mock::given(method("POST")) - .and(path( - "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", - )) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "application/proto") - .set_body_bytes(response.encode_to_vec()), - ) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/resources/definition")) - .respond_with(ResponseTemplate::new(200).set_body_bytes(definition)) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/resources/prompt")) - .respond_with(ResponseTemplate::new(200).set_body_bytes(prompt)) - .mount(&server) - .await; - - let benchmark = resolve_and_download("remote-test", &server.uri()) - .await - .unwrap() - .unwrap(); - assert_eq!(benchmark.version, "v1"); - assert_eq!(benchmark.prompt_template, "{{ row.question }}"); - assert_eq!(benchmark.config.params.dataset.name, "quantiles/example"); -} - -fn resource( - id: &str, - logical_path: &str, - kind: ResourceKind, - download_url: &str, - contents: &[u8], -) -> BenchmarkResource { - BenchmarkResource { - resource_id: id.to_owned(), - logical_path: logical_path.to_owned(), - kind: kind.into(), - download_url: download_url.to_owned(), - sha256: format!("{:x}", Sha256::digest(contents)), - size_bytes: u64::try_from(contents.len()).unwrap(), - content_type: "application/octet-stream".to_owned(), - ..Default::default() - } -} From fd33d7c8745c36c97c92fd2611c5b554be2ce3d7 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:44:11 -0700 Subject: [PATCH 18/18] progress --- .github/workflows/cli-release.yml | 6 ++++++ .github/workflows/cli.yml | 3 ++- mise.toml | 2 ++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cli-release.yml b/.github/workflows/cli-release.yml index 17084fa..3557424 100644 --- a/.github/workflows/cli-release.yml +++ b/.github/workflows/cli-release.yml @@ -4,6 +4,7 @@ on: pull_request: paths: - "cli/**" + - "mise.toml" # This line allows us to test at least the build part of this # workflow definition if/when we change it. Below, we make sure # not to publish the CLI from pull requests, though. @@ -57,6 +58,11 @@ jobs: - uses: Swatinem/rust-cache@v2 with: workspaces: cli + - name: Install Buf + uses: jdx/mise-action@v2 + with: + version: 2026.6.11 + install_args: buf - name: Build release binary run: cargo build --locked --release --target ${{ matrix.target }} - name: Package release artifact diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 7e2bc82..b37d270 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -23,7 +23,8 @@ jobs: - uses: Swatinem/rust-cache@v2 with: workspaces: . - - uses: jdx/mise-action@v2 + - name: Install mise-managed tools, including Buf + uses: jdx/mise-action@v2 with: version: 2026.6.11 # TODO: consider cache: true diff --git a/mise.toml b/mise.toml index c966c9e..c684e2f 100644 --- a/mise.toml +++ b/mise.toml @@ -4,6 +4,8 @@ monorepo_root = true config_roots = [".", "cli", "typescript", "python", "python-examples"] [tools] +# needed by the `.use_buf()` call in cli/build.rs +buf = "1.72.0" # keep this version inline with the version declared in .github/workflows/typescript bun = "1.3.14" # keep this version inline with rust stable