refactor: promote duck-typed capability probes to base-class interfaces - #5967
Conversation
…chg-spin hasattr probes (deepmodeling#5897)
…ling#5897) merge_env_stat's base_obj/link_obj can be a bare Descriptor (se_e2_a, se_r, se_t) as well as a DescriptorBlock, so the same concrete defaults are also declared on make_base_descriptor's BD base - otherwise direct attribute access crashes for descriptors that never set set_stddev_constant themselves.
…ling#5897) Declare type_embd_data/tebd_compress in the tebd-family descriptor __init__s (DescrptDPA1, DescrptBlockSeAtten, DescrptSeTTebd) so their presence is a class property rather than a runtime accident. Convert the three self-probes (hasattr(self, "type_embd_data")) to the equivalent self.type_embd_data is not None check. Fix a regression this uncovers in pt_expt: DescrptDPA1/DescrptSeTTebd compression paths call torch.nn.Module.register_buffer(self, "type_embd_data", ...) directly, bypassing dpmodel_setattr's existing None-slot-to-buffer promotion. Add register_buffer_replacing_slot() in pt_expt/common.py mirroring that existing idiom, and use it at both call sites. Add a negative-assertion test pinning that a non-tebd descriptor (DescrptSeA) never carries either attribute, since the jax restore walker uses hasattr(obj, "tebd_compress") as a family-membership test.
deepmd/pt_expt/descriptor/dpa2.py:_store_type_embd_data still called torch.nn.Module.register_buffer(self, "type_embd_data", ...) directly, the same raw pattern fixed for dpa1/se_t_tebd in the previous commit. DescrptDPA2 (dpmodel) doesn't declare the type_embd_data slot yet, so this was purely defensive, but it removes the landmine for whichever future change declares it there. Swap to the existing register_buffer_replacing_slot() helper in pt_expt/common.py.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (5)
📝 WalkthroughWalkthroughThe PR replaces optional capability probes with explicit interfaces, removes ChangesCapability contract standardization
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/fitting/make_base_fitting.py`:
- Around line 70-93: Update the reinit_exclude method signature to use None
instead of a mutable list default, while preserving the existing truthiness
check so no collection initialization is needed.
In `@source/tests/common/dpmodel/test_zbl_bridging.py`:
- Around line 399-401: Update the assertion comparing get_default_chg_spin() in
the bridged and plain models to compare the returned values directly, rather
than only comparing whether they are non-None, so the test validates the
forwarded default value.
In `@source/tests/universal/common/cases/atomic_model/utils.py`:
- Around line 105-108: Update the assertions around get_pair_exclude_types and
pair_excl to verify the configured excluded-pair values, not only whether both
representations are empty or nonempty. When self.module.pair_excl is present,
compare its mask configuration against pet; retain the existing empty-state
invariant for the None case.
In `@source/tests/universal/common/cases/fitting/utils.py`:
- Around line 111-116: Update the override branch in the fitting test around
reinit_exclude so the final reinit_exclude([]) call asserts that
self.module.exclude_types is empty, verifying the reset clears the prior [0]
exclusion.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0fe30b62-e608-491e-aebf-98408f25a5ad
📒 Files selected for processing (54)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/atomic_model/linear_atomic_model.pydeepmd/dpmodel/descriptor/descriptor.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/dpa2.pydeepmd/dpmodel/descriptor/dpa3.pydeepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/hybrid.pydeepmd/dpmodel/descriptor/make_base_descriptor.pydeepmd/dpmodel/descriptor/se_atten_v2.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/fitting/make_base_fitting.pydeepmd/dpmodel/model/base_model.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/dpmodel/utils/env_mat_stat.pydeepmd/infer/deep_eval.pydeepmd/jax/infer/deep_eval.pydeepmd/jax/jax2tf/serialization.pydeepmd/jax/jax_md/__init__.pydeepmd/jax/train/trainer.pydeepmd/jax/utils/serialization.pydeepmd/kernels/cuda/dpa1/canonical.pydeepmd/pd/infer/deep_eval.pydeepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/descriptor.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pt_expt/common.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/repflows.pydeepmd/pt_expt/descriptor/se_t_tebd.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/network.pydeepmd/pt_expt/utils/serialization.pydeepmd/tf2/model/dp_model.pydeepmd/tf2/train/trainer.pydeepmd/tf2/utils/serialization.pysource/tests/common/dpmodel/test_descriptor_block_defaults.pysource/tests/common/dpmodel/test_descriptor_dpa1.pysource/tests/common/dpmodel/test_make_base_fitting.pysource/tests/common/dpmodel/test_model_compression.pysource/tests/common/dpmodel/test_pair_exclude_contract.pysource/tests/common/dpmodel/test_zbl_bridging.pysource/tests/infer/gen_dpa4_spin_chgspin.pysource/tests/pt_expt/model/test_dpa4_native_spin.pysource/tests/universal/common/cases/atomic_model/utils.pysource/tests/universal/common/cases/descriptor/utils.pysource/tests/universal/common/cases/fitting/utils.pysource/tests/universal/common/cases/model/utils.pysource/tests/universal/dpmodel/descriptor/test_descriptor.py
💤 Files with no reviewable changes (3)
- deepmd/dpmodel/descriptor/dpa3.py
- deepmd/dpmodel/descriptor/dpa4.py
- deepmd/pd/model/descriptor/dpa2.py
get_pair_exclude_types() was declared only on dpmodel's concrete BaseAtomicModel, not on the shared make_base_atomic_model() factory that pt's atomic models also derive from (via make_base_atomic_model(torch.Tensor)). The universal test_pair_exclude_contract case runs for pt too, so it hit AttributeError there. pt atomic models already set self.pair_exclude_types in __init__ via reinit_pair_exclude (deepmd/pt/model/atomic_model/base_atomic_model.py), so the accessor is safe to declare once, concretely, on the BAM base.
The shared universal descriptor/model cases asserted 'not hasattr(type(self.module), "has_default_chg_spin")', which fails for the frozen pt backend: pt still declares the (now-redundant) method on 8 descriptors + make_model, which is allowed since pt is frozen. Move the negative assertion out of the shared case methods (which run for every backend) into dpmodel-only test files, where the merge is authoritative: TestHasDefaultChgSpinAbsentDP in source/tests/universal/dpmodel/descriptor/test_descriptor.py and source/tests/universal/dpmodel/model/test_model.py, asserting absence on both the BD/base_model bases and the concrete dpmodel classes.
hasattr(model, "get_var_name") is always true now that make_base_model declares the method with a concrete None default, so the tf2 SavedModel export unconditionally exported a get_var_name/get_task_dim/get_intensive tf.function trio for every model, including non-property models (where get_var_name() returns None and get_task_dim() would raise). Mirror the jax2tf pattern (deepmd/jax/jax2tf/serialization.py): gate on 'model.get_var_name() is not None' instead.
…erge pd retained a divergent has_default_chg_spin chain (base_atomic_model, dp_atomic_model, make_model, and the se_a/dpa1/dpa3/se_t_tebd descriptors) after its DPA2 override was already dropped elsewhere, and after the rest of the codebase merged has_default_chg_spin into get_default_chg_spin (predicate: get_default_chg_spin() is not None). Delete the has_default_chg_spin definitions and rewrite DPAtomicModel.get_default_chg_spin to the dpmodel pattern (gate purely on add_chg_spin_ebd, not an extra has_default_chg_spin probe on the descriptor). Convert the one caller (pd/train/training.py) to the 'get_default_chg_spin() is not None' predicate, matching pt_expt's get_additional_data_requirement. Paddle is not installed locally; every edited file was verified with 'python -m py_compile' and a repo-wide grep confirms zero remaining has_default_chg_spin references under deepmd/pd/.
…p_eval - SpinModel.forward_common_lower's else branch (backbone_model without a nested .atomic_model) used getattr(self.backbone_model, "pair_excl", None). Per the __init__ annotation, backbone_model in that branch IS a DPAtomicModel, whose BaseAtomicModel.__init__ unconditionally sets self.pair_excl via reinit_pair_exclude -- so this is a guaranteed attribute, not one to defensively probe. - pt_expt's model_type property called self._dpmodel.get_var_name() twice in the same elif condition; hoist to a single local variable alongside model_output_type, set to None in the metadata-only branch.
…e descriptor Extends test_descriptor_block_defaults.py (Task 6's rationale: BD base in make_base_descriptor was extended with set_davg_zero/set_stddev_constant class defaults after ratifying that merge_env_stat's real contract is Union[Descriptor, DescriptorBlock], not "blocks only"): - test_base_descriptor_stat_flags_have_class_defaults: pins the BD base class defaults directly (mirrors the existing DescriptorBlock test). - test_merge_env_stat_on_bare_descriptor_no_attribute_error: constructs a bare DescrptSeA (not a DescriptorBlock), runs compute_input_stats, and calls merge_env_stat on it -- proving no AttributeError when reading the stat-behavior flags on a Descriptor that never sets them itself.
for more information, see https://pre-commit.ci
…ility contracts The frozen pt backend returns torch.Tensor from get_default_chg_spin while dpmodel returns a list; the shared universal assertion must be backend-agnostic. Asserting len(dcs) == get_dim_chg_spin() is also the stronger contract.
DummyModel (jax2tf serialization test) gains get_var_name -> None; the _DescriptorWithStats stub gains the stat-flag class defaults that merge_env_stat now reads directly.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5967 +/- ##
==========================================
- Coverage 79.60% 79.39% -0.22%
==========================================
Files 1085 1085
Lines 126405 126579 +174
Branches 4598 4598
==========================================
- Hits 100631 100503 -128
- Misses 24120 24423 +303
+ Partials 1654 1653 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed independently from correctness, compatibility/testing, and maintainability perspectives. One actionable issue was found; the other two independent reviews found no additional issues.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
- TFModelWrapper: normalize the exported empty default-chg-spin tensor to None at the artifact-decode boundary (its one owner), so the live-model invariant `get_default_chg_spin() is None == no default` holds for SavedModel artifacts too and the jax DeepEval predicate stays a plain `is not None`. Pure decode helper + two-branch test. - test_zbl_bridging: compare the forwarded default chg-spin VALUES, not just their None-ness. - universal atomic-model case: when pair_excl exists, its mask must hold exactly the symmetric closure of get_pair_exclude_types() (both backends symmetrize on reinit; compared as sets of tuples because pt stores a set and dpmodel a list). The closure form is pinned by test_pair_exclude_contract on a (0,1)-only config. - universal fitting case: assert reinit_exclude([]) actually clears exclude_types on the override branch.
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed the current head 7ac154b with three independent passes covering correctness, multi-backend and serialization compatibility, and API and test contracts, followed by consolidated local verification and deduplication against existing review threads.
No new high-confidence actionable findings were found on the current diff, so there are no inline comments to attach.
Validation notes: targeted current-head tests for the new fitting, pair-exclusion, compression, and JAX SavedModel charge-spin decode contracts passed; Ruff passed on the changed backend trees. The remaining Paddle C++ CI failure is an external partial-download failure and is not attributable to this PR.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
njzjz-bot
left a comment
There was a problem hiding this comment.
Approved after three independent review passes and consolidated verification of the current head 7ac154b. No new high-confidence actionable findings were identified.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
adbd6bc
Summary
Promotes duck-typed capability probes (
hasattr/getattr(obj, "x", default)) inventoried in #5897 to concrete-default methods/attributes declared once on the owning base class, so a typo'd or renamed name raises instead of silently degrading to the default. Eight change groups, one commit each (plus one fix-up commit):has_default_chg_spinintoget_default_chg_spin(derive viais not None); drop the charge-spinhasattrprobe family and declare concrete defaults on the owning bases.get_var_name/get_task_dim/get_intensiveonmake_base_modelwith concrete defaults; convert the jax2tf property probe and refresh a stale docstring.get_geo_compress()base accessor and drop consumergetattr/hasattrprobes; pin both the override and default branches with a dedicated compression test.reinit_excludeno-op default onmake_base_fitting, killing thedp_atomic_model.pyprobe.get_pair_exclude_types()accessor onBaseAtomicModeland pinpair_exclas a direct-access contract (verified against every construction path, including deserialize andSpinModel.forward_common_lower).set_davg_zero/set_stddev_constantas class-default stat flags on the descriptor block bases (extended tomake_base_descriptor's BD base after ratifying thatmerge_env_stat's real contract isUnion[Descriptor, DescriptorBlock], not "blocks only").has_message_passing/has_default_fparamprobes now that both are already declared on the base — call directly.tebd_compress/type_embd_datacompression slots in the tebd family__init__s (dpa1 / se_atten_v2 / se_t_tebd + blocks), including a follow-up fix that defuses aregister_buffertrap for the DPA2 tebd compression slot.Every promotion ships a dual-branch (default + override) universal-suite assertion;
pair_exclships a construction-path assertion instead.Closes #5897
Known limitations
py_compileonly — no local Paddle install to run its test suite.get_geo_compress()and graph.pt2export/AOTI paths are untested on this CPU-only box; GPU validation before merge is advisable.DeepEvalcharge-spin unit test (the jaxinfer/deep_eval.pycharge-spin path is a one-line delegation, exercised only indirectly).SpinModel.forward_common_lower(pair_exclude_typesset on a native spin model) is untested — no existing test sets it on a spin model.DescrptDPA2.get_geo_compress()'sTruebranch is unpinned: no DPA2 compression test exists at all (pre-existing gap, not introduced here).atom_exclgetattrprobes (e.g.deepmd/kernels/cuda/dpa1/canonical.py:44) are the same duck-typing pattern but were never part of Promote remaining duck-typed attribute probes (add_chg_spin_ebd, geo_compress, reinit_exclude) to base-class interfaces #5897's inventory — candidate follow-up issue, not addressed here..pt2graph-freeze tests (test_dpa4_export.py,test_dpa4_zbl_parallel.py,test_graph_export_with_comm.py,test_zbl_bridging.py) fail with a torch inductor CPU codegenAssertionError(atomic_addstore on a non-vectorized index) — confirmed pre-existing on a cleanupstream/masterworktree, unrelated to this branch's changes.add_chg_spin_ebdnow derives fromget_dim_chg_spin() > 0(rather than the previous getattr-probe), so a hybrid descriptor with a chg-spin-capable child now actually conditions on it, where it previously never did. This is a latent-bug fix surfaced by the promotion, not an intentional feature change — it has no dedicated end-to-end test.get_var_name()/get_intensive()now returnNone/Falsefor live non-property models instead of raisingNotImplementedError, matching the concrete-default contract onmake_base_model. Callers that relied on the exception (if any exist outside this repo) would need to switch to checking the return value.merge_env_statmulti-task path (multiple descriptors sharing stats across sub-models) has no dedicated unit test; only the single-task path is covered.has_default_chg_spin→get_default_chg_spininterface merge (completing the promotion started elsewhere in the branch) was verified bypython -m py_compileonly — no Paddle install on this box to run pd's test suite.Summary by CodeRabbit