From 3ad1deed2765a1093ca3badb051e843cde9dd1a8 Mon Sep 17 00:00:00 2001 From: Zhenyu Wang Date: Thu, 28 Mar 2024 16:38:19 +0800 Subject: [PATCH] fix: fix caly_evo_step superop --- dpgen2/op/prep_dp_optim.py | 13 +++++++++---- dpgen2/op/run_dp_optim.py | 7 ++++--- dpgen2/superop/caly_evo_step.py | 19 ++++++++++--------- tests/mocked_ops.py | 33 +++++++++++++++++++++++++-------- tests/test_caly_evo_step.py | 15 ++++++++------- 5 files changed, 56 insertions(+), 31 deletions(-) diff --git a/dpgen2/op/prep_dp_optim.py b/dpgen2/op/prep_dp_optim.py index 9c354f0d..271da897 100644 --- a/dpgen2/op/prep_dp_optim.py +++ b/dpgen2/op/prep_dp_optim.py @@ -111,13 +111,17 @@ def execute( finished = ip["finished"] work_dir = Path(ip["task_name"]) + # print("-----in PrepDPOptim====work_dir(task_name)", work_dir) poscar_dir = ip["poscar_dir"] models_dir = ip["models_dir"] + # print("-----in PrepDPOptim====poscar_dir", poscar_dir) + # print("-----in PrepDPOptim====models_dir", models_dir) _caly_run_opt_file = ip["caly_run_opt_file"] _caly_check_opt_file = ip["caly_check_opt_file"] caly_run_opt_file = _caly_run_opt_file.resolve() caly_check_opt_file = _caly_check_opt_file.resolve() poscar_list = [poscar.resolve() for poscar in poscar_dir.rglob("POSCAR_*")] + # print("-----in PrepDPOptim====len(poscar_dir.rglob(POSCAR*))", len(poscar_list)) poscar_list = sorted(poscar_list, key=lambda x: int(x.name.strip("POSCAR_"))) model_name = "frozen_model.pb" model_list = [model.resolve() for model in models_dir.rglob(model_name)] @@ -137,11 +141,12 @@ def execute( ] task_dirs = [] - for idx, poscar_list in enumerate(grouped_poscar_list): + for idx, _poscar_list in enumerate(grouped_poscar_list): opt_path = Path(f"opt_path_{idx}") + # print("-----in PrepDPOptim====opt_path, _poscar_list)", opt_path, _poscar_list) task_dirs.append(work_dir / opt_path) with set_directory(opt_path): - for poscar in poscar_list: + for poscar in _poscar_list: Path(poscar.name).symlink_to(poscar) Path(model_name).symlink_to(model_file) Path(caly_run_opt_file.name).symlink_to(caly_run_opt_file) @@ -157,7 +162,7 @@ def execute( { "task_names": task_names, "task_dirs": task_dirs, - "caly_run_opt_file": work_dir / calypso_run_opt_file, - "caly_check_opt_file": work_dir / calypso_check_opt_file, + "caly_run_opt_file": work_dir / caly_run_opt_file.name, + "caly_check_opt_file": work_dir / caly_check_opt_file.name, } ) diff --git a/dpgen2/op/run_dp_optim.py b/dpgen2/op/run_dp_optim.py index 743ef6f2..48d6e337 100644 --- a/dpgen2/op/run_dp_optim.py +++ b/dpgen2/op/run_dp_optim.py @@ -97,7 +97,10 @@ def execute( cnt_num = ip["cnt_num"] task_path = ip["task_dir"] - input_files = [ii.resolve() for ii in Path(task_path).iterdir()] + if task_path is not None: + input_files = [ii.resolve() for ii in Path(task_path).iterdir()] + else: + input_files = [] config = ip["config"] if ip["config"] is not None else {} command = config.get( @@ -161,7 +164,5 @@ def execute( "task_name": str(work_dir), "optim_results_dir": work_dir / optim_results_dir, "traj_results": work_dir / traj_results_dir, - # "caly_run_opt_file": work_dir / calypso_run_opt_file, - # "caly_check_opt_file": work_dir / calypso_check_opt_file, } ) diff --git a/dpgen2/superop/caly_evo_step.py b/dpgen2/superop/caly_evo_step.py index 333b4c8c..21692e8d 100644 --- a/dpgen2/superop/caly_evo_step.py +++ b/dpgen2/superop/caly_evo_step.py @@ -97,18 +97,19 @@ def __init__( self.collect_run_calypso_keys = [ "%s--collect-run-calypso-%s-%s" % (self.inputs.parameters["block_id"], i, j) - for i in range(50) - for j in range(50) + for i in range(20) + for j in range(20) ] self.prep_dp_optim_keys = [ "%s--prep-dp-optim-%s-%s" % (self.inputs.parameters["block_id"], i, j) - for i in range(50) - for j in range(50) + for i in range(20) + for j in range(20) ] self.run_dp_optim_keys = [ - "%s--run-dp-optim-%s-%s" % (self.inputs.parameters["block_id"], i, j) - for i in range(50) - for j in range(50) + "%s--run-dp-optim-%s-%s-%s" % (self.inputs.parameters["block_id"], i, j, k) + for i in range(20) + for j in range(20) + for k in range(50) ] self._keys = ( self.collect_run_calypso_keys @@ -243,7 +244,7 @@ def _caly_evo_step( run_dp_optim_op, slices=Slices( input_parameter=["task_name"], - input_artifact=["task_path"], + input_artifact=["task_dir"], output_artifact=["traj_results", "optim_results_dir"], ), python_packages=upload_python_packages, @@ -264,7 +265,7 @@ def _caly_evo_step( caly_evo_step_steps.inputs.parameters["iter_num"], caly_evo_step_steps.inputs.parameters["cnt_num"], ), - executor=prep_executor, # cpu is enough to run calypso.x, default step config is c2m4 + executor=run_executor, **run_config, ) caly_evo_step_steps.add(run_dp_optim) diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index a8817d15..a7fede1c 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -972,6 +972,7 @@ def execute( # work_dir name: calypso_task.idx work_dir = Path(ip["task_name"]) + # print("==============----------work_dir in collruncaly:", work_dir) work_dir.mkdir(exist_ok=True, parents=True) qhull_input = ( @@ -1121,20 +1122,38 @@ def execute( finished = ip["finished"] work_dir = Path(ip["task_name"]) cnt_num = ip["cnt_num"] - # print(f"--------=---------task_name: {work_dir}") + # print(f"----in RunDpOptim----task_name: {work_dir}") + # print(f"----cwd in RunDPOptim------cwd: {cwd}") work_dir.mkdir(parents=True, exist_ok=True) config = ip["config"] if ip["config"] is not None else {} command = config.get("run_opt_command", "python -u calypso_run_opt.py") - os.chdir(work_dir) + task_path = ip["task_dir"] + # print(f"----in RunDpOptim----task_dir: {task_path}") + if task_path is not None: + input_files = [ii.resolve() for ii in Path(task_path).iterdir()] + else: + pass - for i in range(1, 6): - Path().joinpath(f"CONTCAR_{str(i)}").write_text(f"CONTCAR_{str(i)}") - Path().joinpath(f"OUTCAR_{str(i)}").write_text(f"OUTCAR_{str(i)}") - Path().joinpath(f"{str(i)}.traj").write_text(f"{str(i)}.traj") + os.chdir(work_dir) if finished == "false": + for ii in input_files: + iname = ii.name + Path(iname).symlink_to(ii) + + poscar_list = sorted(Path().rglob("POSCAR_*")) + # print(f"-----in RunDpOptim----poscar_list: {poscar_list}") + cnt = 0 + for poscar in poscar_list: + cnt += 1 + poscar_name = poscar.name + num = poscar_name.strip("POSCAR_") + Path(poscar_name.replace("POSCAR", "CONTCAR")).write_text("") + Path(poscar_name.replace("POSCAR", "OUTCAR")).write_text("") + Path(f"{num}.traj").write_text("") + optim_results_dir = Path("optim_results_dir") optim_results_dir.mkdir(parents=True, exist_ok=True) for poscar in Path().glob("POSCAR_*"): @@ -1165,8 +1184,6 @@ def execute( "task_name": str(work_dir), "optim_results_dir": work_dir / optim_results_dir, "traj_results": work_dir / traj_results_dir, - "caly_run_opt_file": work_dir / calypso_run_opt_file, - "caly_check_opt_file": work_dir / calypso_check_opt_file, } ) diff --git a/tests/test_caly_evo_step.py b/tests/test_caly_evo_step.py index 41686285..72161083 100644 --- a/tests/test_caly_evo_step.py +++ b/tests/test_caly_evo_step.py @@ -276,12 +276,11 @@ def setUp(self): ) def tearDown(self): - pass - # shutil.rmtree(self.work_dir, ignore_errors=True) - # for i in Path().glob("caly-evo-step-*"): - # shutil.rmtree(i, ignore_errors=True) - # for i in Path().glob("caly_task*"): - # shutil.rmtree(i, ignore_errors=True) + shutil.rmtree(self.work_dir, ignore_errors=True) + for i in Path().glob("caly-evo-step-*"): + shutil.rmtree(i, ignore_errors=True) + for i in Path().glob("caly_task*"): + shutil.rmtree(i, ignore_errors=True) @unittest.skip("only need to run test_01") def test_00(self): @@ -324,7 +323,7 @@ def test_00(self): self.assertEqual(step.phase, "Succeeded") # @unittest.skip("temp skit") - def test_01(self): + def test_caly_evo_step(self): steps = CalyEvoStep( "caly-evo-run", MockedCollRunCaly, @@ -381,6 +380,8 @@ def test_01(self): download_artifact(step.outputs.artifacts["traj_results"]) + cwd = Path().cwd() + print(cwd, self.task_name_list) for idx, name in enumerate(self.task_name_list): cwd = Path().cwd() os.chdir(Path(name))