From 22f394aef9380ec839156324ce49652f136502dc Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Thu, 14 Aug 2025 22:28:47 +0800 Subject: [PATCH 1/4] feat: allow yaml config for training --- deepmd/pd/entrypoints/main.py | 4 ++-- deepmd/pt/entrypoints/main.py | 4 ++-- source/tests/pt/test_training.py | 22 ++++++++++++++++++++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index 3ef075f359..4e47dbfe77 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -22,6 +22,7 @@ ) from deepmd.common import ( expand_sys_str, + j_loader, ) from deepmd.loggers.loggers import ( set_log_handles, @@ -235,8 +236,7 @@ def train( log.info("Configuration path: %s", input_file) if LOCAL_RANK == 0: SummaryPrinter()() - with open(input_file) as fin: - config = json.load(fin) + config = j_loader(input_file) # ensure suffix, as in the command line help, we say "path prefix of checkpoint files" if init_model is not None and not init_model.endswith(".pd"): init_model += ".pd" diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 0e248583ec..630fb6d86f 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -25,6 +25,7 @@ ) from deepmd.common import ( expand_sys_str, + j_loader, ) from deepmd.env import ( GLOBAL_CONFIG, @@ -254,8 +255,7 @@ def train( env.CUSTOM_OP_USE_JIT = True if LOCAL_RANK == 0: SummaryPrinter()() - with open(input_file) as fin: - config = json.load(fin) + config = j_loader(input_file) # ensure suffix, as in the command line help, we say "path prefix of checkpoint files" if init_model is not None and not init_model.endswith(".pt"): init_model += ".pt" diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index c57c896197..2fc60924c2 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -14,6 +14,7 @@ from deepmd.pt.entrypoints.main import ( get_trainer, + train as train_entry, ) from deepmd.pt.utils.finetune import ( get_finetune_rules, @@ -180,8 +181,29 @@ def setUp(self) -> None: self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 + def test_yaml_input(self) -> None: + import yaml + + yaml_file = Path("input.yaml") + with open(yaml_file, "w") as fp: + yaml.safe_dump(self.config, fp) + train_entry( + input_file=str(yaml_file), + init_model=None, + restart=None, + finetune=None, + init_frz_model=None, + model_branch="main", + skip_neighbor_stat=True, + output="out.json", + ) + self.assertTrue(Path("out.json").exists()) + def tearDown(self) -> None: DPTrainTest.tearDown(self) + for ff in ["out.json", "input.yaml"]: + if Path(ff).exists(): + os.remove(ff) class TestDOSModelSeA(unittest.TestCase, DPTrainTest): From c2f83562e290f953e0b6298d1359589a89d4e335 Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Thu, 14 Aug 2025 22:57:01 +0800 Subject: [PATCH 2/4] test: ensure tf training accepts yaml --- deepmd/tf/entrypoints/train.py | 2 +- source/tests/tf/test_training.py | 65 ++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 source/tests/tf/test_training.py diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index b12e4fe1af..5bcca9a4e3 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -13,7 +13,7 @@ Optional, ) -from deepmd.tf.common import ( +from deepmd.common import ( j_loader, ) from deepmd.tf.env import ( diff --git a/source/tests/tf/test_training.py b/source/tests/tf/test_training.py new file mode 100644 index 0000000000..46b6d2d890 --- /dev/null +++ b/source/tests/tf/test_training.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for TensorFlow training entrypoint.""" + +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import yaml + +from deepmd.tf.entrypoints.train import ( # type: ignore + train, +) + +from .common import ( + del_data, + gen_data_type_specific, + j_loader, +) + + +class TestYamlInput: + """Ensure training entrypoint accepts YAML config.""" + + def setup_method(self) -> None: + gen_data_type_specific() + config: dict[str, Any] = j_loader("water_se_atten.json") + config["systems"] = ["system"] + config["stop_batch"] = 1 + config["save_freq"] = 1 + yaml_file = Path("input.yaml") + with open(yaml_file, "w") as fp: + yaml.safe_dump(config, fp) + self.yaml_file = yaml_file + + def teardown_method(self) -> None: + del_data() + for ff in [ + "out.json", + "input.yaml", + "lcurve.out", + "model.ckpt.data-00000-of-00001", + "model.ckpt.index", + "model.ckpt.meta", + ]: + Path(ff).unlink(missing_ok=True) + + def test_yaml_input(self) -> None: + train( + INPUT=str(self.yaml_file), + init_model=None, + restart=None, + output="out.json", + init_frz_model=None, + mpi_log="master", + log_level=0, + log_path=None, + skip_neighbor_stat=True, + ) + assert Path("out.json").exists() + + +__all__ = ["TestYamlInput"] From 5dbdeb411e593be8d6628a59f0110d836362b2bc Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Thu, 14 Aug 2025 23:14:06 +0800 Subject: [PATCH 3/4] test: remove TensorFlow YAML training test --- source/tests/tf/test_training.py | 65 -------------------------------- 1 file changed, 65 deletions(-) delete mode 100644 source/tests/tf/test_training.py diff --git a/source/tests/tf/test_training.py b/source/tests/tf/test_training.py deleted file mode 100644 index 46b6d2d890..0000000000 --- a/source/tests/tf/test_training.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Tests for TensorFlow training entrypoint.""" - -from pathlib import ( - Path, -) -from typing import ( - Any, -) - -import yaml - -from deepmd.tf.entrypoints.train import ( # type: ignore - train, -) - -from .common import ( - del_data, - gen_data_type_specific, - j_loader, -) - - -class TestYamlInput: - """Ensure training entrypoint accepts YAML config.""" - - def setup_method(self) -> None: - gen_data_type_specific() - config: dict[str, Any] = j_loader("water_se_atten.json") - config["systems"] = ["system"] - config["stop_batch"] = 1 - config["save_freq"] = 1 - yaml_file = Path("input.yaml") - with open(yaml_file, "w") as fp: - yaml.safe_dump(config, fp) - self.yaml_file = yaml_file - - def teardown_method(self) -> None: - del_data() - for ff in [ - "out.json", - "input.yaml", - "lcurve.out", - "model.ckpt.data-00000-of-00001", - "model.ckpt.index", - "model.ckpt.meta", - ]: - Path(ff).unlink(missing_ok=True) - - def test_yaml_input(self) -> None: - train( - INPUT=str(self.yaml_file), - init_model=None, - restart=None, - output="out.json", - init_frz_model=None, - mpi_log="master", - log_level=0, - log_path=None, - skip_neighbor_stat=True, - ) - assert Path("out.json").exists() - - -__all__ = ["TestYamlInput"] From 959367fc113009d8c9abc986d59b317edbfaf82c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Aug 2025 08:20:13 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt/test_training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 2fc60924c2..da239212b0 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -14,8 +14,8 @@ from deepmd.pt.entrypoints.main import ( get_trainer, - train as train_entry, ) +from deepmd.pt.entrypoints.main import train as train_entry from deepmd.pt.utils.finetune import ( get_finetune_rules, )