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 1 commit 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.59% 79.35% -0.25%
==========================================
Files 1081 1081
Lines 126244 126259 +15
Branches 4592 4598 +6
==========================================
- Hits 100490 100190 -300
- Misses 24101 24418 +317
+ Partials 1653 1651 -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`)
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:56
- The warning emitted when falling back from neighbor_graph_method='auto' on CUDA is no longer actionable: it names missing optional deps but doesn't tell users how to install/enable them. Elsewhere in this module the training fallback warning and ImportErrors include a
pip install ...hint, so this should too.
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:30
- The PR description/linked issue state that the model-level default (and compiled training eager _forward_graph) should use this shared auto-selection ladder. In the current code, resolve_auto_graph_builder is only called from pt_expt DeepEval (and tests), and the model path still uses its existing default-flip logic in pt_expt/model/make_model.py. Either extend the call sites as described, or adjust the PR description to match the actual scope.
def resolve_auto_graph_builder(
device: torch.device | str,
) -> str:
"""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
default):
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Thanks for picking this up. Before the objections: the part I asked for on #5912 and did not get there, you have delivered here. The five ladder branches are driven deterministically by patching is_nv_available / is_vesin_torch_available / the device rather than recomputing the cascade inside the assertion, so test_resolve_auto_graph_builder_ladder genuinely fails without the change instead of comparing the code to a copy of itself. Extracting the ladder into one function so DeepEval and training stop carrying private copies is also the right direction.
My problem is with what the remaining commit actually changes, and with the description of it. Four comments inline.
The short version: the default flip this PR is named for already shipped -- #5912 landed DeepEval's auto ladder and #5913 landed training's, and both are ancestors of this branch. What is left is one substantive commit, and its real content is that vesin becomes an automatically selected builder. That is precisely the rung that was removed from #5912 in response to review before it merged, and the reasoning that removed it has not changed.
To be clear about what does not conflict: #5913 owns resolve_neighbor_graph_method for training and this adds a separate inference-only function, so the two are complementary and the branch merges cleanly. I am not asking you to rebase or coordinate with it.
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:35
- PR description/issue state that the pt_expt model-level default (None/"auto") and compiled training eager graph path should use the same auto ladder (nv→vesin→dense on CUDA; vesin→dense on CPU). In this PR, the new ladder helper is only used by DeepEval; training still resolves via resolve_neighbor_graph_method and compiled training builds with getattr(_model, "neighbor_graph_method", "dense"), so CPU can never auto-select vesin and the default remains effectively dense unless set elsewhere.
Evidence: deepmd/pt_expt/train/training.py uses resolve_neighbor_graph_method() (dense/nv only) and _forward_graph uses getattr(_model, "neighbor_graph_method", "dense"); deepmd/pt_expt/model/make_model.py::_resolve_graph_method falls back to "dense" when the attribute is absent.
def resolve_auto_graph_builder(
device: torch.device | str,
) -> str:
"""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
default):
* CUDA: ``nv`` if ``nvalchemiops`` is importable, else ``vesin`` if
``vesin.torch`` is importable, else ``dense``.
* CPU: ``vesin`` if ``vesin.torch`` is importable, else ``dense``.
deepmd/pt_expt/utils/graph_builder.py:56
- The CUDA fallback warning for neighbor_graph_method='auto' isn’t very actionable: it mentions both backends are unavailable but doesn’t tell the user how to enable either one (unlike the training-path warning below, which includes an install hint). Consider adding explicit install guidance for both optional backends here so DeepEval users know how to get an O(N) builder.
log.warning(
"nvalchemi-toolkit-ops and vesin[torch] are unavailable; falling "
"back from neighbor_graph_method='auto' to the dense graph builder."
)
e5eebef to
f54e820
Compare
|
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
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. Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>
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.
| if nf == 1 and is_vesin_torch_available(): | ||
| return "vesin" | ||
| if dev.type == "cuda" and not is_nv_available(): | ||
| log.warning( |
There was a problem hiding this comment.
This warning now fires once per batch instead of once per DeepEval.
Moving resolution to call time is right, but it also moved this log line onto the per-batch path. _build_eval_graph is reached from _eval_model_graph, which _eval_func hands to self.auto_batch_size.execute_all(...), so it runs once per batch rather than once at construction. On a CUDA host without nvalchemiops evaluating multi-frame batches -- exactly the configuration this message is written for -- a dp test or dp model-devi run will now repeat it for every batch. Previously _setup_neighbor_backend resolved once and the user saw it a single time.
The message is also the more useful of the two things happening here, so burying it in repetition is a real loss: the point is "install nvalchemi-toolkit-ops", which is a one-time action, not a per-batch one.
A module-level warn-once flag is the smallest fix:
_warned_no_nv = False
...
if dev.type == "cuda" and not nv:
global _warned_no_nv
if not _warned_no_nv:
_warned_no_nv = True
log.warning(...)Caching the whole resolver per (device.type, nf, availability) would work too and would also pick up the second point: is_nv_available() is called twice on this path, once at line 80 and again in this condition. Hoisting it to a local reads better and halves the import probes -- which matter slightly more than they look, because on CPU-only hosts is_nv_available wraps its import in _suppress_native_stderr, so it is doing dup/dup2/close syscalls on every call now rather than once per model.
None of this affects results -- all builders emit the same carry-all neighbor set -- so it is not a correctness concern, just something that belongs with this change rather than after it.
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