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

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

@@ -18,6 +18,93 @@
 
 log = logging.getLogger(__name__)
 
+# Warn once per process: resolve_auto_graph_builder runs per batch after
+# call-time resolution, but "install nvalchemi-toolkit-ops" is a one-shot
+# action for the user.
+_warned_auto_no_nv = False
+
+
+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"``.
+
+    Raises
+    ------
+    ValueError
+        If ``nf`` is not a positive ``int`` (``bool`` is rejected).
+    """
+    global _warned_auto_no_nv
+
+    from deepmd.pt.utils.nv_nlist import (
+        is_nv_available,
+    )
+    from deepmd.pt_expt.utils.vesin_neighbor_list import (
+        is_vesin_torch_available,
+    )
+
+    # ``bool`` is a subclass of ``int``; reject it explicitly.
+    if type(nf) is not int:
+        raise ValueError(f"nf must be a positive int, got {nf!r}")
+    if nf < 1:
+        raise ValueError(f"nf must be >= 1, got {nf}")
+
+    dev = torch.device(device)
+    nv_available = is_nv_available()
+    if dev.type == "cuda" and nv_available:
+        return "nv"
+    if nf == 1 and is_vesin_torch_available():
+        return "vesin"
+    if dev.type == "cuda" and not nv_available:
+        if not _warned_auto_no_nv:
+            _warned_auto_no_nv = True
+            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,
@@ -36,6 +123,8 @@
     -------
     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,27 @@ 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:
+                        # reset warn-once so each case can assert the message
+                        import deepmd.pt_expt.utils.graph_builder as gb
+
+                        gb._warned_auto_no_nv = False
                         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)