feat(pt_expt): include vesin in DeepEval auto graph-builder ladder

Training keeps CPU on dense (vesin loops frames). Inference auto now
shares resolve_auto_graph_builder: CUDA nv→vesin→dense, CPU vesin→dense.

Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>

Author: Shaurya2k06

deepmd/pt_expt/infer/deep_eval.py

@@ -192,16 +192,19 @@ 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`
+        (CUDA: ``nv`` if importable, else ``vesin`` if importable, else
+        ``dense``; CPU: ``vesin`` if 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
-        backend-selection API is deferred to the dense-nlist deprecation.
+        choice is performance-only. Training keeps a separate CPU-dense auto
+        policy (:func:`~deepmd.pt_expt.utils.graph_builder.resolve_neighbor_graph_method`)
+        because vesin's per-frame Python loop is not a multi-frame training
+        default. Consolidating the two knobs into a single backend-selection API
+        is deferred to the dense-nlist deprecation.
     **kwargs : dict
         Keyword arguments.
     """
@@ -281,23 +284,14 @@ def _resolve_neighbor_graph_method(method: str) -> str:
         if method != "auto":
             return method
 
-        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)
 
     def _setup_neighbor_backend(self, nlist_backend: str) -> None:
         """Resolve the graph or neighbor-list construction strategy.

deepmd/pt_expt/utils/graph_builder.py

@@ -19,6 +19,47 @@
 log = logging.getLogger(__name__)
 
 
+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``.
+
+    ``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.
+    """
+    from deepmd.pt.utils.nv_nlist import (
+        is_nv_available,
+    )
+    from deepmd.pt_expt.utils.vesin_neighbor_list import (
+        is_vesin_torch_available,
+    )
+
+    dev = torch.device(device)
+    if dev.type == "cuda":
+        if is_nv_available():
+            return "nv"
+        if is_vesin_torch_available():
+            return "vesin"
+        log.warning(
+            "nvalchemi-toolkit-ops and vesin[torch] are unavailable; falling "
+            "back from neighbor_graph_method='auto' to the dense graph builder."
+        )
+        return "dense"
+    if is_vesin_torch_available():
+        return "vesin"
+    return "dense"
+
+
 def resolve_neighbor_graph_method(
     requested: str,
     device: torch.device,
@@ -36,6 +77,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 ``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`)
+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

@@ -423,15 +423,19 @@ class TestNeighborGraphMethodResolution(unittest.TestCase):
     """Auto graph-builder selection must cover each host policy explicitly."""
 
     def test_auto_resolution(self) -> None:
+        # (device, nv, vesin, expected, warns)
         cases = (
-            ("cpu", False, "dense", False),
-            ("cuda", True, "nv", False),
-            ("cuda", False, "dense", True),
+            ("cpu", False, True, "vesin", False),
+            ("cpu", False, False, "dense", False),
+            ("cuda", True, True, "nv", False),
+            ("cuda", False, True, "vesin", False),
+            ("cuda", False, False, "dense", True),
         )
-        for device_type, nv_available, expected, warns in cases:
+        for device_type, nv_available, vesin_available, expected, warns in cases:
             with self.subTest(
                 device_type=device_type,
                 nv_available=nv_available,
+                vesin_available=vesin_available,
             ):
                 with (
                     mock.patch(
@@ -442,10 +446,14 @@ 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(

source/tests/pt_expt/model/test_graph_builder_dispatch.py

@@ -138,6 +138,33 @@ def test_explicit_nv_rejects_cpu():
         resolve_neighbor_graph_method("nv", torch.device("cpu"))
 
 
+@pytest.mark.parametrize(
+    ("device", "nv", "vesin", "expected"),
+    [
+        ("cpu", False, True, "vesin"),
+        ("cpu", True, False, "dense"),
+        ("cuda", True, True, "nv"),
+        ("cuda", False, True, "vesin"),
+        ("cuda", False, False, "dense"),
+    ],
+)
+def test_resolve_auto_graph_builder_ladder(
+    device: str, nv: bool, vesin: bool, 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) == expected
+
+
 @pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed")
 def test_vesin_matches_dense_energy_force():
     torch.manual_seed(0)