test(lmp): bridged DPA4 2-rank parity with a cross-boundary close pair

Replaces the fails-fast test: the with-comm artifact now completes the
SFPG partials across ranks, so a 2-rank run with the 0.9 A pair
straddling the processors 2 1 1 boundary must (and does) match the
1-rank reference (issue deepmodeling#5906 Task 2 E2E).

Author: Han Wang

source/lmp/tests/test_lammps_dpa4_zbl_pt2.py

@@ -12,18 +12,19 @@
 handling, per-atom virial accumulation and unit conversion), so a regression
 confined to it was invisible.
 
-Single-rank only, deliberately
-------------------------------
+Multi-rank correctness (issue #5906)
+------------------------------------
 Bridging enables the descriptor's Source Freeze Propagation Gate, whose
 per-node ``eta_j = prod_{e: src_e = j} w_e`` folds a node's FULL outgoing-edge
-set.  Edges exist only for owned centres, so eta is incomplete on every rank
-and the freeze exports NO with-comm artifact (``gen_dpa4_zbl.py`` asserts
-``has_comm_artifact is False`` and that no nested
-``forward_lower_with_comm.pt2`` entry exists).  There is therefore no
-correct multi-rank answer to compare against; what this file pins instead is
-that a multi-rank run FAILS LOUDLY rather than silently returning
-wrong-but-plausible numbers -- see
-``test_pair_deepmd_mpi_dpa4_zbl_fails_fast``.
+set.  Edges exist only for owned centres, so the per-node partials are
+rank-incomplete; the with-comm artifact completes them with one
+reverse-accumulate + forward-broadcast border exchange before the gate is
+applied (``gen_dpa4_zbl.py`` asserts ``has_comm_artifact is True`` and that
+the nested ``forward_lower_with_comm.pt2`` entry exists).  This file pins the
+end-to-end contract with a 2-rank vs 1-rank parity run whose sub-``r_outer``
+close pair STRADDLES the ``processors 2 1 1`` x-boundary -- without that
+geometry every cross-rank gate contribution is ``log w = 0`` and the parity
+holds vacuously -- see ``test_pair_deepmd_mpi_dpa4_zbl_close_pair_parity``.
 
 Reference values are computed LIVE at test-setup time via
 ``deepmd.infer.DeepPot.eval`` on the archive itself, mirroring
@@ -72,13 +73,13 @@
 data_file = Path(__file__).parent / "data_dpa4_zbl_pt2.lmp"
 # The MPI runner is backend-agnostic (DATAFILE PB_FILE OUTPUT + flags); reuse
 # the DPA3 driver verbatim rather than duplicate it (same pattern as
-# test_lammps_dpa4_graph_pt2.py).  Only the fail-fast test below uses it.
+# test_lammps_dpa4_graph_pt2.py).  Only the multi-rank tests below use it.
 mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py"
 
-# Ceiling for the mpirun invocation.  The fail-fast under test is expected to
-# throw on EVERY rank before any collective, so a timeout means the guard
-# regressed into a deadlock -- which is a test failure, not a slow machine.
-_MPI_DEFAULT_TIMEOUT = 300.0
+# Ceiling for every mpirun invocation: a with-comm desync hangs the
+# collective forever, so an unbounded should-succeed regression would hang
+# the whole suite -- a timeout is a test failure, not a slow machine.
+_MPI_DEFAULT_TIMEOUT = 600.0
 
 # 6-atom NiO system, coordinates verbatim from
 # ``source/tests/infer/gen_dpa4_zbl.py``'s ``_COORDS``: atoms 0 and 1 sit
@@ -101,6 +102,16 @@
 # -> LAMMPS types [1,1,1,2,2,2] under identity ``pair_coeff * *``.
 type_NiO = np.array([1, 1, 1, 2, 2, 2])
 
+# Close-pair variant for the 2-rank parity test (issue #5906): the SAME
+# 6-atom geometry shifted along x so the 0.9 A Ni-Ni pair (inside the
+# bridging window, r_inner=0.8 < 0.9 < r_outer=1.2) STRADDLES the
+# ``processors 2 1 1`` boundary at lx/2 = 6.5.  Atom 0 lands at x = 6.3
+# (rank 0) and atom 1 at x = 7.2 (rank 1), so the pair's SFPG gate
+# contribution crosses the rank boundary -- the load-bearing geometry.
+_CLOSE_PAIR_X_SHIFT = 5.3
+coord_close_pair = coord + np.array([_CLOSE_PAIR_X_SHIFT, 0.0, 0.0])
+data_file_close_pair = Path(__file__).parent / "data_dpa4_zbl_close_pair_pt2.lmp"
+
 # Reference values, populated by ``_compute_expected`` in ``setup_module``.
 expected_e = None
 expected_ae = None
@@ -192,11 +203,13 @@ def setup_module() -> None:
         )
     _compute_expected()
     write_lmp_data(box, coord, type_NiO, data_file)
+    write_lmp_data(box, coord_close_pair, type_NiO, data_file_close_pair)
 
 
 def teardown_module() -> None:
-    if data_file.exists():
-        os.remove(data_file)
+    for f in (data_file, data_file_close_pair):
+        if f.exists():
+            os.remove(f)
 
 
 def _lammps(data_file, units="metal") -> PyLammps:
@@ -295,17 +308,24 @@ def test_pair_deepmd_atom_energy_and_virial(lammps) -> None:
 
 
 # ---------------------------------------------------------------------------
-# Multi-rank: NOT a correctness test -- a fail-fast test.
+# Multi-rank: 2-rank vs 1-rank close-pair parity (issue #5906 Task 2 E2E).
 # ---------------------------------------------------------------------------
 
 
-def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict:
+def _run_mpi_subprocess(
+    nprocs: int,
+    processors: str,
+    timeout: float = _MPI_DEFAULT_TIMEOUT,
+    data_path: Path | None = None,
+) -> dict:
     """Run the (backend-agnostic) DPA3 MPI runner against the bridged archive
-    and return ``{"returncode", "stdout", "stderr", "timed_out"}``.
+    and return the parsed ``{"pe", "forces", "virials"}`` output.
 
     Always bounded: on expiry the WHOLE mpirun process group is SIGKILLed
     (killing only mpirun can leave orphaned ranks blocking in a collective).
     """
+    if data_path is None:
+        data_path = data_file
     with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f:
         out_path = f.name
     try:
@@ -315,7 +335,7 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict:
             str(nprocs),
             sys.executable,
             str(mpi_runner),
-            str(data_file.resolve()),
+            str(data_path.resolve()),
             str(pb_file.resolve()),
             out_path,
             "--processors",
@@ -329,18 +349,25 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict:
         except sp.TimeoutExpired:
             os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
             stdout, stderr = proc.communicate()
-            return {
-                "returncode": None,
-                "stdout": stdout or "",
-                "stderr": stderr or "",
-                "timed_out": True,
-            }
-        return {
-            "returncode": proc.returncode,
-            "stdout": stdout,
-            "stderr": stderr,
-            "timed_out": False,
-        }
+            raise RuntimeError(
+                f"mpirun timed out after {timeout}s (process group killed); "
+                "a should-succeed MPI regression is deadlocked.\n"
+                f"stdout:\n{(stdout or '')[-2000:]}\n"
+                f"stderr:\n{(stderr or '')[-2000:]}"
+            ) from None
+        if proc.returncode != 0:
+            raise RuntimeError(
+                f"mpirun exited {proc.returncode}.\n"
+                f"stdout:\n{stdout[-2000:]}\nstderr:\n{stderr[-2000:]}"
+            )
+        with open(out_path) as fh:
+            lines = fh.read().strip().splitlines()
+        pe = float(lines[0])
+        rows = np.array(
+            [list(map(float, line.split())) for line in lines[1:]],
+            dtype=np.float64,
+        )
+        return {"pe": pe, "forces": rows[:, :3], "virials": rows[:, 3:]}
     finally:
         if os.path.exists(out_path):
             os.remove(out_path)
@@ -352,40 +379,33 @@ def _run_mpi_subprocess(nprocs: int, processors: str, timeout: float) -> dict:
 @pytest.mark.skipif(
     importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed"
 )
-def test_pair_deepmd_mpi_dpa4_zbl_fails_fast() -> None:
-    """A multi-rank run of a BRIDGED archive must fail loudly, not answer.
-
-    The bridged model is single-rank only by construction (see the module
-    docstring), so its freeze exports no with-comm artifact while still
-    declaring ``has_message_passing``.  ``DeepPotPTExpt::compute_inner``'s
-    dispatch reads exactly that combination -- graph lower + ``nprocs > 1`` +
-    message passing + no with-comm artifact -- and throws before building any
-    tensors.  Without the guard the run would fall through to the plain
-    single-rank artifact on a per-rank subdomain, where the bridging gate's
-    per-node eta is incomplete: wrong, finite, plausible numbers.
-
-    The failure is uniform across ranks (every rank evaluates the same
-    metadata-only predicate before any collective), so a TIMEOUT is a failure
-    of this test: it would mean the guard regressed into a deadlock.
-
-    This is deliberately the ONLY multi-rank test in this file; there is no
-    correct multi-rank reference for a bridged model to compare against.
+def test_pair_deepmd_mpi_dpa4_zbl_close_pair_parity() -> None:
+    """Issue #5906 Task 2 E2E: a bridged model, 2 ranks, with a 0.9 A
+    contact STRADDLING the ``processors 2 1 1`` x-boundary.
+
+    Without the cross-boundary close pair this test is vacuous (edges at
+    ``r >= r_outer`` contribute ``log w = 0``); the geometry guard below
+    pins that the pair actually straddles the split.  The 2-rank run
+    exercises the with-comm artifact's SFPG completion (reverse-accumulate
+    + broadcast of the per-node ``[log_eta, zero_count]`` partials), the
+    per-block ghost exchange, and the reverse-comm force fold; the 1-rank
+    run is the plain-artifact reference on the same trajectory.
     """
-    out = _run_mpi_subprocess(
-        nprocs=2, processors="2 1 1", timeout=_MPI_DEFAULT_TIMEOUT
-    )
-    assert not out["timed_out"], (
-        "Multi-rank run of the bridged archive timed out instead of failing "
-        "promptly; the dispatch guard must throw on every rank BEFORE any "
-        "collective."
+    lx = float(box[1] - box[0])
+    x_lo, x_hi = float(coord_close_pair[0, 0]), float(coord_close_pair[1, 0])
+    assert x_lo < lx / 2.0 < x_hi, (
+        "the close pair no longer straddles the processors 2 1 1 boundary; "
+        "the parity below would be vacuous for the SFPG exchange"
     )
-    assert out["returncode"] != 0, (
-        "Expected the multi-rank run of a bridged (no with-comm artifact) "
-        "archive to fail loudly, but it exited 0.\n"
-        f"stdout:\n{out['stdout'][-2000:]}\nstderr:\n{out['stderr'][-2000:]}"
+    ref = _run_mpi_subprocess(
+        nprocs=1, processors="1 1 1", data_path=data_file_close_pair
     )
-    combined = out["stdout"] + out["stderr"]
-    assert "with-comm artifact" in combined, (
-        "Expected the documented fail-loud message (mentioning the missing "
-        f"'with-comm artifact'), got:\n{combined[-2000:]}"
+    par = _run_mpi_subprocess(
+        nprocs=2, processors="2 1 1", data_path=data_file_close_pair
     )
+    assert par["pe"] == pytest.approx(ref["pe"], rel=1e-8, abs=1e-10)
+    np.testing.assert_allclose(par["forces"], ref["forces"], atol=1e-8, rtol=0)
+    # Same tolerance rationale as test_lammps_dpa4_graph_pt2.py's twin: the
+    # relative component absorbs CUDA atomic-scatter ordering noise without
+    # loosening the CPU-exact case.
+    np.testing.assert_allclose(par["virials"], ref["virials"], atol=1e-8, rtol=1e-8)