From f334f345487e71cd69d5a44b22bc4a89dda2e3f3 Mon Sep 17 00:00:00 2001 From: solarsys Date: Sun, 18 Oct 2026 13:58:58 -0700 Subject: [PATCH 1/4] Base XGBoost's shuffle warning on the fit loader, not the dataset flag get_dataloader's loaders set the dataset's shuffle flag at the start of each pass, so after another model trains with a shuffled loader on the same dataset, the flag stays True. XGBoostModel.fit then warned that an unshuffled loader shuffles. Read the loader's own setting when it has one. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PhbzTePoLMBYwjMLG8R9ek --- pyhealth/models/gradient_boosted_trees.py | 11 ++++++++--- tests/core/test_xgboost_model.py | 12 ++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/pyhealth/models/gradient_boosted_trees.py b/pyhealth/models/gradient_boosted_trees.py index 75de65234..a483ac5af 100644 --- a/pyhealth/models/gradient_boosted_trees.py +++ b/pyhealth/models/gradient_boosted_trees.py @@ -346,9 +346,14 @@ def _collect(self, data, name: str) -> tuple[np.ndarray, np.ndarray]: f"{name} has drop_last=True; tree models must see every " f"sample. Use a loader with drop_last=False." ) - if isinstance(loader.sampler, RandomSampler) or getattr( - loader.dataset, "shuffle", False - ): + # get_dataloader's loaders apply their own shuffle setting at each + # pass, so the dataset's flag may be left over from another loader. + shuffles = getattr(loader, "shuffle_samples", None) + if shuffles is None: + shuffles = isinstance(loader.sampler, RandomSampler) or getattr( + loader.dataset, "shuffle", False + ) + if shuffles: logger.warning( "%s shuffles; row order changes row subsampling, so fits " "are only reproducible with an unshuffled loader " diff --git a/tests/core/test_xgboost_model.py b/tests/core/test_xgboost_model.py index a714b76bb..102ef97a9 100644 --- a/tests/core/test_xgboost_model.py +++ b/tests/core/test_xgboost_model.py @@ -332,6 +332,18 @@ def test_fit_rejects_drop_last(self): with self.assertRaisesRegex(ValueError, "drop_last"): XGBoostModel(ds).fit(loader) + def test_shuffle_warning_follows_the_fit_loader(self): + ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}) + # A shuffled loader on the same dataset, e.g. another model's training. + for _ in get_dataloader(ds, batch_size=32, shuffle=True): + pass + logger = "pyhealth.models.gradient_boosted_trees" + with self.assertNoLogs(logger, level="WARNING"): + XGBoostModel(ds, n_estimators=2).fit(get_dataloader(ds, batch_size=32)) + with self.assertLogs(logger, level="WARNING") as logs: + XGBoostModel(ds, n_estimators=2).fit(get_dataloader(ds, 32, shuffle=True)) + self.assertIn("shuffles", logs.output[0]) + def test_forward_before_fit_raises(self): ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}) with self.assertRaisesRegex(RuntimeError, "model.fit"): From 0e053f08b14ff7b14b056846cc5678be59641370 Mon Sep 17 00:00:00 2001 From: solarsys Date: Sat, 24 Oct 2026 22:07:44 +0100 Subject: [PATCH 2/4] Test GradientBoostedTreeModel without xgboost; warn for shuffled in-memory datasets 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). The base class does not depend on the tree library, so the new tests/core/test_gradient_boosted_trees.py fills its hooks with scikit-learn linear models and covers the feature layout, every label mode, fitting checks, scale_pos_weight, prediction, attributions and checkpoints. Without xgboost the file is now at 96.79%. The tests include a plain torch DataLoader, which has no shuffle_samples attribute. They showed that the fallback missed shuffled in-memory datasets: InMemorySampleDataset, the default of create_sample_dataset, keeps its flag in _shuffle rather than litdata's shuffle, so a plain loader over one never warned. The check now reads both. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015ntMLEmFsQEQJKTmPGGw4T --- pyhealth/models/gradient_boosted_trees.py | 7 +- tests/core/test_gradient_boosted_trees.py | 335 ++++++++++++++++++++++ 2 files changed, 340 insertions(+), 2 deletions(-) create mode 100644 tests/core/test_gradient_boosted_trees.py diff --git a/pyhealth/models/gradient_boosted_trees.py b/pyhealth/models/gradient_boosted_trees.py index a483ac5af..d66428aaf 100644 --- a/pyhealth/models/gradient_boosted_trees.py +++ b/pyhealth/models/gradient_boosted_trees.py @@ -350,8 +350,11 @@ def _collect(self, data, name: str) -> tuple[np.ndarray, np.ndarray]: # pass, so the dataset's flag may be left over from another loader. shuffles = getattr(loader, "shuffle_samples", None) if shuffles is None: - shuffles = isinstance(loader.sampler, RandomSampler) or getattr( - loader.dataset, "shuffle", False + # Iterating a SampleDataset shuffles when its flag is set: + # litdata's ``shuffle`` on disk, ``_shuffle`` in memory. + shuffles = isinstance(loader.sampler, RandomSampler) or bool( + getattr(loader.dataset, "shuffle", False) + or getattr(loader.dataset, "_shuffle", False) ) if shuffles: logger.warning( diff --git a/tests/core/test_gradient_boosted_trees.py b/tests/core/test_gradient_boosted_trees.py new file mode 100644 index 000000000..32e8417b7 --- /dev/null +++ b/tests/core/test_gradient_boosted_trees.py @@ -0,0 +1,335 @@ +"""GradientBoostedTreeModel with a scikit-learn backend, so it runs without xgboost. + +The base class does not depend on the tree library: subclasses only build, +fit, predict with and serialize one estimator. ``LinearBackend`` fills those +hooks with linear models, which keeps the feature layout, fitting, prediction, +attribution and checkpoint code under test in environments without xgboost +(``tests/core/test_xgboost_model.py`` covers XGBoostModel itself). +""" + +import os +import pickle +import tempfile +import unittest + +import numpy as np +import torch +from sklearn.linear_model import LinearRegression, LogisticRegression + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import GradientBoostedTreeModel +from pyhealth.trainer import Trainer + +LOGGER = "pyhealth.models.gradient_boosted_trees" +MIXED = {"labs": "tensor", "flags": "multi_hot", "codes": "sequence", "visits": "nested_sequence"} + + +class LinearBackend(GradientBoostedTreeModel): + """Linear models in place of trees. + + A linear model's contributions (coefficient times value, with the + intercept as the bias) add up to its margin exactly, as TreeSHAP values + add up to a tree ensemble's. + """ + + def __init__(self, dataset, **kwargs): + super().__init__(dataset, **kwargs) + self.eval_sets = [] + + def _make_estimator(self, scale_pos_weight): + if self.mode == "regression": + return LinearRegression() + weight = None if scale_pos_weight is None else {0: 1.0, 1: scale_pos_weight} + return LogisticRegression(class_weight=weight, max_iter=2000) + + def _fit_estimator(self, estimator, X, y, eval_set): + self.eval_sets.append(eval_set) + estimator.fit(X, y) + + def _predict_margin(self, estimator, X): + if self.mode == "regression": + return estimator.predict(X) + return estimator.decision_function(X) + + def _predict_value(self, estimator, X): + if self.mode == "regression": + return estimator.predict(X) + if self.mode == "multilabel": # one estimator per label + return estimator.predict_proba(X)[:, 1] + return estimator.predict_proba(X) + + def _contributions(self, estimator, X): + coef = np.atleast_2d(estimator.coef_) # (K, d) + bias = np.broadcast_to( + np.atleast_1d(estimator.intercept_)[None, :, None], (len(X), len(coef), 1) + ) + values = np.concatenate([X[:, None, :] * coef[None], bias], axis=2) + return values[:, 0] if len(coef) == 1 else values + + def _to_bytes(self, estimator): + return pickle.dumps(estimator) + + def _from_bytes(self, raw): + return pickle.loads(raw) # bytes this test wrote + + +def _samples(n=90, label="binary"): + labels = { + "binary": lambda i: int(i % 5 > 2), + "multiclass": lambda i: i % 3, + "multilabel": lambda i: [["a"], ["b"], ["a", "c"], ["b", "c"]][i % 4], + "regression": lambda i: float(i % 7) / 2, + }[label] + return [ + { + "patient_id": f"p{i}", + "labs": [float(i % 7), float((i * 3) % 5), float(i % 2)], + "flags": ["x", "y", "z"][: i % 4], + "codes": [f"c{(i + j) % 11}" for j in range(1 + i % 6)], + "visits": [[f"d{(i + v + k) % 5}" for k in range(1 + v % 3)] for v in range(1 + i % 3)], + "label": labels(i), + } + for i in range(n) + ] + + +def _plain_loader(dataset, data=None, **kwargs): + """A torch DataLoader without get_dataloader's shuffle_samples attribute.""" + collate = get_dataloader(dataset, batch_size=1).collate_fn + return torch.utils.data.DataLoader( + dataset if data is None else data, batch_size=32, collate_fn=collate, **kwargs + ) + + +class TestEveryMode(unittest.TestCase): + def test_fit_predict_explain_and_checkpoint(self): + width = {"binary": 1, "multiclass": 3, "multilabel": 3, "regression": 1} + for mode in ("binary", "multiclass", "multilabel", "regression"): + with self.subTest(mode=mode), tempfile.TemporaryDirectory() as tmp: + ds = create_sample_dataset(_samples(label=mode), MIXED, {"label": mode}) + loader = get_dataloader(ds, batch_size=32) + model = LinearBackend(ds, bag_of_codes=True).fit(loader, loader) + self.assertTrue(model.is_fitted) + self.assertTrue(all(e is not None for e in model.eval_sets)) + + trainer = Trainer(model=model, enable_logging=False) + _, y_prob, loss = trainer.inference(loader) + self.assertEqual(y_prob.shape, (len(ds), width[mode])) + self.assertTrue(np.isfinite(loss)) + with self.assertRaisesRegex(TypeError, "model.fit"): + trainer.train(loader, epochs=1) + + batch = next(iter(loader)) + inputs = {k: v for k, v in batch.items() if k != "label"} + self.assertEqual(set(model(**inputs)), {"logit", "y_prob"}) + + out = model.explain(**batch) + total = sum(a.sum(1) for a in out["attributions"].values()) + out["bias"] + logit = out["logit"] if width[mode] > 1 else out["logit"][:, 0] + torch.testing.assert_close(total, logit, atol=1e-4, rtol=1e-4) + self.assertEqual( + {k: len(v) for k, v in out["feature_names"].items()}, + {k: v.shape[1] for k, v in out["attributions"].items()}, + ) + ranking = model.mean_abs_shap(ds) + self.assertEqual(sorted(ranking), sorted(model.feature_names)) + self.assertEqual(list(ranking.values()), sorted(ranking.values(), reverse=True)) + + path = os.path.join(tmp, "model.ckpt") + trainer.save_ckpt(path) + fresh = LinearBackend(ds, bag_of_codes=True) + Trainer(model=fresh, enable_logging=False).load_ckpt(path) + self.assertEqual(fresh.feature_layout, model.feature_layout) + a, b = model(**batch), fresh(**batch) + self.assertTrue(torch.equal(a["y_prob"], b["y_prob"])) + self.assertTrue(torch.equal(a["logit"], b["logit"])) + + +class TestInputs(unittest.TestCase): + def test_fixed_width_processors(self): + samples = [ + { + "patient_id": f"p{i}", + "vec": [float(i), float(i % 3)], + "mh": ["u", "v"][: i % 3], + "nmh": [["a", "b"][: 1 + v % 2] for v in range(1 + i % 3)], + "flag": i % 2, + "score": float(i) / 10, + "label": int(i % 4 == 0), + } + for i in range(40) + ] + schema = {"vec": "tensor", "mh": "multi_hot", "nmh": "nested_multihot", + "flag": "binary", "score": "regression"} + ds = create_sample_dataset(samples, schema, {"label": "binary"}) + model = LinearBackend(ds, feature_names={"vec": ["day", "phase"]}).fit(ds) + self.assertEqual( + model.feature_names, + ["day", "phase", "mh=u", "mh=v", "nmh=a", "nmh=b", "flag", "score"], + ) + X = model.build_feature_matrix(**next(iter(get_dataloader(ds, batch_size=40)))) + # nested_multihot: number of visits containing each code + np.testing.assert_array_equal(X[:, 4:6], [[1 + i % 3, (1 + i % 3) // 2] for i in range(40)]) + ranges = {f["key"]: (f["start"], f["stop"]) for f in model.feature_layout} + self.assertEqual(ranges, {"vec": (0, 2), "mh": (2, 4), "nmh": (4, 6), "flag": (6, 7), + "score": (7, 8)}) + + with self.assertRaisesRegex(ValueError, r"feature_names\['vec'\] has 1 names"): + LinearBackend(ds, feature_names={"vec": ["day"]}).fit(ds) + batch = next(iter(get_dataloader(ds, batch_size=4))) + with self.assertRaisesRegex(ValueError, "'nmh' has a vocabulary of 1 entries"): + model.build_feature_matrix(**dict(batch, nmh=torch.zeros(4, 2, 1))) + + def test_bag_of_codes_width_is_fixed(self): + ds = create_sample_dataset(_samples(), MIXED, {"label": "binary"}) + model = LinearBackend(ds, bag_of_codes=True).fit(ds) + # batches of 1 and 32 pad to different lengths; the width must not change + X1 = np.concatenate([model.build_feature_matrix(**b) for b in get_dataloader(ds, batch_size=1)]) + X32 = np.concatenate([model.build_feature_matrix(**b) for b in get_dataloader(ds, batch_size=32)]) + np.testing.assert_array_equal(X1, X32) + self.assertEqual(X1.shape[1], len(model.feature_names)) + + def test_inputs_that_cannot_become_columns_raise(self): + ds = create_sample_dataset(_samples(), MIXED, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "'codes'.*bag_of_codes=True"): + LinearBackend(ds) + + floats = [{"patient_id": f"p{i}", "v": [[float(i)]], "label": i % 2} for i in range(6)] + ds = create_sample_dataset(floats, {"v": "nested_sequence_floats"}, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "'v' uses NestedFloatsProcessor"): + LinearBackend(ds) + + varying = [{"patient_id": f"p{i}", "v": [1.0] * (1 + i // 4), "label": i % 2} + for i in range(8)] + ds = create_sample_dataset(varying, {"v": "tensor"}, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "'v' has per-sample shape"): + LinearBackend(ds).fit(get_dataloader(ds, batch_size=4)) + + def test_batches_that_do_not_match_the_layout_raise(self): + ds = create_sample_dataset(_samples(), MIXED, {"label": "binary"}) + model = LinearBackend(ds, bag_of_codes=True).fit(ds) + batch = next(iter(get_dataloader(ds, batch_size=4))) + cases = { + "missing input field 'labs'": (KeyError, {k: v for k, v in batch.items() if k != "labs"}), + "'labs' is not a tensor": (TypeError, dict(batch, labs=[[1.0, 2.0, 3.0]] * 4)), + r"'flags' has shape \(2,\)": (ValueError, dict(batch, flags=torch.zeros(4, 2))), + } + for message, (error, bad) in cases.items(): + with self.subTest(message), self.assertRaisesRegex(error, message): + model.build_feature_matrix(**bad) + + +class TestFitting(unittest.TestCase): + def test_constructor_checks(self): + ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "'balanced', 'auto' or None"): + LinearBackend(ds, scale_pos_weight="heavy") + multiclass = create_sample_dataset(_samples(label="multiclass"), {"labs": "tensor"}, + {"label": "multiclass"}) + with self.assertRaisesRegex(ValueError, "binary or multilabel"): + LinearBackend(multiclass, scale_pos_weight=2.0) + two = [dict(s, other=s["label"]) for s in _samples()] + ds = create_sample_dataset(two, {"labs": "tensor"}, {"label": "binary", "other": "binary"}) + with self.assertRaisesRegex(ValueError, "exactly one label"): + LinearBackend(ds) + + def test_scale_pos_weight(self): + ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) + y = np.array([s["label"] for s in _samples()]) + model = LinearBackend(ds, scale_pos_weight="balanced").fit(ds) + np.testing.assert_allclose(model.fitted_scale_pos_weight, + [(y == 0).sum() / (y == 1).sum()]) + self.assertEqual(LinearBackend(ds, scale_pos_weight=3).fit(ds).fitted_scale_pos_weight, [3.0]) + self.assertIsNone(LinearBackend(ds).fit(ds).fitted_scale_pos_weight) + + ml = create_sample_dataset(_samples(label="multilabel"), {"labs": "tensor"}, + {"label": "multilabel"}) + model = LinearBackend(ml, scale_pos_weight=[1.0, 2.0, 3.0]).fit(ml) + self.assertEqual(len(model.estimators_), 3) + self.assertEqual(model.fitted_scale_pos_weight, [1.0, 2.0, 3.0]) + with self.assertRaisesRegex(ValueError, "2 values for 3 labels"): + LinearBackend(ml, scale_pos_weight=[1.0, 2.0]).fit(ml) + with self.assertLogs(LOGGER, "WARNING") as logs: + weights = LinearBackend(ml, scale_pos_weight="auto")._resolve_pos_weight( + np.array([[1, 0, 1], [0, 0, 1], [0, 0, 0]]) + ) + self.assertEqual(weights, [2.0, 1.0, 0.5]) + self.assertIn("no positive training examples", logs.output[0]) + + def test_training_sets_that_cannot_be_fit_raise(self): + ds = create_sample_dataset(_samples(label="multiclass"), {"labs": "tensor"}, + {"label": "multiclass"}) + two_classes = ds.subset([i for i in range(len(ds)) if i % 3 != 2]) + with self.assertRaisesRegex(ValueError, r"Classes \[2\] have no training samples"): + LinearBackend(ds).fit(two_classes) + with self.assertRaisesRegex(ValueError, "train_data is empty"): + LinearBackend(ds).fit(torch.utils.data.DataLoader([], batch_size=4)) + with self.assertRaisesRegex(ValueError, "drop_last"): + LinearBackend(ds).fit(_plain_loader(ds, drop_last=True)) + + def test_shuffle_warning_follows_the_fit_loader(self): + ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) + # A shuffled loader on the same dataset, e.g. another model's training, + # leaves the dataset's flag set; get_dataloader's loaders reset it. + for _ in get_dataloader(ds, batch_size=32, shuffle=True): + pass + with self.assertNoLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(get_dataloader(ds, batch_size=32)) + with self.assertLogs(LOGGER, "WARNING") as logs: + LinearBackend(ds).fit(get_dataloader(ds, 32, shuffle=True)) + self.assertIn("shuffles", logs.output[0]) + + def test_shuffle_warning_with_a_plain_dataloader(self): + # Without shuffle_samples, the sampler and the dataset's flag decide. + ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) + rows = [ds[i] for i in range(len(ds))] # map-style, so it can be shuffled + with self.assertNoLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(_plain_loader(ds, rows)) + with self.assertLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(_plain_loader(ds, rows, shuffle=True)) + # Iterating a SampleDataset shuffles when its flag is set; this one is + # in memory, which keeps the flag in _shuffle. + ds.set_shuffle(False) + with self.assertNoLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(_plain_loader(ds)) + ds.set_shuffle(True) + with self.assertLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(_plain_loader(ds)) + ds.set_shuffle(False) + + +class TestCheckpoints(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.ds = create_sample_dataset(_samples(), MIXED, {"label": "binary"}) + cls.model = LinearBackend(cls.ds, bag_of_codes=True).fit(cls.ds) + + def test_forward_before_fit_raises(self): + batch = next(iter(get_dataloader(self.ds, batch_size=4))) + with self.assertRaisesRegex(RuntimeError, "model.fit"): + LinearBackend(self.ds, bag_of_codes=True)(**batch) + + def test_an_unfitted_checkpoint_loads_unfitted(self): + fresh = LinearBackend(self.ds, bag_of_codes=True) + fresh.load_state_dict(LinearBackend(self.ds, bag_of_codes=True).state_dict()) + self.assertFalse(fresh.is_fitted) + + def test_mismatched_checkpoints_raise(self): + state = self.model.get_extra_state() + fresh = LinearBackend(self.ds, bag_of_codes=True) + with self.assertRaisesRegex(ValueError, "state version 99"): + fresh.set_extra_state(dict(state, version=99)) + with self.assertRaisesRegex(ValueError, "fit for mode 'regression'"): + fresh.set_extra_state(dict(state, mode="regression")) + + other = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "feature layout"): + LinearBackend(other).set_extra_state(state) + fewer_codes = [dict(s, codes=["c1"]) for s in _samples()] + smaller = create_sample_dataset(fewer_codes, MIXED, {"label": "binary"}) + with self.assertRaisesRegex(ValueError, "'codes' had"): + LinearBackend(smaller, bag_of_codes=True).set_extra_state(state) + + +if __name__ == "__main__": + unittest.main() From 5992de92a14be20decdd2cbc757412b5ac3a5d80 Mon Sep 17 00:00:00 2001 From: solarsys Date: Sat, 24 Oct 2026 23:26:08 +0100 Subject: [PATCH 3/4] Test the stale shuffle flag on a disk-backed dataset; note the fallback's limit test_shuffle_warning_follows_the_fit_loader built an in-memory dataset, which has no litdata shuffle flag, so its no-warning half passed on the unfixed code too. It now uses a disk-backed dataset, checks that the shuffled pass leaves the flag set, and fails without the fix. The copy of that test in test_gradient_boosted_trees.py is dropped. That file keeps the plain-DataLoader test, which covers in-memory datasets and map-style rows, so each dataset type is tested once. A comment at the dataset-flag fallback notes that get_dataloader sets the flag of other iterable datasets once, when a loader is created, so it follows the most recently created loader. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015ntMLEmFsQEQJKTmPGGw4T --- pyhealth/models/gradient_boosted_trees.py | 5 ++++- tests/core/test_gradient_boosted_trees.py | 17 ++++------------- tests/core/test_xgboost_model.py | 8 ++++++-- 3 files changed, 14 insertions(+), 16 deletions(-) diff --git a/pyhealth/models/gradient_boosted_trees.py b/pyhealth/models/gradient_boosted_trees.py index d66428aaf..211e6dfeb 100644 --- a/pyhealth/models/gradient_boosted_trees.py +++ b/pyhealth/models/gradient_boosted_trees.py @@ -351,7 +351,10 @@ def _collect(self, data, name: str) -> tuple[np.ndarray, np.ndarray]: shuffles = getattr(loader, "shuffle_samples", None) if shuffles is None: # Iterating a SampleDataset shuffles when its flag is set: - # litdata's ``shuffle`` on disk, ``_shuffle`` in memory. + # litdata's ``shuffle`` on disk, ``_shuffle`` in memory. For + # other iterable datasets, get_dataloader sets the flag once, + # when a loader is created, so the flag (and the order of this + # pass) follows the most recently created loader, not this one. shuffles = isinstance(loader.sampler, RandomSampler) or bool( getattr(loader.dataset, "shuffle", False) or getattr(loader.dataset, "_shuffle", False) diff --git a/tests/core/test_gradient_boosted_trees.py b/tests/core/test_gradient_boosted_trees.py index 32e8417b7..c4c5fdfe0 100644 --- a/tests/core/test_gradient_boosted_trees.py +++ b/tests/core/test_gradient_boosted_trees.py @@ -267,20 +267,10 @@ def test_training_sets_that_cannot_be_fit_raise(self): with self.assertRaisesRegex(ValueError, "drop_last"): LinearBackend(ds).fit(_plain_loader(ds, drop_last=True)) - def test_shuffle_warning_follows_the_fit_loader(self): - ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) - # A shuffled loader on the same dataset, e.g. another model's training, - # leaves the dataset's flag set; get_dataloader's loaders reset it. - for _ in get_dataloader(ds, batch_size=32, shuffle=True): - pass - with self.assertNoLogs(LOGGER, "WARNING"): - LinearBackend(ds).fit(get_dataloader(ds, batch_size=32)) - with self.assertLogs(LOGGER, "WARNING") as logs: - LinearBackend(ds).fit(get_dataloader(ds, 32, shuffle=True)) - self.assertIn("shuffles", logs.output[0]) - def test_shuffle_warning_with_a_plain_dataloader(self): # Without shuffle_samples, the sampler and the dataset's flag decide. + # (get_dataloader's loaders with a stale flag on a disk-backed dataset + # are tested in test_xgboost_model.py.) ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) rows = [ds[i] for i in range(len(ds))] # map-style, so it can be shuffled with self.assertNoLogs(LOGGER, "WARNING"): @@ -293,8 +283,9 @@ def test_shuffle_warning_with_a_plain_dataloader(self): with self.assertNoLogs(LOGGER, "WARNING"): LinearBackend(ds).fit(_plain_loader(ds)) ds.set_shuffle(True) - with self.assertLogs(LOGGER, "WARNING"): + with self.assertLogs(LOGGER, "WARNING") as logs: LinearBackend(ds).fit(_plain_loader(ds)) + self.assertIn("shuffles", logs.output[0]) ds.set_shuffle(False) diff --git a/tests/core/test_xgboost_model.py b/tests/core/test_xgboost_model.py index 102ef97a9..a1dd610dd 100644 --- a/tests/core/test_xgboost_model.py +++ b/tests/core/test_xgboost_model.py @@ -333,10 +333,14 @@ def test_fit_rejects_drop_last(self): XGBoostModel(ds).fit(loader) def test_shuffle_warning_follows_the_fit_loader(self): - ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}) - # A shuffled loader on the same dataset, e.g. another model's training. + # Disk-backed, so the dataset's flag is litdata's ``shuffle``. + ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}, + in_memory=False) + # A shuffled loader on the same dataset, e.g. another model's training, + # leaves the flag set; get_dataloader's loaders set it again each pass. for _ in get_dataloader(ds, batch_size=32, shuffle=True): pass + self.assertTrue(ds.shuffle) logger = "pyhealth.models.gradient_boosted_trees" with self.assertNoLogs(logger, level="WARNING"): XGBoostModel(ds, n_estimators=2).fit(get_dataloader(ds, batch_size=32)) From 4e7550b3f3a89e2339da999afae66fa4f086ca9c Mon Sep 17 00:00:00 2001 From: solarsys Date: Sat, 24 Oct 2026 23:30:33 +0100 Subject: [PATCH 4/4] Move the stale-flag shuffle test to the base-class tests, so CI runs it test_shuffle_warning_follows_the_fit_loader needed xgboost, which CI does not install, so CI skipped the PR's main regression test. The check is in GradientBoostedTreeModel, so the disk-backed case now runs with the scikit-learn backend in test_gradient_boosted_trees.py, next to the in-memory and plain-loader test. The XGBoost copy is dropped: it checked nothing XGBoost-specific. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015ntMLEmFsQEQJKTmPGGw4T --- tests/core/test_gradient_boosted_trees.py | 17 +++++++++++++++-- tests/core/test_xgboost_model.py | 16 ---------------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/tests/core/test_gradient_boosted_trees.py b/tests/core/test_gradient_boosted_trees.py index c4c5fdfe0..5aec47758 100644 --- a/tests/core/test_gradient_boosted_trees.py +++ b/tests/core/test_gradient_boosted_trees.py @@ -267,10 +267,23 @@ def test_training_sets_that_cannot_be_fit_raise(self): with self.assertRaisesRegex(ValueError, "drop_last"): LinearBackend(ds).fit(_plain_loader(ds, drop_last=True)) + def test_shuffle_warning_follows_the_fit_loader(self): + # Disk-backed, so the dataset's flag is litdata's ``shuffle``. + ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}, + in_memory=False) + # A shuffled loader on the same dataset, e.g. another model's training, + # leaves the flag set; get_dataloader's loaders set it again each pass. + for _ in get_dataloader(ds, batch_size=32, shuffle=True): + pass + self.assertTrue(ds.shuffle) + with self.assertNoLogs(LOGGER, "WARNING"): + LinearBackend(ds).fit(get_dataloader(ds, batch_size=32)) + with self.assertLogs(LOGGER, "WARNING") as logs: + LinearBackend(ds).fit(get_dataloader(ds, 32, shuffle=True)) + self.assertIn("shuffles", logs.output[0]) + def test_shuffle_warning_with_a_plain_dataloader(self): # Without shuffle_samples, the sampler and the dataset's flag decide. - # (get_dataloader's loaders with a stale flag on a disk-backed dataset - # are tested in test_xgboost_model.py.) ds = create_sample_dataset(_samples(), {"labs": "tensor"}, {"label": "binary"}) rows = [ds[i] for i in range(len(ds))] # map-style, so it can be shuffled with self.assertNoLogs(LOGGER, "WARNING"): diff --git a/tests/core/test_xgboost_model.py b/tests/core/test_xgboost_model.py index a1dd610dd..a714b76bb 100644 --- a/tests/core/test_xgboost_model.py +++ b/tests/core/test_xgboost_model.py @@ -332,22 +332,6 @@ def test_fit_rejects_drop_last(self): with self.assertRaisesRegex(ValueError, "drop_last"): XGBoostModel(ds).fit(loader) - def test_shuffle_warning_follows_the_fit_loader(self): - # Disk-backed, so the dataset's flag is litdata's ``shuffle``. - ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}, - in_memory=False) - # A shuffled loader on the same dataset, e.g. another model's training, - # leaves the flag set; get_dataloader's loaders set it again each pass. - for _ in get_dataloader(ds, batch_size=32, shuffle=True): - pass - self.assertTrue(ds.shuffle) - logger = "pyhealth.models.gradient_boosted_trees" - with self.assertNoLogs(logger, level="WARNING"): - XGBoostModel(ds, n_estimators=2).fit(get_dataloader(ds, batch_size=32)) - with self.assertLogs(logger, level="WARNING") as logs: - XGBoostModel(ds, n_estimators=2).fit(get_dataloader(ds, 32, shuffle=True)) - self.assertIn("shuffles", logs.output[0]) - def test_forward_before_fit_raises(self): ds = create_sample_dataset(_mixed_samples(), {"labs": "tensor"}, {"label": "binary"}) with self.assertRaisesRegex(RuntimeError, "model.fit"):