Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BugenZhao merged 6 commits into from
Aug 21, 2026
Merged

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fcc2216
[Rust][Benchmark] Load HF datasets from parquet shards via hf-hub
esmeetu Aug 9, 2026
99f925b
[Rust][Benchmark] Support native JSON/JSONL data files for private HF…
esmeetu Aug 9, 2026
cf1644e
Merge branch 'main' into rust-bench-hf-parquet-hub-cache
esmeetu Aug 19, 2026
a58b63a
[Rust][Benchmark] Address review: brotli parquet codec, ambiguous config
esmeetu Aug 19, 2026
b004d38
Merge branch 'main' into rust-bench-hf-parquet-hub-cache
BugenZhao Aug 21, 2026
0180cd0
Merge branch 'main' into rust-bench-hf-parquet-hub-cache
BugenZhao Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 96 additions & 0 deletions rust/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions rust/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,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"] }

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

@claude claude Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

paste = "1.0.15"
prometheus-client = "0.24.0"
prometheus-client-derive-encode = "0.5.0"
Expand Down
6 changes: 3 additions & 3 deletions rust/src/bench/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand All @@ -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 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`)
Expand Down
1 change: 1 addition & 0 deletions rust/src/bench/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 10 additions & 5 deletions rust/src/bench/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <model-name> \
--dataset-name hf --dataset-path allenai/WildChat-4.8M \
Expand All @@ -160,7 +161,9 @@ 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 files
# (pass both --hf-subset and --hf-split for private repos).
HF_TOKEN=hf_xxx vllm-bench \
--backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
--dataset-name hf --dataset-path lmsys/lmsys-chat-1m \
Expand Down Expand Up @@ -431,7 +434,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

Expand Down Expand Up @@ -507,7 +510,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 |
Expand Down Expand Up @@ -794,13 +797,15 @@ 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)

## Environment Variables

| 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
Expand Down
19 changes: 10 additions & 9 deletions rust/src/bench/src/benchmark.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -548,17 +548,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
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,
Expand Down
Loading
Loading