Skip to content

Commit dadb45c

Browse files
test(lmp): avoid duplicate pb conversions in tests (#5400)
Problem - `source/lmp/tests` converted several shared `.pbtxt` fixtures at module import time. - Multiple test modules target the same output files such as `graph.pb`, `graph2.pb`, and `deepspin_nlist-2.pb`, so collecting those tests reran the same conversion work multiple times. Change - add a small `ensure_converted_pb()` helper that skips conversion when the output `.pb` already exists and is newer than the source `.pbtxt` - make the helper use a lock file plus atomic replace so repeated imports do not race on the same output file - switch the affected LAMMPS tests to use the helper and add a focused regression test for the cache/refresh behavior Notes - Local validation: `python3 -m unittest source.lmp.tests.test_model_convert` - Local validation: `python3 -m compileall source/lmp/tests/model_convert.py source/lmp/tests/test_model_convert.py source/lmp/tests/test_deeptensor.py source/lmp/tests/test_dplr.py source/lmp/tests/test_lammps.py source/lmp/tests/test_lammps_3types.py source/lmp/tests/test_lammps_dpa_jax.py source/lmp/tests/test_lammps_dpa_pt.py source/lmp/tests/test_lammps_dpa_pt_nopbc.py source/lmp/tests/test_lammps_dpa_sel_pt.py source/lmp/tests/test_lammps_faparam.py source/lmp/tests/test_lammps_jax.py source/lmp/tests/test_lammps_pd.py source/lmp/tests/test_lammps_pt.py source/lmp/tests/test_lammps_spin.py source/lmp/tests/test_lammps_spin_nopbc.py source/lmp/tests/test_lammps_spin_nopbc_pt.py source/lmp/tests/test_lammps_spin_pt.py` Authored by OpenClaw (model: gpt-5.4) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added a centralized test helper for model conversion that performs atomic updates and avoids race conditions. * **Refactor** * Updated many test modules to use the helper and moved conversions into setup routines; conversions are now conditional where applicable, removing prior ad-hoc external invocations. * **Tests** * Added tests verifying skip-when-up-to-date, rebuild-on-stale, and stale-lock handling. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent aa44999 commit dadb45c

18 files changed

Lines changed: 316 additions & 102 deletions

source/lmp/tests/model_convert.py

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Helpers for preparing converted TensorFlow graph files in LAMMPS tests."""
3+
4+
from __future__ import (
5+
annotations,
6+
)
7+
8+
import errno
9+
import os
10+
import subprocess as sp
11+
import sys
12+
import tempfile
13+
import time
14+
from pathlib import (
15+
Path,
16+
)
17+
18+
_LOCK_TIMEOUT_SECONDS = 60.0
19+
_LOCK_POLL_SECONDS = 0.1
20+
21+
22+
def _is_up_to_date(source: Path, output: Path) -> bool:
23+
return output.exists() and output.stat().st_mtime_ns >= source.stat().st_mtime_ns
24+
25+
26+
def _read_lock_pid(lock_file: Path) -> int | None:
27+
try:
28+
for line in lock_file.read_text(encoding="utf-8").splitlines():
29+
if line.startswith("pid="):
30+
return int(line.split("=", maxsplit=1)[1])
31+
except (FileNotFoundError, ValueError):
32+
return None
33+
return None
34+
35+
36+
def _pid_is_running(pid: int) -> bool:
37+
try:
38+
os.kill(pid, 0)
39+
except ProcessLookupError:
40+
return False
41+
except PermissionError:
42+
return True
43+
except OSError as err:
44+
if err.errno == errno.ESRCH:
45+
return False
46+
raise
47+
return True
48+
49+
50+
def _should_break_stale_lock(lock_file: Path) -> bool:
51+
try:
52+
lock_stat = lock_file.stat()
53+
except FileNotFoundError:
54+
return False
55+
56+
lock_pid = _read_lock_pid(lock_file)
57+
if lock_pid is not None:
58+
return not _pid_is_running(lock_pid)
59+
60+
lock_age = time.time() - lock_stat.st_mtime
61+
return lock_age > _LOCK_TIMEOUT_SECONDS
62+
63+
64+
def ensure_converted_pb(source: Path, output: Path) -> Path:
65+
"""Convert ``source`` into ``output`` only when the target is missing or stale.
66+
67+
The conversion is protected by a simple lock file and uses atomic replacement so
68+
repeated imports across multiple test modules do not regenerate the same model
69+
more than once.
70+
"""
71+
source = source.resolve()
72+
output = output.resolve()
73+
output.parent.mkdir(parents=True, exist_ok=True)
74+
lock_file = output.with_name(f".{output.name}.lock")
75+
started = time.monotonic()
76+
77+
while True:
78+
if _is_up_to_date(source, output):
79+
return output
80+
try:
81+
fd = os.open(str(lock_file), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
82+
except FileExistsError as err:
83+
if _should_break_stale_lock(lock_file):
84+
lock_file.unlink(missing_ok=True)
85+
continue
86+
if time.monotonic() - started >= _LOCK_TIMEOUT_SECONDS:
87+
raise TimeoutError(f"Timed out waiting for {lock_file}") from err
88+
time.sleep(_LOCK_POLL_SECONDS)
89+
continue
90+
break
91+
92+
tmp_path: Path | None = None
93+
try:
94+
with os.fdopen(fd, "w", encoding="utf-8") as handle:
95+
handle.write(f"pid={os.getpid()}\n")
96+
97+
if _is_up_to_date(source, output):
98+
return output
99+
100+
tmp_fd, tmp_name = tempfile.mkstemp(
101+
dir=output.parent,
102+
prefix=f".{output.name}.",
103+
)
104+
os.close(tmp_fd)
105+
tmp_path = Path(tmp_name)
106+
sp.run(
107+
[
108+
sys.executable,
109+
"-m",
110+
"deepmd",
111+
"convert-from",
112+
"pbtxt",
113+
"-i",
114+
str(source),
115+
"-o",
116+
str(tmp_path),
117+
],
118+
check=True,
119+
)
120+
tmp_path.replace(output)
121+
tmp_path = None
122+
return output
123+
finally:
124+
if tmp_path is not None:
125+
tmp_path.unlink(missing_ok=True)
126+
lock_file.unlink(missing_ok=True)

source/lmp/tests/test_deeptensor.py

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
22
import os
3-
import subprocess as sp
4-
import sys
53
from pathlib import (
64
Path,
75
)
@@ -12,6 +10,9 @@
1210
from lammps import (
1311
PyLammps,
1412
)
13+
from model_convert import (
14+
ensure_converted_pb,
15+
)
1516
from write_lmp_data import (
1617
write_lmp_data,
1718
)
@@ -56,20 +57,14 @@
5657
# type_HO = np.array([2, 1, 1, 2, 1, 1])
5758

5859

59-
sp.check_output(
60-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file.resolve()} -o {pb_file.resolve()}".split()
61-
)
62-
63-
sp.check_output(
64-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
65-
)
66-
67-
6860
def setup_module() -> None:
6961
if os.environ.get("ENABLE_TENSORFLOW", "1") != "1":
7062
pytest.skip(
7163
"Skip test because TensorFlow support is not enabled.",
7264
)
65+
ensure_converted_pb(pbtxt_file, pb_file)
66+
ensure_converted_pb(pbtxt_file2, pb_file2)
67+
7368
write_lmp_data(box, coord, type_OH, data_file)
7469
# TODO
7570
# write_lmp_data(box, coord, type_HO, data_type_map_file)

source/lmp/tests/test_dplr.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
22
import os
3-
import subprocess as sp
4-
import sys
53
from pathlib import (
64
Path,
75
)
@@ -12,6 +10,9 @@
1210
from lammps import (
1311
PyLammps,
1412
)
13+
from model_convert import (
14+
ensure_converted_pb,
15+
)
1516
from write_lmp_data import (
1617
write_lmp_data_full,
1718
)
@@ -265,16 +266,14 @@
265266
mesh = 10
266267

267268

268-
sp.check_output(
269-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file.resolve()} -o {pb_file.resolve()}".split()
270-
)
271-
272-
273269
def setup_module() -> None:
274270
if os.environ.get("ENABLE_TENSORFLOW", "1") != "1":
275271
pytest.skip(
276272
"Skip test because TensorFlow support is not enabled.",
277273
)
274+
ensure_converted_pb(pbtxt_file, pb_file)
275+
ensure_converted_pb(dipole_pbtxt_file, dipole_pb_file)
276+
278277
write_lmp_data_full(
279278
box, coord, mol_list, type_OH, charge, data_file, bond_list, mass_list
280279
)

source/lmp/tests/test_lammps.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
from lammps import (
1616
PyLammps,
1717
)
18+
from model_convert import (
19+
ensure_converted_pb,
20+
)
1821
from write_lmp_data import (
1922
write_lmp_data,
2023
)
@@ -221,19 +224,14 @@
221224
type_HO = np.array([2, 1, 1, 2, 1, 1])
222225

223226

224-
sp.check_output(
225-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file.resolve()} -o {pb_file.resolve()}".split()
226-
)
227-
sp.check_output(
228-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
229-
)
230-
231-
232227
def setup_module() -> None:
233228
if os.environ.get("ENABLE_TENSORFLOW", "1") != "1":
234229
pytest.skip(
235230
"Skip test because TensorFlow support is not enabled.",
236231
)
232+
ensure_converted_pb(pbtxt_file, pb_file)
233+
ensure_converted_pb(pbtxt_file2, pb_file2)
234+
237235
write_lmp_data(box, coord, type_OH, data_file)
238236
write_lmp_data(box, coord, type_HO, data_type_map_file)
239237
write_lmp_data(

source/lmp/tests/test_lammps_3types.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
22
import os
3-
import subprocess as sp
4-
import sys
53
from pathlib import (
64
Path,
75
)
@@ -11,6 +9,9 @@
119
from lammps import (
1210
PyLammps,
1311
)
12+
from model_convert import (
13+
ensure_converted_pb,
14+
)
1415
from write_lmp_data import (
1516
write_lmp_data,
1617
)
@@ -244,19 +245,15 @@
244245
# https://github.com/lammps/lammps/blob/1e1311cf401c5fc2614b5d6d0ff3230642b76597/src/update.cpp#L193
245246
nktv2p = 1.6021765e6
246247

247-
sp.check_output(
248-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file.resolve()} -o {pb_file.resolve()}".split()
249-
)
250-
sp.check_output(
251-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
252-
)
253-
254248

255249
def setup_module() -> None:
256250
if os.environ.get("ENABLE_TENSORFLOW", "1") != "1":
257251
pytest.skip(
258252
"Skip test because TensorFlow support is not enabled.",
259253
)
254+
ensure_converted_pb(pbtxt_file, pb_file)
255+
ensure_converted_pb(pbtxt_file2, pb_file2)
256+
260257
write_lmp_data(box, coord, type_OH, data_file)
261258
write_lmp_data(box, coord, type_HO, data_type_map_file)
262259

source/lmp/tests/test_lammps_dpa_jax.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
from lammps import (
1616
PyLammps,
1717
)
18+
from model_convert import (
19+
ensure_converted_pb,
20+
)
1821
from write_lmp_data import (
1922
write_lmp_data,
2023
)
@@ -222,16 +225,14 @@
222225
type_HO = np.array([2, 1, 1, 2, 1, 1])
223226

224227

225-
sp.check_output(
226-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
227-
)
228-
229-
230228
def setup_module():
231229
if os.environ.get("ENABLE_JAX", "1") != "1":
232230
pytest.skip(
233231
"Skip test because JAX support is not enabled.",
234232
)
233+
if os.environ.get("ENABLE_TENSORFLOW", "1") == "1":
234+
ensure_converted_pb(pbtxt_file2, pb_file2)
235+
235236
write_lmp_data(box, coord, type_OH, data_file)
236237
write_lmp_data(box, coord, type_HO, data_type_map_file)
237238
write_lmp_data(

source/lmp/tests/test_lammps_dpa_pt.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
from lammps import (
1616
PyLammps,
1717
)
18+
from model_convert import (
19+
ensure_converted_pb,
20+
)
1821
from write_lmp_data import (
1922
write_lmp_data,
2023
)
@@ -220,16 +223,14 @@
220223
type_HO = np.array([2, 1, 1, 2, 1, 1])
221224

222225

223-
sp.check_output(
224-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
225-
)
226-
227-
228226
def setup_module() -> None:
229227
if os.environ.get("ENABLE_PYTORCH", "1") != "1":
230228
pytest.skip(
231229
"Skip test because PyTorch support is not enabled.",
232230
)
231+
if os.environ.get("ENABLE_TENSORFLOW", "1") == "1":
232+
ensure_converted_pb(pbtxt_file2, pb_file2)
233+
233234
write_lmp_data(box, coord, type_OH, data_file)
234235
write_lmp_data(box, coord, type_HO, data_type_map_file)
235236
write_lmp_data(

source/lmp/tests/test_lammps_dpa_pt_nopbc.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
from lammps import (
1616
PyLammps,
1717
)
18+
from model_convert import (
19+
ensure_converted_pb,
20+
)
1821
from write_lmp_data import (
1922
write_lmp_data,
2023
)
@@ -218,16 +221,14 @@
218221
type_HO = np.array([2, 1, 1, 2, 1, 1])
219222

220223

221-
sp.check_output(
222-
f"{sys.executable} -m deepmd convert-from pbtxt -i {pbtxt_file2.resolve()} -o {pb_file2.resolve()}".split()
223-
)
224-
225-
226224
def setup_module() -> None:
227225
if os.environ.get("ENABLE_PYTORCH", "1") != "1":
228226
pytest.skip(
229227
"Skip test because PyTorch support is not enabled.",
230228
)
229+
if os.environ.get("ENABLE_TENSORFLOW", "1") == "1":
230+
ensure_converted_pb(pbtxt_file2, pb_file2)
231+
231232
write_lmp_data(box, coord, type_OH, data_file)
232233
write_lmp_data(box, coord, type_HO, data_type_map_file)
233234
write_lmp_data(

0 commit comments

Comments
 (0)