From b1ece05430671d28a624220944c9ea13135add52 Mon Sep 17 00:00:00 2001 From: zjgemi Date: Mon, 7 Jul 2025 16:50:25 +0800 Subject: [PATCH 1/7] support extra outputs for RunLmp Signed-off-by: zjgemi --- dpgen2/op/run_lmp.py | 13 +++++++++++++ tests/op/test_run_lmp.py | 23 +++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 64c34c9e..23dc1950 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -83,6 +83,7 @@ def get_output_sign(cls): "model_devi": Artifact(Path), "plm_output": Artifact(Path, optional=True), "optional_output": Artifact(Path, optional=True), + "extra_outputs": Artifact(List[Path], optional=True), } ) @@ -213,6 +214,10 @@ def execute( if ele_temp is not None: ret_dict["optional_output"] = work_dir / "job.json" + extra_outputs = [] + for fname in config["extra_output_files"]: + extra_outputs += list(work_dir.glob(fname)) + ret_dict["extra_outputs"] = extra_outputs return OPIO(ret_dict) def get_model_devi(self, model_devi_file): @@ -226,6 +231,7 @@ def lmp_args(): doc_head = "Select a head from multitask" 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" return [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), Argument( @@ -256,6 +262,13 @@ def lmp_args(): default=False, doc=doc_use_hdf5, ), + Argument( + "extra_output_files", + list, + optional=True, + default=[], + doc=doc_extra_output_files, + ), ] @staticmethod diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 5b7f4542..650fd82e 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -126,6 +126,29 @@ def test_error(self, mocked_run): ] mocked_run.assert_has_calls(calls) + def test_extra_outputs(self): + op = RunLmp() + out = op.execute( + OPIO( + { + "config": { + "command": "echo Hello > foo.txt", + "extra_output_files": ["foo.txt"], + }, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + work_dir = Path(self.task_name) + # check output + self.assertEqual(out["extra_outputs"], [work_dir / "foo.txt"]) + self.assertEqual( + (work_dir / "foo.txt").read_text().strip(), + "Hello -i in.lammps -log log.lammps", + ) + class TestRunLmpDist(unittest.TestCase): lmp_config = """variable NSTEPS equal 1000 From 130e5a5aa17bf8f7e8188477498c8418aab9cb1e Mon Sep 17 00:00:00 2001 From: zjgemi Date: Mon, 7 Jul 2025 17:57:01 +0800 Subject: [PATCH 2/7] merge extra_outputs in prep-run-lmp Signed-off-by: zjgemi --- dpgen2/op/run_lmp.py | 2 +- dpgen2/superop/prep_run_lmp.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 23dc1950..a448f605 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -83,7 +83,7 @@ def get_output_sign(cls): "model_devi": Artifact(Path), "plm_output": Artifact(Path, optional=True), "optional_output": Artifact(Path, optional=True), - "extra_outputs": Artifact(List[Path], optional=True), + "extra_outputs": Artifact(List[Path]), } ) diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 48d4028a..c6450ee0 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -76,6 +76,7 @@ def __init__( "model_devis": OutputArtifact(), "plm_output": OutputArtifact(), "optional_outputs": OutputArtifact(), + "extra_outputs": OutputArtifact(), } super().__init__( @@ -179,6 +180,7 @@ def _prep_run_lmp( "model_devi", "plm_output", "optional_output", + "extra_outputs", ], **template_slice_config, ), @@ -217,5 +219,8 @@ def _prep_run_lmp( prep_run_steps.outputs.artifacts[ "optional_outputs" ]._from = run_lmp.outputs.artifacts["optional_output"] + prep_run_steps.outputs.artifacts[ + "extra_outputs" + ]._from = run_lmp.outputs.artifacts["extra_outputs"] return prep_run_steps From f1d965fc015454dc44c0edd625e09583e510a58c Mon Sep 17 00:00:00 2001 From: zjgemi Date: Tue, 8 Jul 2025 15:16:28 +0800 Subject: [PATCH 3/7] add output to dpgen2 download Signed-off-by: zjgemi --- dpgen2/utils/download_dpgen2_artifacts.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dpgen2/utils/download_dpgen2_artifacts.py b/dpgen2/utils/download_dpgen2_artifacts.py index 8ab4cb61..67c2aaf0 100644 --- a/dpgen2/utils/download_dpgen2_artifacts.py +++ b/dpgen2/utils/download_dpgen2_artifacts.py @@ -63,7 +63,8 @@ def add_output( "prep-run-explore": DownloadDefinition() .add_output("logs") .add_output("trajs") - .add_output("model_devis"), + .add_output("model_devis") + .add_output("extra_outputs"), "prep-run-fp": DownloadDefinition() .add_input("confs") .add_output("logs") From 50079701b59d85e75a2ead6ab84e3d54eadc64af Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 9 Jul 2025 01:54:02 +0000 Subject: [PATCH 4/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- dpgen2/superop/prep_run_lmp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index c6450ee0..3e0a0a0f 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -219,8 +219,8 @@ def _prep_run_lmp( prep_run_steps.outputs.artifacts[ "optional_outputs" ]._from = run_lmp.outputs.artifacts["optional_output"] - prep_run_steps.outputs.artifacts[ + prep_run_steps.outputs.artifacts["extra_outputs"]._from = run_lmp.outputs.artifacts[ "extra_outputs" - ]._from = run_lmp.outputs.artifacts["extra_outputs"] + ] return prep_run_steps From 971943e15a4eeca392f0d337689f3005f94da7b4 Mon Sep 17 00:00:00 2001 From: zjgemi Date: Wed, 9 Jul 2025 10:05:33 +0800 Subject: [PATCH 5/7] fix pyright Signed-off-by: zjgemi --- dpgen2/op/run_lmp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index a448f605..60cd9305 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -217,7 +217,7 @@ def execute( extra_outputs = [] for fname in config["extra_output_files"]: extra_outputs += list(work_dir.glob(fname)) - ret_dict["extra_outputs"] = extra_outputs + ret_dict["extra_outputs"] = extra_outputs # type: ignore return OPIO(ret_dict) def get_model_devi(self, model_devi_file): From 6bae65a44d1aec5a1acbbd3e8de8fb9631466ef3 Mon Sep 17 00:00:00 2001 From: zjgemi Date: Wed, 9 Jul 2025 10:18:16 +0800 Subject: [PATCH 6/7] fix UT Signed-off-by: zjgemi --- tests/mocked_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index 9cd13c00..6e6cfadf 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -424,6 +424,7 @@ def execute( "log": work_dir / log, "traj": work_dir / traj, "model_devi": work_dir / model_devi, + "extra_outputs": [], } ) From 93027ae631da25bc2851b764849abea394a57539 Mon Sep 17 00:00:00 2001 From: zjgemi Date: Wed, 9 Jul 2025 10:44:09 +0800 Subject: [PATCH 7/7] fix UT Signed-off-by: zjgemi --- tests/utils/test_dl_dpgen2_arti.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/utils/test_dl_dpgen2_arti.py b/tests/utils/test_dl_dpgen2_arti.py index a9f251ac..037a0d80 100644 --- a/tests/utils/test_dl_dpgen2_arti.py +++ b/tests/utils/test_dl_dpgen2_arti.py @@ -123,6 +123,11 @@ def test_lmp_download(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-extra_outputs", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), ] self.assertEqual(len(mocked_dl.call_args_list), len(expected)) for ii, jj in zip(mocked_dl.call_args_list, expected): @@ -253,6 +258,11 @@ def test_update_finished_steps_exist_steps(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-extra_outputs", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), ] self.assertEqual(len(mocked_dl.call_args_list), len(expected)) for ii, jj in zip(mocked_dl.call_args_list, expected): @@ -315,6 +325,11 @@ def test_update_finished_steps_none_steps(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-extra_outputs", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), ] for ii, jj in zip(mocked_dl.call_args_list, expected): self.assertEqual(ii, jj)