Skip zero-quota splits in get_dataset_samples - #2043
Conversation
📝 WalkthroughWalkthroughDataset sampling now rejects negative sample counts. Split quotas are calculated before loading, and streamed datasets load lazily only for splits with nonzero quotas. Tests verify that zero-quota splits are skipped. ChangesDataset loading
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The lazy loading and n <= 0 guard correctly preserve the existing quota/output ordering while avoiding all zero-quota split opens. The type widening for the default None split is also appropriate. Please add a small mocked unit regression test with fewer samples than splits that asserts load_dataset is called only for the positive-quota split(s); the existing live GPU test demonstrates the performance improvement but does not deterministically protect the behavior and depends on Hub latency.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 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 `@modelopt/torch/utils/dataset_utils.py`:
- Around line 631-639: Validate external num_samples at the function interface
before computing num_per_split, rejecting negative values while preserving valid
zero requests. In the split-loading loop around _load_split, change the quota
skip condition from n <= 0 to n == 0 so negative quotas cannot be silently
treated as zero.
🪄 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: Enterprise
Run ID: 1911d8b4-0af6-4da5-9d36-f7cce25a65bb
📒 Files selected for processing (1)
modelopt/torch/utils/dataset_utils.py
|
When num_samples is smaller than the number of splits, integer division gives most splits a quota of 0 and dumps the remainder on the last one. The old code still opened those splits, and merely starting a streamed split's iterator fetches its first record batch before the `i >= 0` break -- a real download for zero samples. For the `nemotron-post-training-v3` combo (14 samples over 7 datasets, 2 each) this meant 24 split-opens that contribute nothing, including 4 splits of Nemotron-Math-v2 whose parquet shards are ~12.6 GB each. That is most of the runtime of tests/gpu/torch/utils/test_dataset_utils.py, which has been timing out on the 120s tests/gpu cap in CI since 2026-07-31. Load splits lazily and skip zero-quota ones. Sample output is unchanged: the remainder still lands on the last split, so the same rows are returned. Also reject negative num_samples at the interface boundary. Floor division turned num_samples=-1 over 3 splits into quotas [-1, -1, 1], silently returning one sample for a negative request. Locally against the real gated datasets: 120s (timeout, reproduces the CI failure) -> 25s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
e60d534 to
3747a43
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
Re-review complete. The prior critical test-coverage concern is addressed by the new offline unit regression test: with one sample across two splits, it spies on datasets.load_dataset and verifies that only the positive-quota test split is opened. The lazy loading and zero-quota skip preserve quota/output ordering, and the negative num_samples validation resolves the separate negative-quota edge case. No new correctness or licensing concerns found.
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 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 `@tests/unit/torch/utils/test_dataset_utils.py`:
- Line 876: Add a brief inline comment immediately before
pytest.importorskip("datasets") explaining that the local import is intentional
because datasets is an optional dependency.
🪄 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: Enterprise
Run ID: 73761a76-dfe8-43f4-aff8-782b1431dc46
📒 Files selected for processing (2)
modelopt/torch/utils/dataset_utils.pytests/unit/torch/utils/test_dataset_utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/utils/dataset_utils.py
|
|
||
| def test_zero_quota_split_is_not_loaded(self, monkeypatch, make_toy_hf_dataset): | ||
| """1 sample over 2 splits → quotas [0, 1]: the zero-quota split is never opened.""" | ||
| datasets = pytest.importorskip("datasets") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Document the optional dependency import.
At Line 876, pytest.importorskip("datasets") imports an optional dependency inside the test body. Add a brief comment that explains why this local import is allowed.
Suggested fix
- datasets = pytest.importorskip("datasets")
+ # Optional dependency: skip this regression when Hugging Face Datasets is unavailable.
+ datasets = pytest.importorskip("datasets")As per path instructions, in-test imports are allowed for optional dependencies only when a brief comment names the reason.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| datasets = pytest.importorskip("datasets") | |
| # Optional dependency: skip this regression when Hugging Face Datasets is unavailable. | |
| datasets = pytest.importorskip("datasets") |
🤖 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/torch/utils/test_dataset_utils.py` at line 876, Add a brief inline
comment immediately before pytest.importorskip("datasets") explaining that the
local import is intentional because datasets is an optional dependency.
Source: Path instructions
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2043 +/- ##
==========================================
+ Coverage 69.97% 77.40% +7.42%
==========================================
Files 519 519
Lines 59399 59401 +2
==========================================
+ Hits 41565 45979 +4414
+ Misses 17834 13422 -4412
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
What does this PR do?
Type of change: Bug fix
get_dataset_samplesdividesnum_samplesevenly across splits and dumps the remainder on the last one. Whennum_samplesis smaller than the split count, most splits get a quota of0— but the old code still opened them:Merely starting a streamed split's iterator makes
datasetsfetch its first record batch, so a zero-quota split costs a real download and yields nothing.For the
nemotron-post-training-v3combo (14 samples over 7 datasets → 2 each) this is 24 split-opens that contribute zero samples, including 4 splits ofnvidia/Nemotron-Math-v2, whose parquet shards are ~12.6 GB each. Streaming those means reading a footer at byte ~12.6 GB and then pulling multi-MB row groups, all discarded.This PR loads splits lazily and skips zero-quota ones. Sample output is unchanged — the remainder still lands on the last split, so the same rows are returned.
Testing
tests/gpu/torch/utils/test_dataset_utils.py::test_get_dataset_dataloader_nemotron_v3_chat_templatehas been timing out on the 120 stests/gpucap in CI since 2026-07-31 — onmain's nightly and on every PR that runs the GPU job (example). The dumped stack shows it blocked in a socket read insidehf_file_system._fetch_rangewhile streamingNemotron-Math-v2.Reproduced locally against the real gated datasets and measured:
Timeout (>120.0s)as CItests/gpu/torch/utils/test_dataset_utils.py— 7 passedtests/unit/torch/utils/test_dataset_utils.py— 49 passedNote the timeout is not a recent code regression: the test ran 50–72 s for weeks, then stepped to >120 s on 2026-07-31 with no ModelOpt commit, dependency change (
datasets4.8.5 /huggingface_hub1.14.0 /hf-xet1.5.0 /pyarrow24.0.0 identical across the boundary), or dataset-repo commit to explain it — Hub-side latency on Xet-backed range reads is the likely trigger. The wasted downloads were what left the test with no margin to absorb it; it now has ~4× headroom.Also rejects negative
num_samplesat the interface boundary (CodeRabbit) — floor division turnednum_samples=-1over 3 splits into quotas[-1, -1, 1], silently returning one sample. Pre-existing onmain, fixed here since it is adjacent to the quota logic.Before your PR is "Ready for review"
CONTRIBUTING.md: N/Atest_zero_quota_split_is_not_loadedintests/unit/torch/utils/test_dataset_utils.py(offline, uses the toy local dataset fixture); the existingtests/gpu/torch/utils/test_dataset_utils.pyis the test this unblocks🤖 Generated with Claude Code
Summary by CodeRabbit