Back to report index

guardraildetection (ML) b407f47: AI3D-379 Pydantic config models via iolabs-common ConfigModel

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

Commit #78 · 25 snippets

 CLAUDE.md                  |  8 +++++
 pyproject.toml             |  7 ++--
 scripts/train.py           | 38 +++++++++++++++-----
 src/train/config.py        | 90 +++++++++++++++++++++++++---------------------
 test/test_harness_smoke.py | 62 ++++++++++++++++++++++++++++++--
 5 files changed, 152 insertions(+), 53 deletions(-)
Importance #1: src/train/config.py @@ -1,60 +1,70 @@
1"""Task-specific config composed over the shared harness package.1"""Task-specific config composed over the shared harness package.
22
3Plain dataclasses + yaml, no config framework. ``DataConfig`` is3Pydantic models on `iolabs.common.config_loader.ConfigModel` + yaml. ``DataConfig``
4guardrail-specific (dataset/synthetic split); the ``model``/``loss``/``train``4is guardrail-specific (dataset/synthetic split); the ``model``/``loss``/``train``
5sections are the generic dataclasses from ``iolabs_ml_harness.config``, built5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys
6with the same strict-unknown-keys ``build_section`` the shared package uses.6raise, so config typos fail fast instead of silently training with defaults.
7
8Adding a config key = adding one field with its default to the model below.
7"""9"""
8from dataclasses import dataclass, field10import logging
9from pathlib import Path11from pathlib import Path
10from typing import Any12from typing import Literal
13
14import pydantic
15from iolabs.common import config_loader
16from iolabs_ml_harness import config as harness_config
17
18logger = logging.getLogger(__name__)
19
1120
12import yaml21class ConfigError(config_loader.ConfigError):
13from iolabs_ml_harness.config import (22 """Raised when a guardrail harness config holds unknown keys or bad values."""
14 LossConfig, ModelConfig, TrainerConfig, build_section)
1523
1624
17@dataclass25class PairSpec(config_loader.ConfigModel):
18class PairSpec:
19 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""26 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""
20 images: str27 images: str
21 masks: str28 masks: str
2229
2330
24@dataclass31class DataConfig(config_loader.ConfigModel):
25class DataConfig:32 """Guardrail dataset selection and crop/loader sizing."""
26 source: str = "synthetic" # synthetic | pairs33 source: Literal["synthetic", "pairs"] = "synthetic"
27 pairs: list = field(default_factory=list)34 pairs: list[PairSpec] = []
28 crop_size: int = 25635 crop_size: int = pydantic.Field(default=256, gt=0)
29 batch_size: int = 836 batch_size: int = pydantic.Field(default=8, gt=0)
30 num_workers: int = 437 num_workers: int = pydantic.Field(default=4, ge=0)
31 val_fraction: float = 0.1538 val_fraction: float = pydantic.Field(default=0.15, ge=0.0, le=1.0)
32 synthetic_tiles: int = 6439 synthetic_tiles: int = pydantic.Field(default=64, gt=0)
33 synthetic_seed: int = 2026070740 synthetic_seed: int = 20260707
3441
3542
36@dataclass43class HarnessConfig(config_loader.ConfigModel):
37class HarnessConfig:44 """Top-level training config: one YAML file, one instance."""
38 experiment: str = "experiment"45 experiment: str = "experiment"
39 seed: int = 133746 seed: int = 1337
40 data: DataConfig = field(default_factory=DataConfig)47 data: DataConfig = DataConfig()
41 model: ModelConfig = field(default_factory=ModelConfig)48 model: harness_config.ModelConfig = harness_config.ModelConfig()
42 loss: LossConfig = field(default_factory=LossConfig)49 loss: harness_config.LossConfig = harness_config.LossConfig()
43 train: TrainerConfig = field(default_factory=TrainerConfig)50 train: harness_config.TrainerConfig = harness_config.TrainerConfig()
4451
45 @classmethod52 @classmethod
46 def from_yaml(cls, path: str | Path) -> "HarnessConfig":53 def from_yaml(cls, path: str | Path) -> "HarnessConfig":
47 raw: dict[str, Any] = yaml.safe_load(Path(path).read_text()) or {}54 """Loads and validates a harness YAML config.
48 data = build_section(DataConfig, raw.pop("data", {}), "data")55
49 data.pairs = [build_section(PairSpec, p, "data.pairs[]") for p in data.pairs]56 Args:
50 cfg = cls(57 path: Path of the YAML file, read as UTF-8.
51 experiment=raw.pop("experiment", cls.experiment),58
52 seed=raw.pop("seed", cls.seed),59 Returns:
53 data=data,60 The validated, frozen config.
54 model=build_section(ModelConfig, raw.pop("model", {}), "model"),61
55 loss=build_section(LossConfig, raw.pop("loss", {}), "loss"),62 Raises:
56 train=build_section(TrainerConfig, raw.pop("train", {}), "train"),63 FileNotFoundError: If ``path`` does not exist.
57 )64 ConfigError: If the document is not a mapping, holds an unknown key,
58 if raw:65 or holds a value invalid for its field. Derives from
59 raise KeyError(f"unknown top-level config key(s) {sorted(raw)} in {path}")66 ``ValueError``.
60 return cfg67 """
68 raw = harness_config.load_yaml_mapping(path)
69 return config_loader.validate_config(
70 cls, raw, context=str(path), error_cls=ConfigError)
Importance #2: scripts/train.py @@ -18,8 +18,9 @@
18deeplabv3plus, ...) and model.encoder_name any smp encoder. encoder_weights:18deeplabv3plus, ...) and model.encoder_name any smp encoder. encoder_weights:
19imagenet downloads pretrained encoder weights on first use.19imagenet downloads pretrained encoder weights on first use.
20"""20"""
21import argparse21import argparse
22from typing import Any
2223
23import lightning.pytorch as pl24import lightning.pytorch as pl
24from iolabs_ml_harness.losses import build_loss25from iolabs_ml_harness.losses import build_loss
25from iolabs_ml_harness.models import build_model26from iolabs_ml_harness.models import build_model
Importance #3: scripts/train.py @@ -43,23 +44,44 @@
43 parser.add_argument("--cpu", action="store_true", help="force CPU training")44 parser.add_argument("--cpu", action="store_true", help="force CPU training")
44 return parser.parse_args()45 return parser.parse_args()
4546
4647
48def apply_cli_overrides(cfg: config.HarnessConfig,
49 args: argparse.Namespace) -> config.HarnessConfig:
50 """Returns a copy of ``cfg`` with the CLI section overrides applied.
51
52 Config models are frozen, so the overrides are applied by copying each
53 touched section instead of assigning to it.
54
55 Args:
56 cfg: The config parsed from the YAML file.
57 args: Parsed CLI arguments; ``None``/empty values override nothing.
58
59 Returns:
60 ``cfg`` itself when no override was given, otherwise an updated copy.
61 """
62 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
63 if args.max_epochs is not None:
64 sections["train"]["max_epochs"] = args.max_epochs
65 if args.batch_size is not None:
66 sections["data"]["batch_size"] = args.batch_size
67 if args.model:
68 sections["model"]["name"] = args.model
69 if args.encoder:
70 sections["model"]["encoder_name"] = args.encoder
71 updates = {name: getattr(cfg, name).model_copy(update=values)
72 for name, values in sections.items() if values}
73 return cfg.model_copy(update=updates) if updates else cfg
74
75
47def main() -> None:76def main() -> None:
48 args = parse_args()77 args = parse_args()
49 cfg = config.HarnessConfig.from_yaml(args.config)78 cfg = config.HarnessConfig.from_yaml(args.config)
50 if cfg.model.num_classes != len(CLASS_NAMES):79 if cfg.model.num_classes != len(CLASS_NAMES):
51 raise ValueError(80 raise ValueError(
52 f"model.num_classes={cfg.model.num_classes} but the harness tracks "81 f"model.num_classes={cfg.model.num_classes} but the harness tracks "
53 f"{len(CLASS_NAMES)} classes {CLASS_NAMES} -- metrics would mis-bin")82 f"{len(CLASS_NAMES)} classes {CLASS_NAMES} -- metrics would mis-bin")
54 if args.max_epochs is not None:83 cfg = apply_cli_overrides(cfg, args)
55 cfg.train.max_epochs = args.max_epochs
56 if args.batch_size is not None:
57 cfg.data.batch_size = args.batch_size
58 if args.model:
59 cfg.model.name = args.model
60 if args.encoder:
61 cfg.model.encoder_name = args.encoder
6284
63 pl.seed_everything(cfg.seed, workers=True)85 pl.seed_everything(cfg.seed, workers=True)
6486
65 dm = datamodule.GuardrailDataModule(cfg.data)87 dm = datamodule.GuardrailDataModule(cfg.data)
Importance #4: test/test_harness_smoke.py @@ -2,8 +2,11 @@
22
3Requires the ml extra (uv sync --extra ml --extra dev), which also pulls in3Requires the ml extra (uv sync --extra ml --extra dev), which also pulls in
4iolabs-ml-harness (path dep during development); skips cleanly otherwise.4iolabs-ml-harness (path dep during development); skips cleanly otherwise.
5"""5"""
6import argparse
7import importlib.util
8import sys
6from pathlib import Path9from pathlib import Path
710
8import pytest11import pytest
912
Importance #5: test/test_harness_smoke.py @@ -23,8 +26,28 @@
23from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E40226from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E402
24from src.train.config import DataConfig, HarnessConfig # noqa: E40227from src.train.config import DataConfig, HarnessConfig # noqa: E402
25from src.train.datamodule import GuardrailDataModule # noqa: E40228from src.train.datamodule import GuardrailDataModule # noqa: E402
2629
30_REPO_ROOT = Path(__file__).resolve().parents[1]
31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")
32
33
34def _train_script():
35 """Imports scripts/train.py as a module, caching it across tests."""
36 if "train_script" in sys.modules:
37 return sys.modules["train_script"]
38 spec = importlib.util.spec_from_file_location(
39 "train_script", _REPO_ROOT / "scripts" / "train.py")
40 module = importlib.util.module_from_spec(spec)
41 sys.modules["train_script"] = module
42 spec.loader.exec_module(module)
43 return module
44
45
46def _args(**overrides: object) -> argparse.Namespace:
47 """Builds a parsed-CLI namespace where unset flags are None."""
48 return argparse.Namespace(**{name: overrides.get(name) for name in _CLI_FLAGS})
49
2750
28def test_shipped_baseline_config_parses() -> None:51def test_shipped_baseline_config_parses() -> None:
29 cfg = HarnessConfig.from_yaml("configs/unet_baseline.yaml")52 cfg = HarnessConfig.from_yaml("configs/unet_baseline.yaml")
30 assert cfg.experiment == "guardrail_unet_baseline"53 assert cfg.experiment == "guardrail_unet_baseline"
Importance #6: test/test_harness_smoke.py @@ -42,9 +65,9 @@
42 cfg = HarnessConfig.from_yaml(good)65 cfg = HarnessConfig.from_yaml(good)
43 assert cfg.experiment == "t" and cfg.model.name == "unet"66 assert cfg.experiment == "t" and cfg.model.name == "unet"
44 bad = tmp_path / "bad.yaml"67 bad = tmp_path / "bad.yaml"
45 bad.write_text("data:\n soruce: synthetic\n") # typo'd key68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key
46 with pytest.raises(KeyError, match="soruce"):69 with pytest.raises(ValueError, match="soruce"):
47 HarnessConfig.from_yaml(bad)70 HarnessConfig.from_yaml(bad)
4871
4972
50def test_synthetic_dataset_is_deterministic() -> None:73def test_synthetic_dataset_is_deterministic() -> None:
Importance #7: test/test_harness_smoke.py @@ -95,10 +118,10 @@
95 assert len(dm.val_dataset) == 8 # max(8, int(8 * 0.25)) == 8118 assert len(dm.val_dataset) == 8 # max(8, int(8 * 0.25)) == 8
96 batch = next(iter(dm.train_dataloader()))119 batch = next(iter(dm.train_dataloader()))
97 assert batch["image"].shape == (2, 1, 64, 64)120 assert batch["image"].shape == (2, 1, 64, 64)
98 assert batch["mask"].shape == (2, 64, 64)121 assert batch["mask"].shape == (2, 64, 64)
99 with pytest.raises(ValueError, match="unknown data.source"):122 with pytest.raises(ValueError, match="source"):
100 GuardrailDataModule(DataConfig(source="nope")).setup()123 DataConfig(source="nope") # rejected by the config model itself
101 with pytest.raises(NotImplementedError):124 with pytest.raises(NotImplementedError):
102 GuardrailDataModule(DataConfig(source="pairs")).setup()125 GuardrailDataModule(DataConfig(source="pairs")).setup()
103126
104127
Importance #8: test/test_harness_smoke.py @@ -118,4 +141,37 @@
118 trainer, _checkpoint, _val_loss_ckpt = build_trainer(141 trainer, _checkpoint, _val_loss_ckpt = build_trainer(
119 train_cfg, experiment="smoke", fast_dev_run=True, cpu=True)142 train_cfg, experiment="smoke", fast_dev_run=True, cpu=True)
120 trainer.fit(lit, datamodule=dm)143 trainer.fit(lit, datamodule=dm)
121 assert dm.train_dataset is not None and dm.val_dataset is not None144 assert dm.train_dataset is not None and dm.val_dataset is not None
145
146
147def test_apply_cli_overrides_copies_every_touched_section() -> None:
148 cfg = HarnessConfig()
149 args = _args(max_epochs=7, batch_size=3, model="unetplusplus", encoder="resnet34")
150
151 updated = _train_script().apply_cli_overrides(cfg, args)
152
153 assert updated.train.max_epochs == 7 and updated.data.batch_size == 3
154 assert updated.model.name == "unetplusplus"
155 assert updated.model.encoder_name == "resnet34"
156 # untouched fields of the copied sections survive
157 assert updated.train.lr == cfg.train.lr
158 assert updated.data.crop_size == cfg.data.crop_size
159 assert updated.data.synthetic_seed == cfg.data.synthetic_seed
160 assert updated.model.num_classes == cfg.model.num_classes
161 assert updated.loss == cfg.loss and updated.seed == cfg.seed
162 # the frozen input is never mutated
163 assert cfg.train.max_epochs == -1 and cfg.data.batch_size == 8
164
165
166def test_apply_cli_overrides_without_flags_returns_the_input() -> None:
167 cfg = HarnessConfig()
168
169 assert _train_script().apply_cli_overrides(cfg, _args()) is cfg
170
171
172def test_apply_cli_overrides_honours_zero_valued_flags() -> None:
173 cfg = HarnessConfig()
174
175 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))
176
177 assert updated.train.max_epochs == 0
Importance #9: pyproject.toml @@ -1,7 +1,7 @@
1[project]1[project]
2name = "iolabs-image-analyzer-guardrail-detection"2name = "iolabs-image-analyzer-guardrail-detection"
3version = "0.1.0"3version = "0.1.1"
4description = "Detection of highway guardrails (position + type) from point-cloud-derived rasters"4description = "Detection of highway guardrails (position + type) from point-cloud-derived rasters"
5requires-python = ">=3.11,<3.13"5requires-python = ">=3.11,<3.13"
6dependencies = [6dependencies = [
7 "numpy>=1.26",7 "numpy>=1.26",
Importance #10: pyproject.toml @@ -9,8 +9,10 @@
9 "scipy>=1.11",9 "scipy>=1.11",
10 "matplotlib>=3.7",10 "matplotlib>=3.7",
11 "Pillow>=10.0",11 "Pillow>=10.0",
12 "pyyaml>=6.0",12 "pyyaml>=6.0",
13 "pydantic>=2.7",
14 "iolabs-common>=0.9.0",
13]15]
1416
15[project.optional-dependencies]17[project.optional-dependencies]
16dev = [18dev = [
Importance #11: pyproject.toml @@ -23,9 +25,9 @@
23 "albumentations>=1.4",25 "albumentations>=1.4",
24 "lightning>=2.2",26 "lightning>=2.2",
25 "tensorboard>=2.16",27 "tensorboard>=2.16",
26 "torchmetrics>=1.3",28 "torchmetrics>=1.3",
27 "iolabs-ml-harness>=0.1.0",29 "iolabs-ml-harness>=0.2.1",
28]30]
29# Point-cloud cross-section tooling (src/cross_sections/cross_section_pcd.py);31# Point-cloud cross-section tooling (src/cross_sections/cross_section_pcd.py);
30# torch is also required for this path (see the `ml` extra above).32# torch is also required for this path (see the `ml` extra above).
31pointcloud = [33pointcloud = [
Importance #12: pyproject.toml @@ -42,8 +44,9 @@
42url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"44url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"
43authenticate = "always"45authenticate = "always"
4446
45[tool.uv.sources]47[tool.uv.sources]
48iolabs-common = { index = "nexus" }
46iolabs-ml-harness = { index = "nexus" }49iolabs-ml-harness = { index = "nexus" }
47iolabs-geometry-geometry = { index = "nexus" }50iolabs-geometry-geometry = { index = "nexus" }
48iolabs-point-cloud-las-tools = { index = "nexus" }51iolabs-point-cloud-las-tools = { index = "nexus" }
4952
Importance #13: CLAUDE.md @@ -44,8 +44,16 @@
44and datamodule in `src/train`, and dataset code in `src/dataset`. Configs and44and datamodule in `src/train`, and dataset code in `src/dataset`. Configs and
45`scripts/train.py` follow the same shape as45`scripts/train.py` follow the same shape as
46`3dai.iolabs.imageanalyzer.linebitmapsegmentation`.46`3dai.iolabs.imageanalyzer.linebitmapsegmentation`.
4747
48Config sections are pydantic models on `iolabs.common.config_loader.ConfigModel`
49(unknown keys rejected, values coerced by the shared fleet matrix, instances
50frozen). **Adding a config key = adding one field with its default to the model
51in `src/train/config.py`** — nothing else. `HarnessConfig.from_yaml` raises
52`config.ConfigError`, which derives from `ValueError`. Because the models are
53frozen, CLI overrides copy sections (`apply_cli_overrides` in
54`scripts/train.py`) instead of assigning to them.
55
48## Environment56## Environment
4957
50uv-managed: `uv sync`. Extras:58uv-managed: `uv sync`. Extras:
5159
Importance #14: pyproject.toml @@ -1,7 +1,7 @@
1[project]1[project]
2name = "iolabs-image-analyzer-guardrail-detection"2name = "iolabs-image-analyzer-guardrail-detection"
3version = "0.1.0"3version = "0.1.1"
4description = "Detection of highway guardrails (position + type) from point-cloud-derived rasters"4description = "Detection of highway guardrails (position + type) from point-cloud-derived rasters"
5requires-python = ">=3.11,<3.13"5requires-python = ">=3.11,<3.13"
6dependencies = [6dependencies = [
7 "numpy>=1.26",7 "numpy>=1.26",
Importance #15: pyproject.toml @@ -9,8 +9,10 @@
9 "scipy>=1.11",9 "scipy>=1.11",
10 "matplotlib>=3.7",10 "matplotlib>=3.7",
11 "Pillow>=10.0",11 "Pillow>=10.0",
12 "pyyaml>=6.0",12 "pyyaml>=6.0",
13 "pydantic>=2.7",
14 "iolabs-common>=0.9.0",
13]15]
1416
15[project.optional-dependencies]17[project.optional-dependencies]
16dev = [18dev = [
Importance #16: pyproject.toml @@ -23,9 +25,9 @@
23 "albumentations>=1.4",25 "albumentations>=1.4",
24 "lightning>=2.2",26 "lightning>=2.2",
25 "tensorboard>=2.16",27 "tensorboard>=2.16",
26 "torchmetrics>=1.3",28 "torchmetrics>=1.3",
27 "iolabs-ml-harness>=0.1.0",29 "iolabs-ml-harness>=0.2.1",
28]30]
29# Point-cloud cross-section tooling (src/cross_sections/cross_section_pcd.py);31# Point-cloud cross-section tooling (src/cross_sections/cross_section_pcd.py);
30# torch is also required for this path (see the `ml` extra above).32# torch is also required for this path (see the `ml` extra above).
31pointcloud = [33pointcloud = [
Importance #17: pyproject.toml @@ -42,8 +44,9 @@
42url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"44url = "https://nexus.iolabs.ch/repository/pypi-private/simple/"
43authenticate = "always"45authenticate = "always"
4446
45[tool.uv.sources]47[tool.uv.sources]
48iolabs-common = { index = "nexus" }
46iolabs-ml-harness = { index = "nexus" }49iolabs-ml-harness = { index = "nexus" }
47iolabs-geometry-geometry = { index = "nexus" }50iolabs-geometry-geometry = { index = "nexus" }
48iolabs-point-cloud-las-tools = { index = "nexus" }51iolabs-point-cloud-las-tools = { index = "nexus" }
4952
Importance #18: scripts/train.py @@ -18,8 +18,9 @@
18deeplabv3plus, ...) and model.encoder_name any smp encoder. encoder_weights:18deeplabv3plus, ...) and model.encoder_name any smp encoder. encoder_weights:
19imagenet downloads pretrained encoder weights on first use.19imagenet downloads pretrained encoder weights on first use.
20"""20"""
21import argparse21import argparse
22from typing import Any
2223
23import lightning.pytorch as pl24import lightning.pytorch as pl
24from iolabs_ml_harness.losses import build_loss25from iolabs_ml_harness.losses import build_loss
25from iolabs_ml_harness.models import build_model26from iolabs_ml_harness.models import build_model
Importance #19: scripts/train.py @@ -43,23 +44,44 @@
43 parser.add_argument("--cpu", action="store_true", help="force CPU training")44 parser.add_argument("--cpu", action="store_true", help="force CPU training")
44 return parser.parse_args()45 return parser.parse_args()
4546
4647
48def apply_cli_overrides(cfg: config.HarnessConfig,
49 args: argparse.Namespace) -> config.HarnessConfig:
50 """Returns a copy of ``cfg`` with the CLI section overrides applied.
51
52 Config models are frozen, so the overrides are applied by copying each
53 touched section instead of assigning to it.
54
55 Args:
56 cfg: The config parsed from the YAML file.
57 args: Parsed CLI arguments; ``None``/empty values override nothing.
58
59 Returns:
60 ``cfg`` itself when no override was given, otherwise an updated copy.
61 """
62 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
63 if args.max_epochs is not None:
64 sections["train"]["max_epochs"] = args.max_epochs
65 if args.batch_size is not None:
66 sections["data"]["batch_size"] = args.batch_size
67 if args.model:
68 sections["model"]["name"] = args.model
69 if args.encoder:
70 sections["model"]["encoder_name"] = args.encoder
71 updates = {name: getattr(cfg, name).model_copy(update=values)
72 for name, values in sections.items() if values}
73 return cfg.model_copy(update=updates) if updates else cfg
74
75
47def main() -> None:76def main() -> None:
48 args = parse_args()77 args = parse_args()
49 cfg = config.HarnessConfig.from_yaml(args.config)78 cfg = config.HarnessConfig.from_yaml(args.config)
50 if cfg.model.num_classes != len(CLASS_NAMES):79 if cfg.model.num_classes != len(CLASS_NAMES):
51 raise ValueError(80 raise ValueError(
52 f"model.num_classes={cfg.model.num_classes} but the harness tracks "81 f"model.num_classes={cfg.model.num_classes} but the harness tracks "
53 f"{len(CLASS_NAMES)} classes {CLASS_NAMES} -- metrics would mis-bin")82 f"{len(CLASS_NAMES)} classes {CLASS_NAMES} -- metrics would mis-bin")
54 if args.max_epochs is not None:83 cfg = apply_cli_overrides(cfg, args)
55 cfg.train.max_epochs = args.max_epochs
56 if args.batch_size is not None:
57 cfg.data.batch_size = args.batch_size
58 if args.model:
59 cfg.model.name = args.model
60 if args.encoder:
61 cfg.model.encoder_name = args.encoder
6284
63 pl.seed_everything(cfg.seed, workers=True)85 pl.seed_everything(cfg.seed, workers=True)
6486
65 dm = datamodule.GuardrailDataModule(cfg.data)87 dm = datamodule.GuardrailDataModule(cfg.data)
Importance #20: src/train/config.py @@ -1,60 +1,70 @@
1"""Task-specific config composed over the shared harness package.1"""Task-specific config composed over the shared harness package.
22
3Plain dataclasses + yaml, no config framework. ``DataConfig`` is3Pydantic models on `iolabs.common.config_loader.ConfigModel` + yaml. ``DataConfig``
4guardrail-specific (dataset/synthetic split); the ``model``/``loss``/``train``4is guardrail-specific (dataset/synthetic split); the ``model``/``loss``/``train``
5sections are the generic dataclasses from ``iolabs_ml_harness.config``, built5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys
6with the same strict-unknown-keys ``build_section`` the shared package uses.6raise, so config typos fail fast instead of silently training with defaults.
7
8Adding a config key = adding one field with its default to the model below.
7"""9"""
8from dataclasses import dataclass, field10import logging
9from pathlib import Path11from pathlib import Path
10from typing import Any12from typing import Literal
13
14import pydantic
15from iolabs.common import config_loader
16from iolabs_ml_harness import config as harness_config
17
18logger = logging.getLogger(__name__)
19
1120
12import yaml21class ConfigError(config_loader.ConfigError):
13from iolabs_ml_harness.config import (22 """Raised when a guardrail harness config holds unknown keys or bad values."""
14 LossConfig, ModelConfig, TrainerConfig, build_section)
1523
1624
17@dataclass25class PairSpec(config_loader.ConfigModel):
18class PairSpec:
19 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""26 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""
20 images: str27 images: str
21 masks: str28 masks: str
2229
2330
24@dataclass31class DataConfig(config_loader.ConfigModel):
25class DataConfig:32 """Guardrail dataset selection and crop/loader sizing."""
26 source: str = "synthetic" # synthetic | pairs33 source: Literal["synthetic", "pairs"] = "synthetic"
27 pairs: list = field(default_factory=list)34 pairs: list[PairSpec] = []
28 crop_size: int = 25635 crop_size: int = pydantic.Field(default=256, gt=0)
29 batch_size: int = 836 batch_size: int = pydantic.Field(default=8, gt=0)
30 num_workers: int = 437 num_workers: int = pydantic.Field(default=4, ge=0)
31 val_fraction: float = 0.1538 val_fraction: float = pydantic.Field(default=0.15, ge=0.0, le=1.0)
32 synthetic_tiles: int = 6439 synthetic_tiles: int = pydantic.Field(default=64, gt=0)
33 synthetic_seed: int = 2026070740 synthetic_seed: int = 20260707
3441
3542
36@dataclass43class HarnessConfig(config_loader.ConfigModel):
37class HarnessConfig:44 """Top-level training config: one YAML file, one instance."""
38 experiment: str = "experiment"45 experiment: str = "experiment"
39 seed: int = 133746 seed: int = 1337
40 data: DataConfig = field(default_factory=DataConfig)47 data: DataConfig = DataConfig()
41 model: ModelConfig = field(default_factory=ModelConfig)48 model: harness_config.ModelConfig = harness_config.ModelConfig()
42 loss: LossConfig = field(default_factory=LossConfig)49 loss: harness_config.LossConfig = harness_config.LossConfig()
43 train: TrainerConfig = field(default_factory=TrainerConfig)50 train: harness_config.TrainerConfig = harness_config.TrainerConfig()
4451
45 @classmethod52 @classmethod
46 def from_yaml(cls, path: str | Path) -> "HarnessConfig":53 def from_yaml(cls, path: str | Path) -> "HarnessConfig":
47 raw: dict[str, Any] = yaml.safe_load(Path(path).read_text()) or {}54 """Loads and validates a harness YAML config.
48 data = build_section(DataConfig, raw.pop("data", {}), "data")55
49 data.pairs = [build_section(PairSpec, p, "data.pairs[]") for p in data.pairs]56 Args:
50 cfg = cls(57 path: Path of the YAML file, read as UTF-8.
51 experiment=raw.pop("experiment", cls.experiment),58
52 seed=raw.pop("seed", cls.seed),59 Returns:
53 data=data,60 The validated, frozen config.
54 model=build_section(ModelConfig, raw.pop("model", {}), "model"),61
55 loss=build_section(LossConfig, raw.pop("loss", {}), "loss"),62 Raises:
56 train=build_section(TrainerConfig, raw.pop("train", {}), "train"),63 FileNotFoundError: If ``path`` does not exist.
57 )64 ConfigError: If the document is not a mapping, holds an unknown key,
58 if raw:65 or holds a value invalid for its field. Derives from
59 raise KeyError(f"unknown top-level config key(s) {sorted(raw)} in {path}")66 ``ValueError``.
60 return cfg67 """
68 raw = harness_config.load_yaml_mapping(path)
69 return config_loader.validate_config(
70 cls, raw, context=str(path), error_cls=ConfigError)
Importance #21: test/test_harness_smoke.py @@ -2,8 +2,11 @@
22
3Requires the ml extra (uv sync --extra ml --extra dev), which also pulls in3Requires the ml extra (uv sync --extra ml --extra dev), which also pulls in
4iolabs-ml-harness (path dep during development); skips cleanly otherwise.4iolabs-ml-harness (path dep during development); skips cleanly otherwise.
5"""5"""
6import argparse
7import importlib.util
8import sys
6from pathlib import Path9from pathlib import Path
710
8import pytest11import pytest
912
Importance #22: test/test_harness_smoke.py @@ -23,8 +26,28 @@
23from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E40226from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E402
24from src.train.config import DataConfig, HarnessConfig # noqa: E40227from src.train.config import DataConfig, HarnessConfig # noqa: E402
25from src.train.datamodule import GuardrailDataModule # noqa: E40228from src.train.datamodule import GuardrailDataModule # noqa: E402
2629
30_REPO_ROOT = Path(__file__).resolve().parents[1]
31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")
32
33
34def _train_script():
35 """Imports scripts/train.py as a module, caching it across tests."""
36 if "train_script" in sys.modules:
37 return sys.modules["train_script"]
38 spec = importlib.util.spec_from_file_location(
39 "train_script", _REPO_ROOT / "scripts" / "train.py")
40 module = importlib.util.module_from_spec(spec)
41 sys.modules["train_script"] = module
42 spec.loader.exec_module(module)
43 return module
44
45
46def _args(**overrides: object) -> argparse.Namespace:
47 """Builds a parsed-CLI namespace where unset flags are None."""
48 return argparse.Namespace(**{name: overrides.get(name) for name in _CLI_FLAGS})
49
2750
28def test_shipped_baseline_config_parses() -> None:51def test_shipped_baseline_config_parses() -> None:
29 cfg = HarnessConfig.from_yaml("configs/unet_baseline.yaml")52 cfg = HarnessConfig.from_yaml("configs/unet_baseline.yaml")
30 assert cfg.experiment == "guardrail_unet_baseline"53 assert cfg.experiment == "guardrail_unet_baseline"
Importance #23: test/test_harness_smoke.py @@ -42,9 +65,9 @@
42 cfg = HarnessConfig.from_yaml(good)65 cfg = HarnessConfig.from_yaml(good)
43 assert cfg.experiment == "t" and cfg.model.name == "unet"66 assert cfg.experiment == "t" and cfg.model.name == "unet"
44 bad = tmp_path / "bad.yaml"67 bad = tmp_path / "bad.yaml"
45 bad.write_text("data:\n soruce: synthetic\n") # typo'd key68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key
46 with pytest.raises(KeyError, match="soruce"):69 with pytest.raises(ValueError, match="soruce"):
47 HarnessConfig.from_yaml(bad)70 HarnessConfig.from_yaml(bad)
4871
4972
50def test_synthetic_dataset_is_deterministic() -> None:73def test_synthetic_dataset_is_deterministic() -> None:
Importance #24: test/test_harness_smoke.py @@ -95,10 +118,10 @@
95 assert len(dm.val_dataset) == 8 # max(8, int(8 * 0.25)) == 8118 assert len(dm.val_dataset) == 8 # max(8, int(8 * 0.25)) == 8
96 batch = next(iter(dm.train_dataloader()))119 batch = next(iter(dm.train_dataloader()))
97 assert batch["image"].shape == (2, 1, 64, 64)120 assert batch["image"].shape == (2, 1, 64, 64)
98 assert batch["mask"].shape == (2, 64, 64)121 assert batch["mask"].shape == (2, 64, 64)
99 with pytest.raises(ValueError, match="unknown data.source"):122 with pytest.raises(ValueError, match="source"):
100 GuardrailDataModule(DataConfig(source="nope")).setup()123 DataConfig(source="nope") # rejected by the config model itself
101 with pytest.raises(NotImplementedError):124 with pytest.raises(NotImplementedError):
102 GuardrailDataModule(DataConfig(source="pairs")).setup()125 GuardrailDataModule(DataConfig(source="pairs")).setup()
103126
104127
Importance #25: test/test_harness_smoke.py @@ -118,4 +141,37 @@
118 trainer, _checkpoint, _val_loss_ckpt = build_trainer(141 trainer, _checkpoint, _val_loss_ckpt = build_trainer(
119 train_cfg, experiment="smoke", fast_dev_run=True, cpu=True)142 train_cfg, experiment="smoke", fast_dev_run=True, cpu=True)
120 trainer.fit(lit, datamodule=dm)143 trainer.fit(lit, datamodule=dm)
121 assert dm.train_dataset is not None and dm.val_dataset is not None144 assert dm.train_dataset is not None and dm.val_dataset is not None
145
146
147def test_apply_cli_overrides_copies_every_touched_section() -> None:
148 cfg = HarnessConfig()
149 args = _args(max_epochs=7, batch_size=3, model="unetplusplus", encoder="resnet34")
150
151 updated = _train_script().apply_cli_overrides(cfg, args)
152
153 assert updated.train.max_epochs == 7 and updated.data.batch_size == 3
154 assert updated.model.name == "unetplusplus"
155 assert updated.model.encoder_name == "resnet34"
156 # untouched fields of the copied sections survive
157 assert updated.train.lr == cfg.train.lr
158 assert updated.data.crop_size == cfg.data.crop_size
159 assert updated.data.synthetic_seed == cfg.data.synthetic_seed
160 assert updated.model.num_classes == cfg.model.num_classes
161 assert updated.loss == cfg.loss and updated.seed == cfg.seed
162 # the frozen input is never mutated
163 assert cfg.train.max_epochs == -1 and cfg.data.batch_size == 8
164
165
166def test_apply_cli_overrides_without_flags_returns_the_input() -> None:
167 cfg = HarnessConfig()
168
169 assert _train_script().apply_cli_overrides(cfg, _args()) is cfg
170
171
172def test_apply_cli_overrides_honours_zero_valued_flags() -> None:
173 cfg = HarnessConfig()
174
175 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))
176
177 assert updated.train.max_epochs == 0