feat(rollout): add --rollout-shuffle to permute prompts per epoch

docs/user-guide/cli-reference.md

@@ -194,6 +194,7 @@ See [Dtype Control](../advanced/dtype-control.md).
 | `--update-weight-target-module` | str | `transformer` | Modules to train and sync. Wan2.2: `transformer,transformer_2`. |
 | `--update-weight-buffer-size` | int | 512 MiB | Weight-sync chunk size in bytes. |
 | `--rollout-seed` | int | `42` | |
+| `--rollout-shuffle` | flag | off | Shuffle the prompt order once per epoch (seeded by `--rollout-seed`). Off reads the jsonl in file order. |
 | `--over-sampling-batch-size` | int | – | Must equal `--rollout-batch-size` today. |
 | `--sglang-server-concurrency` | int | `512` | Per-engine in-flight request cap. |
 | `--use-distributed-post` | flag | off | Rollout HTTP POSTs go through per-node Ray actors instead of the local client (`MILES_HTTP_POST_ACTORS_PER_NODE` sets the count). |

docs/user-guide/rewards.md

@@ -190,7 +190,8 @@ Override that path with `--custom-reward-post-process-path` — see
 
 ### JSONL format
 
-Training prompts are loaded from `.jsonl` files via `miles/utils/diffusion_data.py`:
+Training prompts are loaded from `.jsonl` files via `miles/utils/diffusion_data.py` and read in
+file order; `--rollout-shuffle` permutes them once per epoch (seeded by `--rollout-seed`):
 
 ```json
 {"input": "A photo of a cat wearing sunglasses"}

miles/rollout/data_source.py

@@ -58,6 +58,8 @@ def __init__(self, args):
                 metadata_key=args.metadata_key,
                 seed=args.rollout_seed,
             )
+            if args.rollout_shuffle:
+                self.dataset.shuffle(self.epoch_id)
         else:
             self.dataset = None
 
@@ -71,6 +73,8 @@ def get_samples(self, num_samples):
                 prompt_samples = self.dataset.samples[self.sample_offset :]
                 num_samples -= len(prompt_samples)
                 self.epoch_id += 1
+                if self.args.rollout_shuffle:
+                    self.dataset.shuffle(self.epoch_id)
                 prompt_samples += self.dataset.samples[:num_samples]
                 self.sample_offset = num_samples
         else:
@@ -124,6 +128,9 @@ def load(self, rollout_id=None):
         self.sample_group_index = state_dict.get("sample_group_index", 0)
         self.sample_index = state_dict.get("sample_index", 0)
         self.metadata = state_dict.get("metadata", {})
+        # the permutation is a function of (seed, epoch), so restoring epoch_id restores the order
+        if self.args.rollout_shuffle:
+            self.dataset.shuffle(self.epoch_id)
 
 
 class RolloutDataSourceWithBuffer(RolloutDataSource):

miles/utils/arguments.py

@@ -421,6 +421,12 @@ def add_rollout_arguments(parser):
                     "This is used to shuffle the prompts and also for the random sampling of the prompts."
                 ),
             )
+            parser.add_argument(
+                "--rollout-shuffle",
+                action="store_true",
+                default=False,
+                help="Shuffle the prompt order once per epoch, seeded by --rollout-seed; off reads the jsonl in file order.",
+            )
 
             # sampling
             parser.add_argument(

tests/fast/rollout/test_rollout_shuffle.py

@@ -0,0 +1,91 @@
+"""--rollout-shuffle permutes the prompt order per epoch, deterministically from --rollout-seed.
+
+Mental model (5 prompts p0..p4, 3 per rollout):
+
+    off : p0 p1 p2 | p3 p4 p0 | p1 p2 p3 ...          file order, wraps around
+    on  : perm(seed, epoch 0) | ... wrap -> perm(seed, epoch 1) | ...
+          save() at the wrap and load() elsewhere resumes at the same offset of the same permutation
+
+Covered: off keeps file order (1); on yields a permutation of every prompt that differs from the
+file order and repeats for the same seed (2); the epoch wrap re-permutes and load() reproduces
+the order after the wrap (3); a save() mid-permutation resumes on the exact next prompt and
+carries through the following wrap (4).
+"""
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
+
+import json
+from argparse import Namespace
+
+from miles.rollout.data_source import RolloutDataSourceWithBuffer
+
+PROMPTS = [f"p{i}" for i in range(5)]
+
+
+def _args(tmp_path, **overrides):
+    path = tmp_path / "train.jsonl"
+    if not path.exists():
+        path.write_text("".join(json.dumps({"input": p}) + "\n" for p in PROMPTS))
+    base = dict(
+        rollout_global_dataset=True,
+        prompt_data=str(path),
+        input_key="input",
+        metadata_key="metadata",
+        rollout_seed=7,
+        rollout_shuffle=True,
+        n_samples_per_prompt=1,
+        save=str(tmp_path / "ckpt"),
+        load=None,
+        buffer_filter_path=None,
+    )
+    return Namespace(**{**base, **overrides})
+
+
+def _prompts(groups):
+    return [group[0].prompt for group in groups]
+
+
+def test_off_reads_the_file_order(tmp_path):
+    source = RolloutDataSourceWithBuffer(_args(tmp_path, rollout_shuffle=False))
+
+    assert _prompts(source.get_samples(3)) == ["p0", "p1", "p2"]
+    assert _prompts(source.get_samples(3)) == ["p3", "p4", "p0"]
+
+
+def test_on_permutes_every_prompt_deterministically(tmp_path):
+    first = _prompts(RolloutDataSourceWithBuffer(_args(tmp_path)).get_samples(5))
+    again = _prompts(RolloutDataSourceWithBuffer(_args(tmp_path)).get_samples(5))
+
+    assert sorted(first) == PROMPTS
+    assert first != PROMPTS
+    assert again == first
+
+
+def test_epoch_wrap_repermutes_and_load_resumes_it(tmp_path):
+    source = RolloutDataSourceWithBuffer(_args(tmp_path))
+    epoch0 = _prompts(source.get_samples(5))
+    source.save(rollout_id=1)
+    epoch1_start = _prompts(source.get_samples(3))
+
+    resumed = RolloutDataSourceWithBuffer(_args(tmp_path, load=str(tmp_path / "ckpt")))
+    resumed.load(rollout_id=1)
+
+    assert source.dataset.epoch_id == 1
+    assert epoch1_start != epoch0[:3]
+    assert _prompts(resumed.get_samples(3)) == epoch1_start
+
+
+def test_resume_mid_epoch_continues_the_same_permutation(tmp_path):
+    """A resumed run must not restart the permutation or re-serve prompts it already used."""
+    source = RolloutDataSourceWithBuffer(_args(tmp_path))
+    source.get_samples(2)
+    source.save(rollout_id=2)
+    continued = [_prompts(source.get_samples(2)) for _ in range(3)]
+
+    resumed = RolloutDataSourceWithBuffer(_args(tmp_path, load=str(tmp_path / "ckpt")))
+    resumed.load(rollout_id=2)
+
+    assert [_prompts(resumed.get_samples(2)) for _ in range(3)] == continued
+    assert resumed.dataset.epoch_id == 1