From f92acfa2d29f54c1de8a9bb40b1a5a8cb1a30700 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Sun, 5 Jul 2026 19:04:08 +0200 Subject: [PATCH 1/4] test: reorganize suite by module under test and remove duplicates Group tests by the unit they cover instead of by bug-fix batch: - add tests/conftest.py (autouse KONFAI env defaults, write_config and image_attributes fixtures), replacing 12 copy-pasted module-level os.environ.setdefault blocks that violated the monkeypatch rule - fold the audit/fixes grab-bags and single-test files into per-module files: config, patching, data_manager, transform, augmentation, network, blocks, schedulers, models, dataset, measure, main_cli - merge 5 duplicated tests into their surviving twin, keeping every assertion (config missing-file/invalid-yaml/invalid-literal errors, Data._split tail dedup) 38 unit files instead of 60; 343 -> 338 collected tests (the 5 exact duplicates); per-module coverage is byte-identical before/after. --- tests/conftest.py | 63 +++ tests/unit/test_audit_fixes.py | 317 ------------- tests/unit/test_augmentation.py | 289 ++++++++++++ tests/unit/test_augmentation_fixes.py | 127 ----- tests/unit/test_augmentation_flip.py | 70 --- .../test_augmentation_label_interpolation.py | 38 -- tests/unit/test_batchnorm_init.py | 30 -- tests/unit/test_blocks.py | 50 ++ tests/unit/test_cli_entry_points.py | 46 -- tests/unit/test_config.py | 371 ++++++++++----- tests/unit/test_config_dict_roundtrip.py | 51 --- tests/unit/test_config_dotted_dict_keys.py | 129 ------ tests/unit/test_config_validation.py | 160 ------- tests/unit/test_criterion_loaders.py | 43 -- tests/unit/test_crop_transform_shape.py | 36 -- tests/unit/test_data_manager.py | 432 ++++++++++++++++++ tests/unit/test_data_pipeline_audit.py | 262 ----------- ...{test_dataset_audit.py => test_dataset.py} | 135 +++++- tests/unit/test_ddp_shard_balance.py | 46 -- tests/unit/test_dict_metric_logging.py | 42 -- tests/unit/test_evaluator_statistics.py | 11 +- tests/unit/test_experimental_models.py | 34 -- tests/unit/test_get_infos_shape_order.py | 75 --- tests/unit/test_hdf5_read_only.py | 57 --- tests/unit/test_inline_augmentations.py | 174 ------- tests/unit/test_loss_weight_scheduling.py | 53 --- tests/unit/test_main_cli.py | 67 +++ tests/unit/test_measure.py | 49 +- tests/unit/test_models.py | 105 +++++ tests/unit/test_named_forward.py | 104 ----- tests/unit/test_named_forward_sibling.py | 51 --- tests/unit/test_network.py | 366 +++++++++++++++ tests/unit/test_network_ddp_fixes.py | 123 ----- tests/unit/test_patch_overlap_border.py | 58 --- tests/unit/test_patching.py | 87 +++- tests/unit/test_perceptual_loss_targets.py | 42 -- tests/unit/test_perf_hot_paths.py | 19 +- tests/unit/test_runtime_guards.py | 84 +--- tests/unit/test_runtime_progress_ddp.py | 6 +- tests/unit/test_schedulers.py | 55 +++ tests/unit/test_transform.py | 395 ++++++++++++++++ tests/unit/test_transform_clip.py | 43 -- tests/unit/test_transform_dilate.py | 66 --- tests/unit/test_transform_fixes.py | 184 -------- tests/unit/test_transform_norm.py | 51 --- tests/unit/test_unet_attention.py | 37 -- tests/unit/test_yaml_model_equivalence.py | 17 +- 47 files changed, 2345 insertions(+), 2805 deletions(-) create mode 100644 tests/conftest.py delete mode 100644 tests/unit/test_audit_fixes.py create mode 100644 tests/unit/test_augmentation.py delete mode 100644 tests/unit/test_augmentation_fixes.py delete mode 100644 tests/unit/test_augmentation_flip.py delete mode 100644 tests/unit/test_augmentation_label_interpolation.py delete mode 100644 tests/unit/test_batchnorm_init.py create mode 100644 tests/unit/test_blocks.py delete mode 100644 tests/unit/test_cli_entry_points.py delete mode 100644 tests/unit/test_config_dict_roundtrip.py delete mode 100644 tests/unit/test_config_dotted_dict_keys.py delete mode 100644 tests/unit/test_config_validation.py delete mode 100644 tests/unit/test_criterion_loaders.py delete mode 100644 tests/unit/test_crop_transform_shape.py create mode 100644 tests/unit/test_data_manager.py delete mode 100644 tests/unit/test_data_pipeline_audit.py rename tests/unit/{test_dataset_audit.py => test_dataset.py} (59%) delete mode 100644 tests/unit/test_ddp_shard_balance.py delete mode 100644 tests/unit/test_dict_metric_logging.py delete mode 100644 tests/unit/test_experimental_models.py delete mode 100644 tests/unit/test_get_infos_shape_order.py delete mode 100644 tests/unit/test_hdf5_read_only.py delete mode 100644 tests/unit/test_inline_augmentations.py delete mode 100644 tests/unit/test_loss_weight_scheduling.py create mode 100644 tests/unit/test_models.py delete mode 100644 tests/unit/test_named_forward.py delete mode 100644 tests/unit/test_named_forward_sibling.py create mode 100644 tests/unit/test_network.py delete mode 100644 tests/unit/test_network_ddp_fixes.py delete mode 100644 tests/unit/test_patch_overlap_border.py delete mode 100644 tests/unit/test_perceptual_loss_targets.py create mode 100644 tests/unit/test_schedulers.py create mode 100644 tests/unit/test_transform.py delete mode 100644 tests/unit/test_transform_clip.py delete mode 100644 tests/unit/test_transform_dilate.py delete mode 100644 tests/unit/test_transform_fixes.py delete mode 100644 tests/unit/test_transform_norm.py delete mode 100644 tests/unit/test_unet_attention.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..0742880e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,63 @@ +# 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 + +"""Shared fixtures for the KonfAI test suite.""" + +from collections.abc import Callable +from pathlib import Path + +import numpy as np +import pytest + + +@pytest.fixture(autouse=True) +def _konfai_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Harmless per-test defaults for the mandatory KONFAI environment variables. + + ``Config()`` requires ``KONFAI_config_file`` and ``KONFAI_CONFIG_MODE`` (AGENTS.md §7). + Tests that exercise the config engine override these with ``monkeypatch.setenv``. + """ + monkeypatch.setenv("KONFAI_config_file", "/tmp/konfai-none.yml") + monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") + + +@pytest.fixture +def write_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Callable[..., Path]: + """Write a YAML config to ``tmp_path`` and point the KONFAI env vars at it.""" + + def write(content: str, *, mode: str = "Done", name: str = "config.yml") -> Path: + config_path = tmp_path / name + config_path.write_text(content, encoding="utf-8") + monkeypatch.setenv("KONFAI_config_file", str(config_path)) + monkeypatch.setenv("KONFAI_CONFIG_MODE", mode) + return config_path + + return write + + +@pytest.fixture +def image_attributes(): + """Factory for an ``Attribute`` carrying Origin/Spacing/Direction geometry.""" + from konfai.utils.dataset import Attribute + + def make(origin: list[float], spacing: list[float]) -> Attribute: + attributes = Attribute() + attributes["Origin"] = np.asarray(origin, dtype=np.float64) + attributes["Spacing"] = np.asarray(spacing, dtype=np.float64) + attributes["Direction"] = np.eye(len(origin), dtype=np.float64).reshape(-1) + return attributes + + return make diff --git a/tests/unit/test_audit_fixes.py b/tests/unit/test_audit_fixes.py deleted file mode 100644 index 4643d717..00000000 --- a/tests/unit/test_audit_fixes.py +++ /dev/null @@ -1,317 +0,0 @@ -# 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 audit bug-fixes (see AUDIT.md §5).""" - -import inspect -import os -import subprocess -import sys - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import numpy as np # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from konfai.data.augmentation import Flip, Rotate # noqa: E402 -from konfai.data.data_manager import DataTrain # noqa: E402 -from konfai.data.transform import Padding, ResampleToShape, Standardize # noqa: E402 -from konfai.network.blocks import Select, Unsqueeze # noqa: E402 -from konfai.utils.dataset import Attribute # noqa: E402 -from konfai.utils.errors import ConfigError, MeasureError # noqa: E402 - - -def test_vae_latent_uses_gaussian_noise(): - """#3 LatentDistributionZ must sample N(0,1), not U[0,1].""" - from konfai.network.blocks import LatentDistribution - - layer = LatentDistribution.LatentDistributionZ() - mu = torch.zeros(200_000) - log_std = torch.zeros(200_000) - z = layer(mu, log_std) # == epsilon - assert abs(float(z.mean())) < 0.05, "mean should be ~0" - assert abs(float(z.std()) - 1.0) < 0.05, "std should be ~1 (Gaussian), not ~0.29 (uniform)" - - -def test_standardize_explicit_scalar_stats(): - """#5 Standardize with explicit scalar mean/std must not crash.""" - t = Standardize(lazy=False, mean=[10.0], std=[2.0]) - x = torch.arange(24, dtype=torch.float32).reshape(1, 2, 3, 4) - out = t("c", x.clone(), Attribute()) - assert torch.allclose(out, (x - 10.0) / 2.0) - - -def test_standardize_explicit_per_channel_stats(): - """#5 Per-channel mean/std broadcast over the channel axis.""" - t = Standardize(lazy=False, mean=[10.0, 20.0], std=[2.0, 4.0]) - x = torch.zeros(2, 3, 4) - x[0] = 10.0 - x[1] = 20.0 - out = t("c", x.clone(), Attribute()) - assert torch.allclose(out, torch.zeros_like(out), atol=1e-6) - - -def test_rotate_converts_degrees_to_radians(): - """#6 A 90-degree rotation must yield [[0,-1],[1,0]], not cos/sin of 90 radians.""" - rot = Rotate(a_min=90.0, a_max=90.0, is_quarter=False) - rot._state_init(0, [[8, 8]], [Attribute()]) - block = rot.matrix[0][0][0, :2, :2] - assert torch.allclose(block, torch.tensor([[0.0, -1.0], [1.0, 0.0]]), atol=1e-5) - - -def test_predict_evaluate_expose_tensorboard_param(): - """#7 CLI -tb/--tensorboard (dest 'tensorboard') must reach predict()/evaluate().""" - from konfai.evaluator import evaluate - from konfai.predictor import predict - - for fn in (predict, evaluate): - params = inspect.signature(fn).parameters - assert "tensorboard" in params, f"{fn.__name__} must accept 'tensorboard'" - assert "tb" not in params, f"{fn.__name__} must not use the old 'tb' name" - - -def test_unsqueeze_forward_accepts_tensor(): - """#8 Unsqueeze.forward(tensor) must work on a single tensor.""" - assert Unsqueeze(dim=1)(torch.randn(3, 4)).shape == (3, 1, 4) - - -def test_resample_to_shape_does_not_mutate_config(): - """#9 transform_shape must not write resolved dims back into the shared instance config.""" - resampler = ResampleToShape(shape=[0, 16, 16]) - before = resampler.shape.clone() - attributes = Attribute() - attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) - out = resampler.transform_shape("CT", "case", [8, 16, 16], attributes) - assert out[0] == 8 # sentinel 0 resolved to the input dim for this call - assert torch.equal(resampler.shape, before), "self.shape must stay [0, 16, 16] for the next case" - - -def test_select_squeezes_size_one_dims_by_size(): - """#12 Select must squeeze dimensions whose size is 1, not the dim at index 1.""" - out = Select([slice(0, 1), slice(None), slice(None)])(torch.randn(1, 5, 6)) - assert out.shape == (5, 6) - # a tensor with no size-1 dims is unchanged - out2 = Select([slice(None), slice(None)])(torch.randn(4, 5)) - assert out2.shape == (4, 5) - - -def test_augmentation_resamples_after_reset_state(): - """#1 Augmentation parameters must be re-sampled each epoch via reset_state. - - Within an epoch ``state_init`` caches the per-case draw so every patch shares - one transform; ``reset_state`` must clear that cache so the next epoch draws - fresh parameters (previously ``who_index`` was never cleared → frozen forever). - """ - aug = Flip([1.0, 1.0, 1.0]) - aug.load(1.0) - - aug.state_init(0, [[4, 4, 4]], [Attribute()]) - first = aug.flip[0] - # Re-running state_init without a reset returns the cached draw unchanged. - aug.state_init(0, [[4, 4, 4]], [Attribute()]) - assert aug.flip[0] is first - - aug.reset_state(0) - assert 0 not in aug.who_index - aug.state_init(0, [[4, 4, 4]], [Attribute()]) - assert 0 in aug.who_index - assert aug.flip[0] is not first # a fresh draw replaced the cached one - - -def test_update_scheduler_empty_raises_config_error(): - """update_scheduler on an empty schedule must raise a clear ConfigError.""" - from konfai.network.network import Measure - - with pytest.raises(ConfigError): - Measure.update_scheduler(None, {}, 0) # type: ignore[arg-type] - - -def test_update_scheduler_past_last_window_clamps_to_last(): - """Past every configured window, the last scheduler is selected (no crash).""" - from konfai.metric.schedulers import Constant - from konfai.network.network import Measure - - s0, s1 = Constant(1.0), Constant(2.0) - schedulers = {s0: 3, s1: 3} # active windows [0,3) and [3,6) - assert Measure.update_scheduler(None, schedulers, 4) is s1 # type: ignore[arg-type] - assert Measure.update_scheduler(None, schedulers, 100) is s1 # type: ignore[arg-type] - - -def test_missing_metric_dependency_raises_actionable_error(): - """Optional criterion deps must surface an actionable MeasureError, not ImportError.""" - from konfai.metric.measure import _require_optional - - with pytest.raises(MeasureError) as excinfo: - _require_optional("konfai_definitely_missing_pkg_zzz", criterion="SSIM", extra="ssim") - message = str(excinfo.value) - assert "SSIM" in message - assert "konfai[ssim]" in message - - -def test_load_state_dict_warm_starts_resized_layer_and_keeps_siblings(): - """#2 A resized layer must warm-start, and sibling layers must still load. - - The bug checked ``isinstance(module, Linear)`` (the parent) instead of the - child, and used an early ``return`` that aborted loading the remaining - siblings of a resized layer. - """ - from konfai.network.network import Network - - class _Net(Network): - def __init__(self, fc_out: int) -> None: - super().__init__(in_channels=1) - self.add_module("fc", torch.nn.Linear(4, fc_out)) - self.add_module("head", torch.nn.Linear(4, 2)) - - old = _Net(fc_out=4) - # Network.state_dict() wraps the flat params under the network name; load_state_dict - # consumes that inner flat dict ("fc.weight", ...). - inner = next(iter(old.state_dict().values())) - checkpoint = {key: value.clone() for key, value in inner.items()} - - new = _Net(fc_out=6) # fc output grows 4 -> 6 (resized); head is unchanged - new.load_state_dict(checkpoint) # must not raise - - fc = new["fc"] - head = new["head"] - assert fc.weight.shape == (6, 4) - assert torch.equal(fc.weight[:4], checkpoint["fc.weight"]) # warm-started rows - # The sibling after the resized layer must still be loaded (old `return` skipped it). - assert torch.equal(head.weight, checkpoint["head.weight"]) - assert torch.equal(head.bias, checkpoint["head.bias"]) - - -def test_adaptation_sets_requires_grad_at_construction(): - """#18 Adaptation must configure requires_grad in __init__, not on every forward.""" - from konfai.models.representation.representation import Adaptation - - adaptation = Adaptation() - # State is correct immediately after construction, before any forward pass. - assert all(not p.requires_grad for p in adaptation.Encoder_1.parameters()) - assert all(p.requires_grad for p in adaptation.FCT_1.parameters()) - - -def test_linear_vae_is_parameterized_and_variational(): - """#17 LinearVAE must be parameterized (no hardcoded dims) and sample a latent.""" - from konfai.models.generation.vae import LinearVAE - - model = LinearVAE(in_features=32, hidden_features=16, latent_dim=4) - x = torch.randn(2, 32) - outputs = dict(model.named_forward(x)) - assert outputs["Head.Tanh"].shape == (2, 32) # reconstruction matches input size - assert "Latent.mu" in outputs and "Latent.log_std" in outputs # KL-ready outputs - # The latent is sampled: the reconstruction differs across RNG draws. - torch.manual_seed(0) - first = dict(model.named_forward(x))["Head.Tanh"] - torch.manual_seed(1) - second = dict(model.named_forward(x))["Head.Tanh"] - assert not torch.allclose(first, second) - - -def _geometry(origin: list[float], spacing: list[float]) -> Attribute: - attributes = Attribute() - attributes["Origin"] = np.asarray(origin, dtype=np.float64) - attributes["Spacing"] = np.asarray(spacing, dtype=np.float64) - attributes["Direction"] = np.eye(len(origin), dtype=np.float64).reshape(-1) - return attributes - - -def test_padding_shifts_origin_along_the_padded_axes(): - """Each F.pad pair (X, Y, Z) must shift the matching (x, y, z) origin component.""" - attributes = _geometry(origin=[10.0, 20.0, 30.0], spacing=[1.0, 2.0, 4.0]) - - padded = Padding(padding=[1, 0, 2, 0, 3, 0])("case", torch.zeros(1, 5, 5, 5), attributes) - - assert list(padded.shape) == [1, 8, 7, 6] - np.testing.assert_allclose( - attributes.get_np_array("Origin"), - [10.0 - 1 * 1.0, 20.0 - 2 * 2.0, 30.0 - 3 * 4.0], - ) - - -def test_padding_after_the_data_keeps_origin(): - """Padding only on the high side of each axis must leave the origin untouched.""" - attributes = _geometry(origin=[10.0, 20.0, 30.0], spacing=[1.0, 2.0, 4.0]) - - padded = Padding(padding=[0, 2, 0, 0, 0, 1])("case", torch.zeros(1, 5, 5, 5), attributes) - - assert list(padded.shape) == [1, 6, 5, 7] - np.testing.assert_allclose(attributes.get_np_array("Origin"), [10.0, 20.0, 30.0]) - - -_SPLIT_PROBE = """ -import os -import random - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -from konfai.data.data_manager import DataTrain - -names = [f"CASE_{i:03d}" for i in range(20)] -data = DataTrain(augmentations=None, validation="0:4") -data._resolve_dataset_sources = lambda: {} -data._resolve_common_names = lambda datasets: ({}, set(names)) -data._get_datasets = lambda case_names, dataset_name, augmentations: ({}, []) -random.seed(1234) -data._prepare_datasets() -print(";".join(data._prepared_train_names)) -print(";".join(data._prepared_validation_names)) -""" - - -def test_train_validation_split_is_reproducible_across_interpreters(): - """Same seed → same split, whatever the interpreter's string-hash randomization.""" - outputs = [] - for hash_seed in ("0", "424242"): - env = dict(os.environ, PYTHONHASHSEED=hash_seed) - result = subprocess.run( - [sys.executable, "-c", _SPLIT_PROBE], - env=env, - capture_output=True, - text=True, - check=True, - ) - outputs.append(result.stdout) - assert outputs[0] == outputs[1] - train_names, validation_names = (line.split(";") for line in outputs[0].splitlines()) - assert len(train_names) == 16 - assert len(validation_names) == 4 - assert set(train_names).isdisjoint(validation_names) - - -def test_train_split_shuffle_draws_from_sorted_names(monkeypatch): - """The seeded shuffle must receive the case names in sorted order and drive the split.""" - captured: dict[str, list[str]] = {} - - def fake_sample(population, k): - captured["population"] = list(population) - assert k == len(population) - return list(reversed(population)) - - monkeypatch.setattr("konfai.data.data_manager.random.sample", fake_sample) - - data = DataTrain(augmentations=None, validation="0:2") - names = {"CASE_010", "CASE_002", "CASE_001", "CASE_005", "CASE_003"} - data._resolve_dataset_sources = lambda: {} - data._resolve_common_names = lambda datasets: ({}, names) - data._get_datasets = lambda case_names, dataset_name, augmentations: ({}, []) - data._prepare_datasets() - - assert captured["population"] == sorted(names) - assert data._prepared_validation_names == ["CASE_010", "CASE_005"] - assert data._prepared_train_names == ["CASE_003", "CASE_002", "CASE_001"] diff --git a/tests/unit/test_augmentation.py b/tests/unit/test_augmentation.py new file mode 100644 index 00000000..3c595971 --- /dev/null +++ b/tests/unit/test_augmentation.py @@ -0,0 +1,289 @@ +# 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 + +"""Tests for ``konfai.data.augmentation``: per-sample state (draw/reset/inverse), +Flip (incl. vector fields), Rotate, Translate, intensity augmentations, and the +SimpleITK-backed Elastix/Mask augmentations.""" + +from pathlib import Path + +import konfai.data.augmentation as augmentation_module +import numpy as np +import pytest +import torch +from konfai.data.augmentation import ( + Brightness, + CutOUT, + Elastix, + Flip, + Mask, + Noise, + Rotate, + Translate, +) +from konfai.utils.dataset import Attribute +from konfai.utils.errors import AugmentationError + +# -------------------------------------------------------------------------------------- +# Per-sample state: draw caching, reset, and inverse slot lookup +# -------------------------------------------------------------------------------------- + + +def test_augmentation_resamples_after_reset_state(): + """#1 Augmentation parameters must be re-sampled each epoch via reset_state. + + Within an epoch ``state_init`` caches the per-case draw so every patch shares + one transform; ``reset_state`` must clear that cache so the next epoch draws + fresh parameters (previously ``who_index`` was never cleared → frozen forever). + """ + aug = Flip([1.0, 1.0, 1.0]) + aug.load(1.0) + + aug.state_init(0, [[4, 4, 4]], [Attribute()]) + first = aug.flip[0] + # Re-running state_init without a reset returns the cached draw unchanged. + aug.state_init(0, [[4, 4, 4]], [Attribute()]) + assert aug.flip[0] is first + + aug.reset_state(0) + assert 0 not in aug.who_index + aug.state_init(0, [[4, 4, 4]], [Attribute()]) + assert 0 in aug.who_index + assert aug.flip[0] is not first # a fresh draw replaced the cached one + + +def test_augmentation_inverse_uses_local_slot_for_global_index(): + """``inverse(a)`` receives the *global* sample index. + + Per-sample state (``flip``/``matrix``) is stored only for the selected + samples, in selection order (local slots). The global index must therefore + be translated to the sample's position within ``who_index`` before indexing + that state, otherwise a selected sample either reads a neighbour's transform + (silent corruption) or overruns the list (IndexError). + """ + aug = Flip([1.0, 0.0, 0.0]) + # Samples 1 and 2 were selected; their flip axes occupy local slots 0 and 1. + aug.who_index[0] = [1, 2] + aug.flip[0] = [[1], [2]] # global 1 -> axis 1, global 2 -> axis 2 + + x = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3) + + # Unselected sample is passed through untouched. + assert torch.equal(aug.inverse(0, 0, x.clone()), x) + # Global 1 -> local slot 0 -> axis 1 (previously read slot 1 -> axis 2). + assert torch.equal(aug.inverse(0, 1, x.clone()), torch.flip(x, [1])) + # Global 2 -> local slot 1 -> axis 2 (previously an out-of-range index). + assert torch.equal(aug.inverse(0, 2, x.clone()), torch.flip(x, [2])) + + +def test_intensity_augmentation_inverses_are_identity(): + """Value-only augmentations must invert to the identity. + + ColorTransform/Noise/CutOUT do not move voxels, so the inverse applied to a + prediction is the tensor itself. They previously returned ``None``, which + crashed the prediction inverse path (``NoneType`` has no ``device``). + """ + x = torch.randn(3, 4, 5, 6) # 3 channels for the ColorTransform path + + color = Brightness(b_std=0.5) + color.who_index[0] = [0] + assert torch.equal(color.inverse(0, 0, x.clone()), x) + + noise = Noise(n_std=0.1) + noise.who_index[0] = [0] + assert torch.equal(noise.inverse(0, 0, x.clone()), x) + + cutout = CutOUT(c_prob=1.0, cutout_size=2, value=0.0) + cutout.who_index[0] = [0] + assert torch.equal(cutout.inverse(0, 0, x.clone()), x) + + +# -------------------------------------------------------------------------------------- +# Flip — vector-field awareness +# -------------------------------------------------------------------------------------- + + +def _flip_all_axes(vector_field: bool) -> Flip: + flip = Flip(f_prob=[1.0, 1.0, 1.0], vector_field=vector_field) + flip._state_init(0, [[4, 5, 6]], [Attribute()]) + return flip + + +def test_flip_vector_field_round_trip_is_identity() -> None: + # TTA un-flips the model output with ``_inverse``: on a displacement field the compose of + # ``_compute`` and ``_inverse`` must be the identity, component signs included. + flip = _flip_all_axes(vector_field=True) + dvf = torch.randn(3, 4, 5, 6) + + augmented = flip._compute("case", 0, [dvf.clone()])[0] + + assert torch.equal(flip._inverse(0, 0, augmented), dvf) + + +def test_flip_vector_field_negates_flipped_components() -> None: + # Mirroring a spatial axis reverses the voxel layout AND the sign of that axis' component channel + # (channels are (dx, dy, dz) while tensor axes are reversed: dim 3 = x -> channel 0, ...). + flip = _flip_all_axes(vector_field=True) + dvf = torch.randn(3, 4, 5, 6) + + augmented = flip._compute("case", 0, [dvf.clone()])[0] + + layout_only = torch.flip(dvf, dims=[1, 2, 3]) + assert torch.equal(augmented, -layout_only) + + +def test_flip_scalar_data_is_layout_only() -> None: + # Single-channel data (images, masks) is mirror-invariant: even with ``vector_field`` enabled the + # shared Flip instance must not negate intensities. + flip = _flip_all_axes(vector_field=True) + volume = torch.randn(1, 4, 5, 6) + + augmented = flip._compute("case", 0, [volume.clone()])[0] + + assert torch.equal(augmented, torch.flip(volume, dims=[1, 2, 3])) + assert torch.equal(flip._inverse(0, 0, augmented), volume) + + +def test_flip_default_stays_layout_only_on_vector_data() -> None: + # ``vector_field`` is opt-in: existing intensity-TTA bundles keep the historical behaviour. + flip = _flip_all_axes(vector_field=False) + dvf = torch.randn(3, 4, 5, 6) + + augmented = flip._compute("case", 0, [dvf.clone()])[0] + + assert torch.equal(augmented, torch.flip(dvf, dims=[1, 2, 3])) + + +# -------------------------------------------------------------------------------------- +# Rotate +# -------------------------------------------------------------------------------------- + + +def test_rotate_converts_degrees_to_radians(): + """#6 A 90-degree rotation must yield [[0,-1],[1,0]], not cos/sin of 90 radians.""" + rot = Rotate(a_min=90.0, a_max=90.0, is_quarter=False) + rot._state_init(0, [[8, 8]], [Attribute()]) + block = rot.matrix[0][0][0, :2, :2] + assert torch.allclose(block, torch.tensor([[0.0, -1.0], [1.0, 0.0]]), atol=1e-5) + + +def test_rotate_quarter_builds_one_signed_permutation_per_sample(): + """``is_quarter`` must draw one rotation per sample/axis. + + The previous flat 9-vector produced 0-d angles that crashed + ``_rotation_3d_matrix`` and ignored the sample count. Each resulting matrix + must be a proper multiple-of-90-degree rotation, i.e. an orthogonal signed + permutation matrix (integer entries, determinant +1). + """ + torch.manual_seed(0) + rot = Rotate(is_quarter=True) + shapes = [[8, 8, 8], [8, 8, 8]] + rot._state_init(0, shapes, [Attribute(), Attribute()]) + + assert len(rot.matrix[0]) == len(shapes) # one matrix per sample, not 9 + identity = torch.eye(3) + for matrix in rot.matrix[0]: + assert matrix.shape == (1, 4, 4) + rotation = matrix[0, :3, :3] + assert torch.allclose(rotation @ rotation.T, identity, atol=1e-5) + assert torch.allclose(rotation, rotation.round(), atol=1e-5) + assert torch.allclose(torch.det(rotation), torch.tensor(1.0), atol=1e-5) + + +def test_rotate_preserves_label_ids_with_nearest() -> None: + # A geometric augmentation applies to every group, including uint8 segmentation targets. + # Bilinear resampling would blend class ids into a non-existent intermediate label (1|3 -> 2); + # nearest-neighbour must keep the label set unchanged. + rotate = Rotate(a_min=45.0, a_max=45.0) # a_max == a_min -> deterministic 45 degrees + labels = torch.zeros(1, 32, 32, dtype=torch.uint8) + labels[:, 8:24, 8:24] = 1 + labels[:, 12:20, 12:20] = 3 + rotate._state_init(0, [[32, 32]], [Attribute()]) + + out = rotate._compute("case", 0, [labels.clone()])[0] + + assert out.dtype == torch.uint8 + assert set(out.unique().tolist()).issubset({0, 1, 3}) + assert 2 not in out.unique().tolist() + + +# -------------------------------------------------------------------------------------- +# Translate +# -------------------------------------------------------------------------------------- + + +def test_translate_scales_voxel_offset_to_normalized_grid(): + """``t_min``/``t_max`` are voxel offsets. + + ``F.affine_grid`` expects normalized coordinates where a full axis spans + [-1, 1] (align_corners=True), so a d-voxel shift becomes d * 2 / (size - 1), + per axis in affine order (x, y, z) = reversed spatial (z, y, x). + """ + aug = Translate(t_min=5.0, t_max=5.0, is_int=False) # deterministic 5-voxel shift + aug._state_init(0, [[4, 6, 10]], [Attribute()]) # spatial (z, y, x) + column = aug.matrix[0][0][0, :3, 3] # affine order (x, y, z) + expected = torch.tensor([5.0 * 2 / (10 - 1), 5.0 * 2 / (6 - 1), 5.0 * 2 / (4 - 1)]) + assert torch.allclose(column, expected, atol=1e-6) + + +def test_translate_is_int_rounds_to_whole_voxels(): + """``is_int`` must round to entire voxels, not to two decimals (0.01).""" + aug = Translate(t_min=5.3, t_max=5.3, is_int=True) + aug._state_init(0, [[9, 9, 9]], [Attribute()]) + column = aug.matrix[0][0][0, :3, 3] + expected = torch.full((3,), 5.0 * 2 / (9 - 1)) # round(5.3) == 5, then normalized + assert torch.allclose(column, expected, atol=1e-6) + # Neither the pre-fix 0.01 rounding (5.3) nor raw-voxel units survive. + assert not torch.allclose(column, torch.full((3,), 5.3), atol=1e-6) + + +# -------------------------------------------------------------------------------------- +# SimpleITK-backed augmentations (Elastix / Mask) +# -------------------------------------------------------------------------------------- + + +def test_simpleitk_augmentations_fail_clearly_when_dependency_is_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(augmentation_module, "sitk", None) + + with pytest.raises(AugmentationError, match="SimpleITK"): + Elastix() + with pytest.raises(AugmentationError, match="SimpleITK"): + Mask("mask.mha", 0) + + +def test_mask_reads_pixels_only_on_first_compute(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sitk = pytest.importorskip("SimpleITK") + mask_path = tmp_path / "mask.mha" + sitk.WriteImage(sitk.GetImageFromArray(np.ones((2, 2), dtype=np.uint8)), str(mask_path)) + + read_count = 0 + original_read_image = sitk.ReadImage + + def counting_read_image(path: str): + nonlocal read_count + read_count += 1 + return original_read_image(path) + + monkeypatch.setattr(augmentation_module.sitk, "ReadImage", counting_read_image) + augmentation = Mask(str(mask_path), 0) + augmentation._state_init(0, [[2, 2]], [Attribute()]) + + assert read_count == 0 + augmentation._compute("case", 0, [torch.ones((1, 2, 2))]) + augmentation._compute("case", 0, [torch.ones((1, 2, 2))]) + assert read_count == 1 diff --git a/tests/unit/test_augmentation_fixes.py b/tests/unit/test_augmentation_fixes.py deleted file mode 100644 index 73be1b2e..00000000 --- a/tests/unit/test_augmentation_fixes.py +++ /dev/null @@ -1,127 +0,0 @@ -# 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 data-augmentation inverse/state-init fixes.""" - -import os - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import torch -from konfai.data.augmentation import ( - Brightness, - CutOUT, - Flip, - Noise, - Rotate, - Translate, -) -from konfai.utils.dataset import Attribute - - -def test_augmentation_inverse_uses_local_slot_for_global_index(): - """``inverse(a)`` receives the *global* sample index. - - Per-sample state (``flip``/``matrix``) is stored only for the selected - samples, in selection order (local slots). The global index must therefore - be translated to the sample's position within ``who_index`` before indexing - that state, otherwise a selected sample either reads a neighbour's transform - (silent corruption) or overruns the list (IndexError). - """ - aug = Flip([1.0, 0.0, 0.0]) - # Samples 1 and 2 were selected; their flip axes occupy local slots 0 and 1. - aug.who_index[0] = [1, 2] - aug.flip[0] = [[1], [2]] # global 1 -> axis 1, global 2 -> axis 2 - - x = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3) - - # Unselected sample is passed through untouched. - assert torch.equal(aug.inverse(0, 0, x.clone()), x) - # Global 1 -> local slot 0 -> axis 1 (previously read slot 1 -> axis 2). - assert torch.equal(aug.inverse(0, 1, x.clone()), torch.flip(x, [1])) - # Global 2 -> local slot 1 -> axis 2 (previously an out-of-range index). - assert torch.equal(aug.inverse(0, 2, x.clone()), torch.flip(x, [2])) - - -def test_rotate_quarter_builds_one_signed_permutation_per_sample(): - """``is_quarter`` must draw one rotation per sample/axis. - - The previous flat 9-vector produced 0-d angles that crashed - ``_rotation_3d_matrix`` and ignored the sample count. Each resulting matrix - must be a proper multiple-of-90-degree rotation, i.e. an orthogonal signed - permutation matrix (integer entries, determinant +1). - """ - torch.manual_seed(0) - rot = Rotate(is_quarter=True) - shapes = [[8, 8, 8], [8, 8, 8]] - rot._state_init(0, shapes, [Attribute(), Attribute()]) - - assert len(rot.matrix[0]) == len(shapes) # one matrix per sample, not 9 - identity = torch.eye(3) - for matrix in rot.matrix[0]: - assert matrix.shape == (1, 4, 4) - rotation = matrix[0, :3, :3] - assert torch.allclose(rotation @ rotation.T, identity, atol=1e-5) - assert torch.allclose(rotation, rotation.round(), atol=1e-5) - assert torch.allclose(torch.det(rotation), torch.tensor(1.0), atol=1e-5) - - -def test_translate_scales_voxel_offset_to_normalized_grid(): - """``t_min``/``t_max`` are voxel offsets. - - ``F.affine_grid`` expects normalized coordinates where a full axis spans - [-1, 1] (align_corners=True), so a d-voxel shift becomes d * 2 / (size - 1), - per axis in affine order (x, y, z) = reversed spatial (z, y, x). - """ - aug = Translate(t_min=5.0, t_max=5.0, is_int=False) # deterministic 5-voxel shift - aug._state_init(0, [[4, 6, 10]], [Attribute()]) # spatial (z, y, x) - column = aug.matrix[0][0][0, :3, 3] # affine order (x, y, z) - expected = torch.tensor([5.0 * 2 / (10 - 1), 5.0 * 2 / (6 - 1), 5.0 * 2 / (4 - 1)]) - assert torch.allclose(column, expected, atol=1e-6) - - -def test_translate_is_int_rounds_to_whole_voxels(): - """``is_int`` must round to entire voxels, not to two decimals (0.01).""" - aug = Translate(t_min=5.3, t_max=5.3, is_int=True) - aug._state_init(0, [[9, 9, 9]], [Attribute()]) - column = aug.matrix[0][0][0, :3, 3] - expected = torch.full((3,), 5.0 * 2 / (9 - 1)) # round(5.3) == 5, then normalized - assert torch.allclose(column, expected, atol=1e-6) - # Neither the pre-fix 0.01 rounding (5.3) nor raw-voxel units survive. - assert not torch.allclose(column, torch.full((3,), 5.3), atol=1e-6) - - -def test_intensity_augmentation_inverses_are_identity(): - """Value-only augmentations must invert to the identity. - - ColorTransform/Noise/CutOUT do not move voxels, so the inverse applied to a - prediction is the tensor itself. They previously returned ``None``, which - crashed the prediction inverse path (``NoneType`` has no ``device``). - """ - x = torch.randn(3, 4, 5, 6) # 3 channels for the ColorTransform path - - color = Brightness(b_std=0.5) - color.who_index[0] = [0] - assert torch.equal(color.inverse(0, 0, x.clone()), x) - - noise = Noise(n_std=0.1) - noise.who_index[0] = [0] - assert torch.equal(noise.inverse(0, 0, x.clone()), x) - - cutout = CutOUT(c_prob=1.0, cutout_size=2, value=0.0) - cutout.who_index[0] = [0] - assert torch.equal(cutout.inverse(0, 0, x.clone()), x) diff --git a/tests/unit/test_augmentation_flip.py b/tests/unit/test_augmentation_flip.py deleted file mode 100644 index ec123108..00000000 --- a/tests/unit/test_augmentation_flip.py +++ /dev/null @@ -1,70 +0,0 @@ -# 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 - -import torch -from konfai.data.augmentation import Flip -from konfai.utils.dataset import Attribute - - -def _flip_all_axes(vector_field: bool) -> Flip: - flip = Flip(f_prob=[1.0, 1.0, 1.0], vector_field=vector_field) - flip._state_init(0, [[4, 5, 6]], [Attribute()]) - return flip - - -def test_flip_vector_field_round_trip_is_identity() -> None: - # TTA un-flips the model output with ``_inverse``: on a displacement field the compose of - # ``_compute`` and ``_inverse`` must be the identity, component signs included. - flip = _flip_all_axes(vector_field=True) - dvf = torch.randn(3, 4, 5, 6) - - augmented = flip._compute("case", 0, [dvf.clone()])[0] - - assert torch.equal(flip._inverse(0, 0, augmented), dvf) - - -def test_flip_vector_field_negates_flipped_components() -> None: - # Mirroring a spatial axis reverses the voxel layout AND the sign of that axis' component channel - # (channels are (dx, dy, dz) while tensor axes are reversed: dim 3 = x -> channel 0, ...). - flip = _flip_all_axes(vector_field=True) - dvf = torch.randn(3, 4, 5, 6) - - augmented = flip._compute("case", 0, [dvf.clone()])[0] - - layout_only = torch.flip(dvf, dims=[1, 2, 3]) - assert torch.equal(augmented, -layout_only) - - -def test_flip_scalar_data_is_layout_only() -> None: - # Single-channel data (images, masks) is mirror-invariant: even with ``vector_field`` enabled the - # shared Flip instance must not negate intensities. - flip = _flip_all_axes(vector_field=True) - volume = torch.randn(1, 4, 5, 6) - - augmented = flip._compute("case", 0, [volume.clone()])[0] - - assert torch.equal(augmented, torch.flip(volume, dims=[1, 2, 3])) - assert torch.equal(flip._inverse(0, 0, augmented), volume) - - -def test_flip_default_stays_layout_only_on_vector_data() -> None: - # ``vector_field`` is opt-in: existing intensity-TTA bundles keep the historical behaviour. - flip = _flip_all_axes(vector_field=False) - dvf = torch.randn(3, 4, 5, 6) - - augmented = flip._compute("case", 0, [dvf.clone()])[0] - - assert torch.equal(augmented, torch.flip(dvf, dims=[1, 2, 3])) diff --git a/tests/unit/test_augmentation_label_interpolation.py b/tests/unit/test_augmentation_label_interpolation.py deleted file mode 100644 index 53ea30ec..00000000 --- a/tests/unit/test_augmentation_label_interpolation.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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 test: geometric augmentations resample integer label maps with nearest-neighbour.""" - -import torch -from konfai.data.augmentation import Rotate -from konfai.utils.dataset import Attribute - - -def test_rotate_preserves_label_ids_with_nearest() -> None: - # A geometric augmentation applies to every group, including uint8 segmentation targets. - # Bilinear resampling would blend class ids into a non-existent intermediate label (1|3 -> 2); - # nearest-neighbour must keep the label set unchanged. - rotate = Rotate(a_min=45.0, a_max=45.0) # a_max == a_min -> deterministic 45 degrees - labels = torch.zeros(1, 32, 32, dtype=torch.uint8) - labels[:, 8:24, 8:24] = 1 - labels[:, 12:20, 12:20] = 3 - rotate._state_init(0, [[32, 32]], [Attribute()]) - - out = rotate._compute("case", 0, [labels.clone()])[0] - - assert out.dtype == torch.uint8 - assert set(out.unique().tolist()).issubset({0, 1, 3}) - assert 2 not in out.unique().tolist() diff --git a/tests/unit/test_batchnorm_init.py b/tests/unit/test_batchnorm_init.py deleted file mode 100644 index b68c9ccf..00000000 --- a/tests/unit/test_batchnorm_init.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 test: BatchNorm gamma must initialise around 1, not 0.""" - -import torch -from konfai.network.network import ModuleArgsDict - - -def test_init_func_centres_batchnorm_gamma_on_one() -> None: - # gamma initialised around 0 scaled the normalised activations to ~0, stalling early training. - batch_norm = torch.nn.BatchNorm2d(128) - - ModuleArgsDict.init_func(batch_norm, "normal", 0.02) - - assert abs(batch_norm.weight.mean().item() - 1.0) < 0.02 - assert batch_norm.bias.abs().max().item() < 1e-6 diff --git a/tests/unit/test_blocks.py b/tests/unit/test_blocks.py new file mode 100644 index 00000000..fdf24cfe --- /dev/null +++ b/tests/unit/test_blocks.py @@ -0,0 +1,50 @@ +# 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 + +"""Tests for individual building blocks in ``konfai.network.blocks``.""" + +import pytest +import torch +from konfai.network.blocks import Exit, LatentDistribution, Select, Unsqueeze + + +def test_vae_latent_uses_gaussian_noise(): + """#3 LatentDistributionZ must sample N(0,1), not U[0,1].""" + layer = LatentDistribution.LatentDistributionZ() + mu = torch.zeros(200_000) + log_std = torch.zeros(200_000) + z = layer(mu, log_std) # == epsilon + assert abs(float(z.mean())) < 0.05, "mean should be ~0" + assert abs(float(z.std()) - 1.0) < 0.05, "std should be ~1 (Gaussian), not ~0.29 (uniform)" + + +def test_unsqueeze_forward_accepts_tensor(): + """#8 Unsqueeze.forward(tensor) must work on a single tensor.""" + assert Unsqueeze(dim=1)(torch.randn(3, 4)).shape == (3, 1, 4) + + +def test_select_squeezes_size_one_dims_by_size(): + """#12 Select must squeeze dimensions whose size is 1, not the dim at index 1.""" + out = Select([slice(0, 1), slice(None), slice(None)])(torch.randn(1, 5, 6)) + assert out.shape == (5, 6) + # a tensor with no size-1 dims is unchanged + out2 = Select([slice(None), slice(None)])(torch.randn(4, 5)) + assert out2.shape == (4, 5) + + +def test_debug_exit_block_raises_runtime_error() -> None: + with pytest.raises(RuntimeError, match="debug Exit block"): + Exit()(torch.ones(1)) diff --git a/tests/unit/test_cli_entry_points.py b/tests/unit/test_cli_entry_points.py deleted file mode 100644 index f25474eb..00000000 --- a/tests/unit/test_cli_entry_points.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Tests verifying that CLI subcommands dispatch to the correct backend functions.""" - -import sys - -import konfai.evaluator as evaluator_module -import konfai.main as main_module -import konfai.trainer as trainer_module -import pytest - - -def test_konfai_help_exits_zero(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(sys, "argv", ["konfai", "--help"]) - - with pytest.raises(SystemExit) as exc_info: - main_module.main() - - assert exc_info.value.code == 0 - - -def test_konfai_train_dispatches_correctly(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - - def fake_train(**kwargs) -> None: - captured.update(kwargs) - - monkeypatch.setattr(trainer_module, "train", fake_train) - monkeypatch.setattr(sys, "argv", ["konfai", "TRAIN", "-c", "Config.yml"]) - - main_module.main() - - assert captured["config"] == "Config.yml" - - -def test_konfai_eval_dispatches_correctly(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - - def fake_evaluate(**kwargs) -> None: - captured.update(kwargs) - - monkeypatch.setattr(evaluator_module, "evaluate", fake_evaluate) - monkeypatch.setattr(sys, "argv", ["konfai", "EVALUATION", "-c", "Evaluation.yml"]) - - main_module.main() - - assert captured["evaluations_file"] == "Evaluation.yml" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 77f1741e..52ee35a1 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,7 +1,32 @@ +# 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 + +"""Tests for the config reflection engine (``konfai.utils.config``). + +Covers ``Config`` file handling and error messages, ``apply_config`` type binding +(literals, unions, dicts, booleans), write-back round-trips (including dotted dict +keys), and the config env-var bookkeeping. +""" + +import os from pathlib import Path from typing import Literal import pytest +import ruamel.yaml from konfai.utils.config import Config, apply_config, config from konfai.utils.errors import ConfigError @@ -10,18 +35,9 @@ def _fail_input(_: str) -> str: raise AssertionError("input should not be used") -def _configure_env( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - content: str, - *, - mode: str = "Done", -) -> Path: - config_path = tmp_path / "config.yml" - config_path.write_text(content, encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", mode) - return config_path +# -------------------------------------------------------------------------------------- +# Config file handling and error messages +# -------------------------------------------------------------------------------------- def test_config_missing_file_raises_clear_error_without_prompt( @@ -33,10 +49,17 @@ def test_config_missing_file_raises_clear_error_without_prompt( monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") monkeypatch.setattr("builtins.input", _fail_input) - with pytest.raises(ConfigError, match="does not exist"): + with pytest.raises(ConfigError) as exc_info: with Config("Trainer"): pass + # The error must name the file, the mode, and hint at the fix. + msg = str(exc_info.value) + assert "missing.yml" in msg + assert "does not exist" in msg + assert "KONFAI_CONFIG_MODE=Done" in msg + assert "konfai TRAINING" in msg + def test_config_default_mode_materializes_missing_file( tmp_path: Path, @@ -47,8 +70,8 @@ def test_config_default_mode_materializes_missing_file( monkeypatch.setenv("KONFAI_CONFIG_MODE", "default") monkeypatch.setattr("builtins.input", _fail_input) - with Config("Trainer") as config: - value = config.get_value("train_name", "default|SMOKE") + with Config("Trainer") as config_obj: + value = config_obj.get_value("train_name", "default|SMOKE") assert config_path.exists() assert value == "SMOKE" @@ -57,14 +80,65 @@ def test_config_default_mode_materializes_missing_file( assert "train_name: SMOKE" in content -def test_apply_config_preserves_none_for_optional_nested_objects( +def test_config_missing_env_var_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KONFAI_config_file", raising=False) + + with pytest.raises(KeyError): + Config("Trainer") + + +def test_get_value_returns_default_when_key_absent( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - config_path = tmp_path / "config.yml" - config_path.write_text("Root:\n child: None\n", encoding="utf-8") + config_path = tmp_path / "empty.yml" + config_path.write_text("", encoding="utf-8") monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") + monkeypatch.setenv("KONFAI_CONFIG_MODE", "default") + monkeypatch.setattr("builtins.input", _fail_input) + + with Config("Root") as cfg: + value = cfg.get_value("missing_key", "default|FALLBACK") + + assert value == "FALLBACK" + + +def test_config_raises_on_invalid_yaml_syntax(write_config) -> None: + write_config("key: {unclosed\n", name="broken.yml") + + with pytest.raises(ConfigError) as exc_info: + with Config("Root"): + pass + + msg = str(exc_info.value) + assert "Invalid YAML syntax" in msg + assert "broken.yml" in msg + + +def test_type_mismatch_error_names_field_and_type(write_config) -> None: + write_config("Root:\n count: hello\n") + + class Root: + def __init__(self, count: int = 0) -> None: + self.count = count + + with pytest.raises(ConfigError) as exc_info: + apply_config("Root")(Root)() + + msg = str(exc_info.value) + assert "count" in msg + assert "int" in msg + + +# -------------------------------------------------------------------------------------- +# apply_config type binding +# -------------------------------------------------------------------------------------- + + +def test_apply_config_preserves_none_for_optional_nested_objects(write_config) -> None: + write_config("Root:\n child: None\n") @config("child") class Child: @@ -80,15 +154,8 @@ def __init__(self, child: Child | None = None) -> None: assert root.child is None -def test_apply_config_accepts_literal_value( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - "Root:\n mode: eval\n", - ) +def test_apply_config_accepts_literal_value(write_config) -> None: + write_config("Root:\n mode: eval\n") class Root: def __init__(self, mode: Literal["train", "eval"] = "train") -> None: @@ -99,33 +166,48 @@ def __init__(self, mode: Literal["train", "eval"] = "train") -> None: assert root.mode == "eval" -def test_apply_config_rejects_invalid_literal_value( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - "Root:\n mode: invalid\n", - ) +def test_apply_config_rejects_invalid_literal_value(write_config) -> None: + write_config("Root:\n mode: invalid\n") class Root: def __init__(self, mode: Literal["train", "eval"] = "train") -> None: self.mode = mode - with pytest.raises(ConfigError, match="Invalid value 'invalid'"): + with pytest.raises(ConfigError, match="Invalid value 'invalid'") as exc_info: apply_config("Root")(Root)() + # The error must mention the valid options. + msg = str(exc_info.value) + assert "train" in msg or "eval" in msg + + +@pytest.mark.parametrize( + ("literal", "expected"), + [("true", True), ("1", True), ("yes", True), ("false", False), ("0", False), ("no", False)], +) +def test_apply_config_parses_boolean_strings(write_config, literal: str, expected: bool) -> None: + write_config(f"Root:\n enabled: '{literal}'\n") + + class Root: + def __init__(self, enabled: bool = True) -> None: + self.enabled = enabled + + assert apply_config("Root")(Root)().enabled is expected + + +def test_apply_config_rejects_unknown_boolean_string(write_config) -> None: + write_config("Root:\n enabled: 'sometimes'\n") + + class Root: + def __init__(self, enabled: bool = True) -> None: + self.enabled = enabled + + with pytest.raises(ConfigError, match="expected bool"): + apply_config("Root")(Root)() -def test_apply_config_instantiates_dict_of_nested_objects( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - ("Root:\n children:\n left:\n value: 3\n right:\n value: 7\n"), - ) + +def test_apply_config_instantiates_dict_of_nested_objects(write_config) -> None: + write_config("Root:\n children:\n left:\n value: 3\n right:\n value: 7\n") class Child: def __init__(self, value: int) -> None: @@ -142,15 +224,8 @@ def __init__(self, children: dict[str, Child]) -> None: assert root.children["right"].value == 7 -def test_apply_config_preserves_dict_of_primitives( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - ("Root:\n weights:\n mae: 1\n ssim: 2\n"), - ) +def test_apply_config_preserves_dict_of_primitives(write_config) -> None: + write_config("Root:\n weights:\n mae: 1\n ssim: 2\n") class Root: def __init__(self, weights: dict[str, int]) -> None: @@ -161,15 +236,8 @@ def __init__(self, weights: dict[str, int]) -> None: assert root.weights == {"mae": 1, "ssim": 2} -def test_apply_config_converts_sequence_of_union_scalars( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - ("Root:\n values:\n - '1'\n - 2\n - '3'\n"), - ) +def test_apply_config_converts_sequence_of_union_scalars(write_config) -> None: + write_config("Root:\n values:\n - '1'\n - 2\n - '3'\n") class Root: def __init__(self, values: list[int | float]) -> None: @@ -181,17 +249,10 @@ def __init__(self, values: list[int | float]) -> None: assert all(isinstance(value, int) for value in root.values) -def test_apply_config_binds_scalar_float_or_str_union( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_apply_config_binds_scalar_float_or_str_union(write_config) -> None: # Mirrors the Clip transform (``min_value``/``max_value: float | str``) which accepts numeric # bounds as well as string sentinels such as ``min`` / ``percentile:99.5``. - _configure_env( - tmp_path, - monkeypatch, - "Root:\n low: min\n high: 'percentile:99.5'\n fixed: 1024\n", - ) + write_config("Root:\n low: min\n high: 'percentile:99.5'\n fixed: 1024\n") class Root: def __init__( @@ -212,15 +273,8 @@ def __init__( assert isinstance(root.fixed, float) -def test_apply_config_honors_konfai_without_for_skipped_parameters( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - ("Root:\n kept: 5\n skipped: 42\n"), - ) +def test_apply_config_honors_konfai_without_for_skipped_parameters(write_config) -> None: + write_config("Root:\n kept: 5\n skipped: 42\n") class Root: def __init__(self, kept: int, skipped: int = 0) -> None: @@ -233,40 +287,133 @@ def __init__(self, kept: int, skipped: int = 0) -> None: assert root.skipped == 0 -def test_config_missing_env_var_raises( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("KONFAI_config_file", raising=False) +# -------------------------------------------------------------------------------------- +# Write-back round-trips +# -------------------------------------------------------------------------------------- - with pytest.raises(KeyError): - Config("Trainer") +class _RoundTripRoot: + def __init__(self, weights: dict[str, int] = {"mae": 1, "ssim": 2}) -> None: + self.weights = weights -def test_get_value_returns_default_when_key_absent( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = tmp_path / "empty.yml" - config_path.write_text("", encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "default") - monkeypatch.setattr("builtins.input", _fail_input) - with Config("Root") as cfg: - value = cfg.get_value("missing_key", "default|FALLBACK") +def test_dict_of_primitives_default_round_trips(write_config) -> None: + config_path = write_config("Root: {}\n") # Root present but no 'weights' - assert value == "FALLBACK" + # Run 1: the default materialises and is written back. + first = apply_config("Root")(_RoundTripRoot)() + assert first.weights == {"mae": 1, "ssim": 2} + # The write-back must persist the values, not collapse the dict to an empty mapping. + written = config_path.read_text(encoding="utf-8") + assert "mae" in written and "ssim" in written -def test_config_raises_on_invalid_yaml_syntax( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = tmp_path / "broken.yml" - config_path.write_text("key: {unclosed\n", encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") + # Run 2: reading the written file must return the same dict, not {} (the pre-fix behaviour). + second = apply_config("Root")(_RoundTripRoot)() + assert second.weights == {"mae": 1, "ssim": 2} - with pytest.raises(ConfigError, match=r"broken\.yml"): - with Config("Root"): - pass + +# A dotted dict key (e.g. a PerceptualLoss module path ``UNetBlock_0.DownConvBlock.Activation_1``) +# must be treated as a single flat config key. Before the fix, ``Config.__init__`` split it on ``.`` +# into separate navigation levels, so the user's value was never found (code defaults were used) +# and the write-back exploded the key into a bogus nested subtree. + + +class _DottedChild: + def __init__(self, value: int = 1) -> None: + self.value = value + + +class _DottedRoot: + def __init__(self, children: dict[str, _DottedChild] = {"a.b.c": _DottedChild(1)}) -> None: + self.children = children + + +def test_apply_config_honors_value_under_dotted_dict_key(write_config) -> None: + write_config("Root:\n children:\n a.b.c:\n value: 99\n") + + root = apply_config("Root")(_DottedRoot)() + + assert list(root.children) == ["a.b.c"] + # Before the fix the dotted key was split and this was silently 1 (the code default). + assert root.children["a.b.c"].value == 99 + + +def test_apply_config_does_not_explode_dotted_dict_key_on_writeback(write_config) -> None: + config_path = write_config("Root:\n children:\n a.b.c:\n value: 99\n") + + apply_config("Root")(_DottedRoot)() + + data = ruamel.yaml.YAML().load(config_path.read_text(encoding="utf-8")) + children = data["Root"]["children"] + # Before the fix, children also contained an exploded ``a: {b: {c: {value: 1}}}`` subtree. + assert set(children) == {"a.b.c"} + assert "a" not in children + assert children["a.b.c"]["value"] == 99 + + +def test_apply_config_dotted_dict_key_round_trips(write_config) -> None: + config_path = write_config("Root:\n children:\n a.b.c:\n value: 99\n") + + first = apply_config("Root")(_DottedRoot)() + assert first.children["a.b.c"].value == 99 + after_first = config_path.read_text(encoding="utf-8") + + # Second run reads the written-back file: value preserved and write-back idempotent. + second = apply_config("Root")(_DottedRoot)() + assert second.children["a.b.c"].value == 99 + assert config_path.read_text(encoding="utf-8") == after_first + + +def test_apply_config_colon_and_plain_dict_keys_unaffected(write_config) -> None: + # Backward-compat guard: keys without ``.`` (``:``-separated module paths, plain + # names) must bind exactly as before and must not be escaped/exploded. + config_path = write_config("R:\n m:\n X:Head:Conv:\n value: 5\n plain:\n value: 8\n") + + class R: + def __init__(self, m: dict[str, _DottedChild] = {"X:Head:Conv": _DottedChild(1), "plain": _DottedChild(1)}): + self.m = m + + root = apply_config("R")(R)() + + assert root.m["X:Head:Conv"].value == 5 + assert root.m["plain"].value == 8 + data = ruamel.yaml.YAML().load(config_path.read_text(encoding="utf-8")) + assert set(data["R"]["m"]) == {"X:Head:Conv", "plain"} + + +# -------------------------------------------------------------------------------------- +# Config env-var bookkeeping +# -------------------------------------------------------------------------------------- + + +def test_apply_config_restores_config_env(write_config, monkeypatch: pytest.MonkeyPatch) -> None: + write_config("Root:\n Child:\n value: 7\n") + monkeypatch.setenv("KONFAI_CONFIG_PATH", "before.path") + monkeypatch.setenv("KONFAI_CONFIG_VARIABLE", "before.variable") + + @config("Child") + class Child: + def __init__(self, value: int = 0) -> None: + self.value = value + + child = apply_config("Root")(Child)() + + assert child.value == 7 + assert os.environ["KONFAI_CONFIG_PATH"] == "before.path" + assert os.environ["KONFAI_CONFIG_VARIABLE"] == "before.variable" + + +def test_apply_config_keeps_config_path_during_constructor_call(write_config) -> None: + write_config("Root:\n Child:\n value: 7\n") + + @config("Child") + class Child: + def __init__(self, value: int = 0) -> None: + self.value = value + self.config_path = os.environ["KONFAI_CONFIG_PATH"] + + child = apply_config("Root")(Child)() + + assert child.value == 7 + assert child.config_path == "Root.Child" diff --git a/tests/unit/test_config_dict_roundtrip.py b/tests/unit/test_config_dict_roundtrip.py deleted file mode 100644 index 151e1efb..00000000 --- a/tests/unit/test_config_dict_roundtrip.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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 test: a dict[str, primitive] default must survive the config write-back.""" - -from pathlib import Path - -import pytest -from konfai.utils.config import apply_config - - -class _Root: - def __init__(self, weights: dict[str, int] = {"mae": 1, "ssim": 2}) -> None: - self.weights = weights - - -def _setup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: - config_path = tmp_path / "config.yml" - config_path.write_text("Root: {}\n", encoding="utf-8") # Root present but no 'weights' - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - return config_path - - -def test_dict_of_primitives_default_round_trips(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - config_path = _setup(tmp_path, monkeypatch) - - # Run 1: the default materialises and is written back. - first = apply_config("Root")(_Root)() - assert first.weights == {"mae": 1, "ssim": 2} - - # The write-back must persist the values, not collapse the dict to an empty mapping. - written = config_path.read_text(encoding="utf-8") - assert "mae" in written and "ssim" in written - - # Run 2: reading the written file must return the same dict, not {} (the pre-fix behaviour). - second = apply_config("Root")(_Root)() - assert second.weights == {"mae": 1, "ssim": 2} diff --git a/tests/unit/test_config_dotted_dict_keys.py b/tests/unit/test_config_dotted_dict_keys.py deleted file mode 100644 index a9fdd0e7..00000000 --- a/tests/unit/test_config_dotted_dict_keys.py +++ /dev/null @@ -1,129 +0,0 @@ -# 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: dict[str, Object] entries keyed by a dotted string. - -A dotted dict key (e.g. a PerceptualLoss module path -``UNetBlock_0.DownConvBlock.Activation_1``) must be treated as a single flat -config key. Before the fix, ``Config.__init__`` split it on ``.`` into separate -navigation levels, so the user's value was never found (code defaults were used) -and the write-back exploded the key into a bogus nested subtree. -""" - -from pathlib import Path - -import pytest -import ruamel.yaml -from konfai.utils.config import apply_config - - -def _configure_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, content: str) -> Path: - config_path = tmp_path / "config.yml" - config_path.write_text(content, encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - return config_path - - -class _Child: - def __init__(self, value: int = 1) -> None: - self.value = value - - -class _Root: - def __init__(self, children: dict[str, _Child] = {"a.b.c": _Child(1)}) -> None: - self.children = children - - -def test_apply_config_honors_value_under_dotted_dict_key( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env( - tmp_path, - monkeypatch, - "Root:\n children:\n a.b.c:\n value: 99\n", - ) - - root = apply_config("Root")(_Root)() - - assert list(root.children) == ["a.b.c"] - # Before the fix the dotted key was split and this was silently 1 (the code default). - assert root.children["a.b.c"].value == 99 - - -def test_apply_config_does_not_explode_dotted_dict_key_on_writeback( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = _configure_env( - tmp_path, - monkeypatch, - "Root:\n children:\n a.b.c:\n value: 99\n", - ) - - apply_config("Root")(_Root)() - - data = ruamel.yaml.YAML().load(config_path.read_text(encoding="utf-8")) - children = data["Root"]["children"] - # Before the fix, children also contained an exploded ``a: {b: {c: {value: 1}}}`` subtree. - assert set(children) == {"a.b.c"} - assert "a" not in children - assert children["a.b.c"]["value"] == 99 - - -def test_apply_config_dotted_dict_key_round_trips( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = _configure_env( - tmp_path, - monkeypatch, - "Root:\n children:\n a.b.c:\n value: 99\n", - ) - - first = apply_config("Root")(_Root)() - assert first.children["a.b.c"].value == 99 - after_first = config_path.read_text(encoding="utf-8") - - # Second run reads the written-back file: value preserved and write-back idempotent. - second = apply_config("Root")(_Root)() - assert second.children["a.b.c"].value == 99 - assert config_path.read_text(encoding="utf-8") == after_first - - -def test_apply_config_colon_and_plain_dict_keys_unaffected( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - # Backward-compat guard: keys without ``.`` (``:``-separated module paths, plain - # names) must bind exactly as before and must not be escaped/exploded. - config_path = _configure_env( - tmp_path, - monkeypatch, - "R:\n m:\n X:Head:Conv:\n value: 5\n plain:\n value: 8\n", - ) - - class R: - def __init__(self, m: dict[str, _Child] = {"X:Head:Conv": _Child(1), "plain": _Child(1)}) -> None: - self.m = m - - root = apply_config("R")(R)() - - assert root.m["X:Head:Conv"].value == 5 - assert root.m["plain"].value == 8 - data = ruamel.yaml.YAML().load(config_path.read_text(encoding="utf-8")) - assert set(data["R"]["m"]) == {"X:Head:Conv", "plain"} diff --git a/tests/unit/test_config_validation.py b/tests/unit/test_config_validation.py deleted file mode 100644 index d9157bbd..00000000 --- a/tests/unit/test_config_validation.py +++ /dev/null @@ -1,160 +0,0 @@ -# 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 - -"""Tests for improved ConfigError messages introduced in Phase 05.""" - -from pathlib import Path -from typing import Literal - -import pytest -from konfai.utils.config import Config, apply_config -from konfai.utils.errors import ConfigError - - -def _configure_env( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - content: str, - *, - mode: str = "Done", -) -> Path: - config_path = tmp_path / "config.yml" - config_path.write_text(content, encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", mode) - return config_path - - -def test_missing_config_file_error_contains_path( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = tmp_path / "missing.yml" - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - - with pytest.raises(ConfigError) as exc_info: - with Config("Trainer"): - pass - - msg = str(exc_info.value) - assert "missing.yml" in msg - assert "does not exist" in msg - - -def test_missing_config_file_error_contains_mode_and_hint( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = tmp_path / "missing.yml" - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - - with pytest.raises(ConfigError) as exc_info: - with Config("Trainer"): - pass - - msg = str(exc_info.value) - assert "KONFAI_CONFIG_MODE=Done" in msg - assert "konfai TRAINING" in msg - - -def test_invalid_yaml_syntax_raises_config_error( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - config_path = tmp_path / "broken.yml" - config_path.write_text("key: {unclosed\n", encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - - with pytest.raises(ConfigError) as exc_info: - with Config("Root"): - pass - - msg = str(exc_info.value) - assert "Invalid YAML syntax" in msg - assert "broken.yml" in msg - - -def test_type_mismatch_error_names_field_and_type( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env(tmp_path, monkeypatch, "Root:\n count: hello\n") - - class Root: - def __init__(self, count: int = 0) -> None: - self.count = count - - with pytest.raises(ConfigError) as exc_info: - apply_config("Root")(Root)() - - msg = str(exc_info.value) - assert "count" in msg - assert "int" in msg - - -@pytest.mark.parametrize( - ("literal", "expected"), - [("true", True), ("1", True), ("yes", True), ("false", False), ("0", False), ("no", False)], -) -def test_apply_config_parses_boolean_strings( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - literal: str, - expected: bool, -) -> None: - _configure_env(tmp_path, monkeypatch, f"Root:\n enabled: '{literal}'\n") - - class Root: - def __init__(self, enabled: bool = True) -> None: - self.enabled = enabled - - assert apply_config("Root")(Root)().enabled is expected - - -def test_apply_config_rejects_unknown_boolean_string( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env(tmp_path, monkeypatch, "Root:\n enabled: 'sometimes'\n") - - class Root: - def __init__(self, enabled: bool = True) -> None: - self.enabled = enabled - - with pytest.raises(ConfigError, match="expected bool"): - apply_config("Root")(Root)() - - -def test_invalid_literal_raises_config_error_with_options( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - _configure_env(tmp_path, monkeypatch, "Root:\n mode: invalid\n") - - class Root: - def __init__(self, mode: Literal["train", "eval"] = "train") -> None: - self.mode = mode - - with pytest.raises(ConfigError) as exc_info: - apply_config("Root")(Root)() - - msg = str(exc_info.value) - assert "invalid" in msg - # The error must mention the valid options - assert "train" in msg or "eval" in msg diff --git a/tests/unit/test_criterion_loaders.py b/tests/unit/test_criterion_loaders.py deleted file mode 100644 index 706d7759..00000000 --- a/tests/unit/test_criterion_loaders.py +++ /dev/null @@ -1,43 +0,0 @@ -from types import SimpleNamespace -from typing import cast - -import pytest - -import konfai.network.network as network_module - - -def test_network_criterion_loader_resets_scheduler_state(monkeypatch: pytest.MonkeyPatch) -> None: - class DummyMeasure: - def __init__(self) -> None: - pass - - class DummySchedulerLoader: - def __init__(self) -> None: - self.nb_step = 3 - - def getschedulers(self, key: str, scheduler_classname: str): - return f"{key}:{scheduler_classname}" - - monkeypatch.setattr(network_module, "apply_config", lambda *args, **kwargs: (lambda cls: cls)) - monkeypatch.setattr(network_module, "konfai_root", lambda: "Trainer") - monkeypatch.setattr( - network_module, - "get_module", - lambda classpath, default: (SimpleNamespace(Measure=DummyMeasure, __name__="torch.optim"), "Measure"), - ) - - attr = network_module.CriterionsAttr( - schedulers=cast( - dict[str, network_module.LossSchedulersLoader], - {"Constant": DummySchedulerLoader()}, - ) - ) - loader = network_module.CriterionsLoader({"dummy:Measure": attr}) - - loader.get_criterions("DemoModel", "Output", "Target") - first_schedulers = dict(attr.schedulers) - loader.get_criterions("DemoModel", "Output", "Target") - - assert attr.isTorchCriterion is True - assert len(attr.schedulers) == 1 - assert attr.schedulers == first_schedulers diff --git a/tests/unit/test_crop_transform_shape.py b/tests/unit/test_crop_transform_shape.py deleted file mode 100644 index 8457ff5b..00000000 --- a/tests/unit/test_crop_transform_shape.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 test: Crop.transform_shape predicts the spatial crop exactly. - -Patch planning consumes transform_shape's output, so it must equal the spatial shape ``__call__`` -actually produces. The pre-fix code treated ``shape[0]`` as a channel dim and paired the crop box -with ``shape[1:]``, shifting every axis by one and returning a wrong shape. -""" - -import numpy as np -from konfai.data.transform import Crop -from konfai.utils.dataset import Attribute - - -def test_crop_transform_shape_matches_spatial_crop() -> None: - attribute = Attribute() - attribute["box"] = np.array([[2, 3], [1, 1], [4, 2]]) # (start, end-distance) per spatial axis - - out = Crop().transform_shape("CT", "CASE_001", [10, 20, 30], attribute) - - # 10-2-3, 20-1-1, 30-4-2 — each spatial axis cropped by its own box row. - assert out == [5, 18, 24] diff --git a/tests/unit/test_data_manager.py b/tests/unit/test_data_manager.py new file mode 100644 index 00000000..88fe8e96 --- /dev/null +++ b/tests/unit/test_data_manager.py @@ -0,0 +1,432 @@ +# 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 + +"""Tests for ``konfai.data.data_manager``: DDP sharding, train/validation split, +cache workers, and DatasetIter (streaming transforms, inline augmentations).""" + +import os +import subprocess +import sys +from typing import cast + +import numpy as np +import pytest +import torch +from konfai.data.augmentation import DataAugmentation, DataAugmentationsList +from konfai.data.data_manager import Data, DatasetIter, DataTrain, Group, GroupTransform, _cache_worker_count +from konfai.data.patching import DatasetManager, DatasetPatch +from konfai.data.transform import TensorCast +from konfai.utils.dataset import Attribute, Dataset +from konfai.utils.runtime import State + +# -------------------------------------------------------------------------------------- +# Data._split — TRAIN/RESUME shards must be equal length to avoid a DDP hang +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("state", [State.TRAIN, State.RESUME]) +def test_train_split_equalises_indivisible_shards(monkeypatch: pytest.MonkeyPatch, state: State) -> None: + # DDP(static_graph=True) needs every rank to run the same number of backward all-reduces per + # epoch. A contiguous split of 7 patches over 3 ranks gives [2, 2, 3] and hangs NCCL on the + # extra step; drop_last equalises to [2, 2, 2]. + monkeypatch.setenv("KONFAI_STATE", str(state)) + + shards = Data._split([(index, 0, 0) for index in range(7)], 3) + + assert [len(shard) for shard in shards] == [2, 2, 2] + flattened = [item for shard in shards for item in shard] + assert len(flattened) == len(set(flattened)) # never duplicated across ranks + + +def test_train_split_two_ranks_indivisible(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("KONFAI_STATE", str(State.TRAIN)) + + mapping = [(i, 0, 0) for i in range(5)] + shards = Data._split(mapping, 2) + + # drop_last semantics: equal-length shards, the tail is dropped, no sample duplicated. + assert [len(shard) for shard in shards] == [2, 2] + flattened = [item for shard in shards for item in shard] + assert len(flattened) == len(set(flattened)) + assert set(flattened).issubset(set(mapping)) + + +def test_train_split_single_process_keeps_everything(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("KONFAI_STATE", str(State.TRAIN)) + mapping = [(i, 0, 0) for i in range(5)] + assert Data._split(mapping, 1) == [mapping] # world_size == 1 is a no-op + + +# -------------------------------------------------------------------------------------- +# DataTrain train/validation split — reproducible and seeded from sorted names +# -------------------------------------------------------------------------------------- + +_SPLIT_PROBE = """ +import os +import random + +os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") +os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") + +from konfai.data.data_manager import DataTrain + +names = [f"CASE_{i:03d}" for i in range(20)] +data = DataTrain(augmentations=None, validation="0:4") +data._resolve_dataset_sources = lambda: {} +data._resolve_common_names = lambda datasets: ({}, set(names)) +data._get_datasets = lambda case_names, dataset_name, augmentations: ({}, []) +random.seed(1234) +data._prepare_datasets() +print(";".join(data._prepared_train_names)) +print(";".join(data._prepared_validation_names)) +""" + + +def test_train_validation_split_is_reproducible_across_interpreters(): + """Same seed → same split, whatever the interpreter's string-hash randomization.""" + outputs = [] + for hash_seed in ("0", "424242"): + env = dict(os.environ, PYTHONHASHSEED=hash_seed) + result = subprocess.run( + [sys.executable, "-c", _SPLIT_PROBE], + env=env, + capture_output=True, + text=True, + check=True, + ) + outputs.append(result.stdout) + assert outputs[0] == outputs[1] + train_names, validation_names = (line.split(";") for line in outputs[0].splitlines()) + assert len(train_names) == 16 + assert len(validation_names) == 4 + assert set(train_names).isdisjoint(validation_names) + + +def test_train_split_shuffle_draws_from_sorted_names(monkeypatch): + """The seeded shuffle must receive the case names in sorted order and drive the split.""" + captured: dict[str, list[str]] = {} + + def fake_sample(population, k): + captured["population"] = list(population) + assert k == len(population) + return list(reversed(population)) + + monkeypatch.setattr("konfai.data.data_manager.random.sample", fake_sample) + + data = DataTrain(augmentations=None, validation="0:2") + names = {"CASE_010", "CASE_002", "CASE_001", "CASE_005", "CASE_003"} + data._resolve_dataset_sources = lambda: {} + data._resolve_common_names = lambda datasets: ({}, names) + data._get_datasets = lambda case_names, dataset_name, augmentations: ({}, []) + data._prepare_datasets() + + assert captured["population"] == sorted(names) + assert data._prepared_validation_names == ["CASE_010", "CASE_005"] + assert data._prepared_train_names == ["CASE_003", "CASE_002", "CASE_001"] + + +# -------------------------------------------------------------------------------------- +# B18 - caching worker count must never fall below one +# -------------------------------------------------------------------------------------- + + +def test_cache_worker_count_never_drops_below_one() -> None: + # 2 CPUs shared across 4 GPUs would be 2 // 4 == 0 without the floor. + assert _cache_worker_count(2, 4) == 1 + assert _cache_worker_count(1, 4) == 1 + assert _cache_worker_count(8, 2) == 4 + assert _cache_worker_count(7, 2) == 3 + assert _cache_worker_count(4, 0) == 4 # no device -> divisor 1 + + +# -------------------------------------------------------------------------------------- +# B3 - patch streaming must persist TensorCast dtype for the inverse +# -------------------------------------------------------------------------------------- + + +def _image_attributes(origin: list[float], spacing: list[float]) -> Attribute: + attributes = Attribute() + attributes["Origin"] = np.asarray(origin, dtype=np.float64) + attributes["Spacing"] = np.asarray(spacing, dtype=np.float64) + attributes["Direction"] = np.eye(len(origin), dtype=np.float64).reshape(-1) + return attributes + + +class _StreamingDatasetStub: + def __init__(self, volume: np.ndarray) -> None: + self.volume = volume + + def get_infos(self, group_src: str, name: str) -> tuple[list[int], Attribute]: + return list(self.volume.shape), _image_attributes([0.0, 0.0], [1.0, 1.0]) + + def read_data(self, group_src: str, name: str) -> tuple[np.ndarray, Attribute]: + return self.volume.copy(), _image_attributes([0.0, 0.0], [1.0, 1.0]) + + def read_data_slice(self, group_src: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: + return self.volume[slices].copy(), _image_attributes([0.0, 0.0], [1.0, 1.0]) + + def read_data_statistics(self, group_src: str, name: str, channels: list[int] | None = None) -> dict[str, float]: + data = self.volume if channels is None else self.volume[channels] + return { + "min": float(data.min()), + "max": float(data.max()), + "mean": float(data.mean()), + "std": float(data.std(ddof=1)), + } + + +def test_streaming_tensorcast_persists_source_dtype_for_inverse() -> None: + volume = np.arange(1 * 4 * 4, dtype=np.int16).reshape(1, 4, 4) + dataset_stub = _StreamingDatasetStub(volume) + manager = DatasetManager( + index=0, + group_src="CT", + group_dest="CT", + name="CASE_000", + dataset=cast(Dataset, dataset_stub), + patch=DatasetPatch([2, 2]), + transforms=[TensorCast(dtype="float32")], + data_augmentations_list=[], + ) + dataset_iter = DatasetIter( + rank=0, + data={"CT": [manager]}, + mapping=[(0, 0, 1)], + groups_src={"CT": Group(groups_dest={"CT": GroupTransform(transforms=None, patch_transforms=None)})}, + inline_augmentations=False, + data_augmentations_list=[], + patch_size=[2, 2], + overlap=None, + buffer_size=1, + use_cache=False, + ) + + sample = dataset_iter[0]["CT"].tensor + + assert sample.dtype == torch.float32 + # The forward cast records the source dtype on the persistent case attribute. + assert "dtype" in manager.cache_attributes[0] + # ... so the write-time inverse can restore the original dtype without crashing. + restored = TensorCast(dtype="float32").inverse("CASE_000", sample, Attribute(manager.cache_attributes[0])) + assert restored.dtype == torch.int16 + + +# -------------------------------------------------------------------------------------- +# DatasetIter — inline augmentations and per-case state draws +# -------------------------------------------------------------------------------------- + + +class _DummyDataset: + def __init__(self, array: np.ndarray) -> None: + self.array = array + + def get_infos(self, group_src: str, name: str) -> tuple[list[int], Attribute]: + return list(self.array.shape), Attribute({"name": name, "group": group_src}) + + def read_data(self, group_src: str, name: str) -> tuple[np.ndarray, Attribute]: + return self.array.copy(), Attribute({"name": name, "group": group_src}) + + +class _CountingOffsetAugmentation(DataAugmentation): + def __init__(self) -> None: + super().__init__() + self.compute_calls = 0 + + def _state_init( + self, + index: int, + shapes: list[list[int]], + caches_attribute: list[Attribute], + ) -> list[list[int]]: + return shapes + + def _compute( + self, + name: str, + index: int, + tensors: list[torch.Tensor], + ) -> list[torch.Tensor]: + self.compute_calls += 1 + return [tensor + (offset + 1) for offset, tensor in enumerate(tensors)] + + def _inverse(self, index: int, a: int, tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +def _make_manager(dataset: Dataset, augmentations: DataAugmentationsList, group_dest: str = "dest") -> DatasetManager: + return DatasetManager( + index=0, + group_src="src", + group_dest=group_dest, + name="case_000", + dataset=dataset, + patch=None, + transforms=[], + data_augmentations_list=[augmentations], + ) + + +def test_inline_augmentations_are_loaded_on_demand() -> None: + base = np.arange(4, dtype=np.float32).reshape(1, 2, 2) + dataset = cast(Dataset, _DummyDataset(base)) + augmentation = _CountingOffsetAugmentation() + augmentation.load(1.0) + + augmentations = DataAugmentationsList(nb=2, data_augmentations={}) + augmentations.data_augmentations = [augmentation] + + manager = _make_manager(dataset, augmentations) + dataset_iter = DatasetIter( + rank=0, + data={"dest": [manager]}, + mapping=[(0, 0, 0), (0, 1, 0), (0, 2, 0)], + groups_src={"src": Group(groups_dest={"dest": GroupTransform(transforms=None, patch_transforms=None)})}, + inline_augmentations=True, + data_augmentations_list=[augmentations], + patch_size=None, + overlap=None, + buffer_size=1, + use_cache=True, + ) + + base_sample = dataset_iter[0]["dest"].tensor + assert augmentation.compute_calls == 0 + assert manager.loaded is True + assert manager.augmentationLoaded is False + assert torch.equal(base_sample, torch.from_numpy(base)) + + first_augmented_sample = dataset_iter[1]["dest"].tensor + assert augmentation.compute_calls == 1 + assert manager.augmentationLoaded is True + assert torch.equal(first_augmented_sample, torch.from_numpy(base) + 1) + + second_augmented_sample = dataset_iter[2]["dest"].tensor + assert augmentation.compute_calls == 1 + assert torch.equal(second_augmented_sample, torch.from_numpy(base) + 2) + + +def test_dataset_iter_can_skip_augmentation_loading_when_validation_disables_them() -> None: + base = np.arange(4, dtype=np.float32).reshape(1, 2, 2) + dataset = cast(Dataset, _DummyDataset(base)) + augmentation = _CountingOffsetAugmentation() + augmentation.load(1.0) + + augmentations = DataAugmentationsList(nb=2, data_augmentations={}) + augmentations.data_augmentations = [augmentation] + + manager = _make_manager(dataset, augmentations) + dataset_iter = DatasetIter( + rank=0, + data={"dest": [manager]}, + mapping=[(0, 0, 0)], + groups_src={"src": Group(groups_dest={"dest": GroupTransform(transforms=None, patch_transforms=None)})}, + inline_augmentations=False, + data_augmentations_list=[augmentations], + patch_size=None, + overlap=None, + buffer_size=1, + apply_augmentations=False, + use_cache=True, + ) + + dataset_iter.load("Validation") + base_sample = dataset_iter[0]["dest"].tensor + + assert augmentation.compute_calls == 0 + assert manager.loaded is True + assert manager.augmentationLoaded is False + assert torch.equal(base_sample, torch.from_numpy(base)) + + +# -------------------------------------------------------------------------------------- +# B11 - reset_augmentation must draw the shared state once per case, not per group +# -------------------------------------------------------------------------------------- + + +class _DrawCountingAugmentation(DataAugmentation): + """Shape-shifting augmentation whose output depends on the draw order.""" + + def __init__(self) -> None: + super().__init__() + self.draws = 0 + + def _state_init(self, index: int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]: + self.draws += 1 + new_shape = [2, 4] if self.draws == 1 else [4, 4] + return [list(new_shape) for _ in shapes] + + def _compute(self, name: str, index: int, tensors: list[torch.Tensor]) -> list[torch.Tensor]: + return tensors + + def _inverse(self, index: int, a: int, tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +def test_reset_augmentation_shares_one_draw_across_destination_groups() -> None: + array = np.zeros((1, 4, 4), dtype=np.float32) + dataset = cast(Dataset, _DummyDataset(array)) + augmentation = _DrawCountingAugmentation() + augmentation.load(1.0) + augmentations = DataAugmentationsList(nb=1, data_augmentations={}) + augmentations.data_augmentations = [augmentation] + + manager_a = DatasetManager( + index=0, + group_src="src", + group_dest="destA", + name="case_000", + dataset=dataset, + patch=DatasetPatch([2, 2]), + transforms=[], + data_augmentations_list=[augmentations], + ) + manager_b = DatasetManager( + index=0, + group_src="src", + group_dest="destB", + name="case_000", + dataset=dataset, + patch=DatasetPatch([2, 2]), + transforms=[], + data_augmentations_list=[augmentations], + ) + dataset_iter = DatasetIter( + rank=0, + data={"destA": [manager_a], "destB": [manager_b]}, + mapping=[(0, 0, 0), (0, 1, 0)], + groups_src={ + "src": Group( + groups_dest={ + "destA": GroupTransform(transforms=None, patch_transforms=None), + "destB": GroupTransform(transforms=None, patch_transforms=None), + } + ) + }, + inline_augmentations=True, + data_augmentations_list=[augmentations], + patch_size=[2, 2], + overlap=None, + buffer_size=1, + use_cache=False, + ) + + augmentation.draws = 0 + dataset_iter.reset_augmentation("Train") + + # A single random draw feeds every destination group of the case. + assert augmentation.draws == 1 + # Both groups therefore rebuild their augmented patch grid from the same shape. + assert manager_a.patch.get_size(1) == manager_b.patch.get_size(1) diff --git a/tests/unit/test_data_pipeline_audit.py b/tests/unit/test_data_pipeline_audit.py deleted file mode 100644 index 24d370fa..00000000 --- a/tests/unit/test_data_pipeline_audit.py +++ /dev/null @@ -1,262 +0,0 @@ -# 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 data-pipeline audit fixes (patching / caching / streaming).""" - -import os -import threading -from typing import cast - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import numpy as np -import torch -from konfai.data.augmentation import DataAugmentation, DataAugmentationsList -from konfai.data.data_manager import DatasetIter, Group, GroupTransform, _cache_worker_count -from konfai.data.patching import Cosinus, DatasetManager, DatasetPatch, Mean -from konfai.data.transform import TensorCast -from konfai.utils.dataset import Attribute, Dataset, _get_h5_file_lock - - -def _image_attributes(origin: list[float], spacing: list[float]) -> Attribute: - attributes = Attribute() - attributes["Origin"] = np.asarray(origin, dtype=np.float64) - attributes["Spacing"] = np.asarray(spacing, dtype=np.float64) - attributes["Direction"] = np.eye(len(origin), dtype=np.float64).reshape(-1) - return attributes - - -# -------------------------------------------------------------------------------------- -# B10 - PathCombine.set_patch_config with overlap == 0 -# -------------------------------------------------------------------------------------- - - -def test_path_combine_overlap_zero_uses_uniform_weights() -> None: - """overlap=0 tiles patches without overlap, so the blend window is all ones.""" - for combine_cls in (Mean, Cosinus): - combine = combine_cls() - combine.set_patch_config([8, 8, 8], 0) # must not raise - assert combine.data.shape == (8, 8, 8) - assert torch.equal(combine.data, torch.ones(8, 8, 8)) - - -def test_path_combine_overlap_zero_leaves_tensor_unchanged() -> None: - combine = Mean() - combine.set_patch_config([4, 4], 0) - tensor = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4) - assert torch.equal(combine(tensor), tensor) - - -# -------------------------------------------------------------------------------------- -# B18 - caching worker count must never fall below one -# -------------------------------------------------------------------------------------- - - -def test_cache_worker_count_never_drops_below_one() -> None: - # 2 CPUs shared across 4 GPUs would be 2 // 4 == 0 without the floor. - assert _cache_worker_count(2, 4) == 1 - assert _cache_worker_count(1, 4) == 1 - assert _cache_worker_count(8, 2) == 4 - assert _cache_worker_count(7, 2) == 3 - assert _cache_worker_count(4, 0) == 4 # no device -> divisor 1 - - -# -------------------------------------------------------------------------------------- -# B3 - patch streaming must persist TensorCast dtype for the inverse -# -------------------------------------------------------------------------------------- - - -class _StreamingDatasetStub: - def __init__(self, volume: np.ndarray) -> None: - self.volume = volume - - def get_infos(self, group_src: str, name: str) -> tuple[list[int], Attribute]: - return list(self.volume.shape), _image_attributes([0.0, 0.0], [1.0, 1.0]) - - def read_data(self, group_src: str, name: str) -> tuple[np.ndarray, Attribute]: - return self.volume.copy(), _image_attributes([0.0, 0.0], [1.0, 1.0]) - - def read_data_slice(self, group_src: str, name: str, slices: tuple[slice, ...]) -> tuple[np.ndarray, Attribute]: - return self.volume[slices].copy(), _image_attributes([0.0, 0.0], [1.0, 1.0]) - - def read_data_statistics(self, group_src: str, name: str, channels: list[int] | None = None) -> dict[str, float]: - data = self.volume if channels is None else self.volume[channels] - return { - "min": float(data.min()), - "max": float(data.max()), - "mean": float(data.mean()), - "std": float(data.std(ddof=1)), - } - - -def test_streaming_tensorcast_persists_source_dtype_for_inverse() -> None: - volume = np.arange(1 * 4 * 4, dtype=np.int16).reshape(1, 4, 4) - dataset_stub = _StreamingDatasetStub(volume) - manager = DatasetManager( - index=0, - group_src="CT", - group_dest="CT", - name="CASE_000", - dataset=cast(Dataset, dataset_stub), - patch=DatasetPatch([2, 2]), - transforms=[TensorCast(dtype="float32")], - data_augmentations_list=[], - ) - dataset_iter = DatasetIter( - rank=0, - data={"CT": [manager]}, - mapping=[(0, 0, 1)], - groups_src={"CT": Group(groups_dest={"CT": GroupTransform(transforms=None, patch_transforms=None)})}, - inline_augmentations=False, - data_augmentations_list=[], - patch_size=[2, 2], - overlap=None, - buffer_size=1, - use_cache=False, - ) - - sample = dataset_iter[0]["CT"].tensor - - assert sample.dtype == torch.float32 - # The forward cast records the source dtype on the persistent case attribute. - assert "dtype" in manager.cache_attributes[0] - # ... so the write-time inverse can restore the original dtype without crashing. - restored = TensorCast(dtype="float32").inverse("CASE_000", sample, Attribute(manager.cache_attributes[0])) - assert restored.dtype == torch.int16 - - -# -------------------------------------------------------------------------------------- -# B11 - reset_augmentation must draw the shared state once per case, not per group -# -------------------------------------------------------------------------------------- - - -class _DummyDataset: - def __init__(self, array: np.ndarray) -> None: - self.array = array - - def get_infos(self, group_src: str, name: str) -> tuple[list[int], Attribute]: - return list(self.array.shape), Attribute({"name": name, "group": group_src}) - - def read_data(self, group_src: str, name: str) -> tuple[np.ndarray, Attribute]: - return self.array.copy(), Attribute({"name": name, "group": group_src}) - - -class _DrawCountingAugmentation(DataAugmentation): - """Shape-shifting augmentation whose output depends on the draw order.""" - - def __init__(self) -> None: - super().__init__() - self.draws = 0 - - def _state_init(self, index: int, shapes: list[list[int]], caches_attribute: list[Attribute]) -> list[list[int]]: - self.draws += 1 - new_shape = [2, 4] if self.draws == 1 else [4, 4] - return [list(new_shape) for _ in shapes] - - def _compute(self, name: str, index: int, tensors: list[torch.Tensor]) -> list[torch.Tensor]: - return tensors - - def _inverse(self, index: int, a: int, tensor: torch.Tensor) -> torch.Tensor: - return tensor - - -def test_reset_augmentation_shares_one_draw_across_destination_groups() -> None: - array = np.zeros((1, 4, 4), dtype=np.float32) - dataset = cast(Dataset, _DummyDataset(array)) - augmentation = _DrawCountingAugmentation() - augmentation.load(1.0) - augmentations = DataAugmentationsList(nb=1, data_augmentations={}) - augmentations.data_augmentations = [augmentation] - - manager_a = DatasetManager( - index=0, - group_src="src", - group_dest="destA", - name="case_000", - dataset=dataset, - patch=DatasetPatch([2, 2]), - transforms=[], - data_augmentations_list=[augmentations], - ) - manager_b = DatasetManager( - index=0, - group_src="src", - group_dest="destB", - name="case_000", - dataset=dataset, - patch=DatasetPatch([2, 2]), - transforms=[], - data_augmentations_list=[augmentations], - ) - dataset_iter = DatasetIter( - rank=0, - data={"destA": [manager_a], "destB": [manager_b]}, - mapping=[(0, 0, 0), (0, 1, 0)], - groups_src={ - "src": Group( - groups_dest={ - "destA": GroupTransform(transforms=None, patch_transforms=None), - "destB": GroupTransform(transforms=None, patch_transforms=None), - } - ) - }, - inline_augmentations=True, - data_augmentations_list=[augmentations], - patch_size=[2, 2], - overlap=None, - buffer_size=1, - use_cache=False, - ) - - augmentation.draws = 0 - dataset_iter.reset_augmentation("Train") - - # A single random draw feeds every destination group of the case. - assert augmentation.draws == 1 - # Both groups therefore rebuild their augmented patch grid from the same shape. - assert manager_a.patch.get_size(1) == manager_b.patch.get_size(1) - - -# -------------------------------------------------------------------------------------- -# B6 - concurrent HDF5 access is serialised per file -# -------------------------------------------------------------------------------------- - - -def test_h5_writes_are_serialised_per_file(tmp_path) -> None: - dataset = Dataset(str(tmp_path / "Volumes"), "h5") - attrs = _image_attributes([0.0, 0.0], [1.0, 1.0]) - dataset.write("CT", "CASE_000", np.zeros((1, 2, 2), dtype=np.float32), attrs) - - lock = _get_h5_file_lock(str(tmp_path / "Volumes") + ".h5") - started = threading.Event() - finished = threading.Event() - - def writer() -> None: - started.set() - dataset.write("CT", "CASE_001", np.ones((1, 2, 2), dtype=np.float32), attrs) - finished.set() - - with lock: # holding the file lock must block any other writer on the same file - thread = threading.Thread(target=writer) - thread.start() - assert started.wait(1.0) - assert not finished.wait(0.2), "a second writer proceeded while the file lock was held" - - thread.join(5.0) - assert finished.is_set() - data, _ = dataset.read_data("CT", "CASE_001") - np.testing.assert_array_equal(data, np.ones((1, 2, 2), dtype=np.float32)) diff --git a/tests/unit/test_dataset_audit.py b/tests/unit/test_dataset.py similarity index 59% rename from tests/unit/test_dataset_audit.py rename to tests/unit/test_dataset.py index 3ae3823d..c34e4a04 100644 --- a/tests/unit/test_dataset_audit.py +++ b/tests/unit/test_dataset.py @@ -14,32 +14,23 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for the dataset-file audit fixes in ``konfai.utils.dataset``.""" +"""Tests for ``konfai.utils.dataset``: the ``Attribute`` sidecar, the SITK/HDF5 storage +backends (modes, locking, transforms, path resolution), and ``get_infos`` shape order.""" import os +import stat +import threading from pathlib import Path -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - import numpy as np import pytest import torch -from konfai.utils.dataset import Attribute, Dataset +from konfai.utils.dataset import Attribute, Dataset, _get_h5_file_lock, get_infos, image_to_data from konfai.utils.errors import DatasetManagerError sitk = pytest.importorskip("SimpleITK") h5py = pytest.importorskip("h5py") - -def _image_attributes(origin: list[float], spacing: list[float]) -> Attribute: - attributes = Attribute() - attributes["Origin"] = np.asarray(origin, dtype=np.float64) - attributes["Spacing"] = np.asarray(spacing, dtype=np.float64) - attributes["Direction"] = np.eye(len(origin), dtype=np.float64).reshape(-1) - return attributes - - # -------------------------------------------------------------------------------------- # B13 - Attribute keys containing '_' are stored raw and must be readable/poppable # -------------------------------------------------------------------------------------- @@ -82,20 +73,69 @@ def test_attribute_repeated_set_returns_latest_version() -> None: # -------------------------------------------------------------------------------------- -# B19 - HDF5 parent directory is created with pathlib (nested paths, OS separators) +# HDF5 backend — directories, modes, and per-file locking # -------------------------------------------------------------------------------------- -def test_h5_dataset_creates_nested_parent_directories(tmp_path: Path) -> None: +def test_h5_dataset_creates_nested_parent_directories(tmp_path: Path, image_attributes) -> None: + # B19 - the parent directory is created with pathlib (nested paths, OS separators). dataset = Dataset(tmp_path / "runs" / "exp" / "Volumes", "h5") volume = np.arange(1 * 2 * 2, dtype=np.float32).reshape(1, 2, 2) - dataset.write("CT", "CASE_000", volume, _image_attributes([0.0, 0.0], [1.0, 1.0])) + dataset.write("CT", "CASE_000", volume, image_attributes([0.0, 0.0], [1.0, 1.0])) assert (tmp_path / "runs" / "exp" / "Volumes.h5").exists() data, _ = dataset.read_data("CT", "CASE_000") np.testing.assert_array_equal(data, volume) +def test_read_data_opens_hdf5_read_only(tmp_path: Path, image_attributes) -> None: + # read_data used to open HDF5 in r+ (stamping a Date attribute on every read), which mutates + # the file and breaks concurrent access across DataLoader/DDP processes. On a read-only file an + # r+ open raises PermissionError, so a successful read here proves the mode is now "r". + volume = np.arange(1 * 3 * 4 * 5, dtype=np.int16).reshape(1, 3, 4, 5) + dataset = Dataset(tmp_path / "H5DS", "h5") + dataset.write("CT", "CASE_001", volume, image_attributes([10.0, 20.0, 30.0], [0.5, 1.5, 2.0])) + + h5_files = list(tmp_path.rglob("*.h5")) + assert h5_files, "the write did not create an .h5 file" + for h5_file in h5_files: + os.chmod(h5_file, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH) + + try: + full, _ = dataset.read_data("CT", "CASE_001") + np.testing.assert_array_equal(full, volume) + finally: + for h5_file in h5_files: + os.chmod(h5_file, stat.S_IRUSR | stat.S_IWUSR) + + +def test_h5_writes_are_serialised_per_file(tmp_path: Path, image_attributes) -> None: + # B6 - concurrent HDF5 access is serialised per file. + dataset = Dataset(str(tmp_path / "Volumes"), "h5") + attrs = image_attributes([0.0, 0.0], [1.0, 1.0]) + dataset.write("CT", "CASE_000", np.zeros((1, 2, 2), dtype=np.float32), attrs) + + lock = _get_h5_file_lock(str(tmp_path / "Volumes") + ".h5") + started = threading.Event() + finished = threading.Event() + + def writer() -> None: + started.set() + dataset.write("CT", "CASE_001", np.ones((1, 2, 2), dtype=np.float32), attrs) + finished.set() + + with lock: # holding the file lock must block any other writer on the same file + thread = threading.Thread(target=writer) + thread.start() + assert started.wait(1.0) + assert not finished.wait(0.2), "a second writer proceeded while the file lock was held" + + thread.join(5.0) + assert finished.is_set() + data, _ = dataset.read_data("CT", "CASE_001") + np.testing.assert_array_equal(data, np.ones((1, 2, 2), dtype=np.float32)) + + # -------------------------------------------------------------------------------------- # B23 - a missing sitk entry raises a clear error instead of UnboundLocalError # -------------------------------------------------------------------------------------- @@ -180,14 +220,14 @@ def test_xml_file_to_data_returns_tuple_with_parsed_values(tmp_path: Path) -> No # -------------------------------------------------------------------------------------- -def test_resolve_data_path_prefers_special_format_like_full_read(tmp_path: Path) -> None: +def test_resolve_data_path_prefers_special_format_like_full_read(tmp_path: Path, image_attributes) -> None: root = tmp_path / "Dataset" dataset = Dataset(root, "mha") dataset.write( "Transf", "CASE_000", np.arange(1 * 2 * 3 * 4, dtype=np.float32).reshape(1, 2, 3, 4), - _image_attributes([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]), + image_attributes([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]), ) euler = sitk.Euler3DTransform() euler.SetParameters((0.1, 0.2, 0.3, 4.0, 5.0, 6.0)) @@ -201,3 +241,60 @@ def test_resolve_data_path_prefers_special_format_like_full_read(tmp_path: Path) assert resolved is not None and resolved.endswith(".itk.txt") full, _ = dataset.read_data("Transf", "CASE_000") assert full.shape == (1, 6) + + +# -------------------------------------------------------------------------------------- +# get_infos returns numpy channel-first order for every rank +# +# Patch planning strips the channel from get_infos' shape and feeds the spatial shape to +# transform_shape and the patch reader; the actual pixel reads (image_to_data / +# _file_to_image_slice) are numpy-order [C, (T), (Z), Y, X]. The pre-fix code reversed sitk +# GetSize() only when len == 3, so 2-D and 4-D images kept sitk (x, y, ...) order and were +# transposed against their own pixel data. +# -------------------------------------------------------------------------------------- + + +def test_get_infos_2d_matches_pixel_data(tmp_path: Path) -> None: + # Non-square 2-D: sitk GetSize() = (x=10, y=4); numpy pixel data is (y=4, x=10). + path = tmp_path / "img2d.nii.gz" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((4, 10), dtype=np.float32)), str(path)) + + size, _ = get_infos(path) + data, _ = image_to_data(sitk.ReadImage(str(path))) + + assert list(size) == list(data.shape) # [1, 4, 10], not [1, 10, 4] + + +def test_get_infos_4d_matches_pixel_data(tmp_path: Path) -> None: + # Genuine 4-D scalar: sitk GetSize() = (5, 4, 3, 2); numpy pixel data is (2, 3, 4, 5). + path = tmp_path / "img4d.nii.gz" + sitk.WriteImage(sitk.Image([5, 4, 3, 2], sitk.sitkFloat32), str(path)) + + size, _ = get_infos(path) + data = sitk.GetArrayFromImage(sitk.ReadImage(str(path))) + + assert list(size) == [1, *data.shape] # [1, 2, 3, 4, 5] + + +def test_get_infos_3d_unchanged(tmp_path: Path) -> None: + # Regression guard: the already-correct 3-D path must stay reversed. + path = tmp_path / "img3d.nii.gz" + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((6, 4, 10), dtype=np.float32)), str(path)) + + size, _ = get_infos(path) + data, _ = image_to_data(sitk.ReadImage(str(path))) + + assert list(size) == list(data.shape) == [1, 6, 4, 10] + + +def test_sitkfile_get_infos_2d_matches_read_data(tmp_path: Path) -> None: + # Same defect in SitkFile.get_infos, reached through the public Dataset API. + ds_dir = str(tmp_path / "ds") + "/" + Path(ds_dir).mkdir(parents=True, exist_ok=True) + sitk.WriteImage(sitk.GetImageFromArray(np.zeros((4, 10), dtype=np.float32)), ds_dir + "case0.mha") + + file = Dataset.SitkFile(ds_dir, read=True, file_format="mha") + size, _ = file.get_infos("", "case0") + data, _ = file.file_to_data("", "case0") + + assert list(size) == list(data.shape) # [1, 4, 10] diff --git a/tests/unit/test_ddp_shard_balance.py b/tests/unit/test_ddp_shard_balance.py deleted file mode 100644 index b3b8ac5a..00000000 --- a/tests/unit/test_ddp_shard_balance.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 test: TRAIN/RESUME shards must be equal length to avoid a DDP hang.""" - -import pytest -from konfai.data.data_manager import Data -from konfai.utils.runtime import State - - -@pytest.mark.parametrize("state", [State.TRAIN, State.RESUME]) -def test_train_split_equalises_indivisible_shards(monkeypatch: pytest.MonkeyPatch, state: State) -> None: - # DDP(static_graph=True) needs every rank to run the same number of backward all-reduces per - # epoch. A contiguous split of 7 patches over 3 ranks gives [2, 2, 3] and hangs NCCL on the - # extra step; drop_last equalises to [2, 2, 2]. - monkeypatch.setenv("KONFAI_STATE", str(state)) - - shards = Data._split([(index, 0, 0) for index in range(7)], 3) - - assert [len(shard) for shard in shards] == [2, 2, 2] - flattened = [item for shard in shards for item in shard] - assert len(flattened) == len(set(flattened)) # never duplicated across ranks - - -def test_train_split_two_ranks_indivisible(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("KONFAI_STATE", str(State.TRAIN)) - assert [len(shard) for shard in Data._split([(i, 0, 0) for i in range(5)], 2)] == [2, 2] - - -def test_train_split_single_process_keeps_everything(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("KONFAI_STATE", str(State.TRAIN)) - mapping = [(i, 0, 0) for i in range(5)] - assert Data._split(mapping, 1) == [mapping] # world_size == 1 is a no-op diff --git a/tests/unit/test_dict_metric_logging.py b/tests/unit/test_dict_metric_logging.py deleted file mode 100644 index 18a2e0d3..00000000 --- a/tests/unit/test_dict_metric_logging.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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 test: dict-payload metrics (Dice, TRE) must not crash the logging windows.""" - -import numpy as np -import pytest -import torch -from konfai.network.network import Measure - - -def test_loss_add_summarises_dict_metric_payload() -> None: - # Dice/TRE return (tensor, {label: value}); the pre-fix code stored the dict in _values, so the - # np.nanmean over _values in get_last_values/format_loss raised TypeError on every batch. - record = Measure.Loss("Dice", "out", "tgt", 0, is_loss=False, accumulation=False) - - record.add(1.0, (torch.tensor([0.7]), {"1": 0.6, "2": 0.8, "3": float("nan")})) - - # The dict is summarised to a scalar (nan-mean of 0.6 and 0.8), and the logging mean is safe. - assert isinstance(record._values[-1], float) - assert record._values[-1] == pytest.approx(0.7) - assert np.nanmean(record._values) == pytest.approx(0.7) - - -def test_loss_add_keeps_plain_scalar_metric() -> None: - # A regular (tensor, float) metric is unchanged. - record = Measure.Loss("MSE", "out", "tgt", 0, is_loss=False, accumulation=False) - record.add(1.0, (torch.tensor([0.5]), 0.5)) - assert record._values[-1] == pytest.approx(0.5) diff --git a/tests/unit/test_evaluator_statistics.py b/tests/unit/test_evaluator_statistics.py index 412cf8ea..1c4add46 100644 --- a/tests/unit/test_evaluator_statistics.py +++ b/tests/unit/test_evaluator_statistics.py @@ -17,15 +17,10 @@ """Aggregation, serialization and summary behaviour of :class:`Statistics`.""" import json -import os -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import numpy as np # noqa: E402 -import pytest # noqa: E402 - -from konfai.evaluator import Statistics # noqa: E402 +import numpy as np +import pytest +from konfai.evaluator import Statistics class TestGetStatisticCount: diff --git a/tests/unit/test_experimental_models.py b/tests/unit/test_experimental_models.py deleted file mode 100644 index 32941517..00000000 --- a/tests/unit/test_experimental_models.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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 - -"""Guard tests: non-functional research models fail fast with a clear message.""" - -import pytest -from konfai.models.generation.ddpm import DDPM -from konfai.models.registration.registration import VoxelMorph - - -def test_ddpm_is_marked_experimental() -> None: - # DDPM cannot execute a forward pass (broken time-embedding wiring); constructing it must raise - # an actionable error instead of crashing opaquely deep in the graph later. - with pytest.raises(NotImplementedError, match="experimental"): - DDPM() - - -def test_voxelmorph_rejects_3d_configuration() -> None: - # VoxelMorph's warping components are 2-D-hardcoded, so its own dim=3 default used to crash. - with pytest.raises(NotImplementedError, match="dim=2"): - VoxelMorph(dim=3) diff --git a/tests/unit/test_get_infos_shape_order.py b/tests/unit/test_get_infos_shape_order.py deleted file mode 100644 index 92b93ee1..00000000 --- a/tests/unit/test_get_infos_shape_order.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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 test: get_infos returns numpy channel-first order for every rank. - -Patch planning strips the channel from get_infos' shape and feeds the spatial shape to -transform_shape and the patch reader; the actual pixel reads (image_to_data / _file_to_image_slice) -are numpy-order [C, (T), (Z), Y, X]. The pre-fix code reversed sitk GetSize() only when len == 3, so -2-D and 4-D images kept sitk (x, y, ...) order and were transposed against their own pixel data. -""" - -from pathlib import Path - -import numpy as np -import SimpleITK as sitk -from konfai.utils.dataset import Dataset, get_infos, image_to_data - - -def test_get_infos_2d_matches_pixel_data(tmp_path: Path) -> None: - # Non-square 2-D: sitk GetSize() = (x=10, y=4); numpy pixel data is (y=4, x=10). - path = tmp_path / "img2d.nii.gz" - sitk.WriteImage(sitk.GetImageFromArray(np.zeros((4, 10), dtype=np.float32)), str(path)) - - size, _ = get_infos(path) - data, _ = image_to_data(sitk.ReadImage(str(path))) - - assert list(size) == list(data.shape) # [1, 4, 10], not [1, 10, 4] - - -def test_get_infos_4d_matches_pixel_data(tmp_path: Path) -> None: - # Genuine 4-D scalar: sitk GetSize() = (5, 4, 3, 2); numpy pixel data is (2, 3, 4, 5). - path = tmp_path / "img4d.nii.gz" - sitk.WriteImage(sitk.Image([5, 4, 3, 2], sitk.sitkFloat32), str(path)) - - size, _ = get_infos(path) - data = sitk.GetArrayFromImage(sitk.ReadImage(str(path))) - - assert list(size) == [1, *data.shape] # [1, 2, 3, 4, 5] - - -def test_get_infos_3d_unchanged(tmp_path: Path) -> None: - # Regression guard: the already-correct 3-D path must stay reversed. - path = tmp_path / "img3d.nii.gz" - sitk.WriteImage(sitk.GetImageFromArray(np.zeros((6, 4, 10), dtype=np.float32)), str(path)) - - size, _ = get_infos(path) - data, _ = image_to_data(sitk.ReadImage(str(path))) - - assert list(size) == list(data.shape) == [1, 6, 4, 10] - - -def test_sitkfile_get_infos_2d_matches_read_data(tmp_path: Path) -> None: - # Same defect in SitkFile.get_infos, reached through the public Dataset API. - ds_dir = str(tmp_path / "ds") + "/" - Path(ds_dir).mkdir(parents=True, exist_ok=True) - sitk.WriteImage(sitk.GetImageFromArray(np.zeros((4, 10), dtype=np.float32)), ds_dir + "case0.mha") - - file = Dataset.SitkFile(ds_dir, read=True, file_format="mha") - size, _ = file.get_infos("", "case0") - data, _ = file.file_to_data("", "case0") - - assert list(size) == list(data.shape) # [1, 4, 10] diff --git a/tests/unit/test_hdf5_read_only.py b/tests/unit/test_hdf5_read_only.py deleted file mode 100644 index 54c2d494..00000000 --- a/tests/unit/test_hdf5_read_only.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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 test: read_data must open HDF5 read-only, not in r+ (write) mode.""" - -import os -import stat -from pathlib import Path - -import numpy as np -import pytest - -pytest.importorskip("h5py") - -from konfai.utils.dataset import Attribute, Dataset # noqa: E402 - - -def _image_attributes() -> Attribute: - attributes = Attribute() - attributes["Origin"] = np.asarray([10.0, 20.0, 30.0]) - attributes["Spacing"] = np.asarray([0.5, 1.5, 2.0]) - attributes["Direction"] = np.eye(3, dtype=np.float64).flatten() - return attributes - - -def test_read_data_opens_hdf5_read_only(tmp_path: Path) -> None: - # read_data used to open HDF5 in r+ (stamping a Date attribute on every read), which mutates - # the file and breaks concurrent access across DataLoader/DDP processes. On a read-only file an - # r+ open raises PermissionError, so a successful read here proves the mode is now "r". - volume = np.arange(1 * 3 * 4 * 5, dtype=np.int16).reshape(1, 3, 4, 5) - dataset = Dataset(tmp_path / "H5DS", "h5") - dataset.write("CT", "CASE_001", volume, _image_attributes()) - - h5_files = list(tmp_path.rglob("*.h5")) - assert h5_files, "the write did not create an .h5 file" - for h5_file in h5_files: - os.chmod(h5_file, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH) - - try: - full, _ = dataset.read_data("CT", "CASE_001") - np.testing.assert_array_equal(full, volume) - finally: - for h5_file in h5_files: - os.chmod(h5_file, stat.S_IRUSR | stat.S_IWUSR) diff --git a/tests/unit/test_inline_augmentations.py b/tests/unit/test_inline_augmentations.py deleted file mode 100644 index d71c09c3..00000000 --- a/tests/unit/test_inline_augmentations.py +++ /dev/null @@ -1,174 +0,0 @@ -from typing import cast -from pathlib import Path - -import numpy as np -import pytest -import torch - -import konfai.data.augmentation as augmentation_module -from konfai.data.augmentation import DataAugmentation, DataAugmentationsList, Elastix, Mask -from konfai.data.data_manager import DatasetIter, Group, GroupTransform -from konfai.data.patching import DatasetManager -from konfai.utils.dataset import Attribute, Dataset -from konfai.utils.errors import AugmentationError - - -class DummyDataset: - def __init__(self, array: np.ndarray) -> None: - self.array = array - - def get_infos(self, group_src: str, name: str) -> tuple[list[int], Attribute]: - return list(self.array.shape), Attribute({"name": name, "group": group_src}) - - def read_data(self, group_src: str, name: str) -> tuple[np.ndarray, Attribute]: - return self.array.copy(), Attribute({"name": name, "group": group_src}) - - -class CountingOffsetAugmentation(DataAugmentation): - def __init__(self) -> None: - super().__init__() - self.compute_calls = 0 - - def _state_init( - self, - index: int, - shapes: list[list[int]], - caches_attribute: list[Attribute], - ) -> list[list[int]]: - return shapes - - def _compute( - self, - name: str, - index: int, - tensors: list[torch.Tensor], - ) -> list[torch.Tensor]: - self.compute_calls += 1 - return [tensor + (offset + 1) for offset, tensor in enumerate(tensors)] - - def _inverse(self, index: int, a: int, tensor: torch.Tensor) -> torch.Tensor: - return tensor - - -def test_inline_augmentations_are_loaded_on_demand() -> None: - base = np.arange(4, dtype=np.float32).reshape(1, 2, 2) - dataset = cast(Dataset, DummyDataset(base)) - augmentation = CountingOffsetAugmentation() - augmentation.load(1.0) - - augmentations = DataAugmentationsList(nb=2, data_augmentations={}) - augmentations.data_augmentations = [augmentation] - - manager = DatasetManager( - index=0, - group_src="src", - group_dest="dest", - name="case_000", - dataset=dataset, - patch=None, - transforms=[], - data_augmentations_list=[augmentations], - ) - dataset_iter = DatasetIter( - rank=0, - data={"dest": [manager]}, - mapping=[(0, 0, 0), (0, 1, 0), (0, 2, 0)], - groups_src={"src": Group(groups_dest={"dest": GroupTransform(transforms=None, patch_transforms=None)})}, - inline_augmentations=True, - data_augmentations_list=[augmentations], - patch_size=None, - overlap=None, - buffer_size=1, - use_cache=True, - ) - - base_sample = dataset_iter[0]["dest"].tensor - assert augmentation.compute_calls == 0 - assert manager.loaded is True - assert manager.augmentationLoaded is False - assert torch.equal(base_sample, torch.from_numpy(base)) - - first_augmented_sample = dataset_iter[1]["dest"].tensor - assert augmentation.compute_calls == 1 - assert manager.augmentationLoaded is True - assert torch.equal(first_augmented_sample, torch.from_numpy(base) + 1) - - second_augmented_sample = dataset_iter[2]["dest"].tensor - assert augmentation.compute_calls == 1 - assert torch.equal(second_augmented_sample, torch.from_numpy(base) + 2) - - -def test_dataset_iter_can_skip_augmentation_loading_when_validation_disables_them() -> None: - base = np.arange(4, dtype=np.float32).reshape(1, 2, 2) - dataset = cast(Dataset, DummyDataset(base)) - augmentation = CountingOffsetAugmentation() - augmentation.load(1.0) - - augmentations = DataAugmentationsList(nb=2, data_augmentations={}) - augmentations.data_augmentations = [augmentation] - - manager = DatasetManager( - index=0, - group_src="src", - group_dest="dest", - name="case_000", - dataset=dataset, - patch=None, - transforms=[], - data_augmentations_list=[augmentations], - ) - dataset_iter = DatasetIter( - rank=0, - data={"dest": [manager]}, - mapping=[(0, 0, 0)], - groups_src={"src": Group(groups_dest={"dest": GroupTransform(transforms=None, patch_transforms=None)})}, - inline_augmentations=False, - data_augmentations_list=[augmentations], - patch_size=None, - overlap=None, - buffer_size=1, - apply_augmentations=False, - use_cache=True, - ) - - dataset_iter.load("Validation") - base_sample = dataset_iter[0]["dest"].tensor - - assert augmentation.compute_calls == 0 - assert manager.loaded is True - assert manager.augmentationLoaded is False - assert torch.equal(base_sample, torch.from_numpy(base)) - - -def test_simpleitk_augmentations_fail_clearly_when_dependency_is_missing( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(augmentation_module, "sitk", None) - - with pytest.raises(AugmentationError, match="SimpleITK"): - Elastix() - with pytest.raises(AugmentationError, match="SimpleITK"): - Mask("mask.mha", 0) - - -def test_mask_reads_pixels_only_on_first_compute(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - sitk = pytest.importorskip("SimpleITK") - mask_path = tmp_path / "mask.mha" - sitk.WriteImage(sitk.GetImageFromArray(np.ones((2, 2), dtype=np.uint8)), str(mask_path)) - - read_count = 0 - original_read_image = sitk.ReadImage - - def counting_read_image(path: str): - nonlocal read_count - read_count += 1 - return original_read_image(path) - - monkeypatch.setattr(augmentation_module.sitk, "ReadImage", counting_read_image) - augmentation = Mask(str(mask_path), 0) - augmentation._state_init(0, [[2, 2]], [Attribute()]) - - assert read_count == 0 - augmentation._compute("case", 0, [torch.ones((1, 2, 2))]) - augmentation._compute("case", 0, [torch.ones((1, 2, 2))]) - assert read_count == 1 diff --git a/tests/unit/test_loss_weight_scheduling.py b/tests/unit/test_loss_weight_scheduling.py deleted file mode 100644 index 582fd2ef..00000000 --- a/tests/unit/test_loss_weight_scheduling.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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 test: the gradient loss uses the current weight, not a stale first one.""" - -import torch -from konfai.network.network import Measure - - -def _loss_record() -> Measure.Loss: - return Measure.Loss("l", "out", "tgt", 0, is_loss=True, accumulation=False) - - -def test_get_loss_uses_current_iteration_weight() -> None: - # reset_loss clears _loss every iteration but _weight keeps growing for the logging windows. - # get_loss must pair the current loss with the current weight; the pre-fix code zipped from the - # front, so a loss-weight scheduler that changes the weight had no effect on the gradient. - record = _loss_record() - - record.reset_loss() - record.add(2.0, torch.tensor([3.0])) - assert record.get_loss().item() == 6.0 # 2 * 3 - - record.reset_loss() # next iteration; _weight is now [2.0, 5.0] - record.add(5.0, torch.tensor([1.0])) - assert record.get_loss().item() == 5.0 # 5 * 1, not the stale 2 * 1 - - -def test_get_loss_handles_multiple_accumulated_patches() -> None: - # Accumulation mode adds several (weight, loss) pairs per iteration; the trailing weights must - # still line up one-to-one with the current losses. - record = _loss_record() - - record.reset_loss() - record.add(1.0, torch.tensor([10.0])) # a previous iteration leaves a weight behind - record.reset_loss() - record.add(0.5, torch.tensor([2.0])) - record.add(0.5, torch.tensor([4.0])) - - assert record.get_loss().item() == 1.5 # mean(0.5 * 2, 0.5 * 4) diff --git a/tests/unit/test_main_cli.py b/tests/unit/test_main_cli.py index 4c509f31..23c143d4 100644 --- a/tests/unit/test_main_cli.py +++ b/tests/unit/test_main_cli.py @@ -1,10 +1,55 @@ +# 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 + +"""Tests for the ``konfai`` CLI (``konfai.main``): subcommand dispatch and the +CLI-facing parameter contract of the backend entry points.""" + +import inspect import sys from pathlib import Path import pytest +import konfai.evaluator as evaluator_module import konfai.main as main_module import konfai.predictor as predictor_module +import konfai.trainer as trainer_module + + +def test_konfai_help_exits_zero(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(sys, "argv", ["konfai", "--help"]) + + with pytest.raises(SystemExit) as exc_info: + main_module.main() + + assert exc_info.value.code == 0 + + +def test_konfai_train_dispatches_correctly(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, object] = {} + + def fake_train(**kwargs) -> None: + captured.update(kwargs) + + monkeypatch.setattr(trainer_module, "train", fake_train) + monkeypatch.setattr(sys, "argv", ["konfai", "TRAIN", "-c", "Config.yml"]) + + main_module.main() + + assert captured["config"] == "Config.yml" def test_main_prediction_dispatches_config_as_prediction_file( @@ -39,3 +84,25 @@ def fake_predict(**kwargs) -> None: assert captured["gpu"] == [] assert captured["models"] == [str(tmp_path / "checkpoint.pt")] assert "config" not in captured + + +def test_konfai_eval_dispatches_correctly(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, object] = {} + + def fake_evaluate(**kwargs) -> None: + captured.update(kwargs) + + monkeypatch.setattr(evaluator_module, "evaluate", fake_evaluate) + monkeypatch.setattr(sys, "argv", ["konfai", "EVALUATION", "-c", "Evaluation.yml"]) + + main_module.main() + + assert captured["evaluations_file"] == "Evaluation.yml" + + +def test_predict_evaluate_expose_tensorboard_param(): + """#7 CLI -tb/--tensorboard (dest 'tensorboard') must reach predict()/evaluate().""" + for fn in (predictor_module.predict, evaluator_module.evaluate): + params = inspect.signature(fn).parameters + assert "tensorboard" in params, f"{fn.__name__} must accept 'tensorboard'" + assert "tb" not in params, f"{fn.__name__} must not use the old 'tb' name" diff --git a/tests/unit/test_measure.py b/tests/unit/test_measure.py index 6d500166..5af76395 100644 --- a/tests/unit/test_measure.py +++ b/tests/unit/test_measure.py @@ -14,18 +14,14 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Numerical behaviour tests for the Dice and SSIM criteria.""" +"""Tests for the criteria in ``konfai.metric.measure`` (Dice, SSIM, Variance, +PerceptualLoss plumbing, and optional-dependency errors).""" -import os - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import numpy as np # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 - -from konfai.metric.measure import SSIM, Dice, Variance # noqa: E402 +import numpy as np +import pytest +import torch +from konfai.metric.measure import SSIM, Dice, PerceptualLoss, Variance, _require_optional +from konfai.utils.errors import MeasureError def _one_hot(target: torch.Tensor, nb_channels: int) -> torch.Tensor: @@ -186,3 +182,34 @@ def test_multi_channel_uses_unbiased_variance(self): # Unbiased var of [1, 3] = ((1-2)^2 + (3-2)^2) / (2 - 1) = 2.0. assert variance.item() == pytest.approx(2.0) assert value == pytest.approx(2.0) + + +def test_perceptual_loss_forward_unpacks_targets() -> None: + # forward(output, *targets) must hand each target to _compute(output, *targets) as its own + # positional tensor; the pre-fix code passed the whole tuple as a single argument, so the + # preprocessing/feature-extraction path received a tuple and crashed. + loss = object.__new__(PerceptualLoss) + loss.shape = [128, 128, 128] # len != 2 -> the non-slice branch is taken + loss.models = {None: object()} # short-circuit the lazy model placement on device index None + + recorded: dict[str, tuple] = {} + + def fake_compute(output, *targets): + recorded["targets"] = targets + return torch.zeros(1) + + loss._compute = fake_compute # type: ignore[method-assign] + + PerceptualLoss.forward(loss, torch.randn(1, 1, 8, 8), torch.randn(1, 1, 8, 8)) + + assert len(recorded["targets"]) == 1 + assert torch.is_tensor(recorded["targets"][0]) + + +def test_missing_metric_dependency_raises_actionable_error(): + """Optional criterion deps must surface an actionable MeasureError, not ImportError.""" + with pytest.raises(MeasureError) as excinfo: + _require_optional("konfai_definitely_missing_pkg_zzz", criterion="SSIM", extra="ssim") + message = str(excinfo.value) + assert "SSIM" in message + assert "konfai[ssim]" in message diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 00000000..fd1e23f9 --- /dev/null +++ b/tests/unit/test_models.py @@ -0,0 +1,105 @@ +# 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 + +"""Tests for the built-in model definitions in ``konfai.models``.""" + +import pytest +import torch +from konfai.models.generation.ddpm import DDPM +from konfai.models.generation.diffusionGan import CycleGanDiscriminator +from konfai.models.generation.vae import LinearVAE +from konfai.models.registration.registration import VoxelMorph +from konfai.models.representation.representation import Adaptation +from konfai.models.segmentation.UNet import UNet + +# -------------------------------------------------------------------------------------- +# UNet +# -------------------------------------------------------------------------------------- + + +def _unet_forward_last(attention: bool) -> torch.Tensor: + net = UNet(dim=2, channels=[1, 8, 16], nb_class=2, attention=attention) + outputs = list(net.named_forward(torch.randn(1, 1, 32, 32))) + return outputs[-1][1] + + +def test_unet_attention_forwards_without_branch_collision() -> None: + # out_branch=[1] collided with the Attention block's internal W_g branch, so the parent captured + # a half-resolution projection and Concat crashed on a size mismatch. The gated skip must reach + # the skip connection, so the network forwards to a full-resolution output. + attended = _unet_forward_last(attention=True) + plain = _unet_forward_last(attention=False) + + assert attended.shape[-2:] == (32, 32) + assert attended.shape == plain.shape + + +# -------------------------------------------------------------------------------------- +# Generation models +# -------------------------------------------------------------------------------------- + + +def test_linear_vae_is_parameterized_and_variational(): + """#17 LinearVAE must be parameterized (no hardcoded dims) and sample a latent.""" + model = LinearVAE(in_features=32, hidden_features=16, latent_dim=4) + x = torch.randn(2, 32) + outputs = dict(model.named_forward(x)) + assert outputs["Head.Tanh"].shape == (2, 32) # reconstruction matches input size + assert "Latent.mu" in outputs and "Latent.log_std" in outputs # KL-ready outputs + # The latent is sampled: the reconstruction differs across RNG draws. + torch.manual_seed(0) + first = dict(model.named_forward(x))["Head.Tanh"] + torch.manual_seed(1) + second = dict(model.named_forward(x))["Head.Tanh"] + assert not torch.allclose(first, second) + + +def test_cyclegan_discriminator_initialized_no_keyerror(): + """#CycleGan: initialized() must not index a missing 'Sample' submodule on load.""" + model = CycleGanDiscriminator() + # Must not raise KeyError('Sample'). + model.initialized() + + +# -------------------------------------------------------------------------------------- +# Representation models +# -------------------------------------------------------------------------------------- + + +def test_adaptation_sets_requires_grad_at_construction(): + """#18 Adaptation must configure requires_grad in __init__, not on every forward.""" + adaptation = Adaptation() + # State is correct immediately after construction, before any forward pass. + assert all(not p.requires_grad for p in adaptation.Encoder_1.parameters()) + assert all(p.requires_grad for p in adaptation.FCT_1.parameters()) + + +# -------------------------------------------------------------------------------------- +# Experimental models fail fast with a clear message +# -------------------------------------------------------------------------------------- + + +def test_ddpm_is_marked_experimental() -> None: + # DDPM cannot execute a forward pass (broken time-embedding wiring); constructing it must raise + # an actionable error instead of crashing opaquely deep in the graph later. + with pytest.raises(NotImplementedError, match="experimental"): + DDPM() + + +def test_voxelmorph_rejects_3d_configuration() -> None: + # VoxelMorph's warping components are 2-D-hardcoded, so its own dim=3 default used to crash. + with pytest.raises(NotImplementedError, match="dim=2"): + VoxelMorph(dim=3) diff --git a/tests/unit/test_named_forward.py b/tests/unit/test_named_forward.py deleted file mode 100644 index dbcef857..00000000 --- a/tests/unit/test_named_forward.py +++ /dev/null @@ -1,104 +0,0 @@ -# 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 - -"""Unit tests for ModuleArgsDict branch routing (named_forward / forward).""" - -import os - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import torch # noqa: E402 - -from konfai.network.blocks import Add # noqa: E402 -from konfai.network.network import ModuleArgsDict # noqa: E402 - - -class _MulConst(torch.nn.Module): - """Deterministic test module: multiplies its input by a fixed constant.""" - - def __init__(self, factor: float) -> None: - super().__init__() - self.factor = factor - - def forward(self, tensor: torch.Tensor) -> torch.Tensor: - return tensor * self.factor - - -class _TwoInputGraph(ModuleArgsDict): - """A(in 0)→branch 0, B(in 1)→branch 1, Sum(in 0,1)→branch 2.""" - - def __init__(self) -> None: - super().__init__() - self.add_module("A", _MulConst(3.0), in_branch=[0], out_branch=[0]) - self.add_module("B", _MulConst(10.0), in_branch=[1], out_branch=[1]) - self.add_module("Sum", Add(), in_branch=[0, 1], out_branch=[2]) - - -class _Inner(ModuleArgsDict): - def __init__(self) -> None: - super().__init__() - self.add_module("Scale", _MulConst(2.0)) - - -class _NestedGraph(ModuleArgsDict): - def __init__(self) -> None: - super().__init__() - self.add_module("Pre", _MulConst(5.0), in_branch=[0], out_branch=[0]) - self.add_module("Block", _Inner(), in_branch=[0], out_branch=[0]) - - -def test_forward_routes_two_inputs_through_branches(): - graph = _TwoInputGraph() - a = torch.ones(1, 1, 2, 2) - b = torch.full((1, 1, 2, 2), 2.0) - out = graph(a, b) # 3*a + 10*b = 3 + 20 = 23 - assert torch.allclose(out, torch.full_like(out, 23.0)) - - -def test_named_forward_exposes_every_intermediate(): - graph = _TwoInputGraph() - a = torch.ones(1, 1, 2, 2) - b = torch.full((1, 1, 2, 2), 2.0) - outputs = {name: float(tensor.flatten()[0]) for name, tensor in graph.named_forward(a, b)} - assert outputs == {"A": 3.0, "B": 20.0, "Sum": 23.0} - - -def test_named_forward_uses_dotted_names_for_nested_graphs(): - graph = _NestedGraph() - x = torch.ones(1, 1, 2, 2) - names = [name for name, _ in graph.named_forward(x)] - assert "Pre" in names - assert "Block.Scale" in names # nested submodule addressable by dotted path - out = graph(x) # 5 then *2 = 10 - assert torch.allclose(out, torch.full_like(out, 10.0)) - - -def test_out_branch_isolation_preserves_a_branch_for_later_use(): - """A branch written by one module must remain available to a later consumer.""" - - class _SkipGraph(ModuleArgsDict): - def __init__(self) -> None: - super().__init__() - # Keep the raw input on branch 1, transform branch 0, then combine. - self.add_module("Identity", torch.nn.Identity(), in_branch=[0], out_branch=[1]) - self.add_module("Scale", _MulConst(4.0), in_branch=[0], out_branch=[0]) - self.add_module("Sum", Add(), in_branch=[0, 1], out_branch=[0]) - - graph = _SkipGraph() - x = torch.ones(1, 1, 2, 2) - out = graph(x) # 4*x + x = 5 - assert torch.allclose(out, torch.full_like(out, 5.0)) diff --git a/tests/unit/test_named_forward_sibling.py b/tests/unit/test_named_forward_sibling.py deleted file mode 100644 index d47ca1fb..00000000 --- a/tests/unit/test_named_forward_sibling.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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 test: a nested sibling's output must not be dropped by a stale inner-match set.""" - -import torch -from konfai.network.network import ModuleArgsDict - - -class _Add(torch.nn.Module): - def __init__(self, value: float) -> None: - super().__init__() - self.value = value - - def forward(self, tensor: torch.Tensor) -> torch.Tensor: - return tensor + self.value - - -def _nested(value: float, inner_out: int | str) -> ModuleArgsDict: - sub = ModuleArgsDict() - sub.add_module("L", _Add(value), in_branch=[0], out_branch=[inner_out]) - return sub - - -def test_later_nested_sibling_output_reaches_downstream() -> None: - # M1 writes branch 0 via inner-match; M2 shares out_branch=[0] but its inner module writes a - # different branch, so it relies on the fallback. A ``tmp`` kept across siblings made the fallback - # see branch 0 as already filled (by M1) and silently drop M2, leaving M1's value downstream. - graph = ModuleArgsDict() - graph.add_module("M1", _nested(1.0, 0), in_branch=[0], out_branch=[0]) - graph.add_module("M2", _nested(10.0, "zz"), in_branch=[0], out_branch=[0]) - graph.add_module("Id", torch.nn.Identity(), in_branch=[0], out_branch=[0]) - - outputs = list(graph.named_forward(torch.zeros(1))) - downstream = [tensor for name, tensor in outputs if name.startswith("Id")][-1] - - # Branch 0: input 0 -> M1 (+1) = 1 -> M2 reads branch 0 (+10) = 11 -> Id. Not M1's stale 1. - assert downstream.item() == 11.0 diff --git a/tests/unit/test_network.py b/tests/unit/test_network.py new file mode 100644 index 00000000..fcb11501 --- /dev/null +++ b/tests/unit/test_network.py @@ -0,0 +1,366 @@ +# 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 + +"""Tests for ``konfai.network.network``: ModuleArgsDict branch routing and init, +Network.load_state_dict, Measure (loss records, backward, scheduler selection), +and CriterionsLoader.""" + +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +import konfai.network.network as network_module +from konfai.metric.schedulers import Constant +from konfai.network.blocks import Add +from konfai.network.network import Measure, ModuleArgsDict, Network +from konfai.utils.dataset import Attribute +from konfai.utils.errors import ConfigError + +# -------------------------------------------------------------------------------------- +# ModuleArgsDict branch routing (named_forward / forward) +# -------------------------------------------------------------------------------------- + + +class _MulConst(torch.nn.Module): + """Deterministic test module: multiplies its input by a fixed constant.""" + + def __init__(self, factor: float) -> None: + super().__init__() + self.factor = factor + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor * self.factor + + +class _TwoInputGraph(ModuleArgsDict): + """A(in 0)→branch 0, B(in 1)→branch 1, Sum(in 0,1)→branch 2.""" + + def __init__(self) -> None: + super().__init__() + self.add_module("A", _MulConst(3.0), in_branch=[0], out_branch=[0]) + self.add_module("B", _MulConst(10.0), in_branch=[1], out_branch=[1]) + self.add_module("Sum", Add(), in_branch=[0, 1], out_branch=[2]) + + +class _Inner(ModuleArgsDict): + def __init__(self) -> None: + super().__init__() + self.add_module("Scale", _MulConst(2.0)) + + +class _NestedGraph(ModuleArgsDict): + def __init__(self) -> None: + super().__init__() + self.add_module("Pre", _MulConst(5.0), in_branch=[0], out_branch=[0]) + self.add_module("Block", _Inner(), in_branch=[0], out_branch=[0]) + + +def test_forward_routes_two_inputs_through_branches(): + graph = _TwoInputGraph() + a = torch.ones(1, 1, 2, 2) + b = torch.full((1, 1, 2, 2), 2.0) + out = graph(a, b) # 3*a + 10*b = 3 + 20 = 23 + assert torch.allclose(out, torch.full_like(out, 23.0)) + + +def test_named_forward_exposes_every_intermediate(): + graph = _TwoInputGraph() + a = torch.ones(1, 1, 2, 2) + b = torch.full((1, 1, 2, 2), 2.0) + outputs = {name: float(tensor.flatten()[0]) for name, tensor in graph.named_forward(a, b)} + assert outputs == {"A": 3.0, "B": 20.0, "Sum": 23.0} + + +def test_named_forward_uses_dotted_names_for_nested_graphs(): + graph = _NestedGraph() + x = torch.ones(1, 1, 2, 2) + names = [name for name, _ in graph.named_forward(x)] + assert "Pre" in names + assert "Block.Scale" in names # nested submodule addressable by dotted path + out = graph(x) # 5 then *2 = 10 + assert torch.allclose(out, torch.full_like(out, 10.0)) + + +def test_out_branch_isolation_preserves_a_branch_for_later_use(): + """A branch written by one module must remain available to a later consumer.""" + + class _SkipGraph(ModuleArgsDict): + def __init__(self) -> None: + super().__init__() + # Keep the raw input on branch 1, transform branch 0, then combine. + self.add_module("Identity", torch.nn.Identity(), in_branch=[0], out_branch=[1]) + self.add_module("Scale", _MulConst(4.0), in_branch=[0], out_branch=[0]) + self.add_module("Sum", Add(), in_branch=[0, 1], out_branch=[0]) + + graph = _SkipGraph() + x = torch.ones(1, 1, 2, 2) + out = graph(x) # 4*x + x = 5 + assert torch.allclose(out, torch.full_like(out, 5.0)) + + +class _AddConst(torch.nn.Module): + def __init__(self, value: float) -> None: + super().__init__() + self.value = value + + def forward(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor + self.value + + +def _nested_adder(value: float, inner_out: int | str) -> ModuleArgsDict: + sub = ModuleArgsDict() + sub.add_module("L", _AddConst(value), in_branch=[0], out_branch=[inner_out]) + return sub + + +def test_later_nested_sibling_output_reaches_downstream() -> None: + # M1 writes branch 0 via inner-match; M2 shares out_branch=[0] but its inner module writes a + # different branch, so it relies on the fallback. A ``tmp`` kept across siblings made the fallback + # see branch 0 as already filled (by M1) and silently drop M2, leaving M1's value downstream. + graph = ModuleArgsDict() + graph.add_module("M1", _nested_adder(1.0, 0), in_branch=[0], out_branch=[0]) + graph.add_module("M2", _nested_adder(10.0, "zz"), in_branch=[0], out_branch=[0]) + graph.add_module("Id", torch.nn.Identity(), in_branch=[0], out_branch=[0]) + + outputs = list(graph.named_forward(torch.zeros(1))) + downstream = [tensor for name, tensor in outputs if name.startswith("Id")][-1] + + # Branch 0: input 0 -> M1 (+1) = 1 -> M2 reads branch 0 (+10) = 11 -> Id. Not M1's stale 1. + assert downstream.item() == 11.0 + + +def test_init_func_centres_batchnorm_gamma_on_one() -> None: + # gamma initialised around 0 scaled the normalised activations to ~0, stalling early training. + batch_norm = torch.nn.BatchNorm2d(128) + + ModuleArgsDict.init_func(batch_norm, "normal", 0.02) + + assert abs(batch_norm.weight.mean().item() - 1.0) < 0.02 + assert batch_norm.bias.abs().max().item() < 1e-6 + + +# -------------------------------------------------------------------------------------- +# Network.load_state_dict +# -------------------------------------------------------------------------------------- + + +def test_load_state_dict_warm_starts_resized_layer_and_keeps_siblings(): + """#2 A resized layer must warm-start, and sibling layers must still load. + + The bug checked ``isinstance(module, Linear)`` (the parent) instead of the + child, and used an early ``return`` that aborted loading the remaining + siblings of a resized layer. + """ + + class _Net(Network): + def __init__(self, fc_out: int) -> None: + super().__init__(in_channels=1) + self.add_module("fc", torch.nn.Linear(4, fc_out)) + self.add_module("head", torch.nn.Linear(4, 2)) + + old = _Net(fc_out=4) + # Network.state_dict() wraps the flat params under the network name; load_state_dict + # consumes that inner flat dict ("fc.weight", ...). + inner = next(iter(old.state_dict().values())) + checkpoint = {key: value.clone() for key, value in inner.items()} + + new = _Net(fc_out=6) # fc output grows 4 -> 6 (resized); head is unchanged + new.load_state_dict(checkpoint) # must not raise + + fc = new["fc"] + head = new["head"] + assert fc.weight.shape == (6, 4) + assert torch.equal(fc.weight[:4], checkpoint["fc.weight"]) # warm-started rows + # The sibling after the resized layer must still be loaded (old `return` skipped it). + assert torch.equal(head.weight, checkpoint["head.weight"]) + assert torch.equal(head.bias, checkpoint["head.bias"]) + + +# -------------------------------------------------------------------------------------- +# Measure.Loss — loss records feeding the gradient and the logging windows +# -------------------------------------------------------------------------------------- + + +def _loss_record() -> Measure.Loss: + return Measure.Loss("l", "out", "tgt", 0, is_loss=True, accumulation=False) + + +def test_get_loss_uses_current_iteration_weight() -> None: + # reset_loss clears _loss every iteration but _weight keeps growing for the logging windows. + # get_loss must pair the current loss with the current weight; the pre-fix code zipped from the + # front, so a loss-weight scheduler that changes the weight had no effect on the gradient. + record = _loss_record() + + record.reset_loss() + record.add(2.0, torch.tensor([3.0])) + assert record.get_loss().item() == 6.0 # 2 * 3 + + record.reset_loss() # next iteration; _weight is now [2.0, 5.0] + record.add(5.0, torch.tensor([1.0])) + assert record.get_loss().item() == 5.0 # 5 * 1, not the stale 2 * 1 + + +def test_get_loss_handles_multiple_accumulated_patches() -> None: + # Accumulation mode adds several (weight, loss) pairs per iteration; the trailing weights must + # still line up one-to-one with the current losses. + record = _loss_record() + + record.reset_loss() + record.add(1.0, torch.tensor([10.0])) # a previous iteration leaves a weight behind + record.reset_loss() + record.add(0.5, torch.tensor([2.0])) + record.add(0.5, torch.tensor([4.0])) + + assert record.get_loss().item() == 1.5 # mean(0.5 * 2, 0.5 * 4) + + +def test_loss_add_summarises_dict_metric_payload() -> None: + # Dice/TRE return (tensor, {label: value}); the pre-fix code stored the dict in _values, so the + # np.nanmean over _values in get_last_values/format_loss raised TypeError on every batch. + record = Measure.Loss("Dice", "out", "tgt", 0, is_loss=False, accumulation=False) + + record.add(1.0, (torch.tensor([0.7]), {"1": 0.6, "2": 0.8, "3": float("nan")})) + + # The dict is summarised to a scalar (nan-mean of 0.6 and 0.8), and the logging mean is safe. + assert isinstance(record._values[-1], float) + assert record._values[-1] == pytest.approx(0.7) + assert np.nanmean(record._values) == pytest.approx(0.7) + + +def test_loss_add_keeps_plain_scalar_metric() -> None: + # A regular (tensor, float) metric is unchanged. + record = Measure.Loss("MSE", "out", "tgt", 0, is_loss=False, accumulation=False) + record.add(1.0, (torch.tensor([0.5]), 0.5)) + assert record._values[-1] == pytest.approx(0.5) + + +# -------------------------------------------------------------------------------------- +# Measure — accumulation backward (AMP scaler vs plain) +# -------------------------------------------------------------------------------------- + + +class _CriterionAttr: + def __init__(self) -> None: + self.start = 0 + self.stop = None + self.schedulers = {Constant(1.0): 1} + self.group = 0 + self.is_loss = True + self.accumulation = True + + +def _make_accumulating_measure(scaler) -> tuple[Measure, torch.Tensor]: + """Build a minimal Measure that triggers the accumulation-backward branch.""" + measure = Measure.__new__(Measure) + criterion = torch.nn.MSELoss() + key = f"out:tgt:{criterion.__class__.__name__}" + measure.outputs_criterions = {"out": {"tgt": {criterion: _CriterionAttr()}}} + measure._loss = {0: {key: Measure.Loss(criterion.__class__.__name__, "out", "tgt", 0, True, True)}} + measure.scaler = scaler + output = torch.zeros(1, 1, 2, 2, requires_grad=True) + return measure, output + + +def test_accumulation_backward_uses_scaler_scale(): + """#AMP: accumulation losses must be scaled before backward when a GradScaler is set.""" + scaler = MagicMock() + scaled = MagicMock() + scaler.scale.return_value = scaled + + measure, output = _make_accumulating_measure(scaler) + target = torch.ones(1, 1, 2, 2) + measure.update("out", output, {"tgt": (target, [Attribute()])}, it=0, nb_patch=1, training=True) + + # The loss must go through scaler.scale(...).backward(), never a bare loss.backward(). + scaler.scale.assert_called_once() + scaled.backward.assert_called_once() + # Bare backward would have populated grads directly; the scaler intercepts it. + assert output.grad is None + + +def test_accumulation_backward_without_scaler_is_plain_backward(): + """Without a scaler the accumulation path must still back-propagate normally.""" + measure, output = _make_accumulating_measure(None) + target = torch.ones(1, 1, 2, 2) + measure.update("out", output, {"tgt": (target, [Attribute()])}, it=0, nb_patch=1, training=True) + + assert output.grad is not None + assert torch.count_nonzero(output.grad) > 0 + + +# -------------------------------------------------------------------------------------- +# Measure.update_scheduler — loss-weight window selection +# -------------------------------------------------------------------------------------- + + +def test_update_scheduler_empty_raises_config_error(): + """update_scheduler on an empty schedule must raise a clear ConfigError.""" + with pytest.raises(ConfigError): + Measure.update_scheduler(None, {}, 0) # type: ignore[arg-type] + + +def test_update_scheduler_past_last_window_clamps_to_last(): + """Past every configured window, the last scheduler is selected (no crash).""" + s0, s1 = Constant(1.0), Constant(2.0) + schedulers = {s0: 3, s1: 3} # active windows [0,3) and [3,6) + assert Measure.update_scheduler(None, schedulers, 4) is s1 # type: ignore[arg-type] + assert Measure.update_scheduler(None, schedulers, 100) is s1 # type: ignore[arg-type] + + +# -------------------------------------------------------------------------------------- +# CriterionsLoader +# -------------------------------------------------------------------------------------- + + +def test_network_criterion_loader_resets_scheduler_state(monkeypatch: pytest.MonkeyPatch) -> None: + class DummyMeasure: + def __init__(self) -> None: + pass + + class DummySchedulerLoader: + def __init__(self) -> None: + self.nb_step = 3 + + def getschedulers(self, key: str, scheduler_classname: str): + return f"{key}:{scheduler_classname}" + + monkeypatch.setattr(network_module, "apply_config", lambda *args, **kwargs: lambda cls: cls) + monkeypatch.setattr(network_module, "konfai_root", lambda: "Trainer") + monkeypatch.setattr( + network_module, + "get_module", + lambda classpath, default: (SimpleNamespace(Measure=DummyMeasure, __name__="torch.optim"), "Measure"), + ) + + attr = network_module.CriterionsAttr( + schedulers=cast( + dict[str, network_module.LossSchedulersLoader], + {"Constant": DummySchedulerLoader()}, + ) + ) + loader = network_module.CriterionsLoader({"dummy:Measure": attr}) + + loader.get_criterions("DemoModel", "Output", "Target") + first_schedulers = dict(attr.schedulers) + loader.get_criterions("DemoModel", "Output", "Target") + + assert attr.isTorchCriterion is True + assert len(attr.schedulers) == 1 + assert attr.schedulers == first_schedulers diff --git a/tests/unit/test_network_ddp_fixes.py b/tests/unit/test_network_ddp_fixes.py deleted file mode 100644 index a8bf2fe2..00000000 --- a/tests/unit/test_network_ddp_fixes.py +++ /dev/null @@ -1,123 +0,0 @@ -# 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 network / DDP / scheduler batch (see AUDIT.md §Training).""" - -import os - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -from unittest.mock import MagicMock # noqa: E402 - -import torch # noqa: E402 - -from konfai.metric.schedulers import Constant, PolyLRScheduler # noqa: E402 -from konfai.network.network import Measure # noqa: E402 -from konfai.utils.dataset import Attribute # noqa: E402 - - -class _Attr: - def __init__(self) -> None: - self.start = 0 - self.stop = None - self.schedulers = {Constant(1.0): 1} - self.group = 0 - self.is_loss = True - self.accumulation = True - - -def _make_accumulating_measure(scaler) -> tuple[Measure, torch.Tensor]: - """Build a minimal Measure that triggers the accumulation-backward branch.""" - measure = Measure.__new__(Measure) - criterion = torch.nn.MSELoss() - key = f"out:tgt:{criterion.__class__.__name__}" - measure.outputs_criterions = {"out": {"tgt": {criterion: _Attr()}}} - measure._loss = {0: {key: Measure.Loss(criterion.__class__.__name__, "out", "tgt", 0, True, True)}} - measure.scaler = scaler - output = torch.zeros(1, 1, 2, 2, requires_grad=True) - return measure, output - - -def test_accumulation_backward_uses_scaler_scale(): - """#AMP: accumulation losses must be scaled before backward when a GradScaler is set.""" - scaler = MagicMock() - scaled = MagicMock() - scaler.scale.return_value = scaled - - measure, output = _make_accumulating_measure(scaler) - target = torch.ones(1, 1, 2, 2) - measure.update("out", output, {"tgt": (target, [Attribute()])}, it=0, nb_patch=1, training=True) - - # The loss must go through scaler.scale(...).backward(), never a bare loss.backward(). - scaler.scale.assert_called_once() - scaled.backward.assert_called_once() - # Bare backward would have populated grads directly; the scaler intercepts it. - assert output.grad is None - - -def test_accumulation_backward_without_scaler_is_plain_backward(): - """Without a scaler the accumulation path must still back-propagate normally.""" - measure, output = _make_accumulating_measure(None) - target = torch.ones(1, 1, 2, 2) - measure.update("out", output, {"tgt": (target, [Attribute()])}, it=0, nb_patch=1, training=True) - - assert output.grad is not None - assert torch.count_nonzero(output.grad) > 0 - - -def test_polylr_resync_resumes_from_last_epoch(): - """#scheduler: PolyLR must honour a resync that sets last_epoch (RESUME fast-forward).""" - param = torch.nn.Parameter(torch.zeros(1)) - opt = torch.optim.SGD([param], lr=0.1) - scheduler = PolyLRScheduler(opt, initial_lr=0.1, max_steps=100) - - # A freshly built PolyLR keeps last_epoch == -1 so the network resync guard fires. - assert scheduler.last_epoch == -1 - - # Network.load() resync: fast-forward to iteration 50. - scheduler.last_epoch = 50 - scheduler.step() - - expected = 0.1 * (1 - 50 / 100) ** 0.9 - assert opt.param_groups[0]["lr"] == expected - assert scheduler.last_epoch == 51 - - -def test_polylr_fresh_run_unchanged(): - """Fresh training (no resync) still steps from the internal counter 0, 1, 2 ...""" - param = torch.nn.Parameter(torch.zeros(1)) - opt = torch.optim.SGD([param], lr=0.1) - scheduler = PolyLRScheduler(opt, initial_lr=0.1, max_steps=100) - - lrs = [] - for _ in range(3): - scheduler.step() - lrs.append(opt.param_groups[0]["lr"]) - - assert lrs[0] == 0.1 * (1 - 1 / 100) ** 0.9 - assert lrs[1] == 0.1 * (1 - 2 / 100) ** 0.9 - assert lrs[2] == 0.1 * (1 - 3 / 100) ** 0.9 - assert scheduler.last_epoch == -1 - - -def test_cyclegan_discriminator_initialized_no_keyerror(): - """#CycleGan: initialized() must not index a missing 'Sample' submodule on load.""" - from konfai.models.generation.diffusionGan import CycleGanDiscriminator - - model = CycleGanDiscriminator() - # Must not raise KeyError('Sample'). - model.initialized() diff --git a/tests/unit/test_patch_overlap_border.py b/tests/unit/test_patch_overlap_border.py deleted file mode 100644 index ce646838..00000000 --- a/tests/unit/test_patch_overlap_border.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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 test: overlap-blended reassembly must not darken the volume border.""" - -import torch -from konfai.data.patching import Accumulator, Cosinus, Mean - - -def test_overlap_blend_is_partition_of_unity_at_the_border() -> None: - # 1-D volume of 20 tiled with patch 8 / overlap 2 -> patches at 0, 6, 12. The border voxels are - # covered by a single patch, whose edge band weights ~0.5, so the pre-fix sum-without-normalise - # reassembled them at 0.5 instead of 1.0. Dividing by the accumulated weight restores unity. - patch_slices = [(slice(0, 8),), (slice(6, 14),), (slice(12, 20),)] - combine = Mean() # Cosinus needs >=2D (SimpleITK distance map), covered below. - combine.set_patch_config([8], 2) - accumulator = Accumulator(patch_slices, patch_size=[8], patch_combine=combine, batch=True) - for index in range(len(patch_slices)): - accumulator.add_layer(index, torch.ones(1, 1, 8)) - - out = accumulator.assemble()[0, 0] - - assert out.shape == (20,) - torch.testing.assert_close(out, torch.ones(20), rtol=0, atol=1e-5) - - -def test_overlap_blend_corner_not_quartered_in_2d() -> None: - # A 2-D corner is covered by one patch on both axes, so the pre-fix output was ~0.25 there. - patch_slices = [ - (slice(0, 8), slice(0, 8)), - (slice(0, 8), slice(6, 14)), - (slice(6, 14), slice(0, 8)), - (slice(6, 14), slice(6, 14)), - ] - for combine_cls in (Mean, Cosinus): - combine = combine_cls() - combine.set_patch_config([8, 8], 2) - accumulator = Accumulator(patch_slices, patch_size=[8, 8], patch_combine=combine, batch=True) - for index in range(len(patch_slices)): - accumulator.add_layer(index, torch.ones(1, 1, 8, 8)) - - out = accumulator.assemble()[0, 0] - - assert out.shape == (14, 14) - torch.testing.assert_close(out, torch.ones(14, 14), rtol=0, atol=1e-5) diff --git a/tests/unit/test_patching.py b/tests/unit/test_patching.py index a72c22a5..83b3e765 100644 --- a/tests/unit/test_patching.py +++ b/tests/unit/test_patching.py @@ -14,19 +14,13 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for patch reconstruction and overlap-blending helpers.""" +"""Unit tests for patch reconstruction and overlap-blending (``konfai.data.patching``).""" -import os - -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import pytest # noqa: E402 -import torch # noqa: E402 - -from konfai.data.patching import Accumulator, Cosinus, Mean # noqa: E402 -from konfai.utils.errors import PatchError # noqa: E402 -from konfai.utils.utils import get_patch_slices_from_shape # noqa: E402 +import pytest +import torch +from konfai.data.patching import Accumulator, Cosinus, Mean +from konfai.utils.errors import PatchError +from konfai.utils.utils import get_patch_slices_from_shape def _tile_2d(full: torch.Tensor, patch_size: list[int], overlap: int): @@ -36,6 +30,11 @@ def _tile_2d(full: torch.Tensor, patch_size: list[int], overlap: int): return patch_slices, patches +# -------------------------------------------------------------------------------------- +# Accumulator reassembly +# -------------------------------------------------------------------------------------- + + def test_accumulator_reconstructs_non_overlapping_tiles(): """Without blending, non-overlapping patches must reassemble exactly.""" full = torch.arange(1 * 1 * 4 * 4, dtype=torch.float32).reshape(1, 1, 4, 4) @@ -90,6 +89,11 @@ def test_assemble_with_missing_first_patch_does_not_crash(): assert torch.equal(out[:, :, 2:4, :], full[:, :, 2:4, :]) +# -------------------------------------------------------------------------------------- +# Blending windows (Mean / Cosinus) +# -------------------------------------------------------------------------------------- + + @pytest.mark.parametrize("combine_cls", [Mean, Cosinus]) def test_path_combine_window_is_bounded_and_unit_at_center(combine_cls): """Blending windows weight each voxel in [0, 1] and reach 1 at the patch centre.""" @@ -120,3 +124,62 @@ def test_path_combine_call_applies_window_and_caches_device(): assert torch.allclose(weighted[0, 0], combine.data) # The per-device window is cached on first use. assert tensor.device in combine._data_per_device + + +def test_path_combine_overlap_zero_uses_uniform_weights() -> None: + """B10: overlap=0 tiles patches without overlap, so the blend window is all ones.""" + for combine_cls in (Mean, Cosinus): + combine = combine_cls() + combine.set_patch_config([8, 8, 8], 0) # must not raise + assert combine.data.shape == (8, 8, 8) + assert torch.equal(combine.data, torch.ones(8, 8, 8)) + + +def test_path_combine_overlap_zero_leaves_tensor_unchanged() -> None: + combine = Mean() + combine.set_patch_config([4, 4], 0) + tensor = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4) + assert torch.equal(combine(tensor), tensor) + + +# -------------------------------------------------------------------------------------- +# Overlap-blended reassembly is a partition of unity (no darkened borders) +# -------------------------------------------------------------------------------------- + + +def test_overlap_blend_is_partition_of_unity_at_the_border() -> None: + # 1-D volume of 20 tiled with patch 8 / overlap 2 -> patches at 0, 6, 12. The border voxels are + # covered by a single patch, whose edge band weights ~0.5, so the pre-fix sum-without-normalise + # reassembled them at 0.5 instead of 1.0. Dividing by the accumulated weight restores unity. + patch_slices = [(slice(0, 8),), (slice(6, 14),), (slice(12, 20),)] + combine = Mean() # Cosinus needs >=2D (SimpleITK distance map), covered below. + combine.set_patch_config([8], 2) + accumulator = Accumulator(patch_slices, patch_size=[8], patch_combine=combine, batch=True) + for index in range(len(patch_slices)): + accumulator.add_layer(index, torch.ones(1, 1, 8)) + + out = accumulator.assemble()[0, 0] + + assert out.shape == (20,) + torch.testing.assert_close(out, torch.ones(20), rtol=0, atol=1e-5) + + +def test_overlap_blend_corner_not_quartered_in_2d() -> None: + # A 2-D corner is covered by one patch on both axes, so the pre-fix output was ~0.25 there. + patch_slices = [ + (slice(0, 8), slice(0, 8)), + (slice(0, 8), slice(6, 14)), + (slice(6, 14), slice(0, 8)), + (slice(6, 14), slice(6, 14)), + ] + for combine_cls in (Mean, Cosinus): + combine = combine_cls() + combine.set_patch_config([8, 8], 2) + accumulator = Accumulator(patch_slices, patch_size=[8, 8], patch_combine=combine, batch=True) + for index in range(len(patch_slices)): + accumulator.add_layer(index, torch.ones(1, 1, 8, 8)) + + out = accumulator.assemble()[0, 0] + + assert out.shape == (14, 14) + torch.testing.assert_close(out, torch.ones(14, 14), rtol=0, atol=1e-5) diff --git a/tests/unit/test_perceptual_loss_targets.py b/tests/unit/test_perceptual_loss_targets.py deleted file mode 100644 index de77a536..00000000 --- a/tests/unit/test_perceptual_loss_targets.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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 test: PerceptualLoss.forward must unpack its targets into _compute.""" - -import torch -from konfai.metric.measure import PerceptualLoss - - -def test_perceptual_loss_forward_unpacks_targets() -> None: - # forward(output, *targets) must hand each target to _compute(output, *targets) as its own - # positional tensor; the pre-fix code passed the whole tuple as a single argument, so the - # preprocessing/feature-extraction path received a tuple and crashed. - loss = object.__new__(PerceptualLoss) - loss.shape = [128, 128, 128] # len != 2 -> the non-slice branch is taken - loss.models = {None: object()} # short-circuit the lazy model placement on device index None - - recorded: dict[str, tuple] = {} - - def fake_compute(output, *targets): - recorded["targets"] = targets - return torch.zeros(1) - - loss._compute = fake_compute # type: ignore[method-assign] - - PerceptualLoss.forward(loss, torch.randn(1, 1, 8, 8), torch.randn(1, 1, 8, 8)) - - assert len(recorded["targets"]) == 1 - assert torch.is_tensor(recorded["targets"][0]) diff --git a/tests/unit/test_perf_hot_paths.py b/tests/unit/test_perf_hot_paths.py index 5a60e1ea..f519d5dc 100644 --- a/tests/unit/test_perf_hot_paths.py +++ b/tests/unit/test_perf_hot_paths.py @@ -16,20 +16,15 @@ """Regression tests for the performance hot-path fixes (see AUDIT.md — Performance backlog).""" -import os +from pathlib import Path +from unittest.mock import MagicMock -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") +import torch -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 +import konfai.utils.dataset as dataset_module +from konfai.data.patching import Accumulator +from konfai.predictor import ModelComposite +from konfai.utils.dataset import Attribute, Dataset def test_accumulator_is_full_counts_without_rescanning(): diff --git a/tests/unit/test_runtime_guards.py b/tests/unit/test_runtime_guards.py index 1b44c416..15eb6c1b 100644 --- a/tests/unit/test_runtime_guards.py +++ b/tests/unit/test_runtime_guards.py @@ -1,18 +1,33 @@ +# 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 + +"""Guard tests for ``konfai.utils.runtime``: workflow preconditions, environment +normalisation, overwrite confirmation, and distributed-launch bookkeeping.""" + import os import sys from pathlib import Path from types import SimpleNamespace import pytest -import torch import konfai as konfai_module -from konfai.data.data_manager import Data from konfai.evaluator import Evaluator -from konfai.network.blocks import Exit from konfai.predictor import Predictor from konfai.trainer import Trainer -from konfai.utils.config import apply_config, config from konfai.utils.errors import ConfigError from konfai.utils.runtime import ( DistributedObject, @@ -53,46 +68,6 @@ def test_configure_workflow_environment_normalizes_paths(monkeypatch: pytest.Mon assert Path(os.environ["KONFAI_STATISTICS_DIRECTORY"]).name == "Statistics" -def test_apply_config_restores_config_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - config_path = tmp_path / "Config.yml" - config_path.write_text("Root:\n Child:\n value: 7\n", encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - monkeypatch.setenv("KONFAI_CONFIG_PATH", "before.path") - monkeypatch.setenv("KONFAI_CONFIG_VARIABLE", "before.variable") - - @config("Child") - class Child: - def __init__(self, value: int = 0) -> None: - self.value = value - - child = apply_config("Root")(Child)() - - assert child.value == 7 - assert os.environ["KONFAI_CONFIG_PATH"] == "before.path" - assert os.environ["KONFAI_CONFIG_VARIABLE"] == "before.variable" - - -def test_apply_config_keeps_config_path_during_constructor_call( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - config_path = tmp_path / "Config.yml" - config_path.write_text("Root:\n Child:\n value: 7\n", encoding="utf-8") - monkeypatch.setenv("KONFAI_config_file", str(config_path)) - monkeypatch.setenv("KONFAI_CONFIG_MODE", "Done") - - @config("Child") - class Child: - def __init__(self, value: int = 0) -> None: - self.value = value - self.config_path = os.environ["KONFAI_CONFIG_PATH"] - - child = apply_config("Root")(Child)() - - assert child.value == 7 - assert child.config_path == "Root.Child" - - def test_confirm_overwrite_or_raise_requires_flag_in_non_interactive(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("KONFAI_OVERWRITE", raising=False) monkeypatch.setattr(sys, "stdin", SimpleNamespace(isatty=lambda: False)) @@ -121,27 +96,6 @@ def test_confirm_overwrite_or_raise_rejects_decline(monkeypatch: pytest.MonkeyPa confirm_overwrite_or_raise(Path("/tmp/output"), "prediction", ConfigError) -def test_debug_exit_block_raises_runtime_error() -> None: - with pytest.raises(RuntimeError, match="debug Exit block"): - Exit()(torch.ones(1)) - - -def test_data_split_does_not_duplicate_tail_samples(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("KONFAI_STATE", str(State.TRAIN)) - - mapping = [(index, 0, 0) for index in range(5)] - shards = Data._split(mapping, 2) - - flattened = [item for shard in shards for item in shard] - - assert len(shards) == 2 - # drop_last semantics: shards are equal length (DDP needs matching step counts), the tail is - # dropped, and no sample is ever duplicated across ranks. - assert {len(shard) for shard in shards} == {2} - assert len(flattened) == len(set(flattened)) - assert set(flattened).issubset(set(mapping)) - - def test_execute_distributed_object_sets_shared_master_port_without_forcing_launch_blocking( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/test_runtime_progress_ddp.py b/tests/unit/test_runtime_progress_ddp.py index 5bf02256..6ea1e8aa 100644 --- a/tests/unit/test_runtime_progress_ddp.py +++ b/tests/unit/test_runtime_progress_ddp.py @@ -17,13 +17,9 @@ """Regression tests for the runtime progress/DDP audit fixes (see AUDIT.md).""" import contextlib -import os import random -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -import konfai.utils.runtime as rt # noqa: E402 +import konfai.utils.runtime as rt def test_synchronize_data_gathers_on_cpu(monkeypatch): diff --git a/tests/unit/test_schedulers.py b/tests/unit/test_schedulers.py new file mode 100644 index 00000000..66b6a7f8 --- /dev/null +++ b/tests/unit/test_schedulers.py @@ -0,0 +1,55 @@ +# 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 + +"""Tests for the learning-rate schedulers in ``konfai.metric.schedulers``.""" + +import torch +from konfai.metric.schedulers import PolyLRScheduler + + +def test_polylr_resync_resumes_from_last_epoch(): + """#scheduler: PolyLR must honour a resync that sets last_epoch (RESUME fast-forward).""" + param = torch.nn.Parameter(torch.zeros(1)) + opt = torch.optim.SGD([param], lr=0.1) + scheduler = PolyLRScheduler(opt, initial_lr=0.1, max_steps=100) + + # A freshly built PolyLR keeps last_epoch == -1 so the network resync guard fires. + assert scheduler.last_epoch == -1 + + # Network.load() resync: fast-forward to iteration 50. + scheduler.last_epoch = 50 + scheduler.step() + + expected = 0.1 * (1 - 50 / 100) ** 0.9 + assert opt.param_groups[0]["lr"] == expected + assert scheduler.last_epoch == 51 + + +def test_polylr_fresh_run_unchanged(): + """Fresh training (no resync) still steps from the internal counter 0, 1, 2 ...""" + param = torch.nn.Parameter(torch.zeros(1)) + opt = torch.optim.SGD([param], lr=0.1) + scheduler = PolyLRScheduler(opt, initial_lr=0.1, max_steps=100) + + lrs = [] + for _ in range(3): + scheduler.step() + lrs.append(opt.param_groups[0]["lr"]) + + assert lrs[0] == 0.1 * (1 - 1 / 100) ** 0.9 + assert lrs[1] == 0.1 * (1 - 2 / 100) ** 0.9 + assert lrs[2] == 0.1 * (1 - 3 / 100) ** 0.9 + assert scheduler.last_epoch == -1 diff --git a/tests/unit/test_transform.py b/tests/unit/test_transform.py new file mode 100644 index 00000000..d8f0e18f --- /dev/null +++ b/tests/unit/test_transform.py @@ -0,0 +1,395 @@ +# 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 + +"""Tests for ``konfai.data.transform``: Clip, Dilate, Norm, Crop, Standardize, Padding, +ResampleToResolution/ResampleToShape, InferenceStack, and KonfAIInference.""" + +import sys +import types +from pathlib import Path + +import numpy as np +import pytest +import torch +import torch.nn.functional as F +from konfai.data.transform import ( + DEFAULT_INFERENCE_MODEL_NAME, + DEFAULT_INFERENCE_REPO_ID, + Clip, + Crop, + Dilate, + InferenceStack, + KonfAIInference, + Norm, + Padding, + ResampleToResolution, + ResampleToShape, + Standardize, +) +from konfai.utils.dataset import Attribute +from konfai.utils.errors import TransformError + +# -------------------------------------------------------------------------------------- +# Clip +# -------------------------------------------------------------------------------------- + + +def test_clip_resolves_min_and_percentile_bounds() -> None: + # ``min`` (torch scalar) and ``percentile:

`` (numpy scalar) bounds must be coerced to float + # so the in-place clip assignments are valid for a torch tensor. + tensor = torch.arange(0, 100, dtype=torch.float32) + clip = Clip(min_value="min", max_value="percentile:90") + + out = clip("case", tensor.clone(), Attribute()) + + assert out.min().item() == pytest.approx(0.0) + assert out.max().item() == pytest.approx(89.1) + assert out.dtype == torch.float32 + + +def test_clip_fixed_numeric_bounds() -> None: + tensor = torch.arange(-50, 50, dtype=torch.float32) + clip = Clip(min_value=-10.0, max_value=10.0) + + out = clip("case", tensor.clone(), Attribute()) + + assert out.min().item() == pytest.approx(-10.0) + assert out.max().item() == pytest.approx(10.0) + + +# -------------------------------------------------------------------------------------- +# Dilate +# -------------------------------------------------------------------------------------- + + +def _dense_cube_dilation(tensor: torch.Tensor, dilate: int) -> torch.Tensor: + """Reference: dilation via a single dense k**n max-pool (the pre-separable implementation).""" + data = (tensor > 0).to(torch.float32) + k = 2 * dilate + 1 + if data.dim() - 1 == 2: + data = F.max_pool2d(data, kernel_size=k, stride=1, padding=dilate) + else: + data = F.max_pool3d(data, kernel_size=k, stride=1, padding=dilate) + return data.to(tensor.dtype) + + +@pytest.mark.parametrize("dilate", [1, 2, 5]) +@pytest.mark.parametrize("shape", [(1, 24, 30), (2, 12, 20, 18)]) +def test_dilate_separable_matches_dense_cube(shape: tuple[int, ...], dilate: int) -> None: + # The separable 1-D max-pool implementation must be bit-identical to the dense k**n cube it replaces, + # for both [C,H,W] and [C,D,H,W] inputs and several radii — this is the correctness guarantee that + # lets the ~14x speedup ship as a transparent optimization. + torch.manual_seed(0) + mask = (torch.rand(shape) > 0.7).to(torch.uint8) + + out = Dilate(dilate)("case", mask.clone(), Attribute()) + ref = _dense_cube_dilation(mask, dilate) + + assert torch.equal(out, ref) + assert out.dtype == mask.dtype + assert out.shape == mask.shape + + +def test_dilate_single_voxel_fills_neighbourhood() -> None: + # A single active voxel dilated by 1 must fill its full 3x3x3 neighbourhood. + mask = torch.zeros(1, 5, 5, 5, dtype=torch.uint8) + mask[0, 2, 2, 2] = 1 + + out = Dilate(1)("case", mask.clone(), Attribute()) + + assert out[0, 1:4, 1:4, 1:4].sum().item() == 27 + assert out.sum().item() == 27 + + +def test_dilate_zero_is_identity() -> None: + mask = (torch.rand(1, 8, 8, 8) > 0.5).to(torch.uint8) + out = Dilate(0)("case", mask.clone(), Attribute()) + assert torch.equal(out, mask) + + +# -------------------------------------------------------------------------------------- +# Norm +# -------------------------------------------------------------------------------------- + + +def _stack_attribute() -> Attribute: + # Geometry of a displacement-field stack: the leading image axis holds the vector components + # (origin 0 / spacing 1 / identity direction row), the remaining axes carry the fixed grid. + attribute = Attribute() + attribute["Origin"] = np.asarray([0.0, 1.0, 2.0, 3.0]) + attribute["Spacing"] = np.asarray([1.0, 2.0, 2.0, 2.0]) + direction = np.eye(4) + direction[1:, 1:] = np.diag([1.0, -1.0, 1.0]) + attribute["Direction"] = direction.flatten() + return attribute + + +def test_norm_reduces_trailing_component_axis_and_geometry() -> None: + # A stack of 2 displacement fields [N=2, D, H, W, C=3] -> per-sample magnitudes [2, D, H, W]. + tensors = torch.randn(2, 4, 5, 6, 3) + attribute = _stack_attribute() + + out = Norm()("case", tensors, attribute) + + assert list(out.shape) == [2, 4, 5, 6] + assert torch.allclose(out, torch.linalg.norm(tensors, dim=-1)) + # The reduced trailing tensor axis is the first geometry axis: it must be dropped. + assert attribute.get_np_array("Origin").tolist() == [1.0, 2.0, 3.0] + assert attribute.get_np_array("Spacing").tolist() == [2.0, 2.0, 2.0] + assert attribute.get_np_array("Direction").tolist() == np.diag([1.0, -1.0, 1.0]).flatten().tolist() + + +def test_norm_transform_shape_drops_trailing_axis() -> None: + assert Norm().transform_shape("group", "case", [4, 5, 6, 3], Attribute()) == [4, 5, 6] + + +# -------------------------------------------------------------------------------------- +# Crop — transform_shape predicts the spatial crop exactly (patch planning depends on it) +# -------------------------------------------------------------------------------------- + + +def test_crop_transform_shape_matches_spatial_crop() -> None: + # The pre-fix code treated ``shape[0]`` as a channel dim and paired the crop box with + # ``shape[1:]``, shifting every axis by one and returning a wrong shape. + attribute = Attribute() + attribute["box"] = np.array([[2, 3], [1, 1], [4, 2]]) # (start, end-distance) per spatial axis + + out = Crop().transform_shape("CT", "CASE_001", [10, 20, 30], attribute) + + # 10-2-3, 20-1-1, 30-4-2 — each spatial axis cropped by its own box row. + assert out == [5, 18, 24] + + +# -------------------------------------------------------------------------------------- +# Standardize +# -------------------------------------------------------------------------------------- + + +def test_standardize_explicit_scalar_stats(): + """#5 Standardize with explicit scalar mean/std must not crash.""" + t = Standardize(lazy=False, mean=[10.0], std=[2.0]) + x = torch.arange(24, dtype=torch.float32).reshape(1, 2, 3, 4) + out = t("c", x.clone(), Attribute()) + assert torch.allclose(out, (x - 10.0) / 2.0) + + +def test_standardize_explicit_per_channel_stats(): + """#5 Per-channel mean/std broadcast over the channel axis.""" + t = Standardize(lazy=False, mean=[10.0, 20.0], std=[2.0, 4.0]) + x = torch.zeros(2, 3, 4) + x[0] = 10.0 + x[1] = 20.0 + out = t("c", x.clone(), Attribute()) + assert torch.allclose(out, torch.zeros_like(out), atol=1e-6) + + +# -------------------------------------------------------------------------------------- +# Padding — origin bookkeeping +# -------------------------------------------------------------------------------------- + + +def test_padding_shifts_origin_along_the_padded_axes(image_attributes): + """Each F.pad pair (X, Y, Z) must shift the matching (x, y, z) origin component.""" + attributes = image_attributes([10.0, 20.0, 30.0], [1.0, 2.0, 4.0]) + + padded = Padding(padding=[1, 0, 2, 0, 3, 0])("case", torch.zeros(1, 5, 5, 5), attributes) + + assert list(padded.shape) == [1, 8, 7, 6] + np.testing.assert_allclose( + attributes.get_np_array("Origin"), + [10.0 - 1 * 1.0, 20.0 - 2 * 2.0, 30.0 - 3 * 4.0], + ) + + +def test_padding_after_the_data_keeps_origin(image_attributes): + """Padding only on the high side of each axis must leave the origin untouched.""" + attributes = image_attributes([10.0, 20.0, 30.0], [1.0, 2.0, 4.0]) + + padded = Padding(padding=[0, 2, 0, 0, 0, 1])("case", torch.zeros(1, 5, 5, 5), attributes) + + assert list(padded.shape) == [1, 6, 5, 7] + np.testing.assert_allclose(attributes.get_np_array("Origin"), [10.0, 20.0, 30.0]) + + +# -------------------------------------------------------------------------------------- +# ResampleToResolution / ResampleToShape +# -------------------------------------------------------------------------------------- + + +def test_resample_to_resolution_transform_shape_missing_spacing_raises(): + """A tensor without 'Spacing' metadata must surface a TransformError, not fall through.""" + with pytest.raises(TransformError): + ResampleToResolution().transform_shape("group", "case", [10, 10, 10], Attribute()) + + +def test_resample_to_shape_transform_shape_missing_spacing_raises(): + """ResampleToShape must also raise when 'Spacing' metadata is absent.""" + with pytest.raises(TransformError): + ResampleToShape().transform_shape("group", "case", [10, 10, 10], Attribute()) + + +def test_resample_to_resolution_transform_shape_dimension_mismatch_message(): + """The dimension-mismatch error is raised and its message interpolates the actual shape.""" + attributes = Attribute() + attributes["Spacing"] = np.asarray([1.0, 1.0], dtype=np.float64) + with pytest.raises(TransformError) as excinfo: + ResampleToResolution(spacing=[1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], attributes) + assert "shape=[10, 10, 10]" in str(excinfo.value) + + +def test_resample_to_shape_transform_shape_dimension_mismatch_message(): + """ResampleToShape raises a formatted (f-string) message on a shape/target mismatch.""" + attributes = Attribute() + attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) + with pytest.raises(TransformError) as excinfo: + ResampleToShape(shape=[4, 4]).transform_shape("group", "case", [10, 10, 10], attributes) + message = str(excinfo.value) + assert "shape=[10, 10, 10]" in message + assert "target_shape" in message + + +def test_resample_to_shape_does_not_mutate_config(): + """#9 transform_shape must not write resolved dims back into the shared instance config.""" + resampler = ResampleToShape(shape=[0, 16, 16]) + before = resampler.shape.clone() + attributes = Attribute() + attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) + out = resampler.transform_shape("CT", "case", [8, 16, 16], attributes) + assert out[0] == 8 # sentinel 0 resolved to the input dim for this call + assert torch.equal(resampler.shape, before), "self.shape must stay [0, 16, 16] for the next case" + + +def test_resample_to_shape_inverse_without_spacing_metadata(): + """Inverting a resample must not pop a 'Spacing' the forward pass never pushed.""" + resampler = ResampleToShape(shape=[4, 4, 4]) + attributes = Attribute() # no image metadata at all + tensor = torch.arange(8 * 8 * 8, dtype=torch.float32).reshape(1, 8, 8, 8) + + forward = resampler("case", tensor, attributes) + assert list(forward.shape) == [1, 4, 4, 4] + + restored = resampler.inverse("case", forward, attributes) + assert list(restored.shape) == [1, 8, 8, 8] + + +def test_resample_to_shape_inverse_pops_pushed_spacing(): + """When 'Spacing' exists, the inverse removes the version the forward pass pushed.""" + resampler = ResampleToShape(shape=[4, 4, 4]) + attributes = Attribute() + attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) + tensor = torch.zeros(1, 8, 8, 8) + + resampler("case", tensor, attributes) + # image_shape / shape == 2 on every axis, so the resampled spacing doubles. + np.testing.assert_allclose(attributes.get_np_array("Spacing"), [2.0, 2.0, 2.0]) + + resampler.inverse("case", torch.zeros(1, 4, 4, 4), attributes) + # The pushed spacing is popped, restoring the original resolution. + np.testing.assert_allclose(attributes.get_np_array("Spacing"), [1.0, 1.0, 1.0]) + + +# -------------------------------------------------------------------------------------- +# InferenceStack / KonfAIInference +# -------------------------------------------------------------------------------------- + + +def test_inference_stack_super_init_enables_dataset_fallback(): + """InferenceStack must inherit Transform's 'datasets' list for the fallback write path.""" + stack = InferenceStack(dataset="", name="pred", mode="mean") + assert stack.datasets == [] + + written: dict[str, object] = {} + + class _FakeDataset: + def write(self, group, name, data, cache_attribute): + written["group"] = group + written["name"] = name + written["shape"] = tuple(data.shape) + + stack.set_datasets([object(), _FakeDataset()]) + + tensors = torch.stack( + [torch.full((1, 2, 2), 3.0), torch.full((1, 2, 2), 5.0)], + dim=0, + ) # shape [2, 1, 2, 2]: two stacked predictions + out = stack("pred", tensors, Attribute()) + + assert written["group"] == "InferenceStack" + assert written["name"] == "pred" + assert written["shape"] == (2, 2, 2) + # 'mean' mode averages the two predictions element-wise. + assert torch.allclose(out, torch.full((1, 2, 2), 4.0)) + + +def test_konfai_inference_reassembles_channels_in_sorted_order(tmp_path, monkeypatch): + """Per-channel outputs must be stacked in deterministic (sorted) case order.""" + sitk = pytest.importorskip("SimpleITK") + + output_dir = tmp_path / "Output" + files = [] + for i in range(3): + case_dir = output_dir / f"P{i:03d}" + case_dir.mkdir(parents=True) + array = np.full((2, 2, 2), float(i * 10), dtype=np.float32) + path = case_dir / "Volume.mha" + sitk.WriteImage(sitk.GetImageFromArray(array), str(path)) + files.append(path) + + # Simulate an arbitrary (here reversed) filesystem enumeration order. + scrambled = list(reversed(files)) + monkeypatch.setattr(Path, "rglob", lambda self, pattern: iter(scrambled)) + + result = KonfAIInference._reassemble_output(output_dir) + + assert list(result.shape) == [3, 2, 2, 2] + assert float(result[0].mean()) == 0.0 + assert float(result[1].mean()) == 10.0 + assert float(result[2].mean()) == 20.0 + + +def test_konfai_inference_default_repo_and_model_preserved(): + """Constructing without arguments keeps the current published repo/model default.""" + transform = KonfAIInference() + + assert transform.repo_id == DEFAULT_INFERENCE_REPO_ID + assert transform.model_name == DEFAULT_INFERENCE_MODEL_NAME + assert transform.repo_id == "VBoussot/MRSegmentator-KonfAI" + assert transform.model_name == "MRSegmentator" + + +def test_konfai_inference_forwards_configured_repo_and_model(monkeypatch): + """A custom repo/model is forwarded verbatim to the KonfAIApp spec, not the default.""" + captured = {} + + class _FakeKonfAIApp: + def __init__(self, spec, *args): + captured["spec"] = spec + + def infer(self, *args, **kwargs): + captured["infer"] = (args, kwargs) + + fake_module = types.ModuleType("konfai_apps") + fake_module.KonfAIApp = _FakeKonfAIApp + monkeypatch.setitem(sys.modules, "konfai_apps", fake_module) + + transform = KonfAIInference( + repo_id="acme/Custom-KonfAI", + model_name="CustomModel", + checkpoints_name=["fold_1"], + ) + transform.infer_entry(Path("dataset"), Path("output"), []) + + assert captured["spec"] == "acme/Custom-KonfAI:CustomModel" diff --git a/tests/unit/test_transform_clip.py b/tests/unit/test_transform_clip.py deleted file mode 100644 index 93d13edc..00000000 --- a/tests/unit/test_transform_clip.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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 - -import pytest -import torch -from konfai.data.transform import Clip -from konfai.utils.dataset import Attribute - - -def test_clip_resolves_min_and_percentile_bounds() -> None: - # ``min`` (torch scalar) and ``percentile:

`` (numpy scalar) bounds must be coerced to float - # so the in-place clip assignments are valid for a torch tensor. - tensor = torch.arange(0, 100, dtype=torch.float32) - clip = Clip(min_value="min", max_value="percentile:90") - - out = clip("case", tensor.clone(), Attribute()) - - assert out.min().item() == pytest.approx(0.0) - assert out.max().item() == pytest.approx(89.1) - assert out.dtype == torch.float32 - - -def test_clip_fixed_numeric_bounds() -> None: - tensor = torch.arange(-50, 50, dtype=torch.float32) - clip = Clip(min_value=-10.0, max_value=10.0) - - out = clip("case", tensor.clone(), Attribute()) - - assert out.min().item() == pytest.approx(-10.0) - assert out.max().item() == pytest.approx(10.0) diff --git a/tests/unit/test_transform_dilate.py b/tests/unit/test_transform_dilate.py deleted file mode 100644 index 80ac3838..00000000 --- a/tests/unit/test_transform_dilate.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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 - -import pytest -import torch -import torch.nn.functional as F -from konfai.data.transform import Dilate -from konfai.utils.dataset import Attribute - - -def _dense_cube_dilation(tensor: torch.Tensor, dilate: int) -> torch.Tensor: - """Reference: dilation via a single dense k**n max-pool (the pre-separable implementation).""" - data = (tensor > 0).to(torch.float32) - k = 2 * dilate + 1 - if data.dim() - 1 == 2: - data = F.max_pool2d(data, kernel_size=k, stride=1, padding=dilate) - else: - data = F.max_pool3d(data, kernel_size=k, stride=1, padding=dilate) - return data.to(tensor.dtype) - - -@pytest.mark.parametrize("dilate", [1, 2, 5]) -@pytest.mark.parametrize("shape", [(1, 24, 30), (2, 12, 20, 18)]) -def test_dilate_separable_matches_dense_cube(shape: tuple[int, ...], dilate: int) -> None: - # The separable 1-D max-pool implementation must be bit-identical to the dense k**n cube it replaces, - # for both [C,H,W] and [C,D,H,W] inputs and several radii — this is the correctness guarantee that - # lets the ~14x speedup ship as a transparent optimization. - torch.manual_seed(0) - mask = (torch.rand(shape) > 0.7).to(torch.uint8) - - out = Dilate(dilate)("case", mask.clone(), Attribute()) - ref = _dense_cube_dilation(mask, dilate) - - assert torch.equal(out, ref) - assert out.dtype == mask.dtype - assert out.shape == mask.shape - - -def test_dilate_single_voxel_fills_neighbourhood() -> None: - # A single active voxel dilated by 1 must fill its full 3x3x3 neighbourhood. - mask = torch.zeros(1, 5, 5, 5, dtype=torch.uint8) - mask[0, 2, 2, 2] = 1 - - out = Dilate(1)("case", mask.clone(), Attribute()) - - assert out[0, 1:4, 1:4, 1:4].sum().item() == 27 - assert out.sum().item() == 27 - - -def test_dilate_zero_is_identity() -> None: - mask = (torch.rand(1, 8, 8, 8) > 0.5).to(torch.uint8) - out = Dilate(0)("case", mask.clone(), Attribute()) - assert torch.equal(out, mask) diff --git a/tests/unit/test_transform_fixes.py b/tests/unit/test_transform_fixes.py deleted file mode 100644 index d29bc589..00000000 --- a/tests/unit/test_transform_fixes.py +++ /dev/null @@ -1,184 +0,0 @@ -# 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 transform-pipeline fixes.""" - -import sys -import types -from pathlib import Path - -import numpy as np -import pytest -import SimpleITK as sitk -import torch -from konfai.data.transform import ( - DEFAULT_INFERENCE_MODEL_NAME, - DEFAULT_INFERENCE_REPO_ID, - InferenceStack, - KonfAIInference, - ResampleToResolution, - ResampleToShape, -) -from konfai.utils.dataset import Attribute -from konfai.utils.errors import TransformError - - -def test_resample_to_resolution_transform_shape_missing_spacing_raises(): - """A tensor without 'Spacing' metadata must surface a TransformError, not fall through.""" - with pytest.raises(TransformError): - ResampleToResolution().transform_shape("group", "case", [10, 10, 10], Attribute()) - - -def test_resample_to_shape_transform_shape_missing_spacing_raises(): - """ResampleToShape must also raise when 'Spacing' metadata is absent.""" - with pytest.raises(TransformError): - ResampleToShape().transform_shape("group", "case", [10, 10, 10], Attribute()) - - -def test_resample_to_resolution_transform_shape_dimension_mismatch_message(): - """The dimension-mismatch error is raised and its message interpolates the actual shape.""" - attributes = Attribute() - attributes["Spacing"] = np.asarray([1.0, 1.0], dtype=np.float64) - with pytest.raises(TransformError) as excinfo: - ResampleToResolution(spacing=[1.0, 1.0]).transform_shape("group", "case", [10, 10, 10], attributes) - assert "shape=[10, 10, 10]" in str(excinfo.value) - - -def test_resample_to_shape_transform_shape_dimension_mismatch_message(): - """ResampleToShape raises a formatted (f-string) message on a shape/target mismatch.""" - attributes = Attribute() - attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) - with pytest.raises(TransformError) as excinfo: - ResampleToShape(shape=[4, 4]).transform_shape("group", "case", [10, 10, 10], attributes) - message = str(excinfo.value) - assert "shape=[10, 10, 10]" in message - assert "target_shape" in message - - -def test_resample_to_shape_inverse_without_spacing_metadata(): - """Inverting a resample must not pop a 'Spacing' the forward pass never pushed.""" - resampler = ResampleToShape(shape=[4, 4, 4]) - attributes = Attribute() # no image metadata at all - tensor = torch.arange(8 * 8 * 8, dtype=torch.float32).reshape(1, 8, 8, 8) - - forward = resampler("case", tensor, attributes) - assert list(forward.shape) == [1, 4, 4, 4] - - restored = resampler.inverse("case", forward, attributes) - assert list(restored.shape) == [1, 8, 8, 8] - - -def test_resample_to_shape_inverse_pops_pushed_spacing(): - """When 'Spacing' exists, the inverse removes the version the forward pass pushed.""" - resampler = ResampleToShape(shape=[4, 4, 4]) - attributes = Attribute() - attributes["Spacing"] = np.asarray([1.0, 1.0, 1.0], dtype=np.float64) - tensor = torch.zeros(1, 8, 8, 8) - - resampler("case", tensor, attributes) - # image_shape / shape == 2 on every axis, so the resampled spacing doubles. - np.testing.assert_allclose(attributes.get_np_array("Spacing"), [2.0, 2.0, 2.0]) - - resampler.inverse("case", torch.zeros(1, 4, 4, 4), attributes) - # The pushed spacing is popped, restoring the original resolution. - np.testing.assert_allclose(attributes.get_np_array("Spacing"), [1.0, 1.0, 1.0]) - - -def test_inference_stack_super_init_enables_dataset_fallback(): - """InferenceStack must inherit Transform's 'datasets' list for the fallback write path.""" - stack = InferenceStack(dataset="", name="pred", mode="mean") - assert stack.datasets == [] - - written: dict[str, object] = {} - - class _FakeDataset: - def write(self, group, name, data, cache_attribute): - written["group"] = group - written["name"] = name - written["shape"] = tuple(data.shape) - - stack.set_datasets([object(), _FakeDataset()]) - - tensors = torch.stack( - [torch.full((1, 2, 2), 3.0), torch.full((1, 2, 2), 5.0)], - dim=0, - ) # shape [2, 1, 2, 2]: two stacked predictions - out = stack("pred", tensors, Attribute()) - - assert written["group"] == "InferenceStack" - assert written["name"] == "pred" - assert written["shape"] == (2, 2, 2) - # 'mean' mode averages the two predictions element-wise. - assert torch.allclose(out, torch.full((1, 2, 2), 4.0)) - - -def test_konfai_inference_reassembles_channels_in_sorted_order(tmp_path, monkeypatch): - """Per-channel outputs must be stacked in deterministic (sorted) case order.""" - output_dir = tmp_path / "Output" - files = [] - for i in range(3): - case_dir = output_dir / f"P{i:03d}" - case_dir.mkdir(parents=True) - array = np.full((2, 2, 2), float(i * 10), dtype=np.float32) - path = case_dir / "Volume.mha" - sitk.WriteImage(sitk.GetImageFromArray(array), str(path)) - files.append(path) - - # Simulate an arbitrary (here reversed) filesystem enumeration order. - scrambled = list(reversed(files)) - monkeypatch.setattr(Path, "rglob", lambda self, pattern: iter(scrambled)) - - result = KonfAIInference._reassemble_output(output_dir) - - assert list(result.shape) == [3, 2, 2, 2] - assert float(result[0].mean()) == 0.0 - assert float(result[1].mean()) == 10.0 - assert float(result[2].mean()) == 20.0 - - -def test_konfai_inference_default_repo_and_model_preserved(): - """Constructing without arguments keeps the current published repo/model default.""" - transform = KonfAIInference() - - assert transform.repo_id == DEFAULT_INFERENCE_REPO_ID - assert transform.model_name == DEFAULT_INFERENCE_MODEL_NAME - assert transform.repo_id == "VBoussot/MRSegmentator-KonfAI" - assert transform.model_name == "MRSegmentator" - - -def test_konfai_inference_forwards_configured_repo_and_model(monkeypatch): - """A custom repo/model is forwarded verbatim to the KonfAIApp spec, not the default.""" - captured = {} - - class _FakeKonfAIApp: - def __init__(self, spec, *args): - captured["spec"] = spec - - def infer(self, *args, **kwargs): - captured["infer"] = (args, kwargs) - - fake_module = types.ModuleType("konfai_apps") - fake_module.KonfAIApp = _FakeKonfAIApp - monkeypatch.setitem(sys.modules, "konfai_apps", fake_module) - - transform = KonfAIInference( - repo_id="acme/Custom-KonfAI", - model_name="CustomModel", - checkpoints_name=["fold_1"], - ) - transform.infer_entry(Path("dataset"), Path("output"), []) - - assert captured["spec"] == "acme/Custom-KonfAI:CustomModel" diff --git a/tests/unit/test_transform_norm.py b/tests/unit/test_transform_norm.py deleted file mode 100644 index 7c87c260..00000000 --- a/tests/unit/test_transform_norm.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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 - -import numpy as np -import torch -from konfai.data.transform import Norm -from konfai.utils.dataset import Attribute - - -def _stack_attribute() -> Attribute: - # Geometry of a displacement-field stack: the leading image axis holds the vector components - # (origin 0 / spacing 1 / identity direction row), the remaining axes carry the fixed grid. - attribute = Attribute() - attribute["Origin"] = np.asarray([0.0, 1.0, 2.0, 3.0]) - attribute["Spacing"] = np.asarray([1.0, 2.0, 2.0, 2.0]) - direction = np.eye(4) - direction[1:, 1:] = np.diag([1.0, -1.0, 1.0]) - attribute["Direction"] = direction.flatten() - return attribute - - -def test_norm_reduces_trailing_component_axis_and_geometry() -> None: - # A stack of 2 displacement fields [N=2, D, H, W, C=3] -> per-sample magnitudes [2, D, H, W]. - tensors = torch.randn(2, 4, 5, 6, 3) - attribute = _stack_attribute() - - out = Norm()("case", tensors, attribute) - - assert list(out.shape) == [2, 4, 5, 6] - assert torch.allclose(out, torch.linalg.norm(tensors, dim=-1)) - # The reduced trailing tensor axis is the first geometry axis: it must be dropped. - assert attribute.get_np_array("Origin").tolist() == [1.0, 2.0, 3.0] - assert attribute.get_np_array("Spacing").tolist() == [2.0, 2.0, 2.0] - assert attribute.get_np_array("Direction").tolist() == np.diag([1.0, -1.0, 1.0]).flatten().tolist() - - -def test_norm_transform_shape_drops_trailing_axis() -> None: - assert Norm().transform_shape("group", "case", [4, 5, 6, 3], Attribute()) == [4, 5, 6] diff --git a/tests/unit/test_unet_attention.py b/tests/unit/test_unet_attention.py deleted file mode 100644 index 8785cddb..00000000 --- a/tests/unit/test_unet_attention.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 test: UNet(attention=True) must build and forward.""" - -import torch -from konfai.models.segmentation.UNet import UNet - - -def _forward_last(attention: bool) -> torch.Tensor: - net = UNet(dim=2, channels=[1, 8, 16], nb_class=2, attention=attention) - outputs = list(net.named_forward(torch.randn(1, 1, 32, 32))) - return outputs[-1][1] - - -def test_unet_attention_forwards_without_branch_collision() -> None: - # out_branch=[1] collided with the Attention block's internal W_g branch, so the parent captured - # a half-resolution projection and Concat crashed on a size mismatch. The gated skip must reach - # the skip connection, so the network forwards to a full-resolution output. - attended = _forward_last(attention=True) - plain = _forward_last(attention=False) - - assert attended.shape[-2:] == (32, 32) - assert attended.shape == plain.shape diff --git a/tests/unit/test_yaml_model_equivalence.py b/tests/unit/test_yaml_model_equivalence.py index d926a529..007fae28 100644 --- a/tests/unit/test_yaml_model_equivalence.py +++ b/tests/unit/test_yaml_model_equivalence.py @@ -22,16 +22,11 @@ hand-written ``konfai.models.segmentation.UNet`` configured identically. """ -import os +from pathlib import Path -os.environ.setdefault("KONFAI_config_file", "/tmp/konfai-none.yml") -os.environ.setdefault("KONFAI_CONFIG_MODE", "Done") - -from pathlib import Path # noqa: E402 - -import torch # noqa: E402 -from konfai.network.blocks import BlockConfig # noqa: E402 -from konfai.utils.model_builder import build_model_from_yaml # noqa: E402 +import torch +from konfai.network.blocks import BlockConfig +from konfai.utils.model_builder import build_model_from_yaml UNET_YML = Path(__file__).resolve().parents[2] / "examples" / "Segmentation" / "UNet.yml" @@ -53,9 +48,7 @@ def test_example_unet_yaml_matches_python_unet_param_count(): from konfai.models.segmentation.UNet import UNet yaml_net = _build_yaml_unet() - block_config = BlockConfig( - kernel_size=3, stride=1, padding=1, bias=True, activation="ReLU", norm_mode="NONE" - ) + block_config = BlockConfig(kernel_size=3, stride=1, padding=1, bias=True, activation="ReLU", norm_mode="NONE") python_net = UNet( dim=2, channels=[1, 32, 64, 128, 256], From 9b4e8d47d0e08fc5553d5142c7a3b59dfa4956b4 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Sun, 5 Jul 2026 19:26:37 +0200 Subject: [PATCH 2/4] test: merge trainer and runtime test files, finish suite lint cleanup - fold test_trainer_checkpoints/test_early_stopping/test_ema_convention/ test_resume_lr_override/test_resume_https_checkpoint into test_trainer.py and test_runtime_guards/test_runtime_progress_ddp into test_runtime.py, all tests moved verbatim (18 and 13 collected IDs unchanged) - fix the remaining ruff violations across tests/ (import order, unused noqa, ClassVar on mutable class defaults, zip strict, raw regex match) and reformat the flagged files - mark tests/integration with pytest.mark.integration so -m selection works --- tests/assets/Workflows/TinySynth.py | 1 - .../integration/test_konfai_core_workflows.py | 7 +- tests/unit/test_dataset_streaming.py | 2 +- tests/unit/test_early_stopping.py | 46 -- tests/unit/test_ema_convention.py | 25 -- tests/unit/test_imaging_roundtrip.py | 13 +- tests/unit/test_main_cli.py | 3 +- tests/unit/test_network.py | 3 +- tests/unit/test_package.py | 3 +- tests/unit/test_perf_hot_paths.py | 19 +- tests/unit/test_predictor_memory.py | 20 +- tests/unit/test_resume_https_checkpoint.py | 46 -- tests/unit/test_resume_lr_override.py | 100 ----- ...test_runtime_guards.py => test_runtime.py} | 99 ++++- tests/unit/test_runtime_progress_ddp.py | 99 ----- tests/unit/test_trainer.py | 415 ++++++++++++++++++ tests/unit/test_trainer_checkpoints.py | 194 -------- 17 files changed, 548 insertions(+), 547 deletions(-) delete mode 100644 tests/unit/test_early_stopping.py delete mode 100644 tests/unit/test_ema_convention.py delete mode 100644 tests/unit/test_resume_https_checkpoint.py delete mode 100644 tests/unit/test_resume_lr_override.py rename tests/unit/{test_runtime_guards.py => test_runtime.py} (64%) delete mode 100644 tests/unit/test_runtime_progress_ddp.py create mode 100644 tests/unit/test_trainer.py delete mode 100644 tests/unit/test_trainer_checkpoints.py diff --git a/tests/assets/Workflows/TinySynth.py b/tests/assets/Workflows/TinySynth.py index 97e6152d..38762ac2 100644 --- a/tests/assets/Workflows/TinySynth.py +++ b/tests/assets/Workflows/TinySynth.py @@ -1,5 +1,4 @@ import torch - from konfai.network import network diff --git a/tests/integration/test_konfai_core_workflows.py b/tests/integration/test_konfai_core_workflows.py index 57ebdf6a..f023ca75 100644 --- a/tests/integration/test_konfai_core_workflows.py +++ b/tests/integration/test_konfai_core_workflows.py @@ -9,11 +9,12 @@ import numpy as np import pytest - from konfai.evaluator import build_evaluate from konfai.predictor import build_predict from konfai.trainer import build_train +pytestmark = pytest.mark.integration + ASSETS_DIR = Path(__file__).resolve().parents[1] / "assets" / "Workflows" REPO_ROOT = Path(__file__).resolve().parents[2] SimpleITK = pytest.importorskip("SimpleITK") @@ -219,9 +220,7 @@ def main() -> None: if __name__ == "__main__": main() - """.replace( - "__TRAIN_NAME__", train_name - ) + """.replace("__TRAIN_NAME__", train_name) ), encoding="utf-8", ) diff --git a/tests/unit/test_dataset_streaming.py b/tests/unit/test_dataset_streaming.py index c9a16489..47c27367 100644 --- a/tests/unit/test_dataset_streaming.py +++ b/tests/unit/test_dataset_streaming.py @@ -444,7 +444,7 @@ class DaemonProcess: monkeypatch.setattr("konfai.data.transform.current_process", lambda: DaemonProcess()) - with pytest.raises(RuntimeError, match="Dataset.num_workers: 0"): + with pytest.raises(RuntimeError, match=r"Dataset\.num_workers: 0"): transform("CASE_000", torch.zeros(1, 4, 4), Attribute()) diff --git a/tests/unit/test_early_stopping.py b/tests/unit/test_early_stopping.py deleted file mode 100644 index 47f43b33..00000000 --- a/tests/unit/test_early_stopping.py +++ /dev/null @@ -1,46 +0,0 @@ -import pytest - -from konfai.trainer import EarlyStopping, EarlyStoppingBase -from konfai.utils.errors import TrainerError - - -def test_early_stopping_base_starts_running_and_can_be_stopped() -> None: - stopper = EarlyStoppingBase() - - assert stopper.is_stopped() is False - - stopper.stop() - - assert stopper.is_stopped() is True - - -def test_early_stopping_inherits_stop_from_base() -> None: - stopper = EarlyStopping(monitor=[], patience=10) - - assert stopper.is_stopped() is False - - stopper.stop() - - assert stopper.is_stopped() is True - - -def test_early_stopping_triggers_after_patience_without_improvement() -> None: - stopper = EarlyStopping(monitor=[], patience=2, mode="min") - - assert stopper(1.0) is False # first score sets the baseline - assert stopper(1.0) is False # no improvement (counter = 1) - assert stopper(1.0) is True # no improvement (counter = 2 >= patience) - assert stopper.is_stopped() is True - - -def test_get_score_reports_missing_metric_and_available_keys() -> None: - stopper = EarlyStopping(monitor=["val_loss"], patience=3) - - with pytest.raises(TrainerError) as exc_info: - stopper.get_score({"train_loss": 1.0, "dice": 0.5}) - - message = str(exc_info.value) - assert "val_loss" in message # the missing monitored metric is named - assert "train_loss" in message # the keys actually available are listed - assert "dice" in message - assert "{}" not in message # the placeholder is interpolated, not left raw diff --git a/tests/unit/test_ema_convention.py b/tests/unit/test_ema_convention.py deleted file mode 100644 index 4e954904..00000000 --- a/tests/unit/test_ema_convention.py +++ /dev/null @@ -1,25 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch -from konfai.trainer import Trainer - - -def test_avg_fn_follows_standard_ema_convention() -> None: - stub = SimpleNamespace(ema_decay=0.9) - averaged = torch.tensor(1.0) - model = torch.tensor(0.0) - - result = Trainer._avg_fn(stub, averaged, model, 0) - - assert result.item() == pytest.approx(0.9) - - -def test_avg_fn_high_decay_keeps_running_average_dominant() -> None: - stub = SimpleNamespace(ema_decay=0.999) - averaged = torch.tensor(10.0) - model = torch.tensor(0.0) - - result = Trainer._avg_fn(stub, averaged, model, 5) - - assert result.item() == pytest.approx(9.99) diff --git a/tests/unit/test_imaging_roundtrip.py b/tests/unit/test_imaging_roundtrip.py index d6c70fbe..cf7139c6 100644 --- a/tests/unit/test_imaging_roundtrip.py +++ b/tests/unit/test_imaging_roundtrip.py @@ -72,10 +72,13 @@ def test_dicom_left_handed_direction_normalizes_like_simpleitk(tmp_path: Path) - # distinct value per slice so ordering is observable vol = np.stack([np.full((4, 5), k, np.float32) for k in range(6)])[np.newaxis] dicom.write_dicom_series( - root, vol, origin=(0.0, 0.0, 30.0), spacing=(1.0, 1.0, 2.0), + root, + vol, + origin=(0.0, 0.0, 30.0), + spacing=(1.0, 1.0, 2.0), direction=np.array([1, 0, 0, 0, 1, 0, 0, 0, -1], float), ) - kvol, kog, ksp, kdir = dicom.read_dicom_series(root) + kvol, kog, _ksp, kdir = dicom.read_dicom_series(root) reader = sitk.ImageSeriesReader() ids = reader.GetGDCMSeriesIDs(str(root)) @@ -194,8 +197,10 @@ def test_ome_zarr_level_reads_coarser_resolution(tmp_path: Path) -> None: root.mkdir() data = (np.arange(1 * 16 * 32 * 32).reshape(1, 16, 32, 32) % 50).astype(np.float32) image = ngff_zarr.to_ngff_image( - data, dims=["c", "z", "y", "x"], - scale={"c": 1.0, "z": 2.0, "y": 0.5, "x": 0.5}, translation={"c": 0.0, "z": 0.0, "y": 0.0, "x": 0.0}, + data, + dims=["c", "z", "y", "x"], + scale={"c": 1.0, "z": 2.0, "y": 0.5, "x": 0.5}, + translation={"c": 0.0, "z": 0.0, "y": 0.0, "x": 0.0}, ) ngff_zarr.to_ngff_zarr( str(root / "CASE0.ome.zarr"), ngff_zarr.to_multiscales(image, scale_factors=[2]), overwrite=True, version="0.4" diff --git a/tests/unit/test_main_cli.py b/tests/unit/test_main_cli.py index 23c143d4..f8a7e851 100644 --- a/tests/unit/test_main_cli.py +++ b/tests/unit/test_main_cli.py @@ -21,12 +21,11 @@ import sys from pathlib import Path -import pytest - import konfai.evaluator as evaluator_module import konfai.main as main_module import konfai.predictor as predictor_module import konfai.trainer as trainer_module +import pytest def test_konfai_help_exits_zero(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/test_network.py b/tests/unit/test_network.py index fcb11501..ebece33b 100644 --- a/tests/unit/test_network.py +++ b/tests/unit/test_network.py @@ -22,11 +22,10 @@ from typing import cast from unittest.mock import MagicMock +import konfai.network.network as network_module import numpy as np import pytest import torch - -import konfai.network.network as network_module from konfai.metric.schedulers import Constant from konfai.network.blocks import Add from konfai.network.network import Measure, ModuleArgsDict, Network diff --git a/tests/unit/test_package.py b/tests/unit/test_package.py index 242fe151..56b9ed83 100644 --- a/tests/unit/test_package.py +++ b/tests/unit/test_package.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Smoke tests verifying package-level import contracts and public API surface.""" -import pytest - import konfai +import pytest from konfai.utils.errors import KonfAIError, TransformError diff --git a/tests/unit/test_perf_hot_paths.py b/tests/unit/test_perf_hot_paths.py index f519d5dc..f7eed8ab 100644 --- a/tests/unit/test_perf_hot_paths.py +++ b/tests/unit/test_perf_hot_paths.py @@ -19,9 +19,8 @@ from pathlib import Path from unittest.mock import MagicMock -import torch - import konfai.utils.dataset as dataset_module +import torch from konfai.data.patching import Accumulator from konfai.predictor import ModelComposite from konfai.utils.dataset import Attribute, Dataset @@ -143,7 +142,7 @@ def test_dicom_slice_info_threading_is_byte_identical_and_removes_rescans(tmp_pa 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): + for a, b in zip(ref, got, strict=False): assert np.array_equal(np.asarray(a), np.asarray(b)) assert len(info["sorted_files"]) == vol.shape[1] @@ -160,8 +159,14 @@ def test_dicom_slice_info_threading_is_byte_identical_and_removes_rescans(tmp_pa # (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]) + 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) @@ -203,7 +208,9 @@ def test_clip_clamp_fast_path_is_byte_identical_on_float32_and_safe_on_int(): 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) + 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 diff --git a/tests/unit/test_predictor_memory.py b/tests/unit/test_predictor_memory.py index 6b157d09..38e9d060 100644 --- a/tests/unit/test_predictor_memory.py +++ b/tests/unit/test_predictor_memory.py @@ -1,9 +1,8 @@ -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import pytest import torch - from konfai.data.data_manager import BatchDataItem, DatasetIter from konfai.network.network import Network from konfai.predictor import Mean, ModelComposite, OutSameAsGroupDataset, _Predictor @@ -11,7 +10,6 @@ class DummyPredictNetwork(Network): - def __init__(self) -> None: super().__init__(in_channels=1) self.scale = 1.0 @@ -60,7 +58,7 @@ def test_model_composite_streams_ensemble_through_a_single_loaded_model() -> Non def test_output_dataset_uses_batch_attributes_when_manager_cache_is_cold() -> None: class DummyPatch: - patch_size = [2, 2] + patch_size: ClassVar[list[int]] = [2, 2] @staticmethod def get_patch_slices(index_augmentation: int): @@ -70,13 +68,13 @@ def get_patch_slices(index_augmentation: int): class DummyManager: name = "CASE_000" patch = DummyPatch() - cache_attributes = [Attribute({"Origin": [0.0, 0.0]})] + cache_attributes: ClassVar[list[Attribute]] = [Attribute({"Origin": [0.0, 0.0]})] class DummyGroupTransform: - patch_transforms: list[object] = [] + patch_transforms: ClassVar[list[object]] = [] class DummyDatasetIter: - groups_src = {"src": {"dest": DummyGroupTransform()}} + groups_src: ClassVar[dict[str, dict[str, object]]] = {"src": {"dest": DummyGroupTransform()}} @staticmethod def get_dataset_from_index(group_dest: str, index: int): @@ -112,7 +110,7 @@ def get_dataset_from_index(group_dest: str, index: int): def test_output_dataset_offloads_patch_predictions_to_cpu_before_accumulating() -> None: class DummyPatch: - patch_size = [2, 2] + patch_size: ClassVar[list[int]] = [2, 2] @staticmethod def get_patch_slices(index_augmentation: int): @@ -122,13 +120,13 @@ def get_patch_slices(index_augmentation: int): class DummyManager: name = "CASE_000" patch = DummyPatch() - cache_attributes = [Attribute({"Origin": [0.0, 0.0]})] + cache_attributes: ClassVar[list[Attribute]] = [Attribute({"Origin": [0.0, 0.0]})] class DummyGroupTransform: - patch_transforms: list[object] = [] + patch_transforms: ClassVar[list[object]] = [] class DummyDatasetIter: - groups_src = {"src": {"dest": DummyGroupTransform()}} + groups_src: ClassVar[dict[str, dict[str, object]]] = {"src": {"dest": DummyGroupTransform()}} @staticmethod def get_dataset_from_index(group_dest: str, index: int): diff --git a/tests/unit/test_resume_https_checkpoint.py b/tests/unit/test_resume_https_checkpoint.py deleted file mode 100644 index 2895878c..00000000 --- a/tests/unit/test_resume_https_checkpoint.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 test: RESUME must keep an https:// checkpoint URL intact through build_train.""" - -from konfai.utils.runtime import State - - -def test_build_train_keeps_https_checkpoint_url(monkeypatch) -> None: - # build_train must not wrap an https:// URL in Path(): that collapses '//' into 'https:/…', - # which then fails both the startswith('https://') check and Path.exists() at load time. - import konfai.trainer as trainer_module - - recorded: dict[str, object] = {} - - class _DummyTrainer: - def set_model(self, path_to_model) -> None: - recorded["model"] = path_to_model - - def set_lr(self, lr) -> None: - recorded["lr"] = lr - - monkeypatch.setattr(trainer_module, "configure_workflow_environment", lambda **kwargs: None) - monkeypatch.setattr( - trainer_module, - "apply_config", - lambda *args, **kwargs: lambda cls: lambda: _DummyTrainer(), - ) - - url = "https://example.com/weights/ckpt.pt" - trainer_module.build_train(command=State.RESUME, model=url) - - assert recorded["model"] == url diff --git a/tests/unit/test_resume_lr_override.py b/tests/unit/test_resume_lr_override.py deleted file mode 100644 index dd4481ef..00000000 --- a/tests/unit/test_resume_lr_override.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Resume/fine-tune learning-rate override semantics for ``Network.load``. - -Without ``override_lr`` a resume must keep the checkpoint (decayed) learning rate and -let the scheduler continue from ``_nb_lr_update``. With ``override_lr`` the learning -rate must restart from the requested value and the scheduler must decay from there. -""" - -import torch -from konfai.metric.schedulers import PolyLRScheduler -from konfai.network.network import Network - -_CONFIG_LR = 0.1 -_GAMMA = 0.5 -_NB_LR_UPDATE = 3 - - -class _LeafNet(Network): - """Minimal concrete network with no sub-networks, driving ``Network.load`` directly.""" - - def __init__(self) -> None: - super().__init__() - - -def _fresh_optimizer() -> torch.optim.Optimizer: - param = torch.nn.Parameter(torch.zeros(1)) - return torch.optim.SGD([param], lr=_CONFIG_LR) - - -def _decayed_optimizer_state() -> tuple[dict, float]: - """Optimizer state as saved by a checkpoint after ``_NB_LR_UPDATE`` StepLR decays.""" - optimizer = _fresh_optimizer() - scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=_GAMMA) - for _ in range(_NB_LR_UPDATE): - scheduler.step() - return optimizer.state_dict(), optimizer.param_groups[0]["lr"] - - -def _make_net(scheduler_factory) -> tuple[_LeafNet, torch.optim.lr_scheduler.LRScheduler, dict]: - net = _LeafNet() - optimizer = _fresh_optimizer() - scheduler = scheduler_factory(optimizer) - net.optimizer = optimizer - net.schedulers = {scheduler: 0} - net._it = 0 - net._nb_lr_update = 0 - optimizer_state, decayed_lr = _decayed_optimizer_state() - state_dict = { - f"{net.get_name()}_optimizer_state_dict": optimizer_state, - f"{net.get_name()}_nb_lr_update": _NB_LR_UPDATE, - } - return net, scheduler, {"state_dict": state_dict, "decayed_lr": decayed_lr} - - -def test_resume_without_override_keeps_decayed_lr_and_restores_scheduler() -> None: - net, scheduler, ctx = _make_net(lambda opt: torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=_GAMMA)) - - net.load(ctx["state_dict"], init=False, ema=False) - - # The decayed learning rate from the checkpoint is preserved (not reset to the config LR). - assert net.optimizer.param_groups[0]["lr"] == ctx["decayed_lr"] - assert net.optimizer.param_groups[0]["lr"] != _CONFIG_LR - # The scheduler continues from where it left off instead of restarting at 0. - assert scheduler.last_epoch == _NB_LR_UPDATE - - -def test_resume_with_override_restarts_lr_and_scheduler() -> None: - override = 0.02 - net, scheduler, ctx = _make_net(lambda opt: torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=_GAMMA)) - - net.load(ctx["state_dict"], init=False, ema=False, override_lr=override) - - # The learning rate is forced to the override, and the schedule restarts from it. - assert net.optimizer.param_groups[0]["lr"] == override - assert net.optimizer.param_groups[0]["initial_lr"] == override - assert scheduler.base_lrs == [override] - assert scheduler.last_epoch == 0 - - # Decaying from the override reproduces a fresh schedule anchored at ``override``. - scheduler.step() - assert net.optimizer.param_groups[0]["lr"] == override * _GAMMA - - -def test_resume_with_override_restarts_polylr_from_value() -> None: - override = 0.05 - max_steps = 100 - exponent = 0.9 - net, scheduler, ctx = _make_net( - lambda opt: PolyLRScheduler(opt, initial_lr=_CONFIG_LR, max_steps=max_steps, exponent=exponent) - ) - - net.load(ctx["state_dict"], init=False, ema=False, override_lr=override) - - assert net.optimizer.param_groups[0]["lr"] == override - assert scheduler.initial_lr == override - assert scheduler.last_epoch == 0 - - scheduler.step() - assert net.optimizer.param_groups[0]["lr"] == override * (1 - 0 / max_steps) ** exponent - scheduler.step() - assert net.optimizer.param_groups[0]["lr"] == override * (1 - 1 / max_steps) ** exponent diff --git a/tests/unit/test_runtime_guards.py b/tests/unit/test_runtime.py similarity index 64% rename from tests/unit/test_runtime_guards.py rename to tests/unit/test_runtime.py index 15eb6c1b..7ec5215e 100644 --- a/tests/unit/test_runtime_guards.py +++ b/tests/unit/test_runtime.py @@ -14,17 +14,20 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Guard tests for ``konfai.utils.runtime``: workflow preconditions, environment -normalisation, overwrite confirmation, and distributed-launch bookkeeping.""" +"""Tests for ``konfai.utils.runtime``: workflow guards, environment normalisation, +overwrite confirmation, distributed-launch bookkeeping, and progress/DDP +synchronisation.""" +import contextlib import os +import random import sys from pathlib import Path from types import SimpleNamespace -import pytest - import konfai as konfai_module +import konfai.utils.runtime as rt +import pytest from konfai.evaluator import Evaluator from konfai.predictor import Predictor from konfai.trainer import Trainer @@ -37,6 +40,11 @@ execute_distributed_object, ) +# --------------------------------------------------------------------------- +# Workflow guards, environment normalisation, overwrite confirmation, and +# distributed-launch bookkeeping +# --------------------------------------------------------------------------- + @pytest.mark.parametrize("factory", [Trainer, Predictor, Evaluator]) def test_core_workflows_raise_config_error_when_mode_is_not_done( @@ -164,3 +172,86 @@ def fake_get_device_name(index: int) -> str: assert devices_index == [3, 5] assert devices_name == ["GPU0", "GPU1"] assert queried_indices == [0, 1] + + +# --------------------------------------------------------------------------- +# Progress/DDP synchronisation (regression tests for the runtime audit fixes, +# see AUDIT.md) +# --------------------------------------------------------------------------- + + +def test_synchronize_data_gathers_on_cpu(monkeypatch): + """gloo/CPU multi-process must still all_gather (not fall back to local rank).""" + calls = {} + + def fake_all_gather_object(outputs, data): + calls["called"] = True + for i in range(len(outputs)): + outputs[i] = data + + def fail_set_device(*_args, **_kwargs): + raise AssertionError("set_device must not be called when CUDA is unavailable") + + monkeypatch.setattr(rt.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rt.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(rt.torch.cuda, "set_device", fail_set_device) + monkeypatch.setattr(rt.dist, "all_gather_object", fake_all_gather_object) + + result = rt.synchronize_data(3, 0, {"a": 1}) + + assert calls.get("called") is True + assert result == [{"a": 1}, {"a": 1}, {"a": 1}] + + +def test_synchronize_data_sets_device_on_cuda(monkeypatch): + """When CUDA is available the target device is selected before gathering.""" + seen = {} + + def fake_all_gather_object(outputs, data): + for i in range(len(outputs)): + outputs[i] = data + + monkeypatch.setattr(rt.dist, "is_initialized", lambda: True) + monkeypatch.setattr(rt.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(rt.torch.cuda, "set_device", lambda gpu: seen.setdefault("gpu", gpu)) + monkeypatch.setattr(rt.dist, "all_gather_object", fake_all_gather_object) + + result = rt.synchronize_data(2, 1, {"b": 2}) + + assert seen.get("gpu") == 1 + assert result == [{"b": 2}, {"b": 2}] + + +def test_synchronize_data_no_dist(monkeypatch): + """Without an active process group the local data is returned as-is.""" + monkeypatch.setattr(rt.dist, "is_initialized", lambda: False) + assert rt.synchronize_data(4, 0, {"a": 1}) == [{"a": 1}] + + +def _run_execute(monkeypatch, obj): + monkeypatch.setattr(rt, "Log", lambda *a, **k: contextlib.nullcontext()) + monkeypatch.setattr(rt, "TensorBoard", lambda *a, **k: contextlib.nullcontext()) + monkeypatch.setattr(rt.mp, "spawn", lambda *a, **k: None) + rt.execute_distributed_object(obj, gpu=None, cpu=1) + + +def test_execute_seeds_parent_before_setup(monkeypatch): + """The parent process (which runs the train/val split) must be seeded.""" + + recorded = [] + + class FakeObject(rt.DistributedObject): + def __init__(self) -> None: + super().__init__("fake-seeded") + self.manual_seed = 123 + + def setup(self, world_size: int) -> None: + recorded.append(random.random()) + + def run_process(self, *args, **kwargs) -> None: # pragma: no cover - not spawned + pass + + _run_execute(monkeypatch, FakeObject()) + _run_execute(monkeypatch, FakeObject()) + + assert recorded[0] == recorded[1] diff --git a/tests/unit/test_runtime_progress_ddp.py b/tests/unit/test_runtime_progress_ddp.py deleted file mode 100644 index 6ea1e8aa..00000000 --- a/tests/unit/test_runtime_progress_ddp.py +++ /dev/null @@ -1,99 +0,0 @@ -# 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 runtime progress/DDP audit fixes (see AUDIT.md).""" - -import contextlib -import random - -import konfai.utils.runtime as rt - - -def test_synchronize_data_gathers_on_cpu(monkeypatch): - """gloo/CPU multi-process must still all_gather (not fall back to local rank).""" - calls = {} - - def fake_all_gather_object(outputs, data): - calls["called"] = True - for i in range(len(outputs)): - outputs[i] = data - - def fail_set_device(*_args, **_kwargs): - raise AssertionError("set_device must not be called when CUDA is unavailable") - - monkeypatch.setattr(rt.dist, "is_initialized", lambda: True) - monkeypatch.setattr(rt.torch.cuda, "is_available", lambda: False) - monkeypatch.setattr(rt.torch.cuda, "set_device", fail_set_device) - monkeypatch.setattr(rt.dist, "all_gather_object", fake_all_gather_object) - - result = rt.synchronize_data(3, 0, {"a": 1}) - - assert calls.get("called") is True - assert result == [{"a": 1}, {"a": 1}, {"a": 1}] - - -def test_synchronize_data_sets_device_on_cuda(monkeypatch): - """When CUDA is available the target device is selected before gathering.""" - seen = {} - - def fake_all_gather_object(outputs, data): - for i in range(len(outputs)): - outputs[i] = data - - monkeypatch.setattr(rt.dist, "is_initialized", lambda: True) - monkeypatch.setattr(rt.torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(rt.torch.cuda, "set_device", lambda gpu: seen.setdefault("gpu", gpu)) - monkeypatch.setattr(rt.dist, "all_gather_object", fake_all_gather_object) - - result = rt.synchronize_data(2, 1, {"b": 2}) - - assert seen.get("gpu") == 1 - assert result == [{"b": 2}, {"b": 2}] - - -def test_synchronize_data_no_dist(monkeypatch): - """Without an active process group the local data is returned as-is.""" - monkeypatch.setattr(rt.dist, "is_initialized", lambda: False) - assert rt.synchronize_data(4, 0, {"a": 1}) == [{"a": 1}] - - -def _run_execute(monkeypatch, obj): - monkeypatch.setattr(rt, "Log", lambda *a, **k: contextlib.nullcontext()) - monkeypatch.setattr(rt, "TensorBoard", lambda *a, **k: contextlib.nullcontext()) - monkeypatch.setattr(rt.mp, "spawn", lambda *a, **k: None) - rt.execute_distributed_object(obj, gpu=None, cpu=1) - - -def test_execute_seeds_parent_before_setup(monkeypatch): - """The parent process (which runs the train/val split) must be seeded.""" - - recorded = [] - - class FakeObject(rt.DistributedObject): - def __init__(self) -> None: - super().__init__("fake-seeded") - self.manual_seed = 123 - - def setup(self, world_size: int) -> None: - recorded.append(random.random()) - - def run_process(self, *args, **kwargs) -> None: # pragma: no cover - not spawned - pass - - _run_execute(monkeypatch, FakeObject()) - _run_execute(monkeypatch, FakeObject()) - - assert recorded[0] == recorded[1] diff --git a/tests/unit/test_trainer.py b/tests/unit/test_trainer.py new file mode 100644 index 00000000..65e6fa55 --- /dev/null +++ b/tests/unit/test_trainer.py @@ -0,0 +1,415 @@ +# 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 + +"""Tests for konfai.trainer: checkpoint save/bootstrap, early stopping, EMA, and RESUME +learning-rate/checkpoint handling.""" + +from collections.abc import Iterator +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import konfai.trainer as trainer_module +import pytest +import torch +from konfai.metric.schedulers import PolyLRScheduler +from konfai.network.network import Network +from konfai.trainer import EarlyStopping, EarlyStoppingBase, Trainer, _Trainer +from konfai.utils.errors import TrainerError +from konfai.utils.runtime import State +from torch import nn +from torch.optim.swa_utils import AveragedModel + +# ---- Checkpoints ---- + + +class _DummySummaryWriter: + def __init__(self, *args, **kwargs) -> None: + pass + + def close(self) -> None: + pass + + +class _DummyModelModule: + @staticmethod + def state_dict() -> dict[str, torch.Tensor]: + return {"weight": torch.tensor([1.0])} + + @staticmethod + def get_networks() -> dict[str, object]: + return {} + + +class _DummyModel: + def __init__(self) -> None: + self.module = _DummyModelModule() + + +def _date_sequence(values: list[str]) -> Iterator[str]: + yield from values + while True: + yield values[-1] + + +def _build_trainer(tmp_path: Path, monkeypatch, date_values: list[str]) -> _Trainer: + checkpoints_dir = tmp_path / "Checkpoints" + statistics_dir = tmp_path / "Statistics" + date_iter = _date_sequence(date_values) + + monkeypatch.setattr(trainer_module, "checkpoints_directory", lambda: checkpoints_dir) + monkeypatch.setattr(trainer_module, "statistics_directory", lambda: statistics_dir) + monkeypatch.setattr(trainer_module, "SummaryWriter", _DummySummaryWriter) + monkeypatch.setattr(trainer_module, "current_date", lambda: next(date_iter)) + + return _Trainer( + world_size=1, + global_rank=0, + local_rank=0, + size=1, + train_name="RUN", + early_stopping=None, + data_log=None, + save_checkpoint_mode="BEST", + epochs=1, + epoch=0, + autocast=False, + it_validation=1, + it_lr_update=1, + it=0, + model=cast(Any, _DummyModel()), + model_ema=None, + dataloader_training=[object()], + dataloader_validation=None, + ) + + +def test_best_checkpoint_save_keeps_only_best_without_rescanning(tmp_path: Path, monkeypatch) -> None: + trainer = _build_trainer(tmp_path, monkeypatch, ["ckpt_a", "ckpt_b", "ckpt_c"]) + original_load = torch.load + + def fail_if_reloaded(*args, **kwargs): + raise AssertionError("BEST checkpoint save unexpectedly rescanned saved checkpoints") + + monkeypatch.setattr(trainer_module.torch, "load", fail_if_reloaded) + + trainer.checkpoint_save(2.0) + trainer.checkpoint_save(1.0) + trainer.checkpoint_save(3.0) + + checkpoints = sorted((tmp_path / "Checkpoints" / "RUN").glob("*.pt")) + assert [path.name for path in checkpoints] == ["ckpt_b.pt"] + assert original_load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 1.0 + + +def test_best_checkpoint_bootstrap_scans_existing_files_once_and_prunes_stale_ones( + tmp_path: Path, + monkeypatch, +) -> None: + checkpoint_dir = tmp_path / "Checkpoints" / "RUN" + checkpoint_dir.mkdir(parents=True) + torch.save({"loss": 5.0}, checkpoint_dir / "old_a.pt") + torch.save({"loss": 3.0}, checkpoint_dir / "old_b.pt") + + original_load = trainer_module.torch.load + load_calls: list[Path] = [] + + def counted_load(path, *args, **kwargs): + load_calls.append(Path(path)) + return original_load(path, *args, **kwargs) + + monkeypatch.setattr(trainer_module.torch, "load", counted_load) + + trainer = _build_trainer(tmp_path, monkeypatch, ["ckpt_new_worse", "ckpt_new_best"]) + + assert [path.name for path in sorted(checkpoint_dir.glob("*.pt"))] == ["old_b.pt"] + assert [path.name for path in load_calls] == ["old_a.pt", "old_b.pt"] + + trainer.checkpoint_save(4.0) + trainer.checkpoint_save(2.0) + + assert [path.name for path in load_calls] == ["old_a.pt", "old_b.pt"] + checkpoints = sorted(checkpoint_dir.glob("*.pt")) + assert [path.name for path in checkpoints] == ["ckpt_new_best.pt"] + assert original_load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 2.0 + + +def test_best_checkpoint_survives_same_second_collision(tmp_path: Path, monkeypatch) -> None: + trainer = _build_trainer(tmp_path, monkeypatch, ["same_stamp", "same_stamp"]) + + trainer.checkpoint_save(1.0) # best + trainer.checkpoint_save(2.0) # worse, produced within the same timestamp + + checkpoints = sorted((tmp_path / "Checkpoints" / "RUN").glob("*.pt")) + assert len(checkpoints) == 1 + assert torch.load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 1.0 + + +def test_exit_checkpoint_loss_does_not_poison_best(tmp_path: Path, monkeypatch) -> None: + trainer = _build_trainer(tmp_path, monkeypatch, ["exit_stamp"]) + + trainer.checkpoint_save(None) # the save emitted on context exit + + saved = torch.load( + tmp_path / "Checkpoints" / "RUN" / "exit_stamp.pt", + map_location="cpu", + weights_only=False, + ) + assert saved["loss"] == float("inf") + + +def test_bootstrap_prefers_real_best_over_exit_checkpoint(tmp_path: Path, monkeypatch) -> None: + checkpoint_dir = tmp_path / "Checkpoints" / "RUN" + checkpoint_dir.mkdir(parents=True) + torch.save({"loss": 3.0}, checkpoint_dir / "real_best.pt") + torch.save({"loss": float("inf")}, checkpoint_dir / "exit.pt") + + trainer = _build_trainer(tmp_path, monkeypatch, ["new_stamp"]) + + assert [path.name for path in sorted(checkpoint_dir.glob("*.pt"))] == ["real_best.pt"] + assert trainer._best_checkpoint_loss == 3.0 + + +def test_checkpoint_persists_ema_n_averaged(tmp_path: Path, monkeypatch) -> None: + base = nn.Linear(2, 2) + ema = AveragedModel(base) + ema.update_parameters(base) + ema.update_parameters(base) + + trainer = _build_trainer(tmp_path, monkeypatch, ["ema_stamp"]) + trainer.model_ema = cast(Any, ema) + + trainer.checkpoint_save(1.0) + + saved = torch.load( + tmp_path / "Checkpoints" / "RUN" / "ema_stamp.pt", + map_location="cpu", + weights_only=False, + ) + assert "Model_EMA" in saved + assert saved["Model_EMA_n_averaged"] == int(ema.n_averaged) == 2 + + +def test_broadcast_stop_returns_local_value_without_distributed(tmp_path: Path, monkeypatch) -> None: + trainer = _build_trainer(tmp_path, monkeypatch, ["stamp"]) + + assert trainer._broadcast_stop(True) is True + assert trainer._broadcast_stop(False) is False + + +def test_broadcast_stop_adopts_rank_zero_decision(tmp_path: Path, monkeypatch) -> None: + trainer = _build_trainer(tmp_path, monkeypatch, ["stamp"]) + + monkeypatch.setattr(trainer_module, "synchronize_data", lambda *_args, **_kwargs: [True, False, False]) + assert trainer._broadcast_stop(False) is True # a non-zero rank still stops when rank 0 does + + monkeypatch.setattr(trainer_module, "synchronize_data", lambda *_args, **_kwargs: [False, True]) + assert trainer._broadcast_stop(True) is False # a non-zero rank keeps going when rank 0 does + + +# ---- EarlyStopping ---- + + +def test_early_stopping_base_starts_running_and_can_be_stopped() -> None: + stopper = EarlyStoppingBase() + + assert stopper.is_stopped() is False + + stopper.stop() + + assert stopper.is_stopped() is True + + +def test_early_stopping_inherits_stop_from_base() -> None: + stopper = EarlyStopping(monitor=[], patience=10) + + assert stopper.is_stopped() is False + + stopper.stop() + + assert stopper.is_stopped() is True + + +def test_early_stopping_triggers_after_patience_without_improvement() -> None: + stopper = EarlyStopping(monitor=[], patience=2, mode="min") + + assert stopper(1.0) is False # first score sets the baseline + assert stopper(1.0) is False # no improvement (counter = 1) + assert stopper(1.0) is True # no improvement (counter = 2 >= patience) + assert stopper.is_stopped() is True + + +def test_get_score_reports_missing_metric_and_available_keys() -> None: + stopper = EarlyStopping(monitor=["val_loss"], patience=3) + + with pytest.raises(TrainerError) as exc_info: + stopper.get_score({"train_loss": 1.0, "dice": 0.5}) + + message = str(exc_info.value) + assert "val_loss" in message # the missing monitored metric is named + assert "train_loss" in message # the keys actually available are listed + assert "dice" in message + assert "{}" not in message # the placeholder is interpolated, not left raw + + +# ---- EMA ---- + + +def test_avg_fn_follows_standard_ema_convention() -> None: + stub = SimpleNamespace(ema_decay=0.9) + averaged = torch.tensor(1.0) + model = torch.tensor(0.0) + + result = Trainer._avg_fn(stub, averaged, model, 0) + + assert result.item() == pytest.approx(0.9) + + +def test_avg_fn_high_decay_keeps_running_average_dominant() -> None: + stub = SimpleNamespace(ema_decay=0.999) + averaged = torch.tensor(10.0) + model = torch.tensor(0.0) + + result = Trainer._avg_fn(stub, averaged, model, 5) + + assert result.item() == pytest.approx(9.99) + + +# ---- RESUME LR override ---- + +# Resume/fine-tune learning-rate override semantics for ``Network.load``. +# +# Without ``override_lr`` a resume must keep the checkpoint (decayed) learning rate and +# let the scheduler continue from ``_nb_lr_update``. With ``override_lr`` the learning +# rate must restart from the requested value and the scheduler must decay from there. + +_CONFIG_LR = 0.1 +_GAMMA = 0.5 +_NB_LR_UPDATE = 3 + + +class _LeafNet(Network): + """Minimal concrete network with no sub-networks, driving ``Network.load`` directly.""" + + def __init__(self) -> None: + super().__init__() + + +def _fresh_optimizer() -> torch.optim.Optimizer: + param = torch.nn.Parameter(torch.zeros(1)) + return torch.optim.SGD([param], lr=_CONFIG_LR) + + +def _decayed_optimizer_state() -> tuple[dict, float]: + """Optimizer state as saved by a checkpoint after ``_NB_LR_UPDATE`` StepLR decays.""" + optimizer = _fresh_optimizer() + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=_GAMMA) + for _ in range(_NB_LR_UPDATE): + scheduler.step() + return optimizer.state_dict(), optimizer.param_groups[0]["lr"] + + +def _make_net(scheduler_factory) -> tuple[_LeafNet, torch.optim.lr_scheduler.LRScheduler, dict]: + net = _LeafNet() + optimizer = _fresh_optimizer() + scheduler = scheduler_factory(optimizer) + net.optimizer = optimizer + net.schedulers = {scheduler: 0} + net._it = 0 + net._nb_lr_update = 0 + optimizer_state, decayed_lr = _decayed_optimizer_state() + state_dict = { + f"{net.get_name()}_optimizer_state_dict": optimizer_state, + f"{net.get_name()}_nb_lr_update": _NB_LR_UPDATE, + } + return net, scheduler, {"state_dict": state_dict, "decayed_lr": decayed_lr} + + +def test_resume_without_override_keeps_decayed_lr_and_restores_scheduler() -> None: + net, scheduler, ctx = _make_net(lambda opt: torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=_GAMMA)) + + net.load(ctx["state_dict"], init=False, ema=False) + + # The decayed learning rate from the checkpoint is preserved (not reset to the config LR). + assert net.optimizer.param_groups[0]["lr"] == ctx["decayed_lr"] + assert net.optimizer.param_groups[0]["lr"] != _CONFIG_LR + # The scheduler continues from where it left off instead of restarting at 0. + assert scheduler.last_epoch == _NB_LR_UPDATE + + +def test_resume_with_override_restarts_lr_and_scheduler() -> None: + override = 0.02 + net, scheduler, ctx = _make_net(lambda opt: torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=_GAMMA)) + + net.load(ctx["state_dict"], init=False, ema=False, override_lr=override) + + # The learning rate is forced to the override, and the schedule restarts from it. + assert net.optimizer.param_groups[0]["lr"] == override + assert net.optimizer.param_groups[0]["initial_lr"] == override + assert scheduler.base_lrs == [override] + assert scheduler.last_epoch == 0 + + # Decaying from the override reproduces a fresh schedule anchored at ``override``. + scheduler.step() + assert net.optimizer.param_groups[0]["lr"] == override * _GAMMA + + +def test_resume_with_override_restarts_polylr_from_value() -> None: + override = 0.05 + max_steps = 100 + exponent = 0.9 + net, scheduler, ctx = _make_net( + lambda opt: PolyLRScheduler(opt, initial_lr=_CONFIG_LR, max_steps=max_steps, exponent=exponent) + ) + + net.load(ctx["state_dict"], init=False, ema=False, override_lr=override) + + assert net.optimizer.param_groups[0]["lr"] == override + assert scheduler.initial_lr == override + assert scheduler.last_epoch == 0 + + scheduler.step() + assert net.optimizer.param_groups[0]["lr"] == override * (1 - 0 / max_steps) ** exponent + scheduler.step() + assert net.optimizer.param_groups[0]["lr"] == override * (1 - 1 / max_steps) ** exponent + + +# ---- RESUME checkpoint URL ---- + + +def test_build_train_keeps_https_checkpoint_url(monkeypatch) -> None: + # build_train must not wrap an https:// URL in Path(): that collapses '//' into 'https:/…', + # which then fails both the startswith('https://') check and Path.exists() at load time. + recorded: dict[str, object] = {} + + class _DummyTrainer: + def set_model(self, path_to_model) -> None: + recorded["model"] = path_to_model + + def set_lr(self, lr) -> None: + recorded["lr"] = lr + + monkeypatch.setattr(trainer_module, "configure_workflow_environment", lambda **kwargs: None) + monkeypatch.setattr( + trainer_module, + "apply_config", + lambda *args, **kwargs: lambda cls: lambda: _DummyTrainer(), + ) + + url = "https://example.com/weights/ckpt.pt" + trainer_module.build_train(command=State.RESUME, model=url) + + assert recorded["model"] == url diff --git a/tests/unit/test_trainer_checkpoints.py b/tests/unit/test_trainer_checkpoints.py deleted file mode 100644 index d84530b8..00000000 --- a/tests/unit/test_trainer_checkpoints.py +++ /dev/null @@ -1,194 +0,0 @@ -from collections.abc import Iterator -from pathlib import Path -from typing import Any, cast - -import torch -from torch import nn -from torch.optim.swa_utils import AveragedModel - -import konfai.trainer as trainer_module -from konfai.trainer import _Trainer - - -class _DummySummaryWriter: - def __init__(self, *args, **kwargs) -> None: - pass - - def close(self) -> None: - pass - - -class _DummyModelModule: - @staticmethod - def state_dict() -> dict[str, torch.Tensor]: - return {"weight": torch.tensor([1.0])} - - @staticmethod - def get_networks() -> dict[str, object]: - return {} - - -class _DummyModel: - def __init__(self) -> None: - self.module = _DummyModelModule() - - -def _date_sequence(values: list[str]) -> Iterator[str]: - yield from values - while True: - yield values[-1] - - -def _build_trainer(tmp_path: Path, monkeypatch, date_values: list[str]) -> _Trainer: - checkpoints_dir = tmp_path / "Checkpoints" - statistics_dir = tmp_path / "Statistics" - date_iter = _date_sequence(date_values) - - monkeypatch.setattr(trainer_module, "checkpoints_directory", lambda: checkpoints_dir) - monkeypatch.setattr(trainer_module, "statistics_directory", lambda: statistics_dir) - monkeypatch.setattr(trainer_module, "SummaryWriter", _DummySummaryWriter) - monkeypatch.setattr(trainer_module, "current_date", lambda: next(date_iter)) - - return _Trainer( - world_size=1, - global_rank=0, - local_rank=0, - size=1, - train_name="RUN", - early_stopping=None, - data_log=None, - save_checkpoint_mode="BEST", - epochs=1, - epoch=0, - autocast=False, - it_validation=1, - it_lr_update=1, - it=0, - model=cast(Any, _DummyModel()), - model_ema=None, - dataloader_training=[object()], - dataloader_validation=None, - ) - - -def test_best_checkpoint_save_keeps_only_best_without_rescanning(tmp_path: Path, monkeypatch) -> None: - trainer = _build_trainer(tmp_path, monkeypatch, ["ckpt_a", "ckpt_b", "ckpt_c"]) - original_load = torch.load - - def fail_if_reloaded(*args, **kwargs): - raise AssertionError("BEST checkpoint save unexpectedly rescanned saved checkpoints") - - monkeypatch.setattr(trainer_module.torch, "load", fail_if_reloaded) - - trainer.checkpoint_save(2.0) - trainer.checkpoint_save(1.0) - trainer.checkpoint_save(3.0) - - checkpoints = sorted((tmp_path / "Checkpoints" / "RUN").glob("*.pt")) - assert [path.name for path in checkpoints] == ["ckpt_b.pt"] - assert original_load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 1.0 - - -def test_best_checkpoint_bootstrap_scans_existing_files_once_and_prunes_stale_ones( - tmp_path: Path, - monkeypatch, -) -> None: - checkpoint_dir = tmp_path / "Checkpoints" / "RUN" - checkpoint_dir.mkdir(parents=True) - torch.save({"loss": 5.0}, checkpoint_dir / "old_a.pt") - torch.save({"loss": 3.0}, checkpoint_dir / "old_b.pt") - - original_load = trainer_module.torch.load - load_calls: list[Path] = [] - - def counted_load(path, *args, **kwargs): - load_calls.append(Path(path)) - return original_load(path, *args, **kwargs) - - monkeypatch.setattr(trainer_module.torch, "load", counted_load) - - trainer = _build_trainer(tmp_path, monkeypatch, ["ckpt_new_worse", "ckpt_new_best"]) - - assert [path.name for path in sorted(checkpoint_dir.glob("*.pt"))] == ["old_b.pt"] - assert [path.name for path in load_calls] == ["old_a.pt", "old_b.pt"] - - trainer.checkpoint_save(4.0) - trainer.checkpoint_save(2.0) - - assert [path.name for path in load_calls] == ["old_a.pt", "old_b.pt"] - checkpoints = sorted(checkpoint_dir.glob("*.pt")) - assert [path.name for path in checkpoints] == ["ckpt_new_best.pt"] - assert original_load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 2.0 - - -def test_best_checkpoint_survives_same_second_collision(tmp_path: Path, monkeypatch) -> None: - trainer = _build_trainer(tmp_path, monkeypatch, ["same_stamp", "same_stamp"]) - - trainer.checkpoint_save(1.0) # best - trainer.checkpoint_save(2.0) # worse, produced within the same timestamp - - checkpoints = sorted((tmp_path / "Checkpoints" / "RUN").glob("*.pt")) - assert len(checkpoints) == 1 - assert torch.load(checkpoints[0], map_location="cpu", weights_only=False)["loss"] == 1.0 - - -def test_exit_checkpoint_loss_does_not_poison_best(tmp_path: Path, monkeypatch) -> None: - trainer = _build_trainer(tmp_path, monkeypatch, ["exit_stamp"]) - - trainer.checkpoint_save(None) # the save emitted on context exit - - saved = torch.load( - tmp_path / "Checkpoints" / "RUN" / "exit_stamp.pt", - map_location="cpu", - weights_only=False, - ) - assert saved["loss"] == float("inf") - - -def test_bootstrap_prefers_real_best_over_exit_checkpoint(tmp_path: Path, monkeypatch) -> None: - checkpoint_dir = tmp_path / "Checkpoints" / "RUN" - checkpoint_dir.mkdir(parents=True) - torch.save({"loss": 3.0}, checkpoint_dir / "real_best.pt") - torch.save({"loss": float("inf")}, checkpoint_dir / "exit.pt") - - trainer = _build_trainer(tmp_path, monkeypatch, ["new_stamp"]) - - assert [path.name for path in sorted(checkpoint_dir.glob("*.pt"))] == ["real_best.pt"] - assert trainer._best_checkpoint_loss == 3.0 - - -def test_checkpoint_persists_ema_n_averaged(tmp_path: Path, monkeypatch) -> None: - base = nn.Linear(2, 2) - ema = AveragedModel(base) - ema.update_parameters(base) - ema.update_parameters(base) - - trainer = _build_trainer(tmp_path, monkeypatch, ["ema_stamp"]) - trainer.model_ema = cast(Any, ema) - - trainer.checkpoint_save(1.0) - - saved = torch.load( - tmp_path / "Checkpoints" / "RUN" / "ema_stamp.pt", - map_location="cpu", - weights_only=False, - ) - assert "Model_EMA" in saved - assert saved["Model_EMA_n_averaged"] == int(ema.n_averaged) == 2 - - -def test_broadcast_stop_returns_local_value_without_distributed(tmp_path: Path, monkeypatch) -> None: - trainer = _build_trainer(tmp_path, monkeypatch, ["stamp"]) - - assert trainer._broadcast_stop(True) is True - assert trainer._broadcast_stop(False) is False - - -def test_broadcast_stop_adopts_rank_zero_decision(tmp_path: Path, monkeypatch) -> None: - trainer = _build_trainer(tmp_path, monkeypatch, ["stamp"]) - - monkeypatch.setattr(trainer_module, "synchronize_data", lambda *_args, **_kwargs: [True, False, False]) - assert trainer._broadcast_stop(False) is True # a non-zero rank still stops when rank 0 does - - monkeypatch.setattr(trainer_module, "synchronize_data", lambda *_args, **_kwargs: [False, True]) - assert trainer._broadcast_stop(True) is False # a non-zero rank keeps going when rank 0 does From 16040e0a21c1c492b572d04991c63bc13c680dc2 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Sun, 5 Jul 2026 19:26:38 +0200 Subject: [PATCH 3/4] ci: lint and format-check the core test suite Extend ruff check/format targets with tests/ in the CI jobs and in both pixi task blocks so test-suite hygiene is gated like source code. --- .github/workflows/konfai_ci.yml | 4 ++-- pyproject.toml | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/konfai_ci.yml b/.github/workflows/konfai_ci.yml index 761afaa4..eff370ff 100644 --- a/.github/workflows/konfai_ci.yml +++ b/.github/workflows/konfai_ci.yml @@ -59,7 +59,7 @@ jobs: run: pip install ruff==0.15.2 - name: Run ruff lint - run: ruff check konfai konfai-apps/konfai_apps + run: ruff check konfai konfai-apps/konfai_apps tests format: runs-on: ubuntu-latest @@ -75,7 +75,7 @@ jobs: run: pip install ruff==0.15.2 - name: Check formatting - run: ruff format --check konfai konfai-apps/konfai_apps + run: ruff format --check konfai konfai-apps/konfai_apps tests build: runs-on: ubuntu-latest diff --git a/pyproject.toml b/pyproject.toml index bdcd4d2b..5e31f0db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -163,9 +163,9 @@ scikit-image = "*" test = { cmd = "pytest -q tests/", description = "Run the test suite" } test-apps = { cmd = "pytest -q konfai-apps/tests", description = "Run the konfai-apps test suite" } test-cov = { cmd = "pytest --cov=konfai --cov-report=term-missing tests/", description = "Run tests with coverage" } -lint = { cmd = "ruff check konfai konfai-apps/konfai_apps", description = "Lint source code" } -format = { cmd = "ruff format konfai konfai-apps/konfai_apps", description = "Format source code" } -format-check = { cmd = "ruff format --check konfai konfai-apps/konfai_apps", description = "Check formatting without modifying files" } +lint = { cmd = "ruff check konfai konfai-apps/konfai_apps tests", description = "Lint source code" } +format = { cmd = "ruff format konfai konfai-apps/konfai_apps tests", description = "Format source code" } +format-check = { cmd = "ruff format --check konfai konfai-apps/konfai_apps tests", description = "Check formatting without modifying files" } typecheck = { cmd = "python -m mypy konfai --ignore-missing-imports --no-site-packages", description = "Type-check the konfai package" } build = { cmd = "python -m build", description = "Build sdist and wheel" } check = { depends-on = ["lint", "format-check", "test", "test-apps"], description = "Run all quality checks" } @@ -181,9 +181,9 @@ ruff = "==0.15.2" pre-commit = "*" [tool.pixi.feature.lint.tasks] -lint = { cmd = "ruff check konfai konfai-apps/konfai_apps", description = "Lint source code" } -format = { cmd = "ruff format konfai konfai-apps/konfai_apps", description = "Format source code" } -format-check = { cmd = "ruff format --check konfai konfai-apps/konfai_apps", description = "Check formatting without modifying files" } +lint = { cmd = "ruff check konfai konfai-apps/konfai_apps tests", description = "Lint source code" } +format = { cmd = "ruff format konfai konfai-apps/konfai_apps tests", description = "Format source code" } +format-check = { cmd = "ruff format --check konfai konfai-apps/konfai_apps tests", description = "Check formatting without modifying files" } pre-commit-install = { cmd = "pre-commit install --hook-type pre-commit --hook-type commit-msg", description = "Install pre-commit and commit-msg hooks" } pre-commit-run = { cmd = "pre-commit run --all-files", description = "Run pre-commit hooks on all repository files" } From f0629940bef73520cae26144c540445af8c14b3f Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Mon, 6 Jul 2026 11:13:09 +0200 Subject: [PATCH 4/4] style: sort imports in test_model_load_device (tests lint) --- tests/unit/test_model_load_device.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_model_load_device.py b/tests/unit/test_model_load_device.py index 555420a3..643d12b3 100644 --- a/tests/unit/test_model_load_device.py +++ b/tests/unit/test_model_load_device.py @@ -16,7 +16,6 @@ import pytest import torch - from konfai.network.network import Network from konfai.predictor import Mean, ModelComposite, _colocate_loaded_modules