diff --git a/dpgen2/op/collect_run_caly.py b/dpgen2/op/collect_run_caly.py index f777fbfa..eeaf87b9 100644 --- a/dpgen2/op/collect_run_caly.py +++ b/dpgen2/op/collect_run_caly.py @@ -67,7 +67,7 @@ def get_input_sign(cls): type=Path, optional=True ), # dir named results for evo "opt_results_dir": Artifact( - type=Path, optional=True + type=List[Path], optional=True ), # dir contains POSCAR* CONTCAR* OUTCAR* "qhull_input": Artifact(type=Path, optional=True), # for vsc } @@ -141,11 +141,16 @@ def execute( results = ( ip["results"].resolve() if ip["results"] is not None else ip["results"] ) - opt_results_dir = ( - ip["opt_results_dir"].resolve() - if ip["opt_results_dir"] is not None - else ip["opt_results_dir"] - ) + # opt_results_dir = ( + # ip["opt_results_dir"].resolve() + # if ip["opt_results_dir"] is not None + # else ip["opt_results_dir"] + # ) + opt_results_dir = [] + if ip["opt_results_dir"] is not None: + for temp in ip["opt_results_dir"]: + opt_results_dir.append(Path(temp).resolve()) + qhull_input = ( ip["qhull_input"].resolve() if ip["qhull_input"] is not None diff --git a/tests/op/test_collect_run_caly.py b/tests/op/test_collect_run_caly.py index aec481ec..22ec37db 100644 --- a/tests/op/test_collect_run_caly.py +++ b/tests/op/test_collect_run_caly.py @@ -102,7 +102,7 @@ def side_effect(*args, **kwargs): "input_file": self.input_file, "step": self.step_file, "results": self.results_dir, - "opt_results_dir": self.opt_results_dir, + "opt_results_dir": [self.opt_results_dir], } ) ) @@ -138,7 +138,7 @@ def side_effect(*args, **kwargs): "input_file": self.input_file, "step": self.step_file, "results": self.results_dir, - "opt_results_dir": self.opt_results_dir, + "opt_results_dir": [self.opt_results_dir], } ) ) @@ -169,7 +169,7 @@ def side_effect(*args, **kwargs): "input_file": self.input_file, "step": self.step_file, "results": self.results_dir, - "opt_results_dir": self.opt_results_dir, + "opt_results_dir": [self.opt_results_dir], } ) ) @@ -203,7 +203,7 @@ def side_effect(*args, **kwargs): "input_file": self.input_file, "step": self.step_file, "results": self.results_dir, - "opt_results_dir": self.opt_results_dir, + "opt_results_dir": [self.opt_results_dir], } ), ) diff --git a/tests/op/test_run_caly_model_devi.py b/tests/op/test_run_caly_model_devi.py index 0718f756..6d45b403 100644 --- a/tests/op/test_run_caly_model_devi.py +++ b/tests/op/test_run_caly_model_devi.py @@ -143,7 +143,7 @@ def side_effect_1(*args, **kwargs): mocked_run_1.side_effect = side_effect_1 def side_effect_2(*args, **kwargs): - return [[1, 1, 1, 1, 1, 1, 1]] + return [[1, 1, 1, 1, 1, 1, 1, 1]] mocked_run_2.side_effect = side_effect_2 diff --git a/tests/op/test_run_dp_optim.py b/tests/op/test_run_dp_optim.py index bed97f88..7709a1f6 100644 --- a/tests/op/test_run_dp_optim.py +++ b/tests/op/test_run_dp_optim.py @@ -102,13 +102,6 @@ def side_effect(*args, **kwargs): Path(self.task_name) / "traj_results" / "0.3.traj" in list_traj_results_dir ) - self.assertEqual( - Path(self.task_name) / calypso_run_opt_file, out["caly_run_opt_file"] - ) - self.assertEqual( - Path(self.task_name) / calypso_check_opt_file, out["caly_check_opt_file"] - ) - @patch("dpgen2.op.run_dp_optim.run_command") def test_01_error(self, mocked_run): def side_effect(*args, **kwargs): diff --git a/tests/test_caly_evo_step.py b/tests/test_caly_evo_step.py index 72161083..54d2b1e5 100644 --- a/tests/test_caly_evo_step.py +++ b/tests/test_caly_evo_step.py @@ -51,6 +51,20 @@ except ModuleNotFoundError: # case of upload everything to argo, no context needed pass +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) +from mocked_ops import ( + MockedCollRunCaly, + MockedPrepDPOptim, + MockedRunDPOptim, + mocked_numb_models, +) + from dpgen2.constants import ( lmp_conf_name, lmp_input_name, @@ -82,20 +96,6 @@ ) from dpgen2.utils.step_config import normalize as normalize_step_dict -from .context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, -) -from .mocked_ops import ( - MockedCollRunCaly, - MockedPrepDPOptim, - MockedRunDPOptim, - mocked_numb_models, -) - default_config = normalize_step_dict( { "template_config": { @@ -109,7 +109,6 @@ ) -@unittest.skip("temp") class TestMockedCollRunCaly(unittest.TestCase): def setUp(self) -> None: self.config = {} @@ -160,7 +159,6 @@ def test_mocked_coll_run_caly_00(self): self.assertTrue(out["results"] == Path(self.task_name).joinpath("results")) -@unittest.skip("temp") class TestMockedRunDPOptim(unittest.TestCase): def setUp(self) -> None: self.config = {} @@ -204,7 +202,7 @@ def test_mocked_run_dp_optim(self): ) self.assertTrue(optim_results_dir, Path(self.task_name) / "optim_results_dir") - self.assertEqual(counts_optim_results_dir, 10) + self.assertEqual(counts_optim_results_dir, 15) self.assertEqual(counts_outcar_in_optim_results_dir, 5) self.assertTrue( Path(self.task_name) / "optim_results_dir" / "CONTCAR_4" @@ -220,13 +218,6 @@ def test_mocked_run_dp_optim(self): Path(self.task_name) / "traj_results" / "3.traj", list_traj_results_dir ) - self.assertEqual( - Path(self.task_name) / calypso_run_opt_file, out["caly_run_opt_file"] - ) - self.assertEqual( - Path(self.task_name) / calypso_check_opt_file, out["caly_check_opt_file"] - ) - # @unittest.skip("temporary pass") @unittest.skipIf(skip_ut_with_dflow, skip_ut_with_dflow_reason) @@ -281,6 +272,7 @@ def tearDown(self): shutil.rmtree(i, ignore_errors=True) for i in Path().glob("caly_task*"): shutil.rmtree(i, ignore_errors=True) + shutil.rmtree("upload", ignore_errors=True) @unittest.skip("only need to run test_01") def test_00(self): diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py index 542cb0e3..5daa600f 100644 --- a/tests/test_prep_run_caly.py +++ b/tests/test_prep_run_caly.py @@ -60,8 +60,8 @@ ) from mocked_ops import ( MockedCollRunCaly, - MockedPrepRunDPOptim, MockedRunCalyModelDevi, + MockedRunDPOptim, mocked_numb_models, ) @@ -72,6 +72,9 @@ from dpgen2.op.prep_caly_input import ( PrepCalyInput, ) +from dpgen2.op.prep_dp_optim import ( + PrepDPOptim, +) from dpgen2.op.run_caly_model_devi import ( RunCalyModelDevi, ) @@ -83,11 +86,22 @@ ) from dpgen2.utils.step_config import normalize as normalize_step_dict -default_config = normalize_step_dict( +prep_default_config = normalize_step_dict( + { + "template_config": { + "image": default_image, + }, + } +) +run_default_config = normalize_step_dict( { "template_config": { "image": default_image, - } + }, + "template_slice_config": { + "group_size": 2, + "pool_size": 1, + }, } ) @@ -105,7 +119,7 @@ def make_task_group_list(njobs): # @unittest.skip("temporary pass") @unittest.skipIf(skip_ut_with_dflow, skip_ut_with_dflow_reason) -class TestCalyEvoStep(unittest.TestCase): +class TestPrepRunCaly(unittest.TestCase): def setUp(self): self.expl_config = {} self.work_dir = Path("storge_files") @@ -114,7 +128,9 @@ def setUp(self): self.nmodels = mocked_numb_models self.model_list = [] for ii in range(self.nmodels): - model = self.work_dir.joinpath(f"model.{ii}.pb") + model_path = self.work_dir.joinpath(f"task.{ii}") + model_path.mkdir(parents=True, exist_ok=True) + model = model_path.joinpath(f"frozen_model.pb") model.write_text(f"model {ii}") self.model_list.append(model) self.models = upload_artifact(self.model_list) @@ -127,15 +143,16 @@ def tearDown(self): shutil.rmtree(self.work_dir, ignore_errors=True) for i in Path().glob("prep-run-caly-step*"): shutil.rmtree(i, ignore_errors=True) - # shutil.rmtree("upload", ignore_errors=True) + shutil.rmtree("upload", ignore_errors=True) def test(self): caly_evo_step_op = CalyEvoStep( "caly-evo-run", MockedCollRunCaly, - MockedPrepRunDPOptim, - prep_config=default_config, - run_config=default_config, + PrepDPOptim, + MockedRunDPOptim, + prep_config=prep_default_config, + run_config=run_default_config, upload_python_packages=upload_python_packages, ) prep_run_caly_op = PrepRunCaly( @@ -143,8 +160,8 @@ def test(self): PrepCalyInput, caly_evo_step_op, MockedRunCalyModelDevi, - prep_config=default_config, - run_config=default_config, + prep_config=prep_default_config, + run_config=run_default_config, upload_python_packages=upload_python_packages, ) prep_run_caly_step = Step(