fix: expl_mode should be delivered into calyevostep

wangzyphysics · Jun 1, 2024 · 837e025 · 837e025
1 parent 0c1d8ab
commit 837e025
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py
@@ -183,6 +183,7 @@ def make_concurrent_learning_op(
                 collect_run_caly=CollRunCaly,
                 prep_dp_optim=PrepCalyDPOptim,
                 run_dp_optim=RunCalyDPOptim,
+                expl_mode=expl_mode,
                 prep_config=prep_explore_config,
                 run_config=run_explore_config,
                 upload_python_packages=None,
@@ -193,6 +194,7 @@ def make_concurrent_learning_op(
                 collect_run_caly=CollRunCaly,
                 prep_dp_optim=PrepCalyDPOptim,
                 run_dp_optim=RunCalyDPOptim,
+                expl_mode=expl_mode,
                 prep_config=prep_explore_config,
                 run_config=run_explore_config,
                 upload_python_packages=upload_python_packages,

diff --git a/dpgen2/superop/caly_evo_step.py b/dpgen2/superop/caly_evo_step.py
@@ -54,10 +54,12 @@ def __init__(
         collect_run_caly: Type[OP],
         prep_dp_optim: Type[OP],
         run_dp_optim: Type[OP],
+        expl_mode: str = "default",
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
     ):
+        self.expl_mode = expl_mode
         self._input_parameters = {
             "iter_num": InputParameter(type=int, value=0),
             "cnt_num": InputParameter(type=int, value=0),
@@ -144,7 +146,7 @@ def _caly_evo_step(
     prep_executor_config = prep_config.pop("executor")
     run_executor_config = run_config.pop("executor")
     template_slice_config = run_config.pop("template_slice_config", {})
-    expl_mode = run_config.pop("mode", "default")
+    expl_mode = caly_evo_step_steps.expl_mode
 
     def wise_executor(expl_mode, origin_executor_config):
         if expl_mode == "default":

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -155,6 +155,7 @@ def test_caly_evo_step_merge_merge_mode(self):
             collect_run_caly=MockedCollRunCaly,
             prep_dp_optim=PrepCalyDPOptim,
             run_dp_optim=MockedRunCalyDPOptim,
+            expl_mode=expl_mode,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=None,
@@ -211,6 +212,7 @@ def test_caly_evo_step_merge_default_mode(self):
             MockedCollRunCaly,
             PrepCalyDPOptim,
             MockedRunCalyDPOptim,
+            expl_mode=expl_mode,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,