forked from deepmodeling/dpgen2
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: merged mode for caly evo step (deepmodeling#224)
one can choose mode from `default` and `merge` in run_explore_config/mode <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Introduced `CalyEvoStepMerge` class to handle merging results from evolutionary steps. - Added support for different execution modes ("merge" and "default") in the concurrent learning operations. - **Improvements** - Updated configuration settings in `input.test.json` to include new image versions and refined atom configurations. - **Bug Fixes** - Adjusted temporary directory path in `prep_caly_dp_optim.py` to correct directory structure. - Modified parameter types and added new parameters in various functions to enhance flexibility and accuracy. - **Tests** - Added new test cases for `CalyEvoStepMerge` and other workflow operations. - Updated existing tests to reflect changes in parameters and paths. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: zjgemi <[email protected]> Co-authored-by: zjgemi <[email protected]>
- Loading branch information
1 parent
aa806e1
commit 9408071
Showing
14 changed files
with
658 additions
and
87 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,129 @@ | ||
import json | ||
import logging | ||
import pickle | ||
import shutil | ||
from pathlib import ( | ||
Path, | ||
) | ||
from typing import ( | ||
List, | ||
Tuple, | ||
) | ||
|
||
from dflow import ( | ||
Step, | ||
Workflow, | ||
download_artifact, | ||
upload_artifact, | ||
) | ||
from dflow.python import ( | ||
OP, | ||
OPIO, | ||
Artifact, | ||
BigParameter, | ||
OPIOSign, | ||
Parameter, | ||
Slices, | ||
TransientError, | ||
) | ||
from dflow.utils import ( | ||
flatten, | ||
) | ||
|
||
from dpgen2.constants import ( | ||
calypso_check_opt_file, | ||
calypso_run_opt_file, | ||
) | ||
from dpgen2.exploration.task import ( | ||
ExplorationTaskGroup, | ||
) | ||
from dpgen2.superop.caly_evo_step import ( | ||
CalyEvoStep, | ||
) | ||
from dpgen2.utils import ( | ||
BinaryFileInput, | ||
set_directory, | ||
) | ||
from dpgen2.utils.run_command import ( | ||
run_command, | ||
) | ||
|
||
|
||
class CalyEvoStepMerge(OP): | ||
def __init__(self, mode="debug", *args, **kwargs): | ||
self.mode = mode | ||
self.args = args | ||
self.kwargs = kwargs | ||
|
||
@classmethod | ||
def get_input_sign(cls): | ||
return OPIOSign( | ||
{ | ||
"iter_num": int, | ||
"cnt_num": Parameter(int, default=0), | ||
"block_id": Parameter(str, default=""), | ||
"task_name": BigParameter(str), | ||
"expl_config": BigParameter(dict), | ||
"models": Artifact(Path), | ||
"input_file": Artifact(Path), | ||
"caly_run_opt_file": Artifact(Path), | ||
"caly_check_opt_file": Artifact(Path), | ||
"results": Artifact(Path, optional=True), | ||
"step": Artifact(Path, optional=True), | ||
"opt_results_dir": Artifact(List[Path], optional=True), | ||
"qhull_input": Artifact(Path, optional=True), | ||
} | ||
) | ||
|
||
@classmethod | ||
def get_output_sign(cls): | ||
return OPIOSign( | ||
{ | ||
"traj_results": Artifact(List[Path]), | ||
} | ||
) | ||
|
||
@OP.exec_sign_check | ||
def execute( | ||
self, | ||
ip: OPIO, | ||
) -> OPIO: | ||
from dflow import ( | ||
config, | ||
) | ||
|
||
config["mode"] = self.mode | ||
wf = Workflow("caly-evo-workflow") | ||
steps = CalyEvoStep(*self.args, **self.kwargs) | ||
step = Step( | ||
"caly-evo-step", | ||
template=steps, | ||
slices=Slices(output_artifact=["traj_results"]), | ||
parameters={k: ip[k] for k in steps.inputs.parameters}, | ||
artifacts={ | ||
k: upload_artifact(ip[k]) if ip[k] is not None else None | ||
for k in steps.inputs.artifacts | ||
}, | ||
with_param=[0], | ||
) | ||
wf.add(step) | ||
wf.submit() | ||
wf.wait() | ||
assert wf.query_status() == "Succeeded" | ||
out = OPIO() | ||
step = wf.query_step("caly-evo-step")[0] | ||
for k in step.outputs.parameters: | ||
out[k] = step.outputs.parameters[k].value | ||
output_sign = self.get_output_sign() | ||
for k in step.outputs.artifacts: | ||
path_list = download_artifact(step.outputs.artifacts[k]) | ||
if output_sign[k].type == List[Path]: | ||
if not isinstance(path_list, list) or any( | ||
[p is not None and not isinstance(p, str) for p in path_list] | ||
): | ||
path_list = list(flatten(path_list).values()) | ||
out[k] = [Path(p) for p in path_list] | ||
elif output_sign[k].type == Path: | ||
assert len(path_list) == 1 | ||
out[k] = Path(path_list[0]) | ||
return out |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.