From d9fa7a96968d84ee19409ef06d6533b81cce51a7 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 21:21:14 +0200 Subject: [PATCH 1/4] fix(data): the correctness pass over what 1.8.0 ships A uint8 label map resampled through a `field` was interpolated and truncated: `_stream_mode` answered nearest for uint8, but only the two paths WITHOUT a field consulted it. Over a source holding {0, 100} the warped path returned 29, 79 and 99 -- labels nobody wrote, in a volume that is still a valid label map. ResampleToReference takes an `interpolation` now, and every sampler asks one method for it: a declaration honoured on one path and ignored on another is worse than none, because the page telling a user to set it is then right about half their chains. A store chunked on the writer's region verbatim, and a slab sweep declares the whole trailing plane: a gigabyte per chunk at 2048x2048, past what zarr holds in one buffer at 4096x4096, paid by every partial read after. An axis the region covers end to end is tiled; the axis the writer advances along keeps its height. Appending pyramid levels rewrote level 0 on ngff-zarr's default chunking, so a streamed write and a declared pyramid cancelled each other out. A statistic after a Reduce was seeded from the fold, so [Reduce, Clip, Normalize] normalised by the unclipped Max and wrote a volume its own header did not describe. `grid: strict` skipped a geometry key no header carried, which is quietest where it costs most -- a missing Direction is a flip that shows in neither extent nor spacing. A second Reduce marker fell past the split and came back diagnosed as an ordinary stage. A reduction is budgeted at each side's own width: members at theirs, the output at its, and whatever the operator builds over the buffer it was handed. Median stacks that buffer and sorts a copy, so a cohort of N peaks at 3N+1 member regions where the plan said N+1; Concat writes N x C where each member holds C, so charging the cohort at the output's width over-stated it by the cohort's size. Vote is added beside them, because Mean and Median both answer with values that were in no input. Warp built its grid on the CPU and met a GPU-resident volume in grid_sample; it also checks its declared bound on the whole-volume path. Chained sweeps stop at the first failure -- on the shared Expand pass as on the per-case one -- so the recorded reason is the cause and not a downstream symptom. A refused plan no longer records the half-folded state it stopped at. _drawn_from restores the CUDA generators it seeded. `from konfai.data import Clip` raised ImportError. The three samplers stated one arithmetic at each site and had already drifted: the CPU-half guard existed in three of them and not the fourth. sampling_dtype, nearest_index and window_index say each rule once, over both gathers. --- konfai/data/__init__.py | 9 +- konfai/data/case_reduction.py | 90 ++++++++--- konfai/data/data_manager.py | 15 +- konfai/data/patching.py | 121 ++++++++------- konfai/data/reduction.py | 42 ++++++ konfai/data/transform.py | 157 ++++++++++++++----- konfai/main.py | 9 +- konfai/predictor.py | 7 +- konfai/transformer.py | 19 ++- konfai/utils/dataset.py | 36 ++++- konfai/utils/ome_zarr.py | 25 +++- tests/unit/test_case_reduction.py | 110 +++++++++++++- tests/unit/test_dataset.py | 2 +- tests/unit/test_ome_zarr_data_surface.py | 26 ++++ tests/unit/test_resample_sampler_rules.py | 175 ++++++++++++++++++++++ tests/unit/test_resample_to_reference.py | 102 ++++++++++++- 16 files changed, 793 insertions(+), 152 deletions(-) create mode 100644 tests/unit/test_resample_sampler_rules.py diff --git a/konfai/data/__init__.py b/konfai/data/__init__.py index 62b852f4..ce2e8ed0 100644 --- a/konfai/data/__init__.py +++ b/konfai/data/__init__.py @@ -19,7 +19,8 @@ This is the surface for driving KonfAI's data machinery from plain Python — no YAML, no environment variables, no workflow. A chain of transforms applied to a dataset, out-of-core, is:: - from konfai.data import Dataset, DatasetManager, Clip, Write + from konfai.data import Dataset, DatasetManager, Write + from konfai.data.transform import Clip # the concrete stages stay in the module that defines them manager = DatasetManager( index=0, group_src="CT", group_dest="CT", name="CASE_000", @@ -41,6 +42,7 @@ from konfai.data.data_manager import DataMetric, DataPrediction, DatasetIter, DataTrain, DataTransform from konfai.data.patching import DatasetManager, DatasetPatch +from konfai.data.reduction import Concat, Mean, Median, Reduction, Vote from konfai.data.transform import ( Expand, LocalityKind, @@ -63,6 +65,7 @@ __all__ = [ "Attribute", + "Concat", "DataMetric", "DataPrediction", "DataTrain", @@ -73,11 +76,15 @@ "DatasetPatch", "Expand", "LocalityKind", + "Mean", + "Median", "PatchLocality", "Reduce", + "Reduction", "Save", "Transform", "TransformInverse", + "Vote", "Write", "append_ome_zarr_levels", "create_ome_zarr_store", diff --git a/konfai/data/case_reduction.py b/konfai/data/case_reduction.py index 6d770d49..2f58c462 100644 --- a/konfai/data/case_reduction.py +++ b/konfai/data/case_reduction.py @@ -37,7 +37,7 @@ from konfai.data.patching import DatasetManager from konfai.data.reduction import Reduction -from konfai.data.transform import LocalityKind, Reduce, Transform +from konfai.data.transform import LocalityKind, PatchLocality, Reduce, Transform from konfai.utils.config import apply_config from konfai.utils.dataset import ( Attribute, @@ -77,6 +77,11 @@ class ReductionPlan: slab_rows: int incremental: bool stat_pass: bool + #: Channels a MEMBER's region carries. Separate from ``channels``, the output's, because an + #: operator may change the count: ``Concat`` writes ``N x C`` where each member holds ``C``, so + #: charging the members at the output's width over-states the peak by the cohort's size. + source_channels: int = 0 + working_multiple: float = 0.0 refusal: str | None = None @property @@ -84,20 +89,36 @@ def streams(self) -> bool: return self.refusal is None @property - def resident_regions(self) -> int: - """Regions held at the peak: two for a running accumulator, else one per case plus the - output's own. This is what the budget multiplies, and why ``slab_rows`` is derived.""" - return 2 if self.incremental else len(self.cases) + 1 + def buffered_regions(self) -> int: + """Member regions resident at once: one for a running accumulator, else the whole cohort.""" + return 1 if self.incremental else len(self.cases) + + @property + def resident_regions(self) -> float: + """Regions held at the peak, counted in MEMBER regions plus the output's own. + + ``Median`` stacks the buffer into a new tensor and sorts a copy of that, so counting the + buffer alone under-states its peak threefold -- and it is what a bare ``Reduce`` gets. + """ + return self.buffered_regions * (1 + self.working_multiple) + 1 + + def _region_bytes(self, channels: int) -> int: + return int(self.slab_rows * np.prod(self.spatial[1:], dtype=np.int64) * channels * _ASSUMED_ITEMSIZE) @property def region_bytes(self) -> int: - return int(self.slab_rows * np.prod(self.spatial[1:], dtype=np.int64) * self.channels * _ASSUMED_ITEMSIZE) + """One OUTPUT region, the unit the written slab is measured in.""" + return self._region_bytes(self.channels) @property def peak_bytes(self) -> int: + # Members at their own width, the output at its, and whatever the operator builds over the + # buffer it is handed -- which is member-sized, since that is what it was handed. + # # A statistics pass is a second traversal, not a second working set: it holds exactly what # one region holds, so the peak is the same whether there are one or two passes. - return self.resident_regions * self.region_bytes + members = self.buffered_regions * self._region_bytes(self.source_channels or self.channels) + return int(members * (1 + self.working_multiple) + self.region_bytes) def describe(self) -> str: verdict = "STREAM" if self.streams else "REFUSED" @@ -107,7 +128,7 @@ def describe(self) -> str: return "\n".join(lines) regime = "incremental accumulator" if self.incremental else "every case resident per region" lines.append( - f" {self.resident_regions} resident region(s) of {self.slab_rows} row(s)" + f" {self.resident_regions:g} resident region(s) of {self.slab_rows} row(s)" f" = {self.peak_bytes / (1 << 30):.2f} GiB ({regime})" ) if self.stat_pass: @@ -204,21 +225,34 @@ def check_post_stages(post: list[Transform], output: str) -> None: -- a halo, a resample, a reorientation -- would take that region for the whole volume and seam at every boundary: a plausible result, and a wrong one. Those are deferred, not forbidden: end the chain, and read the written volume back in a second chain where the ordinary planner can - pull regions through it. That is the boundary that already lets a statistic follow a - value-changing stage. + pull regions through it. + + A statistic may follow the reduction, but only over stages that leave the values alone: the stat + pass measures the FOLD, so an earlier stage that changes the values makes the seed describe a + volume nobody wrote. This mirrors the same refusal in the per-case planner. """ + localities: list[PatchLocality] = [] for index, stage in enumerate(post): - kind = stage.patch_locality(Attribute()).kind - if kind in _POST_KINDS: - continue + locality = stage.patch_locality(Attribute()) + kind = locality.kind name = type(stage).__name__ - raise ReductionError( - f"stage {index} '{name}' follows the Reduce into '{output}' and declares {kind.name}," - " which reads across space -- applied one region at a time it would seam at every" - " region boundary.", - f"Only voxel-local stages can follow a reduction. End this chain, and put '{name}' in a" - f" second chain that reads '{output}' back.", - ) + if kind not in _POST_KINDS: + raise ReductionError( + f"stage {index} '{name}' follows the Reduce into '{output}' and declares {kind.name}," + " which reads across space -- applied one region at a time it would seam at every" + " region boundary.", + f"Only voxel-local stages can follow a reduction. End this chain, and put '{name}' in a" + f" second chain that reads '{output}' back.", + ) + if kind is LocalityKind.GLOBAL_STAT and not all(previous.statistics_preserving for previous in localities): + raise ReductionError( + f"stage {index} '{name}' follows the Reduce into '{output}' and needs whole-volume" + " statistics, but an earlier stage after the Reduce changes the values -- the" + " statistic is measured on the fold, so it would not be this stage's input.", + f"End this chain after the value-changing stage, and put '{name}' in a second chain" + f" that reads '{output}' back, where its statistic is measured on what it receives.", + ) + localities.append(locality) @dataclass @@ -315,8 +349,18 @@ def check_grid(self) -> str | None: for manager in others: attribute = manager.landed_attributes() for key in _GEOMETRY_KEYS: - if key not in expected or key not in attribute: - continue + # ``strict`` is a promise that the geometries WERE compared, and a key nobody + # recorded cannot be. Skipping it is quietest exactly where it costs most: a + # Direction missing from one header is a flip that shows in neither extent nor + # spacing. Fold on extent alone with 'grid: shape_only' if that is what is meant. + absent = [ + name for name, side in ((reference.name, expected), (manager.name, attribute)) if key not in side + ] + if absent: + return ( + f"{' and '.join(repr(name) for name in absent)} lands on no {key}," + f" which 'grid: strict' compares (use 'grid: shape_only' to fold on extent alone)" + ) left = np.asarray(expected.get_np_array(key), dtype=np.float64).ravel() right = np.asarray(attribute.get_np_array(key), dtype=np.float64).ravel() if left.shape != right.shape or not np.allclose(left, right, atol=self.reduce.grid_tolerance): @@ -356,8 +400,10 @@ def plan(self) -> ReductionPlan: # Concat over N cases writes N times the channels, and the plan must probe and size the # shape the run will actually open. channels=self.operator.output_channels(int(reference.base_shape[0]), len(self.managers)), + source_channels=int(reference.base_shape[0]), slab_rows=self.slab_rows, incremental=self.operator.incremental, + working_multiple=float(self.operator.working_multiple), stat_pass=self._needs_stat_pass(), refusal=self._first_refusal(), ) diff --git a/konfai/data/data_manager.py b/konfai/data/data_manager.py index bc6723c3..4d58d132 100755 --- a/konfai/data/data_manager.py +++ b/konfai/data/data_manager.py @@ -1976,7 +1976,7 @@ def _seed_expansions(self) -> None: def _output_destinations(self) -> dict[tuple[str, str], list[tuple[str, str]]]: """Resolved ``(root, group)`` of every Save/Write, keyed by chain — the parse-time view of - what the run would write, resolved exactly as ``_save_destination`` will resolve it.""" + what the run would write, resolved exactly as ``save_destination`` will resolve it.""" destinations: dict[tuple[str, str], list[tuple[str, str]]] = {} for group_src in self.groups_src: for group_dest, group_transform in self.groups_src[group_src].items(): @@ -2058,13 +2058,24 @@ def _validate_expansion(self) -> None: chain = f"groups_src.{group_src}.groups_dest.{group_dest}" transforms = group_transform.transforms expands = [t for t in transforms if isinstance(t, Expand)] + reduces = [t for t in transforms if isinstance(t, Reduce)] if len(expands) > 1: raise TransformerError( f"'{chain}' declares {len(expands)} Expand markers; a chain changes its cardinality at most once.", "Keep one Expand per chain. Successive expansions compose across two" " invocations, the second reading the first one's output back.", ) - if expands and any(isinstance(t, Reduce) for t in transforms): + if len(reduces) > 1: + # Counted here, where a cardinality marker is refused under its own name. The chain + # splits at the FIRST Reduce, so a second one lands among the post stages and comes + # back as an ordinary stage that reads across space: a true sentence about the wrong + # problem, naming a remedy that does not apply. + raise TransformerError( + f"'{chain}' declares {len(reduces)} Reduce markers; a chain changes its cardinality at most once.", + "Keep one Reduce per chain. Successive reductions compose across two" + " invocations, the second reading the first one's output back.", + ) + if expands and reduces: raise TransformerError( f"'{chain}' declares both an Expand and a Reduce.", "One chain changes its cardinality once (1-to-N or N-to-1). Compose the two" diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 0848be1d..9ebf9a81 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -71,11 +71,15 @@ _SWEEP_RESIDENT_SLABS = 2 _SWEEP_ELEMENT_BYTES = 4 -# What a whole-volume fallback holds while a case is in flight -- the assembled tensor plus one -# transform output -- and the bytes each element travels as. The transform plan and the run-time -# budget check (_enforce_fallback_budget) must agree on this figure. -_FALLBACK_INFLIGHT_FACTOR = 2 -_CASE_ELEMENT_BYTES = 4 +#: What a whole-volume fallback holds while a case is in flight -- the assembled tensor plus one +#: transform output -- and the bytes each element travels as. +#: +#: Public because two callers must agree on the figure and neither owns it: the run-time budget check +#: (``_enforce_fallback_budget``) refuses a case against it, and the TRANSFORM plan prints and +#: enforces the same number before a byte is written. A plan estimating differently from the run it +#: describes is worse than no plan. +FALLBACK_INFLIGHT_FACTOR = 2 +CASE_ELEMENT_BYTES = 4 def _halo_radii(halo: tuple[int, ...], n_axes: int) -> list[int]: @@ -143,11 +147,15 @@ def _drawn_from(*key: object) -> Iterator[None]: draw the same copies on every run and on every rank of one run. The state is restored because it is global: reseeding everything downstream would make every - other random decision a function of how many copies were declared. + other random decision a function of how many copies were declared. CUDA's generators are part of + that state -- ``torch.manual_seed`` seeds every device, so restoring the CPU generator alone + would leave every later GPU draw reseeded from here. """ digest = hashlib.blake2b("|".join(str(part) for part in key).encode(), digest_size=4).digest() seed = int.from_bytes(digest, "big") - states = (random.getstate(), np.random.get_state(), torch.random.get_rng_state()) + cpu_state = torch.random.get_rng_state() + cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + states = (random.getstate(), np.random.get_state()) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) @@ -156,7 +164,9 @@ def _drawn_from(*key: object) -> Iterator[None]: finally: random.setstate(states[0]) np.random.set_state(states[1]) - torch.random.set_rng_state(states[2]) + torch.random.set_rng_state(cpu_state) + if cuda_states is not None: + torch.cuda.set_rng_state_all(cuda_states) def _stage_name(stage: Stage) -> str: @@ -277,8 +287,12 @@ class _ReadStagePlan: pull: Callable[[tuple[slice, ...]], list[slice]] | None -def _save_destination(save: Save, default_dataset: Dataset, default_group: str) -> tuple[Dataset, str]: - """The dataset and group a :class:`Save` caches into, the manager's own when it names none.""" +def save_destination(save: Save, default_dataset: Dataset, default_group: str) -> tuple[Dataset, str]: + """The dataset and group a :class:`Save` caches into, the manager's own when it names none. + + Public because a planner has to resolve a destination exactly as the engine will: one that probes + a store the run does not open has verified nothing. + """ if save.dataset: filename, _, file_format = split_path_spec( save.dataset, @@ -1510,7 +1524,7 @@ def _load(self, pre_transform: list[Transform]): data = None for transform_function in reversed(pre_transform): if isinstance(transform_function, Save): - dataset, group_dest = _save_destination(transform_function, self.dataset, self.group_dest) + dataset, group_dest = save_destination(transform_function, self.dataset, self.group_dest) if not self._rewrite_saves and dataset.is_dataset_exist(group_dest, self.name): data, attrib = dataset.read_data(group_dest, self.name) self.cache_attributes[0].update(attrib) @@ -1542,7 +1556,7 @@ def _apply_chain( for transform_function in transforms: tensor = transform_function(self.name, tensor, attribute) if isinstance(transform_function, Save): - dataset, group_dest = _save_destination(transform_function, self.dataset, self.group_dest) + dataset, group_dest = save_destination(transform_function, self.dataset, self.group_dest) dataset.write(group_dest, entry, tensor.numpy(), attribute) return tensor @@ -1740,6 +1754,12 @@ def _plan_stream_region( shape = [int(extent) for extent in source_spatial_shape] localities: list[PatchLocality] = [] plans: list[_ReadStagePlan] = [] + + def refuse(reason: str) -> tuple[bool, tuple[_ReadStagePlan, ...], Attribute, str]: + """A refusal carries the state folded so far, so the caller can still read the geometry + the chain reached before the stage that stopped it.""" + return False, (), evolved, reason + for stage_index, stage in enumerate(stages): loc = stage.patch_locality(Attribute(evolved)) localities.append(loc) @@ -1749,71 +1769,49 @@ def _plan_stream_region( # OUTPUT, which a patch read has no notion of. A stage that is whole-volume only # because something was left undeclared says so itself (PatchLocality.reason), so # the reader is told what to change instead of what happened. - return ( - False, - (), - evolved, - f"{label} declares {loc.kind.name}: {loc.reason or 'it needs the whole volume'}.", - ) + return refuse(f"{label} declares {loc.kind.name}: {loc.reason or 'it needs the whole volume'}.") if loc.kind is LocalityKind.GLOBAL_STAT: # The seed is the STORED volume's statistic, which is this transform's input only when # every earlier stage preserves it; otherwise ([Clip(-200, 400), Standardize()]) every # patch would be standardized by the pre-Clip statistic -- fall back to the whole volume. if not all(previous.statistics_preserving for previous in localities[:-1]): - return ( - False, - (), - evolved, + return refuse( f"{label} needs whole-volume statistics, but an earlier stage changes the values" - " -- the stored volume's statistic is not this stage's input.", + " -- the stored volume's statistic is not this stage's input." ) if seed_statistics and not self._ensure_stream_stats( source_dataset, source_group, source_entry, cache_attribute, set(loc.stat_keys), loc.stat_channels ): - return ( - False, - (), - evolved, - f"{label} needs statistics {sorted(loc.stat_keys)} that the source cannot provide.", - ) + return refuse(f"{label} needs statistics {sorted(loc.stat_keys)} that the source cannot provide.") # The evolving case state carries the seed too: a Save sweep writes it as the cache # header, exactly as the whole-volume pass leaves the statistic in the attribute. for stat_key in loc.stat_keys: if stat_key in cache_attribute and stat_key not in evolved: evolved[stat_key] = cache_attribute[stat_key] if loc.kind is LocalityKind.HALO and not self._affords_halo(a, loc.halo): - return ( - False, - (), - evolved, + return refuse( f"{label} declares a halo of {loc.halo} that is too wide for this grid to be worth" - " reading (over half the patch extent per axis).", + " reading (over half the patch extent per axis)." ) if loc.kind is LocalityKind.RESCALE and (not isinstance(stage, Resample) or "Spacing" not in evolved): # A resample is patch-native only when the source geometry is known: the scale is read # from the evolving 'Spacing' (a free geometry stat, no read_data_statistics). - return ( - False, - (), - evolved, + return refuse( f"{label} declares RESCALE but " + ( "does not inherit from Resample." if not isinstance(stage, Resample) else "the source carries no 'Spacing' to scale from." - ), + ) ) plan = self._plan_read_stage(stage, loc, shape, evolved) plans.append(plan) shape = list(plan.out_shape) expected = landing_shape if landing_shape is not None else self.shapes[a] if shape != [int(extent) for extent in expected]: - return ( - False, - (), - evolved, + return refuse( f"the chain's shapes fold to {shape} but the target grid is" - f" {[int(extent) for extent in expected]} -- a stage's shape map is missing or wrong.", + f" {[int(extent) for extent in expected]} -- a stage's shape map is missing or wrong." ) return True, tuple(plans), evolved, None @@ -1916,7 +1914,7 @@ def _resolve_patch_stream_source(self, a: int, apply_augmentations: bool = True) entry = self.copy_entry(a) continue if isinstance(transform, Save): - dataset, group = _save_destination(transform, self.dataset, self.group_dest) + dataset, group = save_destination(transform, self.dataset, self.group_dest) if not self._rewrite_saves and dataset.is_dataset_exist(group, entry): source_dataset, source_group, source_entry = dataset, group, entry source_shape, boundary_attributes = dataset.get_infos(group, entry) @@ -1998,8 +1996,12 @@ def _resolve_patch_stream_source(self, a: int, apply_augmentations: bool = True) source_dataset, source_group, source_entry, source_shape, stages, stage_plans ) # The state the whole plan lands on, kept for consumers that need the LANDED geometry (a - # reduction seeding its output header) without re-walking the chain. - self._stream_evolved[key] = Attribute(evolved) + # reduction seeding its output header) without re-walking the chain. Recorded only when the + # plan HOLDS: a refused plan folded as far as the stage that refused and no further, and half + # a fold is not a geometry -- it is a Spacing from before the resample meant to change it. + # An unset key is what lets ``landed_attributes`` answer with the stored state instead. + if streamable: + self._stream_evolved[key] = Attribute(evolved) return self._patch_stream_sources[key] def _plan_save_sweep( @@ -2234,14 +2236,14 @@ def peak_case_bytes(self) -> int: spatial = [int(extent) for extent in stage.transform_shape(self.group_src, self.name, source, attributes)] stage.write_stream_cache_attribute(attributes, source) peak = max(peak, channels * int(np.prod(spatial, dtype=np.int64))) - return peak * _CASE_ELEMENT_BYTES + return peak * CASE_ELEMENT_BYTES def _enforce_fallback_budget(self, fallback_budget_bytes: float | None) -> None: if fallback_budget_bytes is None: return # The plan's promise holds at run time too: a case whose sweep failed here (a refusal the # probe could not see) must not assemble a volume the budget cannot hold. - case_bytes = self.peak_case_bytes() * _FALLBACK_INFLIGHT_FACTOR + case_bytes = self.peak_case_bytes() * FALLBACK_INFLIGHT_FACTOR if case_bytes > fallback_budget_bytes: raise PatchError( f"Case '{self.name}' fell back to the whole-volume path at run time and its" @@ -2300,9 +2302,11 @@ def materialize_copies( if first is not None: shared_pending = [sweep for sweep in first.pending_sweeps if sweep.entry == self.name] if shared_pending: - # Each failure records its own reason; _sweep_failed reads them back. for sweep in shared_pending: - self._materialize_save(sweep) + # Chained here as on the per-case path: past a failure the next sweep reads a + # cache nobody wrote and overwrites the recorded reason with its own symptom. + if not self._materialize_save(sweep): + break self._invalidate_stream_plans() shared: list[tuple[int, _PendingSweep]] = [] @@ -2538,7 +2542,12 @@ def _stream_ready(self, a: int, apply_augmentations: bool = True) -> bool: if not source.pending_sweeps: return True for sweep in source.pending_sweeps: - self._materialize_save(sweep) + # Stop at the first failure. The sweeps are CHAINED -- each one's source is the previous + # one's destination -- so past a failure the next reads a cache nobody wrote, fails too, + # and overwrites the recorded reason with its own. _sweep_failure has to keep the cause, + # not the last symptom. + if not self._materialize_save(sweep): + break # Every copy replans: the pending plans pointed at caches that did not exist yet (or, after a # failure, never will -- _sweep_failed reroutes them to the whole-volume path). self._invalidate_stream_plans() @@ -2620,7 +2629,7 @@ def _materialize_save(self, sweep: _PendingSweep) -> bool: # not something a slab can carry an opinion about: refuse, and let the caller's # handler fall back to the whole volume, which reduces correctly. raise PatchError( - f"A stage of the chain writing '{sweep.group}/{self.name}' returned a rank-{block.ndim}" + f"A stage of the chain writing '{sweep.group}/{sweep.entry}' returned a rank-{block.ndim}" f" slab where the channel-first layout needs rank {len(spatial) + 1}" f" (C, {', '.join(str(e) for e in spatial)}).", "A transform that reduces the leading axis must keep it (`keepdim=True`), so a" @@ -2678,6 +2687,12 @@ def _sweep_rows(self, spatial: list[int], channels: int) -> int: landed block (plus the write buffer), so half the budget over that working set is the honest height. Below one row nothing fits; the row is the floor, and the fallback budget check is what refuses a case whose single row cannot fit. + + ``channels`` is the SOURCE's, and ``_SWEEP_ELEMENT_BYTES`` assumes four: the landed block's + channel count and dtype are not known until the first slab has been read, and the height has + to be fixed before that so every slab writes whole chunks. A chain that multiplies channels + (a one-hot, a field synthesis) is therefore under-counted here; the per-case fallback budget + check runs on real shapes and is the authority that refuses. """ cap = max(1, int(_SWEEP_SLAB_ROWS)) budget = self._sweep_budget_bytes diff --git a/konfai/data/reduction.py b/konfai/data/reduction.py index 148a7ddd..bbb85fa9 100644 --- a/konfai/data/reduction.py +++ b/konfai/data/reduction.py @@ -65,6 +65,13 @@ class Reduction(ABC): #: wrong ``False`` only costs memory. incremental: bool = False + #: Regions this operator allocates ON TOP of the ones it is handed, counted in buffers-worth. + #: ``0`` folds what it already holds; an operator that stacks its inputs into a new tensor, or + #: sorts a copy of them, is holding that many buffers again while it runs. The plan multiplies + #: this into the peak it sizes regions against, so leaving it at ``0`` for an operator that + #: copies is a plan promising a working set the run then exceeds. + working_multiple: float = 0.0 + @abstractmethod def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: raise NotImplementedError() @@ -158,9 +165,16 @@ class Median(Reduction): Not incremental, and cannot be: a median needs every case before it can name the middle one. Its working set is therefore every case at one region, which is what bounds the region size rather than the volume. + + **Not for label maps.** Averaging the middle pair means the result can be a value that was in no + input -- over labels 1 and 5 it is 3, a different structure -- and over exactly two cases it is + the mean, so the robustness the name promises is gone. Fold segmentations with :class:`Vote`. """ voxel_local = True + # ``torch.stack`` copies the buffer into a new tensor and ``torch.quantile`` sorts a copy of + # that: two buffers-worth live alongside the one already held, for the duration of the call. + working_multiple = 2.0 def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: if len(tensors) == 1: @@ -169,11 +183,39 @@ def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: return middle.to(_averaged_dtype(tensors[0].dtype)) +class Vote(Reduction): + """The label the most cases agree on, per voxel. The operator for folding SEGMENTATIONS. + + ``Mean`` and ``Median`` both answer with a value that was in no input: the median of labels 1 and + 5 is 3, which is a different structure. A label map made of invented labels is still a label map, + so nothing downstream reports it, and the dtype widens to float32 on top. This one picks and never + blends, and the result keeps the input's dtype. + + A tie goes to the SMALLEST label, so a cohort folds to the same volume on every run and on every + rank. That is arbitrary but it has to be *something*, and an arbitrary rule stated here beats a + stable-sort detail nobody can see. + + Not incremental: a majority needs every case before it can be counted. + """ + + voxel_local = True + # ``torch.mode`` sorts a copy of the stack it is handed, alongside the stack itself. + working_multiple = 2.0 + + def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: + if len(tensors) == 1: + return tensors[0] + return torch.mode(torch.stack(tensors, dim=0), dim=0).values.to(tensors[0].dtype) + + class Concat(Reduction): """Concatenate the cases along the channel dimension.""" # Cats along the channel axis, orthogonal to the spatial axes -- per-voxel, so region-local. voxel_local = True + # Nothing on top of the buffer: the concatenation IS the output region, and the plan charges that + # separately at this operator's own (wider) output width. + working_multiple = 0.0 def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: return torch.cat(tensors, dim=1) diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 4a320745..bfa2ed52 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -897,6 +897,49 @@ def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: dict[str, An return tensor.unsqueeze(self.dim) +# -------------------------------------------------------------------------------------------------- +# The rules every sampler below obeys. There are two gather strategies for one arithmetic -- per-axis +# maps where the coordinate is separable, eight flat corners where a displacement makes it not -- and +# the strategies differ for a measured reason. The RULES must not: written out at each site they +# drift, and the drift is silent because a resampled volume looks right either way. + + +def sampling_dtype(tensor: torch.Tensor) -> torch.dtype: + """The dtype to accumulate a weighted sum of ``tensor``'s voxels in. + + An integer input has no arithmetic of its own to interpolate with. A CPU half does, and it should + not be used: torch's CPU Half kernels are missing from older releases and lossy over a sum of + eight terms, at values a scanner actually produces. A CUDA half keeps its own -- every mode has a + Half kernel there, and upcasting a whole multi-class volume would double its memory for nothing. + """ + if not tensor.is_floating_point(): + return torch.float32 + if tensor.device.type == "cpu" and tensor.dtype in (torch.float16, torch.bfloat16): + return torch.float32 + return tensor.dtype + + +def nearest_index(coordinate: torch.Tensor) -> torch.Tensor: + """ITK's nearest: round half UP on the continuous source index. + + ``torch.round`` breaks a tie to the even index, and ``F.interpolate``'s nearest is + ``floor(o * scale)`` -- a statement about a size RATIO, which says nothing once the target grid + carries an origin of its own. On a label map either wrong rule still yields a label map. + """ + return torch.floor(coordinate + 0.5).to(torch.long) + + +def window_index(index: torch.Tensor, n_in: int, region_start: int, window: int) -> torch.Tensor: + """A global source index as an offset into the sub-region that was actually read. + + Clamped twice, and both matter: to the SOURCE first, so a tap past the volume reproduces the + border value rather than wrapping, and to the WINDOW second, so it stays inside the buffer on + hand. The second clamp is only ever load-bearing where the first already put the sample outside, + which the caller masks to fill -- or, for a halo'd read, where the declared bound was checked. + """ + return torch.clamp(torch.clamp(index, 0, n_in - 1) - region_start, 0, window - 1) + + class Resample(TransformInverse, ABC): def __init__(self, inverse: bool) -> None: super().__init__(inverse) @@ -914,17 +957,7 @@ def _resample(self, tensor: torch.Tensor, size: list[int]) -> torch.Tensor: else: mode = "trilinear" - # Interpolate in the tensor's own float dtype on CUDA. The model output is float16 and CUDA has - # Half kernels for every mode, so upcasting the whole (channels x volume) tensor to float32 would - # double the memory of a multi-class output resample for no argmax benefit. On the CPU, compute in - # float32: Half CPU kernels are missing from older torch releases. Integer inputs (uint8 labels) - # still need a float grid for interpolation. - if not tensor.is_floating_point() or ( - tensor.device.type == "cpu" and tensor.dtype in (torch.float16, torch.bfloat16) - ): - work = tensor.type(torch.float32) - else: - work = tensor + work = tensor.type(sampling_dtype(tensor)) # Return on the input's device (interpolate preserves it): a CPU input stays on the CPU, a # GPU-resident output volume stays on the GPU so the whole finalize runs where the volume is. return F.interpolate(work.unsqueeze(0), size=tuple(size), mode=mode).squeeze(0).type(tensor.dtype) @@ -979,8 +1012,22 @@ def stream_region_target( # Every patch derives its source coordinates from the same global scale (n_in / n_out, from the # truncated integer sizes F.interpolate itself uses), which is what makes the streamed patches # agree with the whole-volume call and with each other across a seam. + #: What this stage interpolates with, or ``None`` to read it off the dtype. A subclass taking an + #: ``interpolation`` argument assigns it here, and every sampler asks the one method below -- + #: a declaration honoured on one path and not another is worse than none, because the page that + #: tells a user to set it is then right about half the chains. + interpolation: str | None = None + def _stream_mode(self, tensor: torch.Tensor) -> str: - if tensor.dtype == torch.uint8: + """``nearest``, or the rank's linear name -- what a sampler asks before it blends anything. + + A dtype cannot settle this on its own: a CT is int16 and so is nothing else about it. The + heuristic therefore claims ``uint8`` and nothing more, and a stage exposing ``interpolation`` + answers for everything it cannot know. Getting it wrong is silent -- two blended labels give + a third that was in no input, in a volume that is still a label map. + """ + declared = self.interpolation or ("nearest" if tensor.dtype == torch.uint8 else "linear") + if declared == "nearest": return "nearest" return "bilinear" if len(tensor.shape) < 4 else "trilinear" @@ -1105,12 +1152,7 @@ def resample_region( # pure coordinate gather, so composing the axes changes no value). return sub_tensor[(slice(None), *torch.meshgrid(*indices, indexing="ij"))] - if not sub_tensor.is_floating_point() or ( - sub_tensor.device.type == "cpu" and sub_tensor.dtype in (torch.float16, torch.bfloat16) - ): - work = sub_tensor.type(torch.float32) - else: - work = sub_tensor + work = sub_tensor.type(sampling_dtype(sub_tensor)) taps: list[tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]] = [] for k in range(ndim): o = torch.arange(target_slices[k].start, target_slices[k].stop, device=dev, dtype=work.dtype) @@ -1182,22 +1224,14 @@ def _resample_offset_region( return torch.full(out_shape, outside, device=device, dtype=torch.float32).type(sub_tensor.dtype) def local(index: torch.Tensor, k: int) -> torch.Tensor: - """A global source index as an offset into the window that was read, kept in range.""" - return torch.clamp(torch.clamp(index, 0, n_in[k] - 1) - region_starts[k], 0, window[k] - 1) + return window_index(index, n_in[k], region_starts[k], window[k]) if self._stream_mode(sub_tensor) == "nearest": - # floor(c + 0.5) is ITK's nearest -- round half up -- not F.interpolate's floor(o * scale), - # which is a statement about a size ratio and says nothing once a grid has its own origin. - picks = [local(torch.floor(axis + 0.5).long(), k) for k, axis in enumerate(coordinates)] + picks = [local(nearest_index(axis), k) for k, axis in enumerate(coordinates)] gathered = sub_tensor[(slice(None), *torch.meshgrid(*picks, indexing="ij"))] out = gathered if gathered.is_floating_point() else gathered.type(torch.float32) else: - if not sub_tensor.is_floating_point() or ( - sub_tensor.device.type == "cpu" and sub_tensor.dtype in (torch.float16, torch.bfloat16) - ): - work = sub_tensor.type(torch.float32) - else: - work = sub_tensor + work = sub_tensor.type(sampling_dtype(sub_tensor)) taps = [] for k, axis in enumerate(coordinates): base = torch.floor(axis) @@ -1415,9 +1449,17 @@ def __init__( field_group: str | None = None, max_displacement: float | str = 0.0, fill: float = 0.0, + interpolation: str | None = None, inverse: bool = True, ) -> None: super().__init__(inverse) + if interpolation is not None and interpolation not in ("linear", "nearest"): + raise TransformError( + f"'ResampleToReference' has an unknown interpolation '{interpolation}'.", + "Use 'linear' for an image or 'nearest' for a label map. Left unset, uint8 is taken" + " for a label map and everything else is interpolated.", + ) + self.interpolation = interpolation if not entry or not str(entry).strip(): raise TransformError( "'ResampleToReference' needs an 'entry': the stored image whose grid to adopt.", @@ -1446,7 +1488,11 @@ def __init__( " this stage resamples onto the grid and nothing more.", ) self.displacement: _DisplacementSource | None = ( - _DisplacementSource("ResampleToReference", field, field_group, max_displacement) if declared else None + _DisplacementSource( + "ResampleToReference", field, field_group, max_displacement, group_keyword="field_group" + ) + if declared + else None ) self._grid: tuple[list[int], np.ndarray, np.ndarray, np.ndarray] | None = None # Each case's map, kept from where its own header was in hand. See _recorded(). @@ -1891,17 +1937,25 @@ def _sample_at( return torch.full(out_shape, self.fill_value, device=sub_tensor.device, dtype=torch.float32).type( sub_tensor.dtype ) - work = sub_tensor if sub_tensor.is_floating_point() else sub_tensor.type(torch.float32) + work = sub_tensor.type(sampling_dtype(sub_tensor)) flat_source = work.reshape(int(work.shape[0]), -1) + + def _tap(offsets: list[torch.Tensor]) -> torch.Tensor: + flat_index = torch.zeros(extent, dtype=torch.long, device=sub_tensor.device) + for axis, offset in enumerate(offsets): + index = window_index(offset, n_in[axis], region_starts[axis], window[axis]) + flat_index = flat_index * window[axis] + index + return flat_source.index_select(1, flat_index.reshape(-1)).reshape(out_shape) + + if self._stream_mode(sub_tensor) == "nearest": + picked = _tap([nearest_index(axis).expand(extent) for axis in coordinates]) + return picked.masked_fill(~inside.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) out = torch.zeros(out_shape, device=sub_tensor.device, dtype=work.dtype) for corner in itertools.product((0, 1), repeat=rank): - flat_index = torch.zeros(extent, dtype=torch.long, device=sub_tensor.device) weight = torch.ones(extent, device=sub_tensor.device, dtype=work.dtype) for axis, step in enumerate(corner): - index = torch.clamp(bases[axis].to(torch.long) + step, 0, n_in[axis] - 1) - region_starts[axis] - flat_index = flat_index * window[axis] + torch.clamp(index, 0, window[axis] - 1) weight = weight * (weights[axis] if step else 1 - weights[axis]).to(work.dtype) - out += flat_source.index_select(1, flat_index.reshape(-1)).reshape(out_shape) * weight + out += _tap([bases[axis].to(torch.long) + step for axis, step in enumerate(corner)]) * weight return out.masked_fill(~inside.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: @@ -2361,11 +2415,22 @@ class _DisplacementSource: exceeds that bound are the same questions, and were answered twice before this existed. ``owner`` is the stage's name, so a refusal reads as coming from the stage the user declared and - not from a helper they have never heard of. + not from a helper they have never heard of. ``group_keyword`` goes with it: the two owners spell + the field's group differently -- ``Warp`` calls it ``group`` because it has no other, and + ``ResampleToReference`` calls it ``field_group`` because ``group`` is already the reference's. A + remedy naming the wrong one sends the user to change the wrong argument. """ - def __init__(self, owner: str, field: str | None, group: str | None, max_displacement: float | str) -> None: + def __init__( + self, + owner: str, + field: str | None, + group: str | None, + max_displacement: float | str, + group_keyword: str = "group", + ) -> None: self.owner = owner + self.group_keyword = group_keyword # A root of its own, or none: with no ``field`` path the fields are a GROUP of the run's own # dataset_filenames, one entry per case — which is how a cohort registered in place stores # them, beside the volumes they were solved on. @@ -2377,7 +2442,7 @@ def __init__(self, owner: str, field: str | None, group: str | None, max_displac raise TransformError( f"'{owner}' has neither a 'field' path nor a group to find the fields in.", f"Name the store — {owner}: {{field: ./DVF:omezarr}} — or, for fields stored beside" - f" the cases, the group they are in: {owner}: {{field_group: DVF}}.", + f" the cases, the group they are in: {owner}: {{{group_keyword}: DVF}}.", ) self.group = group #: The run's own roots, handed over by the owner; only consulted when there is no path. @@ -2455,7 +2520,7 @@ def group_for(self, name: str | None) -> str: if self.dataset is None: # unreachable: a source with no path was given a group to use raise TransformError( f"'{self.owner}' has no field store of its own and no group to look for one in.", - f"Name the group the fields are in: {self.owner}: {{field_group: DVF}}.", + f"Name the group the fields are in: {self.owner}: {{{self.group_keyword}: DVF}}.", ) groups = [str(group) for group in self.dataset.get_group()] if len(groups) == 1: @@ -2463,7 +2528,7 @@ def group_for(self, name: str | None) -> str: where = f"the field for case '{name}'" if name is not None else "the fields" raise TransformError( f"'{self.owner}' cannot tell which group of '{self.dataset.filename}' holds {where}: it has {len(groups)}.", - f"Name it: {self.owner}: {{field: ./DVF:omezarr, group: DVF}}.", + f"Name it: {self.owner}: {{field: ./DVF:omezarr, {self.group_keyword}: DVF}}.", ) def _root_for(self, name: str | None) -> Dataset: @@ -2622,7 +2687,13 @@ def _sample(self, tensor: torch.Tensor, field: torch.Tensor, spacing: list[float the wrong way. """ extent = list(tensor.shape[1:]) - axes = torch.meshgrid(*[torch.arange(size, dtype=torch.float32) for size in extent], indexing="ij") + # Grid and field follow the VOLUME's device: a field is read from disk onto the CPU, the + # volume may be GPU-resident, and grid_sample takes both from one device. + device = tensor.device + field = field.to(device) + axes = torch.meshgrid( + *[torch.arange(size, dtype=torch.float32, device=device) for size in extent], indexing="ij" + ) sample = [] for axis in range(len(extent)): # field component for array axis `axis` (z,y,x) is the reversed one (x,y,z) @@ -2660,6 +2731,10 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) "Use a source whose geometry is readable (mha, nii, h5 or omezarr written by KonfAI).", ) field = self.displacement.read(name, None, len(tensor.shape) - 1) + # The declared bound is checked against every field read, on this path as on the streamed + # one: it is what sizes the halo, and a bound smaller than the field streams a region whose + # edge is missing. + self.displacement.check_bound(field, name) return self._sample(tensor, field, spacing) diff --git a/konfai/main.py b/konfai/main.py index d24f9a24..7f32ff63 100644 --- a/konfai/main.py +++ b/konfai/main.py @@ -37,12 +37,15 @@ def _run(parser: argparse.ArgumentParser) -> None: This function: 1) defines common arguments used by TRAIN / RESUME / PREDICTION / EVALUATION - (config file, overwrite, device selection, quiet, tensorboard) + (config file, overwrite, device selection, quiet, tensorboard) -- TRANSFORM + declares its own set, since it has no TOP-LEVEL model and so nothing to + write scalars about (a chain may still embed a `KonfAIInference` stage) 2) defines subcommands and their command-specific arguments 3) parses CLI args and dispatches to the correct implementation: - `konfai.trainer.train` for TRAIN and RESUME - `konfai.predictor.predict` for PREDICTION - `konfai.evaluator.evaluate` for EVALUATION + - `konfai.transformer.transform` for TRANSFORM Device selection ---------------- @@ -308,11 +311,13 @@ def main(): - RESUME - PREDICTION - EVALUATION + - TRANSFORM Notes ----- The actual execution logic is implemented in `konfai.trainer.train`, - `konfai.predictor.predict`, and `konfai.evaluator.evaluate`. + `konfai.predictor.predict`, `konfai.evaluator.evaluate`, and + `konfai.transformer.transform` -- the one command with no top-level model. """ parser = argparse.ArgumentParser( prog="konfAI", description="KonfAI - Deep learning framework for Medical AI Models", allow_abbrev=False diff --git a/konfai/predictor.py b/konfai/predictor.py index 30d20aff..3ee127c8 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -1703,10 +1703,11 @@ class ModelComposite(Network): Args: model (Network): The base network to replicate. nb_models (int): Number of copies of the model to create. - combine (Reduction): The reduction method used to combine outputs from all model replicas. + combine (konfai.data.reduction.Reduction): The reduction method used to combine outputs from + all model replicas. Attributes: - combine (Reduction): The reduction method used during forward inference. + combine (konfai.data.reduction.Reduction): The reduction used during forward inference. """ def __init__(self, model: Network, combine: Reduction): @@ -1866,7 +1867,7 @@ class Predictor(DistributedObject): Attributes: model (Network): The neural network model to use for prediction. - dataset (DataPrediction): Dataset manager for prediction data. + dataset (konfai.data.data_manager.DataPrediction): Dataset manager for prediction data. combine_classpath (str): Path to the reduction strategy (e.g., "Mean"). autocast (bool): Whether to enable AMP inference. outputs_dataset (dict[str, OutputDataset]): Mapping from layer names to output writers. diff --git a/konfai/transformer.py b/konfai/transformer.py index 0157ad83..df2fe2f7 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -40,7 +40,7 @@ from konfai import config_file, transforms_directory from konfai.data.case_reduction import CaseReduction, split_chain from konfai.data.data_manager import DataTransform, _format_gib, _node_local_ranks -from konfai.data.patching import _CASE_ELEMENT_BYTES, _FALLBACK_INFLIGHT_FACTOR, DatasetManager, _save_destination +from konfai.data.patching import CASE_ELEMENT_BYTES, FALLBACK_INFLIGHT_FACTOR, DatasetManager, save_destination from konfai.data.transform import Reduce, Save, Transform, split_expand from konfai.utils.config import apply_config, config from konfai.utils.dataset import Attribute, Dataset @@ -83,7 +83,7 @@ def working_set_bytes(self) -> int: A reduction's ``case_bytes`` is already its resident regions (it holds one region per case, not one volume), where a whole-volume fallback holds the case plus one in-flight copy. """ - return self.case_bytes if self.reduced else self.case_bytes * _FALLBACK_INFLIGHT_FACTOR + return self.case_bytes if self.reduced else self.case_bytes * FALLBACK_INFLIGHT_FACTOR @dataclass @@ -128,7 +128,7 @@ def report(self) -> str: lines = [ f"[Transformer] plan over {self.world_size} rank(s) | per-rank budget" f" {_format_gib(self.budget_bytes)} ({self.budget_desc}) | fallback working set = case" - f" x {_CASE_ELEMENT_BYTES} B x {_FALLBACK_INFLIGHT_FACTOR} (in-flight copy), headers-only" + f" x {CASE_ELEMENT_BYTES} B x {FALLBACK_INFLIGHT_FACTOR} (in-flight copy), headers-only" f" estimate | output dtype/channels assumed {self.dtype_hypothesis} until the first slab" ] for group_src, dropped in sorted(self.dropped_cases.items()): @@ -248,7 +248,7 @@ def _terminal_destination(manager: DatasetManager) -> tuple[Dataset, str]: f"The chain writing '{manager.group_dest}' does not end with a Write, so it has no destination.", "End every chain with Write: {dataset: [:format]}.", ) - return _save_destination(terminal, manager.dataset, manager.group_dest) + return save_destination(terminal, manager.dataset, manager.group_dest) def _reduction(self, group_dest: str, managers: list[DatasetManager]) -> CaseReduction | None: """The reduction this chain declares, or ``None`` when it is an ordinary per-case chain. @@ -293,7 +293,7 @@ def _build_reduction(self, group_dest: str, managers: list[DatasetManager]) -> C "A reduction writes one entry, so the chain carrying it must say where:" " Write: {dataset: [:format]}.", ) - destination, group = _save_destination(terminal, managers[0].dataset, group_dest) + destination, group = save_destination(terminal, managers[0].dataset, group_dest) reduction = CaseReduction( managers=cases, reduce=reduce, @@ -334,7 +334,7 @@ def _probe_write_destinations( channels = int(manager.base_shape[0]) dtype = self._dtype_hypothesis(manager) for transform, spatial, attributes in manager.write_targets(a): - destination, group = _save_destination(transform, manager.dataset, manager.group_dest) + destination, group = save_destination(transform, manager.dataset, manager.group_dest) key = (str(destination.filename), group) if key in probed: continue @@ -422,7 +422,10 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor group_src = self._group_src_of(group_dest) reduction = self._reduction(group_dest, managers) if reduction is not None: - reduction_dtype = np.dtype("float32") + # Read off the chain, not assumed: a pointwise cast is allowed after the Reduce, and + # the probe below opens the destination with this dtype. A constant here would test + # a write the run never makes -- and mha refuses on dtype. + reduction_dtype = self._dtype_hypothesis(managers[0]) planned_dtypes.add(str(reduction_dtype)) reduction_plan = reduction.plan() skipped = not overwrite and reduction.destination.is_dataset_exist( @@ -593,7 +596,7 @@ def setup(self, world_size: int): for transform in managers[0].transforms if managers else []: if not isinstance(transform, Save): continue - destination, _group = _save_destination(transform, managers[0].dataset, managers[0].group_dest) + destination, _group = save_destination(transform, managers[0].dataset, managers[0].group_dest) # The question here is NOT concurrent_write_safe(): that one asks whether two # entries of one shared store may be written at once, and answers no for omezarr # -- which would refuse the very destination this workflow recommends. Ranks diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 04d278c5..950190e2 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -197,16 +197,35 @@ def _attribute_text(value: Any) -> str: return str(value).replace("\n", "") -def _store_chunks(shape: list[int], region_shape: list[int] | None) -> tuple[int, ...] | None: +def _store_chunks(shape: list[int], region_shape: list[int] | None, dtype: Any) -> tuple[int, ...] | None: """Chunks a store should use, given the region shape its writer declared. A region write that straddles a chunk becomes a read-modify-write of the whole chunk, so the - honest chunking is the writer's own region, clamped to the array. ``None`` when the writer - declared nothing: the store keeps its own default rather than a guess made here. + writer's own region is the honest starting point. Taking it verbatim is not: a slab sweep + declares the whole trailing plane, which at 2048x2048 float32 is a gigabyte in one chunk -- past + what zarr will hold in a single buffer at 4096x4096 -- and every later partial read pays it, + including readers that are not KonfAI. + + So an axis the region covers END TO END is tiled, and an axis where the region is a strict + sub-range keeps the writer's size. Tiling a fully-covered axis cannot split a region write, since + every chunk along it falls inside one region either way; splitting a partial axis is exactly the + read-modify-write this exists to avoid. Innermost axes go first, so the chunk stays long on the + axis the writer actually advances along. + + ``None`` when the writer declared nothing: the store keeps its own default rather than a guess. """ + from konfai.utils.ome_zarr import CHUNK_SPATIAL_TILE, CHUNK_TARGET_BYTES + if region_shape is None or len(region_shape) != len(shape): return None - return tuple(max(1, min(int(region), int(extent))) for region, extent in zip(region_shape, shape, strict=True)) + chunk = [max(1, min(int(region), int(extent))) for region, extent in zip(region_shape, shape, strict=True)] + itemsize = max(1, np.dtype(dtype).itemsize) + covered = [axis for axis, extent in enumerate(shape) if chunk[axis] >= int(extent)] + for axis in reversed(covered): + if int(np.prod(chunk, dtype=np.int64)) * itemsize <= CHUNK_TARGET_BYTES: + break + chunk[axis] = min(chunk[axis], CHUNK_SPATIAL_TILE) + return tuple(chunk) class Attribute(dict[str, Any]): @@ -1636,10 +1655,11 @@ def open_data_stream( origin=attributes.get_np_array("Origin") if "Origin" in attributes else None, attributes=dict(attributes), displacement_field=DISPLACEMENT_FIELD_ATTRIBUTE in attributes, - # Chunk on what the writer says it will write. Guessing costs a read-modify-write on - # every region whose extent straddles a chunk -- measured 1.8x on a slab sweep, paid - # on every byte, and invisible because the bytes are correct either way. - chunks=_store_chunks(shape, region_shape), + # Chunked against what the writer says it will write, capped to something a reader + # can open. Guessing the writer's access pattern costs a read-modify-write on every + # region whose extent straddles a chunk -- measured 1.8x on a slab sweep, paid on + # every byte, and invisible because the bytes are correct either way. + chunks=_store_chunks(shape, region_shape, dtype), ) # The pyramid cannot be created up front -- no level exists until the last region lands -- # so the stream derives it at finalize, on the TEMPORARY store, before the rename. That diff --git a/konfai/utils/ome_zarr.py b/konfai/utils/ome_zarr.py index f3018fd8..2a9d5c95 100644 --- a/konfai/utils/ome_zarr.py +++ b/konfai/utils/ome_zarr.py @@ -95,6 +95,13 @@ def _native_byteorder(array: np.ndarray) -> np.ndarray: _RFC5_VERSION = "0.6" _DEFAULT_VERSION = "0.4" +#: How a chunk is sized, whether the store is created from a shape alone or from the region shape a +#: streamed writer declares (:func:`konfai.utils.dataset._store_chunks`). One rule, two callers: a +#: chunk is the unit a reader decompresses to reach one voxel, so an oversized one is paid by every +#: partial read forever, and by any consumer that is not KonfAI. +CHUNK_SPATIAL_TILE = 128 +CHUNK_TARGET_BYTES = 32 << 20 + def _zarr_v3_available() -> bool: """Whether the installed zarr can write a v3 store, which NGFF >= 0.5 (RFC-5) requires. @@ -419,8 +426,9 @@ def displacement_bound(data: np.ndarray) -> list[float]: field = np.asarray(data) if field.ndim < 2: return [] - flat = field.reshape(field.shape[0], -1) - return [float(np.abs(flat[component]).max(initial=np.float32(0.0))) for component in range(flat.shape[0])] + # Cast before the max, not after: ``initial=np.float32(0.0)`` does not downcast a float64 input. + flat = np.abs(field.reshape(field.shape[0], -1)).astype(np.float32, copy=False) + return [float(flat[component].max(initial=np.float32(0.0))) for component in range(flat.shape[0])] def update_konfai_attributes(store_path: str | Path, extra: dict[str, Any]) -> None: @@ -435,6 +443,10 @@ def update_konfai_attributes(store_path: str | Path, extra: dict[str, Any]) -> N group = zarr.open_group(str(store_path), mode="r+") sidecar = dict(dict(group.attrs).get(_KONFAI_ATTR_KEY, {}).get("attributes", {})) sidecar.update(extra) + # Writing through an `r+` group updates the consolidated copy with it, so a sidecar landing here + # is readable by a consolidated reader without a second consolidation pass. Measured, because the + # opposite is the plausible assumption: a foreign `zarr.open_group(mode="r")` reads back a bound + # written this way on both sides of this line. group.attrs[_KONFAI_ATTR_KEY] = {"attributes": sidecar} clear_ome_zarr_cache() @@ -505,10 +517,10 @@ def create_ome_zarr_store( translation = {"c": 0.0, **dict(zip(spatial_axes, translation_values, strict=True))} if chunks is None: - spatial_chunks = [min(extent, 128) for extent in shape[1:]] - # Keep one chunk around 32 MiB: full 128-wide spatial tiles, channels split to fit the budget. + spatial_chunks = [min(extent, CHUNK_SPATIAL_TILE) for extent in shape[1:]] + # Keep one chunk near the target: full spatial tiles, channels split to fit the budget. tile_bytes = int(np.prod(spatial_chunks, dtype=np.int64)) * np.dtype(dtype).itemsize - chunks = [min(shape[0], max(1, (32 << 20) // max(1, tile_bytes))), *spatial_chunks] + chunks = [min(shape[0], max(1, CHUNK_TARGET_BYTES // max(1, tile_bytes))), *spatial_chunks] chunks = tuple(chunks) stand_in = dask.array.zeros((shape[0], *(1,) * len(spatial_axes)), dtype=np.dtype(dtype)) @@ -589,6 +601,9 @@ def append_ome_zarr_levels( base, scale_factors=[int(f) for f in scale_factors], method=_downsample_method(downsample_method), + # Level 0 is rewritten here, so the base store's own chunking has to be carried across: + # ngff-zarr otherwise re-chunks it on its default. + chunks=tuple(int(size) for size in base.data.chunksize), cache=False, ) version = _DEFAULT_VERSION diff --git a/tests/unit/test_case_reduction.py b/tests/unit/test_case_reduction.py index b40b886b..96fa216a 100644 --- a/tests/unit/test_case_reduction.py +++ b/tests/unit/test_case_reduction.py @@ -20,14 +20,15 @@ never assembled, that a cases which does not agree on its grid is refused before anything is read, and that a chain continues after the reduction.""" +import itertools from pathlib import Path import numpy as np import pytest import torch -from konfai.data.case_reduction import CaseReduction, resolve_operator, split_chain +from konfai.data.case_reduction import CaseReduction, ReductionPlan, resolve_operator, split_chain from konfai.data.patching import DatasetManager -from konfai.data.reduction import Mean, Median, Reduction +from konfai.data.reduction import Concat, Mean, Median, Reduction, Vote from konfai.data.transform import Clip, Dilate, Normalize, Reduce, TensorCast, Transform from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ReductionError, TransformError @@ -134,13 +135,27 @@ def test_averaging_keeps_a_floating_dtype_and_widens_an_integer_one(operator: Re def test_a_non_incremental_operator_holds_the_whole_cohort_per_region(tmp_path: Path) -> None: + """One region per case, plus the output's, plus what the operator allocates over its buffer. + + ``Median`` stacks the buffer into a new tensor and sorts a copy of that, so the buffer alone + under-states its peak threefold — and it is the operator a bare ``Reduce`` gets. + """ engine, _destination, _volumes = _run(tmp_path, [], Reduce(operator="Median", output="t"), []) plan = engine.plan() assert plan.incremental is False - assert plan.resident_regions == CASES + 1 + assert plan.resident_regions == CASES + 1 + Median.working_multiple * CASES assert "resident region" in plan.describe() and "CASE_000" in plan.describe() +def test_an_operator_that_folds_in_place_is_budgeted_for_what_it_holds(tmp_path: Path) -> None: + """The multiplier is the operator's to declare: ``Mean`` accumulates into one running region.""" + engine, _destination, _volumes = _run(tmp_path, [], Reduce(operator="Mean", output="t"), []) + plan = engine.plan() + assert plan.incremental is True + assert Mean.working_multiple == 0.0 + assert plan.resident_regions == 2 + + def test_per_member_stages_run_on_each_member_separately(tmp_path: Path) -> None: """The whole point of cases being cases: a GLOBAL_STAT before the reduction seeds from the case's OWN volume, so a cases of different dynamics normalises per specimen.""" @@ -376,3 +391,92 @@ def test_an_operator_shadowing_a_reduce_key_is_refused() -> None: rather than silently handing the stage's value to the operator.""" with pytest.raises(ReductionError, match="reads for itself"): resolve_operator(Reduce(operator=f"{__name__}:_Shadowing", output="t")) + + +def test_vote_picks_a_label_where_median_would_invent_one() -> None: + """The reason Vote exists. Two label maps have no middle value that is a label. + + Median averages the middle pair, so folding structures 1 and 5 answers 3 -- a third structure, + in a volume that is still a valid label map, which is why nothing downstream reports it. + """ + labels = [torch.full((1, 1, 2, 2), value, dtype=torch.uint8) for value in (1, 5)] + + assert float(Median()(labels).flatten()[0]) == 3.0 + assert Median()(labels).dtype is torch.float32 + + voted = Vote()(labels) + assert float(voted.flatten()[0]) == 1.0, "a tie goes to the smallest label" + assert voted.dtype is torch.uint8, "a vote picks, so the label dtype survives" + + +def test_vote_takes_the_label_the_majority_agrees_on() -> None: + labels = [torch.full((1, 1, 2, 2), value, dtype=torch.uint8) for value in (4, 7, 7, 2, 7)] + assert set(np.unique(Vote()(labels).numpy()).tolist()) == {7} + + +def test_vote_answers_per_voxel_not_per_volume() -> None: + """Each voxel is its own ballot -- a majority somewhere else must not carry it.""" + cases = [ + torch.tensor([[[[1, 2]]]], dtype=torch.uint8), + torch.tensor([[[[1, 3]]]], dtype=torch.uint8), + torch.tensor([[[[9, 3]]]], dtype=torch.uint8), + ] + np.testing.assert_array_equal(Vote()(cases).numpy(), np.array([[[[1, 3]]]], dtype=np.uint8)) + + +def test_a_single_case_is_its_own_vote() -> None: + only = torch.full((1, 1, 2, 2), 6, dtype=torch.uint8) + assert Vote()([only]).dtype is torch.uint8 + np.testing.assert_array_equal(Vote()([only]).numpy(), only.numpy()) + + +@pytest.mark.parametrize( + "operator,output_channels,member_regions,why", + [ + (Mean(), 1, 2, "one running accumulator and the region coming into it"), + (Median(), 1, 3 * CASES + 1, "the cohort, the stack it is copied into, the sort, the output"), + (Vote(), 1, 3 * CASES + 1, "same shape of work: a mode sorts a copy of the stack too"), + (Concat(), CASES, 2 * CASES, "the cohort, and the concatenation that IS the output"), + ], + ids=["mean", "median", "vote", "concat"], +) +def test_the_peak_is_charged_at_each_side_s_own_width( + operator: Reduction, output_channels: int, member_regions: int, why: str +) -> None: + """Members are measured at THEIR channel count, the output at its own. + + Only ``Concat`` tells the two apart -- it writes ``N x C`` where each member holds ``C`` -- and + charging the cohort at the output's width over-states its peak by the cohort's size, which either + shrinks the slab for nothing or refuses a reduction that fits. + """ + spatial, rows, source_channels = [8, 100, 100], 4, 1 + plan = ReductionPlan( + output="t", + cases=[f"CASE_{i:03d}" for i in range(CASES)], + spatial=spatial, + channels=output_channels, + source_channels=source_channels, + slab_rows=rows, + incremental=operator.incremental, + stat_pass=False, + working_multiple=operator.working_multiple, + ) + member = rows * spatial[1] * spatial[2] * source_channels * 4 + assert plan.peak_bytes == member_regions * member, why + + +def test_a_vote_tie_goes_to_the_smallest_label_whatever_the_cohort_order() -> None: + """The reproducibility half of Vote's contract, which its docstring promises out loud. + + A cohort is folded by whichever rank owns it and in whatever order the manager list came out, so + a tie broken by position would write a different template on a rerun and nothing about the volume + would look wrong. ``torch.mode`` documents that the smallest of the most frequent values wins; + this pins that the whole way through, because the promise is ours and not torch's to keep. + """ + for order in itertools.permutations((7, 2, 9)): + cohort = [torch.full((1, 1, 2, 2), value, dtype=torch.uint8) for value in order] + assert int(Vote()(cohort).flatten()[0]) == 2, f"order {order} broke the tie somewhere else" + + # A majority still beats the smallest label: the tie rule is a tie-breaker, not a preference. + counted = [torch.full((1, 1, 2, 2), value, dtype=torch.uint8) for value in (9, 2, 9)] + assert int(Vote()(counted).flatten()[0]) == 9 diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py index 6abf924a..64c82144 100644 --- a/tests/unit/test_dataset.py +++ b/tests/unit/test_dataset.py @@ -496,7 +496,7 @@ def test_get_infos_reads_only_the_header_for_a_mismatched_extension(tmp_path: Pa def test_a_group_written_through_another_dataset_object_is_seen(tmp_path: Path) -> None: """A group can be produced through one Dataset and read through another over the same folder: a - ``Save`` builds its own (data/patching.py:_save_destination) while the reader keeps the DataManager's. + ``Save`` builds its own (data/patching.py:save_destination) while the reader keeps the DataManager's. Membership answered from the reader's memoised listing froze at its first lookup, so every case written after it read as absent -- ImpactSynth masks its own output that way and raised ``NameError: Mask : MASK/P002 not found`` from the third case of a batch on.""" diff --git a/tests/unit/test_ome_zarr_data_surface.py b/tests/unit/test_ome_zarr_data_surface.py index a1479dd4..62d53d51 100644 --- a/tests/unit/test_ome_zarr_data_surface.py +++ b/tests/unit/test_ome_zarr_data_surface.py @@ -34,6 +34,7 @@ read_ome_zarr_data_slice, write_ome_zarr, ) +from konfai.utils.dataset import _store_chunks from konfai.utils.errors import DatasetManagerError @@ -158,3 +159,28 @@ def test_append_levels_without_factors_is_a_no_op(tmp_path: Path) -> None: write_ome_zarr(store, _volume(), spacing=[1.0] * 3) append_ome_zarr_levels(store, []) assert get_ome_zarr_info(store)["n_levels"] == 1 + + +def test_a_chunk_stays_openable_whatever_plane_the_writer_declares() -> None: + """A slab sweep declares the whole trailing plane, and that is not a chunk shape. + + At 2048x2048 float32 the declared region is a gigabyte, and at 4096x4096 it is past what zarr + holds in one buffer. Every axis the region covers end to end can be tiled without splitting a + region write, so the axis the writer advances along keeps its declared height and the rest are + cut down until a chunk is a size a reader can decompress to reach one voxel. + """ + rows = 64 + for spatial in ([400, 512, 512], [400, 1024, 1024], [400, 2048, 2048], [400, 4096, 4096]): + for channels in (1, 3): + shape = [channels, *spatial] + chunk = _store_chunks(shape, [channels, rows, *spatial[1:]], np.float32) + assert chunk is not None + megabytes = int(np.prod(chunk, dtype=np.int64)) * 4 / (1 << 20) + assert megabytes <= 32.0, f"{shape} -> {chunk} is {megabytes:.0f} MiB" + # The sweep axis is the one the writer advances along: shrinking it below the declared + # slab is the read-modify-write this sizing exists to avoid. + assert chunk[1] == rows + + +def test_a_writer_that_declares_nothing_leaves_the_store_its_own_default() -> None: + assert _store_chunks([1, 8, 8, 8], None, np.float32) is None diff --git a/tests/unit/test_resample_sampler_rules.py b/tests/unit/test_resample_sampler_rules.py new file mode 100644 index 00000000..71c0c33e --- /dev/null +++ b/tests/unit/test_resample_sampler_rules.py @@ -0,0 +1,175 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""The rules every sampler in this package obeys, pinned separately from any one sampler. + +There are two gather strategies for one arithmetic. ``Resample._resample_offset_region`` maps each +axis independently, so no coordinate volume is built and one ``index_select`` per axis does the work; +``ResampleToReference._sample_at`` cannot, because a displacement is not separable, so it holds a +coordinate per voxel and gathers eight corners flat. Same rules, different loops, and the loops are +different for a measured reason. + +Rules kept apart from loops is only true while something checks it. These tests are that: they assert +the RULES -- the inside interval, the tap clamp, round-half-up, the working dtype, the fill -- against +SimpleITK and against each other, so a change to one gather cannot quietly stop matching the other. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from konfai.data.transform import Resample + + +class _Sampler(Resample): + """A bare handle on the sampler: these rules belong to `Resample`, not to any stage using it.""" + + def __call__(self, name, tensor, cache_attribute): # pragma: no cover - not the surface tested + raise NotImplementedError + + def write_stream_cache_attribute(self, cache_attribute, source_spatial_shape) -> None: + raise NotImplementedError + + def transform_shape(self, shape, cache_attribute): # pragma: no cover + raise NotImplementedError + + def inverse(self, name, tensor, cache_attribute): # pragma: no cover + raise NotImplementedError + + def patch_locality(self, cache_attribute): # pragma: no cover + raise NotImplementedError + + +_SOURCE = (12, 14, 16) +_SCALES = [1.31, 1.17, 1.23] +_OFFSETS = [0.4, -0.3, 0.2] +_TARGET = (slice(0, 8), slice(0, 9), slice(0, 10)) + + +def _sampler(fill: float = 0.0) -> _Sampler: + sampler = _Sampler(inverse=False) + sampler.fill_value = fill + return sampler + + +def _volume(offset: float = 0.0) -> np.ndarray: + rng = np.random.default_rng(3) + return (rng.random((1, *_SOURCE)) * 400 + offset).astype(np.float32) + + +def _offset_region(tensor: torch.Tensor, fill: float = 0.0, **overrides) -> torch.Tensor: + arguments = { + "target_slices": _TARGET, + "region_starts": [0, 0, 0], + "scales": _SCALES, + "n_in": list(_SOURCE), + "offsets": _OFFSETS, + } + arguments.update(overrides) + return _sampler(fill)._resample_offset_region(tensor, **arguments) # type: ignore[arg-type] + + +def test_the_working_dtype_rule_holds_for_a_cpu_half_volume() -> None: + """A CPU half is accumulated in float32; the eight-corner sum is not done in half. + + torch's CPU half arithmetic is both slow and lossy over a sum of eight terms, and the values a + microscope or a scanner produces sit exactly where float16 spacing is 2. The threshold separates + the two: accumulating in half drifts more than twice as far from the float32 answer. + """ + volume = _volume(offset=2050.0) # 2050..2450, entirely above 2048, where float16 spacing is 2 + guarded = _offset_region(torch.from_numpy(volume).half()).float() + reference = _offset_region(torch.from_numpy(volume)) + + assert guarded.dtype is torch.float32 + drift = float((guarded - reference).abs().max()) + assert drift < 2.0, f"a half accumulation drifts ~4.3 on this fixture; got {drift}" + + +@pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, torch.float16, torch.float32, torch.float64]) +def test_a_volume_comes_back_as_the_dtype_it_went_in_as(dtype: torch.dtype) -> None: + """The sampler computes in whatever it must and casts back once. A store's dtype is the store's.""" + volume = torch.from_numpy((_volume() % 120).astype(np.float32)).to(dtype) + assert _offset_region(volume).dtype is dtype + + +def test_nearest_is_itk_round_half_up_and_not_a_size_ratio() -> None: + """``floor(c + 0.5)``, which is a statement about a coordinate. + + ``F.interpolate``'s nearest is ``floor(o * scale)``, a statement about a size RATIO -- it says + nothing once the target grid carries an origin of its own, which is the whole point of an offset + map. A label map is still a label map under either rule, so only this catches it. + """ + labels = torch.arange(int(np.prod(_SOURCE)), dtype=torch.uint8).reshape(1, *_SOURCE) % 7 + got = _offset_region(labels).numpy()[0] + + expected = np.empty_like(got) + source = labels.numpy()[0] + for z in range(got.shape[0]): + for y in range(got.shape[1]): + for x in range(got.shape[2]): + index = [ + int(np.floor(_SCALES[axis] * position + _OFFSETS[axis] + 0.5)) + for axis, position in enumerate((z, y, x)) + ] + expected[z, y, x] = source[tuple(np.clip(index, 0, np.array(_SOURCE) - 1))] + np.testing.assert_array_equal(got, expected) + + +def test_inside_is_the_half_open_half_voxel_rim() -> None: + """A sample is inside while its source index is in ``[-0.5, n - 0.5)`` -- SimpleITK's interval. + + The rim beyond the outermost voxel CENTRES is inside and reproduces the border value; a hair past + it is fill. Getting this wrong shows as a one-voxel frame, which reads as anatomy. + """ + volume = torch.full((1, 4, 4, 4), 5.0) + fill = -99.0 + + # A target of one voxel per axis, placed by the offset alone. + def at(offset: float) -> float: + one = (slice(0, 1), slice(0, 1), slice(0, 1)) + got = _offset_region( + volume, fill=fill, target_slices=one, scales=[1.0, 1.0, 1.0], n_in=[4, 4, 4], offsets=[offset] * 3 + ) + return float(got.flatten()[0]) + + assert at(-0.5) == 5.0, "the open end of the rim is inside" + assert at(-0.5 - 1e-3) == fill, "a hair before it is not" + assert at(3.5 - 1e-3) == 5.0, "just short of n - 0.5 is inside" + assert at(3.5) == fill, "n - 0.5 itself is outside: the interval is half open" + + +def test_the_separable_sampler_matches_simpleitk() -> None: + """The independent check. Written against SimpleITK because that is what the arithmetic claims. + + The oracle is skipped here and not at module scope: the rules above -- the working dtype, the + inside interval, round-half-up -- are checkable without it, and a net that evaporates when an + optional dependency is missing is the failure mode this file exists to prevent. + """ + sitk = pytest.importorskip("SimpleITK") + volume = _volume() + got = _offset_region(torch.from_numpy(volume)).numpy()[0] + + image = sitk.GetImageFromArray(volume[0]) + image.SetSpacing((1.0, 1.0, 1.0)) + image.SetOrigin((0.0, 0.0, 0.0)) + grid = sitk.Image(*reversed([sl.stop - sl.start for sl in _TARGET]), sitk.sitkFloat32) + # sitk takes geometry in (x, y, z) where the arrays above are (z, y, x). + grid.SetSpacing(tuple(reversed(_SCALES))) + grid.SetOrigin(tuple(reversed(_OFFSETS))) + want = sitk.GetArrayFromImage(sitk.Resample(image, grid, sitk.Transform(), sitk.sitkLinear, 0.0)) + + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-4) diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 47b99f7f..969cd7fa 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -532,12 +532,18 @@ def _warping(images: Dataset, fields: Dataset, **kwargs: object) -> ResampleToRe return stage -def _simpleitk_warp(volume: np.ndarray, field: np.ndarray | None = None) -> np.ndarray: +def _simpleitk_warp( + volume: np.ndarray, + field: np.ndarray | None = None, + interpolator: int = sitk.sitkLinear, + pixel: int = sitk.sitkFloat32, + fill: float = _FILL, +) -> np.ndarray: """``sitk.Resample(image, grid, DisplacementFieldTransform(field))`` — the one-pass authority.""" image = sitk.GetImageFromArray(volume[0]) image.SetOrigin(_SOURCE_ORIGIN) image.SetSpacing(_SOURCE_SPACING) - grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), sitk.sitkFloat32) + grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), pixel) grid.SetOrigin(_REFERENCE_ORIGIN) grid.SetSpacing(_REFERENCE_SPACING) transform: sitk.Transform = sitk.Transform() @@ -547,7 +553,71 @@ def _simpleitk_warp(volume: np.ndarray, field: np.ndarray | None = None) -> np.n vector.SetOrigin(_FIELD_ORIGIN) vector.SetSpacing(_FIELD_SPACING) transform = sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)) - return sitk.GetArrayFromImage(sitk.Resample(image, grid, transform, sitk.sitkLinear, _FILL)) + return sitk.GetArrayFromImage(sitk.Resample(image, grid, transform, interpolator, fill)) + + +def _spaced_labels() -> np.ndarray: + """Two labels far apart, so any value between them can only have come from blending.""" + index = np.arange(int(np.prod(_SOURCE_SPATIAL))).reshape(_SOURCE_SPATIAL) + return np.where((index // 3) % 2 == 0, 0, 100).astype(np.uint8)[None] + + +def test_a_label_map_warped_through_a_field_takes_the_nearest_voxel(tmp_path: Path) -> None: + """The warped path has a sampler of its own, and it has to pick a label rather than blend two. + + Blending is silent here: two labels average into a third that was never in the source, the dtype + is unchanged, and the result is still a label map -- so nothing downstream reports it. The + fixture holds only 0 and 100 so that any other value can only have come from an interpolation. + """ + images = Dataset(tmp_path / "Images", "h5") + labels = _spaced_labels() + images.write("Case", _CASE, labels, _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + images.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + fields = Dataset(tmp_path / "Fields", "h5") + fields.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + + stage = _warping(images, fields, fill=0.0) + got = stage(_CASE, torch.from_numpy(labels.copy()), Attribute(source)).numpy() + want = _simpleitk_warp( + labels, _displacement(), interpolator=sitk.sitkNearestNeighbor, pixel=sitk.sitkUInt8, fill=0.0 + ) + + assert got.dtype == np.uint8 + assert set(np.unique(got).tolist()) <= {0, 100}, "nearest picks a label, it does not blend two" + np.testing.assert_array_equal(got[0], want) + + +@pytest.mark.parametrize("warping", [False, True], ids=["grid_change", "composed_with_a_field"]) +def test_an_explicit_interpolation_is_honoured_on_both_paths(tmp_path: Path, warping: bool) -> None: + """A dtype cannot decide this on its own -- a CT is int16 and so is nothing else about it -- so + the heuristic only ever claims uint8 and ``interpolation`` covers everything it cannot know. + + Parametrised over BOTH gathers on purpose. A declaration honoured by the composed path and + ignored by the plain one is worse than no declaration: the page that tells a user to set it for + a label map is then right about half their chains, and the half it is wrong about says nothing. + """ + images = Dataset(tmp_path / "Images", "h5") + # int16, not uint8: a dtype the heuristic deliberately does not claim, so only the declaration + # can be what decides. + labels = _spaced_labels().astype(np.int16) + images.write("Case", _CASE, labels, _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + images.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + fields = Dataset(tmp_path / "Fields", "h5") + fields.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + + def resample(**kwargs: object) -> set[int]: + stage = _warping(images, fields, fill=0.0, **kwargs) if warping else _stage(images, fill=0.0, **kwargs) + got = stage(_CASE, torch.from_numpy(labels.copy()), Attribute(source)).numpy() + return set(np.unique(got).tolist()) + + assert not resample(interpolation="nearest") - {0, 100}, "nearest picks a label, it does not blend" + assert resample(interpolation="linear") - {0, 100}, "asking for linear must actually interpolate" + assert resample() - {0, 100}, "int16 is not a label map the dtype can claim" + + with pytest.raises(TransformError, match="unknown interpolation"): + _warping(images, fields, interpolation="cubic") def test_it_warps_onto_the_reference_where_simpleitk_does(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: @@ -809,3 +879,29 @@ def test_a_bound_with_no_field_is_refused() -> None: """ with pytest.raises(TransformError, match="no field to apply"): ResampleToReference(entry=_CASE, group="Reference", max_displacement=1.0) + + +def test_the_two_gathers_agree_bit_for_bit_through_an_identity_field(tmp_path: Path) -> None: + """One arithmetic, two loops: per-axis maps, and eight corners at a coordinate volume. + + The separable loop cannot serve a displacement (a displacement is not separable) and the flat + gather is the slower way to do a map that is. So both exist, and both have to obey the same + inside interval, the same tap clamp and the same fill. A ZERO field is where that is checkable: + the composed path reduces to the grid change alone, so the two must land on the same voxels. + + Bit for bit, not close. A tolerance here would hide exactly the drift this guards -- one loop + keeping a rule the other quietly dropped, which is how the CPU-half guard went missing once. + """ + images = Dataset(tmp_path / "Images", "h5") + volume = _high_frequency() + images.write("Case", _CASE, volume, _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + images.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + fields = Dataset(tmp_path / "Fields", "h5") + fields.write("DVF", _CASE, np.zeros((3, *_FIELD_SPATIAL), np.float32), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + + separable = _stage(images, fill=_FILL)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy() + composed = _warping(images, fields, fill=_FILL)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy() + + assert 0 < int((separable == _FILL).sum()) < separable.size, "the fixture must have a rim, and not be all rim" + np.testing.assert_array_equal(composed, separable) From 8931a2ad4a46ed20af961f41dfa9c9ad8ce39dd2 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 21:21:17 +0200 Subject: [PATCH 2/4] feat(examples): a Transform example, and the guard that keeps it running TRANSFORM shipped as the release's headline with no example, and the konfai-cli skill still described three workflows mapping to three files. examples/Transform covers both directions the cardinality can go: a median template folded from a cohort that shares no grid -- which is what ResampleToReference is in the chain for -- and four drawn copies of every case. The cohort is generated locally: six volumes, 3.5 MB, nothing downloaded, a minute on CPU. Written uncompressed, because a compressed .mha cannot serve a disk region and the example would then demonstrate the opposite of streaming. No two members share an extent, a spacing or an origin, so `grid: strict` refuses them as stored and the resample in the chain is what makes it true. Both configs are run AS SHIPPED by an integration test rather than templated into one, so a rename in the grammar fails in CI and not in someone's terminal. The doc-examples test loses two ways of quietly not running: the console script resolves through the harness helper that falls back to `python -c` instead of skipping where the script is named differently, and a missing page is an assertion rather than an empty parametrize set, which pytest reports as a pass. The shared notebook helper asked for `--gpu 0`, which argparse rejects on a machine exporting CUDA_VISIBLE_DEVICES=1, and dropped the last unterminated line from the tail it raises -- usually the exception. Synthesis evaluates with SSIM and installed no scikit-image, so it failed after training rather than before it. TotalSegmentator's notebook advertised an uncertainty step its bundle does not build. --- .claude/skills/konfai-cli/SKILL.md | 16 ++- examples/README.md | 5 +- examples/Synthesis/README.md | 3 +- examples/Synthesis/Synthesis_demo.ipynb | 2 +- .../TotalSegmentator_demo.ipynb | 2 +- examples/Transform/README.md | 118 ++++++++++++++++++ examples/Transform/Transform.yml | 34 +++++ examples/Transform/Transform_expand.yml | 29 +++++ examples/Transform/make_dataset.py | 81 ++++++++++++ examples/konfai_demo.py | 15 ++- .../test_transform_doc_examples.py | 16 +-- tests/integration/test_transform_example.py | 82 ++++++++++++ 12 files changed, 384 insertions(+), 19 deletions(-) create mode 100644 examples/Transform/README.md create mode 100644 examples/Transform/Transform.yml create mode 100644 examples/Transform/Transform_expand.yml create mode 100644 examples/Transform/make_dataset.py create mode 100644 tests/integration/test_transform_example.py diff --git a/.claude/skills/konfai-cli/SKILL.md b/.claude/skills/konfai-cli/SKILL.md index c38ca376..83265c07 100644 --- a/.claude/skills/konfai-cli/SKILL.md +++ b/.claude/skills/konfai-cli/SKILL.md @@ -30,16 +30,24 @@ There are two command-line surfaces: ## The canonical loop (`konfai`) -Three workflows map to three files, each with one mandatory root key: +Each workflow maps to one file with one mandatory root key: | Command | File | Root key | |---|---|---| | `TRAIN` / `RESUME` | `Config.yml` | `Trainer:` | | `PREDICTION` | `Prediction.yml` | `Predictor:` | | `EVALUATION` | `Evaluation.yml` | `Evaluator:` | +| `TRANSFORM` | `Transform.yml` | `Transformer:` | + +`TRANSFORM` sits outside the loop: it runs no model. It reads a dataset, applies a chain, and writes +a dataset — resampling a cohort onto one grid, folding it into a template (`Reduce`, N→1), expanding +each case into drawn copies (`Expand`, 1→N). A chain may still embed a `KonfAIInference` stage, so +"no model" means no top-level one. It takes no `-tb`, and `--plan` prints what a run would do and +stops without writing the deliverable — it does probe each destination with a real region-write it +then removes, so the output store may be created. **Don't write configs from scratch — copy a runnable template from `examples/`** (Segmentation, -Synthesis or Registration) and adapt it. Then: +Synthesis, Registration or Transform) and adapt it. Then: ```bash cd examples/Segmentation # always run from the dir holding the configs + Dataset/ @@ -47,6 +55,10 @@ cd examples/Segmentation # always run from the dir holding the c konfai TRAIN -y --gpu 0 --config Config.yml konfai PREDICTION -y --gpu 0 --config Prediction.yml --models Checkpoints//.pt konfai EVALUATION -y --config Evaluation.yml + +cd ../Transform # no model, no GPU +konfai TRANSFORM --config Transform.yml --plan # what it would do; writes no deliverable +konfai TRANSFORM --config Transform.yml ``` Outputs are namespaced by the `train_name` in the config: `Checkpoints//`, diff --git a/examples/README.md b/examples/README.md index 600080c8..b0a9748b 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,17 +9,20 @@ README, then *Runtime > Run all*. ## Start here -These three are the framework itself: a YAML config, the `konfai` CLI, and nothing else. +These four are the framework itself: a YAML config, the `konfai` CLI, and nothing else. | Example | What you get | Time on a GPU | |---|---|---| | [`Registration`](Registration/) | Train a `VoxelMorph` to align real pelvis CT slices, and measure how much of a *known* deformation it recovered. | ~3 min | | [`Segmentation`](Segmentation/) | Train a UNet on five pelvis CT cases, predict the labels, score them with Dice. | ~7 min | | [`Synthesis`](Synthesis/) | Turn an MR volume into a synthetic CT, scored with MAE / PSNR / SSIM inside the body mask. | ~7 min | +| [`Transform`](Transform/) | Fold a cohort into one template, and expand each case into drawn copies. No model at all. | ~1 min, CPU | `Registration` is the shortest way to see the whole `TRAIN -> PREDICTION -> EVALUATION` loop. `Segmentation` is the best template to copy for your own data. `Synthesis` shows the richer patterns: a custom Python model, a perceptual loss, test-time augmentation, and an optional GAN variant. +`Transform` is the odd one out and the quickest: no network, no checkpoint, no download — it is the +workflow for when what you need is data rather than a prediction. Their training runs are **deliberately short** — enough to see the pipeline work end to end, not enough to produce a usable model. Each README says what score to expect and which knob to raise. diff --git a/examples/Synthesis/README.md b/examples/Synthesis/README.md index 4de2d32b..b4423460 100644 --- a/examples/Synthesis/README.md +++ b/examples/Synthesis/README.md @@ -141,7 +141,8 @@ every command from this directory: ```bash cd examples/Synthesis -pip install "konfai[smp]" # the generator wraps segmentation_models_pytorch +pip install "konfai[smp,ssim]" # smp: the generator wraps segmentation_models_pytorch + # ssim: the SSIM metric in Evaluation.yml needs scikit-image ``` ### 1. Train diff --git a/examples/Synthesis/Synthesis_demo.ipynb b/examples/Synthesis/Synthesis_demo.ipynb index 7b4ab608..2bd3922a 100644 --- a/examples/Synthesis/Synthesis_demo.ipynb +++ b/examples/Synthesis/Synthesis_demo.ipynb @@ -38,7 +38,7 @@ "\n", "from konfai_demo import latest_checkpoint, read, run, setup, show\n", "\n", - "EXAMPLE_DIR, DATASET_DIR, DEVICE = setup(REPO_DIR, \"Synthesis\", (\"konfai\", f\"{REPO_DIR}[imaging,smp]\"), \"huggingface_hub\", \"matplotlib\", (\"segmentation_models_pytorch\", \"segmentation-models-pytorch\"))" + "EXAMPLE_DIR, DATASET_DIR, DEVICE = setup(REPO_DIR, \"Synthesis\", (\"konfai\", f\"{REPO_DIR}[imaging,smp,ssim]\"), \"huggingface_hub\", \"matplotlib\", (\"segmentation_models_pytorch\", \"segmentation-models-pytorch\"))" ] }, { diff --git a/examples/TotalSegmentator/TotalSegmentator_demo.ipynb b/examples/TotalSegmentator/TotalSegmentator_demo.ipynb index 03cc8a88..1db836e8 100644 --- a/examples/TotalSegmentator/TotalSegmentator_demo.ipynb +++ b/examples/TotalSegmentator/TotalSegmentator_demo.ipynb @@ -177,7 +177,7 @@ "- **your own volume** — point `-i` at any `.mha` / `.nii.gz`, or at an OME-Zarr or DICOM directory;\n", " KonfAI detects the store format on read.\n", "- **score it** — `totalsegmentator-konfai eval total -i input.mha --gt reference.mha -o Output`.\n", - "- **everything at once** — `totalsegmentator-konfai pipeline ...` chains inference, evaluation and uncertainty.\n", + "- **everything at once** — `totalsegmentator-konfai pipeline ...` chains inference and evaluation. This bundle ships no uncertainty workflow, so unlike the other apps there is no `-uncertainty` step to chain.\n", "- **see how an app is built** — `apps/totalsegmentator/` in this repo is the whole wrapper." ] } diff --git a/examples/Transform/README.md b/examples/Transform/README.md new file mode 100644 index 00000000..ecd90fbf --- /dev/null +++ b/examples/Transform/README.md @@ -0,0 +1,118 @@ +# Transform — the workflow that makes a dataset + +`konfai TRANSFORM` reads a dataset, applies a chain of transforms, and writes a +dataset. There is no network, no checkpoint and no training loop — `EVALUATION` +has none either, and the difference is the product: an evaluation measures, this +one is the workflow you reach for when the thing you need is *data*. + +Two configs here, one per direction the cardinality can change. + +| File | Shape | What it builds | +| --- | --- | --- | +| `Transform.yml` | N → 1 | One median template from a cohort that does not share a grid | +| `Transform_expand.yml` | 1 → N | Four drawn copies of every case | + +Everything runs on CPU in under a minute, on ~3.5 MB of synthetic data this +directory generates. Nothing is downloaded. + +## Run it + +```bash +pip install "konfai[imaging]" # make_dataset.py writes .mha through SimpleITK +python make_dataset.py # ./Raw//CT.mha, 6 cases + +konfai TRANSFORM --config Transform.yml --plan # read the plan first +konfai TRANSFORM --config Transform.yml # ./Template/template/CT_template.mha + +konfai TRANSFORM --config Transform_expand.yml # ./Augmented/_r01..r04/ +``` + +`--plan` writes none of the deliverable, but it is not read-only: it opens a real region-write on +each destination and removes it, so the output store itself may be created. + +## What the cohort looks like, and why that matters + +`make_dataset.py` writes six volumes that agree about nothing: extents differ by +a few voxels, spacings by up to 30%, origins by more than a voxel. That is the +ordinary state of a cohort as acquired — an acquisition's stage coordinates are +not an anatomical frame — and it is why `Reduce` refuses it as stored: + +```text +case 'CASE_001' lands on extent [44, 60, 52] where 'CASE_000' lands on [48, 56, 56] +``` + +`ResampleToReference` is what makes the agreement true rather than waived. It +puts every member on one named member's grid — extent, spacing, origin and +direction — so `grid: strict` passes because the cohort really is on one grid, +not because the check was relaxed. + +```{warning} +`grid: shape_only` and `grid: reference:` will happily average volumes +that do not overlap. The result still looks like a volume. Put the cohort on one +grid first. +``` + +## Read the plan before you read anything else + +`--plan` prints what the run will do and stops. Nothing is written first, and +the plan is the run's own verdict, not an estimate of it: + +```text + CT -> CT_template: REDUCE 6 case(s) -> 1 output 'template' -- REDUCE + 19 resident region(s) of 64 row(s) = 0.01 GiB (every case resident per region) + peak ~= 0.01 GiB vs per-rank budget 1.86 GiB +``` + +Nineteen regions for six cases is `Median` being honest: it needs every case +resident to name the middle one, then stacks them into a new tensor and sorts a +copy of that. `Mean` folds one case at a time and holds two regions whatever the +cohort's size — swap `operator: Mean` and watch the line change. + +The plan also reports how much of the reference each member actually covers: + +```text +NOTE: case 'CASE_005' covers 82.0% of reference 'CASE_000'; the rest of what it writes is fill (0) +``` + +That is worth reading. A member covering 60% of the reference is contributing +fill to nearly half the template, and nothing about the written volume would +look wrong. + +## Where the cardinality changes + +In `Transform.yml`, everything above `Reduce` runs once per case and everything +below it runs once, on the folded result: + +```yaml +Clip: {min_value: 0.0, max_value: 400.0} # per case +ResampleToReference: {entry: CASE_000, ...} # per case +Reduce: {operator: Median, output: template} # <- N becomes 1 here +Write: {dataset: ./Template:mha} # once +``` + +Only voxel-local stages may follow a `Reduce`, because each is handed one +*region* of the result. Anything reading across space belongs in a second chain +that reads the written template back. + +`Transform_expand.yml` is the mirror image: `Expand` marks where one case +becomes four, and the draws go **after** it. A draw declared before the marker +is applied once per case, which is a random transform rather than an +augmentation, and the run refuses it and says so. + +Because `Brightness` is pointwise, the four copies share a single read pass over +the source — the plan says `4 shared pass, 0 own pass`. A draw that moves voxels +around (`Rotate`, `Flip`) cannot share it, and each copy gets its own pass. + +## Reproducibility + +`manual_seed` in `Transform_expand.yml` is what makes an image chain and its mask +chain draw the *same* copies. The two chains never meet: each derives its draws +from `(seed, case, which copy)`, so they agree without coordinating. A mask +rotated by a different angle than its image is a silently ruined dataset, not an +error — which is why the seed is not optional in practice. + +## Next + +- The full reference: [`config_guide/transform.md`](../../docs/source/config_guide/transform.md) +- Chains that embed a trained model (`KonfAIInference`), streaming rules and the + memory budget are all documented there. diff --git a/examples/Transform/Transform.yml b/examples/Transform/Transform.yml new file mode 100644 index 00000000..e2553c2c --- /dev/null +++ b/examples/Transform/Transform.yml @@ -0,0 +1,34 @@ +# Build one template from a cohort that does not share a grid: N cases in, 1 volume out. +# +# python make_dataset.py +# konfai TRANSFORM --config Transform.yml --plan # read this first +# konfai TRANSFORM --config Transform.yml +# +# Writes ./Template/template/CT_template.mha -- one entry, whatever the cohort's size. +Transformer: + name: TEMPLATE + # A chain that cannot stream is a chain that would hold a whole volume. `error` refuses instead, + # which is what you want while the config is still being written; `warn` is the default. + on_fallback: error + Dataset: + dataset_filenames: + - ./Raw:mha + memory_budget: 2G + groups_src: + CT: + groups_dest: + CT_template: + transforms: + # Per case, in order, before anything is folded. + Clip: {min_value: 0.0, max_value: 400.0} + # The cohort as acquired fails `strict`: extents, spacings and origins all differ. + # This puts every member on CASE_000's grid, which makes the agreement true rather + # than waived. Any member would do -- what matters is that one is named. + ResampleToReference: {entry: CASE_000, group: CT, fill: 0.0} + # The cardinality changes here. Everything above ran once per case; everything below + # runs once, on the folded result. + Reduce: + operator: Median + output: template + grid: strict + Write: {dataset: ./Template:mha} diff --git a/examples/Transform/Transform_expand.yml b/examples/Transform/Transform_expand.yml new file mode 100644 index 00000000..89dee826 --- /dev/null +++ b/examples/Transform/Transform_expand.yml @@ -0,0 +1,29 @@ +# The other direction: 1 case in, N copies out. +# +# konfai TRANSFORM --config Transform_expand.yml --plan +# konfai TRANSFORM --config Transform_expand.yml +# +# Writes ./Augmented/_r01 .. _r04 -- four drawn copies of every case. +Transformer: + name: AUGMENT + on_fallback: error + # Every draw derives from this and the case, so the run is reproducible and two chains asked for + # the same case draw the SAME copies -- which is how an image and its mask stay paired. + manual_seed: 7 + Dataset: + dataset_filenames: + - ./Raw:mha + memory_budget: 2G + groups_src: + CT: + groups_dest: + CT_aug: + transforms: + Clip: {min_value: 0.0, max_value: 400.0} + # Above the marker: once per case. Below it: once per copy. + Expand: {nb: 4, pattern: "{name}_r{a:02d}"} + # A draw belongs after the marker. Declared before one -- or with no marker at all -- + # it is applied once per case, which is a random transform, not an augmentation. + # Brightness is pointwise, so the copies share a single read pass over the source. + Brightness: {b_std: 25.0} + Write: {dataset: ./Augmented:mha} diff --git a/examples/Transform/make_dataset.py b/examples/Transform/make_dataset.py new file mode 100644 index 00000000..00bbb0df --- /dev/null +++ b/examples/Transform/make_dataset.py @@ -0,0 +1,81 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Write a small cohort that does NOT share a grid, which is what makes the example worth running. + +Six volumes, each a smooth blob on its own extent, spacing and origin, with its own intensity +dynamics. That is the ordinary state of a cohort as acquired: nothing about it is wrong, and nothing +about it lets you average two members together. Bringing them onto one grid is the work. + + python make_dataset.py # writes ./Raw//CT.mha + +The whole cohort is about 3 MB and takes a second to generate; there is nothing to download. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import SimpleITK as sitk + +#: Each case gets its own geometry. The spread is deliberate and mild: extents differ by a few +#: voxels, spacings by up to 30%, and origins by more than one voxel -- enough that `grid: strict` +#: refuses the cohort as stored, which is the point the example is making. +_CASES: tuple[tuple[str, tuple[int, int, int], tuple[float, float, float], tuple[float, float, float]], ...] = ( + ("CASE_000", (48, 56, 56), (1.00, 1.00, 1.00), (0.0, 0.0, 0.0)), + ("CASE_001", (44, 60, 52), (1.15, 0.95, 1.05), (2.5, -1.0, 0.5)), + ("CASE_002", (52, 52, 60), (0.90, 1.10, 0.95), (-1.5, 3.0, -2.0)), + ("CASE_003", (46, 58, 54), (1.05, 1.00, 1.20), (1.0, 1.5, 2.5)), + ("CASE_004", (50, 54, 58), (0.95, 1.20, 1.00), (-2.0, 0.5, 1.0)), + ("CASE_005", (45, 59, 53), (1.10, 1.05, 0.90), (3.0, -2.5, -1.5)), +) + + +def _volume(shape: tuple[int, int, int], seed: int) -> np.ndarray: + """A blob with a soft rim and a little noise, in roughly Hounsfield-looking numbers.""" + rng = np.random.default_rng(seed) + grids = np.meshgrid(*[np.linspace(-1.0, 1.0, extent) for extent in shape], indexing="ij") + # A shifted, slightly anisotropic ellipsoid, so the members overlap without coinciding. + centre = rng.uniform(-0.15, 0.15, size=3) + radii = rng.uniform(0.55, 0.75, size=3) + distance = sum(((axis - centre[i]) / radii[i]) ** 2 for i, axis in enumerate(grids)) + blob = 1.0 / (1.0 + np.exp((distance - 1.0) * 8.0)) + intensity = rng.uniform(280.0, 360.0) + return (blob * intensity + rng.normal(0.0, 6.0, size=shape)).astype(np.float32) + + +def write(root: Path, group: str = "CT") -> None: + for seed, (name, shape, spacing, origin) in enumerate(_CASES): + image = sitk.GetImageFromArray(_volume(shape, seed)) + # sitk takes geometry in (x, y, z) where the array is (z, y, x). + image.SetSpacing(tuple(reversed(spacing))) + image.SetOrigin(tuple(reversed(origin))) + case = root / name + case.mkdir(parents=True, exist_ok=True) + # Uncompressed on purpose: a compressed .mha cannot serve a disk region, so every slab would + # decode the whole volume again and the example would demonstrate the opposite of streaming. + sitk.WriteImage(image, str(case / f"{group}.mha"), useCompression=False) + print(f"Wrote {len(_CASES)} cases under {root}") + print("No two share an extent, a spacing or an origin -- 'grid: strict' refuses them as stored.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, default=Path("./Raw"), help="where to write the cohort") + parser.add_argument("--group", default="CT", help="group name each case is stored under") + write(**vars(parser.parse_args())) diff --git a/examples/konfai_demo.py b/examples/konfai_demo.py index 15a0c3fb..05aa92ac 100644 --- a/examples/konfai_demo.py +++ b/examples/konfai_demo.py @@ -41,7 +41,7 @@ def setup(repo_dir: Path, example: str, *packages: str | tuple[str, str]) -> tup """Install what is missing and return the example directory, its dataset directory, and the device flags. `packages` are pip requirements, installed only when already absent. The import name is derived - from the requirement (`konfai[imaging]` -> `konfai`, `scikit-image` -> `scikit_image`); pass a + from the requirement (`konfai[imaging]` -> `konfai`, `huggingface-hub` -> `huggingface_hub`); pass a `(import_name, requirement)` pair when it cannot be, as for a local path. """ global _WORKDIR @@ -59,7 +59,15 @@ def setup(repo_dir: Path, example: str, *packages: str | tuple[str, str]) -> tup import torch _WORKDIR = repo_dir / "examples" / example - device = ["--gpu", "0"] if torch.cuda.is_available() else ["--cpu", "1"] + # The first VISIBLE device, not device 0: `--gpu` is validated against cuda_visible_devices(), + # which reports the raw CUDA_VISIBLE_DEVICES values, so on a workstation exporting + # CUDA_VISIBLE_DEVICES=1 argparse rejects `--gpu 0` outright. + if torch.cuda.is_available(): + from konfai import cuda_visible_devices + + device = ["--gpu", str(cuda_visible_devices()[0])] + else: + device = ["--cpu", "1"] print("KonfAI :", repo_dir) print("Example:", _WORKDIR) print("Device :", " ".join(device)) @@ -92,6 +100,9 @@ def run(*command: str) -> None: print(" ", tail[-1][:140], flush=True) next_print = time.time() - started + interval interval = min(interval * 1.6, 30.0) + if pending.strip(): + # The final line carries no terminator, and on a crash that line is the exception. + tail = [*tail, pending][-40:] if process.wait(): raise RuntimeError("\n".join(tail[-25:])) print(f" done in {time.time() - started:.0f} s\n", flush=True) diff --git a/tests/integration/test_transform_doc_examples.py b/tests/integration/test_transform_doc_examples.py index 953f9128..a3ab1283 100644 --- a/tests/integration/test_transform_doc_examples.py +++ b/tests/integration/test_transform_doc_examples.py @@ -32,7 +32,7 @@ import numpy as np import pytest -from harness import subprocess_env +from harness import konfai_cli_command, subprocess_env from ruamel.yaml import YAML pytest.importorskip("SimpleITK") @@ -142,8 +142,9 @@ def _write_fields(root: Path, cases: int = 2, shape=(4, 6, 6)) -> None: def _doc_examples() -> list[tuple[int, str]]: - if not DOC.is_file(): - return [] + # Asserted, not skipped: an empty parametrize set is a pass, so a page that moved would retire + # this whole guard without a single red test. + assert DOC.is_file(), f"the page this guards is not where it is expected: {DOC}" text = DOC.read_text(encoding="utf-8") return [(line, config) for line, body in _yaml_blocks(text) if (config := _as_config(body)) is not None] @@ -179,7 +180,7 @@ def test_a_documented_example_plans_and_runs(line: int, config: str, tmp_path: P # The plan is computed on the launcher and the run then spawns: planning green proves nothing # about what crosses that boundary. executed = subprocess.run( - [_konfai(), "TRANSFORM", "--config", "Transform.yml", "--transforms-dir", "Transforms"], + [*konfai_cli_command(), "TRANSFORM", "--config", "Transform.yml", "--transforms-dir", "Transforms"], capture_output=True, text=True, env=environment, @@ -187,10 +188,3 @@ def test_a_documented_example_plans_and_runs(line: int, config: str, tmp_path: P timeout=900, ) assert executed.returncode == 0, f"transform.md:{line} plans but does not run:\n{executed.stdout}{executed.stderr}" - - -def _konfai() -> str: - konfai = Path(sys.executable).with_name("konfai") - if not konfai.exists(): - pytest.skip("the konfai console script is not installed in this environment") - return str(konfai) diff --git a/tests/integration/test_transform_example.py b/tests/integration/test_transform_example.py new file mode 100644 index 00000000..4fc61093 --- /dev/null +++ b/tests/integration/test_transform_example.py @@ -0,0 +1,82 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""The shipped ``examples/Transform`` configs, run as shipped. + +An example is the first thing a new user runs, and the only one whose breakage they read as the +framework being broken. These are copied verbatim -- not templated, not rewritten -- so a rename in +the transform grammar fails here rather than in someone's terminal. + +The cohort is generated by the example's own ``make_dataset.py``: six small volumes, no download, +a few seconds each way. +""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +from harness import konfai_cli_command, subprocess_env + +EXAMPLE = Path(__file__).resolve().parents[2] / "examples" / "Transform" + + +def _run(command: list[str], workdir: Path) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, capture_output=True, text=True, env=subprocess_env(), cwd=workdir, timeout=900) + + +@pytest.fixture(scope="module") +def cohort(tmp_path_factory: pytest.TempPathFactory) -> Path: + """The example directory, with the cohort its own script generates.""" + assert EXAMPLE.is_dir(), f"the example this guards is not where it is expected: {EXAMPLE}" + workdir = tmp_path_factory.mktemp("transform_example") + for name in ("Transform.yml", "Transform_expand.yml", "make_dataset.py"): + shutil.copy(EXAMPLE / name, workdir / name) + made = _run([sys.executable, "make_dataset.py"], workdir) + assert made.returncode == 0, f"make_dataset.py failed:\n{made.stdout}{made.stderr}" + assert sorted(p.name for p in (workdir / "Raw").iterdir()) == [f"CASE_{i:03d}" for i in range(6)] + return workdir + + +@pytest.mark.integration +def test_the_template_example_folds_the_cohort_into_one_entry(cohort: Path) -> None: + """N to 1. The README's headline claim, and the reason ResampleToReference is in the chain: + the cases do not share a grid, so ``grid: strict`` would refuse them as stored.""" + planned = _run([*konfai_cli_command(), "TRANSFORM", "--config", "Transform.yml", "--plan"], cohort) + assert planned.returncode == 0, f"the template example does not plan:\n{planned.stdout}{planned.stderr}" + assert "REDUCE 6 case(s) -> 1 output 'template'" in planned.stdout + # --plan probes each destination with a real region-write open, so the store itself may be + # created; what it must not leave behind is an entry, the probe's own included. + assert not (cohort / "Template" / "template").exists(), "--plan must not write the deliverable" + assert [p.name for p in (cohort / "Template").iterdir()] == [], "--plan must remove what it probed with" + + executed = _run([*konfai_cli_command(), "TRANSFORM", "--config", "Transform.yml"], cohort) + assert executed.returncode == 0, f"the template example does not run:\n{executed.stdout}{executed.stderr}" + assert (cohort / "Template" / "template" / "CT_template.mha").is_file() + + +@pytest.mark.integration +def test_the_expand_example_writes_one_entry_per_drawn_copy(cohort: Path) -> None: + """1 to N, and the copies are named by the config's own pattern -- which is what a second chain + reading them back has to agree with.""" + executed = _run([*konfai_cli_command(), "TRANSFORM", "--config", "Transform_expand.yml"], cohort) + assert executed.returncode == 0, f"the expand example does not run:\n{executed.stdout}{executed.stderr}" + + written = sorted(p.name for p in (cohort / "Augmented").iterdir()) + assert written == [f"CASE_{case:03d}_r{copy:02d}" for case in range(6) for copy in range(1, 5)] From f7b2d6a99bd2d1c03551e1ab7450a0e89b4a9839 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 21:21:17 +0200 Subject: [PATCH 3/4] fix(studio,mcp): say what a call does, and refuse the question that cannot be answered /api/quit's loopback check reads the TCP peer, and behind the same-host reverse proxy REMOTE.md documents, every peer IS 127.0.0.1. Uvicorn only rewrites it under --proxy-headers, which is off by default and described as a logging concern, so the guard the docstring called sufficient was a no-op for exactly the deployment the docs recommend. Any forwarding header now refuses without that flag -- Forwarded, X-Real-IP or any X-Forwarded-*, because what matters is not which one arrived but that one did -- and the shutdown task is held so the loop cannot collect it mid-sleep. plan_transform told an agent it "writes no data" two sentences before saying it opens a real region-write on each destination -- and those destinations are the user's own stores. It says so where an agent reads it, along with the 50-entry cap on needs_attention. konfai-mcp bounds konfai, because runner.py imports konfai.transformer at module scope and a 1.7.0 satisfies a bare name: the whole server then dies on import. Spelled `>1.7.0` and not `>=1.8.0` so a source checkout satisfies it -- setuptools_scm builds the unreleased tree as 1.7.1.devN, which PEP 440 sorts below 1.7.1. --- konfai-mcp/konfai_mcp/guide.py | 11 ++++--- konfai-mcp/pyproject.toml | 6 +++- studio/konfai_studio/cli.py | 6 +++- studio/konfai_studio/server.py | 54 ++++++++++++++++++++++++++++++---- studio/tests/test_quit.py | 39 ++++++++++++++++++++++-- 5 files changed, 102 insertions(+), 14 deletions(-) diff --git a/konfai-mcp/konfai_mcp/guide.py b/konfai-mcp/konfai_mcp/guide.py index f3671a0c..81e787b0 100644 --- a/konfai-mcp/konfai_mcp/guide.py +++ b/konfai-mcp/konfai_mcp/guide.py @@ -504,13 +504,16 @@ "Next: wait_for_job." ), "plan_transform": ( - "Use BEFORE run_transform, always: it is the dry run, and it writes no data. " + "Use BEFORE run_transform, always: it is the dry run, and it produces none of the deliverable. " "This plans every (case, chain) from the session Transform.yml. The plan is a measurement, not an " "estimate -- it opens and removes a real region-write on each destination -- so its verdict is the one " "the run will act on: STREAM (bounded memory), WHOLE-VOLUME (the case is assembled whole, with the stage " - "that refused it named), SKIP (already written), REDUCE, REFUSED. " - "Outputs: {ok, report, verdict_counts, budget, needs_attention[], over_budget[]}. A non-empty over_budget " - "means run_transform would refuse before writing anything. " + "that refused it named), SKIP (already written), REDUCE, REFUSED. That probe TOUCHES the output " + "locations, which are the user's own stores: an entry is created and removed, and a single-file store " + "(h5) is created if it did not exist. " + "Outputs: {ok, config_path, world_size, report, verdict_counts, budget, budget_bytes, needs_attention[], " + "over_budget[]}. A non-empty over_budget means run_transform would refuse before writing anything; " + "needs_attention lists at most 50 entries, so trust verdict_counts for the totals. " "Next: fix what needs_attention names, or run_transform." ), "run_transform": ( diff --git a/konfai-mcp/pyproject.toml b/konfai-mcp/pyproject.toml index 0f4e5372..cc309140 100644 --- a/konfai-mcp/pyproject.toml +++ b/konfai-mcp/pyproject.toml @@ -14,8 +14,12 @@ authors = [ { name = "Valentin Boussot", email = "boussot.v@gmail.com" } ] +# konfai carries a bound because runner.py imports konfai.transformer at module scope: a 1.7.0 +# satisfies a bare name, and the server then dies on import -- all of it, not just the transform +# tools. Spelled `>1.7.0` and not `>=1.8.0` so a source checkout satisfies it: setuptools_scm builds +# the unreleased tree as 1.7.1.devN, which PEP 440 sorts BELOW 1.7.1 and above 1.7.0. dependencies = [ - "konfai", + "konfai>1.7.0", "konfai-apps", "fastmcp", "ruamel.yaml", diff --git a/studio/konfai_studio/cli.py b/studio/konfai_studio/cli.py index b3ede21c..3411daf6 100644 --- a/studio/konfai_studio/cli.py +++ b/studio/konfai_studio/cli.py @@ -16,7 +16,8 @@ def main() -> None: parser.add_argument( "--proxy-headers", action="store_true", - help="trust X-Forwarded-* from a reverse proxy (set when behind nginx/Caddy for correct client IP + scheme)", + help="trust X-Forwarded-* from a reverse proxy: correct client IP + scheme, and what lets" + " /api/quit tell a local user from a remote one (set when behind nginx/Caddy)", ) parser.add_argument( "--forwarded-allow-ips", @@ -61,6 +62,9 @@ def main() -> None: ) scheme = "https" if args.ssl_certfile else "http" print(f"KonfAI Studio -> {scheme}://{args.host}:{args.port} (auth {'on' if authed else 'off'})") + # The app is imported by uvicorn from a string, so it cannot read these arguments. /api/quit + # needs to know whether the peer it sees was rewritten from X-Forwarded-For or is a proxy's. + os.environ["KONFAI_STUDIO_PROXY_HEADERS"] = "1" if args.proxy_headers else "0" uvicorn.run( "konfai_studio.server:app", host=args.host, diff --git a/studio/konfai_studio/server.py b/studio/konfai_studio/server.py index 909491de..e23da280 100644 --- a/studio/konfai_studio/server.py +++ b/studio/konfai_studio/server.py @@ -851,15 +851,49 @@ async def apps(session: str = Query("apps")) -> dict[str, Any]: return {"ok": ok, "apps": listed} +#: Tasks kept alive across their own await. The loop holds only a weak reference to a bare +#: ``create_task``, so a task that sleeps before doing its work can be collected before it runs. +_SHUTDOWN_TASKS: set[asyncio.Task] = set() + + +def _forwarded_by_a_proxy(request: Request) -> bool: + """Whether anything in front of this server claims to have forwarded the request. + + Any of them is enough, and the prefix is matched rather than a list: a proxy that sets only + ``X-Real-IP``, or only a vendor's own ``X-Forwarded-Host``, still means the peer address belongs + to the proxy. What matters is not which header arrived but that one did. + """ + return any( + name == "forwarded" or name == "x-real-ip" or name.startswith("x-forwarded-") for name in request.headers + ) + + +def _trusts_proxy_headers() -> bool: + """Whether uvicorn was told to rewrite the peer address from ``X-Forwarded-For``. + + Set by the CLI, because the app is imported from a string and never sees its arguments. Read at + call time rather than at import so a test can set it. + """ + return os.environ.get("KONFAI_STUDIO_PROXY_HEADERS") == "1" + + @app.post("/api/quit") async def quit_server(request: Request) -> dict[str, bool]: """Stop the Studio server (graceful: the lifespan teardown closes agents and reaps TensorBoards). - Two guards, because neither alone is enough. The client must be on this machine, so a remote - user cannot take a shared server down, token or not. And it must send the header below: the - loopback check only proves the TCP peer is local, which any page open in the user's browser - also is, so without it a drive-by form POST to localhost would shut Studio down. A custom - header is unforgeable from a form and, cross-origin, needs a preflight this server never grants. + Three guards, because none of them alone is enough. + + The client must be on this machine, so a remote user cannot take a shared server down, token or + not. That reads the TCP peer, which is the real client only when nothing sits in front: behind + the same-host reverse proxy REMOTE.md documents, every request arrives from 127.0.0.1. Uvicorn + rewrites the peer from ``X-Forwarded-For`` under ``--proxy-headers``, so a forwarding header + arriving WITHOUT that flag means the peer belongs to the proxy and the loopback check answers a + question nobody asked — refuse instead. + + And it must send the header below: the loopback check only proves the TCP peer is local, which + any page open in the user's browser also is, so without it a drive-by form POST to localhost + would shut Studio down. A custom header is unforgeable from a form and, cross-origin, needs a + preflight this server never grants. Slicer's Studio button and the titlebar power button rely on this — the server runs detached, with no terminal to Ctrl+C. @@ -869,6 +903,12 @@ async def quit_server(request: Request) -> dict[str, bool]: client = request.client.host if request.client else "" if client not in ("127.0.0.1", "::1"): raise HTTPException(403, "the Studio server can only be stopped from its own machine") + if _forwarded_by_a_proxy(request) and not _trusts_proxy_headers(): + raise HTTPException( + 403, + "this request came through a proxy and the server was not started with --proxy-headers," + " so it cannot tell which machine it is from", + ) if request.headers.get("x-konfai-studio") != "quit": raise HTTPException(403, "missing the X-KonfAI-Studio header — stop Studio from its own UI") @@ -876,7 +916,9 @@ async def _after_reply() -> None: await asyncio.sleep(0.3) # let this response leave before the shutdown begins os.kill(os.getpid(), signal.SIGTERM) - asyncio.get_running_loop().create_task(_after_reply()) + # Held: the loop keeps only a weak reference and would collect the task mid-sleep. + _SHUTDOWN_TASKS.add(task := asyncio.get_running_loop().create_task(_after_reply())) + task.add_done_callback(_SHUTDOWN_TASKS.discard) return {"ok": True} diff --git a/studio/tests/test_quit.py b/studio/tests/test_quit.py index d40e12a7..90022aea 100644 --- a/studio/tests/test_quit.py +++ b/studio/tests/test_quit.py @@ -15,8 +15,8 @@ import pytest pytest.importorskip("fastapi") -from konfai_studio import server as bff # noqa: E402 -from starlette.testclient import TestClient # noqa: E402 +from konfai_studio import server as bff +from starlette.testclient import TestClient HEADER = {"X-KonfAI-Studio": "quit"} @@ -70,3 +70,38 @@ def test_a_client_off_this_machine_cannot_stop_the_server(killed: list[int]) -> assert response.status_code == 403 assert not killed + + +@pytest.mark.parametrize( + "header", ["X-Forwarded-For", "X-Forwarded-Host", "X-Forwarded-Proto", "X-Real-IP", "Forwarded"] +) +def test_a_request_through_an_untrusted_proxy_cannot_stop_the_server( + killed: list[int], monkeypatch: pytest.MonkeyPatch, header: str +) -> None: + """Behind the same-host reverse proxy REMOTE.md documents, every peer IS 127.0.0.1. + + Without ``--proxy-headers`` uvicorn does not rewrite it, so the loopback check says local about + a request that came from anywhere. A forwarding header is the tell. + """ + monkeypatch.delenv("KONFAI_STUDIO_PROXY_HEADERS", raising=False) + with TestClient(bff.app, client=("127.0.0.1", 54321)) as client: + response = client.post("/api/quit", json={}, headers={**HEADER, header: "203.0.113.7"}) + + assert response.status_code == 403 + assert "proxy" in response.json()["detail"] + assert not killed + + +def test_a_trusted_proxy_leaves_the_loopback_check_meaningful( + killed: list[int], monkeypatch: pytest.MonkeyPatch +) -> None: + """With the flag, uvicorn rewrites the peer from the header before the endpoint sees it, so the + loopback check is a statement about the real client and a forwarding header is not a reason to + refuse.""" + monkeypatch.setenv("KONFAI_STUDIO_PROXY_HEADERS", "1") + with TestClient(bff.app, client=("127.0.0.1", 54321)) as client: + response = client.post("/api/quit", json={}, headers={**HEADER, "X-Forwarded-For": "127.0.0.1"}) + signalled = _wait_for_signal(killed) + + assert response.status_code == 200 + assert signalled From 59ac9a589d2879c45abcbc0ffa95aba990386042 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 21:21:18 +0200 Subject: [PATCH 4/4] ci: publish the changelog section that was written, and pin what runs The workflow re-rendered the release notes from the commits while the committed section is a draft that was then edited -- squash merges collapse to one line, a subject with no conventional prefix is dropped entirely, and a subject written for a reviewer says nothing to a user. So the file and the release page were guaranteed to differ, which is the opposite of what publish.yml, development.md and CHANGELOG.md all claimed. It takes the committed section verbatim now, and a tag whose section is missing fails the job instead of publishing an empty release. Every action is pinned to a commit SHA -- forty references across seven workflows, each carrying the tag it resolved from so a reader can still tell what is pinned. There is a `contents: read` floor, and no checkout persists the token: every one of these jobs installs and runs code a pull request controls. konfai_mcp_ci fetches the tags, without which setuptools_scm builds the core package as 0.1.dev1 and no bound on it can hold. Pre-release detection searched the whole tag for a letter, so `v1.9.0-backport` was a pre-release and `latest` on Docker Hub moved for any tag at all; both now match the version against a numeric pattern. The image's version tags move from semver to pep440, since `v1.8.0rc1` is a valid PEP 440 version and not a valid SemVer one -- semver rules match nothing there and the image would publish with no version tag. Two CI path filters never fired on the changes they guard: the doc examples are extracted and run from config_guide, and Studio's tests guard wiring that lives in konfai/. The reference pages gain ResampleToReference and REGRID, which were documented at length in the transform guide and absent from the tables an extension author reads, and the changelog gains a "Behaviour changes" section: Median's value and dtype moved under a refactor line, and four new refusals replace things that were being done silently and wrongly. The image installs the wheels the run built instead of pulling them back from PyPI. That drops the poll that waited up to five minutes for an index to serve what had just been uploaded, and it drops ARG KONFAI_PYPI_VERSION -- the one version string committed anywhere, which did not follow the tag and sat at 1.6.0 through all of 1.7. A local build now ships the working tree, and an image of a published release is a `docker pull` of its tag rather than a rebuild. The notes extraction moves out of the workflow into .github/scripts/release_notes.py, with a test. It decides what a release publishes, runs once per tag in the job holding contents: write, and was covered by nothing -- so its first version matched the tag with a `\b`, which also matches at a dot: tagging `v1.8` would have published v1.8.0's notes under its own name, silently, because both are real versions and the text reads fine. The tag classifier follows PEP 440 instead of a numeric shape, and both jobs read it from the same tested function: a numeric shape called the stable v1.8.0.post1 a pre-release, which would publish it as one and hand `latest` to nothing. --- .github/scripts/release_notes.py | 80 ++++++++++++ .github/workflows/commit-hygiene.yml | 7 +- .github/workflows/konfai_apps_ci.yml | 10 +- .github/workflows/konfai_ci.yml | 39 ++++-- .github/workflows/konfai_mcp_ci.yml | 9 +- .github/workflows/konfai_studio_ci.yml | 10 +- .github/workflows/pre-commit.yml | 10 +- .github/workflows/publish.yml | 121 +++++++++++------- CHANGELOG.md | 80 +++++++++++- docker/Dockerfile | 13 +- docker/README.md | 21 +-- docs/source/config_guide/transform.md | 69 +++++++--- docs/source/development.md | 34 +++-- docs/source/reference/api/extension-points.md | 13 +- docs/source/reference/components/index.md | 3 +- .../source/reference/components/transforms.md | 3 +- tests/unit/test_release_notes.py | 116 +++++++++++++++++ 17 files changed, 516 insertions(+), 122 deletions(-) create mode 100644 .github/scripts/release_notes.py create mode 100644 tests/unit/test_release_notes.py diff --git a/.github/scripts/release_notes.py b/.github/scripts/release_notes.py new file mode 100644 index 00000000..4c5142a6 --- /dev/null +++ b/.github/scripts/release_notes.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""The body of a GitHub Release: the committed CHANGELOG section for a tag, verbatim. + +Commitizen drafts that section from the commits, but the file is edited afterwards -- a squash merge +collapses to one line, a subject with no conventional prefix is dropped entirely, and a subject +written for a reviewer says nothing to a user. Rendering the commits again at release time would +publish text nobody reviewed, and the file and the release page would then describe one version +differently. + +A file rather than a heredoc in the workflow because this decides what gets published, and a thing +that decides that should be testable. See ``tests/unit/test_release_notes.py``. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + + +def section_for(changelog: str, tag: str) -> str: + """The body under ``## ``, stripped, without its heading. + + The heading has to END where the tag does. A ``\\b`` would match at the dot too, so a ``v1.8`` + tag would take ``v1.8.0``'s notes and publish them under its own release -- silently, since both + are real versions and the text reads fine. + """ + found = re.search(rf"^## {re.escape(tag)}(?=[ \t]|$)[^\n]*\n(.*?)(?=^## |\Z)", changelog, re.M | re.S) + if found is None or not found.group(1).strip(): + raise SystemExit(f"CHANGELOG.md carries no section for {tag}. Write it before tagging.") + return found.group(1).strip() + "\n" + + +#: PEP 440, as the spec itself writes it. Only the groups this module decides on are named. +_VERSION = re.compile( + r"^v?(?:\d+!)?\d+(?:\.\d+)*" + r"(?P
[-_.]?(?:a|b|c|rc|alpha|beta|pre|preview)[-_.]?\d*)?"
+    r"(?:[-_.]?(?:post|rev|r)[-_.]?\d*|-\d+)?"
+    r"(?P[-_.]?dev[-_.]?\d*)?"
+    r"(?:\+[a-z0-9]+(?:[-_.][a-z0-9]+)*)?$",
+    re.IGNORECASE,
+)
+
+
+def is_prerelease(tag: str) -> bool:
+    """Whether ``tag`` names a pre-release — the versions ``latest`` must not follow.
+
+    PEP 440 and not a numeric shape: a POST-release (``v1.8.0.post1``) is stable and must take
+    ``latest``, while ``v1.8.0rc1`` and ``v1.8.0.dev1`` must not. A tag that does not parse counts as
+    a pre-release, because the alternative is handing ``latest`` to something nobody can classify.
+    """
+    matched = _VERSION.match(tag.strip())
+    return matched is None or bool(matched.group("pre") or matched.group("dev"))
+
+
+def main(argv: list[str]) -> None:
+    if argv[1] == "--prerelease":
+        print("true" if is_prerelease(argv[2]) else "false")
+        return
+    tag, source, destination = argv[1], Path(argv[2]), Path(argv[3])
+    destination.write_text(section_for(source.read_text(encoding="utf-8"), tag), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main(sys.argv)
diff --git a/.github/workflows/commit-hygiene.yml b/.github/workflows/commit-hygiene.yml
index fc8d2882..9f9b598f 100644
--- a/.github/workflows/commit-hygiene.yml
+++ b/.github/workflows/commit-hygiene.yml
@@ -11,12 +11,15 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out full history
-        uses: actions/checkout@v5
+        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
           fetch-depth: 0
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.13"
 
diff --git a/.github/workflows/konfai_apps_ci.yml b/.github/workflows/konfai_apps_ci.yml
index df361c84..3a82cb35 100644
--- a/.github/workflows/konfai_apps_ci.yml
+++ b/.github/workflows/konfai_apps_ci.yml
@@ -44,10 +44,14 @@ jobs:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -57,7 +61,7 @@ jobs:
           pip install -e ".[dev]"
 
       - name: Restore Hugging Face cache
-        uses: actions/cache@v5
+        uses: actions/cache@caa296126883cff596d87d8935842f9db880ef25 # v5
         with:
           path: .cache/huggingface
           key: hf-apps-${{ runner.os }}
diff --git a/.github/workflows/konfai_ci.yml b/.github/workflows/konfai_ci.yml
index 92a1091f..56116fc8 100644
--- a/.github/workflows/konfai_ci.yml
+++ b/.github/workflows/konfai_ci.yml
@@ -1,5 +1,7 @@
 name: konfai_ci
 
+# config_guide and examples/Transform are sources here: the integration tests extract the YAML blocks
+# from those pages and run the example's configs as shipped, so editing either is editing a fixture.
 on:
   push:
     branches: [main]
@@ -8,6 +10,8 @@ on:
       - "konfai/**"
       - "konfai-apps/**"
       - "tests/**"
+      - "docs/source/config_guide/**"
+      - "examples/Transform/**"
       - "pyproject.toml"
       - "README.md"
   pull_request:
@@ -16,6 +20,8 @@ on:
       - "konfai/**"
       - "konfai-apps/**"
       - "tests/**"
+      - "docs/source/config_guide/**"
+      - "examples/Transform/**"
       - "pyproject.toml"
       - "README.md"
 
@@ -29,10 +35,14 @@ jobs:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -51,10 +61,14 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.12"
 
@@ -67,10 +81,14 @@ jobs:
   format:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.12"
 
@@ -83,12 +101,15 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
           fetch-depth: 0
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.12"
 
@@ -108,7 +129,7 @@ jobs:
         run: pytest tests/unit/test_packaging.py -m slow
 
       - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
         with:
           name: konfai-wheel
           path: dist/*.whl
diff --git a/.github/workflows/konfai_mcp_ci.yml b/.github/workflows/konfai_mcp_ci.yml
index 841d1589..c5ae51a4 100644
--- a/.github/workflows/konfai_mcp_ci.yml
+++ b/.github/workflows/konfai_mcp_ci.yml
@@ -28,10 +28,15 @@ jobs:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # konfai-mcp bounds konfai, so the tag has to be reachable: without it setuptools_scm
+          # builds the core package as 0.1.dev1 and no bound on it can hold.
+          fetch-depth: 0
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: ${{ matrix.python-version }}
 
diff --git a/.github/workflows/konfai_studio_ci.yml b/.github/workflows/konfai_studio_ci.yml
index 0e0ed2c3..12250213 100644
--- a/.github/workflows/konfai_studio_ci.yml
+++ b/.github/workflows/konfai_studio_ci.yml
@@ -1,5 +1,7 @@
 name: konfai_studio_ci
 
+# konfai/** is here for the same reason konfai_mcp_ci.yml lists it: Studio's tests guard the wiring
+# between Studio and the workflows, so a change confined to konfai/ is exactly what breaks them.
 on:
   push:
     branches: [main]
@@ -7,12 +9,14 @@ on:
       - ".github/workflows/konfai_studio_ci.yml"
       - "studio/**"
       - "konfai-mcp/**"
+      - "konfai/**"
       - "pyproject.toml"
   pull_request:
     paths:
       - ".github/workflows/konfai_studio_ci.yml"
       - "studio/**"
       - "konfai-mcp/**"
+      - "konfai/**"
       - "pyproject.toml"
 
 jobs:
@@ -24,7 +28,7 @@ jobs:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
           # studio/setup.py pins konfai-mcp to the setuptools_scm version -- tags must be reachable.
           fetch-depth: 0
@@ -32,12 +36,12 @@ jobs:
           persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: ${{ matrix.python-version }}
 
       - name: Set up Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4
         with:
           node-version: "20"
 
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 6c40c8b9..8b39f191 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -13,10 +13,14 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # This job runs code the pull request controls; the token has no business staying
+          # in .git/config while it does.
+          persist-credentials: false
 
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.13"
 
-      - uses: pre-commit/action@v3.0.1
+      - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 280c71f9..fd0f945a 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -5,16 +5,25 @@ on:
     tags:
       - "v*"
 
+# The floor for every job. Only `github_release` needs more, and it raises its own: the jobs that
+# build run third-party code -- pip install, npm lifecycle scripts -- and should hold a token that
+# can do nothing.
+permissions:
+  contents: read
+
 jobs:
   test:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
           fetch-depth: 0
+          # This job installs and runs third-party code; the token has no business staying in
+          # .git/config while it does.
+          persist-credentials: false
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.12"
 
@@ -48,7 +57,7 @@ jobs:
           pip install -e ./studio
 
       - name: Set up Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4
         with:
           node-version: "20"
 
@@ -95,17 +104,21 @@ jobs:
             pkg: "studio"
             build_args: "--wheel"
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
+          # Full history because setuptools_scm derives every package's version from the tag; a
+          # shallow clone would build 0.1.dev-something and publish it.
           fetch-depth: 0
+          # `npm ci` runs lifecycle scripts from the lockfile. Nothing here needs the token.
+          persist-credentials: false
 
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.10"
 
       - name: Set up Node
         if: matrix.name == 'konfai-studio'
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4
         with:
           node-version: "20"
 
@@ -124,7 +137,7 @@ jobs:
           python -m build ${{ matrix.build_args }}
 
       - name: Upload dist
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
         with:
           name: dist-${{ matrix.name }}
           path: ${{ matrix.pkg }}/dist/*
@@ -137,7 +150,7 @@ jobs:
       id-token: write
 
     steps:
-      - uses: actions/download-artifact@v4
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
         with:
           path: dist_all
 
@@ -147,7 +160,7 @@ jobs:
           find dist_all -type f \( -name "*.whl" -o -name "*.tar.gz" \) -print0 | xargs -0 -I {} cp {} dist/
 
       - name: Publish to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+        uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # release/v1
         with:
           packages-dir: dist
 
@@ -158,15 +171,14 @@ jobs:
       contents: write  # to create the release
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
         with:
-          fetch-depth: 0  # cz renders a version's section from the history
           # Nothing here pushes -- the release is created through the API by the action below -- and
           # this is the only job that asks for `contents: write`. Left on the default, checkout would
-          # leave that write-scoped token in .git/config while pip and commitizen run.
+          # leave that write-scoped token in .git/config while the steps below run.
           persist-credentials: false
 
-      - uses: actions/download-artifact@v4
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
         with:
           path: dist_all
 
@@ -176,27 +188,30 @@ jobs:
           find dist_all -type f \( -name "*.whl" -o -name "*.tar.gz" \) -print0 | xargs -0 -I {} cp {} dist/
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
         with:
           python-version: "3.13"
 
-      - name: Install Commitizen
-        run: python -m pip install commitizen==4.13.0
+      # The release notes ARE the committed section for this tag; the script says why, and
+      # tests/unit/test_release_notes.py holds it to it. A tag with no section fails the job.
+      - name: Take the release notes from CHANGELOG.md
+        run: python .github/scripts/release_notes.py "$GITHUB_REF_NAME" CHANGELOG.md /tmp/release-notes.md
 
-      # The release notes ARE the changelog section for this tag, rendered from the same commits the
-      # committed CHANGELOG.md was generated from -- so the release page and the file cannot say
-      # different things about one version.
-      - name: Render the release notes for this tag
-        run: cz changelog "${GITHUB_REF_NAME}" --file-name /tmp/release-notes.md
+      # One rule, shared with publish_docker and tested: `contains(ref_name, 'a')` calls
+      # `v1.9.0-backport` a pre-release, and a numeric shape calls the stable `v1.8.0.post1` one.
+      - name: Classify the tag
+        id: kind
+        run: |
+          echo "prerelease=$(python .github/scripts/release_notes.py --prerelease "$GITHUB_REF_NAME")" >> "$GITHUB_OUTPUT"
 
       - name: Create the GitHub Release
-        uses: softprops/action-gh-release@v2
+        uses: softprops/action-gh-release@3bb12739c298aeb8a4eeaf626c5b8d85266b0e65 # v2
         with:
           name: ${{ github.ref_name }}
           body_path: /tmp/release-notes.md
-          # A tag carrying a pre-release segment publishes as one, so `v1.8.0rc1` never lands on
-          # users following the latest release.
-          prerelease: ${{ contains(github.ref_name, 'a') || contains(github.ref_name, 'b') || contains(github.ref_name, 'rc') || contains(github.ref_name, 'dev') }}
+          # A tag carrying anything past the numeric version publishes as a pre-release, so
+          # `v1.8.0rc1` never lands on users following the latest release.
+          prerelease: ${{ steps.kind.outputs.prerelease == 'true' }}
           files: dist/*
           fail_on_unmatched_files: false
 
@@ -207,47 +222,57 @@ jobs:
       contents: read
 
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5
+        with:
+          # The build runs a Dockerfile over this checkout; the token stays out of it.
+          persist-credentials: false
 
-      - name: Derive package version
+      - name: Classify the tag
         id: version
-        run: echo "value=${GITHUB_REF_NAME#v}" >> "$GITHUB_OUTPUT"
+        run: |
+          echo "value=${GITHUB_REF_NAME#v}" >> "$GITHUB_OUTPUT"
+          # `latest` is what a user gets by asking for nothing, so a pre-release must not take it --
+          # and a POST-release is not one. Same rule as github_release, from the same tested script.
+          echo "prerelease=$(python .github/scripts/release_notes.py --prerelease "$GITHUB_REF_NAME")" >> "$GITHUB_OUTPUT"
+
+      # The image installs the wheels this run built, not what PyPI is serving. The job still waits
+      # on `publish`, so no image exists for a release that failed to upload -- but it no longer
+      # polls an index for something it already has on disk.
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4
+        with:
+          path: dist_all
 
-      - name: Wait for PyPI release to be available
+      - name: Flatten dists into the build context
         run: |
-          version="${{ steps.version.outputs.value }}"
-          for attempt in $(seq 1 30); do
-            if curl -fsSL "https://pypi.org/pypi/konfai/${version}/json" >/dev/null; then
-              exit 0
-            fi
-            echo "konfai ${version} is not available on PyPI yet (attempt ${attempt}/30)"
-            sleep 10
-          done
-          echo "konfai ${version} was not available on PyPI after 5 minutes"
-          exit 1
+          mkdir -p dist
+          find dist_all -type f -name "*.whl" -print0 | xargs -0 -I {} cp {} dist/
+          test -n "$(ls dist/konfai-*.whl)" && test -n "$(ls dist/konfai_apps-*.whl)"
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3
 
       - name: Extract Docker metadata
         id: meta
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5
         with:
           images: vboussot/konfai
+          # pep440 and not semver: this project tags the Python way, and `v1.8.0rc1` is a valid
+          # PEP 440 version but not a valid SemVer one -- semver rules match nothing there and the
+          # image would publish with no version tag at all.
           tags: |
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-            type=semver,pattern={{major}}
-            type=raw,value=latest
+            type=pep440,pattern={{version}}
+            type=pep440,pattern={{major}}.{{minor}}
+            type=pep440,pattern={{major}}
+            type=raw,value=latest,enable=${{ steps.version.outputs.prerelease == 'false' }}
 
       - name: Log in to Docker Hub
-        uses: docker/login-action@v3
+        uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
       - name: Build and publish Docker image
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6
         with:
           context: .
           file: docker/Dockerfile
@@ -255,7 +280,5 @@ jobs:
           platforms: linux/amd64
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
-          build-args: |
-            KONFAI_PYPI_VERSION=${{ steps.version.outputs.value }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0ebee3ec..fa86fbde 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,29 +1,58 @@
 # Changelog
 
-Generated from the commit history by [Commitizen](https://commitizen-tools.github.io/commitizen/).
-Each version's section is what the GitHub Release for that tag carries. Regenerate with:
+Drafted from the commit history by [Commitizen](https://commitizen-tools.github.io/commitizen/), then
+edited. Each version's section is what the GitHub Release for that tag carries.
 
 ```bash
-cz changelog --start-rev v1.5.8
+cz changelog --unreleased-version vX.Y.Z --start-rev v1.5.8   # a DRAFT, not the final file
 ```
 
 Conventional Commits started at `v1.5.9`; rendering further back produces empty version headings.
 
+The draft is a starting point, not the answer. It sees only commit subjects, so a squash merge
+collapses to one line, a subject with no conventional prefix is dropped entirely, and a subject
+written for a reviewer ("what reviewing the data surface turned up") tells a reader nothing. Take the
+draft, then say what a user of the package gets that they did not have -- and re-read the section
+against the commits that landed *after* you drafted it. Running the command over a section already
+written replaces it.
+
 ## v1.8.0 (2026-08-04)
 
 ### ✨ Features
 
-- **data**: give a chain a per-component statistic, and a shape update that uses it
+- **transform**: resample a case onto the grid of a declared reference
+- **transform**: compose the grid change and the warp into one pass
+- **data**: give a chain a per-component statistic
 - **transform**: let Warp read the bound the fields recorded, with max_displacement: auto
 - **studio**: show a transform run, and point Browse at the data rather than its log
 - **mcp**: let an agent plan and run a transform, and read the table for what it accepts
 - **transform**: a fifth workflow that reads a dataset, applies a chain, and writes it
+- **transform**: give ResampleToReference an `interpolation`, as Warp already had
+- **data**: Vote, the reduction operator that folds segmentations without inventing a label
 - **data**: declare an OME-Zarr pyramid from a Write, and let a field carry its own bound
 - **impact-reg**: seed the rigid from the centre of mass, not only the frame
 - **studio**: bundle icons through the app interface, and a way to stop Studio (#75)
+- **examples**: a Transform example -- a template folded out of a cohort, and drawn copies of a case
 
 ### 🐛 Bug Fixes
 
+- **transform**: sample a label map by nearest on the warped path too, instead of blending labels
+- **transform**: check Warp's declared bound on the whole-volume path, not only the streamed one
+- **transform**: build Warp's grid on the volume's device, so a GPU-resident case does not raise
+- **data**: cap an OME-Zarr chunk instead of taking the writer's whole trailing plane
+- **data**: keep the store's chunking when a Write appends pyramid levels
+- **data**: refuse a statistic after a Reduce that an earlier post stage invalidates
+- **data**: budget a reduction for what its operator allocates over the buffer it holds
+- **data**: refuse a geometry key `grid: strict` cannot compare, instead of skipping it
+- **data**: stop a chained sweep at the first failure, so the recorded reason is the cause
+- **data**: record the landed state only for a plan that holds
+- **data**: restore the CUDA generators a draw's seeding touched
+- **transform**: name the argument the stage actually takes in a field refusal
+- **transform**: count a chain's Reduce markers where its Expand markers were already counted
+- **konfai-mcp**: require the konfai that has the module it imports at load
+- **studio**: refuse /api/quit when a proxy header arrives and the peer cannot be trusted
+- **examples**: install scikit-image where SSIM is evaluated, and ask for a GPU that exists
+- **transform**: hold the auto bound to the cohort, and plan a resampled case on the grid it lands on
 - **transform**: refuse a stage key that names something not a class
 - **transform**: print the plan's reduction and dropped lines once
 - **transform**: give refusals their true remedy and name
@@ -62,6 +91,49 @@ Conventional Commits started at `v1.5.9`; rendering further back produces empty
 - **data**: share the fallback constants and the Welford kernel
 - **data**: move the reduction operators out of the predictor, into one shared vocabulary
 - **data**: give the resolved memory budget a type that knows its own scope
+- **data**: name what patching shares with the transform workflow, instead of reaching into privates
+- **transform**: state each sampling rule once, over the two gathers that share the one arithmetic
+- **ci**: pin every action to a commit SHA, and take the release notes from the committed changelog
+
+### ⚠️ Behaviour changes
+
+No YAML key, class or default was renamed, but the following answer differently for a config that
+was not touched. Several are new refusals: what they refuse was being done before, silently and
+wrongly.
+
+- **`reduction: Median` returns a different number on an even count.** `torch.median` hands back the
+  lower of the two middle values, which over two tensors is the element-wise minimum; it now averages
+  the middle pair as `numpy.median` does. A 2-model ensemble or a 2-draw TTA that reduced to `1.0`
+  over `[1.0, 3.0]` now reduces to `2.0`. On an odd count the two agree.
+- **`Mean` and `Median` widen an integer input to float32**, including the single-tensor path that
+  previously returned the tensor untouched. Rounding an average back onto an integer grid is a wrong
+  number, not a narrower one -- but a `uint8` prediction output now lands as float32, four times the
+  bytes on disk, and a downstream stage sees the wider dtype.
+- Together those two make **`Median` the wrong operator for a label map**: it can answer with a
+  label that was in no input (over 1 and 5 it gives 3), and over exactly two cases it *is* `Mean`.
+  Fold segmentations with the new **`Vote`**, which picks the label the most cases agree on and
+  keeps the dtype.
+- **A whole-volume statistic after a `Reduce` is refused when an earlier post stage changes the
+  values.** The stat pass measures the fold, so `[Reduce, Clip, Normalize]` normalised by the
+  *unclipped* statistic and wrote a volume its own header did not describe. The per-case planner
+  already refused this; the reduction now does too. Split the chain at the value-changing stage.
+- **Two source groups may no longer declare the same destination group name.** The name is the key
+  everything downstream indexes by, so the second chain used to be built and then silently dropped;
+  it is now refused when the dataset is prepared, in every workflow. Give the chains distinct names
+  and say `group:` on each `Write` to store both under one group.
+- **A `uint8` label map resampled through a `field` now takes the nearest voxel.** The warped path
+  consulted no interpolation rule, so it blended labels and truncated: over a source holding
+  `{0, 100}` it returned 29, 79 and 99. Anything stored as another integer dtype needs
+  `interpolation: nearest` spelled out -- a dtype cannot tell a label map from a CT.
+- **A chain declaring two `Reduce` markers is refused at parse time**, naming the cardinality
+  marker, where the second one used to fall past the split and be reported as an ordinary stage
+  reading across space.
+- **`grid: strict` refuses a geometry key no header records**, rather than skipping the comparison
+  it promised. A missing `Direction` is a flip that shows in neither extent nor spacing. Use
+  `grid: shape_only` for a cohort that means to fold on extent alone.
+- **New OME-Zarr stores are chunked to a size a reader can open.** A streamed writer declaring the
+  whole trailing plane produced a chunk of a gigabyte at 2048², past what zarr holds in one buffer
+  at 4096². Existing stores keep their own chunking; zarr is self-describing.
 
 ## v1.7.0 (2026-07-29)
 
diff --git a/docker/Dockerfile b/docker/Dockerfile
index d93f3a4a..e045bcbc 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -6,9 +6,9 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
     NVIDIA_VISIBLE_DEVICES=all \
     NVIDIA_DRIVER_CAPABILITIES=compute,utility
 
+# What `konfai[...]` pulls: the imaging backends (SimpleITK, h5py, pydicom, zarr, ngff-zarr, dask).
 ARG KONFAI_EXTRAS=imaging
 ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
-ARG KONFAI_PYPI_VERSION=1.8.0
 
 WORKDIR /tmp
 
@@ -18,10 +18,19 @@ RUN apt-get update \
 
 COPY docker/entrypoint.sh /usr/local/bin/konfai-entrypoint
 
+# The wheels this image ships, taken from the build context rather than from PyPI. The version is
+# then whatever was built -- nothing to keep in sync with the tag, and no waiting for an index to
+# serve what was uploaded a moment ago. `python -m build` puts them here; see docker/README.md.
+COPY dist/ /tmp/dist/
+
 RUN chmod +x /usr/local/bin/konfai-entrypoint \
     && pip install --upgrade pip setuptools wheel \
     && pip install --index-url "${TORCH_INDEX_URL}" --extra-index-url https://pypi.org/simple torch \
-    && pip install "konfai[${KONFAI_EXTRAS}]==${KONFAI_PYPI_VERSION}" "konfai-apps==${KONFAI_PYPI_VERSION}"
+    # Resolved through a variable because the shell would read the `[...]` of a glob as a character
+    # class. `konfai-*` cannot match a sibling: a wheel spells `konfai-apps` as `konfai_apps`.
+    && konfai_wheel="$(ls /tmp/dist/konfai-*.whl)" \
+    && pip install "${konfai_wheel}[${KONFAI_EXTRAS}]" /tmp/dist/konfai_apps-*.whl \
+    && rm -rf /tmp/dist
 
 WORKDIR /workspace
 
diff --git a/docker/README.md b/docker/README.md
index 33979805..f89021a4 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -27,24 +27,24 @@ docker run --rm vboussot/konfai
 
 ## Local Build
 
-Build the default image from this repository. It installs the KonfAI release the
-`KONFAI_PYPI_VERSION` argument names — currently `1.8.0`:
+The image installs the wheels it finds in `dist/`, so build them first. It then ships the
+working tree — there is no version to name anywhere, and none to keep in sync:
 
 ```bash
+python -m build --wheel --outdir dist .
+python -m build --wheel --outdir dist ./konfai-apps
 docker build -f docker/Dockerfile -t konfai .
 ```
 
-Build a different KonfAI version from PyPI:
+The release workflow does exactly this with the wheels the tag built, which is why the
+published image never waits on an index to serve what was just uploaded.
+
+**For an image of a published release, pull it** — every version has its own tag:
 
 ```bash
-docker build -f docker/Dockerfile \
-  --build-arg KONFAI_PYPI_VERSION=1.7.0 \
-  -t konfai .
+docker pull vboussot/konfai:1.7.0
 ```
 
-The published images never rely on that default: the release workflow passes the tag's
-version as `--build-arg`, so the image for `vX.Y.Z` always installs `X.Y.Z`.
-
 Build with additional optional dependencies:
 
 ```bash
@@ -146,4 +146,5 @@ docker run --rm -it -p 8000:8000 \
 
 - The image is intended for CLI workflows executed from a mounted workspace.
 - The default image is GPU-oriented; use a custom `TORCH_INDEX_URL` if you want a CPU-only variant.
-- For reproducibility, pin the PyPI version with `KONFAI_PYPI_VERSION` when rebuilding locally.
+- A local rebuild ships the working tree, since it installs the wheels in `dist/`. For a
+  published release, pull its tag rather than rebuilding it.
diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md
index c452d201..e889f859 100644
--- a/docs/source/config_guide/transform.md
+++ b/docs/source/config_guide/transform.md
@@ -1,9 +1,10 @@
 # Transform configuration
 
-Transform configuration lives under the `Transformer` root object. It is the one
-workflow with **no model**: it reads a dataset, applies a chain of transforms,
-and writes the result. If you never train anything, this is the only page you
-need.
+Transform configuration lives under the `Transformer` root object. It runs **no
+model** — nor does `EVALUATION`, and the difference is what comes out: an
+evaluation measures, this one **makes**. It reads a dataset, applies a chain of
+transforms, and writes a dataset. If you never train anything, this is the only
+page you need.
 
 ```yaml
 Transformer:
@@ -47,8 +48,10 @@ workflow emits no scalars.
 
 ## Read the plan before you read anything else
 
-Every run prints its plan first, and writes it to
-`./Transforms//plan.txt`. Nothing is written until the plan is accepted.
+Every run prints its plan first, and a run that proceeds writes it to
+`./Transforms//plan.txt`, next to an `outputs.json` naming where each
+chain's deliverable lands. No *data* is written until the plan is accepted, and
+`--plan` prints and stops without leaving either file behind.
 
 ```text
 [Transformer] plan over 1 rank(s) | per-rank budget 7.45 GiB ('8G') | fallback working set
@@ -143,11 +146,12 @@ a byte is read:
   streaming re-reads the source while writing, so an in-place transform would
   read its own half-written output;
 - a `Save` with no `dataset` of its own, which would write next to the source;
-- **any key this page does not document — a stage's arguments included.** A
-  typo'd `memory_budge:` or `Clip: {min_val: …}` would otherwise be ignored and
-  its default used silently; here it is an error naming the exact path and the
-  legal keys. (A stage that takes `**kwargs` or resolves nowhere is left to the
-  loader's own error.)
+- **any structural key the grammar does not list, and any stage argument its
+  constructor does not take.** A typo'd `memory_budge:` or
+  `Clip: {min_val: …}` would otherwise be ignored and its default used
+  silently; here it is an error naming the exact path and the legal keys. (A
+  stage that takes `**kwargs` or resolves nowhere is left to the loader's own
+  error, and the contents of `subset:` are not walked.)
 
 ## Fields
 
@@ -215,18 +219,29 @@ That writes **one** entry named `template`, whatever the cohort's size.
 
 | Field | Default | Effect |
 | --- | --- | --- |
-| `operator` | `Median` | A classpath resolved against `konfai.data.reduction`: `Mean`, `Median`, `Concat`, or your own `Reduction` subclass. An operator's own parameters go in the same mapping, next to `operator`. |
+| `operator` | `Median` | A classpath resolved against `konfai.data.reduction`: `Mean`, `Median`, `Vote`, `Concat`, or your own `Reduction` subclass. An operator's own parameters go in the same mapping, next to `operator`. |
 | `output` | — | **Required**: the entry name the single result is written under. |
 | `grid` | `strict` | How much agreement between members is demanded before a byte is read (below). |
 | `grid_tolerance` | `1e-6` | The tolerance `strict` compares geometry within. |
 | `provenance` | `true` | Record the operator and the folded case list in the output's header — a cohort that silently changed between two runs writes a different volume under the same name, and nothing about the output would look wrong. |
 
 **Operators.** `Mean` folds one case at a time, so its working set is two
-regions whatever N is. `Median` needs every case per region — N + 1 resident
-regions, which is what `memory_budget` sizes and refuses. `Concat` puts the
-cases side by side: the output carries `N × C` channels. A custom operator must
-declare `voxel_local = True` — one that reads across space cannot stream and is
-refused outright.
+regions whatever N is. `Median` needs every case per region, and stacks and
+sorts them on top — the plan says how many regions that is, and `memory_budget`
+sizes and refuses against it. `Concat` puts the cases side by side: the output
+carries `N × C` channels. A custom operator must declare `voxel_local = True` —
+one that reads across space cannot stream and is refused outright. It should
+also declare `working_multiple` if it allocates over the buffer it is handed,
+or the plan promises a working set the run exceeds.
+
+```{warning}
+`Mean` and `Median` are for intensities. Both answer with values that were in no
+input — the median of labels `1` and `5` is `3`, a different structure — and
+both widen an integer input to float32. Over exactly two cases `Median` *is*
+`Mean`, so the robustness the name promises is not there. Fold segmentations
+with **`Vote`**, which takes the label the most cases agree on, keeps the input
+dtype, and breaks a tie toward the smallest label so the fold is reproducible.
+```
 
 **`grid` decides what counts as "the same space"**, compared on the grid each
 case's chain *lands* on (a `Resample` before the `Reduce` counts):
@@ -344,6 +359,18 @@ taps clamped to the buffer, nearest by round-half-up, and `fill` wherever the
 reference grid reaches past the case.
 ```
 
+**Label maps.** Left unset, `interpolation` is read off the dtype: `uint8` takes
+the nearest voxel, everything else is interpolated. A dtype cannot decide this
+on its own — a CT is `int16` and so is nothing else about it — so a label map
+stored as anything but `uint8` must say so:
+
+```yaml
+ResampleToReference: {entry: case_0, group: Labels, interpolation: nearest}
+```
+
+Getting it wrong is silent. Two labels blended give a third that was never in
+the source, the dtype is unchanged, and the result is still a label map.
+
 **What it refuses**, rather than write something plausible and wrong:
 
 - a case or a reference carrying no `Origin` / `Spacing` / `Direction` — without
@@ -402,6 +429,14 @@ leave: a draw that permutes axes hands the next stage its own extent, and a
 resample between two draws is seen by the second. That is the same contract a
 transform has — a draw is a stage, not a separate phase.
 
+```{warning}
+A bare name resolves against `konfai.data.transform` **first**, and only then
+against `konfai.data.augmentation`. `Flip`, `Permute`, `Mask` and `Foreign`
+exist in both, so `Flip: {f_prob: [0.33, 0.33, 0.33]}` binds the deterministic
+*transform* and fails on an argument it does not take. Spell the draw out:
+`konfai.data.augmentation:Flip`.
+```
+
 `pattern` is a `str.format` template and **both** tokens are required: `{name}`
 keeps cases apart, `{a}` (1-based) keeps a case's copies apart. A pattern missing
 either is refused at parse time, because every copy would otherwise overwrite the
diff --git a/docs/source/development.md b/docs/source/development.md
index 199e80ae..c166b3ff 100644
--- a/docs/source/development.md
+++ b/docs/source/development.md
@@ -238,24 +238,32 @@ the framework, the two sibling packages, and every published App.
 ### Cutting a release
 
 Versions are **tag-derived** — `setuptools_scm` reads the tag, so no *package*
-version is committed anywhere that could drift from it. `CHANGELOG.md` is generated
-from the commit history, and the publish workflow renders the section for the tag it
-is running on as the GitHub Release body, so the file and the release page cannot
-describe a version differently.
-
-```{important}
-One version string **is** committed and does not follow the tag:
-`ARG KONFAI_PYPI_VERSION` in `docker/Dockerfile`. The published image is unaffected —
-the workflow passes the tag's version as `--build-arg` — but it is what a local
-`docker build` without that flag installs, so bump it when you cut a release. It sat
-at `1.6.0` through 1.7.
+version is committed anywhere that could drift from it. `CHANGELOG.md` is drafted
+from the commit history and then edited, and the publish workflow takes the committed
+section for the tag it is running on **verbatim** as the GitHub Release body, so the
+file and the release page cannot describe a version differently. A tag whose section
+is missing — **or present but empty** — fails the job rather than publishing a release
+with nothing in it. A tag carrying a pre-release segment (`v1.8.0rc1`, `v1.8.0.dev1`)
+publishes as a pre-release and does not take `latest`; a post-release (`v1.8.0.post1`)
+is stable and does.
+
+```{note}
+No version string is committed anywhere, the Docker image included: it installs the
+wheels found in `dist/`, which the release workflow fills from the tag's own build. That
+also means the image never waits on PyPI to serve what the run just uploaded.
 ```
 
 That order matters: **the changelog is written before the tag**, because the
-workflow renders what the history already says.
+workflow publishes what the file already says.
+
+The generated draft is a starting point, not the answer. It sees commit subjects
+only, so a squash merge collapses to one line, a subject with no conventional
+prefix is dropped, and a subject written for a reviewer tells a reader nothing.
+Take the draft, then say what a *user* of the package gets that they did not have —
+and re-read it against anything that landed after you drafted it.
 
 ```bash
-# 1. Write the section for the version you are about to cut
+# 1. Draft the section for the version you are about to cut, then edit it
 uvx --from commitizen cz changelog --unreleased-version vX.Y.Z --start-rev v1.5.8
 
 # 2. Commit it
diff --git a/docs/source/reference/api/extension-points.md b/docs/source/reference/api/extension-points.md
index 2becbeb6..d1891b23 100644
--- a/docs/source/reference/api/extension-points.md
+++ b/docs/source/reference/api/extension-points.md
@@ -124,11 +124,20 @@ halo of a geometric draw is that draw's own.
 | `HALO` | bounded neighbourhood, radius `halo` per axis in array order (Z, Y, X) | nothing — the dispatcher reads the enlarged region and crops |
 | `ORIENTATION` | flip or permute | `stream_region_source` |
 | `CROP` | source region is the target region translated | `stream_region_source` |
-| `GLOBAL_STAT` | needs whole-volume statistics, `stat_keys` a subset of Min/Max/Mean/Std | nothing — the dispatcher seeds the statistic from disk |
-| `RESCALE` | resample | subclass `Resample` |
+| `GLOBAL_STAT` | needs whole-volume statistics, `stat_keys` a subset of Min/Max/Mean/Std (or their `…PerChannel` forms) | nothing — the dispatcher seeds the statistic from disk |
+| `RESCALE` | resample by a ratio | subclass `Resample` |
+| `REGRID` | resample onto a grid declared elsewhere — a stored reference, not a ratio — so the source region is computed from the two geometries | subclass `Resample`; declare a halo when a displacement field is composed in |
 | `SLAB` | a per-voxel value map plus a side effect that needs the slab's place in the volume | `stream_slab(name, tensor, region, spatial_shape, cache_attribute)`, and optionally `stream_abort`. The **read** dispatcher has no slab context and treats it as `WHOLE_VOLUME`; the gain is on the write side |
 | `WHOLE_VOLUME` | needs the whole volume | nothing — this is the default |
 
+A sampler you write yourself should take its arithmetic from
+`konfai.data.transform`'s own: `sampling_dtype` (what to accumulate a weighted
+sum in — an integer input and a CPU half both need float32), `nearest_index`
+(ITK's round-half-up, which `torch.round` and `F.interpolate` each get wrong in
+their own way) and `window_index` (a global source index clamped into the
+sub-region that was actually read). A sampler that is only *nearly* the same as
+the ones shipped here makes every comparison against them a negotiation.
+
 A declaration is bound by three rules:
 
 - **read-only** — never write to `cache_attribute`. A declaration is made once for
diff --git a/docs/source/reference/components/index.md b/docs/source/reference/components/index.md
index fc0a44c7..bf4e20d8 100644
--- a/docs/source/reference/components/index.md
+++ b/docs/source/reference/components/index.md
@@ -72,8 +72,7 @@ ways to get the exhaustive list for any component:
    | learning-rate schedulers | `torch.optim.lr_scheduler` **first**, then `konfai/metric/schedulers.py` |
    | loss-weight schedulers | `konfai/metric/schedulers.py` only |
    | patch blending (`patch_combine`) | `konfai/data/patching.py` |
-   | prediction reduction | `konfai/predictor.py` |
-   | case reduction | `konfai/data/reduction.py` |
+   | reduction operators (a prediction's copies *and* a cohort's cases) | `konfai/data/reduction.py` |
 
    So a bare `StepLR` resolves *outside* KonfAI, in torch.
 
diff --git a/docs/source/reference/components/transforms.md b/docs/source/reference/components/transforms.md
index 040bf485..4c3cd1d9 100644
--- a/docs/source/reference/components/transforms.md
+++ b/docs/source/reference/components/transforms.md
@@ -100,6 +100,7 @@ until it declares otherwise.
 | `ResampleToResolution` | Resample to a target voxel spacing (per-axis `<0` = keep). | `spacing=[1,1,1], inverse=True` | **yes** | **yes** | **yes** — resampled from the source region |
 | `ResampleToShape` | Resample to a target shape (per-axis `0/<0` = keep). | `shape=[100,256,256], inverse=True` | **yes** | **yes** | **yes** — resampled from the source region |
 | `ResampleTransform` | Warp by stored SimpleITK transforms read from the dataset. | `transforms`, `inverse=True` | no | no | no — nothing bounds how far the stored displacement reaches |
+| `ResampleToReference` | Resample onto the grid of a **declared reference case** — extent, spacing, origin and direction — so a cohort meets on a grid that is one of its own rather than an invented one. An optional `field` composes the grid change and the warp into one pass, so the intermediate never exists. `interpolation` left unset is nearest for `uint8` and linear otherwise; declare it for a label map stored as anything else. | `entry` (required), `group=None`, `dataset=None`, `field=None`, `field_group=None`, `max_displacement=0.0`, `fill=0.0`, `interpolation=None`, `inverse=True` | **yes** | **yes** — the grid change alone; a composed field is not inverted | **yes** — declares `REGRID`; halo = bound / spacing when a field is composed in |
 | `Warp` | Resample a case through a displacement field on the **same grid** — the shape update of an atlas build. Warping onto a different grid is a resample too, and is not this stage. The declared bound is CHECKED per component against every region read, so a field that exceeds it raises instead of sampling zeros (which would read as a dark rim and nothing else). | `field` (required), `group=None`, `max_displacement=0.0`, `interpolation="linear"` | no | no | **yes** — halo = bound / spacing. `max_displacement: auto` reads the bound the fields recorded when KonfAI wrote them (OME-Zarr only); with no bound at all it declares whole-volume and says which one is missing |
 | `Canonical` | Reorient to canonical direction (3-D); updates Origin/Direction. | `inverse=True` | **yes** — a remap that transposes extents moves the patch grid | **yes** | **yes** — when the case's direction is a signed axis permutation; no on an oblique one (it is resampled) |
 | `Permute` | Permute spatial axes. `dims` is a pipe-separated axis list. | `dims="1\|0\|2", inverse=True` | **yes** | **yes** | **yes** — index remap |
@@ -143,7 +144,7 @@ Operate on a stacked `[N, …]` ensemble axis (prediction post-processing).
 | `Statistics` | Records ImageMin/Max/Mean/Std to the attribute cache and returns the tensor unchanged (feeds the perceptual criteria `SAM_Perceptual`, `IMPACTSynth`, `IMPACTReg`). Order in the transform list matters. | no‡ |
 | `Save` | Writes the preprocessed volume to a cache dataset and passes the tensor through. Once that cache exists it is read instead, and the transforms before it are skipped. A group written this way is also readable, within the same run, by anything that names it (`Mask: {path: }`), including when the write comes from a loader worker — every backend. On `h5` the reader and the writer share one store, which HDF5 does not define for concurrent access without SWMR: the entry is seen, but a read racing a write can raise. One file per case (`mha`/`nii`/…) has no such window. | no — it needs the whole volume to write |
 | `Write` | A `Save` that is a **deliverable**: same boundary semantics, but `dataset` has no default, so a bare `Write:` fails at config time instead of writing into the source tree. The TRANSFORM workflow plans, resumes and reports on `Write` stages and requires every chain to end with one; a `Save` between them is an opportunistic milestone. Args: `dataset` (required), `group=None`, `scale_factors=None`, `downsample_method=None`. | region-writes where the backend allows (uncompressed `mha` / `h5` / `omezarr`) |
-| `Reduce` | Folds every case of a group into one volume at fixed voxel — the stage that makes a chain N-to-1. `operator` is a classpath resolved against `konfai.data.reduction` (`Mean`, `Median`, `Concat`, or your own `Reduction`); `output` is required. `grid` is `strict` / `shape_only` / `reference:`. Args: `operator="Median", output="", grid="strict", grid_tolerance=1e-6, provenance=True`. **TRANSFORM only** — applied to one case it raises. | driven by the reduction engine, one region at a time |
+| `Reduce` | Folds every case of a group into one volume at fixed voxel — the stage that makes a chain N-to-1. `operator` is a classpath resolved against `konfai.data.reduction` (`Mean`, `Median`, `Vote` for label maps, `Concat`, or your own `Reduction`); `output` is required. `grid` is `strict` / `shape_only` / `reference:`. Args: `operator="Median", output="", grid="strict", grid_tolerance=1e-6, provenance=True`. **TRANSFORM only** — applied to one case it raises. | driven by the reduction engine, one region at a time |
 | `Expand` | Turns one case into `nb` copies at a declared point of the chain — `Reduce`'s mirror (1-to-N). Stages before it run once per case, stages after it once per copy, and the draws after it are ordinary stages. `pattern` is a `str.format` template and **both** `{name}` and `{a}` are required. Args: `nb=2, pattern="{name}_{a:02d}", seed=None`. **TRANSFORM only** — applied to one case it raises. | per copy, sharing one read pass when every per-copy stage is pointwise |
 | `KonfAIInference` | Run a nested KonfAI app inference in a spawned subprocess. Needs `konfai-apps` and `num_workers: 0`; defaults to a specific HF repo. | no‡ |
 
diff --git a/tests/unit/test_release_notes.py b/tests/unit/test_release_notes.py
new file mode 100644
index 00000000..17987b82
--- /dev/null
+++ b/tests/unit/test_release_notes.py
@@ -0,0 +1,116 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""What a tag publishes as its release body.
+
+This runs once per release, on a tag, in the job that holds ``contents: write`` -- so its failures
+are expensive and rare, which is exactly the shape of code that goes unchecked. The `v1.8` case
+below is not hypothetical: the first version of this extraction used ``\\b`` and would have taken
+v1.8.0's notes for it, silently, because both are real versions and the text reads fine.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+_SCRIPT = Path(__file__).resolve().parents[2] / ".github" / "scripts" / "release_notes.py"
+_CHANGELOG = Path(__file__).resolve().parents[2] / "CHANGELOG.md"
+
+
+def _section_for():
+    """The shipped script, loaded from where the workflow runs it."""
+    assert _SCRIPT.is_file(), f"the workflow runs a script that is not there: {_SCRIPT}"
+    spec = importlib.util.spec_from_file_location("release_notes", _SCRIPT)
+    assert spec and spec.loader
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.section_for
+
+
+_CHANGELOG_SAMPLE = """# Changelog
+
+## v1.8.0 (2026-08-04)
+
+### Features
+
+- something a user gets
+
+## v1.7.0 (2026-07-29)
+
+### Features
+
+- an older thing
+"""
+
+
+def test_it_takes_the_section_the_tag_names() -> None:
+    body = _section_for()(_CHANGELOG_SAMPLE, "v1.8.0")
+    assert "something a user gets" in body
+    assert "an older thing" not in body, "the section stops at the next heading"
+    assert not body.startswith("## "), "the heading belongs to the release title, not its body"
+
+
+@pytest.mark.parametrize("tag", ["v1.8", "v1", "v1.8.0rc1", "v9.9.9", ""])
+def test_a_tag_with_no_section_of_its_own_fails_the_job(tag: str) -> None:
+    """Refusing is the point: the alternative is an empty release, or -- for a prefix like ``v1.8``
+    -- a release carrying someone else's notes under its own name."""
+    with pytest.raises(SystemExit, match="carries no section"):
+        _section_for()(_CHANGELOG_SAMPLE, tag)
+
+
+def test_the_real_changelog_answers_for_the_version_being_cut() -> None:
+    """The file in the tree, not a fixture: this is what the next tag will actually publish."""
+    body = _section_for()(_CHANGELOG.read_text(encoding="utf-8"), "v1.8.0")
+    assert body.strip(), "v1.8.0 has no notes to publish"
+    assert body.endswith("\n")
+    assert "## v1.7.0" not in body, "the section must not run into the previous release"
+
+
+def _is_prerelease():
+    """The shipped classifier, loaded from where both workflow jobs run it."""
+    spec = importlib.util.spec_from_file_location("release_notes", _SCRIPT)
+    assert spec and spec.loader
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.is_prerelease
+
+
+@pytest.mark.parametrize(
+    "tag,prerelease",
+    [
+        ("v1.8.0", False),
+        ("v1.8.0.post1", False),  # a POST-release is stable and must take `latest`
+        ("v1.8.0-1", False),  # PEP 440's implicit post-release spelling
+        ("v1.8.0rc1", True),
+        ("v1.8.0a1", True),
+        ("v1.8.0b2", True),
+        ("v1.8.0.dev1", True),
+        ("v1.8.0rc1.post1", True),  # a post of a pre is still a pre
+        ("v1.9.0-backport", True),  # unparseable: never hand `latest` to what nobody can classify
+        ("nawak", True),
+    ],
+)
+def test_the_tag_classifier_follows_pep_440(tag: str, prerelease: bool) -> None:
+    """What `latest` follows, on Docker Hub and on the releases page.
+
+    Neither a substring search nor a numeric shape: `contains(tag, 'a')` calls `v1.9.0-backport` a
+    pre-release, and `^[0-9.]+$` calls the stable `v1.8.0.post1` one — which would then be published
+    as a pre-release and lose `latest` to nothing.
+    """
+    assert _is_prerelease()(tag) is prerelease