diff --git a/rust/Cargo.lock b/rust/Cargo.lock index de6c97654b75..ff24484a630c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -15,6 +15,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "const-random", "getrandom 0.3.4", "once_cell", "serde", @@ -31,6 +32,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -379,6 +395,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -450,6 +472,27 @@ dependencies = [ "generic-array", ] +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bstr" version = "1.12.1" @@ -671,6 +714,26 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + [[package]] name = "constant_time_eq" version = "0.4.2" @@ -1584,6 +1647,7 @@ dependencies = [ "bytemuck", "cfg-if", "crunchy", + "num-traits", "zerocopy", ] @@ -1617,6 +1681,12 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -2270,6 +2340,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lz4_flex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" +dependencies = [ + "twox-hash", +] + [[package]] name = "logos" version = "0.15.1" @@ -2648,6 +2727,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -2866,6 +2955,31 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" +dependencies = [ + "ahash", + "base64 0.23.1", + "brotli", + "bytes", + "chrono", + "flate2", + "half", + "hashbrown 0.17.1", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "seq-macro", + "serde_json", + "snap", + "twox-hash", + "zstd", +] + [[package]] name = "paste" version = "1.0.15" @@ -4044,6 +4158,12 @@ version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -4316,6 +4436,12 @@ dependencies = [ "version_check", ] +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + [[package]] name = "socket2" version = "0.6.3" @@ -5341,6 +5467,12 @@ dependencies = [ "utf-8", ] +[[package]] +name = "twox-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" + [[package]] name = "typenum" version = "1.19.0" @@ -5636,6 +5768,7 @@ dependencies = [ "image", "indicatif", "mimalloc", + "parquet", "rand 0.9.2", "rand_distr", "rayon", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 3e0b7b4365f2..3725ebd3c462 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -70,6 +70,7 @@ openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmon openai-protocol = "1.6.0" openssl = "0.10" parking_lot = "0.12.5" +parquet = { version = "59.2.0", default-features = false, features = ["json", "snap", "zstd", "flate2", "flate2-rust_backend", "lz4", "brotli"] } paste = "1.0.15" prometheus-client = "0.24.0" prometheus-client-derive-encode = "0.5.0" diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md index 852138060c61..1f15165fa429 100644 --- a/rust/src/bench/AGENTS.md +++ b/rust/src/bench/AGENTS.md @@ -4,7 +4,7 @@ Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime. -Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend. +Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (hf-hub's async reqwest/native-tls API) instead of hf-hub's ureq backend. ## Build & Test @@ -34,7 +34,7 @@ cargo test -p vllm-bench -- --ignored - `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`) - `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback - `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source) -- `src/hub.rs` — `HubRepo`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans +- `src/hub.rs` — `HubRepo`: thin wrapper over hf-hub's async (reqwest/native-tls) API — file download and repo listing, optional revision; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans - `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up - `src/ready_checker.rs` — Endpoint readiness with retry - `src/backends/` — Backend implementations (enum dispatch, not trait objects) @@ -48,7 +48,7 @@ cargo test -p vllm-bench -- --ignored - `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching - `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`) - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) -- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection) +- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (parquet shards via hf-hub into the standard HF hub cache; datasets-server for config/split discovery; native main-branch parquet/JSON fallback for private datasets; column auto-detection) - `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template) - `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`) - `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`) diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 24a8f12e910c..5a3b14894ab8 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -17,6 +17,7 @@ hf-hub.workspace = true image.workspace = true indicatif.workspace = true mimalloc.workspace = true +parquet.workspace = true rand.workspace = true rand_distr.workspace = true rayon.workspace = true diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index ec58c40e5a3a..ff5b6ebfd55f 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -148,7 +148,8 @@ vllm-bench \ --sonnet-input-len 550 --sonnet-output-len 150 --sonnet-prefix-len 200 \ --num-prompts 500 -# Any public HuggingFace dataset (auto-downloads, auto-detects columns) +# Any HuggingFace dataset (downloads parquet shards into the standard HF hub +# cache — shared with `hf download` and Python tooling — auto-detects columns) vllm-bench \ --backend openai-chat --base-url http://127.0.0.1:8000 --model \ --dataset-name hf --dataset-path allenai/WildChat-4.8M \ @@ -160,7 +161,8 @@ vllm-bench \ --dataset-name hf --dataset-path THUDM/LongBench \ --hf-subset narrativeqa --hf-split test --hf-output-len 512 --num-prompts 200 -# Gated HuggingFace dataset (requires HF_TOKEN) +# Gated HuggingFace dataset (auth via HF_TOKEN or the `hf auth login` token file). +# Private datasets work too when the repo stores native parquet/json/jsonl files. HF_TOKEN=hf_xxx vllm-bench \ --backend openai-chat --base-url http://127.0.0.1:8000 --model \ --dataset-name hf --dataset-path lmsys/lmsys-chat-1m \ @@ -431,7 +433,7 @@ Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Us | `sharegpt` | Real conversations from ShareGPT (auto-downloads from HuggingFace, or use `--dataset-path`) | | `sonnet` | Built-in Shakespeare sonnets; controllable token length + shared prefix, no dataset file needed | | `speed-bench` | NVIDIA SPEED-Bench for speculative decoding evaluation (auto-downloads, 11 categories) | -| `hf` | Any HuggingFace dataset (auto-downloads via datasets-server API, auto-detects chat/text columns) | +| `hf` | Any HuggingFace dataset with parquet data (downloads parquet shards via hf-hub into the standard HF hub cache, auto-detects chat/text columns) | ## Metrics @@ -507,7 +509,7 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below. | `--random-input-len` | `1024` | Input token length | | `--random-output-len` | `128` | Output token length | | `--random-prefix-len` | `0` | Shared prefix length | -| `--random-range-ratio` | `1.0` | Length jitter, range `(0, 1]`. Lengths sampled from `[ratio × target, target]`; `1.0` = fixed length | +| `--random-range-ratio` | `0.0` | Length jitter `r`, range `[0, 1)`. Lengths sampled uniformly from `[len×(1−r), len×(1+r)]`; `0.0` = exact target lengths. Also accepts `'{"input": r1, "output": r2}'` for independent input/output jitter | | `--prompt-token-ids` | `false` | Send prompts as token-ID arrays (skips server-side tokenization, exact counts). Random dataset only | | **Random multimodal** | | | | `--random-mm-base-items-per-request` | `1` | Base number of multimodal items (images) per request | @@ -795,6 +797,7 @@ The Rust implementation matches Python `vllm bench serve` in: - Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up) - Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput) - Sampling parameters merged into the request body via `extra_body` (same precedence rules) +- HF dataset prompt construction (token-identical input totals vs. Python on full GSM8k / MT-Bench splits; Python-only per-dataset chat templating is intentionally not applied — prompts are sent raw) - CLI flags: a superset of Python's, enforced by `tests/cli_parity.rs` against a snapshot of the Python parser (`tests/python_serve_flags.txt`, regenerated by `tests/benchmarks/test_rust_bench_cli_parity.py` in the repo root) with an @@ -806,7 +809,8 @@ The Rust implementation matches Python `vllm bench serve` in: | Variable | Description | | ---------- | ------------- | | `OPENAI_API_KEY` | API key for authenticated endpoints (cached, not read per-request) | -| `HF_TOKEN` | HuggingFace token for gated model tokenizers and gated datasets | +| `HF_TOKEN` | HuggingFace token for gated/private tokenizers and datasets (falls back to the `hf auth login` token file) | +| `HF_HOME` | Overrides the HuggingFace cache location (default `~/.cache/huggingface`) used for tokenizers and dataset parquet shards | | `TOKIO_WORKER_THREADS` | Override tokio worker thread count (default: physical cores) | ## License diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index feb49cf84d91..54b6e8701f27 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -534,17 +534,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let dataset_id = config.dataset_path.as_deref().ok_or_else(|| { BenchError::Config("--dataset-path is required for --dataset-name hf".into()) })?; - let (downloaded_path, _config, _split) = - crate::datasets::hf_dataset::download_hf_dataset( - dataset_id, - config.hf_subset.as_deref(), - config.hf_split.as_deref(), - config.num_prompts, - ) - .await?; + let (entries, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset( + dataset_id, + config.hf_subset.as_deref(), + config.hf_split.as_deref(), + config.num_prompts, + config.seed, + config.disable_shuffle, + ) + .await?; crate::datasets::hf_dataset::load_hf_dataset( tok, - &downloaded_path, + &entries, config.num_prompts, config.hf_output_len, config.seed, diff --git a/rust/src/bench/src/datasets/hf_dataset.rs b/rust/src/bench/src/datasets/hf_dataset.rs index 93c9f4cdbd2a..c0505d076e31 100644 --- a/rust/src/bench/src/datasets/hf_dataset.rs +++ b/rust/src/bench/src/datasets/hf_dataset.rs @@ -1,38 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::path::Path; use std::sync::Arc; +use parquet::file::reader::{FileReader, SerializedFileReader}; use rand::rngs::StdRng; use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use super::SampleRequest; -use super::progress::RowDownloadReporter; use crate::error::{BenchError, Result}; use crate::tokenizer::TokenizerKind; -/// Cache directory for downloaded HF datasets. -fn cache_dir() -> std::path::PathBuf { - dirs::cache_dir() - .unwrap_or_else(|| std::path::PathBuf::from("/tmp")) - .join("vllm-bench") - .join("datasets") -} - -/// Sanitize a dataset name for use in filenames (replace `/` and other unsafe chars). -fn sanitize_name(name: &str) -> String { - name.chars() - .map(|c| { - if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' { - c - } else { - '_' - } - }) - .collect() -} - /// Detected column format for extracting prompts from HF dataset rows. enum ColumnFormat { /// Column contains chat messages (array of {role, content} or {from, value}). @@ -103,28 +83,39 @@ async fn get_with_retry( /// Download an arbitrary HuggingFace dataset via the datasets-server REST API. /// -/// Returns `(cache_path, resolved_config, resolved_split)`. +/// Returns `(rows, resolved_config, resolved_split)`. /// /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization. -/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected -/// or the dataset is exhausted. +/// Rows come from data files downloaded via hf-hub into the standard HF hub cache: the +/// dataset-viewer's `refs/convert/parquet` export when available, otherwise native +/// parquet/JSON/JSONL files on the main branch (e.g. private datasets the viewer does +/// not process). Shard order is shuffled with `seed` (unless `disable_shuffle`) so +/// samples are not always drawn from the start of the split, and up to twice +/// `num_rows_needed` rows are read so downstream validity filtering can backfill from +/// real rows. pub async fn download_hf_dataset( dataset: &str, subset: Option<&str>, split: Option<&str>, num_rows_needed: usize, -) -> Result<(String, String, String)> { + seed: u64, + disable_shuffle: bool, +) -> Result<(Vec, String, String)> { let encoded_dataset: String = url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); let mut client_builder = reqwest::Client::builder().timeout(std::time::Duration::from_secs(120)); - // Add HF_TOKEN auth header if available - if let Ok(token) = std::env::var("HF_TOKEN") { + // Match hf-hub's auth for shard downloads: HF_TOKEN env var, then the hub token file + let token = std::env::var("HF_TOKEN") + .ok() + .filter(|t| !t.is_empty()) + .or_else(|| hf_hub::Cache::from_env().token()); + if let Some(token) = token { let mut headers = reqwest::header::HeaderMap::new(); let header_value = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}")) - .map_err(|e| BenchError::Config(format!("Invalid HF_TOKEN: {e}")))?; + .map_err(|e| BenchError::Config(format!("Invalid HF token: {e}")))?; headers.insert(reqwest::header::AUTHORIZATION, header_value); client_builder = client_builder.default_headers(headers); } @@ -137,12 +128,24 @@ pub async fn download_hf_dataset( let (resolved_config, resolved_split) = if let (Some(cfg), Some(spl)) = (subset, split) { // Both provided — skip /info call (cfg.to_string(), spl.to_string()) - } else { - // Call /info to discover available configs and splits + } else if let Some(info) = { + // Call /info to discover available configs and splits. Private datasets are not + // processed by the dataset viewer; fall back to datasets-library defaults. let info_url = format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}"); - let info = get_with_retry(&client, &info_url, "HF dataset /info").await?; - + match get_with_retry(&client, &info_url, "HF dataset /info").await { + Ok(info) => Some(info), + Err(e) => { + tracing::warn!( + dataset, + error = %e, + "datasets-server /info unavailable (private dataset?); \ + assuming config 'default' and split 'train'" + ); + None + } + } + } { let dataset_info = info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| { BenchError::Config(format!( @@ -202,6 +205,11 @@ pub async fn download_hf_dataset( }; (resolved_config, resolved_split) + } else { + ( + subset.unwrap_or("default").to_string(), + split.unwrap_or("train").to_string(), + ) }; tracing::info!( @@ -211,96 +219,299 @@ pub async fn download_hf_dataset( "resolved Hugging Face dataset" ); - // Check cache - let dir = cache_dir(); - std::fs::create_dir_all(&dir)?; - let cache_path = dir.join(format!( - "hf-{}-{}-{}.json", - sanitize_name(dataset), - sanitize_name(&resolved_config), - sanitize_name(&resolved_split) - )); + let (revision, mut shard_paths) = + match convert_parquet_shards(&client, dataset, &resolved_config, &resolved_split).await { + Ok(paths) => (PARQUET_REVISION, paths), + Err(e) => { + tracing::info!( + dataset, + error = %e, + "no dataset-viewer parquet export; trying native parquet files on main" + ); + ( + "main", + native_data_files(dataset, &resolved_config, &resolved_split).await?, + ) + } + }; - if cache_path.exists() { - let path_str = cache_path.to_string_lossy().to_string(); - tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset"); - return Ok((path_str, resolved_config, resolved_split)); + // Sort for determinism, then shuffle shard order so samples are not always drawn + // from the start of the split (mirrors the datasets library's streaming shuffle). + shard_paths.sort(); + if !disable_shuffle { + shard_paths.shuffle(&mut StdRng::seed_from_u64(seed)); } - tracing::info!( - dataset, - config = resolved_config, - split = resolved_split, - "downloading Hugging Face dataset" - ); + // Read up to 2x the requested rows so validity filtering during sampling backfills + // from real rows instead of duplicating samples. + let target_rows = num_rows_needed.saturating_mul(2); - let encoded_config: String = - url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect(); - let encoded_split: String = - url::form_urlencoded::byte_serialize(resolved_split.as_bytes()).collect(); - - let mut all_rows: Vec = Vec::new(); - let mut offset = 0usize; - let page_size = 100usize; - let mut progress = RowDownloadReporter::new(); - - loop { - let url = format!( - "https://datasets-server.huggingface.co/rows\ - ?dataset={encoded_dataset}\ - &config={encoded_config}\ - &split={encoded_split}\ - &offset={offset}\ - &length={page_size}" - ); + let repo = + crate::hub::HubRepo::dataset_with_revision(dataset.to_string(), revision.to_string()) + .map_err(BenchError::Config)?; + let rows = read_shards(&repo, dataset, shard_paths, target_rows).await?; + Ok((rows, resolved_config, resolved_split)) +} - let data = get_with_retry(&client, &url, "HF dataset /rows").await?; +/// Hub revision holding the dataset-viewer's auto-converted parquet export. +const PARQUET_REVISION: &str = "refs/convert/parquet"; - let rows = data["rows"] - .as_array() - .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?; +/// Extract the shard's repo path within `refs/convert/parquet` from its resolve URL. +fn shard_repo_path(url: &str) -> Option { + url.split_once("refs%2Fconvert%2Fparquet/") + .or_else(|| url.split_once("refs/convert/parquet/")) + .map(|(_, path)| path.to_string()) +} - if rows.is_empty() { - break; +/// List the dataset-viewer's auto-converted parquet shards for a config/split. +async fn convert_parquet_shards( + client: &reqwest::Client, + dataset: &str, + config: &str, + split: &str, +) -> Result> { + let encoded_dataset: String = + url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); + let url = format!("https://datasets-server.huggingface.co/parquet?dataset={encoded_dataset}"); + let listing = get_with_retry(client, &url, "HF dataset /parquet").await?; + + // Partially converted datasets (>5GB) export under a "partial-" split prefix. + let partial_split = format!("partial-{split}"); + let shard_paths: Vec = listing["parquet_files"] + .as_array() + .map(|files| { + files + .iter() + .filter(|f| { + f["config"].as_str() == Some(config) + && matches!(f["split"].as_str(), Some(s) if s == split || s == partial_split) + }) + .filter_map(|f| shard_repo_path(f["url"].as_str()?)) + .collect() + }) + .unwrap_or_default(); + + if shard_paths.is_empty() { + return Err(BenchError::Config(format!( + "no parquet export for config '{config}' split '{split}'" + ))); + } + Ok(shard_paths) +} + +/// File extensions the native main-branch fallback can decode. +const NATIVE_DATA_EXTENSIONS: &[&str] = &[".parquet", ".json", ".jsonl", ".ndjson"]; + +/// Discover native data files on the main branch matching a config/split, following +/// the datasets library's default layout conventions (config directories, split-named +/// files, unlabeled data files = train). Used when the dataset viewer has no parquet +/// export, e.g. for private datasets. +async fn native_data_files(dataset: &str, config: &str, split: &str) -> Result> { + let repo = crate::hub::HubRepo::dataset(dataset.to_string()).map_err(BenchError::Config)?; + let files = repo.list_files().await.map_err(|e| { + BenchError::Config(format!( + "failed to list files of HF dataset '{dataset}': {e}" + )) + })?; + + let matched = select_native_files(files, config, split)?; + if matched.is_empty() { + return Err(BenchError::Config(format!( + "HF dataset '{dataset}' has no usable data for config '{config}' split \ + '{split}': no dataset-viewer parquet export and no native data files matching \ + the split on the main branch. Check the dataset page or the \ + --hf-subset/--hf-split values; private datasets must store parquet, json, or \ + jsonl files." + ))); + } + Ok(matched) +} + +/// Pure selection logic for [`native_data_files`]. +/// +/// Mirrors the datasets library's defaults: files may be scoped under a config +/// directory; split-named files are preferred; a repo whose data files carry no split +/// names at all maps everything to the "train" split. +/// +/// Errors when the config matches no directory but the selection still spans several +/// directories: the config never got resolved (e.g. a private dataset whose `/info` +/// lookup failed falls back to the literal "default"), and pooling those directories +/// would silently mix subsets. +fn select_native_files(files: Vec, config: &str, split: &str) -> Result> { + let data: Vec = files + .into_iter() + .filter(|f| NATIVE_DATA_EXTENSIONS.iter().any(|ext| f.ends_with(ext))) + .collect(); + + // Multi-config layout scopes files under a config directory (e.g. "main/test-*.parquet") + let config_prefix = format!("{config}/"); + let config_scoped = data.iter().any(|f| f.starts_with(&config_prefix)); + let scoped: Vec = if config_scoped { + data.into_iter().filter(|f| f.starts_with(&config_prefix)).collect() + } else { + data + }; + + let matched: Vec = + scoped.iter().filter(|f| path_matches_split(f, split)).cloned().collect(); + + let selected = if !matched.is_empty() { + matched + } else { + // No file carries the requested split name. If no file carries any split name, + // the datasets library maps them all to "train". + let split_keywords = ["train", "test", "validation", "valid", "dev"]; + let any_labeled = + scoped.iter().any(|f| split_keywords.iter().any(|s| path_matches_split(f, s))); + if split == "train" && !any_labeled { + scoped + } else { + Vec::new() } + }; - for row in rows { - if let Some(row_data) = row.get("row") { - all_rows.push(row_data.clone()); - } + if !config_scoped { + let dirs = top_level_dirs(&selected); + if dirs.len() > 1 { + return Err(BenchError::Config(format!( + "config '{config}' matches no directory in this dataset, and split '{split}' \ + data spans multiple config directories {dirs:?}. Pass --hf-subset to pick one." + ))); } + } + Ok(selected) +} - let fetched = rows.len(); - offset += fetched; +/// Distinct leading path segments of files that live in a directory. +fn top_level_dirs(files: &[String]) -> Vec<&str> { + let mut dirs: Vec<&str> = + files.iter().filter_map(|f| f.split_once('/').map(|(d, _)| d)).collect(); + dirs.sort_unstable(); + dirs.dedup(); + dirs +} - let total = data["num_rows_total"].as_u64().unwrap_or(0); - progress.update(offset, total); +/// Split match on a repo file path: a split-named file ("{split}-00000-of-00001.parquet", +/// "{split}.jsonl", "{split}_0.parquet") or a "{split}/" directory segment. +fn path_matches_split(path: &str, split: &str) -> bool { + let base = path.rsplit('/').next().unwrap_or(path); + base.strip_prefix(split) + .is_some_and(|rest| rest.starts_with('.') || rest.starts_with('-') || rest.starts_with('_')) + || path.starts_with(&format!("{split}/")) + || path.contains(&format!("/{split}/")) +} - // Stop if we have enough rows or reached end of dataset - if all_rows.len() >= num_rows_needed || fetched < page_size { +/// Fetch shards via hf-hub (cached under `$HF_HOME/hub`, shared with other Hub tooling) +/// and decode rows in order until `target_rows` are collected or all shards are read. +async fn read_shards( + repo: &crate::hub::HubRepo, + dataset: &str, + shard_paths: Vec, + target_rows: usize, +) -> Result> { + let total_shards = shard_paths.len(); + let mut rows: Vec = Vec::new(); + for (i, shard_path) in shard_paths.into_iter().enumerate() { + if rows.len() >= target_rows { break; } + tracing::info!( + dataset, + shard = i + 1, + total_shards, + path = %shard_path, + "fetching dataset shard" + ); + let local_path = repo.get(&shard_path).await.map_err(BenchError::Config)?; + let limit = target_rows - rows.len(); + let is_parquet = shard_path.ends_with(".parquet"); + let shard_rows = tokio::task::spawn_blocking(move || { + if is_parquet { + read_parquet_rows(&local_path, limit) + } else { + read_json_rows(&local_path, limit) + } + }) + .await + .map_err(|e| BenchError::Config(format!("shard read task failed: {e}")))??; + rows.extend(shard_rows); } - progress.finish(); - if all_rows.is_empty() { + if rows.is_empty() { return Err(BenchError::Config(format!( - "HF dataset '{dataset}' download returned no rows" + "HF dataset '{dataset}' data files contained no rows" ))); } - - // Save to cache - let json_str = serde_json::to_string(&all_rows)?; - std::fs::write(&cache_path, &json_str)?; - - let path_str = cache_path.to_string_lossy().to_string(); tracing::info!( dataset, - rows = all_rows.len(), - path = %path_str, - "saved Hugging Face dataset" + rows = rows.len(), + "loaded HF dataset from data files" ); - Ok((path_str, resolved_config, resolved_split)) + Ok(rows) +} + +/// Read up to `limit` rows from a parquet file as JSON objects. +fn read_parquet_rows(path: &Path, limit: usize) -> Result> { + let file = std::fs::File::open(path)?; + let reader = SerializedFileReader::new(file) + .map_err(|e| BenchError::Config(format!("failed to open parquet shard: {e}")))?; + let iter = reader + .get_row_iter(None) + .map_err(|e| BenchError::Config(format!("failed to read parquet shard: {e}")))?; + let mut rows = Vec::new(); + for row in iter { + if rows.len() >= limit { + break; + } + let row = + row.map_err(|e| BenchError::Config(format!("failed to decode parquet row: {e}")))?; + rows.push(row.to_json_value()); + } + Ok(rows) +} + +/// Read up to `limit` rows from a native JSON data file: a JSON array of objects, an +/// object wrapping a single array field, or JSON Lines. +fn read_json_rows(path: &Path, limit: usize) -> Result> { + let content = std::fs::read_to_string(path)?; + + let doc: Option = serde_json::from_str(&content).ok(); + let mut rows = match doc { + Some(serde_json::Value::Array(rows)) => rows, + // e.g. {"data": [...]} — a single array-of-objects field wrapping the rows + Some(serde_json::Value::Object(mut obj)) => { + let mut arrays: Vec = obj + .iter_mut() + .filter(|(_, v)| { + v.as_array().is_some_and(|a| !a.is_empty() && a.iter().all(|e| e.is_object())) + }) + .map(|(_, v)| v.take()) + .collect(); + match (arrays.len(), arrays.pop()) { + (1, Some(serde_json::Value::Array(rows))) => rows, + _ => { + return Err(BenchError::Config( + "unsupported JSON data file layout: expected an array of rows, \ + an object with a single array-of-objects field, or JSON Lines" + .into(), + )); + } + } + } + // Not a single JSON document — try JSON Lines + _ => content + .lines() + .filter(|l| !l.trim().is_empty()) + .take(limit) + .map(|l| { + serde_json::from_str(l) + .map_err(|e| BenchError::Config(format!("invalid JSON line in data file: {e}"))) + }) + .collect::>>()?, + }; + + rows.truncate(limit); + Ok(rows) } /// Check if a JSON value looks like a chat message (has role+content or from+value). @@ -472,10 +683,10 @@ fn extract_chat_completion(messages: &[serde_json::Value]) -> Option { None } -/// Load an HF dataset from the cached JSON file and convert to SampleRequests. +/// Convert downloaded HF dataset rows to SampleRequests. pub fn load_hf_dataset( tokenizer: &TokenizerKind, - dataset_path: &str, + entries: &[serde_json::Value], num_requests: usize, hf_output_len: Option, seed: u64, @@ -484,19 +695,8 @@ pub fn load_hf_dataset( no_oversample: bool, disable_shuffle: bool, ) -> Result> { - let content = std::fs::read_to_string(dataset_path).map_err(|e| { - BenchError::Config(format!( - "Failed to read HF dataset file '{dataset_path}': {e}" - )) - })?; - - let entries: Vec = serde_json::from_str(&content) - .map_err(|e| BenchError::Config(format!("Invalid JSON in HF dataset file: {e}")))?; - if entries.is_empty() { - return Err(BenchError::Config( - "HF dataset file contains no rows".into(), - )); + return Err(BenchError::Config("HF dataset contains no rows".into())); } // Detect column format from first row @@ -638,7 +838,6 @@ pub fn load_hf_dataset( } else { if !warned_no_output { tracing::warn!( - path = dataset_path, default_output_tokens = 128, "no dataset output column or --hf-output-len; using default output length" ); @@ -649,7 +848,6 @@ pub fn load_hf_dataset( } else { if !warned_no_output { tracing::warn!( - path = dataset_path, default_output_tokens = 128, "no dataset output column or --hf-output-len; using default output length" ); @@ -965,16 +1163,6 @@ mod tests { assert!(!is_chat_array(&serde_json::json!("not an array"))); } - #[test] - fn test_sanitize_name() { - assert_eq!( - sanitize_name("allenai/WildChat-4.8M"), - "allenai_WildChat-4.8M" - ); - assert_eq!(sanitize_name("simple-name"), "simple-name"); - assert_eq!(sanitize_name("org/sub/deep"), "org_sub_deep"); - } - // --- extract_chat_prompt edge cases --- #[test] @@ -1043,13 +1231,6 @@ mod tests { ) } - /// Write JSON data to a unique temp file and return the path string. - fn write_temp_json(name: &str, data: &serde_json::Value) -> String { - let path = std::env::temp_dir().join(format!("vllm-bench-test-{name}.json")); - std::fs::write(&path, serde_json::to_string(data).unwrap()).unwrap(); - path.to_string_lossy().to_string() - } - #[test] fn test_load_hf_dataset_chat_format() { let tok = builtin_tokenizer(); @@ -1067,10 +1248,19 @@ mod tests { ] } ]); - let path = write_temp_json("chat-format", &data); - let result = - load_hf_dataset(&tok, &path, 2, Some(50), 42, "test-", None, false, false).unwrap(); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 2, + Some(50), + 42, + "test-", + None, + false, + false, + ) + .unwrap(); assert_eq!(result.len(), 2); assert!( @@ -1087,9 +1277,19 @@ mod tests { {"prompt": "Explain the difference between machine learning and deep learning in simple terms.", "completion": "Machine learning is a subset of AI."}, {"prompt": "How does photosynthesis work in plants and what role does chlorophyll play?", "completion": "Photosynthesis converts light to energy."} ]); - let path = write_temp_json("plain-text-format", &data); - let result = load_hf_dataset(&tok, &path, 3, None, 0, "req-", None, false, false).unwrap(); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 3, + None, + 0, + "req-", + None, + false, + false, + ) + .unwrap(); assert_eq!(result.len(), 3); // Without hf_output_len override, output len is derived from the completion tokens @@ -1111,10 +1311,19 @@ mod tests { "answers": ["Albert Einstein"] } ]); - let path = write_temp_json("combined-format", &data); - let result = - load_hf_dataset(&tok, &path, 2, Some(64), 1, "comb-", None, false, false).unwrap(); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 2, + Some(64), + 1, + "comb-", + None, + false, + false, + ) + .unwrap(); assert_eq!(result.len(), 2); assert!(result.iter().all(|r| r.expected_output_len == 64)); @@ -1126,9 +1335,18 @@ mod tests { fn test_load_hf_dataset_empty_returns_error() { let tok = builtin_tokenizer(); let data = serde_json::json!([]); - let path = write_temp_json("empty-dataset", &data); - let result = load_hf_dataset(&tok, &path, 5, None, 0, "test-", None, false, false); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 5, + None, + 0, + "test-", + None, + false, + false, + ); assert!(result.is_err(), "empty dataset should return an error"); let err = result.unwrap_err().to_string(); @@ -1147,9 +1365,18 @@ mod tests { {"prompt": "Yes", "completion": "No"}, {"prompt": "Ok", "completion": "Fine"} ]); - let path = write_temp_json("all-filtered", &data); - let result = load_hf_dataset(&tok, &path, 3, None, 0, "test-", None, false, false); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 3, + None, + 0, + "test-", + None, + false, + false, + ); assert!( result.is_err(), @@ -1165,12 +1392,11 @@ mod tests { {"prompt": "What are the key differences between supervised and unsupervised machine learning?", "completion": "Supervised learning uses labeled data while unsupervised does not."}, {"prompt": "Explain how neural networks are inspired by the human brain structure and function.", "completion": "Neural networks mimic brain neurons with layers of nodes."} ]); - let path = write_temp_json("output-len-override", &data); let fixed_len = 77usize; let result = load_hf_dataset( &tok, - &path, + data.as_array().unwrap(), 3, Some(fixed_len), 42, @@ -1196,11 +1422,10 @@ mod tests { {"prompt": "Explain how photosynthesis converts sunlight to energy in plant cells.", "completion": "Plants use chlorophyll to absorb sunlight."}, {"prompt": "What is the difference between a virus and a bacterium in terms of biology?", "completion": "Viruses need host cells while bacteria are self-sufficient."} ]); - let path = write_temp_json("no-oversample", &data); let result = load_hf_dataset( &tok, - &path, + data.as_array().unwrap(), 10, Some(32), 0, @@ -1227,11 +1452,10 @@ mod tests { {"prompt": "Beta prompt about the second topic continuing from our ordered sequence here.", "completion": "Beta answer"}, {"prompt": "Gamma prompt about the third item in our clearly ordered alphabetical sequence.", "completion": "Gamma answer"} ]); - let path = write_temp_json("disable-shuffle", &data); let result = load_hf_dataset( &tok, - &path, + data.as_array().unwrap(), 3, Some(20), 99, @@ -1264,13 +1488,188 @@ mod tests { {"prompt": "What is the boiling point of water at sea level under standard atmospheric pressure?", "completion": "Water boils at 100 degrees Celsius."}, {"prompt": "Describe the structure of DNA and how genetic information is encoded within it.", "completion": "DNA is a double helix with base pairs."} ]); - let path = write_temp_json("request-id-prefix", &data); - let result = - load_hf_dataset(&tok, &path, 2, Some(10), 0, "myprefix-", None, false, false).unwrap(); + let result = load_hf_dataset( + &tok, + data.as_array().unwrap(), + 2, + Some(10), + 0, + "myprefix-", + None, + false, + false, + ) + .unwrap(); assert_eq!(result.len(), 2); assert_eq!(result[0].request_id.as_deref(), Some("myprefix-0")); assert_eq!(result[1].request_id.as_deref(), Some("myprefix-1")); } + + // Network test: the full GSM8k test split must arrive via parquet shards regardless + // of how many rows earlier runs requested (run with --ignored). + #[tokio::test] + #[ignore] + async fn test_download_gsm8k_full_split_via_parquet() { + let (rows, config, split) = + download_hf_dataset("openai/gsm8k", None, Some("test"), 1319, 0, false) + .await + .unwrap(); + assert_eq!(config, "main"); + assert_eq!(split, "test"); + assert_eq!(rows.len(), 1319); + assert!(rows[0].get("question").is_some_and(|q| q.is_string())); + assert!(rows[0].get("answer").is_some_and(|a| a.is_string())); + } + + // Network test: native parquet discovery on the main branch (private-dataset path) + // must find GSM8k's config-dir layout (run with --ignored). + #[tokio::test] + #[ignore] + async fn test_native_data_files_main_branch() { + let shards = native_data_files("openai/gsm8k", "main", "test").await.unwrap(); + assert_eq!(shards, vec!["main/test-00000-of-00001.parquet".to_string()]); + + // Single unlabeled JSON file maps to the train split (private-dataset layout) + let shards = native_data_files("Inferact/codex_swebenchpro_traces", "default", "train") + .await + .unwrap(); + assert_eq!(shards, vec!["codex_swebenchpro.json".to_string()]); + } + + // --- native data file selection --- + + #[test] + fn test_path_matches_split() { + assert!(path_matches_split( + "data/train-00000-of-00042.parquet", + "train" + )); + assert!(path_matches_split( + "main/test-00000-of-00001.parquet", + "test" + )); + assert!(path_matches_split("test.parquet", "test")); + assert!(path_matches_split("test.jsonl", "test")); + assert!(path_matches_split("test_0.parquet", "test")); + assert!(path_matches_split("data/test/0000.parquet", "test")); + assert!(path_matches_split("test/0000.parquet", "test")); + + assert!(!path_matches_split( + "data/train-00000-of-00042.parquet", + "test" + )); + assert!(!path_matches_split("testing-00000.parquet", "test")); + assert!(!path_matches_split("contest/0000.parquet", "test")); + assert!(!path_matches_split("data/latest/0000.parquet", "test")); + } + + #[test] + fn test_select_native_files_split_labeled() { + let files = vec![ + ".gitattributes".to_string(), + "README.md".to_string(), + "data/train-00000-of-00002.parquet".to_string(), + "data/train-00001-of-00002.parquet".to_string(), + "data/test-00000-of-00001.parquet".to_string(), + ]; + let selected = select_native_files(files, "default", "test").unwrap(); + assert_eq!( + selected, + vec!["data/test-00000-of-00001.parquet".to_string()] + ); + } + + #[test] + fn test_select_native_files_unlabeled_maps_to_train() { + let files = vec![ + "README.md".to_string(), + "codex_swebenchpro.json".to_string(), + ]; + let selected = select_native_files(files.clone(), "default", "train").unwrap(); + assert_eq!(selected, vec!["codex_swebenchpro.json".to_string()]); + // ...but only for the train split + assert!(select_native_files(files, "default", "test").unwrap().is_empty()); + } + + #[test] + fn test_select_native_files_labeled_files_do_not_leak_to_train() { + // A repo with only test-labeled files has no train split + let files = vec!["data/test-00000-of-00001.parquet".to_string()]; + assert!(select_native_files(files, "default", "train").unwrap().is_empty()); + } + + #[test] + fn test_select_native_files_config_dir_scoping() { + let files = vec![ + "main/test-00000-of-00001.parquet".to_string(), + "socratic/test-00000-of-00001.parquet".to_string(), + ]; + let selected = select_native_files(files, "socratic", "test").unwrap(); + assert_eq!( + selected, + vec!["socratic/test-00000-of-00001.parquet".to_string()] + ); + } + + #[test] + fn test_select_native_files_unresolved_config_does_not_pool_configs() { + // A private dataset whose /info lookup failed falls back to config "default"; + // pooling both config directories would silently mix subsets. + let files = vec![ + "main/test-00000-of-00001.parquet".to_string(), + "socratic/test-00000-of-00001.parquet".to_string(), + ]; + let err = select_native_files(files, "default", "test").unwrap_err().to_string(); + assert!(err.contains("--hf-subset"), "{err}"); + } + + #[test] + fn test_select_native_files_split_dirs_are_not_ambiguous() { + // Split directories are not config directories + let files = vec![ + "train/0000.parquet".to_string(), + "test/0000.parquet".to_string(), + ]; + let selected = select_native_files(files, "default", "test").unwrap(); + assert_eq!(selected, vec!["test/0000.parquet".to_string()]); + } + + // --- native JSON data file reading --- + + fn write_temp_file(name: &str, content: &str) -> std::path::PathBuf { + let path = std::env::temp_dir().join(format!("vllm-bench-test-json-{name}")); + std::fs::write(&path, content).unwrap(); + path + } + + #[test] + fn test_read_json_rows_array() { + let path = write_temp_file("array.json", r#"[{"a": 1}, {"a": 2}, {"a": 3}]"#); + let rows = read_json_rows(&path, 2).unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0]["a"], 1); + } + + #[test] + fn test_read_json_rows_wrapped_object() { + let path = write_temp_file("wrapped.json", r#"{"data": [{"a": 1}, {"a": 2}]}"#); + let rows = read_json_rows(&path, 10).unwrap(); + assert_eq!(rows.len(), 2); + } + + #[test] + fn test_read_json_rows_jsonl() { + let path = write_temp_file("rows.jsonl", "{\"a\": 1}\n{\"a\": 2}\n\n{\"a\": 3}\n"); + let rows = read_json_rows(&path, 10).unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[2]["a"], 3); + } + + #[test] + fn test_read_json_rows_unsupported_layout() { + let path = write_temp_file("scalar.json", r#"{"a": 1, "b": 2}"#); + assert!(read_json_rows(&path, 10).is_err()); + } } diff --git a/rust/src/bench/src/hub.rs b/rust/src/bench/src/hub.rs index 0b446450fb2d..ff61c907f2cc 100644 --- a/rust/src/bench/src/hub.rs +++ b/rust/src/bench/src/hub.rs @@ -3,8 +3,8 @@ use std::path::PathBuf; -use hf_hub::Repo; use hf_hub::api::tokio::{ApiBuilder, ApiRepo}; +use hf_hub::{Repo, RepoType}; /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache. pub struct HubRepo { @@ -20,6 +20,10 @@ impl HubRepo { Self::new(Repo::dataset(repo_id)) } + pub fn dataset_with_revision(repo_id: String, revision: String) -> Result { + Self::new(Repo::with_revision(repo_id, RepoType::Dataset, revision)) + } + fn new(repo: Repo) -> Result { let mut builder = ApiBuilder::from_env(); if let Ok(token) = std::env::var("HF_TOKEN") { @@ -36,4 +40,10 @@ impl HubRepo { pub async fn get(&self, filename: &str) -> Result { self.repo.get(filename).await.map_err(|e| format!("{e}")) } + + /// List all file paths in the repo at its revision. + pub async fn list_files(&self) -> Result, String> { + let info = self.repo.info().await.map_err(|e| format!("{e}"))?; + Ok(info.siblings.into_iter().map(|s| s.rfilename).collect()) + } }