fix: move caly related params into explore/config

wangzyphysics · May 31, 2024 · 036f3a2 · 036f3a2
1 parent ee0a4a5
commit 036f3a2
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 27 deletions.
diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py
@@ -142,6 +142,7 @@ def make_concurrent_learning_op(
     train_style: str = "dp",
     explore_style: str = "lmp",
     fp_style: str = "vasp",
+    explore_config: dict = {},
     prep_train_config: dict = default_config,
     run_train_config: dict = default_config,
     prep_explore_config: dict = default_config,
@@ -154,7 +155,7 @@ def make_concurrent_learning_op(
     upload_python_packages: Optional[List[os.PathLike]] = None,
     valid_data: Optional[S3Artifact] = None,
 ):
-    expl_mode = run_explore_config.get("mode", "default")
+    expl_mode = explore_config.get("mode", "default")
     if train_style in ("dp", "dp-dist"):
         prep_run_train_op = PrepRunDPTrain(
             "prep-run-dp-train",
@@ -205,6 +206,7 @@ def make_concurrent_learning_op(
             caly_evo_step_op=caly_evo_step_op,
             prep_caly_model_devi_op=PrepCalyModelDevi,
             run_caly_model_devi_op=RunCalyModelDevi,
+            explore_config=explore_config,
             prep_config=prep_explore_config,
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
@@ -453,6 +455,8 @@ def workflow_concurrent_learning(
 ) -> Tuple[Step, Optional[Step]]:
     default_config = config["default_step_config"]
 
+    train_config = config["train"]["config"]
+    explore_config = config["explore"]["config"]
     train_style = config["train"]["type"]
     explore_style = config["explore"]["type"]
     fp_style = config["fp"]["type"]
@@ -506,6 +510,7 @@ def workflow_concurrent_learning(
         train_style,
         explore_style,
         fp_style,
+        explore_config=explore_config,
         prep_train_config=prep_train_config,
         run_train_config=run_train_config,
         prep_explore_config=prep_explore_config,
@@ -527,8 +532,7 @@ def workflow_concurrent_learning(
         template_script = [json.loads(Path(ii).read_text()) for ii in template_script_]
     else:
         template_script = json.loads(Path(template_script_).read_text())
-    train_config = config["train"]["config"]
-    explore_config = config["explore"]["config"]
+
     if (
         "teacher_model_path" in explore_config
         and explore_config["teacher_model_path"] is not None

diff --git a/dpgen2/op/prep_caly_model_devi.py b/dpgen2/op/prep_caly_model_devi.py
@@ -48,7 +48,7 @@ def get_input_sign(cls):
         return OPIOSign(
             {
                 "task_name": Parameter(str),
-                "template_slice_config": Parameter(dict),
+                "model_devi_group_size": Parameter(int),
                 "traj_results": Artifact(List[Path]),
             }
         )
@@ -74,7 +74,7 @@ def execute(
         ip : dict
             Input dict with components:
             - `task_name` : (`str`)
-            - `template_slice_config` : (`dict`)
+            - `model_devi_group_size` : (`int`)
             - `traj_results` : (`Path`)
 
         Returns
@@ -93,9 +93,8 @@ def execute(
             for traj_dir in traj_results_dir
             for traj in Path(traj_dir).rglob("*.traj")
         ]
-        group_size = ip["template_slice_config"].get(
-            "model_devi_group_size", len(trajs)
-        )
+        model_devi_group_size = ip["model_devi_group_size"]
+        group_size = model_devi_group_size if model_devi_group_size != 0 else len(trajs)
 
         with set_directory(work_dir):
             grouped_trajs_list = [

diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
@@ -57,6 +57,7 @@ def __init__(
         caly_evo_step_op: Union[OPTemplate, OP],
         prep_caly_model_devi_op: Type[OP],
         run_caly_model_devi_op: Type[OP],
+        explore_config: dict = {},
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -114,6 +115,7 @@ def __init__(
             caly_evo_step_op,
             prep_caly_model_devi_op,
             run_caly_model_devi_op,
+            explore_config=explore_config,
             prep_config=prep_config,
             run_config=run_config,
             upload_python_packages=upload_python_packages,
@@ -147,6 +149,7 @@ def _prep_run_caly(
     caly_evo_step_op: Union[OPTemplate, OP],
     prep_caly_model_devi_op: Type[OP],
     run_caly_model_devi_op: Type[OP],
+    explore_config: dict = {},
     prep_config: dict = normalize_step_dict({}),
     run_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -158,7 +161,8 @@ def _prep_run_caly(
     prep_executor = init_executor(prep_config.pop("executor"))
     run_executor = init_executor(run_config.pop("executor"))
     template_slice_config = run_config.pop("template_slice_config", {})
-    expl_mode = run_config.pop("mode", "default")
+    expl_mode = explore_config.get("mode", "default")
+    model_devi_group_size = explore_config.get("model_devi_group_size", 0)
 
     # prep caly input files
     prep_caly_input = Step(
@@ -188,7 +192,6 @@ def _prep_run_caly(
         caly_evo_step_config = run_config
         caly_evo_step_executor = run_executor
         caly_evo_step_slice_config = deepcopy(template_slice_config)
-        caly_evo_step_slice_config.pop("model_devi_group_size", None)
         template = PythonOPTemplate(
             caly_evo_step_op,  # type: ignore
             python_packages=upload_python_packages,
@@ -252,7 +255,7 @@ def _prep_run_caly(
         ),
         parameters={
             "task_name": "prep-calypso-model-deviation",
-            "template_slice_config": template_slice_config,
+            "model_devi_group_size": model_devi_group_size,
         },
         artifacts={
             "traj_results": caly_evo_step.outputs.artifacts["traj_results"],

diff --git a/tests/op/test_prep_caly_model_devi.py b/tests/op/test_prep_caly_model_devi.py
@@ -59,8 +59,8 @@ def setUp(self):
 
         self.group_size = 5
         self.ngroup = ntrajs_dir * ntrajs_per_dir / self.group_size
-        self.template_slice_config_1 = {"model_devi_group_size": self.group_size}
-        self.template_slice_config_2 = {}
+        self.model_devi_group_size_1 = self.group_size
+        self.model_devi_group_size_2 = 0
 
     def tearDown(self):
         shutil.rmtree(self.ref_dir)
@@ -72,7 +72,7 @@ def test_00_success(self):
             OPIO(
                 {
                     "task_name": self.run_dir_name,
-                    "template_slice_config": self.template_slice_config_1,
+                    "model_devi_group_size": self.model_devi_group_size_1,
                     "traj_results": self.ref_traj_results,
                 }
             )
@@ -99,7 +99,7 @@ def test_01_success(self):
             OPIO(
                 {
                     "task_name": self.run_dir_name,
-                    "template_slice_config": self.template_slice_config_2,
+                    "model_devi_group_size_2": self.model_devi_group_size_2,
                     "traj_results": self.ref_traj_results,
                 }
             )

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -142,17 +142,13 @@ def tearDown(self):
     def test_caly_evo_step_merge_merge_mode(self):
         run_default_config = normalize_step_dict(
             {
-                "mode": "merge",
                 "template_config": {
                     "image": default_image,
                 },
-                "template_slice_config": {
-                    "group_size": 2,
-                    "pool_size": 1,
-                    "model_devi_group_size": 30,
-                },
+                "template_slice_config": {"group_size": 2, "pool_size": 1},
             }
         )
+        explore_config = {"mode": "merge", "model_devi_group_size": 30}
         caly_evo_step_op = CalyEvoStepMerge(
             name="caly-evo-step",
             collect_run_caly=MockedCollRunCaly,
@@ -168,6 +164,7 @@ def test_caly_evo_step_merge_merge_mode(self):
             caly_evo_step_op,
             PrepCalyModelDevi,
             MockedRunCalyModelDevi,
+            explore_config=explore_config,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,
@@ -200,17 +197,13 @@ def test_caly_evo_step_merge_merge_mode(self):
     def test_caly_evo_step_merge_default_mode(self):
         run_default_config = normalize_step_dict(
             {
-                "mode": "default",
                 "template_config": {
                     "image": default_image,
                 },
-                "template_slice_config": {
-                    "group_size": 2,
-                    "pool_size": 1,
-                    "model_devi_group_size": 30,
-                },
+                "template_slice_config": {"group_size": 2, "pool_size": 1},
             }
         )
+        explore_config = {"mode": "default", "model_devi_group_size": 30}
         caly_evo_step_op = CalyEvoStep(
             "caly-evo-run",
             MockedCollRunCaly,
@@ -226,6 +219,7 @@ def test_caly_evo_step_merge_default_mode(self):
             caly_evo_step_op,
             PrepCalyModelDevi,
             MockedRunCalyModelDevi,
+            explore_config=explore_config,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,