feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate

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>

Author: Shaurya2k06

deepmd/pt_expt/infer/deep_eval.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import json
-import logging
 import warnings
 from collections.abc import (
     Callable,
@@ -80,8 +79,6 @@
         NeighborGraph,
     )
 
-log = logging.getLogger(__name__)
-
 
 # Public output keys emitted by graph-lower forwards, keyed by the
 # output-variable category that ``request_defs`` carries. The graph path is
@@ -192,15 +189,17 @@ class DeepEval(DeepEvalBackend):
     neighbor_graph_method : str, default: "auto"
         Carry-all graph builder for graph-form ``.pt2`` artifacts and
         graph-routed ``.pt`` checkpoints
-        (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects
-        ``"nv"`` on CUDA when nvalchemiops is available and otherwise falls
-        back to ``"dense"``. ``"vesin"`` remains explicit opt-in because it
-        loops over frames in Python. Explicit
+        (``metadata["lower_input_kind"] == "graph"``): ``"auto"`` selects via
+        :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`
+        at each eval call (CUDA: ``nv`` if importable; else ``vesin`` only when
+        ``nf == 1`` and importable; else ``dense``). Explicit
         ``"dense"`` / ``"ase"`` / ``"vesin"`` / ``"nv"`` choices are preserved.
         A non-default value on any other artifact raises at construction because
         the knob would silently do nothing there; use ``nlist_backend`` for the
         nlist path instead. All builders emit the same neighbor set, so the
-        choice is performance-only. Consolidating the two knobs into a single
+        choice is performance-only. Training keeps a separate auto policy
+        (:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`)
+        that never selects ``vesin``. Consolidating the two knobs into a single
         backend-selection API is deferred to the dense-nlist deprecation.
     **kwargs : dict
         Keyword arguments.
@@ -271,33 +270,31 @@ def __init__(
             raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
 
     @staticmethod
-    def _resolve_neighbor_graph_method(method: str) -> str:
-        """Resolve the graph builder once for the active device."""
+    def _resolve_neighbor_graph_method(method: str, nf: int | None = None) -> str:
+        """Validate and optionally resolve the graph builder for the active device.
+
+        ``"auto"`` is left unresolved when ``nf`` is omitted so construction-
+        time setup can defer to :meth:`_build_eval_graph`, where the frame
+        count is known and vesin can be gated on ``nf == 1``.
+        """
         if method not in ("auto", "dense", "ase", "vesin", "nv"):
             raise ValueError(
                 f"Unknown neighbor_graph_method {method!r}; "
                 "expected 'auto', 'dense', 'ase', 'vesin', or 'nv'."
             )
         if method != "auto":
             return method
+        if nf is None:
+            return "auto"
 
-        from deepmd.pt.utils.nv_nlist import (
-            is_nv_available,
-        )
         from deepmd.pt_expt.utils.env import (
             DEVICE,
         )
+        from deepmd.pt_expt.utils.graph_builder import (
+            resolve_auto_graph_builder,
+        )
 
-        if DEVICE.type == "cuda":
-            if is_nv_available():
-                return "nv"
-            log.warning(
-                "nvalchemi-toolkit-ops is unavailable; falling back from "
-                "neighbor_graph_method='auto' to the dense graph builder. "
-                "Install it with `pip install nvalchemi-toolkit-ops` to enable "
-                "the NV graph builder."
-            )
-        return "dense"
+        return resolve_auto_graph_builder(DEVICE, nf)
 
     def _setup_neighbor_backend(self, nlist_backend: str) -> None:
         """Resolve the graph or neighbor-list construction strategy.
@@ -2316,14 +2313,21 @@ def _build_eval_graph(
     ) -> "NeighborGraph":
         """Build the carry-all NeighborGraph for graph-lower inference.
 
-        Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run
-        backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)).
-        All backends emit the SAME neighbor set (carry-all, sel-free), so the
-        selection is a pure performance choice and results are unchanged. The
-        result is canonicalized to the destination-major graph-form ``.pt2``
-        ABI after construction.
+        Dispatches on ``self._neighbor_graph_method``: ``auto`` is resolved
+        call-time via
+        :func:`~deepmd.pt_expt.utils.graph_builder.resolve_auto_graph_builder`
+        using the batch frame count (vesin only when ``nf == 1``);
+        ``dense``/``ase`` run backend-agnostic (numpy); ``vesin``/``nv`` run
+        on-device (torch, O(N)). All backends emit the SAME neighbor set
+        (carry-all, sel-free), so the selection is a pure performance choice
+        and results are unchanged. The result is canonicalized to the
+        destination-major graph-form ``.pt2`` ABI after construction.
         """
         method = self._neighbor_graph_method
+        if method == "auto":
+            coord_arr = np.asarray(coord_input)
+            nf = int(coord_arr.shape[0]) if coord_arr.ndim >= 2 else 1
+            method = self._resolve_neighbor_graph_method("auto", nf=nf)
         # Model-level ``pair_exclude_types`` is a graph-BUILD transform
         # (decision #18): apply it here so the exported ``.pt2`` lower consumes a
         # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++
@@ -2392,7 +2396,7 @@ def _build_eval_graph(
             )
         raise ValueError(
             f"unknown neighbor_graph_method {method!r}; "
-            "use 'dense', 'ase', 'vesin', or 'nv'"
+            "use 'auto', 'dense', 'ase', 'vesin', or 'nv'"
         )
 
     def _model_pair_excl(self) -> "PairExcludeMask | None":

deepmd/pt_expt/utils/graph_builder.py

@@ -19,6 +19,75 @@
 log = logging.getLogger(__name__)
 
 
+def resolve_auto_graph_builder(
+    device: torch.device | str,
+    nf: int = 1,
+) -> 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 never selects ``vesin``.
+
+    Mirrors :func:`deepmd.pt.model.model.sezm_model._select_neighbor_builder`:
+    ``vesin`` is eligible only for a single-frame batch (``nf == 1``), because
+    its API loops frames in Python (~1 ms/frame). Multi-frame batches stay on
+    ``nv`` (CUDA) or ``dense`` so ``auto_batch_size`` / ``dp test`` do not
+    regress to the per-frame loop.
+
+    Policy
+    ------
+    * CUDA + ``nvalchemiops``: ``nv`` (any ``nf``).
+    * ``nf == 1`` + ``vesin.torch``: ``vesin``.
+    * otherwise: ``dense``.
+
+    ``ase`` is never chosen automatically. All builders emit the same carry-all
+    neighbor set; the choice is performance-only. Builders run eagerly outside
+    traced / compiled regions, so this does not change ``.pt2`` artifacts.
+
+    Parameters
+    ----------
+    device : torch.device or str
+        Device the coordinates live on (or will be moved to). Controls whether
+        the CUDA-only ``nv`` builder is eligible.
+    nf : int, default: 1
+        Number of frames in the batch. ``vesin`` is selected only when
+        ``nf == 1`` and ``vesin.torch`` is importable.
+
+    Returns
+    -------
+    str
+        One of ``"nv"``, ``"vesin"``, or ``"dense"``.
+    """
+    from deepmd.pt.utils.nv_nlist import (
+        is_nv_available,
+    )
+    from deepmd.pt_expt.utils.vesin_neighbor_list import (
+        is_vesin_torch_available,
+    )
+
+    if nf < 1:
+        raise ValueError(f"nf must be >= 1, got {nf}")
+
+    dev = torch.device(device)
+    if dev.type == "cuda" and is_nv_available():
+        return "nv"
+    if nf == 1 and is_vesin_torch_available():
+        return "vesin"
+    if dev.type == "cuda" and not is_nv_available():
+        log.warning(
+            "nvalchemi-toolkit-ops is unavailable; falling back from "
+            "neighbor_graph_method='auto' to the dense graph builder"
+            + (
+                ""
+                if nf == 1
+                else " (vesin is not used for nf>1; its API loops frames in Python)"
+            )
+            + ". Install it with `pip install nvalchemi-toolkit-ops` to enable "
+            "the NV graph builder."
+        )
+    return "dense"
+
+
 def resolve_neighbor_graph_method(
     requested: str,
     device: torch.device,
@@ -36,6 +105,8 @@ def resolve_neighbor_graph_method(
     -------
     str
         The concrete builder name, either ``"dense"`` or ``"nv"``.
+        Training auto never selects ``vesin`` (per-frame Python loop); use
+        :func:`resolve_auto_graph_builder` for inference auto selection.
 
     Raises
     ------

deepmd/pt_expt/utils/vesin_graph_builder.py

@@ -8,10 +8,12 @@
 
 Scope note: ``vesin.torch``'s API is single-system, so this builder LOOPS over
 frames in Python (~1 ms/frame call overhead measured on GPU). It is intended
-for ``nf == 1`` inference and CPU use. It is never on a default hot path:
-``neighbor_graph_method=None`` resolves to the ``"dense"`` converter, and
-vesin is explicit opt-in only. For batched multi-frame GPU work prefer
-``nv`` (:mod:`.nv_graph_builder`), which batches all frames in one kernel.
+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 ``nf == 1`` and ``vesin.torch`` is importable (and ``nv`` is
+unavailable on CUDA); multi-frame batches stay on ``nv``/``dense``. Training
+auto never selects vesin. Prefer ``nv`` (:mod:`.nv_graph_builder`) for batched
+multi-frame GPU work, which batches all frames in one kernel.
 """
 
 from __future__ import (

source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py

@@ -422,16 +422,31 @@ def test_unsupported_extension_raises(self) -> None:
 class TestNeighborGraphMethodResolution(unittest.TestCase):
     """Auto graph-builder selection must cover each host policy explicitly."""
 
+    def test_auto_deferred_until_nf_known(self) -> None:
+        """Construction-time resolve leaves ``auto`` unresolved without ``nf``."""
+        self.assertEqual(
+            PtExptDeepEval._resolve_neighbor_graph_method("auto"),
+            "auto",
+        )
+
     def test_auto_resolution(self) -> None:
+        # (device, nv, vesin, nf, expected, warns)
         cases = (
-            ("cpu", False, "dense", False),
-            ("cuda", True, "nv", False),
-            ("cuda", False, "dense", True),
-        )
-        for device_type, nv_available, expected, warns in cases:
+            ("cpu", False, True, 1, "vesin", False),
+            ("cpu", False, True, 4, "dense", False),
+            ("cpu", False, False, 1, "dense", False),
+            ("cuda", True, True, 1, "nv", False),
+            ("cuda", True, True, 4, "nv", False),
+            ("cuda", False, True, 1, "vesin", False),
+            ("cuda", False, True, 4, "dense", True),
+            ("cuda", False, False, 1, "dense", True),
+        )
+        for device_type, nv_available, vesin_available, nf, expected, warns in cases:
             with self.subTest(
                 device_type=device_type,
                 nv_available=nv_available,
+                vesin_available=vesin_available,
+                nf=nf,
             ):
                 with (
                     mock.patch(
@@ -442,17 +457,23 @@ def test_auto_resolution(self) -> None:
                         "deepmd.pt.utils.nv_nlist.is_nv_available",
                         return_value=nv_available,
                     ),
+                    mock.patch(
+                        "deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available",
+                        return_value=vesin_available,
+                    ),
                 ):
                     if warns:
                         with self.assertLogs(
-                            "deepmd.pt_expt.infer.deep_eval",
+                            "deepmd.pt_expt.utils.graph_builder",
                             level="WARNING",
                         ):
                             actual = PtExptDeepEval._resolve_neighbor_graph_method(
-                                "auto"
+                                "auto", nf=nf
                             )
                     else:
-                        actual = PtExptDeepEval._resolve_neighbor_graph_method("auto")
+                        actual = PtExptDeepEval._resolve_neighbor_graph_method(
+                            "auto", nf=nf
+                        )
                 self.assertEqual(actual, expected)
 
 

source/tests/pt_expt/model/test_graph_builder_dispatch.py

@@ -67,13 +67,20 @@ def _make_model():
     return EnergyModel(ds, ft, type_map=["O", "H"]).to(env.DEVICE)
 
 
-def _eval(model, method):
+def _eval(model, method, nf: int = 1):
     rng = np.random.default_rng(0)
     coord = torch.tensor(
-        rng.random((1, 6, 3)) * 4.0, dtype=torch.float64, device=env.DEVICE
+        rng.random((nf, 6, 3)) * 4.0, dtype=torch.float64, device=env.DEVICE
+    )
+    atype = torch.tensor(
+        [[0, 1, 1, 0, 1, 1]] * nf, dtype=torch.int64, device=env.DEVICE
+    )
+    box = (
+        (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6.0)
+        .reshape(1, 3, 3)
+        .expand(nf, 3, 3)
+        .clone()
     )
-    atype = torch.tensor([[0, 1, 1, 0, 1, 1]], dtype=torch.int64, device=env.DEVICE)
-    box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6.0).reshape(1, 3, 3)
     ret = model.forward_common(coord, atype, box, neighbor_graph_method=method)
     # graph path returns the output-agnostic dict (no translated force/virial);
     # energy_redu = total energy, energy_derv_r = d energy / d coord (force parity)
@@ -138,12 +145,43 @@ def test_explicit_nv_rejects_cpu():
         resolve_neighbor_graph_method("nv", torch.device("cpu"))
 
 
+@pytest.mark.parametrize(
+    ("device", "nv", "vesin", "nf", "expected"),
+    [
+        ("cpu", False, True, 1, "vesin"),
+        ("cpu", False, True, 4, "dense"),
+        ("cpu", True, False, 1, "dense"),
+        ("cuda", True, True, 1, "nv"),
+        ("cuda", True, True, 4, "nv"),
+        ("cuda", False, True, 1, "vesin"),
+        ("cuda", False, True, 4, "dense"),
+        ("cuda", False, False, 1, "dense"),
+    ],
+)
+def test_resolve_auto_graph_builder_ladder(
+    device: str, nv: bool, vesin: bool, nf: int, expected: str
+) -> None:
+    from deepmd.pt_expt.utils.graph_builder import (
+        resolve_auto_graph_builder,
+    )
+
+    with (
+        patch("deepmd.pt.utils.nv_nlist.is_nv_available", return_value=nv),
+        patch(
+            "deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available",
+            return_value=vesin,
+        ),
+    ):
+        assert resolve_auto_graph_builder(device, nf=nf) == expected
+
+
 @pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed")
-def test_vesin_matches_dense_energy_force():
+@pytest.mark.parametrize("nf", [1, 4])
+def test_vesin_matches_dense_energy_force(nf: int):
     torch.manual_seed(0)
     model = _make_model()
-    e_d, f_d = _eval(model, "dense")
-    e_v, f_v = _eval(model, "vesin")
+    e_d, f_d = _eval(model, "dense", nf=nf)
+    e_v, f_v = _eval(model, "vesin", nf=nf)
     tol = 1e-12 if env.DEVICE.type == "cpu" else 1e-10
     torch.testing.assert_close(e_v, e_d, rtol=tol, atol=tol)
     torch.testing.assert_close(f_v, f_d, rtol=tol, atol=tol)
@@ -153,11 +191,12 @@ def test_vesin_matches_dense_energy_force():
     not (torch.cuda.is_available() and is_nv_available()),
     reason="nvalchemiops requires CUDA + nvalchemi-toolkit-ops",
 )
-def test_nv_matches_dense_energy_force():
+@pytest.mark.parametrize("nf", [1, 4])
+def test_nv_matches_dense_energy_force(nf: int):
     torch.manual_seed(0)
     model = _make_model()
-    e_d, f_d = _eval(model, "dense")
-    e_n, f_n = _eval(model, "nv")
+    e_d, f_d = _eval(model, "dense", nf=nf)
+    e_n, f_n = _eval(model, "nv", nf=nf)
     tol = 1e-10  # CUDA fp64: absorbs scatter-atomic / index_add nondeterminism
     torch.testing.assert_close(e_n, e_d, rtol=tol, atol=tol)
     torch.testing.assert_close(f_n, f_d, rtol=tol, atol=tol)