Back to report index

guardraildetection (ML) 38464e8: AI3D-379 Review fixes: re-validate CLI overrides, null YAML sections mean defaults

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

Commit #79 · 17 snippets

 CLAUDE.md                  |  6 ++++--
 scripts/train.py           | 12 +++++++-----
 src/train/config.py        | 43 ++++++++++++++++++++++++++++++++++++-------
 test/test_harness_smoke.py | 34 ++++++++++++++++++++++++++++++++--
 4 files changed, 79 insertions(+), 16 deletions(-)
Importance #1: src/train/config.py @@ -5,31 +5,33 @@
5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys
6raise, so config typos fail fast instead of silently training with defaults.6raise, so config typos fail fast instead of silently training with defaults.
77
8Adding a config key = adding one field with its default to the model below.8Adding a config key = adding one field with its default to the model below.
9
10Instances are frozen and validated: derive a changed config with
11:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips
12validation and the fleet coercion matrix.
9"""13"""
10import logging14from collections.abc import Mapping
11from pathlib import Path15from pathlib import Path
12from typing import Literal16from typing import Any, Literal
1317
14import pydantic18import pydantic
15from iolabs.common import config_loader19from iolabs.common import config_loader
16from iolabs_ml_harness import config as harness_config20from iolabs_ml_harness import config as harness_config
1721
18logger = logging.getLogger(__name__)
19
2022
21class ConfigError(config_loader.ConfigError):23class ConfigError(config_loader.ConfigError):
22 """Raised when a guardrail harness config holds unknown keys or bad values."""24 """Raised when a guardrail harness config holds unknown keys or bad values."""
2325
2426
25class PairSpec(config_loader.ConfigModel):27class PairSpec(harness_config.SectionModel):
26 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""28 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""
27 images: str29 images: str
28 masks: str30 masks: str
2931
3032
31class DataConfig(config_loader.ConfigModel):33class DataConfig(harness_config.SectionModel):
32 """Guardrail dataset selection and crop/loader sizing."""34 """Guardrail dataset selection and crop/loader sizing."""
33 source: Literal["synthetic", "pairs"] = "synthetic"35 source: Literal["synthetic", "pairs"] = "synthetic"
34 pairs: list[PairSpec] = []36 pairs: list[PairSpec] = []
35 crop_size: int = pydantic.Field(default=256, gt=0)37 crop_size: int = pydantic.Field(default=256, gt=0)
Importance #2: src/train/config.py @@ -39,9 +41,9 @@
39 synthetic_tiles: int = pydantic.Field(default=64, gt=0)41 synthetic_tiles: int = pydantic.Field(default=64, gt=0)
40 synthetic_seed: int = 2026070742 synthetic_seed: int = 20260707
4143
4244
43class HarnessConfig(config_loader.ConfigModel):45class HarnessConfig(harness_config.SectionModel):
44 """Top-level training config: one YAML file, one instance."""46 """Top-level training config: one YAML file, one instance."""
45 experiment: str = "experiment"47 experiment: str = "experiment"
46 seed: int = 133748 seed: int = 1337
47 data: DataConfig = DataConfig()49 data: DataConfig = DataConfig()
Importance #3: src/train/config.py @@ -67,4 +69,31 @@
67 """69 """
68 raw = harness_config.load_yaml_mapping(path)70 raw = harness_config.load_yaml_mapping(path)
69 return config_loader.validate_config(71 return config_loader.validate_config(
70 cls, raw, context=str(path), error_cls=ConfigError)72 cls, raw, context=str(path), error_cls=ConfigError)
73
74
75def with_overrides(cfg: HarnessConfig,
76 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
77 """Returns a re-validated copy of ``cfg`` with per-section overrides merged in.
78
79 ``model_copy(update=...)`` would store the values unchecked, so a
80 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping through
81 the model keeps CLI overrides on exactly the path YAML values take.
82
83 Args:
84 cfg: The config to derive from; never mutated.
85 sections: Section name -> field name -> override value. Empty sections
86 are ignored.
87
88 Returns:
89 ``cfg`` itself when no override is given, otherwise a validated copy.
90
91 Raises:
92 ConfigError: An override value is invalid for its declared field.
93 """
94 applied = {name: dict(values) for name, values in sections.items() if values}
95 if not applied:
96 return cfg
97 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
98 return config_loader.validate_config(
99 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)
Importance #4: scripts/train.py @@ -48,17 +48,21 @@
48def apply_cli_overrides(cfg: config.HarnessConfig,48def apply_cli_overrides(cfg: config.HarnessConfig,
49 args: argparse.Namespace) -> config.HarnessConfig:49 args: argparse.Namespace) -> config.HarnessConfig:
50 """Returns a copy of ``cfg`` with the CLI section overrides applied.50 """Returns a copy of ``cfg`` with the CLI section overrides applied.
5151
52 Config models are frozen, so the overrides are applied by copying each52 Config models are frozen, so the overrides are re-validated into a copy
53 touched section instead of assigning to it.53 instead of being assigned onto ``cfg``.
5454
55 Args:55 Args:
56 cfg: The config parsed from the YAML file.56 cfg: The config parsed from the YAML file.
57 args: Parsed CLI arguments; ``None``/empty values override nothing.57 args: Parsed CLI arguments; ``None``/empty values override nothing.
5858
59 Returns:59 Returns:
60 ``cfg`` itself when no override was given, otherwise an updated copy.60 ``cfg`` itself when no override was given, otherwise an updated copy.
61
62 Raises:
63 config.ConfigError: An override value is invalid for its field, e.g.
64 ``--max-epochs -2`` or ``--batch-size 0``.
61 """65 """
62 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}66 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
63 if args.max_epochs is not None:67 if args.max_epochs is not None:
64 sections["train"]["max_epochs"] = args.max_epochs68 sections["train"]["max_epochs"] = args.max_epochs
Importance #5: scripts/train.py @@ -67,11 +71,9 @@
67 if args.model:71 if args.model:
68 sections["model"]["name"] = args.model72 sections["model"]["name"] = args.model
69 if args.encoder:73 if args.encoder:
70 sections["model"]["encoder_name"] = args.encoder74 sections["model"]["encoder_name"] = args.encoder
71 updates = {name: getattr(cfg, name).model_copy(update=values)75 return config.with_overrides(cfg, sections)
72 for name, values in sections.items() if values}
73 return cfg.model_copy(update=updates) if updates else cfg
7476
7577
76def main() -> None:78def main() -> None:
77 args = parse_args()79 args = parse_args()
Importance #6: test/test_harness_smoke.py @@ -23,9 +23,9 @@
23from iolabs_ml_harness.trainer import build_trainer # noqa: E40223from iolabs_ml_harness.trainer import build_trainer # noqa: E402
2424
25from src.dataset.rasters import index_raster_pairs # noqa: E40225from src.dataset.rasters import index_raster_pairs # noqa: E402
26from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E40226from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E402
27from src.train.config import DataConfig, HarnessConfig # noqa: E40227from src.train.config import ConfigError, DataConfig, HarnessConfig # noqa: E402
28from src.train.datamodule import GuardrailDataModule # noqa: E40228from src.train.datamodule import GuardrailDataModule # noqa: E402
2929
30_REPO_ROOT = Path(__file__).resolve().parents[1]30_REPO_ROOT = Path(__file__).resolve().parents[1]
31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")
Importance #7: test/test_harness_smoke.py @@ -65,12 +65,28 @@
65 cfg = HarnessConfig.from_yaml(good)65 cfg = HarnessConfig.from_yaml(good)
66 assert cfg.experiment == "t" and cfg.model.name == "unet"66 assert cfg.experiment == "t" and cfg.model.name == "unet"
67 bad = tmp_path / "bad.yaml"67 bad = tmp_path / "bad.yaml"
68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key
69 with pytest.raises(ValueError, match="soruce"):69 with pytest.raises(ConfigError, match="soruce"):
70 HarnessConfig.from_yaml(bad)70 HarnessConfig.from_yaml(bad)
7171
7272
73def test_null_sections_fall_back_to_defaults(tmp_path: Path) -> None:
74 """A bare ``data:`` line parses to None and must mean "all defaults"."""
75 path = tmp_path / "nulls.yaml"
76 path.write_text("experiment: t\ndata:\nmodel:\nloss:\n args:\n")
77 cfg = HarnessConfig.from_yaml(path)
78 assert cfg.data.crop_size == 256 and cfg.data.source == "synthetic"
79 assert cfg.model.name == "unet" and cfg.loss.args == {}
80
81
82def test_config_rejects_out_of_range_values(tmp_path: Path) -> None:
83 path = tmp_path / "bad.yaml"
84 path.write_text("train:\n max_epochs: -2\n")
85 with pytest.raises(ConfigError, match="max_epochs"):
86 HarnessConfig.from_yaml(path)
87
88
73def test_synthetic_dataset_is_deterministic() -> None:89def test_synthetic_dataset_is_deterministic() -> None:
74 ds_a = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)90 ds_a = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)
75 ds_b = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)91 ds_b = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)
76 for i in range(4):92 for i in range(4):
Importance #8: test/test_harness_smoke.py @@ -174,4 +190,18 @@
174190
175 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))191 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))
176192
177 assert updated.train.max_epochs == 0193 assert updated.train.max_epochs == 0
194
195
196@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0}])
197def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:
198 """CLI overrides must hit the same bounds as values coming from YAML."""
199 with pytest.raises(ConfigError):
200 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
201
202
203def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
204 updated = _train_script().apply_cli_overrides(
205 HarnessConfig(), _args(max_epochs="7"))
206
207 assert updated.train.max_epochs == 7
Importance #9: CLAUDE.md @@ -49,10 +49,12 @@
49(unknown keys rejected, values coerced by the shared fleet matrix, instances49(unknown keys rejected, values coerced by the shared fleet matrix, instances
50frozen). **Adding a config key = adding one field with its default to the model50frozen). **Adding a config key = adding one field with its default to the model
51in `src/train/config.py`** — nothing else. `HarnessConfig.from_yaml` raises51in `src/train/config.py`** — nothing else. `HarnessConfig.from_yaml` raises
52`config.ConfigError`, which derives from `ValueError`. Because the models are52`config.ConfigError`, which derives from `ValueError`. Because the models are
53frozen, CLI overrides copy sections (`apply_cli_overrides` in53frozen, CLI overrides go through `config.with_overrides` (used by
54`scripts/train.py`) instead of assigning to them.54`apply_cli_overrides` in `scripts/train.py`), which re-validates the merged
55config — never `model_copy(update=...)`, which stores values unchecked. A bare
56`data:` / `model:` line (YAML `null`) means "use the defaults".
5557
56## Environment58## Environment
5759
58uv-managed: `uv sync`. Extras:60uv-managed: `uv sync`. Extras:
Importance #10: scripts/train.py @@ -48,17 +48,21 @@
48def apply_cli_overrides(cfg: config.HarnessConfig,48def apply_cli_overrides(cfg: config.HarnessConfig,
49 args: argparse.Namespace) -> config.HarnessConfig:49 args: argparse.Namespace) -> config.HarnessConfig:
50 """Returns a copy of ``cfg`` with the CLI section overrides applied.50 """Returns a copy of ``cfg`` with the CLI section overrides applied.
5151
52 Config models are frozen, so the overrides are applied by copying each52 Config models are frozen, so the overrides are re-validated into a copy
53 touched section instead of assigning to it.53 instead of being assigned onto ``cfg``.
5454
55 Args:55 Args:
56 cfg: The config parsed from the YAML file.56 cfg: The config parsed from the YAML file.
57 args: Parsed CLI arguments; ``None``/empty values override nothing.57 args: Parsed CLI arguments; ``None``/empty values override nothing.
5858
59 Returns:59 Returns:
60 ``cfg`` itself when no override was given, otherwise an updated copy.60 ``cfg`` itself when no override was given, otherwise an updated copy.
61
62 Raises:
63 config.ConfigError: An override value is invalid for its field, e.g.
64 ``--max-epochs -2`` or ``--batch-size 0``.
61 """65 """
62 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}66 sections: dict[str, dict[str, Any]] = {"train": {}, "data": {}, "model": {}}
63 if args.max_epochs is not None:67 if args.max_epochs is not None:
64 sections["train"]["max_epochs"] = args.max_epochs68 sections["train"]["max_epochs"] = args.max_epochs
Importance #11: scripts/train.py @@ -67,11 +71,9 @@
67 if args.model:71 if args.model:
68 sections["model"]["name"] = args.model72 sections["model"]["name"] = args.model
69 if args.encoder:73 if args.encoder:
70 sections["model"]["encoder_name"] = args.encoder74 sections["model"]["encoder_name"] = args.encoder
71 updates = {name: getattr(cfg, name).model_copy(update=values)75 return config.with_overrides(cfg, sections)
72 for name, values in sections.items() if values}
73 return cfg.model_copy(update=updates) if updates else cfg
7476
7577
76def main() -> None:78def main() -> None:
77 args = parse_args()79 args = parse_args()
Importance #12: src/train/config.py @@ -5,31 +5,33 @@
5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys5sections are the generic models from ``iolabs_ml_harness.config``. Unknown keys
6raise, so config typos fail fast instead of silently training with defaults.6raise, so config typos fail fast instead of silently training with defaults.
77
8Adding a config key = adding one field with its default to the model below.8Adding a config key = adding one field with its default to the model below.
9
10Instances are frozen and validated: derive a changed config with
11:func:`with_overrides`, not ``model_copy(update=...)`` -- the latter skips
12validation and the fleet coercion matrix.
9"""13"""
10import logging14from collections.abc import Mapping
11from pathlib import Path15from pathlib import Path
12from typing import Literal16from typing import Any, Literal
1317
14import pydantic18import pydantic
15from iolabs.common import config_loader19from iolabs.common import config_loader
16from iolabs_ml_harness import config as harness_config20from iolabs_ml_harness import config as harness_config
1721
18logger = logging.getLogger(__name__)
19
2022
21class ConfigError(config_loader.ConfigError):23class ConfigError(config_loader.ConfigError):
22 """Raised when a guardrail harness config holds unknown keys or bad values."""24 """Raised when a guardrail harness config holds unknown keys or bad values."""
2325
2426
25class PairSpec(config_loader.ConfigModel):27class PairSpec(harness_config.SectionModel):
26 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""28 """One images-dir / masks-dir pair (see src.dataset.rasters.index_raster_pairs)."""
27 images: str29 images: str
28 masks: str30 masks: str
2931
3032
31class DataConfig(config_loader.ConfigModel):33class DataConfig(harness_config.SectionModel):
32 """Guardrail dataset selection and crop/loader sizing."""34 """Guardrail dataset selection and crop/loader sizing."""
33 source: Literal["synthetic", "pairs"] = "synthetic"35 source: Literal["synthetic", "pairs"] = "synthetic"
34 pairs: list[PairSpec] = []36 pairs: list[PairSpec] = []
35 crop_size: int = pydantic.Field(default=256, gt=0)37 crop_size: int = pydantic.Field(default=256, gt=0)
Importance #13: src/train/config.py @@ -39,9 +41,9 @@
39 synthetic_tiles: int = pydantic.Field(default=64, gt=0)41 synthetic_tiles: int = pydantic.Field(default=64, gt=0)
40 synthetic_seed: int = 2026070742 synthetic_seed: int = 20260707
4143
4244
43class HarnessConfig(config_loader.ConfigModel):45class HarnessConfig(harness_config.SectionModel):
44 """Top-level training config: one YAML file, one instance."""46 """Top-level training config: one YAML file, one instance."""
45 experiment: str = "experiment"47 experiment: str = "experiment"
46 seed: int = 133748 seed: int = 1337
47 data: DataConfig = DataConfig()49 data: DataConfig = DataConfig()
Importance #14: src/train/config.py @@ -67,4 +69,31 @@
67 """69 """
68 raw = harness_config.load_yaml_mapping(path)70 raw = harness_config.load_yaml_mapping(path)
69 return config_loader.validate_config(71 return config_loader.validate_config(
70 cls, raw, context=str(path), error_cls=ConfigError)72 cls, raw, context=str(path), error_cls=ConfigError)
73
74
75def with_overrides(cfg: HarnessConfig,
76 sections: Mapping[str, Mapping[str, Any]]) -> HarnessConfig:
77 """Returns a re-validated copy of ``cfg`` with per-section overrides merged in.
78
79 ``model_copy(update=...)`` would store the values unchecked, so a
80 ``--max-epochs -2`` would survive the ``ge=-1`` bound. Round-tripping through
81 the model keeps CLI overrides on exactly the path YAML values take.
82
83 Args:
84 cfg: The config to derive from; never mutated.
85 sections: Section name -> field name -> override value. Empty sections
86 are ignored.
87
88 Returns:
89 ``cfg`` itself when no override is given, otherwise a validated copy.
90
91 Raises:
92 ConfigError: An override value is invalid for its declared field.
93 """
94 applied = {name: dict(values) for name, values in sections.items() if values}
95 if not applied:
96 return cfg
97 merged = config_loader.deep_merge_dicts(cfg.model_dump(), applied)
98 return config_loader.validate_config(
99 type(cfg), merged, context="CLI overrides", error_cls=ConfigError)
Importance #15: test/test_harness_smoke.py @@ -23,9 +23,9 @@
23from iolabs_ml_harness.trainer import build_trainer # noqa: E40223from iolabs_ml_harness.trainer import build_trainer # noqa: E402
2424
25from src.dataset.rasters import index_raster_pairs # noqa: E40225from src.dataset.rasters import index_raster_pairs # noqa: E402
26from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E40226from src.dataset.synthetic import CLASS_NAMES, SyntheticGuardrailDataset # noqa: E402
27from src.train.config import DataConfig, HarnessConfig # noqa: E40227from src.train.config import ConfigError, DataConfig, HarnessConfig # noqa: E402
28from src.train.datamodule import GuardrailDataModule # noqa: E40228from src.train.datamodule import GuardrailDataModule # noqa: E402
2929
30_REPO_ROOT = Path(__file__).resolve().parents[1]30_REPO_ROOT = Path(__file__).resolve().parents[1]
31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")31_CLI_FLAGS = ("max_epochs", "batch_size", "model", "encoder")
Importance #16: test/test_harness_smoke.py @@ -65,12 +65,28 @@
65 cfg = HarnessConfig.from_yaml(good)65 cfg = HarnessConfig.from_yaml(good)
66 assert cfg.experiment == "t" and cfg.model.name == "unet"66 assert cfg.experiment == "t" and cfg.model.name == "unet"
67 bad = tmp_path / "bad.yaml"67 bad = tmp_path / "bad.yaml"
68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key68 bad.write_text("data:\n soruce: synthetic\n") # typo'd key
69 with pytest.raises(ValueError, match="soruce"):69 with pytest.raises(ConfigError, match="soruce"):
70 HarnessConfig.from_yaml(bad)70 HarnessConfig.from_yaml(bad)
7171
7272
73def test_null_sections_fall_back_to_defaults(tmp_path: Path) -> None:
74 """A bare ``data:`` line parses to None and must mean "all defaults"."""
75 path = tmp_path / "nulls.yaml"
76 path.write_text("experiment: t\ndata:\nmodel:\nloss:\n args:\n")
77 cfg = HarnessConfig.from_yaml(path)
78 assert cfg.data.crop_size == 256 and cfg.data.source == "synthetic"
79 assert cfg.model.name == "unet" and cfg.loss.args == {}
80
81
82def test_config_rejects_out_of_range_values(tmp_path: Path) -> None:
83 path = tmp_path / "bad.yaml"
84 path.write_text("train:\n max_epochs: -2\n")
85 with pytest.raises(ConfigError, match="max_epochs"):
86 HarnessConfig.from_yaml(path)
87
88
73def test_synthetic_dataset_is_deterministic() -> None:89def test_synthetic_dataset_is_deterministic() -> None:
74 ds_a = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)90 ds_a = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)
75 ds_b = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)91 ds_b = SyntheticGuardrailDataset(n_tiles=4, crop_size=64, seed=7, train=True)
76 for i in range(4):92 for i in range(4):
Importance #17: test/test_harness_smoke.py @@ -174,4 +190,18 @@
174190
175 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))191 updated = _train_script().apply_cli_overrides(cfg, _args(max_epochs=0))
176192
177 assert updated.train.max_epochs == 0193 assert updated.train.max_epochs == 0
194
195
196@pytest.mark.parametrize("flags", [{"max_epochs": -2}, {"batch_size": 0}])
197def test_apply_cli_overrides_validates_like_the_yaml_path(flags: dict) -> None:
198 """CLI overrides must hit the same bounds as values coming from YAML."""
199 with pytest.raises(ConfigError):
200 _train_script().apply_cli_overrides(HarnessConfig(), _args(**flags))
201
202
203def test_apply_cli_overrides_coerces_like_the_yaml_path() -> None:
204 updated = _train_script().apply_cli_overrides(
205 HarnessConfig(), _args(max_epochs="7"))
206
207 assert updated.train.max_epochs == 7