feat: add TxtSlicesDataset to allow sampling slices from txt file for benchmarking

[jdebache]

Purpose

Sampling randomly directly from a tokenizer for benchmarking creates data that is not ideal to benchmark when using speculative decoding or expert parallelism.

On the other hand, random datasets are very flexible and offer complete control on the input and output sequence lengths, which is desirable to create reproducible benchmarks.

This PR introduces a new type of benchmarking dataset called TxtSlicesDataset which offers a compromise between the flexibility of a random dataset and the fidelity of a real dataset. It allows sampling slices from a user-provided txt file.

Content

• The implementation of TxtSlicesDataset
• Fixes to typing in datasets.py
• A unit test for the new dataset type

[gemini-code-assist]

Code Review

This pull request introduces TxtSlicesDataset for benchmarking, which samples data from a text file. It also includes significant refactoring by moving utility functions from datasets.py to a new dataset_utils.py file and improving typing throughout. The changes are well-structured. My review focuses on improving the robustness and reproducibility of the new TxtSlicesDataset and its tests. I've pointed out a resource leak in the tests and potential for non-reproducible behavior due to the use of the global random module. I've also identified a missing check that could lead to a crash with certain input files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[mergify]

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

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

[cursor]

Cursor Bugbot has reviewed your changes and found 4 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

Comment @cursor review or bugbot run to trigger another review on this PR

[mergify]

Documentation preview: https://vllm--30156.org.readthedocs.build/en/30156/

[mergify]

Hi @hypdeb, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jdebache]

Hey, I'm trying to run the doc generation locally to figure out what is wrong. How long is it supposed to take?

I still can't reproduce locally, but from the error logs, it looks like it's trying to mock vllm over and over again for some reason.

[jdebache]

Cannot reproduce the documentation generation error locally. I get:

INFO    -  Doc file 'features/spec_decode.md' contains an unrecognized relative link '../../tests/spec_decode/e2e', it was left as is.
INFO    -  Doc file 'design/fused_moe_modular_kernel.md' contains a link './moe_kernel_features.md#fused-moe-experts-kernels', but the doc 'design/moe_kernel_features.md' does not contain an anchor '#fused-moe-experts-kernels'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'gpu.md#nvidia-cuda', but the doc 'getting_started/installation/gpu.md' does not contain an anchor '#nvidia-cuda'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'gpu.md#amd-rocm', but the doc 'getting_started/installation/gpu.md' does not contain an anchor '#amd-rocm'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'gpu.md#intel-xpu', but the doc 'getting_started/installation/gpu.md' does not contain an anchor '#intel-xpu'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'cpu.md#intelamd-x86', but the doc 'getting_started/installation/cpu.md' does not contain an anchor '#intelamd-x86'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'cpu.md#arm-aarch64', but the doc 'getting_started/installation/cpu.md' does not contain an anchor '#arm-aarch64'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'cpu.md#apple-silicon', but the doc 'getting_started/installation/cpu.md' does not contain an anchor '#apple-silicon'.
INFO    -  Doc file 'getting_started/installation/README.md' contains a link 'cpu.md#ibm-z-s390x', but the doc 'getting_started/installation/cpu.md' does not contain an anchor '#ibm-z-s390x'.

Aborted with 7 warnings in strict mode!
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

[mergify]

Documentation preview: https://vllm--30156.org.readthedocs.build/en/30156/

[mergify]

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

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

[mergify]

Hi @hypdeb, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @hypdeb, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jdebache]

Hey @hmellor, would you mind having a look? We found this dataset generation useful while benchmarking with features that are data dependent such as EP and speculative decoding.

[jdebache]

Note that I wanted to make additional fixes to the typing in datasets.py, however, it caused a cascade of issues, in particular in the documentation generator. Keeping those would have made the changes too large.

[jdebache]

Okay, I have now addressed all review comments.

[DarkLight1337]

LGTM now, thanks for your patience!

[jdebache]

Thanks for your thorough review!