Make training shuffle a train-only option; migrate config and use samplers for ordering by mtauraso · Pull Request #923 · lincc-frameworks/hyrax

mtauraso · 2026-05-15T19:53:03Z

Motivation

Consolidate dataloader ordering control to a single train-only config key so training sample shuffling is explicit and other verbs remain deterministic.
Remove the legacy behavior of passing shuffle through PyTorch DataLoader kwargs to avoid sampler vs shuffle conflicts and to support explicit subset samplers when split_indices are used.
Provide a migration path so existing data_loader.shuffle entries are preserved and emit deprecation warnings.

Description

Add documentation specs/train_shuffle_config.md describing that train.shuffle is train-only and how it differs from setup_dataset(..., shuffle=...).
Add migration src/hyrax/config_migrations/migrations/002_move_data_loader_shuffle_to_train.py to rename data_loader.shuffle → train.shuffle.
Update default config src/hyrax/hyrax_default_config.toml to include train.shuffle = true and remove data_loader.shuffle.
Change dist_data_loader in src/hyrax/pytorch_ignite.py to accept a shuffle parameter, to warn and ignore legacy data_loader.shuffle, to build deterministic torch.Generator seeding, and to create explicit SubsetRandomSampler or SubsetSequentialSampler via a make_sampler helper instead of passing shuffle into DataLoader.
Wire train verb to read config['train']['shuffle'] and only apply shuffling to the train split; keep validation and test splits deterministic.
Remove per-verb code that forcibly mutated config['data_loader']['shuffle'] for infer and test.
Add and update unit tests to cover migration, default config placement, sampler selection, legacy-key warnings, and runtime behavior across train, infer, and legacy multi-split paths (tests under tests/hyrax/*).

Testing

Ran the test suite with pytest focused on config migrations, dataloader split/sampler behavior, infer/train/test integration, and config utils; all modified and added tests passed.
New/updated tests include tests/hyrax/test_config_migrations.py, tests/hyrax/test_config_utils.py, tests/hyrax/test_infer.py, tests/hyrax/test_split_fraction_integration.py, and tests/hyrax/test_train.py and they succeeded under the CI test run.
Verified that the migration preserves other data_loader keys (e.g. batch_size) while moving shuffle to train.shuffle and that legacy data_loader.shuffle emits a UserWarning/DeprecationWarning as appropriate.

Codex Task

Copilot

Pull request overview

This PR moves dataloader sample-order control from the global data_loader.shuffle option to a train-only train.shuffle setting, keeping non-training verbs deterministic while preserving a migration path for older configs.

Changes:

Adds train.shuffle to the default config and migrates deprecated data_loader.shuffle to it.
Updates dist_data_loader to use explicit subset samplers instead of forwarding shuffle to PyTorch.
Adjusts train/infer/test behavior and expands tests/documentation around shuffle semantics.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`src/hyrax/pytorch_ignite.py`	Adds shuffle-aware sampler selection and ignores legacy dataloader shuffle kwargs.
`src/hyrax/verbs/train.py`	Reads `train.shuffle` and applies it only to training dataloaders.
`src/hyrax/verbs/infer.py`	Removes mutation of legacy dataloader shuffle config for inference.
`src/hyrax/verbs/test.py`	Removes mutation of legacy dataloader shuffle config for testing.
`src/hyrax/hyrax_default_config.toml`	Moves the default shuffle option from `[data_loader]` to `[train]`.
`src/hyrax/config_migrations/migrations/002_move_data_loader_shuffle_to_train.py`	Adds config migration for `data_loader.shuffle` to `train.shuffle`.
`specs/train_shuffle_config.md`	Documents train-only shuffle behavior and migration guidance.
`tests/hyrax/test_config_migrations.py`	Adds migration and deprecated nested-key warning coverage.
`tests/hyrax/test_config_utils.py`	Verifies default shuffle config placement.
`tests/hyrax/test_infer.py`	Confirms inference ordering is unaffected by `train.shuffle`.
`tests/hyrax/test_split_fraction_integration.py`	Adds sampler-selection and legacy shuffle tests.
`tests/hyrax/test_train.py`	Updates training split-fraction test wording for sampler behavior.

codecov · 2026-05-15T19:59:24Z

Codecov Report

❌ Patch coverage is 92.85714% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.72%. Comparing base (f390e09) to head (e89df1a).

Files with missing lines	Patch %	Lines
src/hyrax/pytorch_ignite.py	94.44%	1 Missing ⚠️
src/hyrax/verbs/train.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #923      +/-   ##
==========================================
+ Coverage   67.64%   67.72%   +0.08%     
==========================================
  Files          69       70       +1     
  Lines        6802     6805       +3     
==========================================
+ Hits         4601     4609       +8     
+ Misses       2201     2196       -5

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

mtauraso · 2026-05-15T20:03:35Z

-        # Note that when sampler=None, a default sampler is used. The default config
-        # defines shuffle=False, which should prevent any shuffling of of the data.
-        # We expect that this will be the primary use case when running inference.
-        return idist.auto_dataloader(dataset, sampler=sampler, **loader_kwargs), indexes


On second thought, removing this comment seems okay.

drewoldag

This looks good. A few comments, but nothing overly pressing. Remove the warning, move or maybe remove one of the tests.

Otherwise good to go.

drewoldag · 2026-05-15T20:26:18Z

+
+```toml
+[train]
+shuffle = false


As discussed, the config migration will move the shuffle key and value, which will lead to correct reproducibility of previous runs, but the default value in hyrax_default_config.toml will become true.

drewoldag · 2026-05-15T20:30:26Z

@@ -158,6 +158,7 @@ def dist_data_loader(
    dataset: Dataset,
    config: dict,
    split: Union[str, list[str], bool] = False,


FWIW - I think we might be close to being able to remove this argument. It is only non-False in a couple of legacy call points in train.

drewoldag · 2026-05-15T20:34:28Z

+            "Hyrax controls dataloader ordering with explicit samplers; use config['train']['shuffle'] "
+            "to control training sample shuffling."
+        )
+        warnings.warn(msg, UserWarning, stacklevel=2)


Remove this warning

drewoldag · 2026-05-15T20:39:19Z

+        logger.warning(msg)
+        data_loader_kwargs.pop("shuffle")

    # TODO: Actually DataProvider.collate. Callsites and parameter signature above have not been updated.


Outside of the scope of this PR, but I think this comment can be removed.

drewoldag · 2026-05-15T20:49:27Z

+def test_default_shuffle_config_is_train_only():
+    """Default shuffle config lives under train and is absent from data_loader."""
+    import hyrax
+
+    config = hyrax.Hyrax().config
+
+    assert config["train"]["shuffle"] is True
+    assert "shuffle" not in config["data_loader"]


This test shouldn't be in here. Not sure that it's really necessary, but if we keep it, I think it should be in test_config_migrations.py.

This test suite is for the config_utils.py functionality, no for verifying the contents of the config dictionary.

I'm going with removal.

drewoldag · 2026-05-15T20:49:58Z

    """
    h, dataset = loopback_hyrax
-    h.config["data_loader"]["shuffle"] = shuffle
+    h.config["train"]["shuffle"] = train_shuffle


Yes, I like this very much!

aritraghsh09 · 2026-05-15T21:00:39Z

Agree with the overall design choice here. Deferring to drew for the actual review

Co-authored-by: Drew Oldag <47493171+drewoldag@users.noreply.github.com> Co-authored-by: Michael Tauraso <mtauraso@users.noreply.github.com>

mtauraso · 2026-05-15T22:03:40Z

@copilot Can you change the failing tests so they don't look for the warning that I specifically removed. kthx

Agent-Logs-Url: https://github.com/lincc-frameworks/hyrax/sessions/a471d1fb-1c54-43e9-9671-5f2df78b7c25 Co-authored-by: mtauraso <31012+mtauraso@users.noreply.github.com>

Copilot · 2026-05-15T22:20:59Z

@copilot Can you change the failing tests so they don't look for the warning that I specifically removed. kthx

Updated in 4683548: I removed the two pytest.warns(UserWarning, ...) assertions that were still expecting the removed warning, and validated with targeted tests plus full pytest -m "not slow" and pre-commit.

Warn when ignoring data_loader shuffle

d41ddd4

Copilot AI review requested due to automatic review settings May 15, 2026 19:53

mtauraso added the codex label May 15, 2026 — with ChatGPT Codex Connector

Copilot started reviewing on behalf of mtauraso May 15, 2026 19:53 View session

Copilot AI reviewed May 15, 2026

View reviewed changes

Comment thread src/hyrax/pytorch_ignite.py Outdated

Comment thread src/hyrax/pytorch_ignite.py

mtauraso commented May 15, 2026

View reviewed changes

Fix formatting and config access

3f50c49

mtauraso self-assigned this May 15, 2026

mtauraso requested a review from drewoldag May 15, 2026 20:07

mtauraso and others added 2 commits May 15, 2026 13:10

Fixup edge case

ba114aa

Merge branch 'main' into codex/implement-shuffle-config-for-train-verb

6c8f148

mtauraso requested a review from aritraghsh09 May 15, 2026 20:15

drewoldag approved these changes May 15, 2026

View reviewed changes

mtauraso commented May 15, 2026

View reviewed changes

Comment thread src/hyrax/pytorch_ignite.py Outdated

mtauraso commented May 15, 2026

View reviewed changes

Comment thread tests/hyrax/test_config_utils.py Outdated

Apply suggestions from code review

16293a8

Co-authored-by: Drew Oldag <47493171+drewoldag@users.noreply.github.com> Co-authored-by: Michael Tauraso <mtauraso@users.noreply.github.com>

Copilot started work on behalf of mtauraso May 15, 2026 22:03 View session

test: stop expecting removed legacy shuffle warning

4683548

Agent-Logs-Url: https://github.com/lincc-frameworks/hyrax/sessions/a471d1fb-1c54-43e9-9671-5f2df78b7c25 Co-authored-by: mtauraso <31012+mtauraso@users.noreply.github.com>

Copilot finished work on behalf of mtauraso May 15, 2026 22:21

Merge branch 'main' into codex/implement-shuffle-config-for-train-verb

e89df1a

mtauraso merged commit 7d52a5f into main May 16, 2026
8 checks passed

mtauraso deleted the codex/implement-shuffle-config-for-train-verb branch May 16, 2026 00:25

Uh oh!

Conversation

mtauraso commented May 15, 2026

Motivation

Description

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

drewoldag left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

aritraghsh09 commented May 15, 2026

Uh oh!

Uh oh!

mtauraso commented May 15, 2026

Uh oh!

Copilot AI commented May 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codecov Bot commented May 15, 2026 •

edited

Loading