@tduncan / Nvidia-Optimization-Stack-VLLM / commits / df344bf36f

[Rust][Benchmark] Load HF datasets from parquet shards via hf-hub, fixing truncated-cache sampling (#51570)

Roy Wang committed Aug 21, 2026 at 13:07 UTC df344bf36f107b5f801b2c058c8ea8e922101f6a
8 files changed +721 -172
rust/Cargo.lock
+133
@@ -15,6 +15,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
15 checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
16 dependencies = [
17 "cfg-if",
18 + "const-random",
19 "getrandom 0.3.4",
20 "once_cell",
21 "serde",
@@ -31,6 +32,21 @@ dependencies = [
32 "memchr",
33 ]
34
35 +[[package]]
36 +name = "alloc-no-stdlib"
37 +version = "2.0.4"
38 +source = "registry+https://github.com/rust-lang/crates.io-index"
39 +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3"
40 +
41 +[[package]]
42 +name = "alloc-stdlib"
43 +version = "0.2.4"
44 +source = "registry+https://github.com/rust-lang/crates.io-index"
45 +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195"
46 +dependencies = [
47 + "alloc-no-stdlib",
48 +]
49 +
50 [[package]]
51 name = "android_system_properties"
52 version = "0.1.5"
@@ -379,6 +395,12 @@ version = "0.22.1"
395 source = "registry+https://github.com/rust-lang/crates.io-index"
396 checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
397
398 +[[package]]
399 +name = "base64"
400 +version = "0.23.1"
401 +source = "registry+https://github.com/rust-lang/crates.io-index"
402 +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5"
403 +
404 [[package]]
405 name = "base64ct"
406 version = "1.8.3"
@@ -450,6 +472,27 @@ dependencies = [
472 "generic-array",
473 ]
474
475 +[[package]]
476 +name = "brotli"
477 +version = "8.0.4"
478 +source = "registry+https://github.com/rust-lang/crates.io-index"
479 +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3"
480 +dependencies = [
481 + "alloc-no-stdlib",
482 + "alloc-stdlib",
483 + "brotli-decompressor",
484 +]
485 +
486 +[[package]]
487 +name = "brotli-decompressor"
488 +version = "5.0.3"
489 +source = "registry+https://github.com/rust-lang/crates.io-index"
490 +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583"
491 +dependencies = [
492 + "alloc-no-stdlib",
493 + "alloc-stdlib",
494 +]
495 +
496 [[package]]
497 name = "bstr"
498 version = "1.12.1"
@@ -671,6 +714,26 @@ dependencies = [
714 "windows-sys 0.61.2",
715 ]
716
717 +[[package]]
718 +name = "const-random"
719 +version = "0.1.18"
720 +source = "registry+https://github.com/rust-lang/crates.io-index"
721 +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359"
722 +dependencies = [
723 + "const-random-macro",
724 +]
725 +
726 +[[package]]
727 +name = "const-random-macro"
728 +version = "0.1.16"
729 +source = "registry+https://github.com/rust-lang/crates.io-index"
730 +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
731 +dependencies = [
732 + "getrandom 0.2.17",
733 + "once_cell",
734 + "tiny-keccak",
735 +]
736 +
737 [[package]]
738 name = "constant_time_eq"
739 version = "0.4.2"
@@ -1584,6 +1647,7 @@ dependencies = [
1647 "bytemuck",
1648 "cfg-if",
1649 "crunchy",
1650 + "num-traits",
1651 "zerocopy",
1652 ]
1653
@@ -1617,6 +1681,12 @@ version = "0.16.1"
1681 source = "registry+https://github.com/rust-lang/crates.io-index"
1682 checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
1683
1684 +[[package]]
1685 +name = "hashbrown"
1686 +version = "0.17.1"
1687 +source = "registry+https://github.com/rust-lang/crates.io-index"
1688 +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
1689 +
1690 [[package]]
1691 name = "heck"
1692 version = "0.5.0"
@@ -2270,6 +2340,15 @@ version = "0.4.29"
2340 source = "registry+https://github.com/rust-lang/crates.io-index"
2341 checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
2342
2343 +[[package]]
2344 +name = "lz4_flex"
2345 +version = "0.14.0"
2346 +source = "registry+https://github.com/rust-lang/crates.io-index"
2347 +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226"
2348 +dependencies = [
2349 + "twox-hash",
2350 +]
2351 +
2352 [[package]]
2353 name = "logos"
2354 version = "0.15.1"
@@ -2648,6 +2727,16 @@ dependencies = [
2727 "windows-sys 0.61.2",
2728 ]
2729
2730 +[[package]]
2731 +name = "num-bigint"
2732 +version = "0.5.1"
2733 +source = "registry+https://github.com/rust-lang/crates.io-index"
2734 +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
2735 +dependencies = [
2736 + "num-integer",
2737 + "num-traits",
2738 +]
2739 +
2740 [[package]]
2741 name = "num-complex"
2742 version = "0.4.6"
@@ -2866,6 +2955,31 @@ dependencies = [
2955 "windows-link",
2956 ]
2957
2958 +[[package]]
2959 +name = "parquet"
2960 +version = "59.2.0"
2961 +source = "registry+https://github.com/rust-lang/crates.io-index"
2962 +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c"
2963 +dependencies = [
2964 + "ahash",
2965 + "base64 0.23.1",
2966 + "brotli",
2967 + "bytes",
2968 + "chrono",
2969 + "flate2",
2970 + "half",
2971 + "hashbrown 0.17.1",
2972 + "lz4_flex",
2973 + "num-bigint",
2974 + "num-integer",
2975 + "num-traits",
2976 + "seq-macro",
2977 + "serde_json",
2978 + "snap",
2979 + "twox-hash",
2980 + "zstd",
2981 +]
2982 +
2983 [[package]]
2984 name = "paste"
2985 version = "1.0.15"
@@ -4044,6 +4158,12 @@ version = "1.0.27"
4158 source = "registry+https://github.com/rust-lang/crates.io-index"
4159 checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2"
4160
4161 +[[package]]
4162 +name = "seq-macro"
4163 +version = "0.3.6"
4164 +source = "registry+https://github.com/rust-lang/crates.io-index"
4165 +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc"
4166 +
4167 [[package]]
4168 name = "serde"
4169 version = "1.0.228"
@@ -4316,6 +4436,12 @@ dependencies = [
4436 "version_check",
4437 ]
4438
4439 +[[package]]
4440 +name = "snap"
4441 +version = "1.1.2"
4442 +source = "registry+https://github.com/rust-lang/crates.io-index"
4443 +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886"
4444 +
4445 [[package]]
4446 name = "socket2"
4447 version = "0.6.3"
@@ -5341,6 +5467,12 @@ dependencies = [
5467 "utf-8",
5468 ]
5469
5470 +[[package]]
5471 +name = "twox-hash"
5472 +version = "2.1.3"
5473 +source = "registry+https://github.com/rust-lang/crates.io-index"
5474 +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9"
5475 +
5476 [[package]]
5477 name = "typenum"
5478 version = "1.19.0"
@@ -5636,6 +5768,7 @@ dependencies = [
5768 "image",
5769 "indicatif",
5770 "mimalloc",
5771 + "parquet",
5772 "rand 0.9.2",
5773 "rand_distr",
5774 "rayon",
rust/Cargo.toml
+1
@@ -70,6 +70,7 @@ openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmon
70 openai-protocol = "1.6.0"
71 openssl = "0.10"
72 parking_lot = "0.12.5"
73 +parquet = { version = "59.2.0", default-features = false, features = ["json", "snap", "zstd", "flate2", "flate2-rust_backend", "lz4", "brotli"] }
74 paste = "1.0.15"
75 prometheus-client = "0.24.0"
76 prometheus-client-derive-encode = "0.5.0"
rust/src/bench/AGENTS.md
+3 -3
@@ -4,7 +4,7 @@
4
5 Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime.
6
7 -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.
7 +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.
8
9 ## Build & Test
10
@@ -34,7 +34,7 @@ cargo test -p vllm-bench -- --ignored
34 - `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`)
35 - `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
36 - `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
37 -- `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
37 +- `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
38 - `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up
39 - `src/ready_checker.rs` — Endpoint readiness with retry
40 - `src/backends/` — Backend implementations (enum dispatch, not trait objects)
@@ -48,7 +48,7 @@ cargo test -p vllm-bench -- --ignored
48 - `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching
49 - `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`)
50 - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache)
51 -- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection)
51 +- `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)
52 - `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)
53 - `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`)
54 - `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`)
rust/src/bench/Cargo.toml
+1
@@ -17,6 +17,7 @@ hf-hub.workspace = true
17 image.workspace = true
18 indicatif.workspace = true
19 mimalloc.workspace = true
20 +parquet.workspace = true
21 rand.workspace = true
22 rand_distr.workspace = true
23 rayon.workspace = true
rust/src/bench/README.md
+9 -5
@@ -148,7 +148,8 @@ vllm-bench \
148 --sonnet-input-len 550 --sonnet-output-len 150 --sonnet-prefix-len 200 \
149 --num-prompts 500
150
151 -# Any public HuggingFace dataset (auto-downloads, auto-detects columns)
151 +# Any HuggingFace dataset (downloads parquet shards into the standard HF hub
152 +# cache — shared with `hf download` and Python tooling — auto-detects columns)
153 vllm-bench \
154 --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
155 --dataset-name hf --dataset-path allenai/WildChat-4.8M \
@@ -160,7 +161,8 @@ vllm-bench \
161 --dataset-name hf --dataset-path THUDM/LongBench \
162 --hf-subset narrativeqa --hf-split test --hf-output-len 512 --num-prompts 200
163
163 -# Gated HuggingFace dataset (requires HF_TOKEN)
164 +# Gated HuggingFace dataset (auth via HF_TOKEN or the `hf auth login` token file).
165 +# Private datasets work too when the repo stores native parquet/json/jsonl files.
166 HF_TOKEN=hf_xxx vllm-bench \
167 --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
168 --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
433 | `sharegpt` | Real conversations from ShareGPT (auto-downloads from HuggingFace, or use `--dataset-path`) |
434 | `sonnet` | Built-in Shakespeare sonnets; controllable token length + shared prefix, no dataset file needed |
435 | `speed-bench` | NVIDIA SPEED-Bench for speculative decoding evaluation (auto-downloads, 11 categories) |
434 -| `hf` | Any HuggingFace dataset (auto-downloads via datasets-server API, auto-detects chat/text columns) |
436 +| `hf` | Any HuggingFace dataset with parquet data (downloads parquet shards via hf-hub into the standard HF hub cache, auto-detects chat/text columns) |
437
438 ## Metrics
439
@@ -507,7 +509,7 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below.
509 | `--random-input-len` | `1024` | Input token length |
510 | `--random-output-len` | `128` | Output token length |
511 | `--random-prefix-len` | `0` | Shared prefix length |
510 -| `--random-range-ratio` | `1.0` | Length jitter, range `(0, 1]`. Lengths sampled from `[ratio × target, target]`; `1.0` = fixed length |
512 +| `--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 |
513 | `--prompt-token-ids` | `false` | Send prompts as token-ID arrays (skips server-side tokenization, exact counts). Random dataset only |
514 | **Random multimodal** | | |
515 | `--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:
797 - Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up)
798 - Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput)
799 - Sampling parameters merged into the request body via `extra_body` (same precedence rules)
800 +- 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)
801 - CLI flags: a superset of Python's, enforced by `tests/cli_parity.rs` against a
802 snapshot of the Python parser (`tests/python_serve_flags.txt`, regenerated by
803 `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:
809 | Variable | Description |
810 | ---------- | ------------- |
811 | `OPENAI_API_KEY` | API key for authenticated endpoints (cached, not read per-request) |
809 -| `HF_TOKEN` | HuggingFace token for gated model tokenizers and gated datasets |
812 +| `HF_TOKEN` | HuggingFace token for gated/private tokenizers and datasets (falls back to the `hf auth login` token file) |
813 +| `HF_HOME` | Overrides the HuggingFace cache location (default `~/.cache/huggingface`) used for tokenizers and dataset parquet shards |
814 | `TOKIO_WORKER_THREADS` | Override tokio worker thread count (default: physical cores) |
815
816 ## License
rust/src/bench/src/benchmark.rs
+10 -9
@@ -534,17 +534,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
534 let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
535 BenchError::Config("--dataset-path is required for --dataset-name hf".into())
536 })?;
537 - let (downloaded_path, _config, _split) =
538 - crate::datasets::hf_dataset::download_hf_dataset(
539 - dataset_id,
540 - config.hf_subset.as_deref(),
541 - config.hf_split.as_deref(),
542 - config.num_prompts,
543 - )
544 - .await?;
537 + let (entries, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset(
538 + dataset_id,
539 + config.hf_subset.as_deref(),
540 + config.hf_split.as_deref(),
541 + config.num_prompts,
542 + config.seed,
543 + config.disable_shuffle,
544 + )
545 + .await?;
546 crate::datasets::hf_dataset::load_hf_dataset(
547 tok,
547 - &downloaded_path,
548 + &entries,
549 config.num_prompts,
550 config.hf_output_len,
551 config.seed,
rust/src/bench/src/datasets/hf_dataset.rs
+553 -154
@@ -1,38 +1,18 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +use std::path::Path;
5 use std::sync::Arc;
6
7 +use parquet::file::reader::{FileReader, SerializedFileReader};
8 use rand::rngs::StdRng;
9 use rand::seq::SliceRandom;
10 use rand::{Rng, SeedableRng};
11
12 use super::SampleRequest;
11 -use super::progress::RowDownloadReporter;
13 use crate::error::{BenchError, Result};
14 use crate::tokenizer::TokenizerKind;
15
15 -/// Cache directory for downloaded HF datasets.
16 -fn cache_dir() -> std::path::PathBuf {
17 - dirs::cache_dir()
18 - .unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
19 - .join("vllm-bench")
20 - .join("datasets")
21 -}
22 -
23 -/// Sanitize a dataset name for use in filenames (replace `/` and other unsafe chars).
24 -fn sanitize_name(name: &str) -> String {
25 - name.chars()
26 - .map(|c| {
27 - if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' {
28 - c
29 - } else {
30 - '_'
31 - }
32 - })
33 - .collect()
34 -}
35 -
16 /// Detected column format for extracting prompts from HF dataset rows.
17 enum ColumnFormat {
18 /// Column contains chat messages (array of {role, content} or {from, value}).
@@ -103,28 +83,39 @@ async fn get_with_retry(
83
84 /// Download an arbitrary HuggingFace dataset via the datasets-server REST API.
85 ///
106 -/// Returns `(cache_path, resolved_config, resolved_split)`.
86 +/// Returns `(rows, resolved_config, resolved_split)`.
87 ///
88 /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
109 -/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
110 -/// or the dataset is exhausted.
89 +/// Rows come from data files downloaded via hf-hub into the standard HF hub cache: the
90 +/// dataset-viewer's `refs/convert/parquet` export when available, otherwise native
91 +/// parquet/JSON/JSONL files on the main branch (e.g. private datasets the viewer does
92 +/// not process). Shard order is shuffled with `seed` (unless `disable_shuffle`) so
93 +/// samples are not always drawn from the start of the split, and up to twice
94 +/// `num_rows_needed` rows are read so downstream validity filtering can backfill from
95 +/// real rows.
96 pub async fn download_hf_dataset(
97 dataset: &str,
98 subset: Option<&str>,
99 split: Option<&str>,
100 num_rows_needed: usize,
116 -) -> Result<(String, String, String)> {
101 + seed: u64,
102 + disable_shuffle: bool,
103 +) -> Result<(Vec<serde_json::Value>, String, String)> {
104 let encoded_dataset: String =
105 url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
106
107 let mut client_builder =
108 reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
109
123 - // Add HF_TOKEN auth header if available
124 - if let Ok(token) = std::env::var("HF_TOKEN") {
110 + // Match hf-hub's auth for shard downloads: HF_TOKEN env var, then the hub token file
111 + let token = std::env::var("HF_TOKEN")
112 + .ok()
113 + .filter(|t| !t.is_empty())
114 + .or_else(|| hf_hub::Cache::from_env().token());
115 + if let Some(token) = token {
116 let mut headers = reqwest::header::HeaderMap::new();
117 let header_value = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}"))
127 - .map_err(|e| BenchError::Config(format!("Invalid HF_TOKEN: {e}")))?;
118 + .map_err(|e| BenchError::Config(format!("Invalid HF token: {e}")))?;
119 headers.insert(reqwest::header::AUTHORIZATION, header_value);
120 client_builder = client_builder.default_headers(headers);
121 }
@@ -137,12 +128,24 @@ pub async fn download_hf_dataset(
128 let (resolved_config, resolved_split) = if let (Some(cfg), Some(spl)) = (subset, split) {
129 // Both provided — skip /info call
130 (cfg.to_string(), spl.to_string())
140 - } else {
141 - // Call /info to discover available configs and splits
131 + } else if let Some(info) = {
132 + // Call /info to discover available configs and splits. Private datasets are not
133 + // processed by the dataset viewer; fall back to datasets-library defaults.
134 let info_url =
135 format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
144 - let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
145 -
136 + match get_with_retry(&client, &info_url, "HF dataset /info").await {
137 + Ok(info) => Some(info),
138 + Err(e) => {
139 + tracing::warn!(
140 + dataset,
141 + error = %e,
142 + "datasets-server /info unavailable (private dataset?); \
143 + assuming config 'default' and split 'train'"
144 + );
145 + None
146 + }
147 + }
148 + } {
149 let dataset_info =
150 info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
151 BenchError::Config(format!(
@@ -202,6 +205,11 @@ pub async fn download_hf_dataset(
205 };
206
207 (resolved_config, resolved_split)
208 + } else {
209 + (
210 + subset.unwrap_or("default").to_string(),
211 + split.unwrap_or("train").to_string(),
212 + )
213 };
214
215 tracing::info!(
@@ -211,96 +219,299 @@ pub async fn download_hf_dataset(
219 "resolved Hugging Face dataset"
220 );
221
214 - // Check cache
215 - let dir = cache_dir();
216 - std::fs::create_dir_all(&dir)?;
217 - let cache_path = dir.join(format!(
218 - "hf-{}-{}-{}.json",
219 - sanitize_name(dataset),
220 - sanitize_name(&resolved_config),
221 - sanitize_name(&resolved_split)
222 - ));
222 + let (revision, mut shard_paths) =
223 + match convert_parquet_shards(&client, dataset, &resolved_config, &resolved_split).await {
224 + Ok(paths) => (PARQUET_REVISION, paths),
225 + Err(e) => {
226 + tracing::info!(
227 + dataset,
228 + error = %e,
229 + "no dataset-viewer parquet export; trying native parquet files on main"
230 + );
231 + (
232 + "main",
233 + native_data_files(dataset, &resolved_config, &resolved_split).await?,
234 + )
235 + }
236 + };
237
224 - if cache_path.exists() {
225 - let path_str = cache_path.to_string_lossy().to_string();
226 - tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
227 - return Ok((path_str, resolved_config, resolved_split));
238 + // Sort for determinism, then shuffle shard order so samples are not always drawn
239 + // from the start of the split (mirrors the datasets library's streaming shuffle).
240 + shard_paths.sort();
241 + if !disable_shuffle {
242 + shard_paths.shuffle(&mut StdRng::seed_from_u64(seed));
243 }
244
230 - tracing::info!(
231 - dataset,
232 - config = resolved_config,
233 - split = resolved_split,
234 - "downloading Hugging Face dataset"
235 - );
245 + // Read up to 2x the requested rows so validity filtering during sampling backfills
246 + // from real rows instead of duplicating samples.
247 + let target_rows = num_rows_needed.saturating_mul(2);
248
237 - let encoded_config: String =
238 - url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
239 - let encoded_split: String =
240 - url::form_urlencoded::byte_serialize(resolved_split.as_bytes()).collect();
241 -
242 - let mut all_rows: Vec<serde_json::Value> = Vec::new();
243 - let mut offset = 0usize;
244 - let page_size = 100usize;
245 - let mut progress = RowDownloadReporter::new();
246 -
247 - loop {
248 - let url = format!(
249 - "https://datasets-server.huggingface.co/rows\
250 - ?dataset={encoded_dataset}\
251 - &config={encoded_config}\
252 - &split={encoded_split}\
253 - &offset={offset}\
254 - &length={page_size}"
255 - );
249 + let repo =
250 + crate::hub::HubRepo::dataset_with_revision(dataset.to_string(), revision.to_string())
251 + .map_err(BenchError::Config)?;
252 + let rows = read_shards(&repo, dataset, shard_paths, target_rows).await?;
253 + Ok((rows, resolved_config, resolved_split))
254 +}
255
257 - let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
256 +/// Hub revision holding the dataset-viewer's auto-converted parquet export.
257 +const PARQUET_REVISION: &str = "refs/convert/parquet";
258
259 - let rows = data["rows"]
260 - .as_array()
261 - .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?;
259 +/// Extract the shard's repo path within `refs/convert/parquet` from its resolve URL.
260 +fn shard_repo_path(url: &str) -> Option<String> {
261 + url.split_once("refs%2Fconvert%2Fparquet/")
262 + .or_else(|| url.split_once("refs/convert/parquet/"))
263 + .map(|(_, path)| path.to_string())
264 +}
265
263 - if rows.is_empty() {
264 - break;
266 +/// List the dataset-viewer's auto-converted parquet shards for a config/split.
267 +async fn convert_parquet_shards(
268 + client: &reqwest::Client,
269 + dataset: &str,
270 + config: &str,
271 + split: &str,
272 +) -> Result<Vec<String>> {
273 + let encoded_dataset: String =
274 + url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
275 + let url = format!("https://datasets-server.huggingface.co/parquet?dataset={encoded_dataset}");
276 + let listing = get_with_retry(client, &url, "HF dataset /parquet").await?;
277 +
278 + // Partially converted datasets (>5GB) export under a "partial-" split prefix.
279 + let partial_split = format!("partial-{split}");
280 + let shard_paths: Vec<String> = listing["parquet_files"]
281 + .as_array()
282 + .map(|files| {
283 + files
284 + .iter()
285 + .filter(|f| {
286 + f["config"].as_str() == Some(config)
287 + && matches!(f["split"].as_str(), Some(s) if s == split || s == partial_split)
288 + })
289 + .filter_map(|f| shard_repo_path(f["url"].as_str()?))
290 + .collect()
291 + })
292 + .unwrap_or_default();
293 +
294 + if shard_paths.is_empty() {
295 + return Err(BenchError::Config(format!(
296 + "no parquet export for config '{config}' split '{split}'"
297 + )));
298 + }
299 + Ok(shard_paths)
300 +}
301 +
302 +/// File extensions the native main-branch fallback can decode.
303 +const NATIVE_DATA_EXTENSIONS: &[&str] = &[".parquet", ".json", ".jsonl", ".ndjson"];
304 +
305 +/// Discover native data files on the main branch matching a config/split, following
306 +/// the datasets library's default layout conventions (config directories, split-named
307 +/// files, unlabeled data files = train). Used when the dataset viewer has no parquet
308 +/// export, e.g. for private datasets.
309 +async fn native_data_files(dataset: &str, config: &str, split: &str) -> Result<Vec<String>> {
310 + let repo = crate::hub::HubRepo::dataset(dataset.to_string()).map_err(BenchError::Config)?;
311 + let files = repo.list_files().await.map_err(|e| {
312 + BenchError::Config(format!(
313 + "failed to list files of HF dataset '{dataset}': {e}"
314 + ))
315 + })?;
316 +
317 + let matched = select_native_files(files, config, split)?;
318 + if matched.is_empty() {
319 + return Err(BenchError::Config(format!(
320 + "HF dataset '{dataset}' has no usable data for config '{config}' split \
321 + '{split}': no dataset-viewer parquet export and no native data files matching \
322 + the split on the main branch. Check the dataset page or the \
323 + --hf-subset/--hf-split values; private datasets must store parquet, json, or \
324 + jsonl files."
325 + )));
326 + }
327 + Ok(matched)
328 +}
329 +
330 +/// Pure selection logic for [`native_data_files`].
331 +///
332 +/// Mirrors the datasets library's defaults: files may be scoped under a config
333 +/// directory; split-named files are preferred; a repo whose data files carry no split
334 +/// names at all maps everything to the "train" split.
335 +///
336 +/// Errors when the config matches no directory but the selection still spans several
337 +/// directories: the config never got resolved (e.g. a private dataset whose `/info`
338 +/// lookup failed falls back to the literal "default"), and pooling those directories
339 +/// would silently mix subsets.
340 +fn select_native_files(files: Vec<String>, config: &str, split: &str) -> Result<Vec<String>> {
341 + let data: Vec<String> = files
342 + .into_iter()
343 + .filter(|f| NATIVE_DATA_EXTENSIONS.iter().any(|ext| f.ends_with(ext)))
344 + .collect();
345 +
346 + // Multi-config layout scopes files under a config directory (e.g. "main/test-*.parquet")
347 + let config_prefix = format!("{config}/");
348 + let config_scoped = data.iter().any(|f| f.starts_with(&config_prefix));
349 + let scoped: Vec<String> = if config_scoped {
350 + data.into_iter().filter(|f| f.starts_with(&config_prefix)).collect()
351 + } else {
352 + data
353 + };
354 +
355 + let matched: Vec<String> =
356 + scoped.iter().filter(|f| path_matches_split(f, split)).cloned().collect();
357 +
358 + let selected = if !matched.is_empty() {
359 + matched
360 + } else {
361 + // No file carries the requested split name. If no file carries any split name,
362 + // the datasets library maps them all to "train".
363 + let split_keywords = ["train", "test", "validation", "valid", "dev"];
364 + let any_labeled =
365 + scoped.iter().any(|f| split_keywords.iter().any(|s| path_matches_split(f, s)));
366 + if split == "train" && !any_labeled {
367 + scoped
368 + } else {
369 + Vec::new()
370 }
371 + };
372
267 - for row in rows {
268 - if let Some(row_data) = row.get("row") {
269 - all_rows.push(row_data.clone());
270 - }
373 + if !config_scoped {
374 + let dirs = top_level_dirs(&selected);
375 + if dirs.len() > 1 {
376 + return Err(BenchError::Config(format!(
377 + "config '{config}' matches no directory in this dataset, and split '{split}' \
378 + data spans multiple config directories {dirs:?}. Pass --hf-subset to pick one."
379 + )));
380 }
381 + }
382 + Ok(selected)
383 +}
384
273 - let fetched = rows.len();
274 - offset += fetched;
385 +/// Distinct leading path segments of files that live in a directory.
386 +fn top_level_dirs(files: &[String]) -> Vec<&str> {
387 + let mut dirs: Vec<&str> =
388 + files.iter().filter_map(|f| f.split_once('/').map(|(d, _)| d)).collect();
389 + dirs.sort_unstable();
390 + dirs.dedup();
391 + dirs
392 +}
393
276 - let total = data["num_rows_total"].as_u64().unwrap_or(0);
277 - progress.update(offset, total);
394 +/// Split match on a repo file path: a split-named file ("{split}-00000-of-00001.parquet",
395 +/// "{split}.jsonl", "{split}_0.parquet") or a "{split}/" directory segment.
396 +fn path_matches_split(path: &str, split: &str) -> bool {
397 + let base = path.rsplit('/').next().unwrap_or(path);
398 + base.strip_prefix(split)
399 + .is_some_and(|rest| rest.starts_with('.') || rest.starts_with('-') || rest.starts_with('_'))
400 + || path.starts_with(&format!("{split}/"))
401 + || path.contains(&format!("/{split}/"))
402 +}
403
279 - // Stop if we have enough rows or reached end of dataset
280 - if all_rows.len() >= num_rows_needed || fetched < page_size {
404 +/// Fetch shards via hf-hub (cached under `$HF_HOME/hub`, shared with other Hub tooling)
405 +/// and decode rows in order until `target_rows` are collected or all shards are read.
406 +async fn read_shards(
407 + repo: &crate::hub::HubRepo,
408 + dataset: &str,
409 + shard_paths: Vec<String>,
410 + target_rows: usize,
411 +) -> Result<Vec<serde_json::Value>> {
412 + let total_shards = shard_paths.len();
413 + let mut rows: Vec<serde_json::Value> = Vec::new();
414 + for (i, shard_path) in shard_paths.into_iter().enumerate() {
415 + if rows.len() >= target_rows {
416 break;
417 }
418 + tracing::info!(
419 + dataset,
420 + shard = i + 1,
421 + total_shards,
422 + path = %shard_path,
423 + "fetching dataset shard"
424 + );
425 + let local_path = repo.get(&shard_path).await.map_err(BenchError::Config)?;
426 + let limit = target_rows - rows.len();
427 + let is_parquet = shard_path.ends_with(".parquet");
428 + let shard_rows = tokio::task::spawn_blocking(move || {
429 + if is_parquet {
430 + read_parquet_rows(&local_path, limit)
431 + } else {
432 + read_json_rows(&local_path, limit)
433 + }
434 + })
435 + .await
436 + .map_err(|e| BenchError::Config(format!("shard read task failed: {e}")))??;
437 + rows.extend(shard_rows);
438 }
284 - progress.finish();
439
286 - if all_rows.is_empty() {
440 + if rows.is_empty() {
441 return Err(BenchError::Config(format!(
288 - "HF dataset '{dataset}' download returned no rows"
442 + "HF dataset '{dataset}' data files contained no rows"
443 )));
444 }
291 -
292 - // Save to cache
293 - let json_str = serde_json::to_string(&all_rows)?;
294 - std::fs::write(&cache_path, &json_str)?;
295 -
296 - let path_str = cache_path.to_string_lossy().to_string();
445 tracing::info!(
446 dataset,
299 - rows = all_rows.len(),
300 - path = %path_str,
301 - "saved Hugging Face dataset"
447 + rows = rows.len(),
448 + "loaded HF dataset from data files"
449 );
303 - Ok((path_str, resolved_config, resolved_split))
450 + Ok(rows)
451 +}
452 +
453 +/// Read up to `limit` rows from a parquet file as JSON objects.
454 +fn read_parquet_rows(path: &Path, limit: usize) -> Result<Vec<serde_json::Value>> {
455 + let file = std::fs::File::open(path)?;
456 + let reader = SerializedFileReader::new(file)
457 + .map_err(|e| BenchError::Config(format!("failed to open parquet shard: {e}")))?;
458 + let iter = reader
459 + .get_row_iter(None)
460 + .map_err(|e| BenchError::Config(format!("failed to read parquet shard: {e}")))?;
461 + let mut rows = Vec::new();
462 + for row in iter {
463 + if rows.len() >= limit {
464 + break;
465 + }
466 + let row =
467 + row.map_err(|e| BenchError::Config(format!("failed to decode parquet row: {e}")))?;
468 + rows.push(row.to_json_value());
469 + }
470 + Ok(rows)
471 +}
472 +
473 +/// Read up to `limit` rows from a native JSON data file: a JSON array of objects, an
474 +/// object wrapping a single array field, or JSON Lines.
475 +fn read_json_rows(path: &Path, limit: usize) -> Result<Vec<serde_json::Value>> {
476 + let content = std::fs::read_to_string(path)?;
477 +
478 + let doc: Option<serde_json::Value> = serde_json::from_str(&content).ok();
479 + let mut rows = match doc {
480 + Some(serde_json::Value::Array(rows)) => rows,
481 + // e.g. {"data": [...]} — a single array-of-objects field wrapping the rows
482 + Some(serde_json::Value::Object(mut obj)) => {
483 + let mut arrays: Vec<serde_json::Value> = obj
484 + .iter_mut()
485 + .filter(|(_, v)| {
486 + v.as_array().is_some_and(|a| !a.is_empty() && a.iter().all(|e| e.is_object()))
487 + })
488 + .map(|(_, v)| v.take())
489 + .collect();
490 + match (arrays.len(), arrays.pop()) {
491 + (1, Some(serde_json::Value::Array(rows))) => rows,
492 + _ => {
493 + return Err(BenchError::Config(
494 + "unsupported JSON data file layout: expected an array of rows, \
495 + an object with a single array-of-objects field, or JSON Lines"
496 + .into(),
497 + ));
498 + }
499 + }
500 + }
501 + // Not a single JSON document — try JSON Lines
502 + _ => content
503 + .lines()
504 + .filter(|l| !l.trim().is_empty())
505 + .take(limit)
506 + .map(|l| {
507 + serde_json::from_str(l)
508 + .map_err(|e| BenchError::Config(format!("invalid JSON line in data file: {e}")))
509 + })
510 + .collect::<Result<Vec<serde_json::Value>>>()?,
511 + };
512 +
513 + rows.truncate(limit);
514 + Ok(rows)
515 }
516
517 /// 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<String> {
683 None
684 }
685
475 -/// Load an HF dataset from the cached JSON file and convert to SampleRequests.
686 +/// Convert downloaded HF dataset rows to SampleRequests.
687 pub fn load_hf_dataset(
688 tokenizer: &TokenizerKind,
478 - dataset_path: &str,
689 + entries: &[serde_json::Value],
690 num_requests: usize,
691 hf_output_len: Option<usize>,
692 seed: u64,
@@ -484,19 +695,8 @@ pub fn load_hf_dataset(
695 no_oversample: bool,
696 disable_shuffle: bool,
697 ) -> Result<Vec<SampleRequest>> {
487 - let content = std::fs::read_to_string(dataset_path).map_err(|e| {
488 - BenchError::Config(format!(
489 - "Failed to read HF dataset file '{dataset_path}': {e}"
490 - ))
491 - })?;
492 -
493 - let entries: Vec<serde_json::Value> = serde_json::from_str(&content)
494 - .map_err(|e| BenchError::Config(format!("Invalid JSON in HF dataset file: {e}")))?;
495 -
698 if entries.is_empty() {
497 - return Err(BenchError::Config(
498 - "HF dataset file contains no rows".into(),
499 - ));
699 + return Err(BenchError::Config("HF dataset contains no rows".into()));
700 }
701
702 // Detect column format from first row
@@ -638,7 +838,6 @@ pub fn load_hf_dataset(
838 } else {
839 if !warned_no_output {
840 tracing::warn!(
641 - path = dataset_path,
841 default_output_tokens = 128,
842 "no dataset output column or --hf-output-len; using default output length"
843 );
@@ -649,7 +848,6 @@ pub fn load_hf_dataset(
848 } else {
849 if !warned_no_output {
850 tracing::warn!(
652 - path = dataset_path,
851 default_output_tokens = 128,
852 "no dataset output column or --hf-output-len; using default output length"
853 );
@@ -965,16 +1163,6 @@ mod tests {
1163 assert!(!is_chat_array(&serde_json::json!("not an array")));
1164 }
1165
968 - #[test]
969 - fn test_sanitize_name() {
970 - assert_eq!(
971 - sanitize_name("allenai/WildChat-4.8M"),
972 - "allenai_WildChat-4.8M"
973 - );
974 - assert_eq!(sanitize_name("simple-name"), "simple-name");
975 - assert_eq!(sanitize_name("org/sub/deep"), "org_sub_deep");
976 - }
977 -
1166 // --- extract_chat_prompt edge cases ---
1167
1168 #[test]
@@ -1043,13 +1231,6 @@ mod tests {
1231 )
1232 }
1233
1046 - /// Write JSON data to a unique temp file and return the path string.
1047 - fn write_temp_json(name: &str, data: &serde_json::Value) -> String {
1048 - let path = std::env::temp_dir().join(format!("vllm-bench-test-{name}.json"));
1049 - std::fs::write(&path, serde_json::to_string(data).unwrap()).unwrap();
1050 - path.to_string_lossy().to_string()
1051 - }
1052 -
1234 #[test]
1235 fn test_load_hf_dataset_chat_format() {
1236 let tok = builtin_tokenizer();
@@ -1067,10 +1248,19 @@ mod tests {
1248 ]
1249 }
1250 ]);
1070 - let path = write_temp_json("chat-format", &data);
1251
1072 - let result =
1073 - load_hf_dataset(&tok, &path, 2, Some(50), 42, "test-", None, false, false).unwrap();
1252 + let result = load_hf_dataset(
1253 + &tok,
1254 + data.as_array().unwrap(),
1255 + 2,
1256 + Some(50),
1257 + 42,
1258 + "test-",
1259 + None,
1260 + false,
1261 + false,
1262 + )
1263 + .unwrap();
1264
1265 assert_eq!(result.len(), 2);
1266 assert!(
@@ -1087,9 +1277,19 @@ mod tests {
1277 {"prompt": "Explain the difference between machine learning and deep learning in simple terms.", "completion": "Machine learning is a subset of AI."},
1278 {"prompt": "How does photosynthesis work in plants and what role does chlorophyll play?", "completion": "Photosynthesis converts light to energy."}
1279 ]);
1090 - let path = write_temp_json("plain-text-format", &data);
1280
1092 - let result = load_hf_dataset(&tok, &path, 3, None, 0, "req-", None, false, false).unwrap();
1281 + let result = load_hf_dataset(
1282 + &tok,
1283 + data.as_array().unwrap(),
1284 + 3,
1285 + None,
1286 + 0,
1287 + "req-",
1288 + None,
1289 + false,
1290 + false,
1291 + )
1292 + .unwrap();
1293
1294 assert_eq!(result.len(), 3);
1295 // Without hf_output_len override, output len is derived from the completion tokens
@@ -1111,10 +1311,19 @@ mod tests {
1311 "answers": ["Albert Einstein"]
1312 }
1313 ]);
1114 - let path = write_temp_json("combined-format", &data);
1314
1116 - let result =
1117 - load_hf_dataset(&tok, &path, 2, Some(64), 1, "comb-", None, false, false).unwrap();
1315 + let result = load_hf_dataset(
1316 + &tok,
1317 + data.as_array().unwrap(),
1318 + 2,
1319 + Some(64),
1320 + 1,
1321 + "comb-",
1322 + None,
1323 + false,
1324 + false,
1325 + )
1326 + .unwrap();
1327
1328 assert_eq!(result.len(), 2);
1329 assert!(result.iter().all(|r| r.expected_output_len == 64));
@@ -1126,9 +1335,18 @@ mod tests {
1335 fn test_load_hf_dataset_empty_returns_error() {
1336 let tok = builtin_tokenizer();
1337 let data = serde_json::json!([]);
1129 - let path = write_temp_json("empty-dataset", &data);
1338
1131 - let result = load_hf_dataset(&tok, &path, 5, None, 0, "test-", None, false, false);
1339 + let result = load_hf_dataset(
1340 + &tok,
1341 + data.as_array().unwrap(),
1342 + 5,
1343 + None,
1344 + 0,
1345 + "test-",
1346 + None,
1347 + false,
1348 + false,
1349 + );
1350
1351 assert!(result.is_err(), "empty dataset should return an error");
1352 let err = result.unwrap_err().to_string();
@@ -1147,9 +1365,18 @@ mod tests {
1365 {"prompt": "Yes", "completion": "No"},
1366 {"prompt": "Ok", "completion": "Fine"}
1367 ]);
1150 - let path = write_temp_json("all-filtered", &data);
1368
1152 - let result = load_hf_dataset(&tok, &path, 3, None, 0, "test-", None, false, false);
1369 + let result = load_hf_dataset(
1370 + &tok,
1371 + data.as_array().unwrap(),
1372 + 3,
1373 + None,
1374 + 0,
1375 + "test-",
1376 + None,
1377 + false,
1378 + false,
1379 + );
1380
1381 assert!(
1382 result.is_err(),
@@ -1165,12 +1392,11 @@ mod tests {
1392 {"prompt": "What are the key differences between supervised and unsupervised machine learning?", "completion": "Supervised learning uses labeled data while unsupervised does not."},
1393 {"prompt": "Explain how neural networks are inspired by the human brain structure and function.", "completion": "Neural networks mimic brain neurons with layers of nodes."}
1394 ]);
1168 - let path = write_temp_json("output-len-override", &data);
1395
1396 let fixed_len = 77usize;
1397 let result = load_hf_dataset(
1398 &tok,
1173 - &path,
1399 + data.as_array().unwrap(),
1400 3,
1401 Some(fixed_len),
1402 42,
@@ -1196,11 +1422,10 @@ mod tests {
1422 {"prompt": "Explain how photosynthesis converts sunlight to energy in plant cells.", "completion": "Plants use chlorophyll to absorb sunlight."},
1423 {"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."}
1424 ]);
1199 - let path = write_temp_json("no-oversample", &data);
1425
1426 let result = load_hf_dataset(
1427 &tok,
1203 - &path,
1428 + data.as_array().unwrap(),
1429 10,
1430 Some(32),
1431 0,
@@ -1227,11 +1452,10 @@ mod tests {
1452 {"prompt": "Beta prompt about the second topic continuing from our ordered sequence here.", "completion": "Beta answer"},
1453 {"prompt": "Gamma prompt about the third item in our clearly ordered alphabetical sequence.", "completion": "Gamma answer"}
1454 ]);
1230 - let path = write_temp_json("disable-shuffle", &data);
1455
1456 let result = load_hf_dataset(
1457 &tok,
1234 - &path,
1458 + data.as_array().unwrap(),
1459 3,
1460 Some(20),
1461 99,
@@ -1264,13 +1488,188 @@ mod tests {
1488 {"prompt": "What is the boiling point of water at sea level under standard atmospheric pressure?", "completion": "Water boils at 100 degrees Celsius."},
1489 {"prompt": "Describe the structure of DNA and how genetic information is encoded within it.", "completion": "DNA is a double helix with base pairs."}
1490 ]);
1267 - let path = write_temp_json("request-id-prefix", &data);
1491
1269 - let result =
1270 - load_hf_dataset(&tok, &path, 2, Some(10), 0, "myprefix-", None, false, false).unwrap();
1492 + let result = load_hf_dataset(
1493 + &tok,
1494 + data.as_array().unwrap(),
1495 + 2,
1496 + Some(10),
1497 + 0,
1498 + "myprefix-",
1499 + None,
1500 + false,
1501 + false,
1502 + )
1503 + .unwrap();
1504
1505 assert_eq!(result.len(), 2);
1506 assert_eq!(result[0].request_id.as_deref(), Some("myprefix-0"));
1507 assert_eq!(result[1].request_id.as_deref(), Some("myprefix-1"));
1508 }
1509 +
1510 + // Network test: the full GSM8k test split must arrive via parquet shards regardless
1511 + // of how many rows earlier runs requested (run with --ignored).
1512 + #[tokio::test]
1513 + #[ignore]
1514 + async fn test_download_gsm8k_full_split_via_parquet() {
1515 + let (rows, config, split) =
1516 + download_hf_dataset("openai/gsm8k", None, Some("test"), 1319, 0, false)
1517 + .await
1518 + .unwrap();
1519 + assert_eq!(config, "main");
1520 + assert_eq!(split, "test");
1521 + assert_eq!(rows.len(), 1319);
1522 + assert!(rows[0].get("question").is_some_and(|q| q.is_string()));
1523 + assert!(rows[0].get("answer").is_some_and(|a| a.is_string()));
1524 + }
1525 +
1526 + // Network test: native parquet discovery on the main branch (private-dataset path)
1527 + // must find GSM8k's config-dir layout (run with --ignored).
1528 + #[tokio::test]
1529 + #[ignore]
1530 + async fn test_native_data_files_main_branch() {
1531 + let shards = native_data_files("openai/gsm8k", "main", "test").await.unwrap();
1532 + assert_eq!(shards, vec!["main/test-00000-of-00001.parquet".to_string()]);
1533 +
1534 + // Single unlabeled JSON file maps to the train split (private-dataset layout)
1535 + let shards = native_data_files("Inferact/codex_swebenchpro_traces", "default", "train")
1536 + .await
1537 + .unwrap();
1538 + assert_eq!(shards, vec!["codex_swebenchpro.json".to_string()]);
1539 + }
1540 +
1541 + // --- native data file selection ---
1542 +
1543 + #[test]
1544 + fn test_path_matches_split() {
1545 + assert!(path_matches_split(
1546 + "data/train-00000-of-00042.parquet",
1547 + "train"
1548 + ));
1549 + assert!(path_matches_split(
1550 + "main/test-00000-of-00001.parquet",
1551 + "test"
1552 + ));
1553 + assert!(path_matches_split("test.parquet", "test"));
1554 + assert!(path_matches_split("test.jsonl", "test"));
1555 + assert!(path_matches_split("test_0.parquet", "test"));
1556 + assert!(path_matches_split("data/test/0000.parquet", "test"));
1557 + assert!(path_matches_split("test/0000.parquet", "test"));
1558 +
1559 + assert!(!path_matches_split(
1560 + "data/train-00000-of-00042.parquet",
1561 + "test"
1562 + ));
1563 + assert!(!path_matches_split("testing-00000.parquet", "test"));
1564 + assert!(!path_matches_split("contest/0000.parquet", "test"));
1565 + assert!(!path_matches_split("data/latest/0000.parquet", "test"));
1566 + }
1567 +
1568 + #[test]
1569 + fn test_select_native_files_split_labeled() {
1570 + let files = vec![
1571 + ".gitattributes".to_string(),
1572 + "README.md".to_string(),
1573 + "data/train-00000-of-00002.parquet".to_string(),
1574 + "data/train-00001-of-00002.parquet".to_string(),
1575 + "data/test-00000-of-00001.parquet".to_string(),
1576 + ];
1577 + let selected = select_native_files(files, "default", "test").unwrap();
1578 + assert_eq!(
1579 + selected,
1580 + vec!["data/test-00000-of-00001.parquet".to_string()]
1581 + );
1582 + }
1583 +
1584 + #[test]
1585 + fn test_select_native_files_unlabeled_maps_to_train() {
1586 + let files = vec![
1587 + "README.md".to_string(),
1588 + "codex_swebenchpro.json".to_string(),
1589 + ];
1590 + let selected = select_native_files(files.clone(), "default", "train").unwrap();
1591 + assert_eq!(selected, vec!["codex_swebenchpro.json".to_string()]);
1592 + // ...but only for the train split
1593 + assert!(select_native_files(files, "default", "test").unwrap().is_empty());
1594 + }
1595 +
1596 + #[test]
1597 + fn test_select_native_files_labeled_files_do_not_leak_to_train() {
1598 + // A repo with only test-labeled files has no train split
1599 + let files = vec!["data/test-00000-of-00001.parquet".to_string()];
1600 + assert!(select_native_files(files, "default", "train").unwrap().is_empty());
1601 + }
1602 +
1603 + #[test]
1604 + fn test_select_native_files_config_dir_scoping() {
1605 + let files = vec![
1606 + "main/test-00000-of-00001.parquet".to_string(),
1607 + "socratic/test-00000-of-00001.parquet".to_string(),
1608 + ];
1609 + let selected = select_native_files(files, "socratic", "test").unwrap();
1610 + assert_eq!(
1611 + selected,
1612 + vec!["socratic/test-00000-of-00001.parquet".to_string()]
1613 + );
1614 + }
1615 +
1616 + #[test]
1617 + fn test_select_native_files_unresolved_config_does_not_pool_configs() {
1618 + // A private dataset whose /info lookup failed falls back to config "default";
1619 + // pooling both config directories would silently mix subsets.
1620 + let files = vec![
1621 + "main/test-00000-of-00001.parquet".to_string(),
1622 + "socratic/test-00000-of-00001.parquet".to_string(),
1623 + ];
1624 + let err = select_native_files(files, "default", "test").unwrap_err().to_string();
1625 + assert!(err.contains("--hf-subset"), "{err}");
1626 + }
1627 +
1628 + #[test]
1629 + fn test_select_native_files_split_dirs_are_not_ambiguous() {
1630 + // Split directories are not config directories
1631 + let files = vec![
1632 + "train/0000.parquet".to_string(),
1633 + "test/0000.parquet".to_string(),
1634 + ];
1635 + let selected = select_native_files(files, "default", "test").unwrap();
1636 + assert_eq!(selected, vec!["test/0000.parquet".to_string()]);
1637 + }
1638 +
1639 + // --- native JSON data file reading ---
1640 +
1641 + fn write_temp_file(name: &str, content: &str) -> std::path::PathBuf {
1642 + let path = std::env::temp_dir().join(format!("vllm-bench-test-json-{name}"));
1643 + std::fs::write(&path, content).unwrap();
1644 + path
1645 + }
1646 +
1647 + #[test]
1648 + fn test_read_json_rows_array() {
1649 + let path = write_temp_file("array.json", r#"[{"a": 1}, {"a": 2}, {"a": 3}]"#);
1650 + let rows = read_json_rows(&path, 2).unwrap();
1651 + assert_eq!(rows.len(), 2);
1652 + assert_eq!(rows[0]["a"], 1);
1653 + }
1654 +
1655 + #[test]
1656 + fn test_read_json_rows_wrapped_object() {
1657 + let path = write_temp_file("wrapped.json", r#"{"data": [{"a": 1}, {"a": 2}]}"#);
1658 + let rows = read_json_rows(&path, 10).unwrap();
1659 + assert_eq!(rows.len(), 2);
1660 + }
1661 +
1662 + #[test]
1663 + fn test_read_json_rows_jsonl() {
1664 + let path = write_temp_file("rows.jsonl", "{\"a\": 1}\n{\"a\": 2}\n\n{\"a\": 3}\n");
1665 + let rows = read_json_rows(&path, 10).unwrap();
1666 + assert_eq!(rows.len(), 3);
1667 + assert_eq!(rows[2]["a"], 3);
1668 + }
1669 +
1670 + #[test]
1671 + fn test_read_json_rows_unsupported_layout() {
1672 + let path = write_temp_file("scalar.json", r#"{"a": 1, "b": 2}"#);
1673 + assert!(read_json_rows(&path, 10).is_err());
1674 + }
1675 }
rust/src/bench/src/hub.rs
+11 -1
@@ -3,8 +3,8 @@
3
4 use std::path::PathBuf;
5
6 -use hf_hub::Repo;
6 use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
7 +use hf_hub::{Repo, RepoType};
8
9 /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
10 pub struct HubRepo {
@@ -20,6 +20,10 @@ impl HubRepo {
20 Self::new(Repo::dataset(repo_id))
21 }
22
23 + pub fn dataset_with_revision(repo_id: String, revision: String) -> Result<Self, String> {
24 + Self::new(Repo::with_revision(repo_id, RepoType::Dataset, revision))
25 + }
26 +
27 fn new(repo: Repo) -> Result<Self, String> {
28 let mut builder = ApiBuilder::from_env();
29 if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -36,4 +40,10 @@ impl HubRepo {
40 pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
41 self.repo.get(filename).await.map_err(|e| format!("{e}"))
42 }
43 +
44 + /// List all file paths in the repo at its revision.
45 + pub async fn list_files(&self) -> Result<Vec<String>, String> {
46 + let info = self.repo.info().await.map_err(|e| format!("{e}"))?;
47 + Ok(info.siblings.into_iter().map(|s| s.rfilename).collect())
48 + }
49 }