feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate - #5903
feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate#5903Shaurya2k06 wants to merge 2 commits into
Conversation
|
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:
📝 WalkthroughWalkthroughChangesDeepEval now resolves NeighborGraph auto selection
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant DeepEval
participant resolve_auto_graph_builder
participant NeighborGraph
DeepEval->>resolve_auto_graph_builder: Resolve auto using device and frame count
resolve_auto_graph_builder-->>DeepEval: Return nv, vesin, or dense
DeepEval->>NeighborGraph: Build graph with selected backend
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt_expt/model/test_graph_builder_dispatch.py (1)
146-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert dispatch, not only parity.
All backends are intentionally value-equivalent, so this passes if
None/"auto"incorrectly falls back todense. Mock or spy on the resolver/concrete builder and assert that the resolved backend is invoked.🤖 Prompt for 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. In `@source/tests/pt_expt/model/test_graph_builder_dispatch.py` around lines 146 - 158, Update test_none_and_auto_match_resolved_builder to spy on or mock the resolver/concrete graph builder, then assert that the backend resolved for None and "auto" is actually invoked. Retain the existing output-parity assertions, but ensure the test fails if either input silently falls back to dense.
🤖 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/pt_expt/utils/vesin_graph_builder.py`:
- Around line 11-15: Update the documentation describing the shared
resolve_auto_graph_builder default ladder to state that vesin is selected for
CPU or CUDA fallback only when vesin.torch is importable; otherwise document
that dense is selected.
---
Nitpick comments:
In `@source/tests/pt_expt/model/test_graph_builder_dispatch.py`:
- Around line 146-158: Update test_none_and_auto_match_resolved_builder to spy
on or mock the resolver/concrete graph builder, then assert that the backend
resolved for None and "auto" is actually invoked. Retain the existing
output-parity assertions, but ensure the test fails if either input silently
falls back to dense.
🪄 Autofix (Beta)
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: d02887c0-8595-4f4e-a31f-c9bc7c3749a5
📒 Files selected for processing (9)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/neighbor_graph_method.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/model/test_graph_builder_dispatch.pysource/tests/pt_expt/utils/test_neighbor_graph_method.py
987d303 to
409bb04
Compare
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5903 +/- ##
==========================================
- Coverage 79.60% 79.36% -0.25%
==========================================
Files 1085 1085
Lines 126405 126424 +19
Branches 4598 4598
==========================================
- Hits 100631 100335 -296
- Misses 24120 24437 +317
+ Partials 1654 1652 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
409bb04 to
451f50d
Compare
|
Rebased onto current master after #5912 / #5913 landed the training/eager auto path. What changed in this update
CodeQL empty- |
451f50d to
01253d9
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (3)
deepmd/pt_expt/utils/graph_builder.py:33
- Docstring names the CUDA dependency as
nvalchemiops, but the rest of the module (warnings/errors) refers to the pip packagenvalchemi-toolkit-ops. Using the installable package name here avoids confusion.
* CUDA: ``nv`` if ``nvalchemiops`` is importable, else ``vesin`` if
``vesin.torch`` is importable, else ``dense``.
deepmd/pt_expt/utils/graph_builder.py:56
- When
neighbor_graph_method='auto'falls back todenseon CUDA due to missing optional deps, the warning doesn’t tell users how to enable the faster backends (unlike the training-path warning below). Adding install hints makes the message actionable.
log.warning(
"nvalchemi-toolkit-ops and vesin[torch] are unavailable; falling "
"back from neighbor_graph_method='auto' to the dense graph builder."
)
deepmd/pt_expt/utils/graph_builder.py:31
- PR description/issue state that the model-level/training default should follow the same CPU/CUDA ladder as inference (CPU preferring vesin when available), but the code here explicitly documents (and
resolve_neighbor_graph_methodenforces) a separate training policy that keeps CPU ondense. This means the implementation doesn’t match the stated acceptance criteria unless the PR description/issue closure is adjusted.
This issue also appears in the following locations of the same file:
- line 32
- line 53
Single owner of the inference / DeepEval auto ladder (training uses
:func:`resolve_neighbor_graph_method`, which keeps CPU on ``dense`` because
vesin loops frames in Python and is not safe as a multi-frame training
default):
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt_expt/utils/graph_builder.py (1)
58-60: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider logging the CPU dense fallback for parity with the CUDA branch.
The CUDA branch logs a warning when it falls back to
dense(Lines 53-56). The CPU branch falls back todensesilently at Line 60. Dense is the O(N²) carry-all builder; silently downgrading to it on CPU (e.g.vesin[torch]not installed) can cause an unexplained performance regression on large systems, with no diagnostic for the user to act on.Add a similar
log.warning(orlog.info) call before returning"dense"on the CPU path, mentioning how to installvesin[torch]. If you make this change, updatetest_auto_resolution's("cpu", False, False, "dense", False)case insource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.pytowarns=True.♻️ Proposed fix
if is_vesin_torch_available(): return "vesin" + log.warning( + "vesin[torch] is unavailable; falling back from " + "neighbor_graph_method='auto' to the dense graph builder on CPU. " + "Install it with `pip install vesin[torch]` to enable the O(N) " + "vesin graph builder." + ) return "dense"🤖 Prompt for 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. In `@deepmd/pt_expt/utils/graph_builder.py` around lines 58 - 60, Update the CPU fallback in the graph-builder backend resolution function to log a warning or info message before returning "dense", explicitly mentioning installation of vesin[torch]. Also update the test_auto_resolution case for ("cpu", False, False, "dense") so it expects a warning.
🤖 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.
Nitpick comments:
In `@deepmd/pt_expt/utils/graph_builder.py`:
- Around line 58-60: Update the CPU fallback in the graph-builder backend
resolution function to log a warning or info message before returning "dense",
explicitly mentioning installation of vesin[torch]. Also update the
test_auto_resolution case for ("cpu", False, False, "dense") so it expects a
warning.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 46027a52-d036-4bf6-9812-cd1fa7c5ac10
📒 Files selected for processing (4)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pysource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt_expt/utils/vesin_graph_builder.py
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (2)
deepmd/pt_expt/utils/graph_builder.py:31
- The new helper documents/supports
neighbor_graph_method="auto", but the core builder dispatch (build_neighbor_graph_for_method) still only accepts concrete methods. Outside ofDeepEval._resolve_neighbor_graph_method, passingneighbor_graph_method="auto"into a pt_expt model graph path would still raise aValueErrorfrom the builder dispatcher. Either wire auto-resolution into the model/dispatcher, or clarify here that callers must resolve "auto" before dispatching.
"""Resolve ``neighbor_graph_method="auto"`` to a concrete inference builder.
Single owner of the inference / DeepEval auto ladder (training uses
:func:`resolve_neighbor_graph_method`, which keeps CPU on ``dense`` because
vesin loops frames in Python and is not safe as a multi-frame training
deepmd/pt_expt/utils/vesin_graph_builder.py:16
- This module docstring reads as if
neighbor_graph_method="auto"is a general pt_expt model option, but currently the only in-tree resolver for "auto" is DeepEval (and the graph builder dispatcher itself rejects "auto"). Consider clarifying that "auto" here refers to DeepEval/inference resolution so users don’t try passing "auto" directly into modelneighbor_graph_methodand hit a runtimeValueError.
for ``nf == 1`` inference and CPU use. Inference ``neighbor_graph_method="auto"``
(:func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`) selects
vesin only when ``vesin.torch`` is importable (CPU always; CUDA only when ``nv``
is unavailable); otherwise it falls back to ``dense``. Training auto keeps CPU
on ``dense`` and never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`)
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
f54e820 to
be6d9bb
Compare
Shaurya2k06
changed the title
|
Addressed the review on the current tip (
|
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/pt_expt/utils/graph_builder.py`:
- Around line 68-69: Update the graph-builder resolution flow around the nf
parameter to first reject values that are not integers, explicitly excluding
bool, with ValueError before the existing nf < 1 check. Preserve the current
lower-bound validation for integer frame counts, and add tests covering nf=1.5
and nf=None (as well as bool if consistent with the requirement).
🪄 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: adb029ba-9ffb-4efd-944f-a6e54159a245
📒 Files selected for processing (5)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pysource/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.pysource/tests/pt_expt/model/test_graph_builder_dispatch.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/pt_expt/utils/vesin_graph_builder.py
- source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py
be6d9bb to
abe2e03
Compare
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
All four points are addressed, and this is a real rework rather than a wording change. Checked against abe2e033a.
The vesin rung is now what I was asking for. _resolve_neighbor_graph_method takes nf and returns "auto" unresolved when it is omitted, so construction defers and _build_eval_graph resolves from the batch frame count -- which is the structural change, not just a new parameter. I went looking for the ways that could go wrong and they are all covered: coord_input at both call sites is coords.reshape(nframes, natoms, 3), so np.asarray is safe and shape[0] really is the frame count; "auto" leaks nowhere, since the only reader of self._neighbor_graph_method is the dispatch in _build_eval_graph; and the construction-time raise for a non-default value on a non-graph artifact is intact. Citing _select_neighbor_builder in the docstring is the right call, since that is the policy this now has to stay in step with.
The rewritten summary is accurate, including the explicit note that the model-level and training flips shipped in #5912 / #5913 -- that was the part most likely to mislead a bisect later. The numpydoc sections match the sibling function. And _eval parametrized over nf in {1, 4}, with atype and box expanded to match, closes the parity gap in the batched regime rather than only asserting the resolver's return value.
One thing the move to call-time resolution brings with it, inline.
Worth noting the tests had not run when I read this -- every substantive workflow was sitting at action_required. I have approved them, so Test Python / Test C++ / Test CUDA are going now. Everything above is from reading the head, so the new ladder cases and the nf=4 parity case still need a green run to count as verified.
Extract resolve_auto_graph_builder for inference and select vesin only for single-frame batches, matching _select_neighbor_builder. Multi-frame auto stays on nv/dense so auto_batch_size does not hit the per-frame loop. Warn once per process when CUDA falls back without nvalchemiops. Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>
abe2e03 to
1a8934b
Compare
for more information, see https://pre-commit.ci
| return "vesin" | ||
| if dev.type == "cuda" and not nv_available: | ||
| if not _warned_auto_no_nv: | ||
| _warned_auto_no_nv = True |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
All five points are resolved, checked against 2d386b63d.
The warn-once fix is right: a module-level _warned_auto_no_nv with the global declaration, and the is_nv_available() probe hoisted so a resolve costs one call instead of two. test_resolve_auto_graph_builder_warns_once calls the resolver twice on ("cuda", nf=4) and asserts exactly one message, which fails on the previous head where both calls warned -- so it is a real regression test, not a restatement of the code.
I also closed out the caveat from my last review. vesin[torch] is an unconditional entry in the torch extra, so the skipif not is_vesin_torch_available() guard does not trip in CI, which means test_vesin_matches_dense_energy_force[4] genuinely executed. That was the whole point of the parity gap, and the suite is green across all 54 checks this time rather than sitting at action_required, so the ladder and parity coverage is now verified by a run rather than by reading.
The end state is the one I was arguing for on #5912: vesin is reachable automatically, but only where its per-frame Python loop is not a liability, and the policy lives in one function that cites _select_neighbor_builder as the thing it has to stay in step with. Thanks for working through the iterations on this one.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Three independent full reviews were completed against this exact head. No concrete actionable findings were identified, so no artificial inline comment was added. The call-time frame-count resolution, NV precedence, single-frame Vesin gate, explicit-method behavior, graph semantics, and test coverage are consistent, and the CI matrix is green.
The Codex quota is close to resetting, so I am concentrating the remaining token budget on these reviews.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary
resolve_auto_graph_builder(device, nf)for inference / DeepEval only (training keepsresolve_neighbor_graph_methodfrom perf(pt_expt): use scalable graph builders during training #5913).neighbor_graph_method="auto"at eval call time with the batch frame count: CUDA prefersnv;vesinonly whennf == 1and importable; otherwisedense. Matches_select_neighbor_builder.auto_batch_size/dp testbatches therefore stay off vesin's per-frame Python loop.nf in {1, 4}.This is not a model-level / training default flip — those already shipped in #5912 / #5913. The remaining change is the inference auto ladder: re-introducing vesin only under the
nf == 1gate that review asked for on #5912.Validation
ruff check/ruff formaton touched filespytestresolver ladder + DeepEval resolution + vesin parity fornf=1andnf=4Summary by CodeRabbit
New Features
Bug Fixes
Tests