diff --git a/dpgen2/flow/block.py b/dpgen2/flow/block.py index ef4df133..920a051d 100644 --- a/dpgen2/flow/block.py +++ b/dpgen2/flow/block.py @@ -27,47 +27,117 @@ from typing import Set, List from pathlib import Path -def block_cl( + +class ConcurrentLearningBlock(Steps): + def __init__( + self, + name : str, + prep_run_dp_train_op : OP, + prep_run_lmp_op : OP, + select_confs_op : OP, + prep_run_fp_op : OP, + collect_data_op : OP, + select_confs_image : str = "dflow:v1.0", + collect_data_image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters={ + "block_id" : InputParameter(), + "type_map" : InputParameter(), + "numb_models": InputParameter(type=int), + "template_script" : InputParameter(), + "train_config" : InputParameter(), + "lmp_task_grp" : InputParameter(), + "lmp_config" : InputParameter(), + "conf_selector" : InputParameter(), + "fp_inputs" : InputParameter(), + "fp_config" : InputParameter(), + } + self._input_artifacts={ + "init_models" : InputArtifact(), + "init_data" : InputArtifact(), + "iter_data" : InputArtifact(), + } + self._output_parameters={ + "exploration_report": OutputParameter(), + } + self._output_artifacts={ + "models": OutputArtifact(), + "iter_data" : OutputArtifact(), + "trajs" : OutputArtifact(), + } + + super().__init__( + name = name, + inputs = Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + + self._my_keys = ['select-confs', 'collect-data'] + self._keys = \ + prep_run_dp_train_op.keys + \ + prep_run_lmp_op.keys + \ + self._my_keys[:1] + \ + prep_run_fp_op.keys + \ + self._my_keys[1:2] + self.step_keys = {} + for ii in self._my_keys: + self.step_keys[ii] = os.path.join("%s"%self.inputs.parameters["block_id"], ii) + + self = _block_cl( + self, + self.step_keys, + name, + prep_run_dp_train_op, + prep_run_lmp_op, + select_confs_op, + prep_run_fp_op, + collect_data_op, + select_confs_image = select_confs_image, + collect_data_image = collect_data_image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + +def _block_cl( + block_steps : Steps, + step_keys : List[str], name : str, prep_run_dp_train_op : OP, prep_run_lmp_op : OP, select_confs_op : OP, prep_run_fp_op : OP, collect_data_op : OP, + select_confs_image : str = "dflow:v1.0", + collect_data_image : str = "dflow:v1.0", upload_python_package : str = None, ): - block_steps = Steps( - name = name, - inputs = Inputs( - parameters={ - "block_id" : InputParameter(), - "type_map" : InputParameter(), - "numb_models": InputParameter(type=int), - "template_script" : InputParameter(), - "train_config" : InputParameter(), - "lmp_task_grp" : InputParameter(), - "lmp_config" : InputParameter(), - "conf_selector" : InputParameter(), - "fp_inputs" : InputParameter(), - "fp_config" : InputParameter(), - }, - artifacts={ - "init_models" : InputArtifact(), - "init_data" : InputArtifact(), - "iter_data" : InputArtifact(), - }, - ), - outputs=Outputs( - parameters={ - "exploration_report": OutputParameter(), - }, - artifacts={ - "models": OutputArtifact(), - "iter_data" : OutputArtifact(), - "trajs" : OutputArtifact(), - }, - ), - ) prep_run_dp_train = Step( name + '-prep-run-dp-train', @@ -106,7 +176,7 @@ def block_cl( name = name + '-select-confs', template=PythonOPTemplate( select_confs_op, - image="dflow:v1.0", + image=select_confs_image, output_artifact_archive={ "confs": None }, @@ -121,7 +191,7 @@ def block_cl( "trajs" : prep_run_lmp.outputs.artifacts['trajs'], "model_devis" : prep_run_lmp.outputs.artifacts['model_devis'], }, - key = os.path.join("%s"%block_steps.inputs.parameters["block_id"], "select-conf"), + key = step_keys['select-confs'], ) block_steps.add(select_confs) @@ -144,7 +214,7 @@ def block_cl( name = name + '-collect-data', template=PythonOPTemplate( collect_data_op, - image="dflow:v1.0", + image=collect_data_image, output_artifact_archive={ "iter_data": None }, @@ -157,7 +227,7 @@ def block_cl( "iter_data" : block_steps.inputs.artifacts['iter_data'], "labeled_data" : prep_run_fp.outputs.artifacts['labeled_data'], }, - key = os.path.join("%s"%block_steps.inputs.parameters["block_id"], "collect-data"), + key = step_keys['collect-data'], ) block_steps.add(collect_data) diff --git a/dpgen2/flow/loop.py b/dpgen2/flow/loop.py index 707fb4c4..cc66a594 100644 --- a/dpgen2/flow/loop.py +++ b/dpgen2/flow/loop.py @@ -32,7 +32,7 @@ from dpgen2.exploration.report import ExplorationReport from dpgen2.exploration.task import ExplorationTaskGroup from dpgen2.exploration.selector import ConfSelector -from dpgen2.flow.block import block_cl +from dpgen2.flow.block import ConcurrentLearningBlock class SchedulerWrapper(OP): @@ -100,50 +100,186 @@ def execute( }) +class ConcurrentLearningLoop(Steps): + def __init__( + self, + name : str, + block_op : Steps, + image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters={ + "block_id" : InputParameter(), + "type_map" : InputParameter(), + "numb_models": InputParameter(type=int), + "template_script" : InputParameter(), + "train_config" : InputParameter(), + "lmp_task_grp" : InputParameter(), + "lmp_config" : InputParameter(), + "conf_selector" : InputParameter(), + "fp_inputs" : InputParameter(), + "fp_config" : InputParameter(), + "exploration_scheduler" : InputParameter(), + } + self._input_artifacts={ + "init_models" : InputArtifact(), + "init_data" : InputArtifact(), + "iter_data" : InputArtifact(), + } + self._output_parameters={ + "exploration_scheduler": OutputParameter(), + } + self._output_artifacts={ + "models": OutputArtifact(), + "iter_data" : OutputArtifact(), + } + + super().__init__( + name = name, + inputs = Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + + self._my_keys = ['block', 'scheduler', 'id'] + self._keys = \ + self._my_keys[:1] + \ + block_op.keys + \ + self._my_keys[1:3] + self.step_keys = {} + for ii in self._my_keys: + self.step_keys[ii] = os.path.join("%s"%self.inputs.parameters["block_id"], ii) + + self = _loop( + self, + self.step_keys, + name, + block_op, + image = image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters -def loop ( + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + +class ConcurrentLearning(Steps): + def __init__( + self, + name : str, + loop_op : Steps, + image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters={ + "type_map" : InputParameter(), + "numb_models": InputParameter(type=int), + "template_script" : InputParameter(), + "train_config" : InputParameter(), + "lmp_config" : InputParameter(), + "fp_inputs" : InputParameter(), + "fp_config" : InputParameter(), + "exploration_scheduler" : InputParameter(), + } + self._input_artifacts={ + "init_models" : InputArtifact(), + "init_data" : InputArtifact(), + "iter_data" : InputArtifact(), + } + self._output_parameters={ + "exploration_scheduler": OutputParameter(), + } + self._output_artifacts={ + "models": OutputArtifact(), + "iter_data" : OutputArtifact(), + } + + super().__init__( + name = name, + inputs = Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + + self._init_keys = ['scheduler', 'id'] + self.loop_key = 'loop' + self.step_keys = {} + for ii in self._init_keys: + self.step_keys[ii] = os.path.join('init', ii) + self.loop_op_keys = loop_op.keys + + self = _dpgen( + self, + self.step_keys, + name, + loop_op, + self.loop_key, + image = image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def init_keys(self): + return self._init_keys + + @property + def loop_keys(self): + return [self.loop_key] + self.loop_op_keys + + +def _loop ( + steps, + step_keys, name : str, block_op : OP, + image : str = "dflow:v1.0", upload_python_package : str = None, -): - steps = Steps( - name = name, - inputs = Inputs( - parameters={ - "block_id" : InputParameter(), - "type_map" : InputParameter(), - "numb_models": InputParameter(type=int), - "template_script" : InputParameter(), - "train_config" : InputParameter(), - "lmp_task_grp" : InputParameter(), - "lmp_config" : InputParameter(), - "conf_selector" : InputParameter(), - "fp_inputs" : InputParameter(), - "fp_config" : InputParameter(), - "exploration_scheduler" : InputParameter(), - }, - artifacts={ - "init_models" : InputArtifact(), - "init_data" : InputArtifact(), - "iter_data" : InputArtifact(), - }, - ), - outputs=Outputs( - parameters={ - "exploration_scheduler": OutputParameter(), - }, - artifacts={ - "models": OutputArtifact(), - "iter_data" : OutputArtifact(), - }, - ), - ) - - # suffix = steps.inputs.parameters["name_suffix"].value - suffix='' - +): block_step = Step( - name = name + suffix + '-block', + name = name + '-block', template = block_op, parameters={ "block_id" : steps.inputs.parameters["block_id"], @@ -162,15 +298,15 @@ def loop ( "init_data": steps.inputs.artifacts["init_data"], "iter_data": steps.inputs.artifacts["iter_data"], }, - key = os.path.join("%s"%steps.inputs.parameters["block_id"], "block"), + key = step_keys['block'], ) steps.add(block_step) scheduler_step = Step( - name = name + suffix + '-scheduler', + name = name + '-scheduler', template=PythonOPTemplate( SchedulerWrapper, - image="dflow:v1.0", + image=image, python_packages = upload_python_package, ), parameters={ @@ -180,7 +316,7 @@ def loop ( artifacts={ "trajs" : block_step.outputs.artifacts['trajs'], }, - key = os.path.join("%s"%steps.inputs.parameters["block_id"], "scheduler"), + key = step_keys['scheduler'], ) steps.add(scheduler_step) @@ -188,7 +324,7 @@ def loop ( name = name + '-make-block-id', template=PythonOPTemplate( MakeBlockId, - image="dflow:v1.0", + image=image, python_packages = upload_python_package, ), parameters={ @@ -196,7 +332,7 @@ def loop ( }, artifacts={ }, - key = os.path.join("%s"%steps.inputs.parameters["block_id"], "id"), + key = step_keys['id'], ) steps.add(id_step) @@ -247,46 +383,20 @@ def loop ( return steps -def dpgen( +def _dpgen( + steps, + step_keys, name, loop_op, + loop_key, + image = "dflow:v1.0", upload_python_package : str = None ): - steps = Steps( - name = name, - inputs = Inputs( - parameters={ - "type_map" : InputParameter(), - "numb_models": InputParameter(type=int), - "template_script" : InputParameter(), - "train_config" : InputParameter(), - "lmp_config" : InputParameter(), - "fp_inputs" : InputParameter(), - "fp_config" : InputParameter(), - "exploration_scheduler" : InputParameter(), - }, - artifacts={ - "init_models" : InputArtifact(), - "init_data" : InputArtifact(), - "iter_data" : InputArtifact(), - }, - ), - outputs=Outputs( - parameters={ - "exploration_scheduler": OutputParameter(), - }, - artifacts={ - "models": OutputArtifact(), - "iter_data" : OutputArtifact(), - }, - ), - ) - scheduler_step = Step( name = name + '-scheduler', template=PythonOPTemplate( SchedulerWrapper, - image="dflow:v1.0", + image=image, python_packages = upload_python_package, ), parameters={ @@ -296,7 +406,7 @@ def dpgen( artifacts={ "trajs" : None, }, - key = os.path.join("init", "scheduler"), + key = step_keys['scheduler'], ) steps.add(scheduler_step) @@ -304,7 +414,7 @@ def dpgen( name = name + '-make-block-id', template=PythonOPTemplate( MakeBlockId, - image="dflow:v1.0", + image=image, python_packages = upload_python_package, ), parameters={ @@ -312,7 +422,7 @@ def dpgen( }, artifacts={ }, - key = os.path.join("init", "id"), + key = step_keys['id'], ) steps.add(id_step) @@ -337,6 +447,7 @@ def dpgen( "init_data": steps.inputs.artifacts["init_data"], "iter_data": steps.inputs.artifacts["iter_data"], }, + key = os.path.join("%s"%id_step.outputs.parameters['block_id'], loop_key), ) steps.add(loop_step) diff --git a/dpgen2/flow/prep_run_dp_train.py b/dpgen2/flow/prep_run_dp_train.py index 4679a221..342d31b5 100644 --- a/dpgen2/flow/prep_run_dp_train.py +++ b/dpgen2/flow/prep_run_dp_train.py @@ -24,40 +24,93 @@ import os from typing import Set, List from pathlib import Path +from . import FlowIO +class PrepRunDPTrain(Steps): + def __init__( + self, + name : str, + prep_train_op : OP, + run_train_op : OP, + prep_train_image : str = "dflow:v1.0", + run_train_image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters = { + "block_id" : InputParameter(type=str, value=""), + "numb_models": InputParameter(type=int), + "template_script" : InputParameter(), + "train_config" : InputParameter(), + } + self._input_artifacts = { + "init_models" : InputArtifact(), + "init_data" : InputArtifact(), + "iter_data" : InputArtifact(), + } + self._output_parameters = {} + self._output_artifacts = { + "scripts": OutputArtifact(), + "models": OutputArtifact(), + "logs": OutputArtifact(), + "lcurves": OutputArtifact(), + } -def prep_run_dp_train( - name : str, + super().__init__( + name=name, + inputs=Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + artifacts=self._output_artifacts, + ), + ) + + self._keys = ['prep-train', 'run-train'] + self.step_keys = {} + for ii in self._keys: + self.step_keys[ii] = os.path.join("%s"%self.inputs.parameters["block_id"], ii) + + self = _prep_run_dp_train( + self, + self.step_keys, + prep_train_op, + run_train_op, + prep_train_image = prep_train_image, + run_train_image = run_train_image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + +def _prep_run_dp_train( + train_steps, + step_keys, prep_train_op : OP, run_train_op : OP, prep_train_image : str = "dflow:v1.0", run_train_image : str = "dflow:v1.0", upload_python_package : str = None, ): - train_steps = Steps( - name=name, - inputs=Inputs( - parameters={ - "block_id" : InputParameter(type=str, value=""), - "numb_models": InputParameter(type=int), - "template_script" : InputParameter(), - "train_config" : InputParameter(), - }, - artifacts={ - "init_models" : InputArtifact(), - "init_data" : InputArtifact(), - "iter_data" : InputArtifact(), - }, - ), - outputs=Outputs( - artifacts={ - "scripts": OutputArtifact(), - "models": OutputArtifact(), - "logs": OutputArtifact(), - "lcurves": OutputArtifact(), - }), - ) - prep_train = Step( 'prep-train', template=PythonOPTemplate( @@ -74,7 +127,7 @@ def prep_run_dp_train( }, artifacts={ }, - key = os.path.join("%s"%train_steps.inputs.parameters["block_id"], "prep-train"), + key = step_keys['prep-train'], ) train_steps.add(prep_train) @@ -102,7 +155,7 @@ def prep_run_dp_train( "iter_data": train_steps.inputs.artifacts['iter_data'], }, with_param=argo_range(train_steps.inputs.parameters["numb_models"]), - key = os.path.join("%s"%train_steps.inputs.parameters["block_id"], "run-train"), + key = step_keys['run-train'], ) train_steps.add(run_train) diff --git a/dpgen2/flow/prep_run_fp.py b/dpgen2/flow/prep_run_fp.py index cfda7987..59fda845 100644 --- a/dpgen2/flow/prep_run_fp.py +++ b/dpgen2/flow/prep_run_fp.py @@ -27,39 +27,96 @@ from typing import Set, List from pathlib import Path -def prep_run_fp( - name : str, + +class PrepRunFp(Steps): + def __init__( + self, + name : str, + prep_op : OP, + run_op : OP, + prep_image : str = "dflow:v1.0", + run_image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters = { + "block_id" : InputParameter(type=str, value=""), + "inputs": InputParameter(), + "fp_config" : InputParameter(), + } + self._input_artifacts = { + "confs" : InputArtifact() + } + self._output_parameters = { + "task_names": OutputParameter(), + } + self._output_artifacts = { + "logs": OutputArtifact(), + "labeled_data": OutputArtifact(), + } + + super().__init__( + name=name, + inputs=Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + + self._keys = ['prep-fp', 'run-fp'] + self.step_keys = {} + for ii in self._keys: + self.step_keys[ii] = os.path.join("%s"%self.inputs.parameters["block_id"], ii) + + self = _prep_run_fp( + self, + self.step_keys, + prep_op, + run_op, + prep_image = prep_image, + run_image = run_image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + + +def _prep_run_fp( + prep_run_steps, + step_keys, prep_op : OP, run_op : OP, + prep_image : str = "dflow:v1.0", + run_image : str = "dflow:v1.0", upload_python_package : str = None, ): - prep_run_steps = Steps( - name=name, - inputs=Inputs( - parameters={ - "block_id" : InputParameter(type=str, value=""), - "inputs": InputParameter(), - "fp_config" : InputParameter(), - }, - artifacts={ - "confs" : InputArtifact() - }, - ), - outputs=Outputs( - parameters={ - "task_names": OutputParameter(), - }, - artifacts={ - "logs": OutputArtifact(), - "labeled_data": OutputArtifact(), - }), - ) - prep_fp = Step( 'prep-fp', template=PythonOPTemplate( prep_op, - image="dflow:v1.0", + image=prep_image, output_artifact_archive={ "task_paths": None }, @@ -71,7 +128,7 @@ def prep_run_fp( artifacts={ "confs" : prep_run_steps.inputs.artifacts['confs'], }, - key = os.path.join("%s"%prep_run_steps.inputs.parameters["block_id"], "prep-fp"), + key = step_keys['prep-fp'], ) prep_run_steps.add(prep_fp) @@ -79,7 +136,7 @@ def prep_run_fp( 'run-fp', template=PythonOPTemplate( run_op, - image="dflow:v1.0", + image=run_image, slices = Slices( "{{item}}", input_parameter = ["task_name"], @@ -96,7 +153,7 @@ def prep_run_fp( 'task_path' : prep_fp.outputs.artifacts['task_paths'], }, with_param=argo_range(argo_len(prep_fp.outputs.parameters["task_names"])), - key = os.path.join("%s"%prep_run_steps.inputs.parameters["block_id"], "run-fp"), + key = step_keys['run-fp'], ) prep_run_steps.add(run_fp) diff --git a/dpgen2/flow/prep_run_lmp.py b/dpgen2/flow/prep_run_lmp.py index 8fa5d2a1..4ed8612f 100644 --- a/dpgen2/flow/prep_run_lmp.py +++ b/dpgen2/flow/prep_run_lmp.py @@ -28,40 +28,96 @@ from pathlib import Path -def prep_run_lmp( - name : str, +class PrepRunLmp(Steps): + def __init__( + self, + name : str, + prep_op : OP, + run_op : OP, + prep_image : str = "dflow:v1.0", + run_image : str = "dflow:v1.0", + upload_python_package : str = None, + ): + self._input_parameters = { + "block_id" : InputParameter(type=str, value=""), + "lmp_task_grp": InputParameter(type=int), + "lmp_config" : InputParameter() + } + self._input_artifacts = { + "models" : InputArtifact() + } + self._output_parameters={ + "task_names": OutputParameter(), + } + self._output_artifacts={ + "logs": OutputArtifact(), + "trajs": OutputArtifact(), + "model_devis": OutputArtifact(), + } + + super().__init__( + name=name, + inputs=Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + + self._keys = ['prep-lmp', 'run-lmp'] + self.step_keys = {} + for ii in self._keys: + self.step_keys[ii] = os.path.join("%s"%self.inputs.parameters["block_id"], ii) + + self = _prep_run_lmp( + self, + self.step_keys, + prep_op, + run_op, + prep_image = prep_image, + run_image = run_image, + upload_python_package = upload_python_package, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + +def _prep_run_lmp( + prep_run_steps, + step_keys, prep_op : OP, run_op : OP, + prep_image : str = "dflow:v1.0", + run_image : str = "dflow:v1.0", upload_python_package : str = None, ): - prep_run_steps = Steps( - name=name, - inputs=Inputs( - parameters={ - "block_id" : InputParameter(type=str, value=""), - "lmp_task_grp": InputParameter(type=int), - "lmp_config" : InputParameter() - }, - artifacts={ - "models" : InputArtifact() - }, - ), - outputs=Outputs( - parameters={ - "task_names": OutputParameter(), - }, - artifacts={ - "logs": OutputArtifact(), - "trajs": OutputArtifact(), - "model_devis": OutputArtifact(), - }), - ) prep_lmp = Step( 'prep-lmp', template=PythonOPTemplate( prep_op, - image="dflow:v1.0", + image=prep_image, output_artifact_archive={ "task_paths": None }, @@ -72,7 +128,7 @@ def prep_run_lmp( }, artifacts={ }, - key = os.path.join("%s" % prep_run_steps.inputs.parameters["block_id"], "prep-lmp"), + key = step_keys['prep-lmp'], ) prep_run_steps.add(prep_lmp) @@ -80,7 +136,7 @@ def prep_run_lmp( 'run-lmp', template=PythonOPTemplate( run_op, - image="dflow:v1.0", + image=run_image, slices = Slices( "{{item}}", input_parameter = ["task_name"], @@ -99,7 +155,7 @@ def prep_run_lmp( }, # with_sequence=argo_sequence(argo_len(prep_lmp.outputs.parameters["task_names"])), with_param=argo_range(argo_len(prep_lmp.outputs.parameters["task_names"])), - key = os.path.join("%s" % prep_run_steps.inputs.parameters["block_id"], "run-lmp"), + key = step_keys['run-lmp'], ) prep_run_steps.add(run_lmp) diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py index 22ada85b..7fbada9a 100644 --- a/tests/test_block_cl.py +++ b/tests/test_block_cl.py @@ -36,10 +36,10 @@ pass from context import upload_python_package from dpgen2.op.prep_lmp import PrepLmp -from dpgen2.flow.prep_run_dp_train import prep_run_dp_train -from dpgen2.flow.prep_run_lmp import prep_run_lmp -from dpgen2.flow.prep_run_fp import prep_run_fp -from dpgen2.flow.block import block_cl +from dpgen2.flow.prep_run_dp_train import PrepRunDPTrain +from dpgen2.flow.prep_run_lmp import PrepRunLmp +from dpgen2.flow.prep_run_fp import PrepRunFp +from dpgen2.flow.block import ConcurrentLearningBlock from dpgen2.exploration.task import ExplorationTask, ExplorationTaskGroup from dpgen2.fp.vasp import VaspInputs @@ -78,19 +78,19 @@ class TestBlockCL(unittest.TestCase): def _setUp_ops(self): - self.prep_run_dp_train_op = prep_run_dp_train( + self.prep_run_dp_train_op = PrepRunDPTrain( "prep-run-dp-train", MockedPrepDPTrain, MockedRunDPTrain, upload_python_package = upload_python_package, ) - self.prep_run_lmp_op = prep_run_lmp( + self.prep_run_lmp_op = PrepRunLmp( "prep-run-lmp", PrepLmp, MockedRunLmp, upload_python_package = upload_python_package, ) - self.prep_run_fp_op = prep_run_fp( + self.prep_run_fp_op = PrepRunFp( "prep-run-fp", MockedPrepVasp, MockedRunVasp, @@ -134,7 +134,7 @@ def setUp(self): self.name = 'iter-002' self._setUp_ops() self._setUp_data() - self.block_cl = block_cl( + self.block_cl = ConcurrentLearningBlock( self.name, self.prep_run_dp_train_op, self.prep_run_lmp_op, @@ -142,7 +142,7 @@ def setUp(self): self.prep_run_fp_op, MockedCollectData, upload_python_package = upload_python_package, - ) + ) def tearDown(self): for ii in ['init_data', 'iter_data', 'iter-000', 'iter-001', 'models']: @@ -155,6 +155,13 @@ def tearDown(self): os.remove(name) def test(self): + self.assertEqual( + self.block_cl.keys, + ['prep-train', 'run-train', 'prep-lmp', 'run-lmp', 'select-confs', + 'prep-fp', 'run-fp', 'collect-data', + ] + ) + block_step = Step( 'step', template = self.block_cl, diff --git a/tests/test_loop.py b/tests/test_loop.py index 6c5070b0..8f440a70 100644 --- a/tests/test_loop.py +++ b/tests/test_loop.py @@ -44,13 +44,13 @@ ExplorationScheduler, ) from dpgen2.op.prep_lmp import PrepLmp -from dpgen2.flow.prep_run_dp_train import prep_run_dp_train -from dpgen2.flow.prep_run_lmp import prep_run_lmp -from dpgen2.flow.prep_run_fp import prep_run_fp -from dpgen2.flow.block import block_cl +from dpgen2.flow.prep_run_dp_train import PrepRunDPTrain +from dpgen2.flow.prep_run_lmp import PrepRunLmp +from dpgen2.flow.prep_run_fp import PrepRunFp +from dpgen2.flow.block import ConcurrentLearningBlock from dpgen2.exploration.task import ExplorationTask, ExplorationTaskGroup from dpgen2.fp.vasp import VaspInputs -from dpgen2.flow.loop import dpgen, loop +from dpgen2.flow.loop import ConcurrentLearning, ConcurrentLearningLoop from dpgen2.exploration.report import ExplorationReport from dpgen2.exploration.task import ExplorationTaskGroup, ExplorationStage from dpgen2.exploration.selector import TrustLevelConfSelector, TrustLevel @@ -95,25 +95,25 @@ class TestLoop(unittest.TestCase): def _setUp_ops(self): - self.prep_run_dp_train_op = prep_run_dp_train( + self.prep_run_dp_train_op = PrepRunDPTrain( "prep-run-dp-train", MockedPrepDPTrain, MockedRunDPTrain, upload_python_package = upload_python_package, ) - self.prep_run_lmp_op = prep_run_lmp( + self.prep_run_lmp_op = PrepRunLmp( "prep-run-lmp", PrepLmp, MockedRunLmp, upload_python_package = upload_python_package, ) - self.prep_run_fp_op = prep_run_fp( + self.prep_run_fp_op = PrepRunFp( "prep-run-fp", MockedPrepVasp, MockedRunVasp, upload_python_package = upload_python_package, ) - self.block_cl_op = block_cl( + self.block_cl_op = ConcurrentLearningBlock( self.name+'-block', self.prep_run_dp_train_op, self.prep_run_lmp_op, @@ -122,12 +122,12 @@ def _setUp_ops(self): MockedCollectData, upload_python_package = upload_python_package, ) - self.loop_op = loop( + self.loop_op = ConcurrentLearningLoop( self.name+'-loop', self.block_cl_op, upload_python_package = upload_python_package, ) - self.dpgen_op = dpgen( + self.dpgen_op = ConcurrentLearning( self.name, self.loop_op, upload_python_package = upload_python_package, @@ -202,6 +202,22 @@ def tearDown(self): os.remove(name) def test(self): + self.assertEqual( + self.dpgen_op.loop_keys, + [ + "loop", 'block', + 'prep-train', 'run-train', 'prep-lmp', 'run-lmp', 'select-confs', + 'prep-fp', 'run-fp', 'collect-data', + "scheduler", "id", + ] + ) + self.assertEqual( + self.dpgen_op.init_keys, + [ + "scheduler", "id", + ] + ) + dpgen_step = Step( 'dpgen-step', template = self.dpgen_op, diff --git a/tests/test_prep_run_dp_train.py b/tests/test_prep_run_dp_train.py index 220626dd..fbb4f25a 100644 --- a/tests/test_prep_run_dp_train.py +++ b/tests/test_prep_run_dp_train.py @@ -34,7 +34,7 @@ # case of upload everything to argo, no context needed pass from context import upload_python_package -from dpgen2.flow.prep_run_dp_train import prep_run_dp_train +from dpgen2.flow.prep_run_dp_train import PrepRunDPTrain from dpgen2.constants import train_task_pattern from mocked_ops import ( mocked_template_script, @@ -258,7 +258,7 @@ def tearDown(self): def test_train(self): - steps = prep_run_dp_train( + steps = PrepRunDPTrain( "train-steps", MockedPrepDPTrain, MockedRunDPTrain, diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py index 79b534a6..5719f165 100644 --- a/tests/test_prep_run_lmp.py +++ b/tests/test_prep_run_lmp.py @@ -35,7 +35,7 @@ pass from context import upload_python_package from dpgen2.op.prep_lmp import PrepLmp -from dpgen2.flow.prep_run_lmp import prep_run_lmp +from dpgen2.flow.prep_run_lmp import PrepRunLmp from dpgen2.exploration.task import ExplorationTask, ExplorationTaskGroup from mocked_ops import ( mocked_numb_models, @@ -219,7 +219,7 @@ def check_run_lmp_output( def test(self): - steps = prep_run_lmp( + steps = PrepRunLmp( "prep-run-lmp", PrepLmp, MockedRunLmp, diff --git a/tests/test_prep_run_vasp.py b/tests/test_prep_run_vasp.py index cc7606bf..16aff8c9 100644 --- a/tests/test_prep_run_vasp.py +++ b/tests/test_prep_run_vasp.py @@ -34,7 +34,7 @@ # case of upload everything to argo, no context needed pass from context import upload_python_package -from dpgen2.flow.prep_run_fp import prep_run_fp +from dpgen2.flow.prep_run_fp import PrepRunFp from mocked_ops import ( mocked_incar_template, MockedPrepVasp, @@ -182,7 +182,7 @@ def check_run_vasp_output( os.chdir(cwd) def test(self): - steps = prep_run_fp( + steps = PrepRunFp( "prep-run-vasp", MockedPrepVasp, MockedRunVasp,