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

[esmeetu]

Purpose

Fix a silent under-sampling bug in the Rust bench's hf dataset loader and rework the download path.

The bug: HF datasets were cached as JSON snapshots truncated to the row count of whatever run happened first, keyed only by dataset-config-split. A later run requesting more prompts silently reused the truncated file — e.g. requesting the full 1319-prompt GSM8k test split yielded samples=100 requested=1319 + skipping dataset oversampling, so the benchmark ran on 100 prompts while claiming to be a GSM8k run.

The rework:

• Rows now come from the dataset-viewer's auto-converted parquet shards (refs/convert/parquet), downloaded via hf-hub into the standard HF hub cache ($HF_HOME/hub) — shared bidirectionally with hf download and Python huggingface_hub. The cache unit is the immutable shard, so cache contents can never depend on a previous run's request size; a physically truncated shard fails loudly (missing parquet footer) instead of silently under-sampling.
• The /rows pagination fallback is removed: per the dataset-viewer docs, /rows itself only works for datasets with parquet exports, so it could never succeed where the parquet path fails (verified empirically on script-based/gated datasets — both endpoints return identical errors).
• New fallback for private datasets (which the dataset viewer does not process): discover native *.parquet files on the main branch via hf-hub repo listing, following the datasets library's layout conventions (config dirs, split-named files).
• Sampling semantics aligned with Python vllm bench serve: shard order is shuffled with --seed (mirrors the datasets library's streaming shard shuffle instead of always sampling the split prefix), and up to 2x the requested rows are read so validity filtering backfills from real rows instead of duplicating samples.
• datasets-server discovery calls now use the same auth chain as hf-hub shard downloads: HF_TOKEN env var, then the hf auth login token file (previously only the env var, so gated datasets failed at discovery for logged-in users without the env var).
• parquet crate added with default-features = false + json — no arrow dependency; Row::to_json_value() feeds the existing serde_json-based column-detection pipeline unchanged. New deps in the graph: thrift, snap, lz4_flex (no rustls/ring; cargo metadata clean against deny.toml bans).

Also fixes the stale --random-range-ratio README row, which still documented the removed [len*r, len] / default-1.0 semantics (the flag defaults to 0.0 and samples [len*(1-r), len*(1+r)]), and stale AGENTS.md module descriptions.

Why this is not duplicating an existing PR

Checked open PRs touching this area:

• [Bugfix] Make Rust HF dataset endpoint configurable #50277 (draft) makes the datasets-server endpoint configurable for mirror routing — different problem; this PR incidentally helps that use case since shard downloads (the bulk of traffic) now honor HF_ENDPOINT via hf-hub, with only the two discovery calls remaining on datasets-server.huggingface.co. Will conflict textually in hf_dataset.rs; happy to coordinate.
• [Bugfix] Fix HF benchmark datasets crashing when --hf-split is omitted #49641 fixes the Python loader when --hf-split is omitted — different codebase path.
• [Rust][Benchmark] Support HF ShareGPT datasets in multi-turn mode #51104 is Rust multi-turn ShareGPT — no overlap in hf_dataset.rs.

Test commands and results

cargo test -p vllm-bench                    # 148 passed, 0 failed
cargo clippy -p vllm-bench --all-targets    # 0 warnings
cargo test -p vllm-bench -- --ignored test_download_gsm8k test_native_parquet
                                            # 2 passed (network: full-split regression guard + main-branch discovery)

Live-endpoint verification (Kimi-K3 @ 1M ctx, tiktoken tokenizer):

• GSM8k --num-prompts 1319: 1319/1319 successful (bug scenario previously produced 100); cold-cache dataset load 1.2s (one 419KB shard vs 14 paginated JSON calls before), warm 0.4s.
• Cache interop: pre-populating with hf download openai/gsm8k --repo-type dataset --revision refs/convert/parquet, then running the bench → blob mtimes byte-identical before/after (zero re-download).

Cross-implementation comparison vs Python vllm bench serve (same endpoint/tokenizer/seed, --skip-chat-template on Python):

test	Python	Rust
GSM8k full split input tokens	77527	77527 (identical)
MT-Bench 80-prompt input tokens	5218	5218 (identical)
result JSON schema	34 keys	same 34 + 11 additive

Known intentional divergence (pre-existing, documented in README): the Rust loader sends prompts raw and derives per-row output lengths from the answer column, while Python's per-dataset classes may apply chat templates and fixed default output lengths (GSM8k: 193599 input tokens with template vs 77527 raw).

Model evaluation

N/A — benchmark client change; no model/serving code touched. Note the sampling-semantics alignment (shard shuffle + 2x read) changes which rows are benchmarked for num_prompts < split size; full-split runs are unaffected.

AI assistance

This PR was developed with AI assistance (Claude Code). The submitter has reviewed the changes and run the tests above.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fcc2216690

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Enable Brotli before accepting arbitrary Parquet shards

When a dataset-viewer export or native private-dataset shard uses Parquet's Brotli compression, SerializedFileReader will reject it because default-features are disabled and the brotli codec feature is omitted here. This regresses those datasets from the previous /rows path, where decompression happened server-side, and causes the entire HF benchmark to abort; enable the Parquet Brotli feature or retain a compatible fallback.

Useful? React with 👍 / 👎.

[esmeetu]

@claude review

[claude]

Beyond the inline nits, this run also checked the partial-parquet handling for >5GB datasets (the partial-{split} shard matching in convert_parquet_shards) for silent under-sampling — it correctly includes those shards rather than skipping them, so the split's row count isn't silently truncated there.

Extended reasoning...

This is an informational-only note about one additional area examined and ruled out this run; it is not a guarantee of correctness and doesn't replace review of the two inline nits already posted (missing brotli codec feature, and the 'default' config fallback silently merging multi-config data for private datasets without --hf-subset).

[claude]

🟡 When a private dataset's /info lookup fails and --hf-subset is omitted, resolved_config falls back to the literal string "default". In select_native_files (hf_dataset.rs:297-306), no file in a multi-config layout (e.g. gsm8k's main//socratic/ dirs) starts with "default/", so the config-scoping filter is silently skipped and rows from every config sharing that split name get pooled together with no warning — the opposite of the PR's stated goal of failing loudly instead of silently returning wrong data. A safer fallback would be to error out (or warn explicitly) when multiple config directories exist and none matches the resolved config.

Extended reasoning...

The bug: download_hf_dataset resolves (config, split) via the datasets-server /info endpoint, but private datasets aren't processed by the dataset viewer, so /info (and later /parquet) both fail. In that case the code falls into the else branch at hf_dataset.rs (around line 232): subset.unwrap_or("default").to_string(). If the user didn't pass --hf-subset, resolved_config becomes the literal string "default" — not an actual discovered config name.

That bogus config then flows into native_data_files(dataset, "default", split) → select_native_files (hf_dataset.rs:297-306):

let config_prefix = format!("{config}/");
let scoped: Vec<String> = if data.iter().any(|f| f.starts_with(&config_prefix)) {
    data.into_iter().filter(|f| f.starts_with(&config_prefix)).collect()
} else {
    data
};

For a repo laid out like openai/gsm8k (config directories main/ and socratic/), no file starts with "default/", so the if is false and scoped falls through to all data files, unfiltered by config. The subsequent matched filter only checks path_matches_split(f, split), which matches split-named files in every config directory (e.g. both main/test-*.parquet and socratic/test-*.parquet). The function returns the pooled set with no error and no explicit 'merging configs' warning — the closest the user gets is the earlier, generic assuming config 'default' and split 'train' log line, which doesn't say config resolution failed to disambiguate multiple config directories.

Why nothing else catches this: The scoping logic in select_native_files is specifically designed to isolate a single config directory (see test_select_native_files_config_dir_scoping), but it only activates when the config string actually matches a directory prefix. The fallback literal "default" was never meant to be a real directory name, so it silently defeats the scoping mechanism it was supposed to invoke.

Impact: For a private, multi-config dataset accessed without --hf-subset, the benchmark silently samples prompts from multiple unrelated configs/subsets pooled together, corrupting the intended sample distribution (e.g. mixing gsm8k's main and socratic variants) without any indication to the user that this happened. This directly contradicts this PR's stated design goal: 'a physically truncated shard fails loudly ... instead of silently under-sampling' — here, instead, an ambiguous config fails quietly by merging data instead of erroring or asking the user to disambiguate.

Step-by-step proof:

1. User runs against a private dataset repo laid out like gsm8k (main/test-00000-of-00001.parquet, socratic/test-00000-of-00001.parquet), passing --hf-split test but no --hf-subset.
2. /info fails (401/403, private repo) → resolved_config = "default", resolved_split = "test".
3. /parquet also fails (viewer doesn't process private repos) → falls to native_data_files(dataset, "default", "test").
4. select_native_files filters files by extension → both parquet files remain in data.
5. config_prefix = "default/" — neither main/test-... nor socratic/test-... starts with it, so the if branch is false and scoped = data (both files).
6. matched filters by path_matches_split(f, "test") — both files match (both are named test-...), so both are returned.
7. Rows from main and socratic are downloaded and pooled into one sample set, with no warning that two distinct configs were merged.

Suggested fix: When the config falls back to the literal "default"" and multiple config-like directories exist in the file listing (i.e. scopedends up empty from config filtering but multiple distinct top-level dirs are present), fail loudly with an error asking the user to pass--hf-subset`, rather than silently keeping all files.

[claude]

🟡 rust/Cargo.toml:72 enables the parquet crate with default-features = false and only [json, snap, zstd, flate2, flate2-rust_backend, lz4], omitting brotli. A parquet shard (native private-dataset file, most plausibly) compressed with BROTLI will fail to decode in read_parquet_rows, aborting the whole HF benchmark run instead of the graceful per-codec handling the old /rows path provided. Fix is a one-line addition of "brotli" to the feature list (pure-Rust, no rustls/ring, so it stays deny.toml-clean).

Extended reasoning...

What the bug is. rust/Cargo.toml:72 declares:

parquet = { version = "59.2.0", default-features = false, features = ["json", "snap", "zstd", "flate2", "flate2-rust_backend", "lz4"] }

The parquet (arrow-rs) crate gates each compression codec behind a Cargo feature that is normally pulled in via default-features. By disabling defaults and hand-picking [json, snap, zstd, flate2, flate2-rust_backend, lz4], the brotli codec feature (and lzo) is left out. Cargo.lock confirms this: the compiled parquet dependency set is ahash, base64, bytes, chrono, flate2, half, hashbrown, lz4_flex, num-bigint, num-integer, num-traits, seq-macro, serde_json, snap, twox-hash, zstd — no brotli crate anywhere in the graph.

Code path that triggers it. read_parquet_rows in rust/src/bench/src/datasets/hf_dataset.rs opens a shard with SerializedFileReader::new(file) and iterates rows via reader.get_row_iter(None). When a column's compression codec is BROTLI and the codec feature is compiled out, the parquet crate returns an error at row-decode time (a "disabled feature" / unsupported-codec error) instead of a decoded row. That Result::Err is mapped to BenchError::Config(...) inside read_parquet_rows, which propagates up through read_shards (called per shard inside download_hf_dataset) with no per-codec fallback or retry — the entire HF-dataset benchmark run aborts rather than skipping/reporting just the offending shard.

Why existing code doesn't prevent it. The PR's whole design assumes shards downloaded via hf-hub can always be decoded locally by the embedded parquet reader. There is no codec-capability check before download, no catch-and-skip around get_row_iter/row decode, and no fallback to another shard or dataset path if one shard uses an unsupported codec. Previously, the (now-removed) /rows datasets-server pagination endpoint decompressed server-side for any codec, so this class of failure did not exist for any HF dataset regardless of client capabilities — this is a genuine capability regression introduced by moving to a local parquet reader with a shrunken feature set.

Impact. For a dataset whose parquet columns are Brotli-compressed, vllm-bench --dataset-name hf fails hard with a config error and produces no benchmark results at all, rather than degrading gracefully. This is most likely to bite the new native-parquet fallback path added by this PR (used for private datasets not processed by the HF dataset-viewer), since those files can be written with arbitrary compression by whoever authored the dataset, whereas the primary dataset-viewer-converted path is auto-generated by HuggingFace's parquet converter using pyarrow's default (Snappy) compression — so the common case is unaffected, but the fallback path this PR specifically adds is exactly the one most exposed to this gap.

Step-by-step proof:

1. A user has a private HF dataset repo with a train.parquet file whose columns were written with compression='BROTLI' (e.g. via pyarrow.parquet.write_table(..., compression='BROTLI')).
2. User runs vllm-bench --dataset-name hf --dataset-path <private-org>/<private-repo> --hf-split train --num-prompts 100.
3. download_hf_dataset calls convert_parquet_shards, which fails (dataset-viewer has no export for private datasets) and falls back to native_data_files, which finds train.parquet on main via hf-hub repo listing.
4. read_shards downloads the file and calls read_parquet_rows(&local_path, limit) in a blocking task.
5. SerializedFileReader::new(file) succeeds (footer/metadata parse doesn't need codec support), but reader.get_row_iter(None)'s row iteration hits the BROTLI-compressed column and returns Err(ParquetError::General("Disabled feature at compile time: brotli codec")) (or equivalent).
6. read_parquet_rows maps this to Err(BenchError::Config(...)), which read_shards propagates via ?, which download_hf_dataset returns as the overall Result, which benchmark.rs::run_benchmark propagates out of run_benchmark entirely — the whole benchmark invocation exits with an error and zero results, even though 99% of the dataset's other columns/rows might be perfectly readable.

Fix. Add "brotli" to the feature list on rust/Cargo.toml:72:

parquet = { version = "59.2.0", default-features = false, features = ["json", "snap", "zstd", "flate2", "flate2-rust_backend", "lz4", "brotli"] }

The parquet crate's brotli feature depends on the pure-Rust brotli crate (not brotli-decompressor's C bindings, and critically not rustls/ring/aws-lc/boring/s2n), so this stays clean against the rust/deny.toml crypto/TLS bans documented in rust/src/bench/AGENTS.md. This is a low-risk, mechanical dependency addition.

[BugenZhao]

LGTM!

[BugenZhao]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84794 for commit a58b63a6fd54.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @esmeetu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[BugenZhao]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84922 for commit b004d38f7c4d.

[BugenZhao]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84954 for commit 0180cd0d4c33.