ML configuration system with inheritance and validation

The pain this solves. A research team runs thousands of training jobs. Each job's settings live in a YAML file copied from another one with a few edits. Six months later nobody can say which learning rate produced the best checkpoint, a typo (warmpu: 2000) silently trained with the default, and two "identical" runs differ because a default changed in code. The system must make every run's full configuration explicit, checked and reproducible.

Clarify first. Who writes configs (researchers, in YAML, by hand)? How deep does inheritance go (base model → size variant → experiment → sweep point: 3–4 levels)? Are configs code (Python, like some frameworks) or data (YAML/JSON)? Data here: it can be stored, diffed, hashed and validated without running arbitrary code. Must old runs replay after the schema changes (yes, which drives versioning)?

Functional

• A typed schema: every key has a type, a default or "required", and optional constraints (range, choices).
• Inheritance: a config names a base and overrides parts of it, many levels deep.
• Command-line overrides with dotted paths: optim.lr=1e-4 model.layers=48.
• References between values: scheduler.warmup_steps: ${optim.warmup}.
• Validation that reports every error with its path before a job starts, including unknown keys.
• A resolved, frozen config saved with each run, plus a content hash to group identical runs.

Non-functional

• Fail in seconds on the submitting machine, not 20 minutes later on 512 GPUs.
• Deterministic: resolving the same inputs always produces the same bytes and hash.
• Readable diffs between any two runs.

1. Resolution order (the core of the design).

schema defaults → base chain (root first) → this file → command-line overrides → interpolation → validation → freeze

Each step is a pure function of the previous result, so the output depends only on the inputs and can be recomputed.

2. Merge semantics, chosen on purpose.

• Dicts merge key by key, recursively. An experiment that changes optim.lr keeps the rest of optim.
• Lists replace, never append. Appending makes the result depend on how many ancestors touched the list, which nobody can reason about. If a list needs editing, make it a dict keyed by name.
• An explicit null deletes an inherited key (for optional sections such as eval.extra_sets).

3. Validation is strict. An unknown key is an error, and the message suggests the closest known key: optim.warmpu: unknown key (did you mean optim.warmup?). All errors are collected and shown together, so a user fixes five typos in one round, not five. Types are coerced only from command-line strings ("48" to 48); in files, "48" for an integer is an error.

4. Interpolation runs after merging. So an experiment that changes model.d_model also changes everything that references it. References must form no cycle; a cycle is an error naming the chain.

5. Reproducibility. The frozen config is written as canonical JSON (sorted keys, fixed float format) next to the checkpoint, together with the code commit, the schema version, the seed and the data snapshot id. Its SHA-256 is the run's config id: two runs with the same id ran the same settings. A diff between two frozen configs answers "what changed between these runs" without opening YAML files.

6. Schema evolution. Renaming a key breaks old configs. Keep a list of migrations (v3: rename optim.learning_rate → optim.lr) applied to old configs before validation, and store the schema version in every frozen config.

A compact version of the resolver and validator (configs as Python dicts, which is what a YAML loader produces):

import difflib
import hashlib
import json
import re

REQUIRED = object()
REF = re.compile(r"^\$\{([a-zA-Z0-9_.]+)\}$")


class Field:
    def __init__(self, type_, default=REQUIRED, choices=None, min=None, max=None):
        self.type, self.default, self.choices, self.min, self.max = type_, default, choices, min, max


class ConfigError(Exception):
    def __init__(self, errors):
        super().__init__("\n".join(errors))
        self.errors = errors


def _defaults(schema):
    out = {}
    for key, spec in schema.items():
        if isinstance(spec, dict):
            out[key] = _defaults(spec)
        elif spec.default is not REQUIRED:
            out[key] = spec.default
    return out


def merge(base, over):
    """Dicts merge recursively, everything else replaces; None deletes."""
    out = dict(base)
    for key, value in over.items():
        if value is None:
            out.pop(key, None)
        elif isinstance(value, dict) and isinstance(out.get(key), dict):
            out[key] = merge(out[key], value)
        else:
            out[key] = value
    return out


def _chain(name, registry):
    chain, seen = [], []
    while name is not None:
        if name in seen:
            raise ConfigError([f"base cycle: {' -> '.join(seen + [name])}"])
        if name not in registry:
            raise ConfigError([f"unknown base config {name!r}"])
        seen.append(name)
        chain.append(registry[name])
        name = registry[name].get("base")
    return list(reversed(chain))  # root first


def _set_path(cfg, path, value):
    node = cfg
    *parents, last = path.split(".")
    for p in parents:
        node = node.setdefault(p, {})
    node[last] = value


def _get_path(cfg, path):
    node = cfg
    for p in path.split("."):
        if not isinstance(node, dict) or p not in node:
            raise KeyError(path)
        node = node[p]
    return node


def _spec(schema, path):
    node = schema
    for p in path.split("."):
        if not isinstance(node, dict) or p not in node:
            return None
        node = node[p]
    return node


def _coerce(text, spec):
    if spec is None or isinstance(spec, dict):
        return text
    if spec.type is bool:
        return text.lower() in ("1", "true", "yes")
    try:
        return spec.type(text)
    except ValueError:
        return text  # leave it; validation reports the type error


def _interpolate(cfg):
    def resolve(value, stack):
        if isinstance(value, dict):
            return {k: resolve(v, stack) for k, v in value.items()}
        m = REF.match(value) if isinstance(value, str) else None
        if not m:
            return value
        target = m.group(1)
        if target in stack:
            raise ConfigError([f"reference cycle: {' -> '.join(stack + [target])}"])
        try:
            return resolve(_get_path(cfg, target), stack + [target])
        except KeyError:
            raise ConfigError([f"reference to missing key {target!r}"]) from None
    return resolve(cfg, [])


def validate(cfg, schema, prefix=""):
    errors = []
    for key, value in cfg.items():
        path = prefix + key
        if key == "base" and not prefix:
            continue
        spec = schema.get(key)
        if spec is None:
            close = difflib.get_close_matches(key, list(schema), n=1)
            hint = f" (did you mean {prefix + close[0]}?)" if close else ""
            errors.append(f"{path}: unknown key{hint}")
        elif isinstance(spec, dict):
            if isinstance(value, dict):
                errors += validate(value, spec, path + ".")
            else:
                errors.append(f"{path}: expected a section")
        else:
            ok_type = isinstance(value, spec.type) and not (spec.type is int and isinstance(value, bool))
            if spec.type is float and isinstance(value, int) and not isinstance(value, bool):
                ok_type = True
            if not ok_type:
                errors.append(f"{path}: expected {spec.type.__name__}, got {value!r}")
                continue
            if spec.choices is not None and value not in spec.choices:
                errors.append(f"{path}: {value!r} not in {sorted(spec.choices)}")
            if spec.min is not None and value < spec.min:
                errors.append(f"{path}: {value} < min {spec.min}")
            if spec.max is not None and value > spec.max:
                errors.append(f"{path}: {value} > max {spec.max}")
    for key, spec in schema.items():
        if not isinstance(spec, dict) and spec.default is REQUIRED and key not in cfg:
            errors.append(f"{prefix + key}: required")
    return errors


def resolve(name, registry, schema, overrides=()):
    cfg = _defaults(schema)
    for layer in _chain(name, registry):
        cfg = merge(cfg, {k: v for k, v in layer.items() if k != "base"})
    for item in overrides:
        path, _, text = item.partition("=")
        _set_path(cfg, path, _coerce(text, _spec(schema, path)))
    cfg = _interpolate(cfg)
    errors = validate(cfg, schema)
    if errors:
        raise ConfigError(errors)
    return cfg


def config_hash(cfg):
    canonical = json.dumps(cfg, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode()).hexdigest()[:16]

from solution import ConfigError, Field, config_hash, merge, resolve

SCHEMA = {
    "model": {"d_model": Field(int, 512, min=64), "layers": Field(int, 12, min=1)},
    "optim": {"lr": Field(float), "name": Field(str, "adamw", choices={"adamw", "sgd"}),
              "warmup": Field(int, 1000)},
    "scheduler": {"warmup_steps": Field(int, 0)},
    "tags": Field(list, []),
}

REGISTRY = {
    "base": {"optim": {"lr": 3e-4}, "tags": ["base"]},
    "large": {"base": "base", "model": {"d_model": 2048, "layers": 48}},
    "exp1": {"base": "large", "optim": {"lr": 1e-4}, "scheduler": {"warmup_steps": "${optim.warmup}"}, "tags": ["exp1"]},
}


def test_three_level_inheritance_merges_dicts_and_replaces_lists():
    cfg = resolve("exp1", REGISTRY, SCHEMA)
    assert cfg["model"] == {"d_model": 2048, "layers": 48}
    assert cfg["optim"] == {"lr": 1e-4, "name": "adamw", "warmup": 1000}
    assert cfg["tags"] == ["exp1"]
    assert cfg["scheduler"]["warmup_steps"] == 1000


def test_cli_overrides_are_typed_and_feed_interpolation():
    cfg = resolve("exp1", REGISTRY, SCHEMA, ["optim.warmup=500", "model.layers=24"])
    assert cfg["model"]["layers"] == 24
    assert cfg["scheduler"]["warmup_steps"] == 500


def test_all_errors_are_reported_with_suggestions():
    reg = dict(REGISTRY, bad={"base": "base", "optim": {"warmpu": 10, "name": "adam"}, "model": {"layers": 0}})
    try:
        resolve("bad", reg, SCHEMA)
    except ConfigError as e:
        text = "\n".join(e.errors)
        assert "optim.warmpu: unknown key (did you mean optim.warmup?)" in text
        assert "optim.name: 'adam' not in" in text
        assert "model.layers: 0 < min 1" in text
    else:
        assert False, "expected ConfigError"


def test_required_and_cycles():
    try:
        resolve("x", {"x": {}}, SCHEMA)
    except ConfigError as e:
        assert "optim.lr: required" in e.errors
    try:
        resolve("a", {"a": {"base": "b"}, "b": {"base": "a"}}, SCHEMA)
    except ConfigError as e:
        assert "base cycle" in e.errors[0]
    else:
        assert False


def test_null_deletes_and_hash_ignores_key_order():
    assert merge({"a": 1, "b": 2}, {"b": None}) == {"a": 1}
    assert config_hash({"a": 1, "b": {"c": 2}}) == config_hash({"b": {"c": 2}, "a": 1})
    assert config_hash({"a": 1}) != config_hash({"a": 2})

optim.lr is declared as Field(float) with no default, so it is required.

• Data versus code configs. Python configs (functions returning objects) are flexible and can compute values, but they cannot be validated, hashed or diffed without running them, and they invite hidden logic. YAML plus a schema is less flexible and much more auditable. A middle path: data configs, with a small, named set of computed fields defined in the schema.
• Strict unknown-key errors annoy people during exploration. Offer an explicit scratch: section that is stored but not validated, rather than loosening validation.
• Deep inheritance saves typing and costs readability: a value can come from four files. Mitigate with an explain model.layers command that prints which layer set each value. Cap depth (4 is plenty).
• Interpolation keeps related values in sync but can hide where a value comes from; resolve it before freezing so the stored config has only literal values.
• Sweeps generate many configs. Store the sweep definition and derive each point as overrides on one base, so the frozen configs differ only in the swept keys.

Mid-level (L4)

• Defines a typed schema with defaults and validates configs against it.
• Implements single-level inheritance with a recursive dict merge and dotted command-line overrides.
• Saves the resolved config with each run.

Senior (L5)

• Specifies merge semantics deliberately (lists replace, null deletes), multi-level inheritance with cycle detection, and interpolation after merging.
• Collects all errors with paths and suggestions, and rejects unknown keys.
• Makes runs reproducible: canonical serialization, content hash, code commit and seed stored together.

Staff (L6+)

• Plans for schema evolution with versioned migrations and for sweeps as overrides on a base.
• Weighs code-as-config against data configs and explains the auditing and tooling consequences.
• Adds tooling that makes the system usable at team scale: explain-a-value, diff between runs, validation on submit.