[solarsys]

Problem

XGBoostModel.fit warns when its training loader shuffles, because row order changes row subsampling. The check reads loader.dataset.shuffle. Loaders from get_dataloader (_SeededDataLoader) set that flag on the shared dataset at the start of each pass. So after any shuffled loader on the same dataset has run, for example another model's training loader, the flag stays True. fit with get_dataloader(ds, ..., shuffle=False) then warns that the loader shuffles, which is false.

This showed up in the new Tutorial 00, which trains an RNN and XGBoost on the same split. The notebook currently works around it by training XGBoost first.

Fix

Use the loader's own shuffle_samples setting when it has one. That is what applies when the loader iterates. Other loaders keep the old check: a RandomSampler, or the dataset's flag.

The dataset-flag check also missed in-memory datasets. InMemorySampleDataset, which create_sample_dataset returns by default, keeps its flag in _shuffle, not litdata's shuffle. So a plain DataLoader over a shuffled in-memory dataset never warned. The check now reads both. A comment at the check notes a limitation for other iterable datasets: get_dataloader sets their flag once, when a loader is created, so the flag follows the most recently created loader.

Test

New tests/core/test_gradient_boosted_trees.py, for CI coverage without xgboost. Every test of gradient_boosted_trees.py needed xgboost, which the CI test environment does not install, so CI ran none of them (27.86% coverage of the file). GradientBoostedTreeModel does not depend on the tree library, so this file fills its hooks with scikit-learn linear models. scikit-learn is a core dependency. It covers:

• the feature layout, every label mode, the fitting checks and scale_pos_weight;
• prediction, attributions and checkpoints;
• one shuffle test per dataset type:
  • test_shuffle_warning_follows_the_fit_loader, on a disk-backed dataset. After iterating a shuffled loader, the dataset's shuffle flag is left set. Fitting with an unshuffled loader then logs no warning, and fitting with a shuffled loader still warns. Without the fix, the no-warning assertion fails.
  • test_shuffle_warning_with_a_plain_dataloader, a plain DataLoader over map-style rows and over an in-memory dataset. Without the _shuffle fix, it fails.

None of these tests needs xgboost, so CI runs them all, the regression test included. Run without xgboost, the file's coverage is 96.79% (271 of 280 statements).

All 21 tests in tests/core/test_xgboost_model.py and 13 in tests/core/test_gradient_boosted_trees.py pass, ruff is clean, and tools/check_pr_rules.py passes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PhbzTePoLMBYwjMLG8R9ek

[github-actions]

Test Coverage

Total: 69.80% (17918 / 25671 lines)

Modified file	Coverage	Lines
pyhealth/autoimprove/__init__.py	100.00%	6 / 6
pyhealth/autoimprove/evaluate.py	92.07%	151 / 164
pyhealth/autoimprove/improver.py	95.63%	372 / 389
pyhealth/autoimprove/llm.py	90.62%	29 / 32
pyhealth/autoimprove/proposers.py	93.91%	108 / 115
pyhealth/autoimprove/space.py	95.31%	122 / 128
pyhealth/datasets/omop.py	100.00%	18 / 18
pyhealth/metrics/calibration.py	65.00%	52 / 80
pyhealth/metrics/multiclass.py	25.00%	26 / 104
pyhealth/models/__init__.py	100.00%	52 / 52
pyhealth/models/gradient_boosted_trees.py	96.79%	271 / 280
pyhealth/models/survival_mdn.py	93.83%	213 / 227
pyhealth/processors/deep_nested_sequence_processor.py	91.37%	180 / 197
pyhealth/processors/nested_sequence_processor.py	92.94%	158 / 170
pyhealth/processors/sequence_processor.py	95.77%	68 / 71
pyhealth/processors/stagenet_processor.py	84.08%	169 / 201

Workflow run · 4e7550b

[Logiquo]

/review

[github-actions]

Previous review

AI review: findings

This change makes XGBoostModel._collect (pyhealth/models/gradient_boosted_trees.py, around line 351) decide whether to warn about shuffling from the loader's own shuffle_samples setting when it has one, instead of the dataset's shuffle flag. The flag is only stale for get_dataloader's _SeededDataLoader, which applies its shuffle setting to the dataset at the start of each pass, so after another shuffled loader has run the flag can be left True. Plain DataLoaders have no shuffle_samples attribute and keep the old RandomSampler/dataset-flag check, which is still accurate because iterating a SampleDataset shuffles according to that same flag at iteration time. I checked that SampleDataset.set_shuffle keeps litdata's shuffle attribute in sync, and that the logger name the new test captures matches the module's getLogger(name). The new test test_shuffle_warning_follows_the_fit_loader covers both directions (no false warning after a shuffled loader on the same dataset, and the warning still fires for a shuffled fit loader), and the logic is correct on inspection. I could not run the test suite in this environment, so the claim that it fails without the fix and passes with it is unverified. No blocking problems found; the change is a warning-only edit inside one model's private method, with no duplicate of existing functionality and no effect on the core.

_{Reviewed commit a1f238d}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: need change

Thank you for this fix. Basing XGBoost's shuffle warning on the fit loader's own setting removes a false warning, and the new test covers both the quiet and the warning case.

The change itself looks correct. The one blocker is coverage.

To get this merged:

• Raise test coverage for pyhealth/models/gradient_boosted_trees.py. CI reports 27.86% (78 of 280 statements), which is below the 80% bar for a modified file. The whole repository is at 67.87%, which meets the 60% bar. Please add tests for the uncovered parts of this file. If the low figure comes from XGBoost-dependent tests being skipped in CI, please make sure they run there. If most of the uncovered lines belong to other models in that file, add tests for them, or explain in the PR why they can't be covered.

Optional suggestions:

• Add a test for a plain DataLoader, which has no shuffle_samples attribute. It would cover the fallback to the RandomSampler or dataset-flag check.

_{Reviewed commit a1f238d}

[solarsys]

Thanks for the review. Both points are addressed in 0e053f0.

Coverage of gradient_boosted_trees.py: 27.86% → 96.79% without xgboost

The low figure came from skipped tests. Every test of this file was in tests/core/test_xgboost_model.py and needed xgboost. xgboost is only an optional extra (pyhealth[xgboost]) and isn't in the pixi test environment, so CI skipped all of them except the install-hint test. That left only the module's definitions covered (78 of 280 statements).

GradientBoostedTreeModel doesn't depend on the tree library: subclasses only build, fit, predict with and serialize one estimator. The new tests/core/test_gradient_boosted_trees.py fills those hooks with scikit-learn linear models. scikit-learn is a core dependency, so these tests run in CI as it is. A linear model's contributions (coefficient times value, plus the intercept) add up to its margin exactly, as TreeSHAP values do. That means explain() is tested with the same additivity check as the XGBoost test. The 13 tests cover:

• each label mode (binary, multiclass, multilabel, regression): fitting with a validation set, Trainer.inference, forward with and without the label, explain, mean_abs_shap, and a checkpoint round trip through Trainer.save_ckpt/load_ckpt;
• the feature layout for each processor kind, custom feature_names, and the bag-of-codes width;
• the errors for inputs and batches that don't fit the layout;
• the fitting checks: a missing class, an empty loader, drop_last, the constructor checks, and every scale_pos_weight form;
• the checkpoint version, mode, layout and width checks.

I ran these two test files with xgboost made unimportable. Coverage of the file went from 27.86% (78 of 280), matching CI, to 96.79% (271 of 280). What's left is the abstract hooks that raise NotImplementedError, the unsupported-mode error, and one branch for a dense field with no width yet.

If you'd rather have the XGBoost tests run in CI too, xgboost would need to be added to the pixi test feature, which means re-locking pixi.lock. I've left that to you.

Optional: a plain DataLoader

test_shuffle_warning_with_a_plain_dataloader covers the fallback for loaders without shuffle_samples, and it found a gap. The fallback read loader.dataset.shuffle, litdata's flag. But InMemorySampleDataset, which create_sample_dataset returns by default, keeps its flag in _shuffle, so a plain loader over a shuffled in-memory dataset never warned. The check now reads both flags. The test covers:

• a list (map-style) with and without shuffle=True, which uses the RandomSampler path;
• an in-memory dataset with its flag on and off.

The XGBoost tests (22) and the new tests (13) pass locally, and the PR rules check passes.

[solarsys]

/review

[github-actions]

AI review: findings

The PR changes GradientBoostedTreeModel._collect (pyhealth/models/gradient_boosted_trees.py) so the "shuffles" warning for a training loader reads the loader's own shuffle_samples when it has one. get_dataloader's _SeededDataLoader always sets that attribute, so the stale-dataset-flag false warning is fixed for PyHealth's SampleDataset loaders. Plain loaders keep the old check, now also reading the in-memory dataset's _shuffle. The change only affects a log message, not row order or fit results, and nothing else in the repo performs this check, so it is not redundant. Blocking (medium): the new test_shuffle_warning_follows_the_fit_loader in tests/core/test_xgboost_model.py (line 336) builds an in-memory dataset with create_sample_dataset's default in_memory=True. That dataset has no shuffle attribute, so the "no warning after a shuffled pass" half passes even without the fix, and only the second assertion fails on the old code. The reported disk-backed case is therefore not covered. The same test is duplicated in the new tests/core/test_gradient_boosted_trees.py (line 271) with the same gap. Non-blocking (low): the fallback still trusts a flag that get_dataloader sets only at creation for non-SampleDataset iterables, and the 335-line new base-class test file is not mentioned in the description. I could not run the test suite (read-only tools), so the claimed pass results are unverified. Overall risk is low for the code change and moderate for the test gap; manual review is yes because the change edits a base-class method.

_{Reviewed commit 0e053f0}

[github-actions]

[medium] Shuffle test uses an in-memory dataset, so it does not reproduce the reported bug

create_sample_dataset defaults to in_memory=True, which returns InMemorySampleDataset (pyhealth/datasets/sample_dataset.py:933). That class never calls the litdata base init and keeps its flag in _shuffle (sample_dataset.py:751), so loader.dataset.shuffle does not exist. The original getattr(loader.dataset, 'shuffle', False) check therefore always returns False here, and the first assertNoLogs passes on the unfixed code. The test only fails on the unfixed code because of the second assertLogs, which checks an unrelated behavior. The reported bug needs a disk-backed SampleDataset (in_memory=False), where shuffle is the litdata flag that stays stale. Fix: build the dataset with in_memory=False for the first half, or assert the flag is set after the shuffled pass, so the test fails on the old code for the reason the description gives.

[github-actions]

[medium] Same in-memory gap in the duplicated shuffle test

This test repeats the test in test_xgboost_model.py and uses the default in-memory dataset, so its no-warning half does not exercise the stale flag from the description. Use in_memory=False for the stale-flag case, or drop the duplicate and keep one test per dataset type.

[github-actions]

[low] Fallback still trusts a creation-time flag for non-SampleDataset iterables

The description says get_dataloader's loaders set the flag on each pass. That is only true for SampleDataset (_SeededDataLoader). For other iterables with set_shuffle, get_dataloader sets dataset.shuffle once at creation (pyhealth/datasets/utils.py:414-416), so a later get_dataloader(..., shuffle=True) call on the same dataset can still make an unshuffled loader warn falsely. This is a limitation, not a regression, so it does not block the merge.

[github-actions]

[low] Base-class test file is outside the PR description

The description covers only the XGBoost shuffle warning and its one test, but this PR adds a 335-line test file for GradientBoostedTreeModel with a scikit-learn LinearBackend. The file looks valid and scikit-learn is a core dependency (pyproject.toml:35), so it does not block the merge. Update the description to list it, or split it into a separate PR.

[github-actions]

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change edits _collect in GradientBoostedTreeModel, a base class that XGBoostModel builds on and that a planned LightGBM model will also build on. The edit only changes the shuffle-warning condition, but it modifies a base class directly, so the rules put it in the core.

_{Reviewed commit 0e053f0}

[solarsys]

Thanks for the second review. Everything is addressed in 5992de9 and 4e7550b.

Medium: the shuffle test now reproduces the reported bug, and CI runs it

test_shuffle_warning_follows_the_fit_loader now builds the dataset with in_memory=False, so the flag is litdata's shuffle. It also asserts that the shuffled pass leaves that flag set. It now lives in tests/core/test_gradient_boosted_trees.py, where it runs with the scikit-learn backend. Before, it was in the xgboost-only file, so CI skipped it. I checked it against _collect from the merge base, 7af8779. It fails on the no-warning assertion, with the false "train_data shuffles" warning the description describes. It passes with the fix, and it runs with xgboost unimportable.

Medium: no duplicated test

The XGBoost copy is dropped, since it checked nothing XGBoost-specific. The base-class file now has one shuffle test per dataset type:

• disk-backed: get_dataloader leaves a stale flag (the test above);
• in-memory and map-style rows: test_shuffle_warning_with_a_plain_dataloader. It fails on the PR's first version (a1f238d), which didn't read the in-memory _shuffle flag.

Low: other iterable datasets

A comment at the fallback now notes that get_dataloader sets their flag once, when a loader is created, so the flag follows the most recently created loader. For the pass being fitted, the warning is still accurate, because litdata reads the flag when iteration starts. I checked this with a run on a plain litdata.StreamingDataset. I created an unshuffled loader, then a shuffled one. The unshuffled loader's pass then came out shuffled, so a warning there would be true. I didn't change the behaviour.

Low: PR description

The description now lists tests/core/test_gradient_boosted_trees.py, why it's there, and its shuffle tests, plus the in-memory _shuffle fix. Every test of gradient_boosted_trees.py needed xgboost, which CI doesn't install, so CI covered 27.86% of the file. With xgboost unimportable, the file's coverage is 96.79% (271 of 280).

All 21 tests in test_xgboost_model.py and the 13 in test_gradient_boosted_trees.py pass. With xgboost blocked, the 13 still run, none skipped. Ruff is clean, and the PR rules check passes.

[jhnwu3]

lgtm