Convert lmp_task_group from parameter to artifact (deepmodeling#11)

* remove ExplorationGroup, rename CPTGroup -> NPTTaskGroup Signed-off-by: Han Wang <[email protected]> * convert lmp_task_group from parameter to artifact Signed-off-by: Han Wang <[email protected]> Co-authored-by: Han Wang <[email protected]>
wanghan-iapcm · Feb 28, 2022 · 1541d77 · 1541d77
1 parent 0b08635
commit 1541d77
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 21 deletions.
diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py
@@ -23,7 +23,7 @@
     Artifact,
     Slices,
 )
-import jsonpickle, os
+import pickle, jsonpickle, os
 from typing import (
     List
 )
@@ -49,7 +49,7 @@ def get_output_sign(cls):
         return OPIOSign({
             "exploration_scheduler" : ExplorationScheduler,
             "converged" : bool,
-            "lmp_task_grp" : ExplorationTaskGroup,
+            "lmp_task_grp" : Artifact(Path),
             "conf_selector" : ConfSelector,
         })
 
@@ -63,12 +63,15 @@ def execute(
         trajs = ip['trajs']
 
         conv, lmp_task_grp, selector = scheduler.plan_next_iteration(report, trajs)
-
+
+        with open('lmp_task_grp.dat', 'wb') as fp:
+            pickle.dump(lmp_task_grp, fp)
+
         return OPIO({
             "exploration_scheduler" : scheduler,
             "converged" : conv,
-            "lmp_task_grp" : lmp_task_grp,
             "conf_selector" : selector,
+            "lmp_task_grp" : Path('lmp_task_grp.dat'),
         })
 
 
@@ -114,7 +117,6 @@ def __init__(
             "numb_models": InputParameter(type=int),
             "template_script" : InputParameter(),
             "train_config" : InputParameter(),
-            "lmp_task_grp" : InputParameter(),
             "lmp_config" : InputParameter(),
             "conf_selector" : InputParameter(),
             "fp_inputs" : InputParameter(),
@@ -125,6 +127,7 @@ def __init__(
             "init_models" : InputArtifact(),
             "init_data" : InputArtifact(),
             "iter_data" : InputArtifact(),
+            "lmp_task_grp" : InputArtifact(),
         }
         self._output_parameters={
             "exploration_scheduler": OutputParameter(),
@@ -295,13 +298,13 @@ def _loop (
             "numb_models" : steps.inputs.parameters["numb_models"],
             "template_script" : steps.inputs.parameters["template_script"],
             "train_config" : steps.inputs.parameters["train_config"],
-            "lmp_task_grp" : steps.inputs.parameters["lmp_task_grp"],
             "lmp_config" : steps.inputs.parameters["lmp_config"],
             "conf_selector" : steps.inputs.parameters["conf_selector"],
             "fp_inputs" : steps.inputs.parameters["fp_inputs"],            
             "fp_config" : steps.inputs.parameters["fp_config"],
         },
         artifacts={
+            "lmp_task_grp" : steps.inputs.artifacts["lmp_task_grp"],
             "init_models": steps.inputs.artifacts["init_models"],
             "init_data": steps.inputs.artifacts["init_data"],
             "iter_data": steps.inputs.artifacts["iter_data"],
@@ -353,14 +356,14 @@ def _loop (
             "numb_models" : steps.inputs.parameters["numb_models"],
             "template_script" : steps.inputs.parameters["template_script"],
             "train_config" : steps.inputs.parameters["train_config"],
-            "lmp_task_grp" : scheduler_step.outputs.parameters["lmp_task_grp"],
             "lmp_config" : steps.inputs.parameters["lmp_config"],
             "conf_selector" : scheduler_step.outputs.parameters["conf_selector"],
             "fp_inputs" : steps.inputs.parameters["fp_inputs"],
             "fp_config" : steps.inputs.parameters["fp_config"],
             "exploration_scheduler" : scheduler_step.outputs.parameters["exploration_scheduler"],
         },
         artifacts={
+            "lmp_task_grp" : scheduler_step.outputs.artifacts["lmp_task_grp"],
             "init_models" : block_step.outputs.artifacts['models'],
             "init_data" : steps.inputs.artifacts['init_data'],
             "iter_data" : block_step.outputs.artifacts['iter_data'],
@@ -443,14 +446,14 @@ def _dpgen(
             "numb_models" : steps.inputs.parameters['numb_models'],
             "template_script" : steps.inputs.parameters['template_script'],
             "train_config" : steps.inputs.parameters['train_config'],
-            "lmp_task_grp" : scheduler_step.outputs.parameters['lmp_task_grp'],
             "conf_selector" : scheduler_step.outputs.parameters['conf_selector'],
             "lmp_config" : steps.inputs.parameters['lmp_config'],
             "fp_inputs" : steps.inputs.parameters['fp_inputs'],
             "fp_config" : steps.inputs.parameters['fp_config'],
             "exploration_scheduler" : scheduler_step.outputs.parameters['exploration_scheduler'],
         },
         artifacts={
+            "lmp_task_grp" : scheduler_step.outputs.artifacts['lmp_task_grp'],
             "init_models": steps.inputs.artifacts["init_models"],
             "init_data": steps.inputs.artifacts["init_data"],
             "iter_data": steps.inputs.artifacts["iter_data"],

diff --git a/dpgen2/op/prep_lmp.py b/dpgen2/op/prep_lmp.py
@@ -5,7 +5,7 @@
     Artifact
 )
 
-import json
+import json, pickle
 from typing import Tuple, List
 from pathlib import Path
 from dpgen2.exploration.task import ExplorationTaskGroup
@@ -27,7 +27,7 @@ class PrepLmp(OP):
     @classmethod
     def get_input_sign(cls):
         return OPIOSign({
-            "lmp_task_grp": ExplorationTaskGroup,
+            "lmp_task_grp": Artifact(Path),
         })
 
     @classmethod
@@ -48,8 +48,7 @@ def execute(
         ----------
         ip : dict
             Input dict with components:
-
-            - `lmp_task_grp` : (`ExplorationTaskGroup`) Definitions for LAMMPS tasks
+            - `lmp_task_grp` : (`Artifact(Path)`) Can be pickle loaded as a ExplorationTaskGroup. Definitions for LAMMPS tasks
         
         Returns
         -------
@@ -60,7 +59,8 @@ def execute(
             - `task_paths`: (`Artifact(List[Path])`) The parepared working paths of the tasks. Contains all input files needed to start the LAMMPS simulation. The order fo the Paths should be consistent with `op["task_names"]`
         """
 
-        lmp_task_grp = ip['lmp_task_grp']
+        with open(ip['lmp_task_grp'], 'rb') as fp:
+            lmp_task_grp = pickle.load(fp)
         cc = 0
         task_paths = []
         for tt in lmp_task_grp:

diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py
@@ -47,13 +47,13 @@ def __init__(
             "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={
+            "lmp_task_grp" : InputArtifact(),
             "init_models" : InputArtifact(),
             "init_data" : InputArtifact(),
             "iter_data" : InputArtifact(),
@@ -164,10 +164,10 @@ def _block_cl(
         template = prep_run_lmp_op,
         parameters={
             "block_id" : block_steps.inputs.parameters['block_id'],
-            "lmp_task_grp": block_steps.inputs.parameters['lmp_task_grp'],
             "lmp_config": block_steps.inputs.parameters['lmp_config'],
         },
         artifacts={
+            "lmp_task_grp": block_steps.inputs.artifacts['lmp_task_grp'],
             "models" : prep_run_dp_train.outputs.artifacts['models'],
         },
         key = '--'.join(["%s"%block_steps.inputs.parameters["block_id"], "prep-run-lmp"]),

diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py
@@ -42,10 +42,10 @@ def __init__(
     ):
         self._input_parameters = {
             "block_id" : InputParameter(type=str, value=""),
-            "lmp_task_grp": InputParameter(type=int),
             "lmp_config" : InputParameter()
         }
         self._input_artifacts = {
+            "lmp_task_grp": InputArtifact(),
             "models" : InputArtifact()
         }
         self._output_parameters={
@@ -132,9 +132,9 @@ def _prep_run_lmp(
             python_packages = upload_python_package,
         ),
         parameters={
-            "lmp_task_grp": prep_run_steps.inputs.parameters['lmp_task_grp'],
         },
         artifacts={
+            "lmp_task_grp": prep_run_steps.inputs.artifacts['lmp_task_grp'],
         },
         key = step_keys['prep-lmp'],
     )

diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py
@@ -26,7 +26,7 @@
     upload_packages,
 )
 
-import time, shutil, json, jsonpickle
+import time, shutil, json, jsonpickle, pickle
 from typing import Set, List
 from pathlib import Path
 try:
@@ -119,6 +119,9 @@ def _setUp_data(self):
         self.template_script = mocked_template_script
 
         self.task_group_list = MockedExplorationTaskGroup()
+        with open('lmp_task_grp.dat', 'wb') as fp:
+            pickle.dump(self.task_group_list, fp)
+        self.task_group_list = upload_artifact('lmp_task_grp.dat')
 
         self.conf_selector = MockedConfSelector()
         self.type_map = []
@@ -171,13 +174,13 @@ def test(self):
                 "numb_models" : self.numb_models,
                 "template_script" : self.template_script,
                 "train_config" : {},
-                "lmp_task_grp" : self.task_group_list,
                 "lmp_config" : {},
                 "conf_selector" : self.conf_selector,
                 'fp_inputs' : self.vasp_inputs,
                 "fp_config" : {},
             },
             artifacts = {
+                "lmp_task_grp" : self.task_group_list,
                 "init_models" : self.init_models,
                 "init_data" : self.init_data,
                 "iter_data" : self.iter_data,

diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py
@@ -25,7 +25,7 @@
     Artifact,
 )
 
-import time, shutil, json, jsonpickle
+import time, shutil, json, jsonpickle, pickle
 from typing import Set, List
 from pathlib import Path
 try:
@@ -89,12 +89,18 @@ def setUp(self):
         self.ngrp = 2
         self.ntask_per_grp = 3
         self.task_group_list = make_task_group_list(self.ngrp, self.ntask_per_grp)
+        with open('lmp_task_grp.dat', 'wb') as fp:
+            pickle.dump(self.task_group_list, fp)
+        self.task_group_list = Path('lmp_task_grp.dat')
 
     def tearDown(self):
         for ii in range(self.ngrp * self.ntask_per_grp):
             work_path = Path(lmp_task_pattern % ii)
             if work_path.is_dir():
                 shutil.rmtree(work_path)
+        task_grp = Path('lmp_task_grp.dat')
+        if task_grp.is_file():
+            os.remove(task_grp)
 
     def test(self):
         op = PrepLmp()
@@ -178,6 +184,9 @@ def setUp(self):
         self.ngrp = 2
         self.ntask_per_grp = 3
         self.task_group_list = make_task_group_list(self.ngrp, self.ntask_per_grp)
+        with open('lmp_task_grp.dat', 'wb') as fp:
+            pickle.dump(self.task_group_list, fp)
+        self.task_group_list = upload_artifact('lmp_task_grp.dat')
         self.nmodels = mocked_numb_models
         self.model_list = []
         for ii in range(self.nmodels):
@@ -195,6 +204,9 @@ def tearDown(self):
             work_path = Path(f'task.{ii:06d}')
             if work_path.is_dir():
                 shutil.rmtree(work_path)
+        task_grp = Path('lmp_task_grp.dat')
+        if task_grp.is_file():
+            os.remove(task_grp)
 
 
     def check_run_lmp_output(
@@ -229,10 +241,10 @@ def test(self):
             'prep-run-step', 
             template = steps,
             parameters = {
-                "lmp_task_grp" : self.task_group_list,
                 "lmp_config" : {},
             },
             artifacts = {
+                "lmp_task_grp" : self.task_group_list,
                 "models" : self.models,
             },
         )