fix(train): derive train/val split boundaries from one expression by eldarkurtic · Pull Request #868 · vllm-project/speculators

eldarkurtic · 2026-07-26T21:01:41Z

Summary

The ArrowDataset split logic computed the train boundary as int(n * ratio) and the val boundary as int(n * (ratio - 1.0)). These two expressions disagree in floating point at certain ratios — e.g. at ratio=0.2 with n=12501: int(12501 * 0.2) == 2500 but int(12501 * -0.8) → int(12501 * 0.2000…3) == 2501 — putting one row in both splits (silent train/val leakage).
The fix replaces the split_ratio: float parameter with train_ratio: float + split: Literal["train", "val"]. Both splits derive their boundary from the single expression int(len(data) * train_ratio), guaranteeing exact complementarity. Input validation rejects out-of-range ratios and val splits at train_ratio=1.0.

Test plan

34 unit tests (test_split_boundaries.py): exact complementarity across 10 ratios (including 0.2 which was broken), no row in both splits, union covers full dataset, default takes whole dataset, rejects out-of-range, rejects val at ratio 1.0
Verified the 0.2 ratio case produces off-by-one on unfixed code, exact complement on fixed code

🤖 Generated with Claude Code

coderabbitai · 2026-07-26T21:01:59Z

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 0804b1c3-3140-494b-888b-61cbed14517b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

Dataset split pipeline

Layer / File(s)	Summary
ArrowDataset split contract `src/speculators/train/data.py`	`ArrowDataset` replaces `split_ratio` with `train_ratio` and `split`, validates inputs, and selects contiguous train or validation ranges.
Train and validation loader wiring `src/speculators/train/dataloader.py`	Loader construction passes the shared ratio with explicit `train` and `val` split values.
Split boundary and error coverage `tests/unit/train/test_split_boundaries.py`, `tests/unit/train/test_data.py`	Tests verify complementary and disjoint partitions, full-dataset defaults, ratio validation, rejection of empty validation splits, and updated default-ratio indexing behavior.

Possibly related PRs

vllm-project/speculators#528: Both changes modify ArrowDataset split handling and self.start_file_idx.
vllm-project/speculators#709: Both changes update the train/validation ratio through the loader and Arrow dataset split pipeline.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main train/val split boundary fix.
Description check	✅ Passed	The description clearly matches the code changes and test coverage in the pull request.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

mergify · 2026-07-26T21:02:18Z

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require approval from approved reviewers list

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

mergify · 2026-07-26T21:03:35Z

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

coderabbitai

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/train/data.py`:
- Around line 212-230: Validate the split argument in the initializer before
calculating or selecting the slice, accepting only "train" and "val" and raising
a clear ValueError for any other value. Update the conditional assigning start
and stop so it relies on this validated input and preserves the existing
complementary partition behavior.

In `@tests/unit/train/test_split_boundaries.py`:
- Around line 30-36: Update the _split helper’s split parameter annotation to
use the same Literal["train", "val"] type as ArrowDataset.split, preserving the
existing train and validation call behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 1146c2fb-e3b7-4ce6-8910-4ca78564771d

📥 Commits

Reviewing files that changed from the base of the PR and between f7ec341 and 63f6f89.

📒 Files selected for processing (3)

src/speculators/train/data.py
src/speculators/train/dataloader.py
tests/unit/train/test_split_boundaries.py

eldarkurtic · 2026-07-26T21:25:06Z

@coderabbitai full review

coderabbitai · 2026-07-26T21:25:11Z

✅ Action performed

Full review finished.

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unit/train/test_split_boundaries.py (1)
36-47: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add tiny-dataset boundary coverage.

All parameterized datasets are large enough to keep the selected partitions non-empty. Add cases such as n=1 or n=2 and document whether an empty partition should raise or be supported.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/train/test_split_boundaries.py` around lines 36 - 47, Extend
test_splits_are_exactly_complementary with tiny dataset sizes such as n=1 and
n=2, and explicitly assert the intended behavior when the ratio produces an
empty train or validation partition. Document through the test expectations
whether empty partitions are supported or should raise, while preserving the
existing complementary split assertions for valid non-empty cases.
Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/train/data.py`:
- Around line 227-232: Define the empty-partition behavior in the split logic
around split_idx, start, stop, and self.data.select: either reject empty
selected ranges with a clear validation error or explicitly support them
consistently. Update src/speculators/train/data.py lines 227-232 accordingly,
and add small-dataset boundary cases in
tests/unit/train/test_split_boundaries.py lines 36-47 asserting the chosen
contract.

---

Nitpick comments:
In `@tests/unit/train/test_split_boundaries.py`:
- Around line 36-47: Extend test_splits_are_exactly_complementary with tiny
dataset sizes such as n=1 and n=2, and explicitly assert the intended behavior
when the ratio produces an empty train or validation partition. Document through
the test expectations whether empty partitions are supported or should raise,
while preserving the existing complementary split assertions for valid non-empty
cases.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d67e30d1-b831-41cb-ba12-a5af73146ec8

📥 Commits

Reviewing files that changed from the base of the PR and between f7ec341 and 1f40845.

📒 Files selected for processing (4)

src/speculators/train/data.py
src/speculators/train/dataloader.py
tests/unit/train/test_data.py
tests/unit/train/test_split_boundaries.py

`ArrowDataset` encoded the split in the sign of `split_ratio`: the caller passed `train_data_ratio` for train and `train_data_ratio - 1.0` for val, and the val branch reconstructed the boundary as `int(n * (1.0 + split_ratio))`. That round trip is not exact in binary floating point, so the two branches could disagree on where the split falls. At `train_data_ratio=0.2` over 12501 rows, train ended at `int(12501 * 0.2) == 2500` while val started at `int(12501 * (1.0 + (0.2 - 1.0))) == int(12501 * 0.19999999999999996) == 2499`. Row 2499 therefore appeared in *both* splits, and the val split was one row larger than requested. The trainer's default 0.9 happens to agree, so the bug was latent, but any ratio can hit it and the failure is silent -- train/val leakage with no error. Take the split side explicitly (`train_ratio` plus `split="train"|"val"`) and compute the boundary once with a single expression, so the two splits are exactly complementary by construction. Also reject out-of-range ratios and a val split requested at `train_ratio=1.0`, which would silently be empty. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R1mQaSLBM1djDWD6herKZw Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com>

## Summary - The `ArrowDataset` split logic computed the train boundary as `int(n * ratio)` and the val boundary as `int(n * (ratio - 1.0))`. These two expressions disagree in floating point at certain ratios — e.g. at `ratio=0.2` with `n=12501`: `int(12501 * 0.2) == 2500` but `int(12501 * -0.8) → int(12501 * 0.2000…3) == 2501` — putting **one row in both splits** (silent train/val leakage). - The fix replaces the `split_ratio: float` parameter with `train_ratio: float` + `split: Literal["train", "val"]`. Both splits derive their boundary from the single expression `int(len(data) * train_ratio)`, guaranteeing exact complementarity. Input validation rejects out-of-range ratios and val splits at `train_ratio=1.0`. ## Test plan - [x] 34 unit tests (`test_split_boundaries.py`): exact complementarity across 10 ratios (including 0.2 which was broken), no row in both splits, union covers full dataset, default takes whole dataset, rejects out-of-range, rejects val at ratio 1.0 - [x] Verified the 0.2 ratio case produces off-by-one on unfixed code, exact complement on fixed code 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Roderick-Wu <rowu@redhat.com>

eldarkurtic requested review from fynnsu and shanjiaz as code owners July 26, 2026 21:01

mergify Bot added the quality-failed label Jul 26, 2026

coderabbitai Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread src/speculators/train/data.py

Comment thread tests/unit/train/test_split_boundaries.py Outdated

eldarkurtic force-pushed the fix-split-ratio-rounding branch 4 times, most recently from 9cc70d5 to 1f40845 Compare July 26, 2026 21:23

mergify Bot removed the quality-failed label Jul 26, 2026

coderabbitai Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread src/speculators/train/data.py

eldarkurtic force-pushed the fix-split-ratio-rounding branch from 1f40845 to c2efa18 Compare July 26, 2026 21:43

fynnsu approved these changes Jul 27, 2026

View reviewed changes

fynnsu added the ready This PR is ready for review label Jul 27, 2026

fynnsu enabled auto-merge (squash) July 27, 2026 13:22

fynnsu merged commit a3b5c75 into vllm-project:main Jul 27, 2026
9 of 10 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(train): derive train/val split boundaries from one expression - #868

fix(train): derive train/val split boundaries from one expression#868
fynnsu merged 1 commit into
vllm-project:mainfrom
eldarkurtic:fix-split-ratio-rounding

eldarkurtic commented Jul 26, 2026

Uh oh!

coderabbitai Bot commented Jul 26, 2026 •

edited

Loading

Review skipped

Walkthrough

Changes

Uh oh!

mergify Bot commented Jul 26, 2026 •

edited

Loading

🟢 Require approval from approved reviewers list

Uh oh!

mergify Bot commented Jul 26, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

eldarkurtic commented Jul 26, 2026

Uh oh!

coderabbitai Bot commented Jul 26, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

eldarkurtic commented Jul 26, 2026

Summary

Test plan

Uh oh!

coderabbitai Bot commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

Uh oh!

mergify Bot commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Protections

🟢 Require approval from approved reviewers list

Uh oh!

mergify Bot commented Jul 26, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

eldarkurtic commented Jul 26, 2026

Uh oh!

coderabbitai Bot commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderabbitai Bot commented Jul 26, 2026 •

edited

Loading

mergify Bot commented Jul 26, 2026 •

edited

Loading

coderabbitai Bot commented Jul 26, 2026 •

edited

Loading