feat(models): harden artifact loading and registry verification

[zubatyuk]

AIMNetCentral Pull Request

Description

Secure model loading across official registry, direct v2, legacy TorchScript, and Hugging Face sources. Registry artifacts are now digest-verified before use, v2 artifacts are restricted and schema-validated before construction, and Python references in model YAML are checked against a documented trusted set.

This prevents corrupted or replaced registry artifacts, same-named implicit local files, unsafe pickle fallback, sidecar YAML expansion, and unapproved constructors from bypassing the source-specific trust boundary. Direct custom artifacts retain explicit controls for trusted custom code, while registry policy cannot be weakened by caller options.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Changes Made

• Require every registry model entry to provide a lowercase SHA-256 digest, verify bundled, cached, and downloaded bytes, and install verified artifacts atomically.
• Resolve registry names and aliases before implicit local paths while preserving explicit absolute, ./, and ../ local paths.
• Load v2 artifacts once with torch.load(..., weights_only=True), validate their envelope, metadata, tensor-only state dict, and model YAML, and disable sidecar YAML references before model construction.
• Route .jpt files explicitly through torch.jit.load() as a trusted-code boundary without weakening the existing load_model() compatibility path.
• Add the public immutable ALLOWED_MODEL_IMPORT_PATHS set and keyword-only model_import_paths and model_import_mode controls for direct local and complete Hugging Face custom artifacts.
• Keep registry names, registry Hugging Face fallback, raw modules, and .jpt sources on fixed default import settings, including validation before remote weight download or model construction.
• Document loader dispatch, cache recovery, custom import modes, unsafe-loading risk, registry publication requirements, and Hugging Face fallback behavior.

Compatibility

• load_model() continues to accept both v2 .pt and legacy .jpt files; .jpt suffix matching is case-insensitive and now selects the trusted TorchScript loader directly.
• Bare registry names now take precedence over same-named implicit local files; callers that intend a local artifact must use an explicit relative or absolute path.
• Direct custom v2 and complete Hugging Face artifacts use the default trusted import set unless callers select extend, replace, or explicit unsafe behavior.
• Registry access now rejects entries without a valid digest and cached or downloaded bytes that do not match it.

Testing

• Unit tests pass locally.
• Ruff checks pass.
• New tests cover the changed behavior.
• Documentation builds without warnings.

Checklist

• I have performed a self-review of the code and documentation.
• I have checked that the code and documentation follow the project style.
• I have checked that the documentation is clear and readable.
• I have updated the changelog.

[isayev]

Reviewed the full diff and verified the fleet directly: all 24 registry entries carry valid digests that match the published GCS bytes (and the bundled assets), every registry artifact loads end-to-end under the strict policy with weights_only=True, and the embedded model_yaml of every published model stays exactly within the default allowlist. The security design is sound — validation ordering at each trust boundary (config/YAML/metadata before weight download, digest before install, atomic replace) is right, and the test suite is genuinely adversarial.

Three compatibility regressions on documented non-registry paths need a decision before merge:

1. HF loader rejects repos that follow its own documented contract (hf_hub.py:222). validate_model_metadata(..., require_cross_field_consistency=True) runs on raw config.json before the SRCoulomb rc/envelope extraction fallback (hf_hub.py:275) can fill those fields. A third-party HF repo with coulomb_mode="sr_embedded" and rc/env omitted — derivable from model_yaml, and still documented as optional in the tutorial table — loads on main but is rejected on this branch (confirmed empirically). This also leaves _extract_sr_coulomb_from_config dead for its intended purpose. Either run the extraction before the cross-field check, or update docs/tutorials/hugging_face.md and drop the fallback. The 5 official isayevlab/* repos carry explicit fields, so only third-party repos break.

2. Load-time metadata rejection removes the calculator's documented override escape (artifact_validation.py:353). An artifact with needs_dispersion=True, d3_params=None used to load and could be rescued by AIMNet2Calculator(..., needs_dispersion=False) ("Explicit flags override metadata"). It is now rejected inside load_model, before any override applies, and no import mode relaxes metadata validation. Either gate the strict cross-field rules to the registry policy or add a caveat to the calculator docs.

3. aimnet export can produce artifacts its own loader refuses (aimnet/train/export_model.py, untouched by this PR). aimnet export --no-coulomb on a YAML containing LRCoulomb writes coulomb_mode="sr_embedded" with needs_coulomb=False and null rc/env — rejected by three of the new rules (confirmed on this branch); --needs-dispersion without D3 params likewise. The export CLI should call validate_model_metadata before saving so it cannot emit unloadable files.

Also important:

• The fleet-compatibility guarantee never runs in CI. test_registry_digests_match, test_every_registry_artifact_loads_with_strict_policy, and the role-defaults ABI test are all @pytest.mark.network, and both the Makefile and the GitHub workflow deselect them — the "no official artifact is rejected" invariant is only verified manually. A scheduled network CI job would close that.
• Security-relevant predicates are duplicated and already diverging. resolve.py:106 uses model.lower().endswith(".jpt") while base.py:126 uses Path(path).suffix.lower() == ".jpt" (they disagree for a file literally named .jpt); the explicit_local check is copy-pasted verbatim in model_registry.py:236 and resolve.py:108; the default-import-settings guard appears in four modules. One shared helper each (is_legacy_jit_path, is_explicit_local_path, ModelImportPolicy.is_default) removes the drift risk.
• Security-critical symbols are imported cross-package under private names. resolve.py imports _load_registry_model from aimnet.models.base; hf_hub.py imports _REGISTRY_IMPORT_POLICY; tests import/monkeypatch both. With three external consumers these should be public (load_registry_model, REGISTRY_IMPORT_POLICY). Relatedly, hf_hub.py:236-300 re-implements the v2 assembly pipeline from base.py:134-180 and the two have already drifted (model._metadata = ... vs model.__dict__["_metadata"] = ...) — worth extracting a single assembly helper in the models layer.
• load_legacy_jit metadata violates the new schema it ships with (base.py:61-75): coulomb_mode="full_embedded" without has_embedded_lr=True would fail the new cross-field rule if that metadata is ever validated; True is also the physically correct value for embedded-LR TorchScript models.
• A stale bundled asset blocks the download fallback by design (_acquire_asset): if a release ever bumps registry digests without rebuilding bundled assets, cold-cache loads hard-fail even though the URL serves correct bytes. Fail-closed is defensible, but this belongs on the release checklist: bundled assets and registry digests must be bumped together.

[isayev]

Any thoughts on the three compatibility items from the review — the sr-coulomb extraction fallback ordering, the load-time metadata rejection vs the calculator's explicit-flag override, and export-side validation? Happy to discuss trade-offs on any of them; the rest of the findings are non-blocking. Would like to get this merged soon since it closes real gaps.

[zubatyuk]

Thank you for the detailed review. I kept the source-specific security boundaries intact and addressed the three compatibility issues you identified. I also addressed the additional maintainability, CI, legacy-model, and release-policy concerns.

1. Hugging Face SRCoulomb metadata ordering

Request: A complete third-party HF repository could omit coulomb_sr_rc and coulomb_sr_envelope when those values were present in model_yaml, but the loader validated the incomplete config.json before attempting that derivation. This rejected repositories that followed the documented format.

Solution: The loader now parses and authorizes model_yaml, recursively collects SRCoulomb parameter pairs, derives omitted values, and only then performs structural metadata validation. Derivation is accepted only when there is exactly one distinct complete cutoff/envelope pair. Ambiguous definitions, incomplete pairs, and conflicts between explicit metadata and YAML fail before weights are resolved or read.

Registry-backed HF fallback remains stricter: digest-verified registry metadata and YAML are authoritative, family-level HF configuration is used only for routing, and duplicated artifact fields must exactly match the registry values.

Evidence: aimnet/calculators/hf_hub.py:121-148, aimnet/calculators/hf_hub.py:314-376, tests/test_hf_hub.py:277-475, docs/tutorials/hugging_face.md:23, and docs/tutorials/hugging_face.md:100-121.

2. Calculator overrides versus load-time rejection

Request: A structurally valid direct artifact with incomplete optional external-component metadata could previously be loaded and used with an explicit calculator override such as needs_dispersion=False. Strict load-time cross-field validation removed that documented escape path.

Solution: Metadata validation is now separated into four levels:

1. Envelope/schema validation checks the artifact container, safe YAML, authorized imports, scalar types, and tensor-only state values.
2. Structural validation checks intrinsic facts that caller flags cannot change.
3. Canonical validation additionally requires action flags and external-component metadata to agree for registry artifacts, registry-backed HF fallback, and exported artifacts.
4. Effective validation runs after calculator defaults and explicit flags have been resolved.

Direct local artifacts and complete third-party HF repositories use structural validation, so an explicit False may disable an otherwise valid external component. Registry and export paths remain canonical and fail closed. Overrides cannot bypass intrinsic errors, enable external Coulomb alongside fully embedded Coulomb, or enable external dispersion without complete D3 parameters or alongside embedded D3TS.

Evidence: aimnet/models/artifact_validation.py:228-250, aimnet/models/artifact_validation.py:291-430, aimnet/calculators/calculator.py:195-247, tests/test_calculator.py:1959-2144, docs/model_format.md:57-66, and docs/calculator.md:132-154.

3. Export could create artifacts rejected by the loader

Request: aimnet export could accept contradictory flags and serialize an artifact that the hardened loader immediately rejected.

Solution: Export now:

• rejects --no-coulomb when sr_embedded Coulomb was detected;
• rejects enabled dispersion unless s8, a1, and a2 are available;
• includes the embedded-D3TS metadata flag;
• validates safe YAML, import authorization, the complete v2 envelope, and canonical metadata consistency before replacing the output;
• supports explicitly trusted local custom constructors through repeatable --model-import-path options;
• serializes to an exclusively created sibling temporary file and atomically replaces the destination only after validation and serialization succeed;
• preserves existing destination permissions and applies the process umask to new files.

Validation or serialization failures therefore leave an existing output untouched.

Evidence: aimnet/train/export_model.py:83-116, aimnet/train/export_model.py:185-204, aimnet/train/export_model.py:267-285, tests/test_train_utils.py:148-318, and docs/model_format.md:234-265.

4. Official registry fleet verification in CI

Request: The network tests proving that all official models retain valid digests, load under strict policy, and use exact role-specific import defaults were excluded from normal CI.

Solution: A dedicated registry-fleet.yml workflow now runs weekly and supports manual dispatch. It uses an isolated cache and executes exactly these invariants:

• every published registry digest matches the downloaded bytes;
• every registry artifact loads end-to-end through the strict registry loader;
• digest-verified YAML uses the exact role-specific default import policy.

A fresh isolated-cache run of those exact test nodes passed all 49 cases.

Evidence: .github/workflows/registry-fleet.yml:1-35, tests/test_model_registry.py:575-601, and tests/test_serialization_abi.py:223-262.

5. Duplicated security-routing predicates

Request: .jpt detection, explicit-local-path detection, and default-import-setting checks were duplicated across modules and had already diverged on edge cases.

Solution: Shared helpers now define:

• case-insensitive legacy .jpt routing, including a file literally named .jpt;
• explicit absolute, ./, and ../ local paths;
• the exact default import configuration (paths is None with mode="extend").

The loaders, resolver, registry path handling, and relevant test fixtures now use these shared predicates.

Evidence: aimnet/models/artifact_validation.py:78-94, aimnet/calculators/resolve.py:107-121, aimnet/calculators/model_registry.py:231-243, and tests/test_model_artifact_security.py:43-80.

6. Private cross-package contracts and duplicated model assembly

Request: Multiple packages imported private registry symbols, while local/registry and HF loaders maintained separate v2 construction pipelines that had begun to drift.

Solution: Cross-package consumers now use the module-public REGISTRY_IMPORT_POLICY and load_registry_model contracts. Backward-compatible private aliases remain for existing internal consumers, but these internals are not promoted as new top-level aimnet.models APIs.

All v2 sources now share one assembly path that:

• constructs on CPU with file references disabled;
• enforces role-aware runtime import authorization;
• converts every AtomicShift instance to float64 before state loading;
• applies source-specific missing/unexpected-key policy;
• moves the completed model to the destination device once;
• attaches a copied metadata mapping consistently.

Direct custom sources may warn about unexpected keys, while registry and registry-backed HF sources reject them.

Evidence: aimnet/models/artifact_validation.py:70-75, aimnet/models/base.py:65-89, aimnet/models/base.py:165-214, aimnet/calculators/hf_hub.py:388-398, and tests/test_model_artifact_security.py:83-157.

7. Legacy TorchScript metadata

Request: Legacy .jpt metadata declared fully embedded Coulomb without has_embedded_lr=True, contradicting the new schema and the model's physical behavior.

Solution: The legacy loader now sets has_embedded_lr=True. A regression test loads a real ScriptModule through a mixed-case .jpt suffix and verifies version-1 metadata, long-range neighbor handling, no external modules by default, and rejection of an external Coulomb override that would double count the embedded interaction.

Legacy embedded Coulomb is also no longer misclassified as embedded dispersion: unknown legacy long-range Coulomb uses all pairs rather than inheriting the finite D3 cutoff.

Evidence: aimnet/models/base.py:92-113, aimnet/calculators/calculator.py:249-271, tests/test_calculator.py:2298-2336, and docs/model_format.md:224-232.

8. Stale bundled assets and release policy

Request: If a future package contains a stale bundled model, fail-closed acquisition prevents falling back to a valid network download. Bundled bytes and registry digests therefore need coordinated release handling.

Solution: Fail-closed behavior is retained and now has a synthetic regression test proving that a stale bundled artifact raises a checksum error without attempting a download.

Official wheels and source distributions currently do not bundle model artifacts, so adding a release verifier now would pass vacuously. The documentation instead makes future bundling a separate design and release decision requiring maintainer approval, immutable artifact identity, digest matching for actual wheel/source contents, and non-vacuous release checks.

Evidence: aimnet/calculators/model_registry.py:160-228, tests/test_model_registry.py:491-540, docs/train.md:164-168, and docs/model_format.md:326.

Additional safety fixes found during remediation

• HF configuration and weights are pinned to one immutable snapshot, preventing a mutable branch from mixing validated configuration from one commit with weights from another (aimnet/calculators/hf_hub.py:367-410, tests/test_hf_hub.py:674-716).
• Registry loading explicitly performs canonical validation and rejects unexpected state keys (aimnet/models/base.py:203-211, tests/test_model_registry.py:248-284).
• Schema-less raw nn.Module metadata remains compatible but still receives effective runtime checks that prevent Coulomb or dispersion double counting (aimnet/models/artifact_validation.py:400-430, tests/test_calculator.py:2135-2151).
• Atomic export preserves existing file modes and respects restrictive process umasks (aimnet/train/export_model.py:83-116, tests/test_train_utils.py:226-250).

[isayev]

All three compatibility regressions from the review are addressed the right way — YAML-derived SRCoulomb metadata before structural validation, the four-level validation split restoring the explicit-flag escape for direct artifacts while keeping registry/export fail-closed, and export-time validation with atomic replacement. The registry-fleet weekly workflow and the consolidated routing predicates close the two structural concerns as well. CI green across the matrix.