Use permutation-cycle index sampling in diffusion trainers and guard non-object run records

Replace the with-replacement per-batch index draw in the SDXL and DiT LoRA
trainers with a shared PermutationBatchSampler that visits every image once per
cycle before repeating, so short runs cover the whole dataset. The sampler
reshuffles from the run's rng so the index stream stays seed-deterministic.

Guard the diffusion run detail route against a valid-JSON non-object record,
which previously raised TypeError and returned a 500; it now 404s like the list
path's shape check.

Add regression tests for both.

studio/backend/core/training/diffusion_dit_trainer.py

@@ -49,6 +49,7 @@ from core.training.diffusion_train_common import (
     _restore_perf_flags,
     discover_image_caption_pairs,
     has_functional_torchao,
+    PermutationBatchSampler,
     repo_is_prequantized,
     resolve_train_steps,
 )
@@ -1158,6 +1159,10 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
     transformer.train()
     n_images = len(image_paths)
     batch_size = cfg.train_batch_size
+    # Permutation-cycle index sampler (shared with the SDXL trainer): visits every image once
+    # per cycle before repeating, so a short run covers the whole dataset instead of the old
+    # with-replacement draw. Uses the loop's own rng to stay seed-deterministic.
+    index_sampler = PermutationBatchSampler(n_images, rng)
     stopped = False
     running_loss = 0.0
     peak_gb = 0.0
@@ -1177,7 +1182,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
         optimizer.zero_grad(set_to_none = True)
         step_loss = 0.0
         for _ in range(cfg.gradient_accumulation_steps):
-            idxs = [rng.randrange(n_images) for _ in range(batch_size)]
+            idxs = index_sampler.next_batch(batch_size)
             if latent_cache is not None:
                 latents = _sample_cached_latents(
                     latent_cache, idxs, variant_rng, device, weight_dtype

studio/backend/core/training/diffusion_lora_trainer.py

@@ -59,6 +59,7 @@ from core.training.diffusion_train_common import (  # noqa: F401
     _restore_perf_flags,
     discover_image_caption_pairs,
     get_trainer,
+    PermutationBatchSampler,
     resolve_train_steps,
 )
 
@@ -436,8 +437,14 @@ def run_diffusion_lora_training(
 
         _emit(on_event, "model_load_completed")
 
+        # Permutation-cycle index sampler (shared with the DiT trainer): each dataset image is
+        # visited once per cycle before any repeat, so a short run does not leave part of a
+        # small dataset unseen. Draws from the loop's own rng so the sequence stays
+        # seed-deterministic.
+        index_sampler = PermutationBatchSampler(len(pairs), rng)
+
         def _next_batch() -> tuple[list[int], list[str], list[str]]:
-            idx = rng.sample(range(len(pairs)), k = min(cfg.train_batch_size, len(pairs)))
+            idx = index_sampler.next_batch(min(cfg.train_batch_size, len(pairs)))
             chosen = [pairs[i] for i in idx]
             return idx, [c[0] for c in chosen], [c[1] for c in chosen]
 

studio/backend/core/training/diffusion_train_common.py

@@ -431,6 +431,47 @@ def resolve_train_steps(cfg: "DiffusionLoraConfig", n_images: int) -> int:
     return cfg.train_steps
 
 
+class PermutationBatchSampler:
+    """Yields batch indices as consecutive slices of a reshuffled permutation of
+    ``range(n)``, so every index is visited exactly once per cycle before any repeats --
+    an epoch-style full pass instead of the with-replacement draw that leaves part of a
+    small dataset unseen at low step counts (num_epochs converts to a step budget, but the
+    per-batch index draw is what decides coverage). When a cycle is exhausted the order is
+    reshuffled from the run's own ``rng`` so the index stream stays seed-deterministic and
+    each cycle differs.
+
+    Both trainers share this so the SDXL ``_next_batch`` path and the DiT per-sample draw
+    select indices the same way. Only the index selection changes (with-replacement ->
+    permutation cycles); step count and batch shapes are unchanged.
+    """
+
+    def __init__(self, n: int, rng: random.Random) -> None:
+        if n <= 0:
+            raise ValueError("PermutationBatchSampler needs at least one item")
+        self._n = n
+        self._rng = rng
+        self._order: list[int] = []
+        self._pos = 0
+
+    def _reshuffle(self) -> None:
+        self._order = list(range(self._n))
+        self._rng.shuffle(self._order)
+        self._pos = 0
+
+    def next_batch(self, k: int) -> list[int]:
+        # k may exceed n (batch larger than the dataset): the permutation is refilled across
+        # as many cycles as needed so the caller always gets exactly k indices and the batch
+        # never shrinks, matching the old sampler's fixed batch shape.
+        out: list[int] = []
+        while len(out) < k:
+            if self._pos >= len(self._order):
+                self._reshuffle()
+            take = min(k - len(out), len(self._order) - self._pos)
+            out.extend(self._order[self._pos : self._pos + take])
+            self._pos += take
+        return out
+
+
 def discover_image_caption_pairs(
     data_dir: str | os.PathLike[str],
     *,

studio/backend/routes/training.py

@@ -1337,7 +1337,11 @@ async def get_diffusion_training_run(
     from core.training.diffusion_training_service import get_diffusion_run
 
     rec = get_diffusion_run(job_id)
-    if rec is None:
+    # A valid-JSON file that is not an object (a truncated / hand-edited [] record) would make
+    # DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below
+    # -- and 500 the endpoint. Treat any non-dict record as absent, matching the list route's
+    # shape check.
+    if not isinstance(rec, dict):
         raise HTTPException(status_code = 404, detail = "No such training run.")
     try:
         return DiffusionTrainingRunDetail(**rec)

studio/backend/tests/test_diffusion_training.py

@@ -427,6 +427,45 @@ def test_config_from_dict_epoch_mode_drops_max_steps_sentinel():
     assert cfg_explicit.train_steps == 25
 
 
+def test_permutation_sampler_covers_dataset_once_per_cycle():
+    # Every index must appear exactly once per cycle before any repeat (epoch-style pass),
+    # so a short run over a small dataset never leaves images unseen the way the old
+    # with-replacement draw did. Consecutive cycles must be reshuffled (differ).
+    import random
+
+    from core.training.diffusion_train_common import PermutationBatchSampler
+
+    n = 100
+    sampler = PermutationBatchSampler(n, random.Random(0))
+
+    # Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans
+    # the cycle boundary); the first n indices must be a permutation of range(n).
+    drawn: list[int] = []
+    while len(drawn) < n:
+        drawn.extend(sampler.next_batch(3))
+    first_cycle = drawn[:n]
+    assert sorted(first_cycle) == list(range(n))  # each index once, none missing
+
+    # The next full cycle is also a permutation, and it is reshuffled (order differs).
+    fresh = PermutationBatchSampler(n, random.Random(0))
+    cycle_a = fresh.next_batch(n)
+    cycle_b = fresh.next_batch(n)
+    assert sorted(cycle_a) == list(range(n))
+    assert sorted(cycle_b) == list(range(n))
+    assert cycle_a != cycle_b  # cycles are reshuffled, not repeated in the same order
+
+    # A seed replays the exact index stream (determinism for reproducible runs).
+    replay = PermutationBatchSampler(n, random.Random(0))
+    assert replay.next_batch(n) == cycle_a
+
+    # A batch larger than the dataset refills across cycles so it never shrinks (batch shape
+    # preserved), even though it must then repeat indices within the batch.
+    big = PermutationBatchSampler(4, random.Random(1))
+    batch = big.next_batch(10)
+    assert len(batch) == 10
+    assert set(batch) == {0, 1, 2, 3}
+
+
 def test_route_start_accepts_zero_max_grad_norm(client):
     # 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
     # the request model must not reject it.
@@ -937,3 +976,16 @@ def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir):
     assert r.status_code == 200, r.text
     adapters = [x["adapter"] for x in r.json()["runs"]]
     assert adapters == ["good"]  # the bad-field record was skipped, the good one remained
+
+
+def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir):
+    # A valid-JSON but non-object record (a truncated / hand-edited [] file named with a real
+    # job id) makes DiffusionTrainingRunDetail(**rec) raise TypeError, not ValidationError; the
+    # detail route must shape-check like the list path and 404 instead of 500.
+    import json
+
+    job_id = "a" * 32
+    (_isolated_runs_dir / f"{job_id}.json").write_text(json.dumps([]))
+
+    r = client.get(f"/api/train/diffusion/runs/{job_id}")
+    assert r.status_code == 404, r.text