Conversation
ACEsuit#1556) `mace_core.config.ReforgeBaseConfig` is the pydantic base every v1 config schema derives from, and `mace_core.metadata.ModelMetadata` the versioned record every trained model carries. Both are framework-free; a test in a fresh interpreter asserts neither imports torch, jax or e3nn. Config: - A schema is a tree of `ConfigSection`s; `load()` reads one TOML, YAML or JSON file and applies dotted CLI overrides (`--model.radial.cutoff 5.0`, `--seed=7`), precedence defaults < file < CLI. Nothing else feeds a config: no environment variables, no dotenv files. - Overrides are parsed by a plain tokenizer, not pydantic-settings or argparse. A value starting with `[` or `{`, or `null`, is JSON, so a whole section, list or dict can be given. Each override merges onto the file values in order by one rule: dict-valued fields gain or replace entries, list-valued fields are replaced whole. - Unknown keys are `ConfigError`s that name the key by its dotted path and the nearest valid neighbour: CLI keys are checked against the schema's dotted paths before anything is built, file keys and keys inside a JSON-valued override by pydantic (`extra="forbid"` on every level). Wrong types re-raise pydantic's `ValidationError`. - `to_resolved_dict()` exports every field, defaults filled in, in schema order, as a fixed point: loading it back and resolving again gives the same dict. `to_user_dict()` exports only what the file and CLI set. - Field shapes the contract cannot keep are rejected at class definition: sets (order is not stable across runs), aliases and computed fields (the export would not validate back), a union of section classes (a value must not become whichever alternative accepts it; offer alternatives as one optional field each), and sections that do not subclass `ConfigSection`. Metadata: - `ModelMetadata` holds the config as written and as resolved (`ConfigRecord.from_config`), provenance (code version, git commit), one `DataSourceSummary` per data source with the `<method>_<quantity>` reference-key convention, `E0Details` per head (explicit or estimated, with method and parameters), the model's DOI, structured citations, notes, and `parents`: the models it was built from, each with a role (`initial_weights` for fine-tuning and continued training, `teacher` for distillation) and its own record embedded, so the lineage travels with the model. - `to_json()` refuses a record that would not read back equal, so a tuple, datetime or NaN is never stored as something else; infinity round-trips. `from_json()` checks `schema_version` before any field is read and rejects a newer version with a message naming both versions. - `format_citations()` renders a numbered block for printing. Floors: Python 3.10, pydantic 2.7; the suite passes at both floors and at current releases. Tests live in `test_mace_core_config.py` and `test_mace_core_metadata.py`.
…suit#1556) `ModelMetadata.heads` replaces the `e0` dict: a `HeadSummary` per head holds its `E0Details` and the names of the entries in `data.sources` it was fitted on, so the source-to-head wiring can be read from the record without the resolved config. A validator checks that every named source is summarised and that each source is summarised once, so a source shared by two heads is not counted twice. Docstrings: a source's `reference_keys` are the keys it provides, not what a head reproduces; `to_json` claims only that a text reading back to a different record is refused; a version mismatch inside an embedded parent record is pydantic's error, since a record only embeds parents it could read.
Both tickets now have pull requests from the people assigned to them, so the versions I wrote to keep the integration branch moving are gone and theirs are in: ACEsuit#1732 from steffen-wedig for ARCH-1 and ACEsuit#1733 from arnon-1 for CORE-2. His ARCH-1 is more complete than mine was in every direction: eight radial classes rather than four, the embedding blocks, and a closed-form spherical harmonics evaluated by recurrence rather than my recursion through the coefficients. His is also more robust where it matters, since the homogeneous form keeps a zero vector finite. The cross-ticket agreement held, and that was the thing worth checking. We arrived independently at the same fact about the convention: his `E3NN_AXIS_ORDER` and the `AXIS_PERMUTATION` derived here are both (2, 0, 1), which is the whole difference between e3nn's harmonics and the textbook ones. So his harmonics and this branch's Clebsch-Gordan basis are in the same basis, and the convention test written against mine passes unchanged against his. His Chebyshev fixes the defect this branch recorded, also independently. It evaluates the three-term recurrence instead of calling `torch.special.chebyshev_polynomial_t`, which returns a tensor carrying no gradient, so a force computed through the frozen tree's basis is missing the term through the polynomial. Measured on his: requires_grad true, gradient norm 30.19. Two mechanical consequences of the merge. `backends/reference` becomes the package his layout wants, with this branch's backend moved into it as `backend.py` so the entry point still resolves; and `ReferenceRadialBasis` now builds ARCH-1's classes instead of the simplified ones written here, so there is no second implementation of a basis anywhere. His Chebyshev takes no r_max, correctly, because the frozen tree stored one and never used it, so the lookup passes it only to the bases that have it. 1062 tests, `tests/parity` included, which his branch creates.
…(CORE-2 follow-up, ACEsuit#1556) A config field may be a discriminated union of sections, "kinds": loss: Annotated[WeightedLoss | HuberLoss, Field(discriminator="kind")] = WeightedLoss() Code sees the union. A file or override never writes the tag; it writes the section under its kind as the key, `loss: {huber: {delta: 0.1}}` or `--loss.huber.delta 0.1`, or as a bare name, `loss: huber`, which is that kind with nothing set under it. The file and each override merge by the plain deep update of CORE-2; the loader records which layer last wrote each path and at each kinds field keeps the kind written last, dropping the others with a ConfigWarning that names both layers. Two kinds written in one place, and null at or under a kind, are ConfigErrors. A field with no kind written takes the kind of its default. The resolved and user dicts write the kind as key, so the resolved dict stays a fixed point. A before-validator on ConfigSection turns the key form into pydantic's tagged form and a wrap serializer turns it back; the tagged form is not a file format and load() refuses it. The schema check rejects at class definition what the contract cannot keep: a kinds field typed `| None` ("none of the kinds" is an empty variant, `kind: Literal["none"]`), a default that is None or a factory or not a variant, a non-variant arm, a kind named like the tag, and a union of sections inside a collection. Error paths name the kind as written, `loss.huber.delta`, and did-you-mean suggests kinds and their keys.
ef1ad1e to
e5339d4
Compare
| """A nested section of a configuration: a table in the file, a dotted | ||
| prefix on the command line. Unknown keys are errors here too.""" | ||
|
|
||
| model_config = ConfigDict(extra="forbid") |
There was a problem hiding this comment.
inf and nan come out as null which would put the wrong value into the model metadata
| model_config = ConfigDict(extra="forbid") | |
| model_config = ConfigDict(extra="forbid", ser_json_inf_nan="constants") |
| ) | ||
|
|
||
|
|
||
| def _check_schema(model: type[BaseModel]) -> None: |
There was a problem hiding this comment.
I think these checks run before pydantic resolves forward references, so a forward-ref field gets an error that contradicts itself:
class CA(ConfigSection):
q: "CB | None" = None
-> TypeError: CA.q defaults to None but its type does not admit None; add | None
The other half is the one that worries me. A forward-ref field with a non-None default slips past every guard and then a plain BaseModel behind it silently drops unknown keys:
class Root(ConfigSection):
leaky: "Leaky" = Field(default_factory=lambda: Leaky())
class Leaky(BaseModel): # not a ConfigSection, so it allows extra keys
a: int = 1
Root.model_rebuild()
class Demo(ReforgeBaseConfig):
root: Root = Field(default_factory=Root)
# demo.json is {"root": {"leaky": {"a": 2, "TYPO": 9}}}
print(Demo.load("demo.json", []).to_resolved_dict())
# -> {'root': {'leaky': {'a': 2}}}. # TYPO gone
With the class defined before Root, the same shape is caught properly, so it's the unresolved ForwardRef that opens the hole. Running the checks at model_rebuild() time or re-checking once the ForwardRef resolves would cover both halves.
…t (CORE-2 follow-up, ACEsuit#1556) The schema check ran once, at class definition. A field whose annotation still named a class defined later was an unresolved forward reference then, opaque to every rule: `q: "Later | None" = None` was rejected with an error that contradicted itself, and a plain BaseModel behind such a name passed the ConfigSection rule and silently dropped unknown keys. Pydantic's rebuild of an outer model does not rebuild an inner one, so a rebuild hook alone would miss a nested section. An incomplete class is now skipped at definition; `load` completes and checks every section of the tree before it reads anything. Found in review of ACEsuit#1733.
|
@cursor review verbose=true |
|
Bugbot request id: serverGenReqId_da7a7463-4743-4c1a-8fca-79119896cff3 |
Bugbot rules debugNo rules were used for this review. https://cursor.com/docs/bugbot#team-rules Bugbot request id: serverGenReqId_da7a7463-4743-4c1a-8fca-79119896cff3 |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit e5339d4. Configure here.
| reject( | ||
| name, | ||
| f"has a default that is not one of its variants; write e.g. {example}", | ||
| ) |
There was a problem hiding this comment.
Schema checks skip forward references
Medium Severity
_check_schema runs in __pydantic_init_subclass__ on field.annotation before pydantic resolves forward references, and it is not run again after rebuild. A quoted X | None defaulting to None is rejected with a message that says to add | None, which is already there. A quoted nested plain BaseModel with a non-None default bypasses the ConfigSection requirement, so unknown keys on that nested model are ignored instead of being errors.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit e5339d4. Configure here.
…te (CORE-2, ACEsuit#1556) The tests now pin what `mace_core.config.base` must do after its clean-room rewrite: any number of files in order, settings of every kind kept across layers with the selection separate (`kind:`, `--loss.kind`, bare name), a warning for every override without effect and never for a file, `kind` as the tag only under a kinds field, sections never optional, and unknown keys reported under their file or override. They fail against the module on the branch on purpose; the rewrite follows in the next commit. The kinds error is pinned as one sentence naming the fix and the source of each kind key; the same override twice warns for the first; a variant subclass instance exports under its kind; `.YAML` loads and `load(None)` is `load(())`. Both config modules turn warnings into errors (`pytestmark`), so a spurious `ConfigWarning` fails the suite.
…ion (CORE-2, ACEsuit#1556) `load` flattens every file (any number, in order) and every override into leaf updates carrying their source, walks each path once through the schema (unknown keys and wrong shapes with full dotted paths and neighbours, a bare kind name rewritten to its `kind`, the kinds fields a path enters recorded on the update), merges them set-at-path into one dict, validates once, and then warns for each override that a later override wrote over or that wrote under a kind that does not run (C11): two rules over the override list, the running kind decided from the merged dict by the rule the validator uses. Settings of every kind are kept across layers; the selection is a separate scalar, so a file may tune several kinds and a later layer picks one. `kind` is the tag only by position under a kinds field: a variant class used as a plain section field keeps it as a key on input and export. The kinds error is one `PydanticCustomError` (two kind keys unselected, a required kinds field with nothing written, the tag under a kind), rendered by `load` with the path, the fix and the source of each kind key; it ends the validation of its class, so that class's other errors come on the next load (rebuilding them beside it would need pydantic to know every custom error type). The exports accept a variant subclass instance under the variant's kind. A variant field named like a kind is rejected at class definition (F6). The file extension is case-folded, `load(None)` is `load(())`, and `tomli` is required below Python 3.11 only. The definition-time schema rules (F1-F7, G5) and the annotation introspection move to `config/_schema_rules.py`, which binds `base` as a module so that the import cycle resolves whichever side is imported first. `typing-extensions` is declared for `Self` (the floor is 3.10).
ACEsuit#1556) pydantic's JSON mode writes inf and nan as null, so a config holding either stored a different value in the model metadata and could not be loaded back (a null is not a float). `ConfigSection` now sets `ser_json_inf_nan="constants"`, as `metadata._Record` already does: both exports keep them as floats, JSON writes Infinity and NaN, and JSON, YAML and TOML each read their own spelling back to the same config. From the PR ACEsuit#1733 review on e5339d4. The review's other comment, a schema behind a forward reference escaping the checks, is covered by rule F7 and its tests.
…te (CORE-2 follow-up, ACEsuit#1556) The config base is reduced to one file, one pydantic validation and two exports, and the command-line half moves to its own module. The tests now pin exactly that. `test_mace_core_config.py`: `load(config_file)` parses one TOML, YAML or JSON file and validates it; `from_dict(document)` is the same validation for a caller that edits the parsed dict first and does not write into it; everything the schema rejects, unknown keys included, is pydantic's `ValidationError` at its dotted location, and `ConfigError` is raised only for a file that cannot be read or parsed. `test_mace_core_config_kinds.py`: a kinds field is written in pydantic's tagged form and its errors land under the tag; both exports write the tag back, with the documented caveat that a variant built in code without its tag exports without `kind`. The definition-time checks keep the set, alias, excluded and computed-field rules, the lenient section rule (also behind a forward reference) and the `extra="forbid"` guard. `test_mace_core_config_cli.py` is new and pins the command-line half: `parse_overrides` reads `--a.b value` and `--a.b=value` tokens into a mapping of dotted paths (`null`, `[` and `{` values are JSON, anything else a string for pydantic to coerce; a bad token, a missing value or bad JSON is a `ConfigError`), and `apply_overrides` writes such a mapping into a copy of the parsed file set-at-path, creating mappings on the way, replacing a parent that is not a mapping, merging a mapping value into a mapping and replacing otherwise, without writing into either argument, so a YAML anchor's two keys stop sharing one object and an anchor that contains itself is a `ConfigError`. Composed with `read_config_file` and `from_dict`, an override beats the file, which beats the defaults, and an unknown key in an override is pydantic's error at the path the override named. Gone with the contract: several files, merge and warning rules, `ConfigWarning`, the kind-as-key wire form, bare kind names, kind switching, the union-shape schema rules and the unknown-key neighbour messages. Tests revision 7 does not contradict are kept, some renamed or reworded. These tests fail on purpose against the module on the branch; the next commit replaces it.
…ine beside it (CORE-2 follow-up, ACEsuit#1556) `base.py` knows nothing of a command line. `load(config_file)` reads one TOML, YAML or JSON file (`read_config_file`) and validates it once (`from_dict`, the same validation for a caller that edits the parsed dict first; it does not write into the dict). Everything the schema rejects is pydantic's `ValidationError` at its dotted location, unknown keys included; `ConfigError` is raised only for a file that cannot be read or parsed. The two exports are the plain `model_dump` calls, with inf and nan kept as JSON constants. A kinds field is pydantic's discriminated union in its tagged form; the module has no knowledge of it. `cli.py` is the command-line half, importing only `ConfigError` from the base. `parse_overrides` reads `--a.b value` and `--a.b=value` tokens into a mapping of dotted path to value (`null`, `[` and `{` values are JSON, anything else a string for pydantic to coerce; a token that is not an option, a missing value or bad JSON is a `ConfigError`). `apply_overrides` returns a copy of the parsed file with each override written at its path, creating mappings on the way, replacing a parent that is not a mapping, merging a mapping value into a mapping already there and replacing otherwise. The copy takes dicts and lists apart, so a YAML anchor's two keys stop sharing one object and an anchor that contains itself is a `ConfigError`. A command line builds its config as `Config.from_dict(apply_overrides(read_config_file(path), parse_overrides(rest)))`. `ConfigSection` keeps the definition-time checks only: `extra="forbid"` cannot be reopened, every model a field reaches is a section, no set of any kind (abstract ones included, since pydantic validates them to a frozenset), no aliases, excluded or computed fields. Every `from_dict` walks the sections its root reaches, rebuilds one a forward reference left incomplete, and checks each; re-checking is a few attribute reads per class and needs no record of what was checked. No validator, serializer or schema walk runs at load time. `_schema_rules.py` is deleted and `ConfigWarning` no longer exists; the public names are `ReforgeBaseConfig`, `ConfigSection`, `ConfigError` and `read_config_file` from the base, `parse_overrides` and `apply_overrides` from `cli`. The package README describes the new surface.
|
Ok so as requested, I dropped most of the requirements in the original ticket. So the current cli is pretty basic but the config itself is basically the same. I do still need to review this more carefully. |
steffen-wedig
left a comment
There was a problem hiding this comment.
Very nice changes to make the config scaffold much slimmer. I think most of the remaining issues are only there, because the rest of the package scaffold doesnt exist yet particularly unclear to me at the moment is how do we route the configs to the module builders? We also apparently have two construction methods at the moment planned, one for dispatched backends, and then another one for the remaining torch modules.
… share the test schema (CORE-2 follow-up, ACEsuit#1556) Decisions from the reforge meeting: - ModelMetadata no longer records a head's E0s (E0Details is gone; the head's parameters in the model hold them, so nothing is lost) nor the records of the models it was built from (ParentModel and `parents` are gone for now). HeadSummary keeps the sources a head consumed. - Provenance records `versions`, a version per distribution involved (`{"mace-core": ..., "mace-torch": ...}`), in place of the single `code_version` of mace-core: the packages version independently, so one number does not identify the code. - A data source's `elements` are atomic numbers, not chemical symbols. - ReforgeBaseConfig is BaseConfig: the class name outlives the project name. - A ConfigSection is frozen once validated: what validation produced is what the run uses, and a change is a new validation of an edited dict (`from_dict`), never an assignment behind it. `frozen=True` cannot be reopened by a subclass, like `extra="forbid"`. - The two config test modules defined the same demo schema twice; it now lives once in tests/mace_core_demo.py (the union of both), and the three CLI tests that only repeated a base test through parse/apply are folded into the composed override test. Coverage of mace_core is unchanged (98%, the same five uncovered lines: the tomli fallback, the uninstalled-version fallback and two citation-rendering branches). - The exports' JSON mode and their inf setting are now pinned under a free field; either could be removed with every test passing. The config tests no longer promise that NaN survives: the metadata refuses to store one. - The config docstrings and the README entry are cut down to what the code does not say by itself.
… what they do (CORE-2 follow-up, ACEsuit#1556) - read_config_file no longer refuses a YAML anchor that contains itself. Nobody writes one, and without the check it still fails loudly: a pydantic ValidationError under a typed section, a ValueError at export under a free field. - _leaf_types is _types_in: it returned every class an annotation mentions, containers included, and now filters out non-classes itself. - _check_field_declarations is _check_section_fields and _check_sections_reached_by is _check_schema; their docstrings say why the check runs at definition and again on load (a forward reference leaves a section unchecked until then). - The docstrings of from_dict, to_resolved_dict and to_user_dict are one sentence each.
…th write from the merge (CORE-2 follow-up, ACEsuit#1556) - parse_overrides no longer refuses an empty key in a path (`--a..b`): it reaches pydantic as the unknown key "" and fails there. Only under a free dict field is it stored, where it shows in the export. - A value that contains itself, or JSON nested past the recursion limit, is no longer turned into a ConfigError, in apply_overrides, parse_overrides and read_config_file alike: nobody writes one, and it still fails, as a RecursionError. - _write did two things and is two functions: _set_at_path walks the dotted path and _set_or_merge sets the value or merges a dict into a dict already there. - The docstrings of parse_overrides and apply_overrides open with one sentence that says what the function does.
…s only the version (CORE-2 follow-up, ACEsuit#1556) - The schema version was written twice, in SCHEMA_VERSION and in the `Literal[1]` of the field. The field is now an int defaulting to the constant, so a bump is one edit. Only from_json checks the version; to_json still refuses a wrong one because it reads its own output back. Validating a dict directly no longer checks it. - from_json keeps the one check nothing else makes, the version, ahead of validation so that a newer record reports its version and not its unknown fields. Invalid JSON is json's own error, and a missing or non-integer version fails the same comparison. A top level that is not an object is now an AttributeError; to_json never writes one. - The round-trip error of to_json names what does not survive (a tuple, NaN) instead of the docstring. - The docstrings are cut down; the reference-key convention sits above `reference_keys`.
…on (CORE-2 follow-up, ACEsuit#1556) The check moves from from_json to a before-validator on ModelMetadata, so the constructor, model_validate and model_validate_json all refuse a record of another version, and do so before any field is read: a newer record reports its version, not its unknown fields. - from_json is model_validate_json. Invalid JSON and a top level that is not an object are pydantic's errors again. - A wrong version is a pydantic ValidationError carrying the message; MetadataSchemaError is left for a value that does not survive JSON. - A record with no schema_version key reads as the current version: the validator cannot tell it from a record built in code. to_json always writes the key.


Hi,
Fixes #1556
Adds
mace_core.config.ReforgeBaseConfigandmace_core.metadata.ModelMetadata.Deviations from the ticket:
No pydantic-settings as it doesn't entirely cover the use-case and writing a wrapper ends up in more convoluted code than just writing a simple parser.
test files use the package prefix e.g.
test_mace_core_config.pyso not literally the same as the verify command in the ticket but adheres to the packages README rule.Some choices not mentioned in the ticket (non-exhaustive).
Some edge-cases can still raise errors different from ConfigError or MetadataSchemaError, but it doesn't make sense to handle every single edge case explicitly as the code would become quite large.
I am rejecting a number of patterns that results in non-deterministic json (like sets) or are bug-prone (like overloading parameters so that pydantic implicitly selects classes).
Field groups adapted a bit to take into account multi-dataset multi-head models. Adds a bit of overhead for single-head model data. Can change it if someone disagrees with the way I did it here.
Also added a parent model field for fine-tuning or distillation (which uses ModelMetadata again recursively), I thought this was nice for provenance, but can remove it if seen as unneeded.
Note
Medium Risk
New foundational config and checkpoint-metadata contracts will shape all future v1 training and loading; behavior is well-tested but downstream adopters must match precedence, kinds syntax, and metadata schema version 1.
Overview
Replaces the mace-core scaffold with the v1 configuration and model metadata contracts (issue #1556).
ReforgeBaseConfig/ConfigSectionadd a custom Pydantic loading path: one TOML/YAML/JSON file plus--dotted.clioverrides (defaults → file → CLI, no env vars). Unknown keys raiseConfigErrorwith typo suggestions; discriminated “kinds” sections use kind-as-key in files/CLI with merge rules andConfigWarningwhen a later layer switches variant. Schemas are guarded at class definition (no sets, aliases, undiscriminated section unions). Exports includeto_resolved_dict()/to_user_dict()for JSON-stable, round-trippable configs.ModelMetadatadefines schema version 1 checkpoint sidecar data: user vs resolved config (ConfigRecord.from_config), provenance, per-source and per-head summaries (including E0s), citations, optional recursiveparentsfor fine-tune/distill lineage, strictextra=forbid, andto_json()/from_json()with version mismatch and lossy-value rejection.format_citations()renders a numbered bibliography block.Package dependencies are now
pydantic,pyyaml, andtomli; the top-levelmace_corepackage re-exports the public types. Large pytest suites cover file formats, CLI/kinds behavior, and metadata round-trips (including no torch/jax import leak).Reviewed by Cursor Bugbot for commit e5339d4. Bugbot is set up for automated code reviews on this repo. Configure here.