Skip to content

Commit

Permalink
fix: split merge and default op
Browse files Browse the repository at this point in the history
  • Loading branch information
@wangzyphysics
wangzyphysics committed May 28, 2024
1 parent ba1e85f commit ccbadd1
Show file tree
Hide file tree
Showing 6 changed files with 206 additions and 94 deletions.
35 changes: 25 additions & 10 deletions dpgen2/entrypoint/submit.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,9 @@
PrepRunFp,
PrepRunLmp,
)
from dpgen2.superop.caly_evo_step import (
CalyEvoStep,
)
from dpgen2.utils import (
BinaryFileInput,
bohrium_config_from_dict,
Expand Down Expand Up @@ -174,16 +177,28 @@ def make_concurrent_learning_op(
upload_python_packages=upload_python_packages,
)
elif explore_style == "calypso":
caly_evo_step_op = CalyEvoStepMerge(
mode=expl_mode,
name="caly-evo-step",
collect_run_caly=CollRunCaly,
prep_dp_optim=PrepCalyDPOptim,
run_dp_optim=RunCalyDPOptim,
prep_config=prep_explore_config,
run_config=run_explore_config,
upload_python_packages=upload_python_packages,
)
if expl_mode == "merge":
caly_evo_step_op = CalyEvoStepMerge(
name="caly-evo-step",
collect_run_caly=CollRunCaly,
prep_dp_optim=PrepCalyDPOptim,
run_dp_optim=RunCalyDPOptim,
prep_config=prep_explore_config,
run_config=run_explore_config,
upload_python_packages=None,
)
elif expl_mode == "default":
caly_evo_step_op = CalyEvoStep(
name="caly-evo-step",
collect_run_caly=CollRunCaly,
prep_dp_optim=PrepCalyDPOptim,
run_dp_optim=RunCalyDPOptim,
prep_config=prep_explore_config,
run_config=run_explore_config,
upload_python_packages=upload_python_packages,
)
else:
raise KeyError(f"Unknown key: {expl_mode}, support `default` and `merge`.")
prep_run_explore_op = PrepRunCaly(
"prep-run-calypso",
prep_caly_input_op=PrepCalyInput,
Expand Down
9 changes: 4 additions & 5 deletions dpgen2/op/caly_evo_step_merge.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@
Slices,
TransientError,
)
from dflow.utils import (
flatten,
)

from dpgen2.constants import (
calypso_check_opt_file,
Expand All @@ -47,7 +50,7 @@


class CalyEvoStepMerge(OP):
def __init__(self, mode="default", *args, **kwargs):
def __init__(self, mode="debug", *args, **kwargs):
self.mode = mode
self.args = args
self.kwargs = kwargs
Expand Down Expand Up @@ -115,10 +118,6 @@ def execute(
for k in step.outputs.artifacts:
path_list = download_artifact(step.outputs.artifacts[k])
if output_sign[k].type == List[Path]:
from dflow.utils import (
flatten,
)

if not isinstance(path_list, list) or any(
[p is not None and not isinstance(p, str) for p in path_list]
):
Expand Down
5 changes: 2 additions & 3 deletions dpgen2/superop/caly_evo_step.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,6 @@ def _caly_evo_step(
run_config: dict = normalize_step_dict({}),
upload_python_packages: Optional[List[os.PathLike]] = None,
):
print("run_config. = == = = =", run_config)
prep_config = deepcopy(prep_config)
run_config = deepcopy(run_config)
prep_template_config = prep_config.pop("template_config")
Expand All @@ -150,11 +149,11 @@ def _caly_evo_step(
def wise_executor(expl_mode, origin_executor):
if expl_mode == "default":
return init_executor(origin_executor)
elif expl_mode == "debug":
elif expl_mode == "merge":
return None
else:
raise NotImplementedError(
f"Unknown expl_mode {expl_mode}, only support `default` and `debug`."
f"Unknown expl_mode {expl_mode}, only support `default` and `merge`."
)

# collect the last step files and run calypso.x to generate structures
Expand Down
156 changes: 101 additions & 55 deletions dpgen2/superop/prep_run_calypso.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def __init__(
self,
name: str,
prep_caly_input_op: Type[OP],
caly_evo_step_op: Type[OP],
caly_evo_step_op,
prep_caly_model_devi_op: Type[OP],
run_caly_model_devi_op: Type[OP],
prep_config: dict = normalize_step_dict({}),
Expand Down Expand Up @@ -143,7 +143,7 @@ def _prep_run_caly(
prep_run_caly_steps: Steps,
step_keys: Dict[str, Any],
prep_caly_input_op: Type[OP],
caly_evo_step_op: Type[OP],
caly_evo_step_op,
prep_caly_model_devi_op: Type[OP],
run_caly_model_devi_op: Type[OP],
prep_config: dict = normalize_step_dict({}),
Expand Down Expand Up @@ -179,60 +179,106 @@ def _prep_run_caly(

temp_value = None
if expl_mode == "default":
caly_evo_step_merge_executor = prep_executor
caly_evo_step_merge_config = prep_config
elif expl_mode == "debug":
caly_evo_step_merge_executor = run_executor
caly_evo_step_merge_config = run_config
caly_evo_step_executor = prep_executor
caly_evo_step_config = prep_config
elif expl_mode == "merge":
caly_evo_step_executor = run_executor
caly_evo_step_config = run_config
else:
raise NotImplementedError(f"Unknown expl mode {expl_mode}")
raise NotImplementedError(f"Unknown expl mode `{expl_mode}`")

caly_evo_step_merge = Step(
"caly-evo-step",
template=PythonOPTemplate(
caly_evo_step_op,
python_packages=upload_python_packages,
**run_template_config,
),
slices=Slices(
input_parameter=[
"task_name",
],
input_artifact=[
"input_file",
"results",
"step",
"opt_results_dir",
"caly_run_opt_file",
"caly_check_opt_file",
],
output_artifact=["traj_results"],
),
parameters={
"block_id": prep_run_caly_steps.inputs.parameters["block_id"],
"expl_config": prep_run_caly_steps.inputs.parameters["explore_config"],
"task_name": prep_caly_input.outputs.parameters["task_names"],
"iter_num": "{{item}}",
},
artifacts={
"models": prep_run_caly_steps.inputs.artifacts["models"],
"input_file": prep_caly_input.outputs.artifacts["input_dat_files"],
"caly_run_opt_file": prep_caly_input.outputs.artifacts[
"caly_run_opt_files"
],
"caly_check_opt_file": prep_caly_input.outputs.artifacts[
"caly_check_opt_files"
],
"results": temp_value,
"step": temp_value,
"opt_results_dir": temp_value,
"qhull_input": temp_value,
},
key=step_keys["caly-evo-step-{{item}}"],
executor=caly_evo_step_merge_executor,
**caly_evo_step_merge_config,
)
prep_run_caly_steps.add(caly_evo_step_merge)
if expl_mode == "merge":
caly_evo_step = Step(
"caly-evo-step",
template=PythonOPTemplate(
caly_evo_step_op,
python_packages=upload_python_packages,
**run_template_config,
),
slices=Slices(
input_parameter=[
"task_name",
],
input_artifact=[
"input_file",
"results",
"step",
"opt_results_dir",
"caly_run_opt_file",
"caly_check_opt_file",
],
output_artifact=["traj_results"],
),
parameters={
"block_id": prep_run_caly_steps.inputs.parameters["block_id"],
"expl_config": prep_run_caly_steps.inputs.parameters["explore_config"],
"task_name": prep_caly_input.outputs.parameters["task_names"],
"iter_num": "{{item}}",
},
artifacts={
"models": prep_run_caly_steps.inputs.artifacts["models"],
"input_file": prep_caly_input.outputs.artifacts["input_dat_files"],
"caly_run_opt_file": prep_caly_input.outputs.artifacts[
"caly_run_opt_files"
],
"caly_check_opt_file": prep_caly_input.outputs.artifacts[
"caly_check_opt_files"
],
"results": temp_value,
"step": temp_value,
"opt_results_dir": temp_value,
"qhull_input": temp_value,
},
key=step_keys["caly-evo-step-{{item}}"],
executor=caly_evo_step_executor,
**caly_evo_step_config,
)
prep_run_caly_steps.add(caly_evo_step)
elif expl_mode == "default":
caly_evo_step = Step(
name="caly-evo-step",
template=caly_evo_step_op,
slices=Slices(
input_parameter=[
"task_name",
],
input_artifact=[
"input_file",
"results",
"step",
"opt_results_dir",
"caly_run_opt_file",
"caly_check_opt_file",
],
output_artifact=["traj_results"],
),
parameters={
"block_id": prep_run_caly_steps.inputs.parameters["block_id"],
"expl_config": prep_run_caly_steps.inputs.parameters["explore_config"],
"task_name": prep_caly_input.outputs.parameters["task_names"],
"iter_num": "{{item}}",
},
artifacts={
"models": prep_run_caly_steps.inputs.artifacts["models"],
"input_file": prep_caly_input.outputs.artifacts["input_dat_files"],
"caly_run_opt_file": prep_caly_input.outputs.artifacts[
"caly_run_opt_files"
],
"caly_check_opt_file": prep_caly_input.outputs.artifacts[
"caly_check_opt_files"
],
"results": temp_value,
"step": temp_value,
"opt_results_dir": temp_value,
"qhull_input": temp_value,
},
key=step_keys["caly-evo-step-{{item}}"],
executor=caly_evo_step_executor,
**caly_evo_step_config,
)
prep_run_caly_steps.add(caly_evo_step)
else:
raise KeyError(f"Unknown key: `{expl_mode}`, support `default` and `merge`.")

# prep_caly_model_devi
prep_caly_model_devi = Step(
Expand All @@ -247,7 +293,7 @@ def _prep_run_caly(
"template_slice_config": template_slice_config,
},
artifacts={
"traj_results": caly_evo_step_merge.outputs.artifacts["traj_results"],
"traj_results": caly_evo_step.outputs.artifacts["traj_results"],
},
key="%s--prep-caly-model-devi"
% (prep_run_caly_steps.inputs.parameters["block_id"],),
Expand Down
3 changes: 1 addition & 2 deletions tests/test_merge_caly_evo_step.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,14 +278,13 @@ def tearDown(self):

def test_caly_evo_step(self):
steps = CalyEvoStepMerge(
mode="debug",
name="caly-evo-step",
collect_run_caly=MockedCollRunCaly,
prep_dp_optim=PrepCalyDPOptim,
run_dp_optim=MockedRunCalyDPOptim,
prep_config=default_config,
run_config=default_config,
upload_python_packages=upload_python_packages,
upload_python_packages=None,
)
caly_evo_step = Step(
"caly-evo-step",
Expand Down
Loading

0 comments on commit ccbadd1

Please sign in to comment.