Trust the embedded-dispersion modules the CPCM artifact needs - #115
Merged
Conversation
The artifact trust boundary shipped with a default allowlist missing two
first-party classes: aimnet.modules.D3TS and aimnet.modules.lr.DispParam,
both nn.Module subclasses defined in aimnet/modules/lr.py. The shipped
CPCM(water) solvation artifact references both, so it stopped loading with
ValueError: Invalid v2 artifact field 'model_yaml':
Untrusted import path for 'class': 'aimnet.modules.lr.DispParam'
Nothing loads the solvation model until it is actually needed, so downstream
this did not surface as a startup error. It surfaced as a multi-hour
production run aborting partway with every GPU worker raising on that model,
having already completed all of its gas-phase work.
Both classes are enumerated from the artifact's own model_yaml rather than
from the error message: fixing only DispParam, the class the exception names,
leaves D3TS to fail on the next load.
Tests assert membership and that each entry resolves to a real nn.Module --
an allowlist entry that does not resolve would trade a load failure for an
import error.
test_allowed_model_import_paths_are_shared_and_immutable asserts the EXACT contents of ALLOWED_MODEL_IMPORT_PATHS, so adding the two embedded-dispersion entries broke it. Update the pin. The pre-existing state of this file is itself the argument for the change: _FROZEN_CLASS_PATHS already listed aimnet.modules.D3TS, under a comment saying those are paths released checkpoints reference. So the serialization ABI knew D3TS was load-bearing while the runtime trust allowlist did not, and that inconsistency is exactly what stopped the shipped CPCM(water) artifact loading. aimnet.modules.lr.DispParam joins _FROZEN_CLASS_PATHS for the same reason D3TS is already there: a released artifact references it, so it is ABI, and test_frozen_class_path_resolves now covers it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The artifact trust boundary shipped with a default allowlist missing two first-party classes:
aimnet.modules.D3TSaimnet.modules.lr.DispParamBoth are
nn.Modulesubclasses defined inaimnet/modules/lr.py. The shipped CPCM(water) solvation artifact references both, so it stopped loading:Why it was not caught sooner
Nothing loads the solvation model until it is actually needed, so this does not surface at import or at startup. It surfaced as a multi-hour production run aborting partway, with every GPU worker raising on that model after having already completed all of its gas-phase work.
Both classes, not just the one in the message
The two entries were enumerated from the artifact's own
model_yaml, not from the exception text. Fixing onlyDispParam— the class the error names — leavesD3TSto fail on the very next load.Verification
With this change the CPCM model loads and computes: a water optimization under
optimize_and_thermoconverges to a finite energy on an L40S. Without it, the same call fails at load.New tests assert membership in the allowlist and that each entry resolves to a real
nn.Module, since an entry that does not resolve would trade a load failure for an import error.Full suite: 160 passed.