Back to report index

mlharness bf85ad2: AI3D-379 Pydantic config models via iolabs-common ConfigModel

Miroslav Simko <ms@iolabs.ch> 2026-09-02T08:45:50+02:00

Commit #75 · 24 snippets

 README.md                         | 18 +++++++-
 pyproject.toml                    | 12 +++++-
 src/iolabs_ml_harness/__init__.py |  2 +
 src/iolabs_ml_harness/config.py   | 88 +++++++++++++++++++++------------------
 tests/test_config_registry.py     | 23 ++++++----
 tests/test_ml_harness.py          |  2 +-
 6 files changed, 92 insertions(+), 53 deletions(-)
Importance #1: src/iolabs_ml_harness/config.py @@ -53,62 +61,60 @@
53 The parsed mapping. An empty document yields an empty dict.61 The parsed mapping. An empty document yields an empty dict.
5462
55 Raises:63 Raises:
56 FileNotFoundError: If ``path`` does not exist.64 FileNotFoundError: If ``path`` does not exist.
57 ValueError: If the document's top level is not a mapping.65 ConfigError: If the document's top level is not a mapping. The class
66 derives from ``ValueError``, so legacy handlers keep working.
58 """67 """
59 config_path = Path(path)68 config_path = Path(path)
60 with config_path.open("r", encoding="utf-8") as handle:69 with config_path.open("r", encoding="utf-8") as handle:
61 loaded = yaml.safe_load(handle)70 loaded = yaml.safe_load(handle)
62 if loaded is None:71 if loaded is None:
63 return {}72 return {}
64 if not isinstance(loaded, Mapping):73 if not isinstance(loaded, Mapping):
65 raise ValueError(74 raise ConfigError(
66 f"config {str(config_path)!r} must contain a top-level mapping, "75 f"config {str(config_path)!r} must contain a top-level mapping, "
67 f"got {type(loaded).__name__}")76 f"got {type(loaded).__name__}")
68 return dict(loaded)77 return dict(loaded)
6978
7079
71@dataclass80class FactoryConfig(config_loader.ConfigModel):
72class FactoryConfig:
73 """Modality-neutral registry selection: a name plus factory keyword args."""81 """Modality-neutral registry selection: a name plus factory keyword args."""
74 name: str82 name: str
75 args: dict[str, Any] = field(default_factory=dict)83 args: dict[str, Any] = {}
7684
7785
78@dataclass86class ModelConfig(config_loader.ConfigModel):
79class ModelConfig:
80 """Legacy image-model schema (segmentation-models-pytorch shaped)."""87 """Legacy image-model schema (segmentation-models-pytorch shaped)."""
81 name: str = "unet"88 name: str = "unet"
82 encoder_name: str = "resnet18"89 encoder_name: str = "resnet18"
83 encoder_weights: str | None = "imagenet" # None = train from scratch90 encoder_weights: str | None = "imagenet" # None = train from scratch
84 in_channels: int = 191 in_channels: int = pydantic.Field(default=1, gt=0)
85 num_classes: int = 392 num_classes: int = pydantic.Field(default=3, gt=0)
86 extra: dict[str, Any] = field(default_factory=dict) # passed to the factory93 extra: dict[str, Any] = {} # passed to the factory
8794
8895
89@dataclass96class LossConfig(config_loader.ConfigModel):
90class LossConfig:
91 """Loss selection by registry name plus factory keyword arguments."""97 """Loss selection by registry name plus factory keyword arguments."""
92 name: str = "dice_focal"98 name: str = "dice_focal"
93 args: dict[str, Any] = field(default_factory=dict)99 args: dict[str, Any] = {}
94100
95101
96@dataclass102class TrainerConfig(config_loader.ConfigModel):
97class TrainerConfig:
98 """Harness-visible Lightning trainer, logger, and callback knobs."""103 """Harness-visible Lightning trainer, logger, and callback knobs."""
99 max_epochs: int = -1 # -1 = no cap; early stopping ends training instead104 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
100 lr: float = 3.0e-4105 lr: float = pydantic.Field(default=3.0e-4, gt=0)
101 weight_decay: float = 1.0e-4106 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
102 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU107 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU
103 accumulate_grad_batches: int = 1 # match effective batch across experiments108 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
104 accelerator: str = "auto"109 accelerator: str = "auto"
105 devices: int | str = 1110 devices: int | str = 1
106 viz_every_n_epochs: int = 2111 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
107 viz_samples: int = 4112 viz_samples: int = pydantic.Field(default=4, ge=0)
108 monitor: str = "val/f1_mean_fg"113 monitor: str = "val/f1_mean_fg"
109 monitor_mode: str = "max"114 monitor_mode: Literal["max", "min"] = "max"
110 early_stop_monitor: str = "val/loss" # stop when this stops improving115 early_stop_monitor: str = "val/loss" # stop when this stops improving
111 early_stop_mode: str = "min"116 early_stop_mode: Literal["min", "max"] = "min"
112 early_stop_patience: int = 4 # epochs without improvement before stopping; 0 disables117 # epochs without improvement before stopping; 0 disables
118 early_stop_patience: int = pydantic.Field(default=4, ge=0)
113 log_dir: str = "runs"119 log_dir: str = "runs"
114 log_every_n_steps: int = 10120 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
Importance #2: src/iolabs_ml_harness/config.py @@ -1,50 +1,58 @@
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.3Task-specific data specs and top-level harness config stay in consumers.
4Compose local dataclasses with ``build_section`` so unknown keys still fail fast.4Compose local `iolabs.common.config_loader.ConfigModel` sections with
5``build_section`` so unknown keys still fail fast.
6
7Adding a knob is a one-line change: add the field (with its default) to the
8model below; nothing else has to be touched.
59
6``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D10``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D
7consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict11consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict
8dataclass.12`ConfigModel`.
9"""13"""
14import logging
10from collections.abc import Mapping15from collections.abc import Mapping
11from dataclasses import dataclass, field, fields
12from pathlib import Path16from pathlib import Path
13from typing import Any, TypeVar17from typing import Any, Literal, TypeVar
1418
19import pydantic
15import yaml20import yaml
21from iolabs.common import config_loader
22
23logger = logging.getLogger(__name__)
24
25T = TypeVar("T", bound=config_loader.ConfigModel)
26
1627
17T = TypeVar("T")28class ConfigError(config_loader.ConfigError):
29 """Raised when a harness config section holds unknown keys or bad values."""
1830
1931
20def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:32def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:
21 """Builds a config dataclass from a mapping, rejecting unknown keys.33 """Builds a config model from a mapping, rejecting unknown keys.
2234
23 Args:35 Args:
24 cls: Dataclass to instantiate.36 cls: `ConfigModel` subclass to instantiate.
25 data: Mapping of field name to value, or ``None`` for all defaults.37 data: Mapping of field name to value, or ``None`` for all defaults.
26 where: Section name used in the error message, e.g. ``"model"``.38 where: Section name used in the error message, e.g. ``"model"``.
2739
28 Returns:40 Returns:
29 An instance of ``cls`` built from ``data``.41 An instance of ``cls`` built from ``data``.
3042
31 Raises:43 Raises:
32 KeyError: If ``data`` contains keys that are not fields of ``cls``.44 ConfigError: If ``data`` holds keys that are not fields of ``cls``, or a
45 value that is invalid for its declared field type.
33 """46 """
34 data = dict(data or {})47 return config_loader.validate_config(
35 known = {f.name for f in fields(cls)}48 cls, dict(data or {}), context=where, error_cls=ConfigError)
36 unknown = sorted(set(data) - known)
37 if unknown:
38 raise KeyError(f"unknown key(s) {unknown} in config section {where!r}; "
39 f"known keys: {sorted(known)}")
40 return cls(**data)
4149
4250
43def load_yaml_mapping(path: str | Path) -> dict[str, Any]:51def load_yaml_mapping(path: str | Path) -> dict[str, Any]:
44 """Loads a YAML file whose top level must be a mapping.52 """Loads a YAML file whose top level must be a mapping.
4553
46 Consumers keep ownership of their top-level and task dataclasses; this only54 Consumers keep ownership of their top-level and task models; this only
47 removes the repeated open/parse/validate boilerplate around them.55 removes the repeated open/parse/validate boilerplate around them.
4856
49 Args:57 Args:
50 path: Path of the YAML file, read as UTF-8.58 path: Path of the YAML file, read as UTF-8.
Importance #3: src/iolabs_ml_harness/__init__.py @@ -7,8 +7,9 @@
7"""7"""
8from typing import TYPE_CHECKING, Any8from typing import TYPE_CHECKING, Any
99
10from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,
11 FactoryConfig,12 FactoryConfig,
12 LossConfig,13 LossConfig,
13 ModelConfig,14 ModelConfig,
14 TrainerConfig,15 TrainerConfig,
Importance #4: src/iolabs_ml_harness/__init__.py @@ -65,8 +66,9 @@
65)66)
6667
67__all__ = [68__all__ = [
68 "DEFAULT_CLASS_COLORS_RGB",69 "DEFAULT_CLASS_COLORS_RGB",
70 "ConfigError",
69 "FactoryConfig",71 "FactoryConfig",
70 "FactoryRegistry",72 "FactoryRegistry",
71 "LossConfig",73 "LossConfig",
72 "MaskOverlayWriter",74 "MaskOverlayWriter",
Importance #5: tests/test_config_registry.py @@ -1,13 +1,15 @@
1"""Strict config construction, YAML loading, and generic registry unit tests."""1"""Strict config construction, YAML loading, and generic registry unit tests."""
2from dataclasses import dataclass, field
3from pathlib import Path2from pathlib import Path
4from types import MappingProxyType3from types import MappingProxyType
5from typing import Any4from typing import Any
65
7import pytest6import pytest
87
8from iolabs.common import config_loader
9
9from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,
10 FactoryConfig,12 FactoryConfig,
11 LossConfig,13 LossConfig,
12 ModelConfig,14 ModelConfig,
13 TrainerConfig,15 TrainerConfig,
Importance #6: tests/test_config_registry.py @@ -16,13 +18,12 @@
16)18)
17from iolabs_ml_harness.registry import FactoryRegistry19from iolabs_ml_harness.registry import FactoryRegistry
1820
1921
20@dataclass22class _Section(config_loader.ConfigModel):
21class _Section:
22 alpha: int = 123 alpha: int = 1
23 beta: str = "b"24 beta: str = "b"
24 extra: dict[str, Any] = field(default_factory=dict)25 extra: dict[str, Any] = {}
2526
2627
27def test_build_section_accepts_any_mapping() -> None:28def test_build_section_accepts_any_mapping() -> None:
28 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")29 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")
Importance #7: tests/test_config_registry.py @@ -34,13 +35,19 @@
34 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})35 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})
3536
3637
37def test_build_section_reports_section_and_known_keys() -> None:38def test_build_section_reports_section_and_known_keys() -> None:
38 with pytest.raises(KeyError) as excinfo:39 with pytest.raises(ConfigError) as excinfo:
39 build_section(_Section, {"gamma": 1, "delta": 2}, "train")40 build_section(_Section, {"gamma": 1, "delta": 2}, "train")
40 message = str(excinfo.value)41 message = str(excinfo.value)
41 assert "unknown key(s) ['delta', 'gamma']" in message42 assert "Unknown train key(s): delta, gamma" in message
42 assert "'train'" in message and "'alpha'" in message43 assert "alpha" in message and "beta" in message
44
45
46def test_config_error_is_a_value_error() -> None:
47 assert issubclass(ConfigError, config_loader.ConfigError)
48 with pytest.raises(ValueError):
49 build_section(_Section, {"alpha": "not-an-int"}, "train")
4350
4451
45def test_legacy_dataclass_defaults_are_unchanged() -> None:52def test_legacy_dataclass_defaults_are_unchanged() -> None:
46 model = ModelConfig()53 model = ModelConfig()
Importance #8: tests/test_config_registry.py @@ -65,9 +72,9 @@
65 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",72 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",
66 "args": {"alpha": 0.8}}, "loss")73 "args": {"alpha": 0.8}}, "loss")
67 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}74 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}
68 assert FactoryConfig(name="x").args == {}75 assert FactoryConfig(name="x").args == {}
69 with pytest.raises(KeyError, match="unknown key"):76 with pytest.raises(ConfigError, match="Unknown loss key"):
70 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")77 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")
7178
7279
73def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:80def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:
Importance #9: tests/test_ml_harness.py @@ -42,9 +42,9 @@
4242
43def test_build_section_rejects_unknown_keys() -> None:43def test_build_section_rejects_unknown_keys() -> None:
44 cfg = build_section(ModelConfig, {"name": "unet", "num_classes": 2}, "model")44 cfg = build_section(ModelConfig, {"name": "unet", "num_classes": 2}, "model")
45 assert cfg.name == "unet" and cfg.num_classes == 245 assert cfg.name == "unet" and cfg.num_classes == 2
46 with pytest.raises(KeyError, match="unknown key.*encoder"):46 with pytest.raises(ValueError, match="Unknown model key.*encoder"):
47 build_section(ModelConfig, {"encoder": "oops"}, "model")47 build_section(ModelConfig, {"encoder": "oops"}, "model")
4848
4949
50def test_model_registry_smp_fallback_and_unknown_name() -> None:50def test_model_registry_smp_fallback_and_unknown_name() -> None:
Importance #10: pyproject.toml @@ -3,15 +3,17 @@
3build-backend = "hatchling.build"3build-backend = "hatchling.build"
44
5[project]5[project]
6name = "iolabs-ml-harness"6name = "iolabs-ml-harness"
7version = "0.2.0"7version = "0.2.1"
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 = [
12 "numpy>=1.26",12 "numpy>=1.26",
13 "pyyaml>=6.0",13 "pyyaml>=6.0",
14 "pydantic>=2.7",
15 "iolabs-common>=0.9.0",
14 "torch>=2.2.0",16 "torch>=2.2.0",
15 "lightning>=2.2",17 "lightning>=2.2",
16 "tensorboard>=2.16",18 "tensorboard>=2.16",
17 "torchmetrics>=1.3",19 "torchmetrics>=1.3",
Importance #11: pyproject.toml @@ -33,8 +35,16 @@
3335
34[tool.uv]36[tool.uv]
35publish-url = "https://nexus.iolabs.ch/repository/pypi-private/"37publish-url = "https://nexus.iolabs.ch/repository/pypi-private/"
3638
39[[tool.uv.index]]
40name = "nexus"
41url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"
42authenticate = "always"
43
44[tool.uv.sources]
45iolabs-common = { index = "nexus" }
46
37[tool.pytest.ini_options]47[tool.pytest.ini_options]
38testpaths = ["tests"]48testpaths = ["tests"]
39markers = [49markers = [
40 "compat_gate: clean-environment 2D installation gate; slow, needs uv and network",50 "compat_gate: clean-environment 2D installation gate; slow, needs uv and network",
Importance #12: README.md @@ -3,13 +3,13 @@
3Task-agnostic PyTorch-Lightning core extracted from the line bitmap segmentation training harness. Consumers such as `linebitmapsegmentation` (2D rasters) and `pointcloud.mlsegmentation` (3D corridors) provide their own datasets, task config, task-specific metrics, and any domain-specific losses.3Task-agnostic PyTorch-Lightning core extracted from the line bitmap segmentation training harness. Consumers such as `linebitmapsegmentation` (2D rasters) and `pointcloud.mlsegmentation` (3D corridors) provide their own datasets, task config, task-specific metrics, and any domain-specific losses.
44
5## Core versus the `image` extra5## Core versus the `image` extra
66
7Core is modality neutral and depends only on `numpy`, `pyyaml`, `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`, `FactoryConfig`, and the legacy model/loss/trainer dataclasses. |11| `config` | core | Strict `build_section`, `load_yaml_mapping`, `ConfigError`, `FactoryConfig`, and the legacy model/loss/trainer `ConfigModel` 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,10 +25,24 @@
25```bash25```bash
26pip install 'iolabs-ml-harness[image]'26pip install 'iolabs-ml-harness[image]'
27```27```
2828
29## Config
30
31Config sections are pydantic models derived from
32`iolabs.common.config_loader.ConfigModel` (unknown keys rejected, values coerced
33by the shared fleet matrix, instances frozen). **Adding a config key = adding one
34field with its default to the model in `config.py`** — there is no separate
35allow-list, coercion helper, or dataclass to keep in sync. `build_section`
36raises `ConfigError`, which derives from `ValueError`.
37
29## Compatibility38## Compatibility
3039
400.2.1 keeps every public name; `build_section` now raises `ConfigError`
41(a `ValueError`) instead of `KeyError` for unknown keys, and the section classes
42are `ConfigModel`s rather than dataclasses — attribute access is unchanged, but
43instances are frozen (use `model_copy(update=...)` instead of assignment).
44
310.2.0 is additive. Existing top-level imports, `iolabs_ml_harness.visualize`, `models.build_model`, `losses.build_loss`, `SegmentationStats`, `SegmentationModule`, and `build_trainer` are unchanged; see `docs/0.2.0-migration.md`.450.2.0 is additive. Existing top-level imports, `iolabs_ml_harness.visualize`, `models.build_model`, `losses.build_loss`, `SegmentationStats`, `SegmentationModule`, and `build_trainer` are unchanged; see `docs/0.2.0-migration.md`.
3246
33## Development47## Development
3448
Importance #14: pyproject.toml @@ -3,15 +3,17 @@
3build-backend = "hatchling.build"3build-backend = "hatchling.build"
44
5[project]5[project]
6name = "iolabs-ml-harness"6name = "iolabs-ml-harness"
7version = "0.2.0"7version = "0.2.1"
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 = [
12 "numpy>=1.26",12 "numpy>=1.26",
13 "pyyaml>=6.0",13 "pyyaml>=6.0",
14 "pydantic>=2.7",
15 "iolabs-common>=0.9.0",
14 "torch>=2.2.0",16 "torch>=2.2.0",
15 "lightning>=2.2",17 "lightning>=2.2",
16 "tensorboard>=2.16",18 "tensorboard>=2.16",
17 "torchmetrics>=1.3",19 "torchmetrics>=1.3",
Importance #15: pyproject.toml @@ -33,8 +35,16 @@
3335
34[tool.uv]36[tool.uv]
35publish-url = "https://nexus.iolabs.ch/repository/pypi-private/"37publish-url = "https://nexus.iolabs.ch/repository/pypi-private/"
3638
39[[tool.uv.index]]
40name = "nexus"
41url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"
42authenticate = "always"
43
44[tool.uv.sources]
45iolabs-common = { index = "nexus" }
46
37[tool.pytest.ini_options]47[tool.pytest.ini_options]
38testpaths = ["tests"]48testpaths = ["tests"]
39markers = [49markers = [
40 "compat_gate: clean-environment 2D installation gate; slow, needs uv and network",50 "compat_gate: clean-environment 2D installation gate; slow, needs uv and network",
Importance #16: src/iolabs_ml_harness/__init__.py @@ -7,8 +7,9 @@
7"""7"""
8from typing import TYPE_CHECKING, Any8from typing import TYPE_CHECKING, Any
99
10from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,
11 FactoryConfig,12 FactoryConfig,
12 LossConfig,13 LossConfig,
13 ModelConfig,14 ModelConfig,
14 TrainerConfig,15 TrainerConfig,
Importance #17: src/iolabs_ml_harness/__init__.py @@ -65,8 +66,9 @@
65)66)
6667
67__all__ = [68__all__ = [
68 "DEFAULT_CLASS_COLORS_RGB",69 "DEFAULT_CLASS_COLORS_RGB",
70 "ConfigError",
69 "FactoryConfig",71 "FactoryConfig",
70 "FactoryRegistry",72 "FactoryRegistry",
71 "LossConfig",73 "LossConfig",
72 "MaskOverlayWriter",74 "MaskOverlayWriter",
Importance #18: src/iolabs_ml_harness/config.py @@ -1,50 +1,58 @@
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.3Task-specific data specs and top-level harness config stay in consumers.
4Compose local dataclasses with ``build_section`` so unknown keys still fail fast.4Compose local `iolabs.common.config_loader.ConfigModel` sections with
5``build_section`` so unknown keys still fail fast.
6
7Adding a knob is a one-line change: add the field (with its default) to the
8model below; nothing else has to be touched.
59
6``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D10``ModelConfig`` is the legacy image schema and is kept unchanged for the 2D
7consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict11consumer. Modality-neutral consumers use ``FactoryConfig`` or their own strict
8dataclass.12`ConfigModel`.
9"""13"""
14import logging
10from collections.abc import Mapping15from collections.abc import Mapping
11from dataclasses import dataclass, field, fields
12from pathlib import Path16from pathlib import Path
13from typing import Any, TypeVar17from typing import Any, Literal, TypeVar
1418
19import pydantic
15import yaml20import yaml
21from iolabs.common import config_loader
22
23logger = logging.getLogger(__name__)
24
25T = TypeVar("T", bound=config_loader.ConfigModel)
26
1627
17T = TypeVar("T")28class ConfigError(config_loader.ConfigError):
29 """Raised when a harness config section holds unknown keys or bad values."""
1830
1931
20def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:32def build_section(cls: type[T], data: Mapping[str, Any] | None, where: str) -> T:
21 """Builds a config dataclass from a mapping, rejecting unknown keys.33 """Builds a config model from a mapping, rejecting unknown keys.
2234
23 Args:35 Args:
24 cls: Dataclass to instantiate.36 cls: `ConfigModel` subclass to instantiate.
25 data: Mapping of field name to value, or ``None`` for all defaults.37 data: Mapping of field name to value, or ``None`` for all defaults.
26 where: Section name used in the error message, e.g. ``"model"``.38 where: Section name used in the error message, e.g. ``"model"``.
2739
28 Returns:40 Returns:
29 An instance of ``cls`` built from ``data``.41 An instance of ``cls`` built from ``data``.
3042
31 Raises:43 Raises:
32 KeyError: If ``data`` contains keys that are not fields of ``cls``.44 ConfigError: If ``data`` holds keys that are not fields of ``cls``, or a
45 value that is invalid for its declared field type.
33 """46 """
34 data = dict(data or {})47 return config_loader.validate_config(
35 known = {f.name for f in fields(cls)}48 cls, dict(data or {}), context=where, error_cls=ConfigError)
36 unknown = sorted(set(data) - known)
37 if unknown:
38 raise KeyError(f"unknown key(s) {unknown} in config section {where!r}; "
39 f"known keys: {sorted(known)}")
40 return cls(**data)
4149
4250
43def load_yaml_mapping(path: str | Path) -> dict[str, Any]:51def load_yaml_mapping(path: str | Path) -> dict[str, Any]:
44 """Loads a YAML file whose top level must be a mapping.52 """Loads a YAML file whose top level must be a mapping.
4553
46 Consumers keep ownership of their top-level and task dataclasses; this only54 Consumers keep ownership of their top-level and task models; this only
47 removes the repeated open/parse/validate boilerplate around them.55 removes the repeated open/parse/validate boilerplate around them.
4856
49 Args:57 Args:
50 path: Path of the YAML file, read as UTF-8.58 path: Path of the YAML file, read as UTF-8.
Importance #19: src/iolabs_ml_harness/config.py @@ -53,62 +61,60 @@
53 The parsed mapping. An empty document yields an empty dict.61 The parsed mapping. An empty document yields an empty dict.
5462
55 Raises:63 Raises:
56 FileNotFoundError: If ``path`` does not exist.64 FileNotFoundError: If ``path`` does not exist.
57 ValueError: If the document's top level is not a mapping.65 ConfigError: If the document's top level is not a mapping. The class
66 derives from ``ValueError``, so legacy handlers keep working.
58 """67 """
59 config_path = Path(path)68 config_path = Path(path)
60 with config_path.open("r", encoding="utf-8") as handle:69 with config_path.open("r", encoding="utf-8") as handle:
61 loaded = yaml.safe_load(handle)70 loaded = yaml.safe_load(handle)
62 if loaded is None:71 if loaded is None:
63 return {}72 return {}
64 if not isinstance(loaded, Mapping):73 if not isinstance(loaded, Mapping):
65 raise ValueError(74 raise ConfigError(
66 f"config {str(config_path)!r} must contain a top-level mapping, "75 f"config {str(config_path)!r} must contain a top-level mapping, "
67 f"got {type(loaded).__name__}")76 f"got {type(loaded).__name__}")
68 return dict(loaded)77 return dict(loaded)
6978
7079
71@dataclass80class FactoryConfig(config_loader.ConfigModel):
72class FactoryConfig:
73 """Modality-neutral registry selection: a name plus factory keyword args."""81 """Modality-neutral registry selection: a name plus factory keyword args."""
74 name: str82 name: str
75 args: dict[str, Any] = field(default_factory=dict)83 args: dict[str, Any] = {}
7684
7785
78@dataclass86class ModelConfig(config_loader.ConfigModel):
79class ModelConfig:
80 """Legacy image-model schema (segmentation-models-pytorch shaped)."""87 """Legacy image-model schema (segmentation-models-pytorch shaped)."""
81 name: str = "unet"88 name: str = "unet"
82 encoder_name: str = "resnet18"89 encoder_name: str = "resnet18"
83 encoder_weights: str | None = "imagenet" # None = train from scratch90 encoder_weights: str | None = "imagenet" # None = train from scratch
84 in_channels: int = 191 in_channels: int = pydantic.Field(default=1, gt=0)
85 num_classes: int = 392 num_classes: int = pydantic.Field(default=3, gt=0)
86 extra: dict[str, Any] = field(default_factory=dict) # passed to the factory93 extra: dict[str, Any] = {} # passed to the factory
8794
8895
89@dataclass96class LossConfig(config_loader.ConfigModel):
90class LossConfig:
91 """Loss selection by registry name plus factory keyword arguments."""97 """Loss selection by registry name plus factory keyword arguments."""
92 name: str = "dice_focal"98 name: str = "dice_focal"
93 args: dict[str, Any] = field(default_factory=dict)99 args: dict[str, Any] = {}
94100
95101
96@dataclass102class TrainerConfig(config_loader.ConfigModel):
97class TrainerConfig:
98 """Harness-visible Lightning trainer, logger, and callback knobs."""103 """Harness-visible Lightning trainer, logger, and callback knobs."""
99 max_epochs: int = -1 # -1 = no cap; early stopping ends training instead104 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
100 lr: float = 3.0e-4105 lr: float = pydantic.Field(default=3.0e-4, gt=0)
101 weight_decay: float = 1.0e-4106 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
102 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU107 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU
103 accumulate_grad_batches: int = 1 # match effective batch across experiments108 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
104 accelerator: str = "auto"109 accelerator: str = "auto"
105 devices: int | str = 1110 devices: int | str = 1
106 viz_every_n_epochs: int = 2111 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
107 viz_samples: int = 4112 viz_samples: int = pydantic.Field(default=4, ge=0)
108 monitor: str = "val/f1_mean_fg"113 monitor: str = "val/f1_mean_fg"
109 monitor_mode: str = "max"114 monitor_mode: Literal["max", "min"] = "max"
110 early_stop_monitor: str = "val/loss" # stop when this stops improving115 early_stop_monitor: str = "val/loss" # stop when this stops improving
111 early_stop_mode: str = "min"116 early_stop_mode: Literal["min", "max"] = "min"
112 early_stop_patience: int = 4 # epochs without improvement before stopping; 0 disables117 # epochs without improvement before stopping; 0 disables
118 early_stop_patience: int = pydantic.Field(default=4, ge=0)
113 log_dir: str = "runs"119 log_dir: str = "runs"
114 log_every_n_steps: int = 10120 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
Importance #20: tests/test_config_registry.py @@ -1,13 +1,15 @@
1"""Strict config construction, YAML loading, and generic registry unit tests."""1"""Strict config construction, YAML loading, and generic registry unit tests."""
2from dataclasses import dataclass, field
3from pathlib import Path2from pathlib import Path
4from types import MappingProxyType3from types import MappingProxyType
5from typing import Any4from typing import Any
65
7import pytest6import pytest
87
8from iolabs.common import config_loader
9
9from iolabs_ml_harness.config import (10from iolabs_ml_harness.config import (
11 ConfigError,
10 FactoryConfig,12 FactoryConfig,
11 LossConfig,13 LossConfig,
12 ModelConfig,14 ModelConfig,
13 TrainerConfig,15 TrainerConfig,
Importance #21: tests/test_config_registry.py @@ -16,13 +18,12 @@
16)18)
17from iolabs_ml_harness.registry import FactoryRegistry19from iolabs_ml_harness.registry import FactoryRegistry
1820
1921
20@dataclass22class _Section(config_loader.ConfigModel):
21class _Section:
22 alpha: int = 123 alpha: int = 1
23 beta: str = "b"24 beta: str = "b"
24 extra: dict[str, Any] = field(default_factory=dict)25 extra: dict[str, Any] = {}
2526
2627
27def test_build_section_accepts_any_mapping() -> None:28def test_build_section_accepts_any_mapping() -> None:
28 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")29 cfg = build_section(_Section, MappingProxyType({"alpha": 7}), "section")
Importance #22: tests/test_config_registry.py @@ -34,13 +35,19 @@
34 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})35 assert (cfg.alpha, cfg.beta, cfg.extra) == (1, "b", {})
3536
3637
37def test_build_section_reports_section_and_known_keys() -> None:38def test_build_section_reports_section_and_known_keys() -> None:
38 with pytest.raises(KeyError) as excinfo:39 with pytest.raises(ConfigError) as excinfo:
39 build_section(_Section, {"gamma": 1, "delta": 2}, "train")40 build_section(_Section, {"gamma": 1, "delta": 2}, "train")
40 message = str(excinfo.value)41 message = str(excinfo.value)
41 assert "unknown key(s) ['delta', 'gamma']" in message42 assert "Unknown train key(s): delta, gamma" in message
42 assert "'train'" in message and "'alpha'" in message43 assert "alpha" in message and "beta" in message
44
45
46def test_config_error_is_a_value_error() -> None:
47 assert issubclass(ConfigError, config_loader.ConfigError)
48 with pytest.raises(ValueError):
49 build_section(_Section, {"alpha": "not-an-int"}, "train")
4350
4451
45def test_legacy_dataclass_defaults_are_unchanged() -> None:52def test_legacy_dataclass_defaults_are_unchanged() -> None:
46 model = ModelConfig()53 model = ModelConfig()
Importance #23: tests/test_config_registry.py @@ -65,9 +72,9 @@
65 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",72 cfg = build_section(FactoryConfig, {"name": "masked_focal_tversky",
66 "args": {"alpha": 0.8}}, "loss")73 "args": {"alpha": 0.8}}, "loss")
67 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}74 assert cfg.name == "masked_focal_tversky" and cfg.args == {"alpha": 0.8}
68 assert FactoryConfig(name="x").args == {}75 assert FactoryConfig(name="x").args == {}
69 with pytest.raises(KeyError, match="unknown key"):76 with pytest.raises(ConfigError, match="Unknown loss key"):
70 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")77 build_section(FactoryConfig, {"name": "x", "kwargs": {}}, "loss")
7178
7279
73def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:80def test_load_yaml_mapping_round_trip(tmp_path: Path) -> None:
Importance #24: tests/test_ml_harness.py @@ -42,9 +42,9 @@
4242
43def test_build_section_rejects_unknown_keys() -> None:43def test_build_section_rejects_unknown_keys() -> None:
44 cfg = build_section(ModelConfig, {"name": "unet", "num_classes": 2}, "model")44 cfg = build_section(ModelConfig, {"name": "unet", "num_classes": 2}, "model")
45 assert cfg.name == "unet" and cfg.num_classes == 245 assert cfg.name == "unet" and cfg.num_classes == 2
46 with pytest.raises(KeyError, match="unknown key.*encoder"):46 with pytest.raises(ValueError, match="Unknown model key.*encoder"):
47 build_section(ModelConfig, {"encoder": "oops"}, "model")47 build_section(ModelConfig, {"encoder": "oops"}, "model")
4848
4949
50def test_model_registry_smp_fallback_and_unknown_name() -> None:50def test_model_registry_smp_fallback_and_unknown_name() -> None: