diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 8e177c59..535d4ad1 100755 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -176,12 +176,17 @@ def __init__( self.patch_size = patch_size self.patch_combine = patch_combine self.batch = batch + self._filled = 0 def add_layer(self, index: int, layer: torch.Tensor) -> None: + if self._layer_accumulator[index] is None: + self._filled += 1 self._layer_accumulator[index] = layer def is_full(self) -> bool: - return len(self.patch_slices) == len([v for v in self._layer_accumulator if v is not None]) + # O(1): a running counter avoids re-scanning every slot after each added patch + # (the completion check ran once per patch, i.e. O(P^2) per case). + return self._filled == len(self.patch_slices) def assemble(self) -> torch.Tensor: n = 2 if self.batch else 1 @@ -195,7 +200,8 @@ def assemble(self) -> torch.Tensor: result = torch.zeros( (list(reference.shape[:n]) + list(max([[v.stop for v in patch] for patch in self.patch_slices]))), dtype=reference.dtype, - ).to(reference.device) + device=reference.device, + ) # Overlap blending weights each patch (edge bands < 1 so interior overlaps sum to unity). # A voxel covered by fewer patches (a volume border without whole-image padding) would sum # to < 1 and come out darkened (x0.5 edges, x0.25 corners), so divide by the accumulated @@ -224,6 +230,7 @@ def assemble(self) -> torch.Tensor: result = result[tuple([slice(None, None)] + [slice(0, s) for s in self.shape])] self._layer_accumulator.clear() + self._filled = 0 return result diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 0f512c2b..79ca2d30 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -163,8 +163,17 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) min_value = float(min_value) max_value = float(max_value) - tensor[torch.where(tensor.float() < min_value)] = min_value - tensor[torch.where(tensor.float() > max_value)] = max_value + # Fast path: one fused in-place clamp instead of two float()-copy + where-scatter passes. + # Restricted to float32 (integer tensors reject float bounds; float16/float64 would compare + # at a different precision than the legacy float()-cast scatter) and to non-NaN bounds: a + # NaN bound — from a dynamic min/max/percentile over data containing NaN — makes clamp_ + # propagate NaN to the whole tensor, whereas the legacy scatter no-ops on it (NaN + # comparisons are False). All other cases keep the exact original behaviour byte-for-byte. + if tensor.dtype == torch.float32 and min_value == min_value and max_value == max_value: + tensor.clamp_(min=min_value, max=max_value) + else: + tensor[torch.where(tensor.float() < min_value)] = min_value + tensor[torch.where(tensor.float() > max_value)] = max_value if self.save_clip_min: cache_attribute["Min"] = min_value if self.save_clip_max: diff --git a/konfai/predictor.py b/konfai/predictor.py index cf8ea25c..dce45e0f 100644 --- a/konfai/predictor.py +++ b/konfai/predictor.py @@ -712,6 +712,10 @@ def __init__(self, model: Network, combine: Reduction): self._base_model_name = model.get_name() self._state_sources: list[dict[str, Any] | Path | str] = [] self._loaded_state_index: int | None = None + # Cache the CPU state_dict per index so a local-path ensemble is read from + # disk once, not re-read + re-unpickled on every batch (the index cycles + # 0..N-1 each forward, so the next batch would otherwise reload all N). + self._state_cache: dict[int, dict[str, Any]] = {} self.add_module( self._model_name, copy.deepcopy(model), @@ -732,10 +736,14 @@ def _read_state_source(self, source: dict[str, Any] | Path | str) -> dict[str, A def _ensure_model_loaded(self, index: int) -> Network: model = self._get_model() if self._loaded_state_index != index: + state = self._state_cache.get(index) + if state is None: + state = self._read_state_source(self._state_sources[index]) + self._state_cache[index] = state # Checkpoints are keyed by the base model name, not by the streamed # ensemble suffix added after the previous load. model.set_name(self._base_model_name) - model.load(self._read_state_source(self._state_sources[index]), init=False) + model.load(state, init=False) model.set_name(f"{self._base_model_name}_{index}") self._loaded_state_index = index return model @@ -749,6 +757,7 @@ def load(self, state_sources: list[dict[str, Any] | Path | str]): """ self._state_sources = state_sources self._loaded_state_index = None + self._state_cache = {} if len(self._state_sources) == 1: self._ensure_model_loaded(0) diff --git a/konfai/utils/dataset.py b/konfai/utils/dataset.py index 22215815..eb9597d7 100644 --- a/konfai/utils/dataset.py +++ b/konfai/utils/dataset.py @@ -976,7 +976,9 @@ def file_to_data_slice(self, group: str, name: str, slices: tuple[slice, ...]) - path = self._path(name) info = get_dicom_info(path) - data, origin, spacing, direction = read_dicom_series_slice(path, slices, series_uid=info["series_uid"]) + data, origin, spacing, direction = read_dicom_series_slice( + path, slices, series_uid=info["series_uid"], info=info + ) info.update(origin=origin, spacing=spacing, direction=direction) return data, self._attributes(info) @@ -986,13 +988,18 @@ def file_to_data_statistics( name: str, channels: list[int] | None = None, ) -> dict[str, float]: - shape, _ = self.get_infos(group, name) + from konfai.utils.dicom import get_dicom_info, read_dicom_series_slice + + path = self._path(name) + info = get_dicom_info(path) + shape = info["shape"] state: dict[str, float] | None = None for index in range(shape[1]): - chunk, _ = self.file_to_data_slice( - group, - name, + chunk, _, _, _ = read_dicom_series_slice( + path, (slice(None), slice(index, index + 1), slice(None), slice(None)), + series_uid=info["series_uid"], + info=info, ) if channels is not None: chunk = chunk[channels] @@ -1088,6 +1095,7 @@ def __init__(self, filename: str | Path, file_format: str) -> None: self.filename = str(filename) self.file_format = file_format self._names_cache: dict[str, list[str]] = {} + self._infos_cache: dict[tuple[str, str], tuple[list[int], Attribute]] = {} def _exists_on_disk(self) -> bool: if os.path.exists(self.filename): @@ -1102,6 +1110,7 @@ def write( attributes: Attribute | None = None, ) -> None: self._names_cache.clear() + self._infos_cache.clear() if attributes is None: attributes = Attribute() if self.is_directory: @@ -1290,6 +1299,15 @@ def get_group(self) -> list[str]: return list(groups) def get_infos(self, groups: str, name: str) -> tuple[list[int], Attribute]: + # Memoize the header read (SITK reader + ReadImageInformation, or the HDF5/Zarr + # metadata parse): get_infos is called once per name per group per build-pass at + # setup, so caching it (like get_names) avoids re-parsing the same header N times. + # Cache and hand back copies so a caller mutating the geometry cannot poison it. + cache_key = (groups, name) + cached = self._infos_cache.get(cache_key) + if cached is not None: + shape, attr = cached + return list(shape), Attribute(attr) if self.is_directory: for sub_directory in self._get_sub_directories(groups): group = groups.split("/")[-1] @@ -1300,10 +1318,14 @@ def get_infos(self, groups: str, name: str) -> tuple[list[int], Attribute]: self.file_format, self.level, ) as file: - return file.get_infos("", group) + result = file.get_infos("", group) + self._infos_cache[cache_key] = (list(result[0]), Attribute(result[1])) + return result else: with Dataset.File(self.filename, True, self.file_format, self.level) as file: - return file.get_infos(groups, name) + result = file.get_infos(groups, name) + self._infos_cache[cache_key] = (list(result[0]), Attribute(result[1])) + return result raise NameError(f"Dataset entry '{groups}/{name}' not found in {self.filename}.") def get_statistics(self, groups: str) -> dict[str, dict[str, dict[str, float | list[float]]]]: diff --git a/konfai/utils/dicom.py b/konfai/utils/dicom.py index b4adf0d8..cf3b7261 100644 --- a/konfai/utils/dicom.py +++ b/konfai/utils/dicom.py @@ -387,6 +387,7 @@ def get_dicom_info( return { "series_uid": selected_uid, "files": files, + "sorted_files": [Path(ds.filename) for ds in datasets], "shape": [1, len(datasets), rows, columns], "origin": origin, "spacing": spacing, @@ -400,9 +401,15 @@ def read_dicom_series_slice( *, series_uid: str | None = None, apply_rescale: bool = True, + info: dict[str, Any] | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Read only the selected DICOM slices and return updated patch geometry.""" - info = get_dicom_info(directory, series_uid=series_uid) + if info is None: + info = get_dicom_info(directory, series_uid=series_uid) + elif series_uid is not None and series_uid != info["series_uid"]: + raise DatasetManagerError( + f"series_uid '{series_uid}' does not match the provided info series '{info['series_uid']}'." + ) shape = info["shape"] if len(slices) != len(shape): raise DatasetManagerError(f"Expected {len(shape)} slices, got {len(slices)}.") @@ -411,10 +418,8 @@ def read_dicom_series_slice( if list(channel_indices) not in ([0], []): raise DatasetManagerError("DICOM stores scalar data and supports only channel 0.") - _selected_uid, files = _select_series_files(directory, series_uid or info["series_uid"]) - headers = sort_series(files, stop_before_pixels=True) z_indices = list(range(*normalized[1].indices(shape[1]))) - selected_files = [Path(headers[index].filename) for index in z_indices] + selected_files = [info["sorted_files"][index] for index in z_indices] datasets = sort_series(selected_files) volume = read_volume(datasets, apply_rescale=apply_rescale) volume = volume[normalized[0], :, normalized[2], normalized[3]] diff --git a/tests/unit/test_perf_hot_paths.py b/tests/unit/test_perf_hot_paths.py new file mode 100644 index 00000000..5a60e1ea --- /dev/null +++ b/tests/unit/test_perf_hot_paths.py @@ -0,0 +1,240 @@ +# 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 + +"""Regression tests for the performance hot-path fixes (see AUDIT.md — Performance backlog).""" + +import os + +os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") +os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") + +from pathlib import Path # noqa: E402 +from unittest.mock import MagicMock # noqa: E402 + +import torch # noqa: E402 + +import konfai.utils.dataset as dataset_module # noqa: E402 +from konfai.data.patching import Accumulator # noqa: E402 +from konfai.predictor import ModelComposite # noqa: E402 +from konfai.utils.dataset import Attribute, Dataset # noqa: E402 + + +def test_accumulator_is_full_counts_without_rescanning(): + """P6: is_full() is O(1) — a running counter, idempotent adds, reset on assemble.""" + patch_slices = [(slice(0, 2), slice(0, 2)), (slice(0, 2), slice(2, 4))] + acc = Accumulator(patch_slices, [0, 0], None, batch=False) + + assert acc.is_full() is False + acc.add_layer(0, torch.ones(1, 2, 2)) + assert acc.is_full() is False + acc.add_layer(1, torch.ones(1, 2, 2) * 2) + assert acc.is_full() is True + + # Re-adding the same index must not double-count. + acc.add_layer(1, torch.ones(1, 2, 2) * 3) + assert acc._filled == 2 + assert acc.is_full() is True + + acc.assemble() + assert acc._filled == 0 + assert acc.is_full() is False + + +def test_ensemble_reads_each_checkpoint_once_across_batches(): + """P1: a local-path ensemble reads/unpickles each checkpoint once, not once per batch.""" + mc = ModelComposite.__new__(ModelComposite) + mc._base_model_name = "Model" + mc._state_sources = [Path("/fake/ckpt_0.pt"), Path("/fake/ckpt_1.pt"), Path("/fake/ckpt_2.pt")] + mc._loaded_state_index = None + mc._state_cache = {} + mc._get_model = lambda: MagicMock() + + reads: list[str] = [] + + def fake_read(src): + reads.append(str(src)) + return {"w": str(src)} + + mc._read_state_source = fake_read + + # Four forward passes, each looping over all three sub-models (as forward() does). + for _batch in range(4): + for idx in range(3): + mc._ensure_model_loaded(idx) + + assert len(reads) == 3, f"expected 3 disk reads (one per index), got {len(reads)}" + # Compare via str(Path(...)) so the expected separators match the platform + # (the reads store str(src); Windows renders these with backslashes). + assert set(reads) == {str(Path(f"/fake/ckpt_{i}.pt")) for i in range(3)} + + # load() must invalidate the stale cache when the sources change. + mc.load([Path("/other.pt")]) + assert 1 not in mc._state_cache and 2 not in mc._state_cache + assert mc._state_cache.get(0) == {"w": str(Path("/other.pt"))} + + +def test_get_infos_is_memoized_and_returns_independent_copies(monkeypatch): + """P4: the header read is cached per (group, name); results are copies (no aliasing).""" + ds = Dataset.__new__(Dataset) + ds.is_directory = False + ds.filename = "/fake/ds" + ds.file_format = "sitk" + ds.level = 0 + ds._names_cache = {} + ds._infos_cache = {} + + opens = {"n": 0} + + class _FakeFile: + def __init__(self, *args, **kwargs): + opens["n"] += 1 + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def get_infos(self, groups, name): + return [4, 5, 6], Attribute({"Spacing": "1.0 1.0 1.0"}) + + monkeypatch.setattr(dataset_module.Dataset, "File", _FakeFile) + + first = ds.get_infos("g", "n") + assert opens["n"] == 1 + second = ds.get_infos("g", "n") + assert opens["n"] == 1, "second call must be served from cache, not re-open the file" + assert first[0] == second[0] == [4, 5, 6] + + # Copies, not aliases: mutating a returned result must not poison the cache. + first[0].append(999) + third = ds.get_infos("g", "n") + assert third[0] == [4, 5, 6] + + # A write invalidates the cache (mirrors get_names). + ds._infos_cache.clear() # write() calls this + ds.get_infos("g", "n") + assert opens["n"] == 2, "after invalidation the header is read again" + + +def test_dicom_slice_info_threading_is_byte_identical_and_removes_rescans(tmp_path, monkeypatch): + """P2: threading get_dicom_info's sorted files stays byte-identical while removing re-scans/re-sorts.""" + import numpy as np + import pytest + + pytest.importorskip("pydicom") + from konfai.utils import dicom + from konfai.utils.errors import DatasetManagerError + + root = tmp_path / "case" + vol = (np.arange(1 * 4 * 5 * 6).reshape(1, 4, 5, 6) % 97).astype(np.float32) + dicom.write_dicom_series(root, vol, origin=(1.0, 2.0, 3.0), spacing=(0.7, 0.8, 2.5), direction=np.eye(3).flatten()) + + # (a) byte-identical: standalone (info=None) path == threaded (info precomputed) path + sl = (slice(None), slice(1, 3), slice(None), slice(None)) + ref = dicom.read_dicom_series_slice(root, sl) + info = dicom.get_dicom_info(root) + got = dicom.read_dicom_series_slice(root, sl, series_uid=info["series_uid"], info=info) + for a, b in zip(ref, got): + assert np.array_equal(np.asarray(a), np.asarray(b)) + + assert len(info["sorted_files"]) == vol.shape[1] + assert all(isinstance(p, Path) for p in info["sorted_files"]) + + # a mismatching series_uid must not silently read the wrong series + with pytest.raises(DatasetManagerError): + dicom.read_dicom_series_slice(root, sl, series_uid="9.9.9.mismatch", info=info) + + # arity-mismatch path (info=None) still raises before any file selection + with pytest.raises(DatasetManagerError): + dicom.read_dicom_series_slice(root, (slice(None), slice(0, 2))) + + # (b) redundant work is gone: spy discover_series / sort_series call counts + calls = {"discover": 0, "sort": 0} + real_discover, real_sort = dicom.discover_series, dicom.sort_series + monkeypatch.setattr(dicom, "discover_series", lambda *a, **k: (calls.__setitem__("discover", calls["discover"] + 1), real_discover(*a, **k))[1]) + monkeypatch.setattr(dicom, "sort_series", lambda *a, **k: (calls.__setitem__("sort", calls["sort"] + 1), real_sort(*a, **k))[1]) + + dataset_file = Dataset.DicomFile(str(tmp_path), read=True) + + # one patch read: 1 discovery + 2 sorts (pre-fix: 3 discoveries + 4 sorts) + calls["discover"] = calls["sort"] = 0 + data, _attr = dataset_file.file_to_data_slice("", "case", sl) + assert np.array_equal(np.asarray(data), np.asarray(ref[0])) + assert calls["discover"] == 1 + assert calls["sort"] == 2 + + # statistics over Z: 1 discovery (pre-fix: O(Z)); numerics preserved (Welford, ddof=1) + calls["discover"] = calls["sort"] = 0 + stats = dataset_file.file_to_data_statistics("", "case") + assert calls["discover"] == 1 + assert np.isclose(stats["mean"], float(vol.mean()), atol=1e-4) + assert np.isclose(stats["std"], float(np.std(vol, ddof=1)), atol=1e-4) + assert np.isclose(stats["min"], float(vol.min()), atol=1e-4) + assert np.isclose(stats["max"], float(vol.max()), atol=1e-4) + + +def _old_clip(tensor, lo, hi): + """The pre-fix Clip inner logic (float()-cast where-scatter), for byte-identity checks.""" + t = tensor.clone() + t[torch.where(t.float() < lo)] = lo + t[torch.where(t.float() > hi)] = hi + return t + + +def _eq_nan(a, b): + return bool(((a == b) | (a.isnan() & b.isnan())).all()) + + +def test_clip_clamp_fast_path_is_byte_identical_on_float32_and_safe_on_int(): + """Perf batch: float32 clamp_ fast path is byte-identical; int/float64 keep the old scatter.""" + from konfai.data.transform import Clip + from konfai.utils.dataset import Attribute + + clip = Clip(min_value=-5.0, max_value=5.0) + lo, hi = -5.0, 5.0 + + # float32 with the hostile edge cases the red-team flagged: NaN, +/-inf, exact bounds + f32 = torch.tensor([-1e9, -5.0, -2.0, 0.0, 2.0, 5.0, 1e9, float("nan"), float("inf"), float("-inf")], dtype=torch.float32) + got = clip("x", f32.clone(), Attribute()) + assert _eq_nan(got, _old_clip(f32, lo, hi)) + assert got.dtype == torch.float32 + + # int16 (CT-style) must NOT crash and must equal the old scatter (else-branch) + i16 = torch.tensor([[-2000, -5, 0, 5, 2000]], dtype=torch.int16) + got_i = clip("x", i16.clone(), Attribute()) + assert torch.equal(got_i, _old_clip(i16, lo, hi)) + assert got_i.dtype == torch.int16 + + # float64 keeps the legacy float()-cast comparison path (else-branch), unchanged + f64 = torch.tensor([-9.0, -5.0, 0.0, 5.0, 9.0], dtype=torch.float64) + got_d = clip("x", f64.clone(), Attribute()) + assert _eq_nan(got_d, _old_clip(f64, lo, hi)) + assert got_d.dtype == torch.float64 + + +def test_clip_float32_nan_dynamic_bound_does_not_corrupt_volume(): + """Review catch: a dynamic bound resolving to NaN must NOT turn the whole float32 volume to NaN.""" + from konfai.data.transform import Clip + from konfai.utils.dataset import Attribute + + # data contains a NaN voxel -> min()/max() resolve to NaN bounds + data = torch.tensor([1.0, 2.0, float("nan"), 3.0], dtype=torch.float32) + got = Clip(min_value="min", max_value="max")("x", data.clone(), Attribute()) + # legacy behaviour: NaN comparisons are False, so the scatter is a no-op (values preserved) + assert not bool(got.isnan().all()), "must not become all-NaN" + assert got[0] == 1.0 and got[1] == 2.0 and got[3] == 3.0 + assert bool(got[2].isnan())