feat(train): implement checkpointing and robust data loading

Enhance the training pipeline with stateful checkpointing and improve
the resilience of the data loading process against filesystem latency
and transient I/O errors.

- Implement auto-resuming in `train_ddp` by loading model, optimizer,
  and scheduler states from `state.pt`.
- Add atomic state saving using temporary files to prevent corruption.
- Introduce `_read_nii` with exponential backoff retries to handle
  transient NFS/filesystem failures during NIfTI reading.
- Add explicit error handling for missing or unreadable label files in
  `PatchDataset`.
- Update `sliding_window_probs` to conditionally apply Test-Time
  Augmentation (TTA) based on the `tta` parameter.
- Add `scripts/test_dataloader.py` for verifying dataset integrity.

scripts/test_dataloader.py

@@ -0,0 +1,24 @@
+"""Test: iterate the full train dataloader with workers, report None-label batches."""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def main():
+    from src import training
+    from src.common import load_jsonl
+    rows = load_jsonl(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "manifests", "split_train.jsonl"))
+    for seed in (0, 1, 2):
+        ds, dl = training.make_dataloader(rows, (96, 96, 96), 3, True, num_workers=4, seed=seed)
+        n_none, n = 0, 0
+        for i, (img, lab, w) in enumerate(dl):
+            if lab is None:
+                n_none += 1
+                print("  seed", seed, "batch", i, "None label", flush=True)
+            n += 1
+        print("seed", seed, "batches", n, "none-label", n_none, flush=True)
+    print("TEST DONE")
+
+
+if __name__ == "__main__":
+    main()

src/dataset.py

@@ -1,13 +1,25 @@
 import os
+import time
 import numpy as np
 import torch
 from torch.utils.data import Dataset, DataLoader
 
 
+def _read_nii(path, tries=10, delay=2.0):
+    import SimpleITK as sitk
+    for i in range(tries):
+        try:
+            return sitk.GetArrayFromImage(sitk.ReadImage(path))
+        except Exception:
+            if i == tries - 1:
+                raise
+            time.sleep(delay * (i + 1))
+
+
 class PatchDataset(Dataset):
     """Random 3D patches from preprocessed nifti volumes.
 
-    row: {key, pimg, plabel?, w (loss weight, default 1.0)}
+    row: {key, pimg, plabel?, is_neg?, w (loss weight, default 1.0)}
     yields (img(1,H,W,D), label(H,W,D) or None, weight)
     """
 
@@ -20,36 +32,34 @@ class PatchDataset(Dataset):
     def __len__(self):
         return len(self.rows)
 
-    @staticmethod
-    def _read(path):
-        import SimpleITK as sitk
-        return sitk.GetArrayFromImage(sitk.ReadImage(path))
-
     def __getitem__(self, idx):
         row = self.rows[idx]
         w = float(row.get("w", 1.0))
-        img = self._read(row["pimg"]).astype(np.float32)
+        img = _read_nii(row["pimg"]).astype(np.float32)
         img = np.nan_to_num(img, nan=0.0, posinf=1.0, neginf=0.0)
         img = np.clip(img, 0.0, 1.5)
-        # pad to multiple of 8 and at least the patch size (training) / 256 max (eval irrelevant)
         p = tuple(max((img.shape[i] + 7) // 8 * 8, self.patch[i] if self.train else 8) for i in range(3))
         if any(a < b for a, b in zip(img.shape, p)):
             img = np.pad(img, [(0, b - a) for a, b in zip(img.shape, p)])
         label = None
-        if row.get("plabel") and os.path.exists(row["plabel"]):
+        plabel = row.get("plabel")
-            label = self._read(row["plabel"]).astype(np.int64)
+        if plabel:
-            # pad/trim label to the same shape as img
+            if not os.path.exists(plabel):
+                for i in range(10):  # transient NFS handle failures
+                    time.sleep(2 * (i + 1))
+                    if os.path.exists(plabel):
+                        break
+                else:
+                    raise RuntimeError(f"label file vanished for {row.get('key')}: {plabel}")
+            label = _read_nii(plabel).astype(np.int64)
             if label.shape != img.shape:
-                lab_p = []
+                label = np.pad(label, [(0, max(img.shape[i] - label.shape[i], 0)) for i in range(3)])
-                for i in range(3):
-                    pad_i = 0
-                    pad_j = max(img.shape[i] - label.shape[i], 0)
-                    lab_p.append((pad_i, pad_j))
-                label = np.pad(label, lab_p)
                 label = label[: img.shape[0], : img.shape[1], : img.shape[2]]
             label = np.clip(np.rint(label).astype(np.int64), 0, 1)
         elif row.get("is_neg"):
             label = np.zeros(img.shape, dtype=np.int64)
+        elif row.get("plabel"):
+            raise RuntimeError(f"label file unreadable for {row.get('key')}: {row['plabel']}")
 
         if self.train:
             s0 = [self.rng.randint(0, img.shape[i] - self.patch[i] + 1) if img.shape[i] >= self.patch[i] else 0
@@ -97,6 +107,7 @@ def _collate(items):
     if all(l is not None for l in labs):
         labs = torch.stack(labs)
     else:
+        # rows without labels (should not mix with labeled rows in a training batch)
         labs = None
     wts = torch.stack([it[2] for it in items])
     return imgs, labs, wts
@@ -106,5 +117,6 @@ def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persi
     ds = PatchDataset(rows, patch=patch, train=train, seed=seed)
     dl = DataLoader(ds, batch_size=batch_size, shuffle=train, num_workers=num_workers,
                     collate_fn=_collate, drop_last=train and len(ds) > batch_size, pin_memory=True,
-                    persistent_workers=(train and num_workers > 0 and persistent), timeout=900)
+                    persistent_workers=(train and num_workers > 0 and persistent),
+                    timeout=0 if num_workers == 0 else 900)
     return ds, dl

src/training.py

@@ -31,9 +31,10 @@ def sliding_window_probs(model, vol, device, win=96, step=48, tta=True):
     cs = (coords(h), coords(w), coords(d))
     acc = np.zeros((h, w, d), dtype=np.float32)
     cnt = np.zeros((h, w, d), dtype=np.float32)
-    flips = [lambda x: x, lambda x: x[:, ::-1].copy(), lambda x: x[:, :, ::-1].copy(),
+    flips = [lambda x: x]
-             lambda x: x[::-1, :, :].copy()]
     if tta:
+        flips += [lambda x: x[:, ::-1].copy(), lambda x: x[:, :, ::-1].copy(), lambda x: x[::-1, :, :].copy()]
+    if tta == 2:
         flips.append(lambda x: x[::-1, ::-1].copy())
         flips.append(lambda x: x[::-1, :, ::-1].copy())
         flips.append(lambda x: x[:, ::-1, ::-1].copy())
@@ -127,14 +128,37 @@ def train_ddp(rows, val_rows, args):
     ckpt_dir = args.ckpt_dir
     if rank == 0:
         os.makedirs(ckpt_dir, exist_ok=True)
+    start_epoch = 0
+    state_f = os.path.join(ckpt_dir, "state.pt")
+    if world > 1:
+        dist.barrier()
+    if os.path.exists(state_f):
+        st = torch.load(state_f, map_location="cpu", weights_only=True)
+        model.load_state_dict(st["model"])
+        opt.load_state_dict(st["opt"])
+        sched.load_state_dict(st["sched"])
+        best, best_epoch, start_epoch = st["best"], st["best_epoch"], st["epoch"]
+        if rank == 0:
+            print(f"[train] auto-resuming from state.pt @epoch {start_epoch} (best={best:.4f})", flush=True)
+    if rank == 0:
         print(f"[train] rows={len(rows)} val={len(val_rows)} epochs={args.epochs} steps/epoch={steps_per_epoch}", flush=True)
+    if world > 1:
+        dist.barrier()
 
-    for epoch in range(args.epochs):
+    def save_state():
+        tmp = state_f + ".tmp"
+        torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(),
+                    "best": best, "best_epoch": best_epoch, "epoch": epoch + 1}, tmp)
+        os.replace(tmp, state_f)
+
+    for epoch in range(start_epoch, args.epochs):
         if hasattr(dl.sampler, "set_epoch"):
             dl.sampler.set_epoch(epoch)
         model.train()
         run_loss, run_n = 0.0, 0
         for img, lab, wts in dl:
+            if lab is None:
+                raise RuntimeError("training batch without labels; mixing labeled/unlabeled rows is not supported")
             img = img.to(device, non_blocking=True)
             lab = lab.to(device, non_blocking=True)
             wts = wts.to(device, non_blocking=True)
@@ -155,6 +179,8 @@ def train_ddp(rows, val_rows, args):
         if rank == 0:
             print(f"  epoch {epoch+1}/{args.epochs} loss={run_loss/max(run_n,1):.4f} "
                   f"lr={opt.param_groups[0]['lr']:.2e}", flush=True)
+        if rank == 0 and (epoch + 1) % max(1, args.val_every // 2) == 0:
+            save_state()
         if val_rows and rank == 0 and (epoch + 1) % args.val_every == 0:
             d, _ = evaluate(model, val_rows[: args.val_limit], device, patch[0], tta=True)
             print(f"  [val] epoch {epoch+1} dice={d:.4f}", flush=True)
@@ -163,6 +189,7 @@ def train_ddp(rows, val_rows, args):
                 best_epoch = epoch + 1
                 torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best},
                            os.path.join(ckpt_dir, "best.pt"))
+                save_state()
         if world > 1:
             dist.barrier()
     if rank == 0: