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/.gitignore b/.gitignore index 76bb1e4..6e05153 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ .DS_Store /.quantiles node_modules +.fastembed_cache diff --git a/README.md b/README.md index 2ce5994..36ead5c 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,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 3c83a5f..762b6a8 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" @@ -852,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" @@ -1033,6 +1081,70 @@ 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" +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" @@ -2226,6 +2338,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -2558,6 +2671,7 @@ dependencies = [ "allocator-api2", "equivalent", "foldhash 0.1.5", + "serde", ] [[package]] @@ -4080,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" @@ -4337,8 +4471,11 @@ dependencies = [ "assert_cmd", "async-trait", "axum", + "buffa", "clap", "comfy-table", + "connectrpc", + "connectrpc-build", "datafusion", "dirs", "fastembed", @@ -4350,6 +4487,7 @@ dependencies = [ "rand 0.8.6", "reqwest 0.13.4", "rstest", + "rustls-platform-verifier", "sea-orm", "seq-macro", "serde", @@ -5454,6 +5592,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" @@ -6230,6 +6377,7 @@ dependencies = [ "pin-project-lite", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -7362,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 631f074..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"] } @@ -38,6 +41,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/README.md b/cli/README.md index 0f02450..7e8d299 100644 --- a/cli/README.md +++ b/cli/README.md @@ -43,6 +43,7 @@ Add a `quantiles.toml` or `.quantiles.toml` file to configure an evaluation. For ```toml [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" samples = 50 model = "openai:gpt-5.6" max_workers = 100 @@ -50,6 +51,14 @@ 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). +### 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. + +>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/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/examples/configs/anthropic/quantiles.toml b/cli/examples/configs/anthropic/quantiles.toml index 66d0544..950c435 100644 --- a/cli/examples/configs/anthropic/quantiles.toml +++ b/cli/examples/configs/anthropic/quantiles.toml @@ -4,9 +4,12 @@ # 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 1,000 samples. Because Anthropic charges for -# usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the -# full benchmark. +# usage, this configuration entry limits the run to 50 samples to control costs. +# +# Omit the below `samples` key to run the full benchmark. samples = 50 # This example uses a relatively low-cost, fast model. Model details are available at: 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 42dee7e..e017700 100644 --- a/cli/examples/configs/cloudflare/quantiles.toml +++ b/cli/examples/configs/cloudflare/quantiles.toml @@ -12,6 +12,8 @@ # https://dash.cloudflare.com//ai/ai-gateway [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because Cloudflare charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/gemini/quantiles.toml b/cli/examples/configs/gemini/quantiles.toml index d10b1fd..82f851a 100644 --- a/cli/examples/configs/gemini/quantiles.toml +++ b/cli/examples/configs/gemini/quantiles.toml @@ -6,6 +6,8 @@ # https://aistudio.google.com/prompts/new_chat [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because Google charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/openai/quantiles.toml b/cli/examples/configs/openai/quantiles.toml index eef0f56..ba41ed6 100644 --- a/cli/examples/configs/openai/quantiles.toml +++ b/cli/examples/configs/openai/quantiles.toml @@ -7,6 +7,8 @@ # https://platform.openai.com/home [benchmarks.pubmedqa] +dataset = "hf://quantiles/PubMedQA" + # The full expert-labeled PubMedQA benchmark has 1,000 samples. Because OpenAI charges for # usage, this example limits the run to 50 samples to control costs. Omit `samples` to run the # full benchmark. diff --git a/cli/examples/configs/simple/quantiles.toml b/cli/examples/configs/simple/quantiles.toml index f6ff43e..1a0b772 100644 --- a/cli/examples/configs/simple/quantiles.toml +++ b/cli/examples/configs/simple/quantiles.toml @@ -8,7 +8,9 @@ # 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/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; +} diff --git a/cli/src/benchmark_registry/benchmark.rs b/cli/src/benchmark_registry/benchmark.rs new file mode 100644 index 0000000..cb501ce --- /dev/null +++ b/cli/src/benchmark_registry/benchmark.rs @@ -0,0 +1,84 @@ +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, +} + +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() + ) + })?; + + 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..a9623b9 --- /dev/null +++ b/cli/src/benchmark_registry/client.rs @@ -0,0 +1,144 @@ +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")) + }, + ) +} + +/// 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}`"))?; + 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) +} + +/// Resolve benchmark metadata from the remote `ConnectRPC` service. +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..303c434 --- /dev/null +++ b/cli/src/benchmark_registry/download.rs @@ -0,0 +1,71 @@ +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; + +/// Download validated resources into memory and verify their sizes and digests. +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..e88daec --- /dev/null +++ b/cli/src/benchmark_registry/manifest.rs @@ -0,0 +1,158 @@ +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}; + +/// 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, +) -> 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(()) +} + +/// Validate a manifest's resource identities, locations, kinds, digests, and sizes. +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()) +} + +/// 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"); + } + 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()) +} + +/// 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}`"))?; + 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 new file mode 100644 index 0000000..6a373f0 --- /dev/null +++ b/cli/src/benchmark_registry/mod.rs @@ -0,0 +1,12 @@ +//! Client support for resolving remotely hosted benchmark definitions. + +pub use self::benchmark::RemoteBenchmark; +pub use self::client::select_remote_url; +pub use self::resolver::resolve_and_download; + +mod benchmark; +mod client; +mod download; +mod manifest; +mod proto; +mod resolver; 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/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/builtins/common.rs b/cli/src/builtins/common.rs index 374fc26..75fc999 100644 --- a/cli/src/builtins/common.rs +++ b/cli/src/builtins/common.rs @@ -20,6 +20,11 @@ pub(crate) struct BuiltinConfig { /// Number of dataset rows to evaluate. If omitted, the entire dataset is used. #[serde(default)] pub(crate) limit: Option, + /// 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. #[serde(default)] pub(crate) model: Option, diff --git a/cli/src/builtins/custom_nocode/mod.rs b/cli/src/builtins/custom_nocode/mod.rs index a109cda..191e412 100644 --- a/cli/src/builtins/custom_nocode/mod.rs +++ b/cli/src/builtins/custom_nocode/mod.rs @@ -6,7 +6,8 @@ use self::data::{DatasetRow, prepare_row}; use self::evaluation::{EvaluateRowArgs, evaluate_row}; use self::metrics::emit_default_aggregate_metrics; use self::runtime::{ - LoadedTemplate, load_template, parse_input, resolve_dataset_limit, resolve_sampler_for_style, + 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; @@ -20,14 +21,30 @@ 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, } 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), + } } } @@ -42,7 +59,10 @@ impl BuiltinWorkflow for CustomNoCodeBuiltin { let LoadedTemplate { template: template_str, environment: env, - } = load_template(&config.prompt_template_file)?; + } = 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 f0eef9f..7d1fad8 100644 --- a/cli/src/builtins/custom_nocode/runtime.rs +++ b/cli/src/builtins/custom_nocode/runtime.rs @@ -7,7 +7,8 @@ use crate::dataset::DatasetManager; use crate::llm::random::RandomSampler; use crate::llm::random_label::RandomLabelSampler; -/// A validated prompt template and its Jinja environment. +/// A validated prompt template and the Jinja environment +/// used to render it. pub(super) struct LoadedTemplate { /// The prompt template source. pub(super) template: String, @@ -55,6 +56,15 @@ pub(super) fn load_template(path: &str) -> Result { }) } +/// 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 mut env = jinja::Environment::new(); env.add_template_owned(source.to_owned(), template_str.to_owned()) 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 4e0df77..4385097 100644 --- a/cli/src/builtins/pubmedqa/mod.rs +++ b/cli/src/builtins/pubmedqa/mod.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::random_label::RandomLabelSampler; use config::{PubMedQAConfig, RowOutput}; @@ -23,6 +23,9 @@ 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 { fn name(&self) -> String { @@ -48,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?; @@ -61,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 1c8fe0f..b4e3a01 100644 --- a/cli/src/builtins/similarity.rs +++ b/cli/src/builtins/similarity.rs @@ -9,7 +9,8 @@ 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::random::RandomSampler; use crate::similarity::{ SimilarityMetric, SimilarityMetricName, levenshtein::LevenshteinSimilarity, @@ -44,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, } @@ -52,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", }; @@ -60,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", }; @@ -94,7 +95,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 @@ -107,6 +114,7 @@ impl BuiltinWorkflow for SimilarityBenchmark { set_builtin_run_input( db, run_id, + dataset_source, config.base.model.as_ref(), limit, config.base.max_workers, @@ -118,7 +126,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/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/resume.rs b/cli/src/commands/resume.rs index ef8033b..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 } @@ -147,6 +148,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: None, + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }); @@ -161,6 +163,7 @@ mod tests { let bench = qt::config::BenchmarkConfig::Builtin(qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + 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 aeb0f07..195cdc9 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); @@ -100,13 +105,18 @@ pub async fn run( input: Some(&input), json, process_start, + remote_hash: None, }) .await } } } 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 +132,57 @@ 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 remote_hash = remote.manifest_sha256.clone(); + let input = assemble_custom_nocode_input(&remote.config, cli_input)?; + let effective_params: qt::config::CustomNoCodeParams = serde_json::from_str(&input) + .context("failed to parse assembled remote custom_nocode input")?; + + let 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, + remote_hash: Some(&remote_hash), + }) + .await +} + fn assemble_builtin_input( bench: Option<&qt::config::BuiltinBenchmarkConfig>, cli_input: Option<&str>, @@ -131,12 +192,9 @@ fn assemble_builtin_input( } if let Some(bench) = bench { - if bench.samples.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, }; @@ -259,6 +317,7 @@ async fn run_builtin_workflow( input, json, process_start, + remote_hash: None, }) .await } @@ -273,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<()> { @@ -334,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 { @@ -468,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.*]`. @@ -480,6 +546,7 @@ struct BuiltinRunJsonOutput { struct BuiltinConfigInput { #[serde(skip_serializing_if = "Option::is_none")] limit: Option, + dataset: String, #[serde(skip_serializing_if = "Option::is_none")] model: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -661,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] @@ -721,6 +816,7 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(10), + dataset: "hf://quantiles/PubMedQA".to_owned(), model: None, max_workers: None, }; @@ -735,28 +831,32 @@ mod tests { let bench = qt::config::BuiltinBenchmarkConfig { type_: "builtin".to_owned(), samples: Some(5), + dataset: "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); } - /// 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: "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/mod.rs b/cli/src/config/mod.rs index f51de57..686df37 100644 --- a/cli/src/config/mod.rs +++ b/cli/src/config/mod.rs @@ -96,6 +96,8 @@ pub struct BuiltinBenchmarkConfig { pub type_: String, /// Number of samples (rows) to evaluate. pub samples: Option, + /// Dataset source for this benchmark. + pub dataset: String, /// Which model sampler to use for this benchmark. pub model: Option, /// Maximum number of concurrent workers for this benchmark. @@ -182,6 +184,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(); @@ -190,6 +193,7 @@ mod tests { if let BenchmarkConfig::Builtin(b) = bench { assert_eq!(b.type_, "builtin"); assert_eq!(b.samples, Some(10)); + assert_eq!(b.dataset, "hf://quantiles/demo"); assert!(b.model.is_none()); } } @@ -200,11 +204,25 @@ 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, "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] diff --git a/cli/src/dataset/mod.rs b/cli/src/dataset/mod.rs index e726d87..d144681 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; @@ -156,3 +156,65 @@ 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. +/// +/// # Errors +/// +/// If the given `source` is missing a `hf://` / `huggingface://` prefix, or the +/// source is otherwise unresolvable, this function returns a descriptive error. +pub fn resolve_hf_dataset_source(source: &str) -> Result<&str> { + if let Some(dataset_id) = source + .strip_prefix("hf://") + .or_else(|| source.strip_prefix("huggingface://")) + { + if dataset_id.is_empty() { + 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://...` or `huggingface://...`"); + } else { + bail!("dataset source `{source}` is missing required `hf://` or `huggingface://` 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_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(); + 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://` or `huggingface://` prefix") + ); + } +} 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..62deb0d 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -7,6 +7,13 @@ 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, 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"))?; + let process_start = Instant::now(); // TODO: allow number of total threads to be configurable, and possibly default @@ -33,8 +40,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 } 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); 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