From 0cf6550fe5a3f740e39ced34603bcc1670b00aea Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 11:12:26 +0800 Subject: [PATCH] test(pt): reuse PT2 artifacts for metadata checks Build the base AOTInductor package once in a module-scoped fixture, derive metadata variants by editing ZIP extras, and share the artifact between metadata and inference test classes.\n\nCoding-Agent: Codex\nCodex-Version: codex-cli 0.144.4\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh --- source/tests/pt_expt/infer/test_deep_eval.py | 220 +++++++++++-------- 1 file changed, 132 insertions(+), 88 deletions(-) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 856bc0724a..510abe1104 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -6,12 +6,14 @@ """ import importlib +import json import os import tempfile import unittest import zipfile import numpy as np +import pytest import torch from deepmd.infer import ( @@ -55,6 +57,46 @@ def _assert_repeatable(a, b) -> None: np.testing.assert_allclose(a, b, rtol=1e-10, atol=1e-10) +def _deserialize_pt2_with_default_device_cleared( + path: str, + data: dict, + *, + do_atomic_virial: bool = False, +) -> None: + """Compile a PT2 archive without inheriting a leaked default device.""" + previous_default_device = torch.get_default_device() + torch.set_default_device(None) + try: + deserialize_to_file(path, data, do_atomic_virial=do_atomic_virial) + finally: + torch.set_default_device(previous_default_device) + + +def _copy_pt2_with_extra_json( + source: str, + destination: str, + *, + remove: set[str] | None = None, + replace: dict[str, object] | None = None, +) -> None: + """Copy a compiled PT2 archive while editing only its extra JSON files. + + The compiled AOTInductor artifact is byte-for-byte reused. This lets + metadata, ZIP-layout, and accessor tests exercise alternate metadata + without paying for another code-generation pass. + """ + remove = remove or set() + replace = replace or {} + edited_names = remove | set(replace) + with zipfile.ZipFile(source, "r") as zin, zipfile.ZipFile(destination, "w") as zout: + for info in zin.infolist(): + if info.filename in edited_names: + continue + zout.writestr(info, zin.read(info.filename)) + for name, value in replace.items(): + zout.writestr(name, json.dumps(value)) + + class TestDeepEvalEner(unittest.TestCase): """Test pt_expt inference for energy models.""" @@ -657,71 +699,86 @@ def test_ase_nlist_multiple_frames(self) -> None: np.testing.assert_allclose(v1, v2, rtol=1e-10, atol=1e-10, err_msg="virial") -class TestDeepEvalEnerPt2(unittest.TestCase): - """Test pt_expt inference for energy models via .pt2 (AOTInductor).""" - - @classmethod - def setUpClass(cls) -> None: - cls.rcut = 4.0 - cls.rcut_smth = 0.5 - cls.sel = [8, 6] - cls.nt = 2 - cls.type_map = ["foo", "bar"] - - # Build pt_expt model - ds = DescrptSeA(cls.rcut, cls.rcut_smth, cls.sel) - ft = EnergyFittingNet( - cls.nt, - ds.get_dim_out(), - mixed_types=ds.mixed_types(), - seed=GLOBAL_SEED, - ) - cls.model = EnergyModel(ds, ft, type_map=cls.type_map) - cls.model = cls.model.to(torch.float64) - cls.model.eval() - - # Serialize and save to .pt2 - cls.model_data = {"model": cls.model.serialize()} - cls.tmpfile = tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) - cls.tmpfile.close() - # Temporarily clear default device to avoid poisoning AOTInductor - # compilation (tests/pt/__init__.py sets it to "cuda:9999999"). - torch.set_default_device(None) - try: - deserialize_to_file(cls.tmpfile.name, cls.model_data, do_atomic_virial=True) - finally: - torch.set_default_device("cuda:9999999") - - cls.meta_tmpfile = tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) - cls.meta_tmpfile.close() - with ( - zipfile.ZipFile(cls.tmpfile.name, "r") as zin, - zipfile.ZipFile(cls.meta_tmpfile.name, "w") as zout, - ): - for info in zin.infolist(): - if info.filename == "model/extra/model.json": - continue - zout.writestr(info, zin.read(info.filename)) - - # Also save to .pte for cross-format comparison - cls.pte_tmpfile = tempfile.NamedTemporaryFile(suffix=".pte", delete=False) - cls.pte_tmpfile.close() - deserialize_to_file(cls.pte_tmpfile.name, cls.model_data, do_atomic_virial=True) +@pytest.fixture(scope="module") +def deep_eval_ener_pt2_artifacts( + tmp_path_factory: pytest.TempPathFactory, +) -> dict[str, object]: + """Build one base PT2 artifact and derive all metadata variants from it.""" + root = tmp_path_factory.mktemp("deep_eval_ener_pt2") + rcut = 4.0 + rcut_smth = 0.5 + sel = [8, 6] + nt = 2 + type_map = ["foo", "bar"] + + ds = DescrptSeA(rcut, rcut_smth, sel) + ft = EnergyFittingNet( + nt, + ds.get_dim_out(), + mixed_types=ds.mixed_types(), + seed=GLOBAL_SEED, + ) + model = EnergyModel(ds, ft, type_map=type_map).to(torch.float64).eval() + model_data = {"model": model.serialize()} + + pt2_path = str(root / "model.pt2") + _deserialize_pt2_with_default_device_cleared( + pt2_path, + model_data, + do_atomic_virial=True, + ) - # Create DeepPot for .pt2 - cls.dp = DeepPot(cls.tmpfile.name) - # Create DeepPot for metadata-only .pt2 - cls.dp_meta = DeepPot(cls.meta_tmpfile.name) - # Create DeepPot for .pte reference - cls.dp_pte = DeepPot(cls.pte_tmpfile.name) + metadata_only_pt2_path = str(root / "metadata_only.pt2") + _copy_pt2_with_extra_json( + pt2_path, + metadata_only_pt2_path, + remove={"model/extra/model.json"}, + ) - @classmethod - def tearDownClass(cls) -> None: - import os + training_config = {"type_map": type_map, "descriptor": {"type": "se_e2_a"}} + configured_pt2_path = str(root / "configured.pt2") + _copy_pt2_with_extra_json( + pt2_path, + configured_pt2_path, + replace={"model/extra/model_def_script.json": training_config}, + ) - os.unlink(cls.tmpfile.name) - os.unlink(cls.meta_tmpfile.name) - os.unlink(cls.pte_tmpfile.name) + pte_path = str(root / "model.pte") + deserialize_to_file(pte_path, model_data, do_atomic_virial=True) + + return { + "rcut": rcut, + "rcut_smth": rcut_smth, + "sel": sel, + "nt": nt, + "type_map": type_map, + "model": model, + "model_data": model_data, + "pt2_path": pt2_path, + "metadata_only_pt2_path": metadata_only_pt2_path, + "pte_path": pte_path, + "dp": DeepPot(pt2_path), + "dp_meta": DeepPot(metadata_only_pt2_path), + "dp_pte": DeepPot(pte_path), + "dp_with_config": DeepPot(configured_pt2_path), + "training_config": training_config, + } + + +@pytest.fixture(scope="class") +def _bind_deep_eval_ener_pt2_artifacts( + request: pytest.FixtureRequest, + deep_eval_ener_pt2_artifacts: dict[str, object], +) -> None: + """Expose the shared module fixture through the legacy unittest class.""" + assert request.cls is not None + for name, value in deep_eval_ener_pt2_artifacts.items(): + setattr(request.cls, name, value) + + +@pytest.mark.usefixtures("_bind_deep_eval_ener_pt2_artifacts") +class TestDeepEvalEnerPt2Metadata(unittest.TestCase): + """Test PT2 metadata, ZIP layout, and metadata-only inference.""" def test_get_rcut(self) -> None: self.assertAlmostEqual(self.dp.deep_eval.get_rcut(), self.rcut) @@ -756,27 +813,9 @@ def test_get_model_def_script(self) -> None: self.assertEqual(mds, {}) def test_get_model_def_script_with_params(self) -> None: - """Export with model_params → get_model_def_script returns them.""" - training_config = {"type_map": self.type_map, "descriptor": {"type": "se_e2_a"}} - with tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) as f: - tmpfile2 = f.name - try: - torch.set_default_device(None) - try: - data_with_config = { - **self.model_data, - "model_def_script": training_config, - } - deserialize_to_file(tmpfile2, data_with_config) - finally: - torch.set_default_device("cuda:9999999") - dp2 = DeepPot(tmpfile2) - mds = dp2.deep_eval.get_model_def_script() - self.assertEqual(mds, training_config) - finally: - import os - - os.unlink(tmpfile2) + """A metadata-edited variant exposes its replaced model definition.""" + mds = self.dp_with_config.deep_eval.get_model_def_script() + self.assertEqual(mds, self.training_config) def test_model_api_delegation(self) -> None: """Verify that model API calls are delegated to the deserialized dpmodel.""" @@ -790,11 +829,11 @@ def test_model_api_delegation(self) -> None: def test_pt2_file_is_zip(self) -> None: """The .pt2 file should be a valid ZIP archive.""" - self.assertTrue(zipfile.is_zipfile(self.tmpfile.name)) + self.assertTrue(zipfile.is_zipfile(self.pt2_path)) def test_pt2_has_metadata(self) -> None: """The .pt2 ZIP should contain metadata entries under ``model/extra/``.""" - with zipfile.ZipFile(self.tmpfile.name, "r") as zf: + with zipfile.ZipFile(self.pt2_path, "r") as zf: names = zf.namelist() self.assertIn("model/extra/metadata.json", names) self.assertIn("model/extra/model_def_script.json", names) @@ -804,7 +843,7 @@ def test_pt2_has_metadata(self) -> None: def test_metadata_only_pt2_has_no_model_json(self) -> None: """The metadata-only .pt2 keeps metadata but drops model.json.""" - with zipfile.ZipFile(self.meta_tmpfile.name, "r") as zf: + with zipfile.ZipFile(self.metadata_only_pt2_path, "r") as zf: names = zf.namelist() self.assertIn("model/extra/metadata.json", names) self.assertNotIn("model/extra/model.json", names) @@ -840,6 +879,11 @@ def test_metadata_only_pt2_eval_parity(self) -> None: for ref, test in zip(full_out, meta_out, strict=True): np.testing.assert_array_equal(test, ref) + +@pytest.mark.usefixtures("_bind_deep_eval_ener_pt2_artifacts") +class TestDeepEvalEnerPt2(unittest.TestCase): + """Test numerical inference through the shared PT2 AOTInductor artifact.""" + def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" rng = np.random.default_rng(GLOBAL_SEED) @@ -979,7 +1023,7 @@ def test_oversized_nlist(self) -> None: implementation that truncates without sorting) gives a different energy. This proves the distance sort is necessary. """ - exported = torch.export.load(self.pte_tmpfile.name) + exported = torch.export.load(self.pte_path) exported_mod = exported.module() nnei = sum(self.sel) # model's expected neighbor count @@ -1064,7 +1108,7 @@ def test_oversized_nlist(self) -> None: def test_serialize_round_trip(self) -> None: """Test .pt2 → serialize_from_file → deserialize → model gives same outputs.""" - loaded_data = serialize_from_file(self.tmpfile.name) + loaded_data = serialize_from_file(self.pt2_path) model2 = EnergyModel.deserialize(loaded_data["model"]) model2 = model2.to(torch.float64)