fix: make the random-choose bucket usable, and seed every draw with torch

smauglab/transforms/gpu/contrast.py

@@ -1,5 +1,4 @@
 import math
-import random
 from collections.abc import Callable
 from typing import Any, Union
 
@@ -9,6 +8,7 @@
 from torch.nn import functional as F
 
 from smauglab.transforms.gpu.base import ImageOnlyTransform
+from smauglab.transforms.rng import shared_choice
 
 
 def _choose_region_mode(p_in: float, p_out: float, seg_mask: torch.Tensor | None) -> str:  # noqa: ARG001 -- seg_mask kept for signature symmetry with _apply_region_mode
@@ -216,7 +216,7 @@ def get_kernel(self, device: torch.device) -> Union[Tensor, list[Tensor]]:
             kernel = get_gaussian_kernel3d(kernel_size, sigma, torch.float32, device)
         elif self.kernel_type == "RandConv":
             # choose random odd kernel size e.g. [1,3,5,7]
-            k = int(random.choice(self.kernel_sizes))  # define kernel_sizes in __init__
+            k = int(shared_choice(self.kernel_sizes))  # define kernel_sizes in __init__
 
             std = 1.0 / math.sqrt(k * k)
             kernel = torch.randn((k, k, k), device=device) * std  # for 3D

smauglab/transforms/gpu/fromSeg.py

@@ -1,12 +1,11 @@
-import random
 from typing import Any
 
 import torch
-import torch.distributed as dist
 from torch import Tensor, nn
 from torch.nn import functional as F
 
 from smauglab.transforms.gpu.base import ImageOnlyTransform
+from smauglab.transforms.rng import shared_choice
 
 # ── PALETTE AUG helpers ──────────────────────────────────────────────────
 
@@ -432,7 +431,7 @@ def apply_transform(
         synth = torch.stack(synth_list)  # (B, N)
         synth_01 = synth.reshape(B, 1, D, H, W)
 
-        sigma = random.choice(self.blur_sigmas)
+        sigma = shared_choice(self.blur_sigmas)
         if sigma > 0.0:
             synth_01 = _gaussian_blur_3d(synth_01, sigma)
             synth = synth_01.reshape(B, N)
@@ -472,7 +471,7 @@ def apply_transform(
 
         # ── Step 3: optional second blur, then foreground z-score ─────────────
         synth_01 = synth.reshape(B, 1, D, H, W)
-        sigma2 = random.choice(self.blur_sigmas)
+        sigma2 = shared_choice(self.blur_sigmas)
         if sigma2 > 0.0:
             synth_01 = _gaussian_blur_3d(synth_01, sigma2)
             synth = synth_01.reshape(B, N)
@@ -489,20 +488,6 @@ def apply_transform(
         return out
 
 
-_SHARED_RNG_COUNTER = 0
-
-
-def _next_shared_seed() -> int:
-    global _SHARED_RNG_COUNTER  # noqa: PLW0603 -- module-level counter is the point: it makes successive seeds distinct
-    _SHARED_RNG_COUNTER += 1
-    seed = (int(torch.initial_seed()) + _SHARED_RNG_COUNTER) % (2**63 - 1)
-    if dist.is_available() and dist.is_initialized():
-        seed_tensor = torch.tensor([seed], dtype=torch.long)
-        dist.broadcast(seed_tensor, src=0)
-        seed = int(seed_tensor.item())
-    return seed
-
-
 def _minmax_norm(x: torch.Tensor, eps: float = 1e-8) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Per-sample min-max normalise to [0, 1]. Returns (normed, min, max)."""
     B = x.shape[0]
@@ -532,19 +517,6 @@ def _zscore_renorm(x: torch.Tensor, bg_threshold: float = 1e-6) -> torch.Tensor:
     return torch.where(fg, (x - mean) / std, torch.zeros_like(x))
 
 
-def _shared_cpu_generator() -> torch.Generator:
-    generator = torch.Generator(device="cpu")
-    generator.manual_seed(_next_shared_seed())
-    return generator
-
-
-def _shared_rand(shape: tuple[int, ...], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    if not (dist.is_available() and dist.is_initialized()):
-        return torch.rand(shape, device=device, dtype=dtype)
-    rand_cpu = torch.rand(shape, generator=_shared_cpu_generator(), device="cpu", dtype=dtype)
-    return rand_cpu.to(device=device, dtype=dtype)
-
-
 def collapse_onehot_to_index(seg_raw: torch.Tensor) -> torch.Tensor:
     """
     Convert a one-hot segmentation mask to a single-channel integer index mask.

smauglab/transforms/gpu/spatial.py

@@ -232,12 +232,17 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[
 
         scales = params["scale"]  # shape [B, 3]
 
-        if flags["data_keys"][0] is DataKey.IMAGE:
+        # Only MaskSequentialOpsCustom injects "data_keys" (see gpu/base.py), so a bare
+        # `flags["data_keys"]` raised KeyError for every other caller -- calling this
+        # transform standalone, or from inside RandomChooseXTransformsGPU, which passes
+        # the transform's own `flags`. Defaulting to IMAGE is what those callers mean.
+        data_keys = flags.get("data_keys") or [DataKey.INPUT]
+        if data_keys[0] in (DataKey.INPUT, DataKey.IMAGE):
             resample = "trilinear"
-        elif flags["data_keys"][0] is DataKey.MASK:
+        elif data_keys[0] is DataKey.MASK:
             resample = "nearest"
         else:
-            raise ValueError(f"Unsupported data key {flags['data_keys'][0]} for RandomLowResTransformGPU. Expected IMAGE or MASK.")
+            raise ValueError(f"Unsupported data key {data_keys[0]} for RandomLowResTransformGPU. Expected IMAGE or MASK.")
 
         # Define interpolation modes
         interp_down = resample

smauglab/transforms/gpu/transforms_list.py

@@ -723,9 +723,18 @@ def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor:
                 continue
             if not hasattr(t, "apply_transform"):
                 raise TypeError(f"All transforms must implement apply_transform like ImageOnlyTransform. Got {type(t)}")
-            # Most contrast transforms perform their random sampling inside apply_transform.
+            # Most contrast transforms perform their random sampling inside
+            # apply_transform, so an empty params dict is all they need. The ones with a
+            # kornia `_param_generator` (the spatial transforms) read their draw out of
+            # `params` instead, and calling apply_transform directly skips the
+            # forward_parameters step that fills it -- they used to raise
+            # "params must contain 'scale'" from inside a bucket. Sampling here keeps
+            # the bucket usable for both kinds.
+            t_params = child_params
+            if getattr(t, "_param_generator", None) is not None:
+                t_params = {**child_params, **t.forward_parameters(x.shape)}
             t_flags = getattr(t, "flags", {})
-            x = t.apply_transform(x, child_params, t_flags, transform=None)
+            x = t.apply_transform(x, t_params, t_flags, transform=None)
         return x
 
     @torch.no_grad()  # disable gradients for efficiency
@@ -736,7 +745,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[
             return self._apply_mix(input, seg)
 
         batch_size = input.shape[0]
-        out = input
+        # A clone, not `out = input`: the loop writes back through `out[i:i+1]`, so
+        # without it the caller's batch is modified in place. Every sibling transform
+        # in gpu/spatial.py clones.
+        out = input.clone()
         for i in range(batch_size):
             xi = out[i : i + 1]
             seg_i = None

smauglab/transforms/rng.py

@@ -0,0 +1,70 @@
+"""Random draws that `torch.manual_seed` actually reaches, and that DDP ranks agree on.
+
+Several GPU transforms reached for Python's `random.choice` to pick a blur sigma or a
+kernel size, inside an `apply_transform` that was otherwise entirely `torch.rand`
+driven. Two consequences:
+
+* `torch.manual_seed(...)` does not seed Python's `random`, so a "seeded" run was not
+  reproducible. The test suite hid this -- `unit_tests/helpers.py::seed_everything`
+  seeds torch, numpy *and* random -- but training does not call that.
+* Under DistributedDataParallel each rank has its own `random` state, so ranks picked
+  different sigmas for the same batch.
+
+`gpu/fromSeg.py` already contained `_next_shared_seed` / `_shared_rand` written for
+exactly this, and never called them. That machinery lives here now, with the `choice`
+helper the call sites actually needed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TypeVar
+
+import torch
+import torch.distributed as dist
+
+T = TypeVar("T")
+
+_SHARED_RNG_COUNTER = 0
+
+
+def next_shared_seed() -> int:
+    """A seed every rank agrees on, different on each call."""
+    global _SHARED_RNG_COUNTER  # noqa: PLW0603 -- module-level counter is the point: it makes successive seeds distinct
+    _SHARED_RNG_COUNTER += 1
+    seed = (int(torch.initial_seed()) + _SHARED_RNG_COUNTER) % (2**63 - 1)
+    if dist.is_available() and dist.is_initialized():
+        seed_tensor = torch.tensor([seed], dtype=torch.long)
+        dist.broadcast(seed_tensor, src=0)
+        seed = int(seed_tensor.item())
+    return seed
+
+
+def shared_cpu_generator() -> torch.Generator:
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(next_shared_seed())
+    return generator
+
+
+def shared_rand(shape: tuple[int, ...], device: torch.device, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    """Uniform [0, 1) draws; identical across ranks when running under DDP.
+
+    Outside DDP this is just `torch.rand`, so it stays on whatever device and generator
+    the caller has already seeded.
+    """
+    if not (dist.is_available() and dist.is_initialized()):
+        return torch.rand(shape, device=device, dtype=dtype)
+    rand_cpu = torch.rand(shape, generator=shared_cpu_generator(), device="cpu", dtype=dtype)
+    return rand_cpu.to(device=device, dtype=dtype)
+
+
+def shared_choice(options: Sequence[T]) -> T:
+    """Pick one element of `options`, using torch's RNG rather than Python's.
+
+    The drop-in replacement for `random.choice` in a transform.
+    """
+    if len(options) == 0:
+        raise ValueError("cannot choose from an empty sequence")
+    draw = float(shared_rand((1,), torch.device("cpu")).item())
+    # torch.rand is [0, 1), so the index is already in range; the clamp is belt-and-braces.
+    return options[min(int(draw * len(options)), len(options) - 1)]

unit_tests/test_bucket_and_rng.py

@@ -0,0 +1,152 @@
+"""`RandomChooseXTransformsGPU`, and the draws that `torch.manual_seed` did not reach.
+
+* The bucket wrote into the caller's batch, and could not run any transform with a
+  kornia parameter generator: calling `apply_transform` directly skips the
+  `forward_parameters` step that fills `params`, so those raised
+  "params must contain 'scale'".
+* `RandomLowResTransformGPU` read `flags["data_keys"]` unguarded, which only the mask
+  path injects -- so it raised `KeyError` standalone and inside a bucket.
+* Blur sigmas and kernel sizes were drawn with Python's `random`, which
+  `torch.manual_seed` does not reach and which diverges across DDP ranks.
+"""
+
+import torch
+
+from smauglab.transforms.gpu.contrast import RandomConvTransformGPU
+from smauglab.transforms.gpu.spatial import RandomLowResTransformGPU
+from smauglab.transforms.gpu.transforms_list import RandomChooseXTransformsGPU
+from smauglab.transforms.rng import shared_choice, shared_rand
+from unit_tests.helpers import SmaugLabTestCase, first_output
+
+
+class TestLowResRunsOutsideTheMaskPath(SmaugLabTestCase):
+    def test_it_runs_standalone(self):
+        """flags['data_keys'] is only injected by MaskSequentialOpsCustom."""
+        torch.manual_seed(0)
+        transform = RandomLowResTransformGPU(p=1.0)
+        volume = self.tiny_volume()
+
+        out = first_output(transform(volume))
+
+        self.assertIsImageLike(out, volume, "RandomLowResTransformGPU")
+
+    def test_it_runs_inside_a_random_choose_bucket(self):
+        """The bucket calls apply_transform with the transform's own flags, which
+        carry no data_keys either."""
+        torch.manual_seed(0)
+        bucket = RandomChooseXTransformsGPU(transforms_list=[RandomLowResTransformGPU(p=1.0)], num_transforms=1, p=1.0)
+        volume = self.tiny_volume()
+
+        out = bucket.apply_transform(volume.clone(), {}, {}, transform=None)
+
+        self.assertIsImageLike(out, volume, "RandomLowResTransformGPU in a bucket")
+
+    def test_an_explicit_mask_key_still_selects_nearest(self):
+        """The branch that does exist must keep working."""
+        from kornia.constants import DataKey
+
+        torch.manual_seed(0)
+        transform = RandomLowResTransformGPU(p=1.0)
+        seg = self.tiny_seg()
+        params = transform.forward_parameters(seg.shape)
+
+        out = transform.apply_transform(seg.clone(), params, {"data_keys": [DataKey.MASK]}, transform=None)
+
+        self.assertEqual(out.shape, seg.shape)
+        self.assertTrue(bool(torch.isin(out, torch.tensor([0.0, 1.0])).all()), "a mask was resampled with interpolation")
+
+
+class TestBucketDoesNotMutateItsInput(SmaugLabTestCase):
+    def test_the_callers_tensor_is_left_alone(self):
+        torch.manual_seed(0)
+        bucket = RandomChooseXTransformsGPU(
+            transforms_list=[RandomLowResTransformGPU(p=1.0)],
+            num_transforms=1,
+            p=1.0,
+            same_on_batch=False,
+        )
+        volume = torch.rand(3, 1, 12, 12, 12)
+        before = volume.clone()
+
+        bucket.apply_transform(volume, {}, {}, transform=None)
+
+        self.assertTrue(torch.equal(volume, before), "RandomChooseXTransformsGPU wrote into the caller's batch")
+
+    def test_it_still_returns_something_transformed(self):
+        """Cloning must not turn the bucket into a no-op."""
+        torch.manual_seed(0)
+        bucket = RandomChooseXTransformsGPU(
+            transforms_list=[RandomConvTransformGPU(kernel_type="Laplace", p=1.0)],
+            num_transforms=1,
+            p=1.0,
+            same_on_batch=False,
+        )
+        volume = torch.rand(2, 1, 10, 10, 10)
+
+        out = bucket.apply_transform(volume.clone(), {}, {}, transform=None)
+
+        self.assertFalse(torch.allclose(out, volume, atol=1e-6))
+
+    def test_an_empty_bucket_is_a_no_op(self):
+        bucket = RandomChooseXTransformsGPU(transforms_list=[], num_transforms=0, p=1.0)
+        volume = self.tiny_volume()
+        self.assertTrue(torch.equal(bucket.apply_transform(volume.clone(), {}, {}, transform=None), volume))
+
+    def test_the_same_on_batch_path_also_runs_a_generator_transform(self):
+        torch.manual_seed(0)
+        bucket = RandomChooseXTransformsGPU(
+            transforms_list=[RandomLowResTransformGPU(p=1.0)],
+            num_transforms=1,
+            p=1.0,
+            same_on_batch=True,
+        )
+        volume = self.tiny_volume()
+
+        out = bucket.apply_transform(volume.clone(), {}, {}, transform=None)
+
+        self.assertIsImageLike(out, volume, "bucket with same_on_batch")
+
+
+class TestTorchSeedReachesEveryDraw(SmaugLabTestCase):
+    """`torch.manual_seed` alone must be enough.
+
+    The suite's own `seed_everything` seeds torch, numpy *and* Python's `random`, which
+    is exactly why this went unnoticed -- training does not call it. These tests seed
+    only torch.
+    """
+
+    def test_randconv_is_reproducible_under_torch_seed_alone(self):
+        outputs = []
+        for _ in range(2):
+            torch.manual_seed(1234)
+            transform = RandomConvTransformGPU(kernel_type="RandConv", p=1.0, kernel_sizes=[1, 3, 5, 7])
+            outputs.append(transform.apply_transform(self.tiny_volume(), {}, {}, transform=None).clone())
+
+        self.assertTrue(torch.equal(outputs[0], outputs[1]), "RandomConvTransformGPU drew its kernel size from an unseeded generator")
+
+    def test_shared_choice_covers_the_whole_sequence(self):
+        torch.manual_seed(0)
+        options = (1, 3, 5, 7)
+
+        seen = {shared_choice(options) for _ in range(200)}
+
+        self.assertEqual(seen, set(options))
+
+    def test_shared_choice_is_reproducible(self):
+        def draw():
+            torch.manual_seed(7)
+            return [shared_choice((1, 3, 5, 7)) for _ in range(20)]
+
+        self.assertEqual(draw(), draw())
+
+    def test_shared_choice_rejects_an_empty_sequence(self):
+        with self.assertRaises(ValueError):
+            shared_choice([])
+
+    def test_shared_rand_is_plain_torch_rand_outside_ddp(self):
+        """No process group initialised, so it must stay on the caller's generator."""
+        torch.manual_seed(11)
+        expected = torch.rand((4,))
+        torch.manual_seed(11)
+
+        self.assertTrue(torch.equal(shared_rand((4,), torch.device("cpu")), expected))