Back to report index

mlharness 2ce8912: AI3D-379 Align config module with fleet pattern

Miroslav Simko <ms@iolabs.ch> 2026-09-02T09:39:03+02:00

Commit #77 · 25 snippets

 README.md                         |  38 +++++---
 pyproject.toml                    |   2 +-
 src/iolabs_ml_harness/__init__.py |   4 +
 src/iolabs_ml_harness/config.py   |  70 ++++++++++++---
 tests/test_config.py              | 182 ++++++++++++++++++++++++++++++++++++++
 tests/test_config_registry.py     | 177 ------------------------------------
 tests/test_registry.py            |  41 +++++++++
 7 files changed, 309 insertions(+), 205 deletions(-)
Importance #1: src/iolabs_ml_harness/config.py @@ -1,17 +1,23 @@
1"""Shared configuration sections for the training harness.1"""Shared configuration sections for the training harness.
22
3Task-specific data specs and top-level harness config stay in consumers.3The schema is a set of `<Name>Config` sections (each a
4Compose local `iolabs.common.config_loader.ConfigModel` sections with4`config_loader.ConfigModel`) that consumers compose into their own top-level
5``build_section`` so unknown keys still fail fast.5model; this package owns no config file of its own, so there is no packaged
6JSON to mirror.
67
7Adding a knob is a one-line change: add the field (with its default) to the8Adding a config key means adding the field (with its type, default and any
8model below; nothing else has to be touched.9`Field` range) to the section model below -- nothing else. Unknown keys are
10rejected.
11
12`build_section` and `with_overrides` return the frozen section model;
13`load_yaml_mapping` returns a plain `dict`.
914
10``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D15``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D
11consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict16consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict
12`ConfigModel`.17`ConfigModel`.
13"""18"""
19import logging
14from collections.abc import Mapping20from collections.abc import Mapping
15from pathlib import Path21from pathlib import Path
16from typing import Any, Literal, TypeVar, get_args22from typing import Any, Literal, TypeVar, get_args
1723
Importance #2: src/iolabs_ml_harness/config.py @@ -19,8 +25,13 @@
19import yaml25import yaml
20from iolabs.common import config_loader26from iolabs.common import config_loader
21from pydantic import fields as pydantic_fields27from pydantic import fields as pydantic_fields
2228
29logger = logging.getLogger(__name__)
30
31_CONTEXT = "harness config"
32_OVERRIDES_CONTEXT = "harness config overrides"
33
23T = TypeVar("T", bound=config_loader.ConfigModel)34T = TypeVar("T", bound=config_loader.ConfigModel)
2435
2536
26class SectionModel(config_loader.ConfigModel):37class SectionModel(config_loader.ConfigModel):
Importance #3: src/iolabs_ml_harness/config.py @@ -64,10 +75,13 @@
64 raise ValueError(f"{name} must be an int or a str, got bool {value!r}")75 raise ValueError(f"{name} must be an int or a str, got bool {value!r}")
65 return value76 return value
6677
6778
68class ConfigError(config_loader.ConfigError):79class HarnessConfigError(config_loader.ConfigError):
69 """Raised when a harness config section holds unknown keys or bad values."""80 """Raised when harness config contains unsupported keys or values."""
81
82
83ConfigError = HarnessConfigError # pre-0.2.2 name, kept for consumers
7084
7185
72def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:86def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:
73 """Builds a config model from a mapping, rejecting unknown keys.87 """Builds a config model from a mapping, rejecting unknown keys.
Importance #4: src/iolabs_ml_harness/config.py @@ -80,13 +94,40 @@
80 Returns:94 Returns:
81 An instance of ``cls`` built from ``data``.95 An instance of ``cls`` built from ``data``.
8296
83 Raises:97 Raises:
84 ConfigError: If ``data`` holds keys that are not fields of ``cls``, or a98 HarnessConfigError: If ``data`` holds keys that are not fields of
85 value that is invalid for its declared field type.99 ``cls``, or a value that is invalid for its declared field type.
100 """
101 return config_loader.validate_config(
102 cls, dict(data or {}), context=where, error_cls=HarnessConfigError)
103
104
105def with_overrides(config: T, sections: Mapping[str, Mapping[str, Any]]) -> T:
106 """Returns a re-validated copy of *config* with per-section overrides merged in.
107
108 ``model_copy(update=...)`` would store the values unchecked, so a
109 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping
110 through the model keeps CLI overrides on exactly the path YAML values take.
111
112 Args:
113 config: The config to derive from; never mutated.
114 sections: Section name -> field name -> override value. Empty sections
115 are ignored.
116
117 Returns:
118 ``config`` itself when no override is given, otherwise a validated copy.
119
120 Raises:
121 HarnessConfigError: An override value is invalid for its declared field.
86 """122 """
123 applied = {name: dict(values) for name, values in sections.items() if values}
124 if not applied:
125 return config
126 logger.info("Config overrides applied: %s", ", ".join(sorted(applied)))
127 merged = config_loader.deep_merge_dicts(config.model_dump(), applied)
87 return config_loader.validate_config(128 return config_loader.validate_config(
88 cls, dict(data or {}), context=where, error_cls=ConfigError)129 type(config), merged, context=_OVERRIDES_CONTEXT, error_cls=HarnessConfigError)
89130
90131
91def load_yaml_mapping(path: str | Path) -> dict[str, Any]:132def load_yaml_mapping(path: str | Path) -> dict[str, Any]:
92 """Loads a YAML file whose top level must be a mapping.133 """Loads a YAML file whose top level must be a mapping.
Importance #5: src/iolabs_ml_harness/config.py @@ -101,19 +142,20 @@
101 The parsed mapping. An empty document yields an empty dict.142 The parsed mapping. An empty document yields an empty dict.
102143
103 Raises:144 Raises:
104 FileNotFoundError: If ``path`` does not exist.145 FileNotFoundError: If ``path`` does not exist.
105 ConfigError: If the document's top level is not a mapping. The class146 HarnessConfigError: If the document's top level is not a mapping. The
106 derives from ``ValueError``, so legacy handlers keep working.147 class derives from ``ValueError``, so legacy handlers keep working.
107 """148 """
108 config_path = Path(path)149 config_path = Path(path)
109 with config_path.open("r", encoding="utf-8") as handle:150 with config_path.open("r", encoding="utf-8") as handle:
110 loaded = yaml.safe_load(handle)151 loaded = yaml.safe_load(handle)
152 logger.info("Config file applied: %s", config_path)
111 if loaded is None:153 if loaded is None:
112 return {}154 return {}
113 if not isinstance(loaded, Mapping):155 if not isinstance(loaded, Mapping):
114 raise ConfigError(156 raise HarnessConfigError(
115 f"config {str(config_path)!r} must contain a top-level mapping, "157 f"{_CONTEXT} {str(config_path)!r} must contain a top-level mapping, "
116 f"got {type(loaded).__name__}")158 f"got {type(loaded).__name__}")
117 return dict(loaded)159 return dict(loaded)
118160
119161
Importance #6: src/iolabs_ml_harness/__init__.py @@ -9,14 +9,16 @@
99
10from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,11 ConfigError,
12 FactoryConfig,12 FactoryConfig,
13 HarnessConfigError,
13 LossConfig,14 LossConfig,
14 ModelConfig,15 ModelConfig,
15 SectionModel,16 SectionModel,
16 TrainerConfig,17 TrainerConfig,
17 build_section,18 build_section,
18 load_yaml_mapping,19 load_yaml_mapping,
20 with_overrides,
19)21)
20from iolabs_ml_harness.losses import (22from iolabs_ml_harness.losses import (
21 WeightedSum,23 WeightedSum,
22 available_losses,24 available_losses,
Importance #7: src/iolabs_ml_harness/__init__.py @@ -70,8 +72,9 @@
70 "DEFAULT_CLASS_COLORS_RGB",72 "DEFAULT_CLASS_COLORS_RGB",
71 "ConfigError",73 "ConfigError",
72 "FactoryConfig",74 "FactoryConfig",
73 "FactoryRegistry",75 "FactoryRegistry",
76 "HarnessConfigError",
74 "LossConfig",77 "LossConfig",
75 "MaskOverlayWriter",78 "MaskOverlayWriter",
76 "ModelConfig",79 "ModelConfig",
77 "SectionModel",80 "SectionModel",
Importance #8: src/iolabs_ml_harness/__init__.py @@ -100,8 +103,9 @@
100 "register_loss",103 "register_loss",
101 "register_model",104 "register_model",
102 "render_overlay",105 "render_overlay",
103 "sha256_file",106 "sha256_file",
107 "with_overrides",
104 "write_run_provenance",108 "write_run_provenance",
105]109]
106110
107111
Importance #9: tests/test_config.py @@ -0,0 +1,182 @@
1"""Strict config construction, override merging, and YAML loading unit tests."""
2from pathlib import Path
3from types import MappingProxyType
4from typing import Any
5
6import pytest
7
8from iolabs.common import config_loader
9
10from iolabs_ml_harness.config import (
11 ConfigError,
12 FactoryConfig,
13 HarnessConfigError,
14 LossConfig,
15 ModelConfig,
16 SectionModel,
17 TrainerConfig,
18 build_section,
19 load_yaml_mapping,
20 with_overrides,
21)
22
23
24class _Section(config_loader.ConfigModel):
25 alpha: int = 1
26 beta: str = "b"
27 extra: dict[str, Any] = {}
28
29
30def test_build_section_accepts_any_mapping() -> None:
31 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")
32 assert cfg.alpha == 7 and cfg.beta == "b"
33
34
35def test_build_section_none_yields_defaults() -> None:
36 cfg = build_section(_Section, None, "section")
37 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})
38
39
40def test_unknown_top_level_key_is_rejected() -> None:
41 with pytest.raises(HarnessConfigError) as excinfo:
42 build_section(_Section, {"gamma": 1, "delta": 2}, "train")
43 message = str(excinfo.value)
44 assert "Unknown train key(s): delta, gamma" in message
45 assert "alpha" in message and "beta" in message
46
47
48def test_error_class_is_config_error() -> None:
49 assert issubclass(HarnessConfigError, config_loader.ConfigError)
50 assert issubclass(HarnessConfigError, ValueError)
51 assert ConfigError is HarnessConfigError, "pre-0.2.2 alias must stay bound"
52 with pytest.raises(ValueError):
53 build_section(_Section, {"alpha": "not-an-int"}, "train")
54
55
56def test_shared_section_defaults_are_unchanged() -> None:
57 model = ModelConfig()
58 assert (model.name, model.encoder_name, model.encoder_weights) == (
59 "unet", "resnet18", "imagenet")
60 assert (model.in_channels, model.num_classes, model.extra) == (1, 3, {})
61 loss = LossConfig()
62 assert (loss.name, loss.args) == ("dice_focal", {})
63 train = TrainerConfig()
64 assert train.max_epochs == -1 and train.lr == pytest.approx(3.0e-4)
65 assert train.monitor == "val/f1_mean_fg" and train.monitor_mode == "max"
66 assert train.early_stop_monitor == "val/loss" and train.early_stop_patience == 4
67 assert train.log_dir == "runs" and train.log_every_n_steps == 10
68 assert train.viz_every_n_epochs == 2 and train.viz_samples == 4
69 assert train.precision == "auto" and train.accelerator == "auto"
70 assert train.devices == 1 and train.accumulate_grad_batches == 1
71 assert train.weight_decay == pytest.approx(1.0e-4)
72 assert train.early_stop_mode == "min"
73
74
75def test_null_section_value_falls_back_to_the_default() -> None:
76 """A bare ``args:``/``extra:`` line parses to None and must mean "default"."""
77 assert build_section(LossConfig, {"name": "dice", "args": None}, "loss").args == {}
78 assert build_section(ModelConfig, {"extra": None}, "model").extra == {}
79 # a field that really accepts None keeps it
80 assert build_section(
81 ModelConfig, {"encoder_weights": None}, "model").encoder_weights is None
82 # required fields still fail on None
83 with pytest.raises(HarnessConfigError, match="name"):
84 build_section(FactoryConfig, {"name": None}, "loss")
85
86
87def test_null_means_default_only_applies_to_section_models() -> None:
88 class _Plain(config_loader.ConfigModel):
89 alpha: int = 1
90
91 assert issubclass(LossConfig, SectionModel)
92 with pytest.raises(HarnessConfigError):
93 build_section(_Plain, {"alpha": None}, "plain")
94
95
96def test_precision_accepts_lightning_ints_and_strings() -> None:
97 assert build_section(TrainerConfig, {"precision": 16}, "train").precision == 16
98 assert build_section(
99 TrainerConfig, {"precision": "bf16-mixed"}, "train").precision == "bf16-mixed"
100
101
102@pytest.mark.parametrize("field", ["precision", "devices"])
103def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str) -> None:
104 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""
105 with pytest.raises(HarnessConfigError, match="got bool"):
106 build_section(TrainerConfig, {field: True}, "train")
107
108
109def test_factory_config_is_name_plus_args() -> None:
110 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",
111 "args": {"alpha": 0.8}}, "loss")
112 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}
113 assert FactoryConfig(name="x").args == {}
114 with pytest.raises(HarnessConfigError, match="Unknown loss key"):
115 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")
116
117
118def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:
119 path = tmp_path / "cfg.yaml"
120 path.write_text("experiment: e01\nseed: 1337\nnested:\n a: [1, 2]\n",
121 encoding="utf-8")
122 loaded = load_yaml_mapping(path)
123 assert loaded == {"experiment": "e01", "seed": 1337,
124 "nested": {"a": [1, 2]}}
125 assert load_yaml_mapping(str(path)) == loaded
126
127
128def test_load_yaml_mapping_empty_document(tmp_path: Path) -> None:
129 path = tmp_path / "empty.yaml"
130 path.write_text("", encoding="utf-8")
131 assert load_yaml_mapping(path) == {}
132
133
134def test_load_yaml_mapping_rejects_non_mapping(tmp_path: Path) -> None:
135 path = tmp_path / "list.yaml"
136 path.write_text("- a\n- b\n", encoding="utf-8")
137 with pytest.raises(ValueError, match="top-level mapping"):
138 load_yaml_mapping(path)
139
140
141def test_load_yaml_mapping_missing_file(tmp_path: Path) -> None:
142 with pytest.raises(FileNotFoundError):
143 load_yaml_mapping(tmp_path / "absent.yaml")
144
145
146def test_unknown_nested_key_is_rejected() -> None:
147 class _Root(SectionModel):
148 train: TrainerConfig = TrainerConfig()
149
150 with pytest.raises(HarnessConfigError, match="lr_schedule"):
151 build_section(_Root, {"train": {"lr_schedule": "cosine"}}, "root")
152
153
154def test_overrides_deep_merge_onto_defaults() -> None:
155 class _Root(SectionModel):
156 experiment: str = "experiment"
157 model: ModelConfig = ModelConfig()
158 train: TrainerConfig = TrainerConfig()
159
160 base = build_section(_Root, {"model": {"num_classes": 5}}, "root")
161 merged = with_overrides(base, {"train": {"max_epochs": 7}, "loss": {}})
162 assert merged.train.max_epochs == 7
163 assert merged.model.num_classes == 5, "untouched sections must survive"
164 assert merged.model.encoder_name == "resnet18"
165 assert merged.experiment == "experiment"
166 assert base.train.max_epochs == -1, "the source config must not be mutated"
167 assert with_overrides(base, {"train": {}}) is base
168
169
170def test_set_override_coercion_and_rejection() -> None:
171 """Overrides take the fleet coercion matrix, not raw pydantic narrowing."""
172 class _Root(SectionModel):
173 augment: bool = False
174 train: TrainerConfig = TrainerConfig()
175
176 coerced = with_overrides(_Root(), {"train": {"log_every_n_steps": "1e3"}})
177 assert coerced.train.log_every_n_steps == 1000
178 assert build_section(_Root, {"augment": "on"}, "root").augment is True
179 with pytest.raises(HarnessConfigError, match="augment"):
180 build_section(_Root, {"augment": "flase"}, "root")
181 with pytest.raises(HarnessConfigError, match="max_epochs"):
182 with_overrides(_Root(), {"train": {"max_epochs": -2}})
0
Importance #10: tests/test_registry.py @@ -0,0 +1,41 @@
1"""Generic factory-registry unit tests."""
2import pytest
3
4from iolabs_ml_harness.registry import FactoryRegistry
5
6
7def test_registry_registers_lists_and_builds() -> None:
8 registry: FactoryRegistry[str] = FactoryRegistry("widget")
9
10 @registry.register("beta")
11 def beta(scale: int = 1) -> str:
12 return f"beta-{scale}"
13
14 @registry.register("alpha")
15 def alpha() -> str:
16 return "alpha"
17
18 assert registry.names() == ["alpha", "beta"]
19 assert registry.build("beta", scale=3) == "beta-3"
20 assert registry.build("alpha") == "alpha"
21 assert "alpha" in registry and "gamma" not in registry
22 assert registry.get("alpha") is alpha
23 assert beta(2) == "beta-2", "the decorator must return the factory unchanged"
24
25
26def test_registry_unknown_name_names_kind_and_options() -> None:
27 registry: FactoryRegistry[str] = FactoryRegistry("loss")
28 registry.register("cross_entropy")(lambda: "ce")
29 with pytest.raises(KeyError) as excinfo:
30 registry.build("soft_cldice")
31 message = str(excinfo.value)
32 assert "unknown loss 'soft_cldice'" in message
33 assert "['cross_entropy']" in message
34
35
36def test_registry_registration_is_last_one_wins() -> None:
37 registry: FactoryRegistry[str] = FactoryRegistry("widget")
38 registry.register("dup")(lambda: "first")
39 registry.register("dup")(lambda: "second")
40 assert registry.build("dup") == "second"
41 assert registry.names() == ["dup"]
0
Importance #11: pyproject.toml @@ -3,9 +3,9 @@
3build-backend = "hatchling.build"3build-backend = "hatchling.build"
44
5[project]5[project]
6name = "iolabs-ml-harness"6name = "iolabs-ml-harness"
7version = "0.2.1"7version = "0.2.2"
8description = "Task-agnostic PyTorch-Lightning segmentation training harness."8description = "Task-agnostic PyTorch-Lightning segmentation training harness."
9readme = "README.md"9readme = "README.md"
10requires-python = ">=3.11,<3.13"10requires-python = ">=3.11,<3.13"
11dependencies = [11dependencies = [
Importance #12: README.md @@ -7,9 +7,9 @@
7Core is modality neutral and depends only on `numpy`, `pyyaml`, `pydantic`, `iolabs-common` (the shared config layer), `torch`, `lightning`, `tensorboard`, and `torchmetrics`. OpenCV and `segmentation-models-pytorch` moved into the `image` extra in 0.2.0, so a point-cloud consumer never installs an image stack.7Core is modality neutral and depends only on `numpy`, `pyyaml`, `pydantic`, `iolabs-common` (the shared config layer), `torch`, `lightning`, `tensorboard`, and `torchmetrics`. OpenCV and `segmentation-models-pytorch` moved into the `image` extra in 0.2.0, so a point-cloud consumer never installs an image stack.
88
9| Module | Extra | Purpose |9| Module | Extra | Purpose |
10| --- | --- | --- |10| --- | --- | --- |
11| `config` | core | Strict `build_section`, `load_yaml_mapping`, `ConfigError`, `SectionModel`, `FactoryConfig`, and the legacy model/loss/trainer sections. |11| `config` | core | Strict `build_section`, `with_overrides`, `load_yaml_mapping`, `HarnessConfigError`, `SectionModel`, `FactoryConfig`, and the legacy model/loss/trainer sections. |
12| `registry` | core | Typed `FactoryRegistry` behind the model and loss registries. |12| `registry` | core | Typed `FactoryRegistry` behind the model and loss registries. |
13| `models` | core | Model registry, `build_registered_model`, and the legacy lazy smp fallback in `build_model`. |13| `models` | core | Model registry, `build_registered_model`, and the legacy lazy smp fallback in `build_model`. |
14| `losses` | core | Task-neutral registry, `WeightedSum`, and rank-agnostic `cross_entropy` / `focal_cross_entropy` / `masked_focal_tversky`. |14| `losses` | core | Task-neutral registry, `WeightedSum`, and rank-agnostic `cross_entropy` / `focal_cross_entropy` / `masked_focal_tversky`. |
15| `metrics` | core | Rank-agnostic, ignore-aware confusion matrix and `SegmentationStats`. |15| `metrics` | core | Rank-agnostic, ignore-aware confusion matrix and `SegmentationStats`. |
Importance #13: README.md @@ -25,20 +25,28 @@
25```bash25```bash
26pip install 'iolabs-ml-harness[image]'26pip install 'iolabs-ml-harness[image]'
27```27```
2828
29## Config29## Configuration
3030
31Config sections are pydantic models derived from31This package owns no config file: it publishes the shared sections a consumer
32`iolabs.common.config_loader.ConfigModel` (unknown keys rejected, values coerced32composes into its own top-level model, so there is no packaged `*.default.json`.
33by the shared fleet matrix, instances frozen). **Adding a config key = adding one33The schema is `SectionModel`, `FactoryConfig`, `ModelConfig`, `LossConfig` and
34field with its default to the model in `config.py`** — there is no separate34`TrainerConfig` in `iolabs_ml_harness.config` (each a
35allow-list, coercion helper, or dataclass to keep in sync. `build_section`35`config_loader.ConfigModel`); nested YAML sections are nested models and unknown
36raises `ConfigError`, which derives from `ValueError`. Constructing a model36keys are rejected. **To add a config key: add the field (with its type, default
37directly (`TrainerConfig(max_epochs=-2)`) raises `pydantic.ValidationError`;37and any `Field` range) to the model — nothing else.** `build_section` and
38only the `build_section` / `validate_config` entry points wrap it in38`with_overrides` return the frozen model; `load_yaml_mapping` returns a plain
39`ConfigError`, so mutate configs through them rather than `model_copy(update=)`,39`dict`. Runtime overrides come from the consumer's CLI flags through
40which skips validation entirely.40`with_overrides`, never repo-local YAML.
41
42`build_section` raises `HarnessConfigError` (a `config_loader.ConfigError`, so a
43`ValueError`); `ConfigError` stays bound to it as the pre-0.2.2 alias.
44Constructing a model directly (`TrainerConfig(max_epochs=-2)`) raises
45`pydantic.ValidationError`; only the `build_section` / `with_overrides` /
46`validate_config` entry points wrap it in `HarnessConfigError`, so mutate
47configs through them rather than `model_copy(update=)`, which skips validation
48entirely.
4149
42Sections derive from `SectionModel`, so a bare `loss:` / `args:` line (YAML50Sections derive from `SectionModel`, so a bare `loss:` / `args:` line (YAML
43`null`) means "use the defaults" — as the pre-pydantic `data or {}` coalescing51`null`) means "use the defaults" — as the pre-pydantic `data or {}` coalescing
44did — while a field that genuinely accepts `None` (`model.encoder_weights`)52did — while a field that genuinely accepts `None` (`model.encoder_weights`)
Importance #14: README.md @@ -46,8 +54,12 @@
46`precision`/`devices` reject YAML `yes`/`on` instead of narrowing them to `1`.54`precision`/`devices` reject YAML `yes`/`on` instead of narrowing them to `1`.
4755
48## Compatibility56## Compatibility
4957
580.2.2 keeps every public name: the config error class is now
59`HarnessConfigError` (fleet naming), with `ConfigError` kept as an alias, and
60`with_overrides` is new.
61
500.2.1 keeps every public name; `build_section` now raises `ConfigError`620.2.1 keeps every public name; `build_section` now raises `ConfigError`
51(a `ValueError`) instead of `KeyError` for unknown keys, and the section classes63(a `ValueError`) instead of `KeyError` for unknown keys, and the section classes
52are `ConfigModel`s rather than dataclasses — attribute access is unchanged, but64are `ConfigModel`s rather than dataclasses — attribute access is unchanged, but
53instances are frozen (use `model_copy(update=...)` instead of assignment).65instances are frozen (use `model_copy(update=...)` instead of assignment).
Importance #15: pyproject.toml @@ -3,9 +3,9 @@
3build-backend = "hatchling.build"3build-backend = "hatchling.build"
44
5[project]5[project]
6name = "iolabs-ml-harness"6name = "iolabs-ml-harness"
7version = "0.2.1"7version = "0.2.2"
8description = "Task-agnostic PyTorch-Lightning segmentation training harness."8description = "Task-agnostic PyTorch-Lightning segmentation training harness."
9readme = "README.md"9readme = "README.md"
10requires-python = ">=3.11,<3.13"10requires-python = ">=3.11,<3.13"
11dependencies = [11dependencies = [
Importance #16: src/iolabs_ml_harness/__init__.py @@ -9,14 +9,16 @@
99
10from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,11 ConfigError,
12 FactoryConfig,12 FactoryConfig,
13 HarnessConfigError,
13 LossConfig,14 LossConfig,
14 ModelConfig,15 ModelConfig,
15 SectionModel,16 SectionModel,
16 TrainerConfig,17 TrainerConfig,
17 build_section,18 build_section,
18 load_yaml_mapping,19 load_yaml_mapping,
20 with_overrides,
19)21)
20from iolabs_ml_harness.losses import (22from iolabs_ml_harness.losses import (
21 WeightedSum,23 WeightedSum,
22 available_losses,24 available_losses,
Importance #17: src/iolabs_ml_harness/__init__.py @@ -70,8 +72,9 @@
70 "DEFAULT_CLASS_COLORS_RGB",72 "DEFAULT_CLASS_COLORS_RGB",
71 "ConfigError",73 "ConfigError",
72 "FactoryConfig",74 "FactoryConfig",
73 "FactoryRegistry",75 "FactoryRegistry",
76 "HarnessConfigError",
74 "LossConfig",77 "LossConfig",
75 "MaskOverlayWriter",78 "MaskOverlayWriter",
76 "ModelConfig",79 "ModelConfig",
77 "SectionModel",80 "SectionModel",
Importance #18: src/iolabs_ml_harness/__init__.py @@ -100,8 +103,9 @@
100 "register_loss",103 "register_loss",
101 "register_model",104 "register_model",
102 "render_overlay",105 "render_overlay",
103 "sha256_file",106 "sha256_file",
107 "with_overrides",
104 "write_run_provenance",108 "write_run_provenance",
105]109]
106110
107111
Importance #19: src/iolabs_ml_harness/config.py @@ -1,17 +1,23 @@
1"""Shared configuration sections for the training harness.1"""Shared configuration sections for the training harness.
22
3Task-specific data specs and top-level harness config stay in consumers.3The schema is a set of `<Name>Config` sections (each a
4Compose local `iolabs.common.config_loader.ConfigModel` sections with4`config_loader.ConfigModel`) that consumers compose into their own top-level
5``build_section`` so unknown keys still fail fast.5model; this package owns no config file of its own, so there is no packaged
6JSON to mirror.
67
7Adding a knob is a one-line change: add the field (with its default) to the8Adding a config key means adding the field (with its type, default and any
8model below; nothing else has to be touched.9`Field` range) to the section model below -- nothing else. Unknown keys are
10rejected.
11
12`build_section` and `with_overrides` return the frozen section model;
13`load_yaml_mapping` returns a plain `dict`.
914
10``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D15``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D
11consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict16consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict
12`ConfigModel`.17`ConfigModel`.
13"""18"""
19import logging
14from collections.abc import Mapping20from collections.abc import Mapping
15from pathlib import Path21from pathlib import Path
16from typing import Any, Literal, TypeVar, get_args22from typing import Any, Literal, TypeVar, get_args
1723
Importance #20: src/iolabs_ml_harness/config.py @@ -19,8 +25,13 @@
19import yaml25import yaml
20from iolabs.common import config_loader26from iolabs.common import config_loader
21from pydantic import fields as pydantic_fields27from pydantic import fields as pydantic_fields
2228
29logger = logging.getLogger(__name__)
30
31_CONTEXT = "harness config"
32_OVERRIDES_CONTEXT = "harness config overrides"
33
23T = TypeVar("T", bound=config_loader.ConfigModel)34T = TypeVar("T", bound=config_loader.ConfigModel)
2435
2536
26class SectionModel(config_loader.ConfigModel):37class SectionModel(config_loader.ConfigModel):
Importance #21: src/iolabs_ml_harness/config.py @@ -64,10 +75,13 @@
64 raise ValueError(f"{name} must be an int or a str, got bool {value!r}")75 raise ValueError(f"{name} must be an int or a str, got bool {value!r}")
65 return value76 return value
6677
6778
68class ConfigError(config_loader.ConfigError):79class HarnessConfigError(config_loader.ConfigError):
69 """Raised when a harness config section holds unknown keys or bad values."""80 """Raised when harness config contains unsupported keys or values."""
81
82
83ConfigError = HarnessConfigError # pre-0.2.2 name, kept for consumers
7084
7185
72def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:86def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:
73 """Builds a config model from a mapping, rejecting unknown keys.87 """Builds a config model from a mapping, rejecting unknown keys.
Importance #22: src/iolabs_ml_harness/config.py @@ -80,13 +94,40 @@
80 Returns:94 Returns:
81 An instance of ``cls`` built from ``data``.95 An instance of ``cls`` built from ``data``.
8296
83 Raises:97 Raises:
84 ConfigError: If ``data`` holds keys that are not fields of ``cls``, or a98 HarnessConfigError: If ``data`` holds keys that are not fields of
85 value that is invalid for its declared field type.99 ``cls``, or a value that is invalid for its declared field type.
100 """
101 return config_loader.validate_config(
102 cls, dict(data or {}), context=where, error_cls=HarnessConfigError)
103
104
105def with_overrides(config: T, sections: Mapping[str, Mapping[str, Any]]) -> T:
106 """Returns a re-validated copy of *config* with per-section overrides merged in.
107
108 ``model_copy(update=...)`` would store the values unchecked, so a
109 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping
110 through the model keeps CLI overrides on exactly the path YAML values take.
111
112 Args:
113 config: The config to derive from; never mutated.
114 sections: Section name -> field name -> override value. Empty sections
115 are ignored.
116
117 Returns:
118 ``config`` itself when no override is given, otherwise a validated copy.
119
120 Raises:
121 HarnessConfigError: An override value is invalid for its declared field.
86 """122 """
123 applied = {name: dict(values) for name, values in sections.items() if values}
124 if not applied:
125 return config
126 logger.info("Config overrides applied: %s", ", ".join(sorted(applied)))
127 merged = config_loader.deep_merge_dicts(config.model_dump(), applied)
87 return config_loader.validate_config(128 return config_loader.validate_config(
88 cls, dict(data or {}), context=where, error_cls=ConfigError)129 type(config), merged, context=_OVERRIDES_CONTEXT, error_cls=HarnessConfigError)
89130
90131
91def load_yaml_mapping(path: str | Path) -> dict[str, Any]:132def load_yaml_mapping(path: str | Path) -> dict[str, Any]:
92 """Loads a YAML file whose top level must be a mapping.133 """Loads a YAML file whose top level must be a mapping.
Importance #23: src/iolabs_ml_harness/config.py @@ -101,19 +142,20 @@
101 The parsed mapping. An empty document yields an empty dict.142 The parsed mapping. An empty document yields an empty dict.
102143
103 Raises:144 Raises:
104 FileNotFoundError: If ``path`` does not exist.145 FileNotFoundError: If ``path`` does not exist.
105 ConfigError: If the document's top level is not a mapping. The class146 HarnessConfigError: If the document's top level is not a mapping. The
106 derives from ``ValueError``, so legacy handlers keep working.147 class derives from ``ValueError``, so legacy handlers keep working.
107 """148 """
108 config_path = Path(path)149 config_path = Path(path)
109 with config_path.open("r", encoding="utf-8") as handle:150 with config_path.open("r", encoding="utf-8") as handle:
110 loaded = yaml.safe_load(handle)151 loaded = yaml.safe_load(handle)
152 logger.info("Config file applied: %s", config_path)
111 if loaded is None:153 if loaded is None:
112 return {}154 return {}
113 if not isinstance(loaded, Mapping):155 if not isinstance(loaded, Mapping):
114 raise ConfigError(156 raise HarnessConfigError(
115 f"config {str(config_path)!r} must contain a top-level mapping, "157 f"{_CONTEXT} {str(config_path)!r} must contain a top-level mapping, "
116 f"got {type(loaded).__name__}")158 f"got {type(loaded).__name__}")
117 return dict(loaded)159 return dict(loaded)
118160
119161
Importance #24: tests/test_config.py @@ -0,0 +1,182 @@
1"""Strict config construction, override merging, and YAML loading unit tests."""
2from pathlib import Path
3from types import MappingProxyType
4from typing import Any
5
6import pytest
7
8from iolabs.common import config_loader
9
10from iolabs_ml_harness.config import (
11 ConfigError,
12 FactoryConfig,
13 HarnessConfigError,
14 LossConfig,
15 ModelConfig,
16 SectionModel,
17 TrainerConfig,
18 build_section,
19 load_yaml_mapping,
20 with_overrides,
21)
22
23
24class _Section(config_loader.ConfigModel):
25 alpha: int = 1
26 beta: str = "b"
27 extra: dict[str, Any] = {}
28
29
30def test_build_section_accepts_any_mapping() -> None:
31 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")
32 assert cfg.alpha == 7 and cfg.beta == "b"
33
34
35def test_build_section_none_yields_defaults() -> None:
36 cfg = build_section(_Section, None, "section")
37 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})
38
39
40def test_unknown_top_level_key_is_rejected() -> None:
41 with pytest.raises(HarnessConfigError) as excinfo:
42 build_section(_Section, {"gamma": 1, "delta": 2}, "train")
43 message = str(excinfo.value)
44 assert "Unknown train key(s): delta, gamma" in message
45 assert "alpha" in message and "beta" in message
46
47
48def test_error_class_is_config_error() -> None:
49 assert issubclass(HarnessConfigError, config_loader.ConfigError)
50 assert issubclass(HarnessConfigError, ValueError)
51 assert ConfigError is HarnessConfigError, "pre-0.2.2 alias must stay bound"
52 with pytest.raises(ValueError):
53 build_section(_Section, {"alpha": "not-an-int"}, "train")
54
55
56def test_shared_section_defaults_are_unchanged() -> None:
57 model = ModelConfig()
58 assert (model.name, model.encoder_name, model.encoder_weights) == (
59 "unet", "resnet18", "imagenet")
60 assert (model.in_channels, model.num_classes, model.extra) == (1, 3, {})
61 loss = LossConfig()
62 assert (loss.name, loss.args) == ("dice_focal", {})
63 train = TrainerConfig()
64 assert train.max_epochs == -1 and train.lr == pytest.approx(3.0e-4)
65 assert train.monitor == "val/f1_mean_fg" and train.monitor_mode == "max"
66 assert train.early_stop_monitor == "val/loss" and train.early_stop_patience == 4
67 assert train.log_dir == "runs" and train.log_every_n_steps == 10
68 assert train.viz_every_n_epochs == 2 and train.viz_samples == 4
69 assert train.precision == "auto" and train.accelerator == "auto"
70 assert train.devices == 1 and train.accumulate_grad_batches == 1
71 assert train.weight_decay == pytest.approx(1.0e-4)
72 assert train.early_stop_mode == "min"
73
74
75def test_null_section_value_falls_back_to_the_default() -> None:
76 """A bare ``args:``/``extra:`` line parses to None and must mean "default"."""
77 assert build_section(LossConfig, {"name": "dice", "args": None}, "loss").args == {}
78 assert build_section(ModelConfig, {"extra": None}, "model").extra == {}
79 # a field that really accepts None keeps it
80 assert build_section(
81 ModelConfig, {"encoder_weights": None}, "model").encoder_weights is None
82 # required fields still fail on None
83 with pytest.raises(HarnessConfigError, match="name"):
84 build_section(FactoryConfig, {"name": None}, "loss")
85
86
87def test_null_means_default_only_applies_to_section_models() -> None:
88 class _Plain(config_loader.ConfigModel):
89 alpha: int = 1
90
91 assert issubclass(LossConfig, SectionModel)
92 with pytest.raises(HarnessConfigError):
93 build_section(_Plain, {"alpha": None}, "plain")
94
95
96def test_precision_accepts_lightning_ints_and_strings() -> None:
97 assert build_section(TrainerConfig, {"precision": 16}, "train").precision == 16
98 assert build_section(
99 TrainerConfig, {"precision": "bf16-mixed"}, "train").precision == "bf16-mixed"
100
101
102@pytest.mark.parametrize("field", ["precision", "devices"])
103def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str) -> None:
104 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""
105 with pytest.raises(HarnessConfigError, match="got bool"):
106 build_section(TrainerConfig, {field: True}, "train")
107
108
109def test_factory_config_is_name_plus_args() -> None:
110 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",
111 "args": {"alpha": 0.8}}, "loss")
112 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}
113 assert FactoryConfig(name="x").args == {}
114 with pytest.raises(HarnessConfigError, match="Unknown loss key"):
115 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")
116
117
118def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:
119 path = tmp_path / "cfg.yaml"
120 path.write_text("experiment: e01\nseed: 1337\nnested:\n a: [1, 2]\n",
121 encoding="utf-8")
122 loaded = load_yaml_mapping(path)
123 assert loaded == {"experiment": "e01", "seed": 1337,
124 "nested": {"a": [1, 2]}}
125 assert load_yaml_mapping(str(path)) == loaded
126
127
128def test_load_yaml_mapping_empty_document(tmp_path: Path) -> None:
129 path = tmp_path / "empty.yaml"
130 path.write_text("", encoding="utf-8")
131 assert load_yaml_mapping(path) == {}
132
133
134def test_load_yaml_mapping_rejects_non_mapping(tmp_path: Path) -> None:
135 path = tmp_path / "list.yaml"
136 path.write_text("- a\n- b\n", encoding="utf-8")
137 with pytest.raises(ValueError, match="top-level mapping"):
138 load_yaml_mapping(path)
139
140
141def test_load_yaml_mapping_missing_file(tmp_path: Path) -> None:
142 with pytest.raises(FileNotFoundError):
143 load_yaml_mapping(tmp_path / "absent.yaml")
144
145
146def test_unknown_nested_key_is_rejected() -> None:
147 class _Root(SectionModel):
148 train: TrainerConfig = TrainerConfig()
149
150 with pytest.raises(HarnessConfigError, match="lr_schedule"):
151 build_section(_Root, {"train": {"lr_schedule": "cosine"}}, "root")
152
153
154def test_overrides_deep_merge_onto_defaults() -> None:
155 class _Root(SectionModel):
156 experiment: str = "experiment"
157 model: ModelConfig = ModelConfig()
158 train: TrainerConfig = TrainerConfig()
159
160 base = build_section(_Root, {"model": {"num_classes": 5}}, "root")
161 merged = with_overrides(base, {"train": {"max_epochs": 7}, "loss": {}})
162 assert merged.train.max_epochs == 7
163 assert merged.model.num_classes == 5, "untouched sections must survive"
164 assert merged.model.encoder_name == "resnet18"
165 assert merged.experiment == "experiment"
166 assert base.train.max_epochs == -1, "the source config must not be mutated"
167 assert with_overrides(base, {"train": {}}) is base
168
169
170def test_set_override_coercion_and_rejection() -> None:
171 """Overrides take the fleet coercion matrix, not raw pydantic narrowing."""
172 class _Root(SectionModel):
173 augment: bool = False
174 train: TrainerConfig = TrainerConfig()
175
176 coerced = with_overrides(_Root(), {"train": {"log_every_n_steps": "1e3"}})
177 assert coerced.train.log_every_n_steps == 1000
178 assert build_section(_Root, {"augment": "on"}, "root").augment is True
179 with pytest.raises(HarnessConfigError, match="augment"):
180 build_section(_Root, {"augment": "flase"}, "root")
181 with pytest.raises(HarnessConfigError, match="max_epochs"):
182 with_overrides(_Root(), {"train": {"max_epochs": -2}})
0
Importance #25: tests/test_registry.py @@ -0,0 +1,41 @@
1"""Generic factory-registry unit tests."""
2import pytest
3
4from iolabs_ml_harness.registry import FactoryRegistry
5
6
7def test_registry_registers_lists_and_builds() -> None:
8 registry: FactoryRegistry[str] = FactoryRegistry("widget")
9
10 @registry.register("beta")
11 def beta(scale: int = 1) -> str:
12 return f"beta-{scale}"
13
14 @registry.register("alpha")
15 def alpha() -> str:
16 return "alpha"
17
18 assert registry.names() == ["alpha", "beta"]
19 assert registry.build("beta", scale=3) == "beta-3"
20 assert registry.build("alpha") == "alpha"
21 assert "alpha" in registry and "gamma" not in registry
22 assert registry.get("alpha") is alpha
23 assert beta(2) == "beta-2", "the decorator must return the factory unchanged"
24
25
26def test_registry_unknown_name_names_kind_and_options() -> None:
27 registry: FactoryRegistry[str] = FactoryRegistry("loss")
28 registry.register("cross_entropy")(lambda: "ce")
29 with pytest.raises(KeyError) as excinfo:
30 registry.build("soft_cldice")
31 message = str(excinfo.value)
32 assert "unknown loss 'soft_cldice'" in message
33 assert "['cross_entropy']" in message
34
35
36def test_registry_registration_is_last_one_wins() -> None:
37 registry: FactoryRegistry[str] = FactoryRegistry("widget")
38 registry.register("dup")(lambda: "first")
39 registry.register("dup")(lambda: "second")
40 assert registry.build("dup") == "second"
41 assert registry.names() == ["dup"]
0