From bf091f85066955b90d3fd664ed6d3b022761f133 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 21 Aug 2026 08:51:26 +0800 Subject: [PATCH 01/14] feat: support DPA4 and DPA4C deployment --- docs/input.md | 28 ++++++ dpgen2/constants.py | 3 +- dpgen2/op/__init__.py | 1 + dpgen2/op/run_dp_train.py | 15 ++- dpgen2/op/run_lmp.py | 176 +++++++++++++++++++++++++++++++-- dpgen2/superop/prep_run_lmp.py | 29 +++++- tests/op/test_run_dp_train.py | 27 +++++ tests/op/test_run_lmp.py | 87 ++++++++++++++++ 8 files changed, 351 insertions(+), 15 deletions(-) diff --git a/docs/input.md b/docs/input.md index 6fc03bd1..f83aab2d 100644 --- a/docs/input.md +++ b/docs/input.md @@ -59,6 +59,34 @@ The `"type" : "dp"` tell the traning method is {dargs:argument}`"dp" `, i The `"config"` key defines the training configs, see {ref}`the full documentation`. The {dargs:argument}`"template_script" ` provides the template training script in `json` format. +For DPA4, use the regular PyTorch training backend and deploy `.pt2` models in +LAMMPS exploration: + +```json +"train": {"type": "dp", "config": {"impl": "pytorch"}}, +"explore": { + "type": "lmp", + "config": {"model_devi_backend": "pytorch", "model_format": "pt2"} +} +``` + +For DPA4C, both training and deployment use the PyTorch Exportable backend. +Compression is optional: + +```json +"train": {"type": "dp", "config": {"impl": "pytorch-exportable"}}, +"explore": { + "type": "lmp", + "config": { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + "dp_compress": true + } +} +``` + +TensorFlow remains the default when `impl` is omitted. + ### Exploration diff --git a/dpgen2/constants.py b/dpgen2/constants.py index 6d5d0197..7647ef16 100644 --- a/dpgen2/constants.py +++ b/dpgen2/constants.py @@ -4,7 +4,8 @@ train_log_name = "train.log" model_name_pattern = "model.%03d.pb" pytorch_model_name_pattern = "model.%03d.pth" -model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth)" +pt2_model_name_pattern = "model.%03d.pt2" +model_name_match_pattern = r"model\.[0-9]{3,}(\.pb|\.pth|\.pt2)" lmp_index_pattern = "%06d" lmp_task_pattern = "task." + lmp_index_pattern lmp_conf_name = "conf.lmp" diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index f4fec3a2..015e9c58 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -35,6 +35,7 @@ RunDPTrain, ) from .run_lmp import ( + PrepareDPModels, RunLmp, RunLmpHDF5, ) diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 5a9782f4..518de9f4 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -61,7 +61,10 @@ def _make_train_command( # find checkpoint if impl == "tensorflow" and os.path.isfile("checkpoint"): checkpoint = "model.ckpt" - elif impl == "pytorch" and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0: + elif ( + impl in ["pytorch", "pytorch-exportable"] + and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0 + ): checkpoint = "model.ckpt-%s.pt" % max( [int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")] ) @@ -184,10 +187,14 @@ def execute( finetune_mode = ip["optional_parameter"]["finetune_mode"] config = ip["config"] if ip["config"] is not None else {} impl = ip["config"].get("impl", "tensorflow") + if impl == "pt-expt": + impl = "pytorch-exportable" dp_command = ip["config"].get("command", "dp").split() - assert impl in ["tensorflow", "pytorch"] + assert impl in ["tensorflow", "pytorch", "pytorch-exportable"] if impl == "pytorch": dp_command.append("--pt") + elif impl == "pytorch-exportable": + dp_command.append("--pt-expt") finetune_args = config.get("finetune_args", "") train_args = config.get("train_args", "") config = RunDPTrain.normalize_config(config) @@ -318,7 +325,7 @@ def clean_before_quit(): shutil.copy2("input_v2_compat.json", train_script_name) # freeze model - if impl == "pytorch": + if impl in ["pytorch", "pytorch-exportable"]: model_file = "model.ckpt.pt" else: ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"]) @@ -503,7 +510,7 @@ def decide_init_model( @staticmethod def training_args(): doc_command = "The command for DP, 'dp' for default" - doc_impl = "The implementation/backend of DP. It can be 'tensorflow' or 'pytorch'. 'tensorflow' for default." + doc_impl = "The implementation/backend of DP. It can be 'tensorflow', 'pytorch', or 'pytorch-exportable' (alias 'pt-expt'). 'tensorflow' for default." doc_init_model_policy = "The policy of init-model training. It can be\n\n\ - 'no': No init-model training. Traing from scratch.\n\n\ - 'yes': Do init-model training.\n\n\ diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 60cd9305..7fe49b90 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -41,6 +41,7 @@ model_name_match_pattern, model_name_pattern, plm_output_name, + pt2_model_name_pattern, pytorch_model_name_pattern, ) from dpgen2.utils import ( @@ -52,6 +53,35 @@ ) +_MODEL_BACKEND_ALIASES = {"pt-expt": "pytorch-exportable"} +_MODEL_BACKEND_FLAGS = { + "pytorch": "--pt", + "pytorch-exportable": "--pt-expt", +} + + +class PrepareDPModels(OP): + """Freeze DP checkpoints once before exploration tasks are fanned out.""" + + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "config": BigParameter(dict), + "models": Artifact(List[Path]), + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign({"models": Artifact(List[Path])}) + + @OP.exec_sign_check + def execute(self, ip: OPIO) -> OPIO: + config = RunLmp.normalize_config(ip["config"] or {}) + return OPIO({"models": prepare_dp_models(ip["models"], config)}) + + class RunLmp(OP): r"""Execute a LAMMPS task. @@ -132,9 +162,9 @@ def execute( work_dir = Path(task_name) if teacher_model is not None: - assert ( - len(model_files) == 1 - ), "One model is enough in knowledge distillation" + assert len(model_files) == 1, ( + "One model is enough in knowledge distillation" + ) ext = os.path.splitext(teacher_model.file_name)[-1] teacher_model_file = "teacher_model" + ext teacher_model.save_as_file(teacher_model_file) @@ -152,10 +182,26 @@ def execute( if ext == ".pb": mname = model_name_pattern % (idx) Path(mname).symlink_to(mm) + elif ext == ".pth": + mname = pytorch_model_name_pattern % (idx) + Path(mname).symlink_to(mm) + elif ext == ".pt2": + mname = pt2_model_name_pattern % (idx) + Path(mname).symlink_to(mm) elif ext == ".pt": # freeze model - mname = pytorch_model_name_pattern % (idx) - freeze_model(mm, mname, config.get("model_frozen_head")) + backend = _model_backend(config) + mname = _model_name(idx, config["model_format"]) + freeze_model( + mm, + mname, + config.get("model_frozen_head"), + backend, + ) + if config["dp_compress"]: + compressed = _compressed_model_name(idx, config["model_format"]) + compress_model(mname, compressed, backend) + mname = compressed else: raise RuntimeError( "Model file with extension '%s' is not supported" % ext @@ -232,6 +278,11 @@ def lmp_args(): doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature" doc_use_hdf5 = "Use HDF5 to store trajs and model_devis" doc_extra_output_files = "Extra output file names, support wildcards" + doc_model_devi_backend = ( + "The DeePMD backend used to freeze models for exploration" + ) + doc_model_format = "The frozen model format. Use 'pt2' for DPA4 and DPA4C" + doc_dp_compress = "Compress the frozen model before exploration" return [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), Argument( @@ -269,6 +320,27 @@ def lmp_args(): default=[], doc=doc_extra_output_files, ), + Argument( + "model_devi_backend", + str, + optional=True, + default="pytorch", + doc=doc_model_devi_backend, + ), + Argument( + "model_format", + str, + optional=True, + default="pth", + doc=doc_model_format, + ), + Argument( + "dp_compress", + bool, + optional=True, + default=False, + doc=doc_dp_compress, + ), ] @staticmethod @@ -306,7 +378,7 @@ def set_models(lmp_input_name: str, model_names: List[str]): break if match_first == -1: raise RuntimeError( - f"cannot file model pattern {pattern} in line " f" {lmp_input_lines[idx]}" + f"cannot file model pattern {pattern} in line {lmp_input_lines[idx]}" ) if match_last == -1: raise RuntimeError(f"last matching index should not be -1, terribly wrong ") @@ -362,11 +434,81 @@ def get_ele_temp(lmp_log_name): return None -def freeze_model(input_model, frozen_model, head=None): +def _model_backend(config): + backend = _MODEL_BACKEND_ALIASES.get( + config["model_devi_backend"], config["model_devi_backend"] + ) + model_format = config["model_format"] + if backend not in _MODEL_BACKEND_FLAGS: + raise RuntimeError(f"Unsupported model-deviation backend '{backend}'") + if model_format not in ["pth", "pt2"]: + raise RuntimeError(f"Unsupported model format '{model_format}'") + if model_format == "pth" and backend != "pytorch": + raise RuntimeError("The pth model format requires the pytorch backend") + if config["dp_compress"] and not ( + backend == "pytorch-exportable" and model_format == "pt2" + ): + raise RuntimeError( + "Compressed pt2 models require the pytorch-exportable backend" + ) + return backend + + +def _model_name(index, model_format): + if model_format == "pt2": + return pt2_model_name_pattern % index + return pytorch_model_name_pattern % index + + +def _compressed_model_name(index, model_format): + return "model.%03d.compressed.%s" % (index, model_format) + + +def prepare_dp_models(models, config): + """Return frozen models, exporting checkpoints once when needed.""" + backend = _model_backend(config) + prepared = [] + output_dir = Path("prepared_models") + for idx, model in enumerate(models): + model = Path(model).resolve() + ext = model.suffix + if ext != ".pt": + if ext not in [".pb", ".pth", ".pt2"]: + raise RuntimeError( + "Model file with extension '%s' is not supported" % ext + ) + prepared.append(model) + continue + output_dir.mkdir(exist_ok=True) + frozen_model = output_dir / _model_name(idx, config["model_format"]) + freeze_model( + model, + frozen_model, + config.get("model_frozen_head"), + backend, + ) + if config["dp_compress"]: + compressed_model = output_dir / _compressed_model_name( + idx, config["model_format"] + ) + compress_model(frozen_model, compressed_model, backend) + frozen_model = compressed_model + prepared.append(frozen_model) + return prepared + + +def freeze_model(input_model, frozen_model, head=None, backend="pytorch"): + backend = _MODEL_BACKEND_ALIASES.get(backend, backend) freeze_args = "-o %s" % frozen_model if head is not None: freeze_args += " --head %s" % head - freeze_cmd = "dp --pt freeze -c %s %s" % (input_model, freeze_args) + if backend == "pytorch-exportable" and Path(frozen_model).suffix == ".pt2": + freeze_args += " --lower-kind graph" + freeze_cmd = "dp %s freeze -c %s %s" % ( + _MODEL_BACKEND_FLAGS[backend], + input_model, + freeze_args, + ) ret, out, err = run_command(freeze_cmd, shell=True) if ret != 0: logging.error( @@ -387,6 +529,24 @@ def freeze_model(input_model, frozen_model, head=None): raise TransientError("freeze failed") +def compress_model(input_model, output_model, backend="pytorch-exportable"): + backend = _MODEL_BACKEND_ALIASES.get(backend, backend) + compress_cmd = "dp %s compress -i %s -o %s" % ( + _MODEL_BACKEND_FLAGS[backend], + input_model, + output_model, + ) + ret, out, err = run_command(compress_cmd, shell=True) + if ret != 0: + logging.error( + "compress failed\ncommand was%s\nout msg%s\nerr msg%s\n", + compress_cmd, + out, + err, + ) + raise TransientError("compress failed") + + def merge_pimd_files(): traj_files = glob.glob("traj.*.dump") if len(traj_files) > 0: diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 3e0a0a0f..77402294 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -40,6 +40,9 @@ from dpgen2.constants import ( lmp_index_pattern, ) +from dpgen2.op.run_lmp import ( + PrepareDPModels, +) from dpgen2.utils.step_config import ( init_executor, ) @@ -91,10 +94,12 @@ def __init__( ), ) - self._keys = ["prep-lmp", "run-lmp"] + self._keys = ["prep-lmp", "prepare-models", "run-lmp"] self.step_keys = {} ii = "prep-lmp" self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + ii = "prepare-models" + self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) ii = "run-lmp" self.step_keys[ii] = "--".join( ["%s" % self.inputs.parameters["block_id"], ii + "-{{item}}"] @@ -166,6 +171,26 @@ def _prep_run_lmp( ) prep_run_steps.add(prep_lmp) + prepare_models = Step( + "prepare-models", + template=PythonOPTemplate( + PrepareDPModels, + output_artifact_archive={"models": None}, + python_packages=upload_python_packages, + **run_template_config, + ), + parameters={ + "config": prep_run_steps.inputs.parameters["explore_config"], + }, + artifacts={ + "models": prep_run_steps.inputs.artifacts["models"], + }, + key=step_keys["prepare-models"], + executor=run_executor, + **run_config, + ) + prep_run_steps.add(prepare_models) + run_lmp = Step( "run-lmp", template=PythonOPTemplate( @@ -193,7 +218,7 @@ def _prep_run_lmp( }, artifacts={ "task_path": prep_lmp.outputs.artifacts["task_paths"], - "models": prep_run_steps.inputs.artifacts["models"], + "models": prepare_models.outputs.artifacts["models"], }, with_sequence=argo_sequence( argo_len(prep_lmp.outputs.parameters["task_names"]), diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 7649b520..cc56e1e4 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -434,6 +434,33 @@ def test_exec_v1(self, mocked_run): jdata = json.load(fp) self.assertDictEqual(jdata, self.expected_odict_v1) + @patch("dpgen2.op.run_dp_train.run_command") + def test_exec_pytorch_exportable(self, mocked_run): + mocked_run.return_value = (0, "foo\n", "") + config = self.config.copy() + config.update({"impl": "pt-expt", "init_model_policy": "no"}) + Path(self.task_path).mkdir(exist_ok=True) + with open(Path(self.task_path) / train_script_name, "w") as fp: + json.dump(self.idict_v2, fp, indent=4) + + out = RunDPTrain().execute( + OPIO( + { + "config": config, + "task_name": self.task_name, + "task_path": Path(self.task_path), + "init_model": Path(self.init_model), + "init_data": [Path(ii) for ii in self.init_data], + "iter_data": [Path(ii) for ii in self.iter_data], + } + ) + ) + + self.assertEqual(out["model"], Path(self.task_name) / "model.ckpt.pt") + mocked_run.assert_called_once_with( + ["dp", "--pt-expt", "train", train_script_name] + ) + @patch("dpgen2.op.run_dp_train.run_command") def test_exec_v2(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", ""), (0, "bar\n", "")] diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 650fd82e..45f05dbe 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -32,8 +32,10 @@ lmp_model_devi_name, lmp_traj_name, model_name_pattern, + pt2_model_name_pattern, ) from dpgen2.op.run_lmp import ( + PrepareDPModels, RunLmp, get_ele_temp, merge_pimd_files, @@ -258,6 +260,81 @@ def swap_element(arg): arg[0] = bk[1] +class TestPrepareDPModels(unittest.TestCase): + def setUp(self): + self.model_dir = Path("checkpoint_models") + self.model_dir.mkdir() + self.models = [] + for idx in range(2): + model = self.model_dir / f"model.{idx}.pt" + model.write_text("checkpoint") + self.models.append(model) + + def tearDown(self): + shutil.rmtree(self.model_dir, ignore_errors=True) + shutil.rmtree("prepared_models", ignore_errors=True) + + @patch("dpgen2.op.run_lmp.run_command") + def test_dpa4_pt2(self, mocked_run): + mocked_run.return_value = (0, "", "") + models = PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pytorch", + "model_format": "pt2", + }, + "models": self.models, + } + ) + )["models"] + self.assertEqual( + models, + [ + Path("prepared_models/model.000.pt2"), + Path("prepared_models/model.001.pt2"), + ], + ) + mocked_run.assert_has_calls( + [ + call( + f"dp --pt freeze -c {model.resolve()} -o {Path('prepared_models') / f'model.{idx:03d}.pt2'}", + shell=True, + ) + for idx, model in enumerate(self.models) + ] + ) + + @patch("dpgen2.op.run_lmp.run_command") + def test_dpa4c_compressed_pt2(self, mocked_run): + mocked_run.return_value = (0, "", "") + models = PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pt-expt", + "model_format": "pt2", + "dp_compress": True, + }, + "models": self.models[:1], + } + ) + )["models"] + self.assertEqual(models, [Path("prepared_models/model.000.compressed.pt2")]) + mocked_run.assert_has_calls( + [ + call( + f"dp --pt-expt freeze -c {self.models[0].resolve()} -o {Path('prepared_models/model.000.pt2')} --lower-kind graph", + shell=True, + ), + call( + f"dp --pt-expt compress -i {Path('prepared_models/model.000.pt2')} -o {Path('prepared_models/model.000.compressed.pt2')}", + shell=True, + ), + ] + ) + + class TestSetModels(unittest.TestCase): def setUp(self): self.input_name = Path("lmp.input") @@ -274,6 +351,16 @@ def test(self): set_models(input_name, self.model_names) self.assertEqual(input_name.read_text(), expected_output) + def test_pt2(self): + lmp_config = "pair_style deepmd model.000.pb model.001.pb out_freq 10\n" + expected_output = "pair_style deepmd model.000.pt2 model.001.pt2 out_freq 10\n" + self.input_name.write_text(lmp_config) + set_models( + self.input_name, + [pt2_model_name_pattern % 0, pt2_model_name_pattern % 1], + ) + self.assertEqual(self.input_name.read_text(), expected_output) + def test_failed(self): lmp_config = "pair_style deepmd model.000.pb model.001.pb out_freq 10 out_file model_devi.out model.002.pb\n" input_name = self.input_name From 454d1886c9e2266b0b38f5df048e445ec93f1ba9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:52:23 +0000 Subject: [PATCH 02/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- dpgen2/op/run_lmp.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 7fe49b90..2b82628d 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -52,7 +52,6 @@ run_command, ) - _MODEL_BACKEND_ALIASES = {"pt-expt": "pytorch-exportable"} _MODEL_BACKEND_FLAGS = { "pytorch": "--pt", @@ -162,9 +161,9 @@ def execute( work_dir = Path(task_name) if teacher_model is not None: - assert len(model_files) == 1, ( - "One model is enough in knowledge distillation" - ) + assert ( + len(model_files) == 1 + ), "One model is enough in knowledge distillation" ext = os.path.splitext(teacher_model.file_name)[-1] teacher_model_file = "teacher_model" + ext teacher_model.save_as_file(teacher_model_file) From aab734b345012895f6ba981c95d89819b71f1a24 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 21 Aug 2026 08:57:36 +0800 Subject: [PATCH 03/14] test: account for model preparation step --- tests/test_block_cl.py | 2 ++ tests/test_dpgen_loop.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py index fa0571a7..f1392984 100644 --- a/tests/test_block_cl.py +++ b/tests/test_block_cl.py @@ -225,6 +225,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", @@ -405,6 +406,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py index 49b4c873..bf49cd9d 100644 --- a/tests/test_dpgen_loop.py +++ b/tests/test_dpgen_loop.py @@ -270,6 +270,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", @@ -502,6 +503,7 @@ def test(self): "prep-train", "run-train", "prep-lmp", + "prepare-models", "run-lmp", "select-confs", "prep-fp", From 22ef000697ec9f93625e89c96c931d0346832a99 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 21 Aug 2026 09:32:01 +0800 Subject: [PATCH 04/14] fix: address deployment review findings --- dpgen2/op/run_dp_train.py | 8 +++--- dpgen2/op/run_lmp.py | 52 +++++++++++++++++------------------ tests/op/test_run_dp_train.py | 6 ++++ tests/op/test_run_lmp.py | 35 +++++++++++++++++++---- 4 files changed, 64 insertions(+), 37 deletions(-) diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 518de9f4..e660603d 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -346,10 +346,10 @@ def clean_before_quit(): ) raise FatalError("dp freeze failed") model_file = "frozen_model.pb" - fplog.write("#=================== freeze std out ===================\n") - fplog.write(out) - fplog.write("#=================== freeze std err ===================\n") - fplog.write(err) + fplog.write("#=================== freeze std out ===================\n") + fplog.write(out) + fplog.write("#=================== freeze std err ===================\n") + fplog.write(err) clean_before_quit() diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 2b82628d..b49a4401 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -498,44 +498,42 @@ def prepare_dp_models(models, config): def freeze_model(input_model, frozen_model, head=None, backend="pytorch"): backend = _MODEL_BACKEND_ALIASES.get(backend, backend) - freeze_args = "-o %s" % frozen_model + freeze_cmd = [ + "dp", + _MODEL_BACKEND_FLAGS[backend], + "freeze", + "-c", + str(input_model), + "-o", + str(frozen_model), + ] if head is not None: - freeze_args += " --head %s" % head + freeze_cmd.extend(["--head", str(head)]) if backend == "pytorch-exportable" and Path(frozen_model).suffix == ".pt2": - freeze_args += " --lower-kind graph" - freeze_cmd = "dp %s freeze -c %s %s" % ( - _MODEL_BACKEND_FLAGS[backend], - input_model, - freeze_args, - ) - ret, out, err = run_command(freeze_cmd, shell=True) + freeze_cmd.extend(["--lower-kind", "graph"]) + ret, out, err = run_command(freeze_cmd) if ret != 0: logging.error( - "".join( - ( - "freeze failed\n", - "command was", - freeze_cmd, - "out msg", - out, - "\n", - "err msg", - err, - "\n", - ) - ) + "freeze failed\ncommand was %s\nout msg%s\nerr msg%s\n", + freeze_cmd, + out, + err, ) raise TransientError("freeze failed") def compress_model(input_model, output_model, backend="pytorch-exportable"): backend = _MODEL_BACKEND_ALIASES.get(backend, backend) - compress_cmd = "dp %s compress -i %s -o %s" % ( + compress_cmd = [ + "dp", _MODEL_BACKEND_FLAGS[backend], - input_model, - output_model, - ) - ret, out, err = run_command(compress_cmd, shell=True) + "compress", + "-i", + str(input_model), + "-o", + str(output_model), + ] + ret, out, err = run_command(compress_cmd) if ret != 0: logging.error( "compress failed\ncommand was%s\nout msg%s\nerr msg%s\n", diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index cc56e1e4..d9a03800 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -460,6 +460,12 @@ def test_exec_pytorch_exportable(self, mocked_run): mocked_run.assert_called_once_with( ["dp", "--pt-expt", "train", train_script_name] ) + self.assertEqual( + out["log"].read_text(), + "#=================== train std out ===================\n" + "foo\n" + "#=================== train std err ===================\n", + ) @patch("dpgen2.op.run_dp_train.run_command") def test_exec_v2(self, mocked_run): diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 45f05dbe..45469213 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -298,8 +298,15 @@ def test_dpa4_pt2(self, mocked_run): mocked_run.assert_has_calls( [ call( - f"dp --pt freeze -c {model.resolve()} -o {Path('prepared_models') / f'model.{idx:03d}.pt2'}", - shell=True, + [ + "dp", + "--pt", + "freeze", + "-c", + str(model.resolve()), + "-o", + str(Path("prepared_models") / f"model.{idx:03d}.pt2"), + ] ) for idx, model in enumerate(self.models) ] @@ -324,12 +331,28 @@ def test_dpa4c_compressed_pt2(self, mocked_run): mocked_run.assert_has_calls( [ call( - f"dp --pt-expt freeze -c {self.models[0].resolve()} -o {Path('prepared_models/model.000.pt2')} --lower-kind graph", - shell=True, + [ + "dp", + "--pt-expt", + "freeze", + "-c", + str(self.models[0].resolve()), + "-o", + str(Path("prepared_models/model.000.pt2")), + "--lower-kind", + "graph", + ] ), call( - f"dp --pt-expt compress -i {Path('prepared_models/model.000.pt2')} -o {Path('prepared_models/model.000.compressed.pt2')}", - shell=True, + [ + "dp", + "--pt-expt", + "compress", + "-i", + str(Path("prepared_models/model.000.pt2")), + "-o", + str(Path("prepared_models/model.000.compressed.pt2")), + ] ), ] ) From 45b09a03a27d3280a25a543aec8bfdd0af48b135 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 13:41:58 +0800 Subject: [PATCH 05/14] fix: prepare pt2 exploration inputs safely --- docs/input.md | 5 +++ dpgen2/entrypoint/submit.py | 3 ++ dpgen2/op/__init__.py | 1 + dpgen2/op/run_lmp.py | 79 +++++++++++++++++++++++++++++++++++++ tests/op/test_run_lmp.py | 52 ++++++++++++++++++++++++ 5 files changed, 140 insertions(+) diff --git a/docs/input.md b/docs/input.md index f83aab2d..026b1a2c 100644 --- a/docs/input.md +++ b/docs/input.md @@ -86,6 +86,11 @@ Compression is optional: ``` TensorFlow remains the default when `impl` is omitted. +For PyTorch and PyTorch Exportable checkpoints, `model_devi_backend` must match +the training `impl`; checkpoints cannot be frozen across these backends. PT2 +export runs with `run_explore_config`, which must select hardware, the libtorch +version, and other runtime libraries compatible with the LAMMPS exploration +environment. ### Exploration diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 9e07374f..737b2ca2 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -107,6 +107,7 @@ RunRelax, RunRelaxHDF5, SelectConfs, + validate_model_backend, ) from dpgen2.op.caly_evo_step_merge import ( CalyEvoStepMerge, @@ -472,6 +473,8 @@ def workflow_concurrent_learning( train_style = config["train"]["type"] explore_style = config["explore"]["type"] fp_style = config["fp"]["type"] + if train_style in ["dp", "dp-dist"] and explore_style == "lmp": + validate_model_backend(train_config["impl"], explore_config) prep_train_config = config["step_configs"]["prep_train_config"] run_train_config = config["step_configs"]["run_train_config"] prep_explore_config = config["step_configs"]["prep_explore_config"] diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index 015e9c58..74cbd5bc 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -38,6 +38,7 @@ PrepareDPModels, RunLmp, RunLmpHDF5, + validate_model_backend, ) from .run_relax import ( RunRelax, diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index b49a4401..5f474f54 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -211,6 +211,8 @@ def execute( random.shuffle(model_names) set_models(lmp_input_name, model_names) + if any(Path(name).suffix == ".pt2" for name in model_names): + ensure_pt2_atom_map(lmp_input_name) # run lmp command = " ".join([command, "-i", lmp_input_name, "-log", lmp_log_name]) @@ -394,6 +396,56 @@ def set_models(lmp_input_name: str, model_names: List[str]): f.write("".join(lmp_input_lines)) +def ensure_pt2_atom_map(lmp_input_name: str): + """Ensure a PT2 LAMMPS input enables the atom map before reading atoms. + + Parameters + ---------- + lmp_input_name : str + Path to the LAMMPS input file. + + Raises + ------ + RuntimeError + If an existing ``atom_modify map yes`` follows the first + ``read_data`` or ``read_restart`` command, or neither read command is + present. + """ + with open(lmp_input_name, encoding="utf8") as f: + lmp_input_lines = f.readlines() + + read_index = next( + ( + index + for index, line in enumerate(lmp_input_lines) + if re.search(r"\bread_(?:data|restart)\b", line.partition("#")[0]) + ), + None, + ) + if read_index is None: + raise RuntimeError("PT2 LAMMPS inputs require read_data or read_restart") + + atom_map_index = next( + ( + index + for index, line in enumerate(lmp_input_lines) + if re.match(r"^\s*atom_modify\s+.*\bmap\s+yes\b", line.partition("#")[0]) + ), + None, + ) + if atom_map_index is not None: + if atom_map_index > read_index: + raise RuntimeError( + "PT2 LAMMPS inputs require 'atom_modify map yes' before " + "read_data or read_restart" + ) + return + + lmp_input_lines.insert(read_index, "atom_modify map yes\n") + with open(lmp_input_name, "w", encoding="utf8") as f: + f.write("".join(lmp_input_lines)) + + def find_only_one_key(lmp_lines, key, raise_not_found=True): found = [] for idx in range(len(lmp_lines)): @@ -453,6 +505,33 @@ def _model_backend(config): return backend +def validate_model_backend(train_backend, config): + """Validate that a PyTorch checkpoint is frozen by its training backend. + + Parameters + ---------- + train_backend : str + DeePMD training backend. + config : dict + LAMMPS exploration configuration. + + Raises + ------ + RuntimeError + If PyTorch training and deployment backends differ. + """ + train_backend = _MODEL_BACKEND_ALIASES.get(train_backend, train_backend) + if train_backend not in _MODEL_BACKEND_FLAGS: + return + model_backend = _model_backend(RunLmp.normalize_config(config)) + if model_backend != train_backend: + raise RuntimeError( + f"The model-deviation backend '{model_backend}' cannot freeze a " + f"checkpoint trained by '{train_backend}'; use the same backend " + "for training and model deployment" + ) + + def _model_name(index, model_format): if model_format == "pt2": return pt2_model_name_pattern % index diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 45469213..77618326 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -40,6 +40,7 @@ get_ele_temp, merge_pimd_files, set_models, + validate_model_backend, ) from dpgen2.utils import ( BinaryFileInput, @@ -104,6 +105,40 @@ def test_success(self, mocked_run): (work_dir / (model_name_pattern % ii)).read_text(), f"model{ii}" ) + @patch("dpgen2.op.run_lmp.run_command") + def test_pt2_enables_atom_map_before_read(self, mocked_run): + mocked_run.return_value = (0, "", "") + (self.task_path / lmp_input_name).write_text( + "atom_style atomic\n" + 'if "${restart} > 0" then "read_restart dpgen.restart.*" ' + 'else "read_data conf.lmp"\n' + "pair_style deepmd model.000.pb model.001.pb out_freq 10\n" + ) + models = [self.model_path / f"model_{index}.pt2" for index in range(2)] + for model in models: + model.write_text("model") + + def copy_link(source, target, target_is_directory=False): + shutil.copyfile(source, target) + + with patch("os.symlink", side_effect=copy_link): + RunLmp().execute( + OPIO( + { + "config": {"command": "mylmp"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": models, + } + ) + ) + + lmp_input = (Path(self.task_name) / lmp_input_name).read_text() + atom_map = "atom_modify map yes" + self.assertEqual(lmp_input.count(atom_map), 1) + self.assertLess(lmp_input.index(atom_map), lmp_input.index("read_restart")) + self.assertLess(lmp_input.index(atom_map), lmp_input.index("read_data")) + @patch("dpgen2.op.run_lmp.run_command") def test_error(self, mocked_run): mocked_run.side_effect = [(1, "foo\n", "")] @@ -357,6 +392,23 @@ def test_dpa4c_compressed_pt2(self, mocked_run): ] ) + def test_training_and_deployment_backends_must_match(self): + with self.assertRaisesRegex(RuntimeError, "cannot freeze a checkpoint"): + validate_model_backend( + "pytorch", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + ) + validate_model_backend( + "pt-expt", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + ) + class TestSetModels(unittest.TestCase): def setUp(self): From 0aec2edbfe98a37031c78743ac5e1f506a3dadf9 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 13:48:05 +0800 Subject: [PATCH 06/14] fix: preserve default training backend --- dpgen2/entrypoint/submit.py | 2 +- tests/entrypoint/test_submit_args.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 737b2ca2..5fb99d7b 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -474,7 +474,7 @@ def workflow_concurrent_learning( explore_style = config["explore"]["type"] fp_style = config["fp"]["type"] if train_style in ["dp", "dp-dist"] and explore_style == "lmp": - validate_model_backend(train_config["impl"], explore_config) + validate_model_backend(train_config.get("impl", "tensorflow"), explore_config) prep_train_config = config["step_configs"]["prep_train_config"] run_train_config = config["step_configs"]["run_train_config"] prep_explore_config = config["step_configs"]["prep_explore_config"] diff --git a/tests/entrypoint/test_submit_args.py b/tests/entrypoint/test_submit_args.py index ec7e4582..518a154c 100644 --- a/tests/entrypoint/test_submit_args.py +++ b/tests/entrypoint/test_submit_args.py @@ -126,6 +126,7 @@ def test(self): # self.assertEqual(old_data['default_training_param'], new_data['train']['template_script']) self.assertEqual(new_data["train"]["template_script"], "dp_input_template") self.assertEqual(RunDPTrain.normalize_config({}), new_data["train"]["config"]) + self.assertEqual(new_data["train"]["config"]["impl"], "tensorflow") self.assertEqual( RunLmp.normalize_config(old_data.get("lmp_config", {})), new_data["explore"]["config"], From f8a9fab2e335316fd78cfbed854385e4245c8e41 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 13:36:53 +0000 Subject: [PATCH 07/14] fix: validate DPA4 backend-specific training config --- docs/input.md | 36 +++++++++++ dpgen2/entrypoint/args.py | 2 + dpgen2/entrypoint/submit.py | 99 ++++++++++++++++++++++++++++--- dpgen2/superop/prep_run_lmp.py | 6 +- tests/entrypoint/test_submit.py | 90 ++++++++++++++++++++++++++++ tests/test_prep_run_lmp_config.py | 48 +++++++++++++++ 6 files changed, 273 insertions(+), 8 deletions(-) create mode 100644 tests/test_prep_run_lmp_config.py diff --git a/docs/input.md b/docs/input.md index 026b1a2c..aa52c38b 100644 --- a/docs/input.md +++ b/docs/input.md @@ -70,6 +70,17 @@ LAMMPS exploration: } ``` +Put DPA4 training acceleration controls in the DeePMD training template under +`model`, not in the DPGEN2 workflow file: + +```json +"model": { + "type": "dpa4", + "use_compile": true, + "enable_tf32": true +} +``` + For DPA4C, both training and deployment use the PyTorch Exportable backend. Compression is optional: @@ -85,6 +96,31 @@ Compression is optional: } ``` +Put DPA4C training acceleration controls in the DeePMD training template file +referenced by `train.template_script`, under `training`: + +```json +"model": {"descriptor": {"type": "dpa4c"}}, +"training": { + "training_data": { + "systems": [], + "batch_size": "auto:512" + }, + "numb_steps": 1000000, + "enable_compile": true, + "enable_tf32": true +} +``` + +Do not copy the DPA4 paths `model.use_compile` or `model.enable_tf32` into a +DPA4C template. Conversely, DPA4 does not use the DPA4C paths +`training.enable_compile` or `training.enable_tf32`. DPGEN2 validates these +backend-specific placements before creating the workflow, but it does not inject +or change performance and numerical-policy settings. Therefore setting only +`train.config.impl` is not sufficient to enable compilation: the DeePMD template +must contain `training.enable_compile: true` explicitly. Run a bounded smoke test +and inspect the generated `task.*/input.json` before launching a long campaign. + TensorFlow remains the default when `impl` is omitted. For PyTorch and PyTorch Exportable checkpoints, `model_devi_backend` must match the training `impl`; checkpoints cannot be frozen across these backends. PT2 diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..62f543d8 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -45,6 +45,7 @@ def make_link(content, ref_key): def dp_dist_train_args(): + doc_numb_models = "Number of student models trained for model deviation" doc_config = "Configuration of training" doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. " dock_student_model_path = "The path of student model" @@ -63,6 +64,7 @@ def dp_dist_train_args(): Argument( "template_script", [List[str], str], optional=False, doc=doc_template_script ), + Argument("numb_models", int, optional=True, default=1, doc=doc_numb_models), Argument("student_model_path", str, optional=True, doc=dock_student_model_path), Argument( "student_model_uri", diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 5fb99d7b..a73c657c 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -463,6 +463,82 @@ def get_systems_from_data(data, data_prefix=None): return data +def _normalize_model_backend(backend: str) -> str: + return {"pt-expt": "pytorch-exportable"}.get(backend, backend) + + +def _iter_model_sections(template_script: dict): + model = template_script.get("model", {}) + yield "model", model + for name, branch in (model.get("model_dict") or {}).items(): + yield f"model.model_dict.{name}", branch + + +def _model_family(template_script: dict) -> Optional[str]: + for _, model in _iter_model_sections(template_script): + model_type = model.get("type") + descriptor = model.get("descriptor", {}) + descriptor_type = descriptor.get("type") if isinstance(descriptor, dict) else None + model_type = model_type.lower() if isinstance(model_type, str) else model_type + descriptor_type = ( + descriptor_type.lower() + if isinstance(descriptor_type, str) + else descriptor_type + ) + if descriptor_type == "dpa4c": + return "dpa4c" + if model_type == "dpa4" or descriptor_type in {"dpa4", "sezm"}: + return "dpa4" + return None + + +def validate_dpa_training_template( + train_backend: str, + explore_config: dict, + template_script: dict, +) -> None: + """Validate DPA4/DPA4C backend and compile-option placement.""" + family = _model_family(template_script) + if family is None: + return + + train_backend = _normalize_model_backend(train_backend) + expected_backend = "pytorch" if family == "dpa4" else "pytorch-exportable" + if train_backend != expected_backend: + raise RuntimeError( + f"{family.upper()} training requires impl='{expected_backend}', " + f"not '{train_backend}'" + ) + + normalized_explore = RunLmp.normalize_config(explore_config) + if normalized_explore["model_format"] != "pt2": + raise RuntimeError(f"{family.upper()} LAMMPS exploration requires model_format='pt2'") + + misplaced = [] + if family == "dpa4c": + for scope, model in _iter_model_sections(template_script): + for key in ("use_compile", "enable_tf32"): + if key in model: + misplaced.append(f"{scope}.{key}") + if misplaced: + raise RuntimeError( + "DPA4C uses training.enable_compile and training.enable_tf32; " + f"remove misplaced {', '.join(misplaced)}" + ) + else: + training = template_script.get("training", {}) + misplaced = [ + f"training.{key}" + for key in ("enable_compile", "enable_tf32") + if key in training + ] + if misplaced: + raise RuntimeError( + "DPA4 uses model.use_compile and model.enable_tf32; " + f"remove misplaced {', '.join(misplaced)}" + ) + + def workflow_concurrent_learning( config: Dict, ) -> Step: @@ -473,8 +549,17 @@ def workflow_concurrent_learning( train_style = config["train"]["type"] explore_style = config["explore"]["type"] fp_style = config["fp"]["type"] + template_script_ = config["train"]["template_script"] + if isinstance(template_script_, list): + template_script = [json.loads(Path(ii).read_text()) for ii in template_script_] + else: + template_script = json.loads(Path(template_script_).read_text()) if train_style in ["dp", "dp-dist"] and explore_style == "lmp": - validate_model_backend(train_config.get("impl", "tensorflow"), explore_config) + train_backend = train_config.get("impl", "tensorflow") + validate_model_backend(train_backend, explore_config) + templates = template_script if isinstance(template_script, list) else [template_script] + for template in templates: + validate_dpa_training_template(train_backend, explore_config, template) prep_train_config = config["step_configs"]["prep_train_config"] run_train_config = config["step_configs"]["run_train_config"] prep_explore_config = config["step_configs"]["prep_explore_config"] @@ -496,12 +581,17 @@ def workflow_concurrent_learning( "not match numb_models={numb_models}" ) elif train_style == "dp-dist": + numb_models = config["train"]["numb_models"] + if "student_model_path" in config["train"] and numb_models != 1: + raise RuntimeError( + "student_model_path initializes one model; omit it for multiple " + "from-scratch students or set numb_models=1" + ) init_models_paths = ( [config["train"]["student_model_path"]] if "student_model_path" in config["train"] else None ) - config["train"]["numb_models"] = 1 else: raise RuntimeError(f"unknown params, train_style: {train_style}") @@ -557,11 +647,6 @@ def workflow_concurrent_learning( type_map = config["inputs"]["type_map"] numb_models = config["train"]["numb_models"] - template_script_ = config["train"]["template_script"] - if isinstance(template_script_, list): - template_script = [json.loads(Path(ii).read_text()) for ii in template_script_] - else: - template_script = json.loads(Path(template_script_).read_text()) if ( "teacher_model_path" in explore_config diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 77402294..699716c1 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -153,6 +153,10 @@ def _prep_run_lmp( run_executor = init_executor(run_config.pop("executor")) template_slice_config = run_config.pop("template_slice_config", {}) + prepare_models_config = deepcopy(run_config) + prepare_models_config.pop("continue_on_num_success", None) + prepare_models_config.pop("continue_on_success_ratio", None) + prep_lmp = Step( "prep-lmp", template=PythonOPTemplate( @@ -187,7 +191,7 @@ def _prep_run_lmp( }, key=step_keys["prepare-models"], executor=run_executor, - **run_config, + **prepare_models_config, ) prep_run_steps.add(prepare_models) diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py index bc551a2c..04e31a68 100644 --- a/tests/entrypoint/test_submit.py +++ b/tests/entrypoint/test_submit.py @@ -23,6 +23,7 @@ print_list_steps, submit_concurrent_learning, update_reuse_step_scheduler, + validate_dpa_training_template, ) from dpgen2.exploration.render import ( TrajRenderLammps, @@ -106,6 +107,95 @@ def modify_output_parameter(self, key, scheduler): class TestSubmit(unittest.TestCase): + def test_validate_dpa4_training_template(self): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + { + "model": { + "type": "dpa4", + "descriptor": {"type": "dpa4"}, + "use_compile": True, + "enable_tf32": True, + }, + "training": {}, + }, + ) + with self.assertRaisesRegex(RuntimeError, "model.use_compile"): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + { + "model": {"type": "dpa4", "descriptor": {"type": "dpa4"}}, + "training": {"enable_compile": True}, + }, + ) + + def test_validate_dpa4c_training_template(self): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": {"descriptor": {"type": "dpa4c"}}, + "training": {"enable_compile": True, "enable_tf32": True}, + }, + ) + with self.assertRaisesRegex(RuntimeError, "training.enable_compile"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "descriptor": {"type": "dpa4c"}, + "use_compile": True, + }, + "training": {}, + }, + ) + + with self.assertRaisesRegex(RuntimeError, "training.enable_tf32"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "descriptor": {"type": "DPA4C"}, + "enable_tf32": True, + }, + "training": {"enable_compile": True}, + }, + ) + + def test_validate_dpa_training_backend_and_format(self): + dpa4c = { + "model": {"descriptor": {"type": "dpa4c"}}, + "training": {}, + } + with self.assertRaisesRegex(RuntimeError, "requires impl='pytorch-exportable'"): + validate_dpa_training_template( + "pytorch", + {"model_devi_backend": "pytorch", "model_format": "pt2"}, + dpa4c, + ) + with self.assertRaisesRegex(RuntimeError, "requires model_format='pt2'"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pth", + }, + dpa4c, + ) + def test_expand_idx(self): ilist = ["1", "3-5", "10-20:2"] olist = expand_idx(ilist) diff --git a/tests/test_prep_run_lmp_config.py b/tests/test_prep_run_lmp_config.py new file mode 100644 index 00000000..afb9abfc --- /dev/null +++ b/tests/test_prep_run_lmp_config.py @@ -0,0 +1,48 @@ +import unittest + +from dflow import ( + Step, +) + +from dpgen2.op.prep_lmp import ( + PrepLmp, +) +from dpgen2.op.run_lmp import ( + RunLmp, +) +from dpgen2.superop.prep_run_lmp import ( + PrepRunLmp, +) +from dpgen2.utils.step_config import ( + normalize as normalize_step_dict, +) + + +class TestPrepRunLmpConfig(unittest.TestCase): + def test_non_sliced_model_preparation_ignores_slice_success_controls(self): + config = normalize_step_dict( + { + "continue_on_num_success": 1, + "continue_on_success_ratio": 0.5, + } + ) + steps = PrepRunLmp( + "prep-run-lmp", + PrepLmp, + RunLmp, + prep_config=normalize_step_dict({}), + run_config=config, + ) + prepare_models = next( + step for step in steps.steps if step.name == "prepare-models" + ) + run_lmp = next(step for step in steps.steps if step.name == "run-lmp") + + self.assertIsNone(prepare_models.continue_on_num_success) + self.assertIsNone(prepare_models.continue_on_success_ratio) + self.assertEqual(run_lmp.continue_on_num_success, 1) + self.assertEqual(run_lmp.continue_on_success_ratio, 0.5) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From fbf266d0370b8555e252f43783e10567864141af Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Aug 2026 13:37:08 +0000 Subject: [PATCH 08/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- dpgen2/entrypoint/submit.py | 12 +++++++++--- tests/test_prep_run_lmp_config.py | 6 ++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index a73c657c..d751cfed 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -478,7 +478,9 @@ def _model_family(template_script: dict) -> Optional[str]: for _, model in _iter_model_sections(template_script): model_type = model.get("type") descriptor = model.get("descriptor", {}) - descriptor_type = descriptor.get("type") if isinstance(descriptor, dict) else None + descriptor_type = ( + descriptor.get("type") if isinstance(descriptor, dict) else None + ) model_type = model_type.lower() if isinstance(model_type, str) else model_type descriptor_type = ( descriptor_type.lower() @@ -512,7 +514,9 @@ def validate_dpa_training_template( normalized_explore = RunLmp.normalize_config(explore_config) if normalized_explore["model_format"] != "pt2": - raise RuntimeError(f"{family.upper()} LAMMPS exploration requires model_format='pt2'") + raise RuntimeError( + f"{family.upper()} LAMMPS exploration requires model_format='pt2'" + ) misplaced = [] if family == "dpa4c": @@ -557,7 +561,9 @@ def workflow_concurrent_learning( if train_style in ["dp", "dp-dist"] and explore_style == "lmp": train_backend = train_config.get("impl", "tensorflow") validate_model_backend(train_backend, explore_config) - templates = template_script if isinstance(template_script, list) else [template_script] + templates = ( + template_script if isinstance(template_script, list) else [template_script] + ) for template in templates: validate_dpa_training_template(train_backend, explore_config, template) prep_train_config = config["step_configs"]["prep_train_config"] diff --git a/tests/test_prep_run_lmp_config.py b/tests/test_prep_run_lmp_config.py index afb9abfc..7da46cb8 100644 --- a/tests/test_prep_run_lmp_config.py +++ b/tests/test_prep_run_lmp_config.py @@ -13,9 +13,7 @@ from dpgen2.superop.prep_run_lmp import ( PrepRunLmp, ) -from dpgen2.utils.step_config import ( - normalize as normalize_step_dict, -) +from dpgen2.utils.step_config import normalize as normalize_step_dict class TestPrepRunLmpConfig(unittest.TestCase): @@ -45,4 +43,4 @@ def test_non_sliced_model_preparation_ignores_slice_success_controls(self): if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() From 6045f3bfbe75256274200ab85eea47c9798acf49 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Mon, 24 Aug 2026 13:51:46 +0000 Subject: [PATCH 09/14] fix: reject mixed DPA4 backend templates --- docs/input.md | 7 ++++--- dpgen2/entrypoint/submit.py | 12 +++++++++--- tests/entrypoint/test_submit.py | 22 ++++++++++++++++++++++ 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/docs/input.md b/docs/input.md index aa52c38b..badafb73 100644 --- a/docs/input.md +++ b/docs/input.md @@ -117,9 +117,10 @@ DPA4C template. Conversely, DPA4 does not use the DPA4C paths `training.enable_compile` or `training.enable_tf32`. DPGEN2 validates these backend-specific placements before creating the workflow, but it does not inject or change performance and numerical-policy settings. Therefore setting only -`train.config.impl` is not sufficient to enable compilation: the DeePMD template -must contain `training.enable_compile: true` explicitly. Run a bounded smoke test -and inspect the generated `task.*/input.json` before launching a long campaign. +`train.config.impl` is not sufficient to enable compilation: a DPA4 template +must contain `model.use_compile: true`, while a DPA4C template must contain +`training.enable_compile: true`. Run a bounded smoke test and inspect the +generated `task.*/input.json` before launching a long campaign. TensorFlow remains the default when `impl` is omitted. For PyTorch and PyTorch Exportable checkpoints, `model_devi_backend` must match diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index d751cfed..6a97da41 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -475,6 +475,7 @@ def _iter_model_sections(template_script: dict): def _model_family(template_script: dict) -> Optional[str]: + families = set() for _, model in _iter_model_sections(template_script): model_type = model.get("type") descriptor = model.get("descriptor", {}) @@ -488,10 +489,15 @@ def _model_family(template_script: dict) -> Optional[str]: else descriptor_type ) if descriptor_type == "dpa4c": - return "dpa4c" + families.add("dpa4c") if model_type == "dpa4" or descriptor_type in {"dpa4", "sezm"}: - return "dpa4" - return None + families.add("dpa4") + if len(families) > 1: + raise RuntimeError( + "A training template cannot mix DPA4 and DPA4C branches because " + "they require different DeePMD backends" + ) + return next(iter(families), None) def validate_dpa_training_template( diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py index 04e31a68..88f2d548 100644 --- a/tests/entrypoint/test_submit.py +++ b/tests/entrypoint/test_submit.py @@ -196,6 +196,28 @@ def test_validate_dpa_training_backend_and_format(self): dpa4c, ) + def test_validate_mixed_dpa4_and_dpa4c_branches(self): + with self.assertRaisesRegex(RuntimeError, "cannot mix DPA4 and DPA4C"): + validate_dpa_training_template( + "pytorch-exportable", + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + { + "model": { + "model_dict": { + "teacher": {"descriptor": {"type": "dpa4"}}, + "student": {"descriptor": {"type": "dpa4c"}}, + } + }, + "training": { + "enable_compile": True, + "enable_tf32": True, + }, + }, + ) + def test_expand_idx(self): ilist = ["1", "3-5", "10-20:2"] olist = expand_idx(ilist) From 3aa43499e47fd8776a6bf55c0de404ff3746d033 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Wed, 26 Aug 2026 07:33:37 +0000 Subject: [PATCH 10/14] fix: preflight PT2 export compiler --- docs/input.md | 8 ++++++++ dpgen2/op/run_lmp.py | 24 +++++++++++++++++++++++ tests/op/test_run_lmp.py | 41 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 71 insertions(+), 2 deletions(-) diff --git a/docs/input.md b/docs/input.md index badafb73..c0423de8 100644 --- a/docs/input.md +++ b/docs/input.md @@ -129,6 +129,14 @@ export runs with `run_explore_config`, which must select hardware, the libtorch version, and other runtime libraries compatible with the LAMMPS exploration environment. +PT2 export uses PyTorch AOTInductor and therefore also requires a working C++ +compiler on `run_explore_config`. Install `g++` or set `CXX` to an executable +compiler path in that runtime. DPGEN2 checks this capability before freezing a +checkpoint and reports a non-retryable configuration error when it is absent. +For remote wrappers, run the same compiler preflight before long training jobs +when training and export share one runtime, so a missing compiler cannot waste a +completed training campaign. + ### Exploration diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 5f474f54..ce29f596 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -4,6 +4,7 @@ import os import random import re +import shutil from pathlib import ( Path, ) @@ -542,9 +543,32 @@ def _compressed_model_name(index, model_format): return "model.%03d.compressed.%s" % (index, model_format) +def check_pt2_export_environment(): + """Fail fast when AOTInductor cannot find a C++ compiler.""" + configured = os.environ.get("CXX") + compiler = configured or shutil.which("g++") or shutil.which("c++") + if compiler is None or not Path(compiler).is_file(): + raise FatalError( + "PT2 export requires a working C++ compiler on the model-deviation " + "resource. Install g++ or set CXX to an executable compiler path " + "in run_explore_config before submitting the workflow." + ) + ret, out, err = run_command([str(compiler), "--version"]) + if ret != 0: + raise FatalError( + "PT2 export cannot execute the configured C++ compiler " + f"'{compiler}'. Set CXX to a working compiler on the " + f"model-deviation resource. Output: {out}{err}" + ) + + def prepare_dp_models(models, config): """Return frozen models, exporting checkpoints once when needed.""" backend = _model_backend(config) + if config["model_format"] == "pt2" and any( + Path(model).suffix == ".pt" for model in models + ): + check_pt2_export_environment() prepared = [] output_dir = Path("prepared_models") for idx, model in enumerate(models): diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 77618326..c905d493 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -12,6 +12,7 @@ OP, OPIO, Artifact, + FatalError, OPIOSign, TransientError, ) @@ -309,8 +310,9 @@ def tearDown(self): shutil.rmtree(self.model_dir, ignore_errors=True) shutil.rmtree("prepared_models", ignore_errors=True) + @patch("dpgen2.op.run_lmp.check_pt2_export_environment") @patch("dpgen2.op.run_lmp.run_command") - def test_dpa4_pt2(self, mocked_run): + def test_dpa4_pt2(self, mocked_run, mocked_preflight): mocked_run.return_value = (0, "", "") models = PrepareDPModels().execute( OPIO( @@ -346,9 +348,11 @@ def test_dpa4_pt2(self, mocked_run): for idx, model in enumerate(self.models) ] ) + mocked_preflight.assert_called_once_with() + @patch("dpgen2.op.run_lmp.check_pt2_export_environment") @patch("dpgen2.op.run_lmp.run_command") - def test_dpa4c_compressed_pt2(self, mocked_run): + def test_dpa4c_compressed_pt2(self, mocked_run, mocked_preflight): mocked_run.return_value = (0, "", "") models = PrepareDPModels().execute( OPIO( @@ -391,6 +395,39 @@ def test_dpa4c_compressed_pt2(self, mocked_run): ), ] ) + mocked_preflight.assert_called_once_with() + + @patch.dict("os.environ", {}, clear=True) + @patch("dpgen2.op.run_lmp.shutil.which", return_value=None) + def test_pt2_export_requires_compiler(self, mocked_which): + with self.assertRaisesRegex(FatalError, "set CXX"): + PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + "models": self.models[:1], + } + ) + ) + mocked_which.assert_has_calls([call("g++"), call("c++")]) + + @patch.dict("os.environ", {"CXX": "/missing/compiler"}, clear=True) + def test_pt2_export_rejects_missing_configured_compiler(self): + with self.assertRaisesRegex(FatalError, "set CXX"): + PrepareDPModels().execute( + OPIO( + { + "config": { + "model_devi_backend": "pytorch-exportable", + "model_format": "pt2", + }, + "models": self.models[:1], + } + ) + ) def test_training_and_deployment_backends_must_match(self): with self.assertRaisesRegex(RuntimeError, "cannot freeze a checkpoint"): From 6bf4d515d1612713cae37eb3b37401ec78583478 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Wed, 26 Aug 2026 07:44:18 +0000 Subject: [PATCH 11/14] Revert "fix: preflight PT2 export compiler" This reverts commit 3aa43499e47fd8776a6bf55c0de404ff3746d033. --- docs/input.md | 8 -------- dpgen2/op/run_lmp.py | 24 ----------------------- tests/op/test_run_lmp.py | 41 ++-------------------------------------- 3 files changed, 2 insertions(+), 71 deletions(-) diff --git a/docs/input.md b/docs/input.md index c0423de8..badafb73 100644 --- a/docs/input.md +++ b/docs/input.md @@ -129,14 +129,6 @@ export runs with `run_explore_config`, which must select hardware, the libtorch version, and other runtime libraries compatible with the LAMMPS exploration environment. -PT2 export uses PyTorch AOTInductor and therefore also requires a working C++ -compiler on `run_explore_config`. Install `g++` or set `CXX` to an executable -compiler path in that runtime. DPGEN2 checks this capability before freezing a -checkpoint and reports a non-retryable configuration error when it is absent. -For remote wrappers, run the same compiler preflight before long training jobs -when training and export share one runtime, so a missing compiler cannot waste a -completed training campaign. - ### Exploration diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index ce29f596..5f474f54 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -4,7 +4,6 @@ import os import random import re -import shutil from pathlib import ( Path, ) @@ -543,32 +542,9 @@ def _compressed_model_name(index, model_format): return "model.%03d.compressed.%s" % (index, model_format) -def check_pt2_export_environment(): - """Fail fast when AOTInductor cannot find a C++ compiler.""" - configured = os.environ.get("CXX") - compiler = configured or shutil.which("g++") or shutil.which("c++") - if compiler is None or not Path(compiler).is_file(): - raise FatalError( - "PT2 export requires a working C++ compiler on the model-deviation " - "resource. Install g++ or set CXX to an executable compiler path " - "in run_explore_config before submitting the workflow." - ) - ret, out, err = run_command([str(compiler), "--version"]) - if ret != 0: - raise FatalError( - "PT2 export cannot execute the configured C++ compiler " - f"'{compiler}'. Set CXX to a working compiler on the " - f"model-deviation resource. Output: {out}{err}" - ) - - def prepare_dp_models(models, config): """Return frozen models, exporting checkpoints once when needed.""" backend = _model_backend(config) - if config["model_format"] == "pt2" and any( - Path(model).suffix == ".pt" for model in models - ): - check_pt2_export_environment() prepared = [] output_dir = Path("prepared_models") for idx, model in enumerate(models): diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index c905d493..77618326 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -12,7 +12,6 @@ OP, OPIO, Artifact, - FatalError, OPIOSign, TransientError, ) @@ -310,9 +309,8 @@ def tearDown(self): shutil.rmtree(self.model_dir, ignore_errors=True) shutil.rmtree("prepared_models", ignore_errors=True) - @patch("dpgen2.op.run_lmp.check_pt2_export_environment") @patch("dpgen2.op.run_lmp.run_command") - def test_dpa4_pt2(self, mocked_run, mocked_preflight): + def test_dpa4_pt2(self, mocked_run): mocked_run.return_value = (0, "", "") models = PrepareDPModels().execute( OPIO( @@ -348,11 +346,9 @@ def test_dpa4_pt2(self, mocked_run, mocked_preflight): for idx, model in enumerate(self.models) ] ) - mocked_preflight.assert_called_once_with() - @patch("dpgen2.op.run_lmp.check_pt2_export_environment") @patch("dpgen2.op.run_lmp.run_command") - def test_dpa4c_compressed_pt2(self, mocked_run, mocked_preflight): + def test_dpa4c_compressed_pt2(self, mocked_run): mocked_run.return_value = (0, "", "") models = PrepareDPModels().execute( OPIO( @@ -395,39 +391,6 @@ def test_dpa4c_compressed_pt2(self, mocked_run, mocked_preflight): ), ] ) - mocked_preflight.assert_called_once_with() - - @patch.dict("os.environ", {}, clear=True) - @patch("dpgen2.op.run_lmp.shutil.which", return_value=None) - def test_pt2_export_requires_compiler(self, mocked_which): - with self.assertRaisesRegex(FatalError, "set CXX"): - PrepareDPModels().execute( - OPIO( - { - "config": { - "model_devi_backend": "pytorch-exportable", - "model_format": "pt2", - }, - "models": self.models[:1], - } - ) - ) - mocked_which.assert_has_calls([call("g++"), call("c++")]) - - @patch.dict("os.environ", {"CXX": "/missing/compiler"}, clear=True) - def test_pt2_export_rejects_missing_configured_compiler(self): - with self.assertRaisesRegex(FatalError, "set CXX"): - PrepareDPModels().execute( - OPIO( - { - "config": { - "model_devi_backend": "pytorch-exportable", - "model_format": "pt2", - }, - "models": self.models[:1], - } - ) - ) def test_training_and_deployment_backends_must_match(self): with self.assertRaisesRegex(RuntimeError, "cannot freeze a checkpoint"): From 08edc826d616cb4eac6e616ab4b2f11f118fc91f Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Wed, 26 Aug 2026 11:07:58 +0000 Subject: [PATCH 12/14] test: cover model-backend validation and export branches --- tests/op/test_run_lmp.py | 147 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 77618326..e80ec7ea 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -37,8 +37,12 @@ from dpgen2.op.run_lmp import ( PrepareDPModels, RunLmp, + _model_backend, + compress_model, + ensure_pt2_atom_map, get_ele_temp, merge_pimd_files, + prepare_dp_models, set_models, validate_model_backend, ) @@ -534,3 +538,146 @@ def tearDown(self): ]: if os.path.exists(f): os.remove(f) + + +class TestModelBackendValidation(unittest.TestCase): + def test_unsupported_backend(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported model-deviation backend"): + _model_backend({"model_devi_backend": "unknown", "model_format": "pt2", "dp_compress": False}) + + def test_unsupported_format(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): + _model_backend({"model_devi_backend": "pytorch", "model_format": "onnx", "dp_compress": False}) + + def test_pth_requires_pytorch(self): + with self.assertRaisesRegex(RuntimeError, "pth model format requires the pytorch backend"): + _model_backend({"model_devi_backend": "pytorch-exportable", "model_format": "pth", "dp_compress": False}) + + def test_compress_requires_pt_expt_pt2(self): + with self.assertRaisesRegex(RuntimeError, "pytorch-exportable backend"): + _model_backend({"model_devi_backend": "pytorch", "model_format": "pt2", "dp_compress": True}) + + def test_tensorflow_backend_skips_validation(self): + validate_model_backend("tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"}) + + +class TestEnsurePt2AtomMap(unittest.TestCase): + def test_existing_map_before_read_is_kept(self): + lines = "atom_modify map yes\nread_data conf.lmp\npair_style deepmd\n" + with open("_test_lmp.in", "w") as f: + f.write(lines) + ensure_pt2_atom_map("_test_lmp.in") + result = Path("_test_lmp.in").read_text() + self.assertEqual(result.count("atom_modify map yes"), 1) + os.remove("_test_lmp.in") + + def test_map_inserted_before_read(self): + lines = "units metal\nread_data conf.lmp\npair_style deepmd\n" + with open("_test_lmp.in", "w") as f: + f.write(lines) + ensure_pt2_atom_map("_test_lmp.in") + result = Path("_test_lmp.in").read_text() + self.assertIn("map yes", result) + self.assertLess(result.index("map yes"), result.index("read_data")) + os.remove("_test_lmp.in") + + def test_no_read_raises(self): + lines = "units metal\npair_style deepmd\n" + with open("_test_lmp.in", "w") as f: + f.write(lines) + with self.assertRaisesRegex(RuntimeError, "read_data or read_restart"): + ensure_pt2_atom_map("_test_lmp.in") + os.remove("_test_lmp.in") + + def test_map_after_read_raises(self): + lines = "read_data conf.lmp\natom_modify map yes\n" + with open("_test_lmp.in", "w") as f: + f.write(lines) + with self.assertRaisesRegex(RuntimeError, "atom_modify map yes"): + ensure_pt2_atom_map("_test_lmp.in") + os.remove("_test_lmp.in") + + +class TestCompressModelFailure(unittest.TestCase): + @patch("dpgen2.op.run_lmp.run_command") + def test_compress_failure_raises(self, mocked_run): + mocked_run.return_value = (1, "out", "compress error") + with self.assertRaises(TransientError): + compress_model("frozen.pt2", "compressed.pt2", "pytorch-exportable") + + +class TestPrepareDPModelsPassthrough(unittest.TestCase): + def setUp(self): + self.model_dir = Path("_test_models") + self.model_dir.mkdir(exist_ok=True) + + def tearDown(self): + shutil.rmtree(self.model_dir, ignore_errors=True) + shutil.rmtree("prepared_models", ignore_errors=True) + + def test_pth_passthrough(self): + model = self.model_dir / "model.000.pth" + model.write_text("frozen") + config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pth"}) + result = prepare_dp_models([model], config) + self.assertEqual(result, [model.resolve()]) + + def test_pt2_passthrough(self): + model = self.model_dir / "model.000.pt2" + model.write_text("frozen") + config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pt2"}) + result = prepare_dp_models([model], config) + self.assertEqual(result, [model.resolve()]) + + def test_unsupported_extension_raises(self): + model = self.model_dir / "model.onnx" + model.write_text("bad") + config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pt2"}) + with self.assertRaisesRegex(RuntimeError, "not supported"): + prepare_dp_models([model], config) + + +class TestModelBackendValidation(unittest.TestCase): + def test_unsupported_backend(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported model-deviation backend"): + _model_backend({"model_devi_backend": "bogus", "model_format": "pt2", "dp_compress": False}) + + def test_unsupported_format(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): + _model_backend({"model_devi_backend": "pytorch", "model_format": "xyz", "dp_compress": False}) + + def test_pth_requires_pytorch(self): + with self.assertRaisesRegex(RuntimeError, "pth model format requires the pytorch"): + _model_backend({"model_devi_backend": "pytorch-exportable", "model_format": "pth", "dp_compress": False}) + + def test_compress_requires_exportable_pt2(self): + with self.assertRaisesRegex(RuntimeError, "Compressed pt2"): + _model_backend({"model_devi_backend": "pytorch", "model_format": "pt2", "dp_compress": True}) + + def test_validate_non_pytorch_backend_skips(self): + validate_model_backend("tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"}) + + +class TestEnsurePt2AtomMap(unittest.TestCase): + def test_map_already_present_before_read(self): + lines = "atom_modify map yes\nread_data conf.lmp\n" + Path("test_input.lammps").write_text(lines) + ensure_pt2_atom_map("test_input.lammps") + result = Path("test_input.lammps").read_text() + self.assertIn("atom_modify map yes", result) + self.assertEqual(result.count("atom_modify"), 1) + os.remove("test_input.lammps") + + def test_no_read_command_raises(self): + Path("test_input.lammps").write_text("atom_modify map yes\npair_style deepmd\n") + with self.assertRaisesRegex(RuntimeError, "read_data or read_restart"): + ensure_pt2_atom_map("test_input.lammps") + os.remove("test_input.lammps") + + +class TestCompressModelFailure(unittest.TestCase): + @patch("dpgen2.op.run_lmp.run_command") + def test_compress_failure_raises(self, mocked_run): + mocked_run.return_value = (1, "", "compress error") + with self.assertRaisesRegex(TransientError, "compress failed"): + compress_model("input.pt2", "output.pt2", "pytorch-exportable") From 02e5898d9da8813110c3fde9acaf4f3b8c9103b7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:08:11 +0000 Subject: [PATCH 13/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/op/test_run_lmp.py | 100 ++++++++++++++++++++++++++++++++------- 1 file changed, 83 insertions(+), 17 deletions(-) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index e80ec7ea..e22f5d37 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -542,23 +542,53 @@ def tearDown(self): class TestModelBackendValidation(unittest.TestCase): def test_unsupported_backend(self): - with self.assertRaisesRegex(RuntimeError, "Unsupported model-deviation backend"): - _model_backend({"model_devi_backend": "unknown", "model_format": "pt2", "dp_compress": False}) + with self.assertRaisesRegex( + RuntimeError, "Unsupported model-deviation backend" + ): + _model_backend( + { + "model_devi_backend": "unknown", + "model_format": "pt2", + "dp_compress": False, + } + ) def test_unsupported_format(self): with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): - _model_backend({"model_devi_backend": "pytorch", "model_format": "onnx", "dp_compress": False}) + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "onnx", + "dp_compress": False, + } + ) def test_pth_requires_pytorch(self): - with self.assertRaisesRegex(RuntimeError, "pth model format requires the pytorch backend"): - _model_backend({"model_devi_backend": "pytorch-exportable", "model_format": "pth", "dp_compress": False}) + with self.assertRaisesRegex( + RuntimeError, "pth model format requires the pytorch backend" + ): + _model_backend( + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pth", + "dp_compress": False, + } + ) def test_compress_requires_pt_expt_pt2(self): with self.assertRaisesRegex(RuntimeError, "pytorch-exportable backend"): - _model_backend({"model_devi_backend": "pytorch", "model_format": "pt2", "dp_compress": True}) + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "pt2", + "dp_compress": True, + } + ) def test_tensorflow_backend_skips_validation(self): - validate_model_backend("tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"}) + validate_model_backend( + "tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) class TestEnsurePt2AtomMap(unittest.TestCase): @@ -618,44 +648,80 @@ def tearDown(self): def test_pth_passthrough(self): model = self.model_dir / "model.000.pth" model.write_text("frozen") - config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pth"}) + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pth"} + ) result = prepare_dp_models([model], config) self.assertEqual(result, [model.resolve()]) def test_pt2_passthrough(self): model = self.model_dir / "model.000.pt2" model.write_text("frozen") - config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pt2"}) + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) result = prepare_dp_models([model], config) self.assertEqual(result, [model.resolve()]) def test_unsupported_extension_raises(self): model = self.model_dir / "model.onnx" model.write_text("bad") - config = RunLmp.normalize_config({"model_devi_backend": "pytorch", "model_format": "pt2"}) + config = RunLmp.normalize_config( + {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) with self.assertRaisesRegex(RuntimeError, "not supported"): prepare_dp_models([model], config) class TestModelBackendValidation(unittest.TestCase): def test_unsupported_backend(self): - with self.assertRaisesRegex(RuntimeError, "Unsupported model-deviation backend"): - _model_backend({"model_devi_backend": "bogus", "model_format": "pt2", "dp_compress": False}) + with self.assertRaisesRegex( + RuntimeError, "Unsupported model-deviation backend" + ): + _model_backend( + { + "model_devi_backend": "bogus", + "model_format": "pt2", + "dp_compress": False, + } + ) def test_unsupported_format(self): with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): - _model_backend({"model_devi_backend": "pytorch", "model_format": "xyz", "dp_compress": False}) + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "xyz", + "dp_compress": False, + } + ) def test_pth_requires_pytorch(self): - with self.assertRaisesRegex(RuntimeError, "pth model format requires the pytorch"): - _model_backend({"model_devi_backend": "pytorch-exportable", "model_format": "pth", "dp_compress": False}) + with self.assertRaisesRegex( + RuntimeError, "pth model format requires the pytorch" + ): + _model_backend( + { + "model_devi_backend": "pytorch-exportable", + "model_format": "pth", + "dp_compress": False, + } + ) def test_compress_requires_exportable_pt2(self): with self.assertRaisesRegex(RuntimeError, "Compressed pt2"): - _model_backend({"model_devi_backend": "pytorch", "model_format": "pt2", "dp_compress": True}) + _model_backend( + { + "model_devi_backend": "pytorch", + "model_format": "pt2", + "dp_compress": True, + } + ) def test_validate_non_pytorch_backend_skips(self): - validate_model_backend("tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"}) + validate_model_backend( + "tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"} + ) class TestEnsurePt2AtomMap(unittest.TestCase): From 2d8233929c9ab3570f1f0e7ae6d40c88c31876c0 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Wed, 26 Aug 2026 11:55:35 +0000 Subject: [PATCH 14/14] fix: deduplicate test classes --- tests/op/test_run_lmp.py | 96 ---------------------------------------- 1 file changed, 96 deletions(-) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index e22f5d37..743fb307 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -540,102 +540,6 @@ def tearDown(self): os.remove(f) -class TestModelBackendValidation(unittest.TestCase): - def test_unsupported_backend(self): - with self.assertRaisesRegex( - RuntimeError, "Unsupported model-deviation backend" - ): - _model_backend( - { - "model_devi_backend": "unknown", - "model_format": "pt2", - "dp_compress": False, - } - ) - - def test_unsupported_format(self): - with self.assertRaisesRegex(RuntimeError, "Unsupported model format"): - _model_backend( - { - "model_devi_backend": "pytorch", - "model_format": "onnx", - "dp_compress": False, - } - ) - - def test_pth_requires_pytorch(self): - with self.assertRaisesRegex( - RuntimeError, "pth model format requires the pytorch backend" - ): - _model_backend( - { - "model_devi_backend": "pytorch-exportable", - "model_format": "pth", - "dp_compress": False, - } - ) - - def test_compress_requires_pt_expt_pt2(self): - with self.assertRaisesRegex(RuntimeError, "pytorch-exportable backend"): - _model_backend( - { - "model_devi_backend": "pytorch", - "model_format": "pt2", - "dp_compress": True, - } - ) - - def test_tensorflow_backend_skips_validation(self): - validate_model_backend( - "tensorflow", {"model_devi_backend": "pytorch", "model_format": "pt2"} - ) - - -class TestEnsurePt2AtomMap(unittest.TestCase): - def test_existing_map_before_read_is_kept(self): - lines = "atom_modify map yes\nread_data conf.lmp\npair_style deepmd\n" - with open("_test_lmp.in", "w") as f: - f.write(lines) - ensure_pt2_atom_map("_test_lmp.in") - result = Path("_test_lmp.in").read_text() - self.assertEqual(result.count("atom_modify map yes"), 1) - os.remove("_test_lmp.in") - - def test_map_inserted_before_read(self): - lines = "units metal\nread_data conf.lmp\npair_style deepmd\n" - with open("_test_lmp.in", "w") as f: - f.write(lines) - ensure_pt2_atom_map("_test_lmp.in") - result = Path("_test_lmp.in").read_text() - self.assertIn("map yes", result) - self.assertLess(result.index("map yes"), result.index("read_data")) - os.remove("_test_lmp.in") - - def test_no_read_raises(self): - lines = "units metal\npair_style deepmd\n" - with open("_test_lmp.in", "w") as f: - f.write(lines) - with self.assertRaisesRegex(RuntimeError, "read_data or read_restart"): - ensure_pt2_atom_map("_test_lmp.in") - os.remove("_test_lmp.in") - - def test_map_after_read_raises(self): - lines = "read_data conf.lmp\natom_modify map yes\n" - with open("_test_lmp.in", "w") as f: - f.write(lines) - with self.assertRaisesRegex(RuntimeError, "atom_modify map yes"): - ensure_pt2_atom_map("_test_lmp.in") - os.remove("_test_lmp.in") - - -class TestCompressModelFailure(unittest.TestCase): - @patch("dpgen2.op.run_lmp.run_command") - def test_compress_failure_raises(self, mocked_run): - mocked_run.return_value = (1, "out", "compress error") - with self.assertRaises(TransientError): - compress_model("frozen.pt2", "compressed.pt2", "pytorch-exportable") - - class TestPrepareDPModelsPassthrough(unittest.TestCase): def setUp(self): self.model_dir = Path("_test_models")