Back to report index

linebitmapimagerasterizer (ML) 6a3391c: AI3D-379 Review fixes: re-validate CLI overrides, null YAML sections mean defaults, int precision, bool guard

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

Commit #82 · 21 snippets

 CLAUDE.md                    |  5 ++-
 scripts/train.py             | 12 +++---
 src/train/config.py          | 94 ++++++++++++++++++++++++++++++++++++++------
 test/test_train_overrides.py | 56 +++++++++++++++++++++++++-
 4 files changed, 148 insertions(+), 19 deletions(-)
Importance #1: src/train/config.py @@ -4,20 +4,20 @@
4raise, so config typos fail fast instead of silently training with defaults, and4raise, so config typos fail fast instead of silently training with defaults, and
5values are coerced by the shared fleet matrix.5values are coerced by the shared fleet matrix.
66
7Adding a config key = adding one field with its default to the model below.7Adding a config key = adding one field with its default to the model below.
8Instances are frozen: derive a changed config with ``model_copy(update=...)``.8Instances are frozen and validated: derive a changed config with
9:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips
10validation and the fleet coercion matrix.
9"""11"""
10import logging
11from collections.abc import Mapping12from collections.abc import Mapping
12from pathlib import Path13from pathlib import Path
13from typing import Any, Literal14from typing import Any, Literal, get_args
1415
15import pydantic16import pydantic
16import yaml17import yaml
17from iolabs.common import config_loader18from iolabs.common import config_loader
1819from pydantic import fields as pydantic_fields
19logger = logging.getLogger(__name__)
2020
21_STROKE_KINDS = frozenset({"solid", "dashed"})21_STROKE_KINDS = frozenset({"solid", "dashed"})
22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")
23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")
Importance #2: src/train/config.py @@ -26,15 +26,47 @@
26class ConfigError(config_loader.ConfigError):26class ConfigError(config_loader.ConfigError):
27 """Raised when a harness config holds unknown keys or invalid values."""27 """Raised when a harness config holds unknown keys or invalid values."""
2828
2929
30class PairSpec(config_loader.ConfigModel):30class SectionModel(config_loader.ConfigModel):
31 """`ConfigModel` in which a YAML ``null`` means "use this field's default".
32
33 A bare ``data:`` / ``args:`` line parses to ``None``; the hand-written
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):
31 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""63 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""
32 images: str64 images: str
33 masks: str65 masks: str
3466
3567
36class DataConfig(config_loader.ConfigModel):68class DataConfig(SectionModel):
37 """Tile corpus, split, crop sampling and label rasterization knobs."""69 """Tile corpus, split, crop sampling and label rasterization knobs."""
38 pairs: list[PairSpec] = []70 pairs: list[PairSpec] = []
39 # Optional explicit, pre-split directories (e.g. the symlink folders under71 # Optional explicit, pre-split directories (e.g. the symlink folders under
40 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_72 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_
Importance #3: src/train/config.py @@ -134,9 +166,9 @@
134 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)166 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)
135 return data_cfg.model_copy(update=updates)167 return data_cfg.model_copy(update=updates)
136168
137169
138class ModelConfig(config_loader.ConfigModel):170class ModelConfig(SectionModel):
139 """Segmentation-models-pytorch architecture/encoder selection."""171 """Segmentation-models-pytorch architecture/encoder selection."""
140 name: str = "unet"172 name: str = "unet"
141 encoder_name: str = "resnet18"173 encoder_name: str = "resnet18"
142 encoder_weights: str | None = "imagenet" # None = train from scratch174 encoder_weights: str | None = "imagenet" # None = train from scratch
Importance #4: src/train/config.py @@ -144,20 +176,22 @@
144 num_classes: int = pydantic.Field(default=3, gt=0)176 num_classes: int = pydantic.Field(default=3, gt=0)
145 extra: dict[str, Any] = {} # passed through to the model factory177 extra: dict[str, Any] = {} # passed through to the model factory
146178
147179
148class LossConfig(config_loader.ConfigModel):180class LossConfig(SectionModel):
149 """Loss selection by registry name plus factory keyword arguments."""181 """Loss selection by registry name plus factory keyword arguments."""
150 name: str = "dice_focal"182 name: str = "dice_focal"
151 args: dict[str, Any] = {}183 args: dict[str, Any] = {}
152184
153185
154class TrainerConfig(config_loader.ConfigModel):186class TrainerConfig(SectionModel):
155 """Lightning trainer, logger, and callback knobs."""187 """Lightning trainer, logger, and callback knobs."""
156 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap188 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
157 lr: float = pydantic.Field(default=3.0e-4, gt=0)189 lr: float = pydantic.Field(default=3.0e-4, gt=0)
158 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)190 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
159 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU191 # "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"
160 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)194 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
161 accelerator: str = "auto"195 accelerator: str = "auto"
162 devices: int | str = 1196 devices: int | str = 1
163 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)197 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
Importance #5: src/train/config.py @@ -170,10 +204,19 @@
170 early_stop_patience: int = pydantic.Field(default=4, ge=0)204 early_stop_patience: int = pydantic.Field(default=4, ge=0)
171 log_dir: str = "runs"205 log_dir: str = "runs"
172 log_every_n_steps: int = pydantic.Field(default=10, ge=1)206 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
173207
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
174217
175class HarnessConfig(config_loader.ConfigModel):218class HarnessConfig(SectionModel):
176 """Top-level training config: one YAML file, one instance."""219 """Top-level training config: one YAML file, one instance."""
177 experiment: str = "experiment"220 experiment: str = "experiment"
178 seed: int = 1337221 seed: int = 1337
179 data: DataConfig = DataConfig()222 data: DataConfig = DataConfig()
Importance #6: src/train/config.py @@ -203,4 +246,31 @@
203 f"config {str(path)!r} must contain a top-level mapping, "246 f"config {str(path)!r} must contain a top-level mapping, "
204 f"got {type(raw).__name__}")247 f"got {type(raw).__name__}")
205 return config_loader.validate_config(248 return config_loader.validate_config(
206 cls, raw, context=str(path), error_cls=ConfigError)249 cls, raw, context=str(path), error_cls=ConfigError)
250
251
252def with_overrides(cfg: HarnessConfig,
253 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
254 """Returns a re-validated copy of ``cfg`` with per-section overrides merged in.
255
256 ``model_copy(update=...)`` would store the values unchecked, so a
257 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping through
258 the model keeps CLI overrides on exactly the path YAML values take.
259
260 Args:
261 cfg: The config to derive from; never mutated.
262 sections: Section name -> field name -> override value. Empty sections
263 are ignored.
264
265 Returns:
266 ``cfg`` itself when no override is given, otherwise a validated copy.
267
268 Raises:
269 ConfigError: An override value is invalid for its declared field.
270 """
271 applied = {name: dict(values) for name, values in sections.items() if values}
272 if not applied:
273 return cfg
274 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
275 return config_loader.validate_config(
276 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)
Importance #7: scripts/train.py @@ -58,17 +58,21 @@
58def apply_cli_overrides(cfg: config.HarnessConfig,58def apply_cli_overrides(cfg: config.HarnessConfig,
59 args: argparse.Namespace) -> config.HarnessConfig:59 args: argparse.Namespace) -> config.HarnessConfig:
60 """Returns a copy of ``cfg`` with the CLI overrides applied.60 """Returns a copy of ``cfg`` with the CLI overrides applied.
6161
62 Config models are frozen, so overrides are applied by copying each touched62 Config models are frozen, so overrides are re-validated into a copy instead
63 section instead of assigning to it.63 of being assigned onto ``cfg``.
6464
65 Args:65 Args:
66 cfg: The config parsed from the YAML file.66 cfg: The config parsed from the YAML file.
67 args: Parsed CLI arguments; ``None``/empty values override nothing.67 args: Parsed CLI arguments; ``None``/empty values override nothing.
6868
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.
71
72 Raises:
73 config.ConfigError: An override value is invalid for its field, e.g.
74 ``--max-epochs -2`` or ``--batch-size 0``.
71 """75 """
72 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
73 if args.log_dir:77 if args.log_dir:
74 sections["train"]["log_dir"] = args.log_dir78 sections["train"]["log_dir"] = args.log_dir
Importance #8: scripts/train.py @@ -81,11 +85,9 @@
81 if args.model:85 if args.model:
82 sections["model"]["name"] = args.model86 sections["model"]["name"] = args.model
83 if args.encoder:87 if args.encoder:
84 sections["model"]["encoder_name"] = args.encoder88 sections["model"]["encoder_name"] = args.encoder
85 updates = {name: getattr(cfg, name).model_copy(update=values)89 return config.with_overrides(cfg, sections)
86 for name, values in sections.items() if values}
87 return cfg.model_copy(update=updates) if updates else cfg
8890
8991
90def main() -> None:92def main() -> None:
91 args = parse_args()93 args = parse_args()
Importance #9: test/test_train_overrides.py @@ -5,9 +5,11 @@
5from pathlib import Path5from pathlib import Path
66
7import pytest7import pytest
88
9from src.train.config import DataConfig, HarnessConfig, PairSpec, reroot_data_paths9from src.train.config import (
10 ConfigError, DataConfig, HarnessConfig, PairSpec, TrainerConfig,
11 reroot_data_paths)
1012
11_REPO_ROOT = Path(__file__).resolve().parents[1]13_REPO_ROOT = Path(__file__).resolve().parents[1]
12_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")
1315
Importance #10: test/test_train_overrides.py @@ -145,4 +147,56 @@
145147
146 updated = _train_script().apply_cli_overrides(cfg, _args(num_workers=0, max_epochs=0))148 updated = _train_script().apply_cli_overrides(cfg, _args(num_workers=0, max_epochs=0))
147149
148 assert updated.data.num_workers == 0 and updated.train.max_epochs == 0150 assert updated.data.num_workers == 0 and updated.train.max_epochs == 0
151
152
153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},
154 {"num_workers": -1}])
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."""
157 with pytest.raises(ConfigError):
158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
159
160
161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
162 updated = _train_script().apply_cli_overrides(
163 HarnessConfig(), _args(max_epochs="7", batch_size="4"))
164
165 assert updated.train.max_epochs == 7 and updated.data.batch_size == 4
166
167
168def test_apply_cli_overrides_survives_a_stroke_mapping() -> None:
169 """The dict-valued label_stroke_px must round-trip through re-validation."""
170 cfg = HarnessConfig(data=DataConfig(label_stroke_px={"solid": 5, "dashed": 3}))
171
172 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=2))
173
174 assert updated.data.label_stroke_px == {"solid": 5, "dashed": 3}
175
176
177def test_null_sections_fall_back_to_defaults(tmp_path: Path) -> None:
178 """A bare ``data:`` line parses to None and must mean "all defaults"."""
179 path = tmp_path / "nulls.yaml"
180 path.write_text("experiment: t\ndata:\nmodel:\nloss:\n args:\n"
181 " name: dice_focal\n", encoding="utf-8")
182
183 cfg = HarnessConfig.from_yaml(path)
184
185 assert cfg.data.crop_size == 512 and cfg.data.label_stroke_px == 4
186 assert cfg.model.name == "unet" and cfg.loss.args == {}
187 # a field that really accepts None keeps it
188 assert HarnessConfig.from_yaml(path).model.encoder_weights == "imagenet"
189
190
191def test_precision_accepts_lightning_ints() -> None:
192 assert TrainerConfig(precision=16).precision == 16
193 assert TrainerConfig(precision="bf16-mixed").precision == "bf16-mixed"
194
195
196@pytest.mark.parametrize("field", ["precision", "devices"])
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."""
199 path = tmp_path / "bool.yaml"
200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")
201 with pytest.raises(ConfigError, match="got bool"):
202 HarnessConfig.from_yaml(path)
Importance #11: CLAUDE.md @@ -43,9 +43,12 @@
43- `src/train/` — LightningModule/DataModule, YAML config (pydantic models on43- `src/train/` — LightningModule/DataModule, YAML config (pydantic models on
44 `iolabs.common.config_loader.ConfigModel`: unknown keys rejected, values44 `iolabs.common.config_loader.ConfigModel`: unknown keys rejected, values
45 coerced, instances frozen — **adding a config key = adding one field with45 coerced, instances frozen — **adding a config key = adding one field with
46 its default to the model in `src/train/config.py`**; `from_yaml` raises46 its default to the model in `src/train/config.py`**; `from_yaml` raises
47 `config.ConfigError`, a `ValueError`), `MaskOverlayWriter`47 `config.ConfigError`, a `ValueError`; CLI overrides go through
48 `config.with_overrides`, which re-validates — never `model_copy(update=...)`,
49 which stores values unchecked; a bare `data:` line means "use the defaults"),
50 `MaskOverlayWriter`
48 (`intensity | label | prediction` sheets → TensorBoard + `runs/.../overlays/`)51 (`intensity | label | prediction` sheets → TensorBoard + `runs/.../overlays/`)
49- `scripts/train.py --config configs/<experiment>.yaml` — train (from repo52- `scripts/train.py --config configs/<experiment>.yaml` — train (from repo
50 root, needs `--extra ml`); `tensorboard --logdir runs` to monitor. Configs:53 root, needs `--extra ml`); `tensorboard --logdir runs` to monitor. Configs:
51 `unet_baseline`, `unet_vector_labels`, `unet_confirmed_good` (reviewer-`ok`54 `unet_baseline`, `unet_vector_labels`, `unet_confirmed_good` (reviewer-`ok`
Importance #12: scripts/train.py @@ -58,17 +58,21 @@
58def apply_cli_overrides(cfg: config.HarnessConfig,58def apply_cli_overrides(cfg: config.HarnessConfig,
59 args: argparse.Namespace) -> config.HarnessConfig:59 args: argparse.Namespace) -> config.HarnessConfig:
60 """Returns a copy of ``cfg`` with the CLI overrides applied.60 """Returns a copy of ``cfg`` with the CLI overrides applied.
6161
62 Config models are frozen, so overrides are applied by copying each touched62 Config models are frozen, so overrides are re-validated into a copy instead
63 section instead of assigning to it.63 of being assigned onto ``cfg``.
6464
65 Args:65 Args:
66 cfg: The config parsed from the YAML file.66 cfg: The config parsed from the YAML file.
67 args: Parsed CLI arguments; ``None``/empty values override nothing.67 args: Parsed CLI arguments; ``None``/empty values override nothing.
6868
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.
71
72 Raises:
73 config.ConfigError: An override value is invalid for its field, e.g.
74 ``--max-epochs -2`` or ``--batch-size 0``.
71 """75 """
72 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}76 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
73 if args.log_dir:77 if args.log_dir:
74 sections["train"]["log_dir"] = args.log_dir78 sections["train"]["log_dir"] = args.log_dir
Importance #13: scripts/train.py @@ -81,11 +85,9 @@
81 if args.model:85 if args.model:
82 sections["model"]["name"] = args.model86 sections["model"]["name"] = args.model
83 if args.encoder:87 if args.encoder:
84 sections["model"]["encoder_name"] = args.encoder88 sections["model"]["encoder_name"] = args.encoder
85 updates = {name: getattr(cfg, name).model_copy(update=values)89 return config.with_overrides(cfg, sections)
86 for name, values in sections.items() if values}
87 return cfg.model_copy(update=updates) if updates else cfg
8890
8991
90def main() -> None:92def main() -> None:
91 args = parse_args()93 args = parse_args()
Importance #14: src/train/config.py @@ -4,20 +4,20 @@
4raise, so config typos fail fast instead of silently training with defaults, and4raise, so config typos fail fast instead of silently training with defaults, and
5values are coerced by the shared fleet matrix.5values are coerced by the shared fleet matrix.
66
7Adding a config key = adding one field with its default to the model below.7Adding a config key = adding one field with its default to the model below.
8Instances are frozen: derive a changed config with ``model_copy(update=...)``.8Instances are frozen and validated: derive a changed config with
9:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips
10validation and the fleet coercion matrix.
9"""11"""
10import logging
11from collections.abc import Mapping12from collections.abc import Mapping
12from pathlib import Path13from pathlib import Path
13from typing import Any, Literal14from typing import Any, Literal, get_args
1415
15import pydantic16import pydantic
16import yaml17import yaml
17from iolabs.common import config_loader18from iolabs.common import config_loader
1819from pydantic import fields as pydantic_fields
19logger = logging.getLogger(__name__)
2020
21_STROKE_KINDS = frozenset({"solid", "dashed"})21_STROKE_KINDS = frozenset({"solid", "dashed"})
22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")22_PAIR_LIST_FIELDS = ("pairs", "val_pairs", "test_pairs")
23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")23_PATH_FIELD_SUFFIXES = ("path", "paths", "dir", "dirs", "root", "roots")
Importance #15: src/train/config.py @@ -26,15 +26,47 @@
26class ConfigError(config_loader.ConfigError):26class ConfigError(config_loader.ConfigError):
27 """Raised when a harness config holds unknown keys or invalid values."""27 """Raised when a harness config holds unknown keys or invalid values."""
2828
2929
30class PairSpec(config_loader.ConfigModel):30class SectionModel(config_loader.ConfigModel):
31 """`ConfigModel` in which a YAML ``null`` means "use this field's default".
32
33 A bare ``data:`` / ``args:`` line parses to ``None``; the hand-written
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):
31 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""63 """One images-dir / masks-dir pair (see src.dataset.index_tile_pairs)."""
32 images: str64 images: str
33 masks: str65 masks: str
3466
3567
36class DataConfig(config_loader.ConfigModel):68class DataConfig(SectionModel):
37 """Tile corpus, split, crop sampling and label rasterization knobs."""69 """Tile corpus, split, crop sampling and label rasterization knobs."""
38 pairs: list[PairSpec] = []70 pairs: list[PairSpec] = []
39 # Optional explicit, pre-split directories (e.g. the symlink folders under71 # Optional explicit, pre-split directories (e.g. the symlink folders under
40 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_72 # data/02_processed/<ds>/{train,val,test} built by scripts/build_processed_
Importance #16: src/train/config.py @@ -134,9 +166,9 @@
134 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)166 updates[name] = _reroot_path_value(getattr(data_cfg, name), root)
135 return data_cfg.model_copy(update=updates)167 return data_cfg.model_copy(update=updates)
136168
137169
138class ModelConfig(config_loader.ConfigModel):170class ModelConfig(SectionModel):
139 """Segmentation-models-pytorch architecture/encoder selection."""171 """Segmentation-models-pytorch architecture/encoder selection."""
140 name: str = "unet"172 name: str = "unet"
141 encoder_name: str = "resnet18"173 encoder_name: str = "resnet18"
142 encoder_weights: str | None = "imagenet" # None = train from scratch174 encoder_weights: str | None = "imagenet" # None = train from scratch
Importance #17: src/train/config.py @@ -144,20 +176,22 @@
144 num_classes: int = pydantic.Field(default=3, gt=0)176 num_classes: int = pydantic.Field(default=3, gt=0)
145 extra: dict[str, Any] = {} # passed through to the model factory177 extra: dict[str, Any] = {} # passed through to the model factory
146178
147179
148class LossConfig(config_loader.ConfigModel):180class LossConfig(SectionModel):
149 """Loss selection by registry name plus factory keyword arguments."""181 """Loss selection by registry name plus factory keyword arguments."""
150 name: str = "dice_focal"182 name: str = "dice_focal"
151 args: dict[str, Any] = {}183 args: dict[str, Any] = {}
152184
153185
154class TrainerConfig(config_loader.ConfigModel):186class TrainerConfig(SectionModel):
155 """Lightning trainer, logger, and callback knobs."""187 """Lightning trainer, logger, and callback knobs."""
156 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap188 max_epochs: int = pydantic.Field(default=-1, ge=-1) # -1 = no cap
157 lr: float = pydantic.Field(default=3.0e-4, gt=0)189 lr: float = pydantic.Field(default=3.0e-4, gt=0)
158 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)190 weight_decay: float = pydantic.Field(default=1.0e-4, ge=0)
159 precision: str = "auto" # auto -> 16-mixed on CUDA, 32-true on CPU191 # "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"
160 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)194 accumulate_grad_batches: int = pydantic.Field(default=1, ge=1)
161 accelerator: str = "auto"195 accelerator: str = "auto"
162 devices: int | str = 1196 devices: int | str = 1
163 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)197 viz_every_n_epochs: int = pydantic.Field(default=2, ge=0)
Importance #18: src/train/config.py @@ -170,10 +204,19 @@
170 early_stop_patience: int = pydantic.Field(default=4, ge=0)204 early_stop_patience: int = pydantic.Field(default=4, ge=0)
171 log_dir: str = "runs"205 log_dir: str = "runs"
172 log_every_n_steps: int = pydantic.Field(default=10, ge=1)206 log_every_n_steps: int = pydantic.Field(default=10, ge=1)
173207
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
174217
175class HarnessConfig(config_loader.ConfigModel):218class HarnessConfig(SectionModel):
176 """Top-level training config: one YAML file, one instance."""219 """Top-level training config: one YAML file, one instance."""
177 experiment: str = "experiment"220 experiment: str = "experiment"
178 seed: int = 1337221 seed: int = 1337
179 data: DataConfig = DataConfig()222 data: DataConfig = DataConfig()
Importance #19: src/train/config.py @@ -203,4 +246,31 @@
203 f"config {str(path)!r} must contain a top-level mapping, "246 f"config {str(path)!r} must contain a top-level mapping, "
204 f"got {type(raw).__name__}")247 f"got {type(raw).__name__}")
205 return config_loader.validate_config(248 return config_loader.validate_config(
206 cls, raw, context=str(path), error_cls=ConfigError)249 cls, raw, context=str(path), error_cls=ConfigError)
250
251
252def with_overrides(cfg: HarnessConfig,
253 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
254 """Returns a re-validated copy of ``cfg`` with per-section overrides merged in.
255
256 ``model_copy(update=...)`` would store the values unchecked, so a
257 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping through
258 the model keeps CLI overrides on exactly the path YAML values take.
259
260 Args:
261 cfg: The config to derive from; never mutated.
262 sections: Section name -> field name -> override value. Empty sections
263 are ignored.
264
265 Returns:
266 ``cfg`` itself when no override is given, otherwise a validated copy.
267
268 Raises:
269 ConfigError: An override value is invalid for its declared field.
270 """
271 applied = {name: dict(values) for name, values in sections.items() if values}
272 if not applied:
273 return cfg
274 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
275 return config_loader.validate_config(
276 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)
Importance #20: test/test_train_overrides.py @@ -5,9 +5,11 @@
5from pathlib import Path5from pathlib import Path
66
7import pytest7import pytest
88
9from src.train.config import DataConfig, HarnessConfig, PairSpec, reroot_data_paths9from src.train.config import (
10 ConfigError, DataConfig, HarnessConfig, PairSpec, TrainerConfig,
11 reroot_data_paths)
1012
11_REPO_ROOT = Path(__file__).resolve().parents[1]13_REPO_ROOT = Path(__file__).resolve().parents[1]
12_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")
1315
Importance #21: test/test_train_overrides.py @@ -145,4 +147,56 @@
145147
146 updated = _train_script().apply_cli_overrides(cfg, _args(num_workers=0, max_epochs=0))148 updated = _train_script().apply_cli_overrides(cfg, _args(num_workers=0, max_epochs=0))
147149
148 assert updated.data.num_workers == 0 and updated.train.max_epochs == 0150 assert updated.data.num_workers == 0 and updated.train.max_epochs == 0
151
152
153@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0},
154 {"num_workers": -1}])
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."""
157 with pytest.raises(ConfigError):
158 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
159
160
161def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
162 updated = _train_script().apply_cli_overrides(
163 HarnessConfig(), _args(max_epochs="7", batch_size="4"))
164
165 assert updated.train.max_epochs == 7 and updated.data.batch_size == 4
166
167
168def test_apply_cli_overrides_survives_a_stroke_mapping() -> None:
169 """The dict-valued label_stroke_px must round-trip through re-validation."""
170 cfg = HarnessConfig(data=DataConfig(label_stroke_px={"solid": 5, "dashed": 3}))
171
172 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=2))
173
174 assert updated.data.label_stroke_px == {"solid": 5, "dashed": 3}
175
176
177def test_null_sections_fall_back_to_defaults(tmp_path: Path) -> None:
178 """A bare ``data:`` line parses to None and must mean "all defaults"."""
179 path = tmp_path / "nulls.yaml"
180 path.write_text("experiment: t\ndata:\nmodel:\nloss:\n args:\n"
181 " name: dice_focal\n", encoding="utf-8")
182
183 cfg = HarnessConfig.from_yaml(path)
184
185 assert cfg.data.crop_size == 512 and cfg.data.label_stroke_px == 4
186 assert cfg.model.name == "unet" and cfg.loss.args == {}
187 # a field that really accepts None keeps it
188 assert HarnessConfig.from_yaml(path).model.encoder_weights == "imagenet"
189
190
191def test_precision_accepts_lightning_ints() -> None:
192 assert TrainerConfig(precision=16).precision == 16
193 assert TrainerConfig(precision="bf16-mixed").precision == "bf16-mixed"
194
195
196@pytest.mark.parametrize("field", ["precision", "devices"])
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."""
199 path = tmp_path / "bool.yaml"
200 path.write_text(f"train:\n {field}: yes\n", encoding="utf-8")
201 with pytest.raises(ConfigError, match="got bool"):
202 HarnessConfig.from_yaml(path)