Back to report index

linebitmapimagerasterizer (ML) 9b68a51: AI3D-379 Align config module with fleet pattern

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

Commit #83 · 23 snippets

 README.md                    |  14 ++++
 pyproject.toml               |   3 +
 scripts/train.py             |   2 +-
 src/train/config.py          | 155 ++++++++++++-------------------------------
 test/test_config.py          |  99 +++++++++++++++++++++++++++
 test/test_train_overrides.py |   8 +--
 6 files changed, 164 insertions(+), 117 deletions(-)
Importance #1: src/train/config.py @@ -1,72 +1,51 @@
1"""YAML-backed configuration for the training harness.1"""YAML-backed configuration for the line-bitmap training harness.
22
3Pydantic models on `iolabs.common.config_loader.ConfigModel` + yaml. Unknown keys3The schema is `HarnessConfig` (a `config_loader.ConfigModel`), mirroring a
4raise, so config typos fail fast instead of silently training with defaults, and4harness YAML file (`configs/*.yaml`) key for key. The task-specific `DataConfig`
5values are coerced by the shared fleet matrix.5is owned here; the generic ``model``/``loss``/``train`` sections are the shared
66models from `iolabs_ml_harness.config`.
7Adding a config key = adding one field with its default to the model below.7
8Instances are frozen and validated: derive a changed config with8Adding a config key means adding the field (with its type, default and any
9:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips9`Field` range) to the model below and the same key to the YAML -- nothing else.
10validation and the fleet coercion matrix.10Unknown keys are rejected, so config typos fail fast instead of silently
11training with defaults.
12
13`HarnessConfig.from_yaml` and `with_overrides` return the frozen
14`HarnessConfig`; instances are frozen and validated, so a changed config is
15derived with :func:`with_overrides`, never ``model_copy(update=...)`` -- the
16latter skips validation and the fleet coercion matrix.
11"""17"""
18import logging
12from collections.abc import Mapping19from collections.abc import Mapping
13from pathlib import Path20from pathlib import Path
14from typing import Any, Literal, get_args21from typing import Any, Literal
1522
16import pydantic23import pydantic
17import yaml
18from iolabs.common import config_loader24from iolabs.common import config_loader
19from pydantic import fields as pydantic_fields25from iolabs_ml_harness import config as harness_config
26
27logger = logging.getLogger(__name__)
28
29_CONTEXT = "line bitmap training config"
30_OVERRIDES_CONTEXT = "line bitmap training config overrides"
2031
21_STROKE_KINDS = frozenset({"solid", "dashed"})32_STROKE_KINDS = frozenset({"solid", "dashed"})
22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")33_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")
23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")34_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")
2435
2536
26class ConfigError(config_loader.ConfigError):37class HarnessConfigError(config_loader.ConfigError):
27 """Raised when a harness config holds unknown keys or invalid values."""38 """Raised when line bitmap training config contains unsupported keys or values."""
28
2939
30class SectionModel(config_loader.ConfigModel):
31 """`ConfigModel` in which a YAML ``null`` means "use this field's default".
3240
33 A bare ``data:`` / ``args:`` line parses to ``None``; the hand-written41class PairSpec(harness_config.SectionModel):
34 ``raw.pop(..., {})`` coalescing that predates the pydantic layer treated
35 that as "section omitted". Fields that accept ``None`` (e.g.
36 ``encoder_weights``) keep it as a real value, and required fields still
37 fail. Mirrors ``iolabs_ml_harness.config.SectionModel``, which this repo
38 cannot import (no harness dependency).
39 """
40
41 @pydantic.model_validator(mode="before")
42 @classmethod
43 def _null_means_default(cls, data: Any) -> Any:
44 """Drop ``None`` entries whose field has a default and forbids ``None``."""
45 if not isinstance(data, Mapping):
46 return data
47 dropped = {
48 name for name, value in data.items()
49 if value is None and _null_means_default_for(cls.model_fields.get(name))}
50 if not dropped:
51 return data
52 return {name: value for name, value in data.items() if name not in dropped}
53
54
55def _null_means_default_for(field: pydantic_fields.FieldInfo | None) -> bool:
56 """True when *field* has a default and its annotation does not accept None."""
57 if field is None or field.is_required():
58 return False
59 return type(None) not in get_args(field.annotation)
60
61
62class PairSpec(SectionModel):
63 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""42 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""
64 images: str43 images: str
65 masks: str44 masks: str
6645
6746
68class DataConfig(SectionModel):47class DataConfig(harness_config.SectionModel):
69 """Tile corpus, split, crop sampling and label rasterization knobs."""48 """Tile corpus, split, crop sampling and label rasterization knobs."""
70 pairs: list[PairSpec] = []49 pairs: list[PairSpec] = []
71 # Optional explicit, pre-split directories (e.g. the symlink folders under50 # Optional explicit, pre-split directories (e.g. the symlink folders under
72 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_51 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_
Importance #2: src/train/config.py @@ -166,64 +145,16 @@
166 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)145 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)
167 return data_cfg.model_copy(update=updates)146 return data_cfg.model_copy(update=updates)
168147
169148
170class ModelConfig(SectionModel):149class HarnessConfig(harness_config.SectionModel):
171 """Segmentation-models-pytorch architecture/encoder selection."""
172 name: str = "unet"
173 encoder_name: str = "resnet18"
174 encoder_weights: str | None = "imagenet" # None = train from scratch
175 in_channels: int = pydantic.Field(default=1, gt=0)
176 num_classes: int = pydantic.Field(default=3, gt=0)
177 extra: dict[str, Any] = {} # passed through to the model factory
178
179
180class LossConfig(SectionModel):
181 """Loss selection by registry name plus factory keyword arguments."""
182 name: str = "dice_focal"
183 args: dict[str, Any] = {}
184
185
186class TrainerConfig(SectionModel):
187 """Lightning trainer, logger, and callback knobs."""
188 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
189 lr: float = pydantic.Field(default=3.0e-4, gt=0)
190 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
191 # "auto" -> 16-mixed on CUDA, 32-true on CPU; else a Lightning precision
192 # (16, "16-mixed", "bf16-mixed", 32, "32-true", ...)
193 precision: int | str = "auto"
194 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
195 accelerator: str = "auto"
196 devices: int | str = 1
197 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
198 viz_samples: int = pydantic.Field(default=4, ge=0)
199 monitor: str = "val/f1_mean_fg"
200 monitor_mode: Literal["max", "min"] = "max"
201 early_stop_monitor: str = "val/loss" # stop when this stops improving
202 early_stop_mode: Literal["min", "max"] = "min"
203 # epochs without improvement before stopping; 0 disables
204 early_stop_patience: int = pydantic.Field(default=4, ge=0)
205 log_dir: str = "runs"
206 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
207
208 @pydantic.field_validator("precision", "devices", mode="before")
209 @classmethod
210 def _reject_bool(cls, value: Any, info: pydantic.ValidationInfo) -> Any:
211 """Keep YAML ``yes``/``on`` from silently narrowing to ``1``."""
212 if isinstance(value, bool):
213 raise ValueError(
214 f"{info.field_name} must be an int or a str, got bool {value!r}")
215 return value
216
217
218class HarnessConfig(SectionModel):
219 """Top-level training config: one YAML file, one instance."""150 """Top-level training config: one YAML file, one instance."""
220 experiment: str = "experiment"151 experiment: str = "experiment"
221 seed: int = 1337152 seed: int = 1337
222 data: DataConfig = DataConfig()153 data: DataConfig = DataConfig()
223 model: ModelConfig = ModelConfig()154 model: harness_config.ModelConfig = harness_config.ModelConfig()
224 loss: LossConfig = LossConfig()155 loss: harness_config.LossConfig = harness_config.LossConfig()
225 train: TrainerConfig = TrainerConfig()156 train: harness_config.TrainerConfig = harness_config.TrainerConfig()
226157
227 @classmethod158 @classmethod
228 def from_yaml(cls, path: str | Path) -> "HarnessConfig":159 def from_yaml(cls, path: str | Path) -> "HarnessConfig":
229 """Loads and validates a harness YAML config.160 """Loads and validates a harness YAML config.
Importance #3: src/train/config.py @@ -235,19 +166,18 @@
235 The validated, frozen config.166 The validated, frozen config.
236167
237 Raises:168 Raises:
238 FileNotFoundError: If ``path`` does not exist.169 FileNotFoundError: If ``path`` does not exist.
239 ConfigError: If the document is not a mapping, holds an unknown key,170 config_loader.ConfigError: If the document is not a mapping, holds
240 or holds a value invalid for its field. Derives from171 an unknown key, or holds a value invalid for its field. Unknown
241 ``ValueError``.172 keys and bad values raise ``HarnessConfigError``; a non-mapping
173 document raises the shared harness error class. Both derive
174 from ``ValueError``.
242 """175 """
243 raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}176 raw = harness_config.load_yaml_mapping(path)
244 if not isinstance(raw, Mapping):
245 raise ConfigError(
246 f"config {str(path)!r} must contain a top-level mapping, "
247 f"got {type(raw).__name__}")
248 return config_loader.validate_config(177 return config_loader.validate_config(
249 cls, raw, context=str(path), error_cls=ConfigError)178 cls, raw, context=f"{_CONTEXT} {str(path)!r}",
179 error_cls=HarnessConfigError)
250180
251181
252def with_overrides(cfg: HarnessConfig,182def with_overrides(cfg: HarnessConfig,
253 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:183 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
Importance #4: src/train/config.py @@ -265,12 +195,13 @@
265 Returns:195 Returns:
266 ``cfg`` itself when no override is given, otherwise a validated copy.196 ``cfg`` itself when no override is given, otherwise a validated copy.
267197
268 Raises:198 Raises:
269 ConfigError: An override value is invalid for its declared field.199 HarnessConfigError: An override value is invalid for its declared field.
270 """200 """
271 applied = {name: dict(values) for name, values in sections.items() if values}201 applied = {name: dict(values) for name, values in sections.items() if values}
272 if not applied:202 if not applied:
273 return cfg203 return cfg
204 logger.info("Config overrides applied: %s", ", ".join(sorted(applied)))
274 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)205 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
275 return config_loader.validate_config(206 return config_loader.validate_config(
276 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)207 type(cfg), merged, context=_OVERRIDES_CONTEXT, error_cls=HarnessConfigError)
Importance #5: scripts/train.py @@ -69,9 +69,9 @@
69 Returns:69 Returns:
70 ``cfg`` itself when no override was given, otherwise an updated copy.70 ``cfg`` itself when no override was given, otherwise an updated copy.
7171
72 Raises:72 Raises:
73 config.ConfigError: An override value is invalid for its field, e.g.73 config.HarnessConfigError: An override value is invalid for its field, e.g.
74 ``--max-epochs -2`` or ``--batch-size 0``.74 ``--max-epochs -2`` or ``--batch-size 0``.
75 """75 """
76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
77 if args.log_dir:77 if args.log_dir:
Importance #6: test/test_config.py @@ -0,0 +1,99 @@
1"""Config-schema tests for the line-bitmap training harness.
2
3The harness config is YAML-backed and has no packaged JSON defaults, so the
4fleet parity test is replaced by a defaults-vs-model check on the shipped
5baseline YAML's schema-level behaviour.
6"""
7from pathlib import Path
8
9import pytest
10import yaml
11from iolabs.common import config_loader
12from iolabs_ml_harness import config as harness_config
13
14from src.train import config
15
16
17def _write(tmp_path: Path, text: str) -> Path:
18 path = tmp_path / "cfg.yaml"
19 path.write_text(text, encoding="utf-8")
20 return path
21
22
23def test_error_class_is_config_error() -> None:
24 assert issubclass(config.HarnessConfigError, config_loader.ConfigError)
25 assert issubclass(config.HarnessConfigError, ValueError)
26
27
28def test_shared_sections_come_from_the_harness_package() -> None:
29 """model/loss/train are the shared models, not repo-local copies."""
30 fields = config.HarnessConfig.model_fields
31 assert fields["model"].annotation is harness_config.ModelConfig
32 assert fields["loss"].annotation is harness_config.LossConfig
33 assert fields["train"].annotation is harness_config.TrainerConfig
34 assert issubclass(config.DataConfig, harness_config.SectionModel)
35
36
37def test_model_defaults_match_documented_defaults() -> None:
38 dumped = config.HarnessConfig().model_dump()
39
40 assert dumped == {
41 "experiment": "experiment",
42 "seed": 1337,
43 "data": {
44 "pairs": [], "val_pairs": [], "test_pairs": [],
45 "crop_size": 512, "batch_size": 16, "num_workers": 4,
46 "val_fraction": 0.15, "crops_per_tile": 4, "pos_crop_prob": 0.7,
47 "min_valid_fraction": 0.10, "augment": True,
48 "label_source": "rendered", "label_stroke_px": 4,
49 "review_statuses": ["ok"]},
50 "model": harness_config.ModelConfig().model_dump(),
51 "loss": harness_config.LossConfig().model_dump(),
52 "train": harness_config.TrainerConfig().model_dump()}
53
54
55def test_from_yaml_returns_defaults_for_an_empty_document(tmp_path: Path) -> None:
56 assert config.HarnessConfig.from_yaml(_write(tmp_path, "")) == config.HarnessConfig()
57
58
59def test_unknown_top_level_key_is_rejected(tmp_path: Path) -> None:
60 with pytest.raises(config.HarnessConfigError, match="experimnt"):
61 config.HarnessConfig.from_yaml(_write(tmp_path, "experimnt: t\n"))
62
63
64def test_unknown_nested_key_is_rejected(tmp_path: Path) -> None:
65 with pytest.raises(config.HarnessConfigError, match="crop_sizes"):
66 config.HarnessConfig.from_yaml(_write(tmp_path, "data:\n crop_sizes: 64\n"))
67
68
69def test_non_mapping_document_is_rejected(tmp_path: Path) -> None:
70 with pytest.raises(config_loader.ConfigError, match="mapping"):
71 config.HarnessConfig.from_yaml(_write(tmp_path, "- a\n- b\n"))
72
73
74def test_overrides_deep_merge_onto_defaults() -> None:
75 cfg = config.HarnessConfig(experiment="e")
76
77 updated = config.with_overrides(
78 cfg, {"data": {"batch_size": 2}, "train": {"max_epochs": 3}, "loss": {}})
79
80 assert updated.data.batch_size == 2 and updated.train.max_epochs == 3
81 assert updated.experiment == "e" and updated.data.crop_size == 512
82 assert cfg.data.batch_size == 16 # source untouched
83
84
85def test_override_coercion_and_rejection() -> None:
86 updated = config.with_overrides(config.HarnessConfig(),
87 {"data": {"batch_size": "4"}})
88 assert updated.data.batch_size == 4
89
90 with pytest.raises(config.HarnessConfigError, match="batch_size"):
91 config.with_overrides(config.HarnessConfig(), {"data": {"batch_size": 0}})
92
93
94def test_shipped_configs_parse() -> None:
95 for path in sorted(Path("configs").glob("*.yaml")):
96 raw = yaml.safe_load(path.read_text(encoding="utf-8"))
97 if not isinstance(raw, dict) or "experiment" not in raw:
98 continue # non-harness config in the same folder
99 assert config.HarnessConfig.from_yaml(path).experiment
0
Importance #7: test/test_train_overrides.py @@ -5,11 +5,11 @@
5from pathlib import Path5from pathlib import Path
66
7import pytest7import pytest
88
9from iolabs_ml_harness.config import TrainerConfig
9from src.train.config import (10from src.train.config import (
10 ConfigError, DataConfig, HarnessConfig, PairSpec, TrainerConfig,11 DataConfig, HarnessConfig, HarnessConfigError, PairSpec, reroot_data_paths)
11 reroot_data_paths)
1212
13_REPO_ROOT = Path(__file__).resolve().parents[1]13_REPO_ROOT = Path(__file__).resolve().parents[1]
14_CLI_FLAGS = ("log_dir", "num_workers", "max_epochs", "batch_size", "model", "encoder")14_CLI_FLAGS = ("log_dir", "num_workers", "max_epochs", "batch_size", "model", "encoder")
1515
Importance #8: test/test_train_overrides.py @@ -153,9 +153,9 @@
153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},
154 {"num_workers": -1}])154 {"num_workers": -1}])
155def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:155def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:
156 """CLI overrides must hit the same bounds as values coming from YAML."""156 """CLI overrides must hit the same bounds as values coming from YAML."""
157 with pytest.raises(ConfigError):157 with pytest.raises(HarnessConfigError):
158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
159159
160160
161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
Importance #9: test/test_train_overrides.py @@ -197,6 +197,6 @@
197def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str, tmp_path: Path) -> None:197def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str, tmp_path: Path) -> None:
198 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""198 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""
199 path = tmp_path / "bool.yaml"199 path = tmp_path / "bool.yaml"
200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")
201 with pytest.raises(ConfigError, match="got bool"):201 with pytest.raises(HarnessConfigError, match="got bool"):
202 HarnessConfig.from_yaml(path)202 HarnessConfig.from_yaml(path)
Importance #10: pyproject.toml @@ -33,8 +33,10 @@
33 "lightning>=2.2",33 "lightning>=2.2",
34 "tensorboard>=2.16",34 "tensorboard>=2.16",
35 "torchmetrics>=1.3",35 "torchmetrics>=1.3",
36 "pyyaml>=6.0",36 "pyyaml>=6.0",
37 # Shared training-harness config sections (SectionModel, model/loss/train).
38 "iolabs-ml-harness>=0.2.1",
37 # Single source of truth for inference (shared FlipTTA, predict, vectorize).39 # Single source of truth for inference (shared FlipTTA, predict, vectorize).
38 "iolabs-image-analyzer-line-bitmap-inference>=0.1.1",40 "iolabs-image-analyzer-line-bitmap-inference>=0.1.1",
39]41]
4042
Importance #11: pyproject.toml @@ -52,8 +54,9 @@
5254
53[tool.uv.sources]55[tool.uv.sources]
54iolabs-image-analyzer-line-bitmap-inference = { index = "nexus" }56iolabs-image-analyzer-line-bitmap-inference = { index = "nexus" }
55iolabs-common = { index = "nexus" }57iolabs-common = { index = "nexus" }
58iolabs-ml-harness = { index = "nexus" }
56iolabs-logstash = { index = "nexus" }59iolabs-logstash = { index = "nexus" }
5760
58[build-system]61[build-system]
59requires = ["hatchling"]62requires = ["hatchling"]
Importance #12: README.md @@ -101,8 +101,22 @@
101101
102Still planned from the harness plan: `labels`, `review`, `registry`, `eval`,102Still planned from the harness plan: `labels`, `review`, `registry`, `eval`,
103`infer`, `vectorize`, `backproject`, `raster_frame`.103`infer`, `vectorize`, `backproject`, `raster_frame`.
104104
105## Configuration
106
107Defaults live in the harness YAML files under `configs/` (e.g.
108`configs/unet_baseline.yaml`). The schema is `HarnessConfig` in
109`src/train/config.py` (a `config_loader.ConfigModel`); nested YAML sections are
110nested models and unknown keys are rejected. The `data` section is repo-owned;
111`model`, `loss` and `train` are the shared sections from
112`iolabs_ml_harness.config`. **To add a config key: add the field (with its type,
113default and any `Field` range) to the model and the same key with the same
114default to the YAML — nothing else.** `HarnessConfig.from_yaml` and
115`with_overrides` return the frozen `HarnessConfig`. Runtime overrides come from
116`scripts/train.py` CLI flags (`--model`, `--encoder`, `--max-epochs`, …), never
117repo-local edits to a shipped config.
118
105## Setup119## Setup
106120
107### Prerequisites121### Prerequisites
108122
Importance #13: pyproject.toml @@ -33,8 +33,10 @@
33 "lightning>=2.2",33 "lightning>=2.2",
34 "tensorboard>=2.16",34 "tensorboard>=2.16",
35 "torchmetrics>=1.3",35 "torchmetrics>=1.3",
36 "pyyaml>=6.0",36 "pyyaml>=6.0",
37 # Shared training-harness config sections (SectionModel, model/loss/train).
38 "iolabs-ml-harness>=0.2.1",
37 # Single source of truth for inference (shared FlipTTA, predict, vectorize).39 # Single source of truth for inference (shared FlipTTA, predict, vectorize).
38 "iolabs-image-analyzer-line-bitmap-inference>=0.1.1",40 "iolabs-image-analyzer-line-bitmap-inference>=0.1.1",
39]41]
4042
Importance #14: pyproject.toml @@ -52,8 +54,9 @@
5254
53[tool.uv.sources]55[tool.uv.sources]
54iolabs-image-analyzer-line-bitmap-inference = { index = "nexus" }56iolabs-image-analyzer-line-bitmap-inference = { index = "nexus" }
55iolabs-common = { index = "nexus" }57iolabs-common = { index = "nexus" }
58iolabs-ml-harness = { index = "nexus" }
56iolabs-logstash = { index = "nexus" }59iolabs-logstash = { index = "nexus" }
5760
58[build-system]61[build-system]
59requires = ["hatchling"]62requires = ["hatchling"]
Importance #15: scripts/train.py @@ -69,9 +69,9 @@
69 Returns:69 Returns:
70 ``cfg`` itself when no override was given, otherwise an updated copy.70 ``cfg`` itself when no override was given, otherwise an updated copy.
7171
72 Raises:72 Raises:
73 config.ConfigError: An override value is invalid for its field, e.g.73 config.HarnessConfigError: An override value is invalid for its field, e.g.
74 ``--max-epochs -2`` or ``--batch-size 0``.74 ``--max-epochs -2`` or ``--batch-size 0``.
75 """75 """
76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
77 if args.log_dir:77 if args.log_dir:
Importance #16: src/train/config.py @@ -1,72 +1,51 @@
1"""YAML-backed configuration for the training harness.1"""YAML-backed configuration for the line-bitmap training harness.
22
3Pydantic models on `iolabs.common.config_loader.ConfigModel` + yaml. Unknown keys3The schema is `HarnessConfig` (a `config_loader.ConfigModel`), mirroring a
4raise, so config typos fail fast instead of silently training with defaults, and4harness YAML file (`configs/*.yaml`) key for key. The task-specific `DataConfig`
5values are coerced by the shared fleet matrix.5is owned here; the generic ``model``/``loss``/``train`` sections are the shared
66models from `iolabs_ml_harness.config`.
7Adding a config key = adding one field with its default to the model below.7
8Instances are frozen and validated: derive a changed config with8Adding a config key means adding the field (with its type, default and any
9:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips9`Field` range) to the model below and the same key to the YAML -- nothing else.
10validation and the fleet coercion matrix.10Unknown keys are rejected, so config typos fail fast instead of silently
11training with defaults.
12
13`HarnessConfig.from_yaml` and `with_overrides` return the frozen
14`HarnessConfig`; instances are frozen and validated, so a changed config is
15derived with :func:`with_overrides`, never ``model_copy(update=...)`` -- the
16latter skips validation and the fleet coercion matrix.
11"""17"""
18import logging
12from collections.abc import Mapping19from collections.abc import Mapping
13from pathlib import Path20from pathlib import Path
14from typing import Any, Literal, get_args21from typing import Any, Literal
1522
16import pydantic23import pydantic
17import yaml
18from iolabs.common import config_loader24from iolabs.common import config_loader
19from pydantic import fields as pydantic_fields25from iolabs_ml_harness import config as harness_config
26
27logger = logging.getLogger(__name__)
28
29_CONTEXT = "line bitmap training config"
30_OVERRIDES_CONTEXT = "line bitmap training config overrides"
2031
21_STROKE_KINDS = frozenset({"solid", "dashed"})32_STROKE_KINDS = frozenset({"solid", "dashed"})
22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")33_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")
23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")34_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")
2435
2536
26class ConfigError(config_loader.ConfigError):37class HarnessConfigError(config_loader.ConfigError):
27 """Raised when a harness config holds unknown keys or invalid values."""38 """Raised when line bitmap training config contains unsupported keys or values."""
28
2939
30class SectionModel(config_loader.ConfigModel):
31 """`ConfigModel` in which a YAML ``null`` means "use this field's default".
3240
33 A bare ``data:`` / ``args:`` line parses to ``None``; the hand-written41class PairSpec(harness_config.SectionModel):
34 ``raw.pop(..., {})`` coalescing that predates the pydantic layer treated
35 that as "section omitted". Fields that accept ``None`` (e.g.
36 ``encoder_weights``) keep it as a real value, and required fields still
37 fail. Mirrors ``iolabs_ml_harness.config.SectionModel``, which this repo
38 cannot import (no harness dependency).
39 """
40
41 @pydantic.model_validator(mode="before")
42 @classmethod
43 def _null_means_default(cls, data: Any) -> Any:
44 """Drop ``None`` entries whose field has a default and forbids ``None``."""
45 if not isinstance(data, Mapping):
46 return data
47 dropped = {
48 name for name, value in data.items()
49 if value is None and _null_means_default_for(cls.model_fields.get(name))}
50 if not dropped:
51 return data
52 return {name: value for name, value in data.items() if name not in dropped}
53
54
55def _null_means_default_for(field: pydantic_fields.FieldInfo | None) -> bool:
56 """True when *field* has a default and its annotation does not accept None."""
57 if field is None or field.is_required():
58 return False
59 return type(None) not in get_args(field.annotation)
60
61
62class PairSpec(SectionModel):
63 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""42 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""
64 images: str43 images: str
65 masks: str44 masks: str
6645
6746
68class DataConfig(SectionModel):47class DataConfig(harness_config.SectionModel):
69 """Tile corpus, split, crop sampling and label rasterization knobs."""48 """Tile corpus, split, crop sampling and label rasterization knobs."""
70 pairs: list[PairSpec] = []49 pairs: list[PairSpec] = []
71 # Optional explicit, pre-split directories (e.g. the symlink folders under50 # Optional explicit, pre-split directories (e.g. the symlink folders under
72 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_51 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_
Importance #17: src/train/config.py @@ -166,64 +145,16 @@
166 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)145 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)
167 return data_cfg.model_copy(update=updates)146 return data_cfg.model_copy(update=updates)
168147
169148
170class ModelConfig(SectionModel):149class HarnessConfig(harness_config.SectionModel):
171 """Segmentation-models-pytorch architecture/encoder selection."""
172 name: str = "unet"
173 encoder_name: str = "resnet18"
174 encoder_weights: str | None = "imagenet" # None = train from scratch
175 in_channels: int = pydantic.Field(default=1, gt=0)
176 num_classes: int = pydantic.Field(default=3, gt=0)
177 extra: dict[str, Any] = {} # passed through to the model factory
178
179
180class LossConfig(SectionModel):
181 """Loss selection by registry name plus factory keyword arguments."""
182 name: str = "dice_focal"
183 args: dict[str, Any] = {}
184
185
186class TrainerConfig(SectionModel):
187 """Lightning trainer, logger, and callback knobs."""
188 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
189 lr: float = pydantic.Field(default=3.0e-4, gt=0)
190 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
191 # "auto" -> 16-mixed on CUDA, 32-true on CPU; else a Lightning precision
192 # (16, "16-mixed", "bf16-mixed", 32, "32-true", ...)
193 precision: int | str = "auto"
194 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
195 accelerator: str = "auto"
196 devices: int | str = 1
197 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
198 viz_samples: int = pydantic.Field(default=4, ge=0)
199 monitor: str = "val/f1_mean_fg"
200 monitor_mode: Literal["max", "min"] = "max"
201 early_stop_monitor: str = "val/loss" # stop when this stops improving
202 early_stop_mode: Literal["min", "max"] = "min"
203 # epochs without improvement before stopping; 0 disables
204 early_stop_patience: int = pydantic.Field(default=4, ge=0)
205 log_dir: str = "runs"
206 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
207
208 @pydantic.field_validator("precision", "devices", mode="before")
209 @classmethod
210 def _reject_bool(cls, value: Any, info: pydantic.ValidationInfo) -> Any:
211 """Keep YAML ``yes``/``on`` from silently narrowing to ``1``."""
212 if isinstance(value, bool):
213 raise ValueError(
214 f"{info.field_name} must be an int or a str, got bool {value!r}")
215 return value
216
217
218class HarnessConfig(SectionModel):
219 """Top-level training config: one YAML file, one instance."""150 """Top-level training config: one YAML file, one instance."""
220 experiment: str = "experiment"151 experiment: str = "experiment"
221 seed: int = 1337152 seed: int = 1337
222 data: DataConfig = DataConfig()153 data: DataConfig = DataConfig()
223 model: ModelConfig = ModelConfig()154 model: harness_config.ModelConfig = harness_config.ModelConfig()
224 loss: LossConfig = LossConfig()155 loss: harness_config.LossConfig = harness_config.LossConfig()
225 train: TrainerConfig = TrainerConfig()156 train: harness_config.TrainerConfig = harness_config.TrainerConfig()
226157
227 @classmethod158 @classmethod
228 def from_yaml(cls, path: str | Path) -> "HarnessConfig":159 def from_yaml(cls, path: str | Path) -> "HarnessConfig":
229 """Loads and validates a harness YAML config.160 """Loads and validates a harness YAML config.
Importance #18: src/train/config.py @@ -235,19 +166,18 @@
235 The validated, frozen config.166 The validated, frozen config.
236167
237 Raises:168 Raises:
238 FileNotFoundError: If ``path`` does not exist.169 FileNotFoundError: If ``path`` does not exist.
239 ConfigError: If the document is not a mapping, holds an unknown key,170 config_loader.ConfigError: If the document is not a mapping, holds
240 or holds a value invalid for its field. Derives from171 an unknown key, or holds a value invalid for its field. Unknown
241 ``ValueError``.172 keys and bad values raise ``HarnessConfigError``; a non-mapping
173 document raises the shared harness error class. Both derive
174 from ``ValueError``.
242 """175 """
243 raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}176 raw = harness_config.load_yaml_mapping(path)
244 if not isinstance(raw, Mapping):
245 raise ConfigError(
246 f"config {str(path)!r} must contain a top-level mapping, "
247 f"got {type(raw).__name__}")
248 return config_loader.validate_config(177 return config_loader.validate_config(
249 cls, raw, context=str(path), error_cls=ConfigError)178 cls, raw, context=f"{_CONTEXT} {str(path)!r}",
179 error_cls=HarnessConfigError)
250180
251181
252def with_overrides(cfg: HarnessConfig,182def with_overrides(cfg: HarnessConfig,
253 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:183 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
Importance #19: src/train/config.py @@ -265,12 +195,13 @@
265 Returns:195 Returns:
266 ``cfg`` itself when no override is given, otherwise a validated copy.196 ``cfg`` itself when no override is given, otherwise a validated copy.
267197
268 Raises:198 Raises:
269 ConfigError: An override value is invalid for its declared field.199 HarnessConfigError: An override value is invalid for its declared field.
270 """200 """
271 applied = {name: dict(values) for name, values in sections.items() if values}201 applied = {name: dict(values) for name, values in sections.items() if values}
272 if not applied:202 if not applied:
273 return cfg203 return cfg
204 logger.info("Config overrides applied: %s", ", ".join(sorted(applied)))
274 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)205 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
275 return config_loader.validate_config(206 return config_loader.validate_config(
276 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)207 type(cfg), merged, context=_OVERRIDES_CONTEXT, error_cls=HarnessConfigError)
Importance #20: test/test_config.py @@ -0,0 +1,99 @@
1"""Config-schema tests for the line-bitmap training harness.
2
3The harness config is YAML-backed and has no packaged JSON defaults, so the
4fleet parity test is replaced by a defaults-vs-model check on the shipped
5baseline YAML's schema-level behaviour.
6"""
7from pathlib import Path
8
9import pytest
10import yaml
11from iolabs.common import config_loader
12from iolabs_ml_harness import config as harness_config
13
14from src.train import config
15
16
17def _write(tmp_path: Path, text: str) -> Path:
18 path = tmp_path / "cfg.yaml"
19 path.write_text(text, encoding="utf-8")
20 return path
21
22
23def test_error_class_is_config_error() -> None:
24 assert issubclass(config.HarnessConfigError, config_loader.ConfigError)
25 assert issubclass(config.HarnessConfigError, ValueError)
26
27
28def test_shared_sections_come_from_the_harness_package() -> None:
29 """model/loss/train are the shared models, not repo-local copies."""
30 fields = config.HarnessConfig.model_fields
31 assert fields["model"].annotation is harness_config.ModelConfig
32 assert fields["loss"].annotation is harness_config.LossConfig
33 assert fields["train"].annotation is harness_config.TrainerConfig
34 assert issubclass(config.DataConfig, harness_config.SectionModel)
35
36
37def test_model_defaults_match_documented_defaults() -> None:
38 dumped = config.HarnessConfig().model_dump()
39
40 assert dumped == {
41 "experiment": "experiment",
42 "seed": 1337,
43 "data": {
44 "pairs": [], "val_pairs": [], "test_pairs": [],
45 "crop_size": 512, "batch_size": 16, "num_workers": 4,
46 "val_fraction": 0.15, "crops_per_tile": 4, "pos_crop_prob": 0.7,
47 "min_valid_fraction": 0.10, "augment": True,
48 "label_source": "rendered", "label_stroke_px": 4,
49 "review_statuses": ["ok"]},
50 "model": harness_config.ModelConfig().model_dump(),
51 "loss": harness_config.LossConfig().model_dump(),
52 "train": harness_config.TrainerConfig().model_dump()}
53
54
55def test_from_yaml_returns_defaults_for_an_empty_document(tmp_path: Path) -> None:
56 assert config.HarnessConfig.from_yaml(_write(tmp_path, "")) == config.HarnessConfig()
57
58
59def test_unknown_top_level_key_is_rejected(tmp_path: Path) -> None:
60 with pytest.raises(config.HarnessConfigError, match="experimnt"):
61 config.HarnessConfig.from_yaml(_write(tmp_path, "experimnt: t\n"))
62
63
64def test_unknown_nested_key_is_rejected(tmp_path: Path) -> None:
65 with pytest.raises(config.HarnessConfigError, match="crop_sizes"):
66 config.HarnessConfig.from_yaml(_write(tmp_path, "data:\n crop_sizes: 64\n"))
67
68
69def test_non_mapping_document_is_rejected(tmp_path: Path) -> None:
70 with pytest.raises(config_loader.ConfigError, match="mapping"):
71 config.HarnessConfig.from_yaml(_write(tmp_path, "- a\n- b\n"))
72
73
74def test_overrides_deep_merge_onto_defaults() -> None:
75 cfg = config.HarnessConfig(experiment="e")
76
77 updated = config.with_overrides(
78 cfg, {"data": {"batch_size": 2}, "train": {"max_epochs": 3}, "loss": {}})
79
80 assert updated.data.batch_size == 2 and updated.train.max_epochs == 3
81 assert updated.experiment == "e" and updated.data.crop_size == 512
82 assert cfg.data.batch_size == 16 # source untouched
83
84
85def test_override_coercion_and_rejection() -> None:
86 updated = config.with_overrides(config.HarnessConfig(),
87 {"data": {"batch_size": "4"}})
88 assert updated.data.batch_size == 4
89
90 with pytest.raises(config.HarnessConfigError, match="batch_size"):
91 config.with_overrides(config.HarnessConfig(), {"data": {"batch_size": 0}})
92
93
94def test_shipped_configs_parse() -> None:
95 for path in sorted(Path("configs").glob("*.yaml")):
96 raw = yaml.safe_load(path.read_text(encoding="utf-8"))
97 if not isinstance(raw, dict) or "experiment" not in raw:
98 continue # non-harness config in the same folder
99 assert config.HarnessConfig.from_yaml(path).experiment
0
Importance #21: test/test_train_overrides.py @@ -5,11 +5,11 @@
5from pathlib import Path5from pathlib import Path
66
7import pytest7import pytest
88
9from iolabs_ml_harness.config import TrainerConfig
9from src.train.config import (10from src.train.config import (
10 ConfigError, DataConfig, HarnessConfig, PairSpec, TrainerConfig,11 DataConfig, HarnessConfig, HarnessConfigError, PairSpec, reroot_data_paths)
11 reroot_data_paths)
1212
13_REPO_ROOT = Path(__file__).resolve().parents[1]13_REPO_ROOT = Path(__file__).resolve().parents[1]
14_CLI_FLAGS = ("log_dir", "num_workers", "max_epochs", "batch_size", "model", "encoder")14_CLI_FLAGS = ("log_dir", "num_workers", "max_epochs", "batch_size", "model", "encoder")
1515
Importance #22: test/test_train_overrides.py @@ -153,9 +153,9 @@
153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},
154 {"num_workers": -1}])154 {"num_workers": -1}])
155def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:155def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:
156 """CLI overrides must hit the same bounds as values coming from YAML."""156 """CLI overrides must hit the same bounds as values coming from YAML."""
157 with pytest.raises(ConfigError):157 with pytest.raises(HarnessConfigError):
158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
159159
160160
161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
Importance #23: test/test_train_overrides.py @@ -197,6 +197,6 @@
197def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str, tmp_path: Path) -> None:197def test_yaml_yes_is_rejected_instead_of_becoming_one(field: str, tmp_path: Path) -> None:
198 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""198 """PyYAML turns ``devices: yes`` into True; int|str would narrow it to 1."""
199 path = tmp_path / "bool.yaml"199 path = tmp_path / "bool.yaml"
200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")
201 with pytest.raises(ConfigError, match="got bool"):201 with pytest.raises(HarnessConfigError, match="got bool"):
202 HarnessConfig.from_yaml(path)202 HarnessConfig.from_yaml(path)