test(jax): isolate training test in subprocess

Run the JAX training end-to-end test in a child Python process so CUDA/XLA teardown failures cannot poison the parent pytest process. Shrink the model and use the single-frame water data to minimize memory use.\n\nAuthored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

Author: njzjz-bot

source/tests/jax/test_training.py

@@ -2,27 +2,17 @@
 """End-to-end tests for the local JAX training entrypoint."""
 
 import argparse
-import functools
 import json
 import os
 import shutil
-import signal
+import subprocess
+import sys
 import tempfile
+import textwrap
 import unittest
-from collections.abc import (
-    Callable,
-)
-from copy import (
-    deepcopy,
-)
 from pathlib import (
     Path,
 )
-from typing import (
-    Any,
-    TypeVar,
-    cast,
-)
 from unittest.mock import (
     patch,
 )
@@ -33,65 +23,49 @@
 from deepmd.jax.entrypoints.main import (
     main,
 )
-from deepmd.jax.entrypoints.train import (
-    train,
-)
 from deepmd.utils.compat import (
     convert_optimizer_v31_to_v32,
 )
 
-_F = TypeVar("_F", bound=Callable[..., Any])
-
-
-def _training_timeout(seconds: int) -> Callable[[_F], _F]:
-    """Limit real training tests on platforms that support SIGALRM."""
-
-    def decorate(func: _F) -> _F:
-        if not hasattr(signal, "SIGALRM"):
-            return func
-
-        @functools.wraps(func)
-        def wrapped(*args: Any, **kwargs: Any) -> Any:
-            def raise_timeout(signum: int, frame: Any) -> None:
-                raise TimeoutError(f"training test exceeded {seconds} seconds")
-
-            previous_handler = signal.signal(signal.SIGALRM, raise_timeout)
-            signal.alarm(seconds)
-            try:
-                return func(*args, **kwargs)
-            finally:
-                signal.alarm(0)
-                signal.signal(signal.SIGALRM, previous_handler)
-
-        return cast("_F", wrapped)
-
-    return decorate
-
-
-TRAINING_TEST_TIMEOUT = _training_timeout(60)
-
 MODEL_SE_E2_A = {
     "type_map": ["O", "H", "B"],
     "descriptor": {
         "type": "se_e2_a",
-        "sel": [46, 92, 4],
+        "sel": [6, 12, 1],
         "rcut_smth": 0.50,
         "rcut": 4.00,
-        "neuron": [25, 50, 100],
+        "neuron": [2, 4, 8],
         "resnet_dt": False,
-        "axis_neuron": 16,
+        "axis_neuron": 2,
         "type_one_side": True,
         "seed": 1,
     },
     "fitting_net": {
-        "neuron": [24, 24, 24],
+        "neuron": [4, 4, 4],
         "resnet_dt": True,
         "seed": 1,
     },
-    "data_stat_nbatch": 20,
+    "data_stat_nbatch": 1,
 }
 
 
+TRAINING_SCRIPT = """
+from pathlib import Path
+from unittest.mock import patch
+
+from deepmd.main import main
+
+with patch("deepmd.jax.entrypoints.train.SummaryPrinter.__call__"):
+    main(["--jax", "train", "input.json", "--log-level", "2"])
+
+for path in ["out.json", "lcurve.out", "checkpoint", "model-1.jax"]:
+    if not Path(path).exists():
+        raise FileNotFoundError(path)
+if "1" not in Path("lcurve.out").read_text():
+    raise AssertionError("lcurve.out does not contain the first training step")
+"""
+
+
 class TestJAXTraining(unittest.TestCase):
     """Regression tests for complete JAX training runs."""
 
@@ -103,12 +77,12 @@ def setUp(self) -> None:
 
         source_dir = Path(__file__).resolve().parents[1] / "pt" / "water"
         shutil.copytree(source_dir, self.work_dir / "water")
-        data_file = [str(self.work_dir / "water" / "data" / "data_0")]
+        data_file = [str(self.work_dir / "water" / "data" / "single")]
 
         with (self.work_dir / "water" / "se_atten.json").open() as f:
             self.config = json.load(f)
         self.config = convert_optimizer_v31_to_v32(self.config, warning=False)
-        self.config["model"] = deepcopy(MODEL_SE_E2_A)
+        self.config["model"] = MODEL_SE_E2_A
         self.config["model"]["data_stat_nbatch"] = 1
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
@@ -126,26 +100,18 @@ def tearDown(self) -> None:
         os.chdir(self.cwd)
         shutil.rmtree(self.work_dir)
 
-    @TRAINING_TEST_TIMEOUT
-    @patch("deepmd.jax.entrypoints.train.SummaryPrinter.__call__")
-    def test_train_entrypoint_runs_one_step_from_scratch(self, _summary) -> None:
-        """Run local JAX training and check that expected artifacts are written."""
-        train(
-            INPUT=str(self.input_file),
-            init_model=None,
-            restart=None,
-            output="out.json",
-            init_frz_model=None,
-            mpi_log="master",
-            log_level=2,
-            log_path=None,
+    def test_train_entrypoint_runs_one_step_from_scratch(self) -> None:
+        """Run local JAX training in a child process and check artifacts."""
+        proc = subprocess.run(
+            [sys.executable, "-c", textwrap.dedent(TRAINING_SCRIPT)],
+            cwd=self.work_dir,
+            text=True,
+            capture_output=True,
+            timeout=60,
+            check=False,
         )
 
-        self.assertTrue(Path("out.json").is_file())
-        self.assertTrue(Path("lcurve.out").is_file())
-        self.assertTrue(Path("checkpoint").is_file())
-        self.assertTrue(Path("model-1.jax").is_dir())
-        self.assertIn("1", Path("lcurve.out").read_text())
+        self.assertEqual(proc.returncode, 0, proc.stdout + proc.stderr)
 
     @patch("deepmd.jax.entrypoints.freeze.deserialize_to_file")
     @patch("deepmd.jax.entrypoints.freeze.serialize_from_file")