From e65147b110576f708d29899b74d7c1b918ce8b57 Mon Sep 17 00:00:00 2001
From: zhenyu wang <wzy@calypso.cn>
Date: Tue, 28 May 2024 10:59:12 +0800
Subject: [PATCH] fix: make executor flexible accorind to expl_mode (debug or
 default)

---
 dpgen2/superop/caly_evo_step.py    | 19 ++++++++++++++++---
 dpgen2/superop/prep_run_calypso.py | 14 ++++++++++++--
 tests/test_prep_run_caly.py        |  2 +-
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/dpgen2/superop/caly_evo_step.py b/dpgen2/superop/caly_evo_step.py
index 547cd1ea..072a4b2c 100644
--- a/dpgen2/superop/caly_evo_step.py
+++ b/dpgen2/superop/caly_evo_step.py
@@ -144,6 +144,17 @@ def _caly_evo_step(
     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")
+
+    def wise_executor(expl_mode, origin_executor):
+        if expl_mode == "default":
+            return origin_executor
+        elif expl_mode == "debug":
+            return None
+        else:
+            raise NotImplementedError(
+                f"Unknown expl_mode {expl_mode}, only support `default` and `debug`."
+            )
 
     # collect the last step files and run calypso.x to generate structures
     collect_run_calypso = Step(
@@ -171,7 +182,7 @@ def _caly_evo_step(
             caly_evo_step_steps.inputs.parameters["iter_num"],
             caly_evo_step_steps.inputs.parameters["cnt_num"],
         ),
-        executor=prep_executor,
+        executor=wise_executor(expl_mode, prep_executor),
         **run_config,
     )
     caly_evo_step_steps.add(collect_run_calypso)
@@ -205,7 +216,9 @@ def _caly_evo_step(
             caly_evo_step_steps.inputs.parameters["iter_num"],
             caly_evo_step_steps.inputs.parameters["cnt_num"],
         ),
-        executor=prep_executor,  # cpu is enough to run calypso.x, default step config is c2m4
+        executor=wise_executor(
+            expl_mode, prep_executor
+        ),  # cpu is enough to run calypso.x, default step config is c2m4
         **run_config,
     )
     caly_evo_step_steps.add(prep_dp_optim)
@@ -238,7 +251,7 @@ def _caly_evo_step(
             caly_evo_step_steps.inputs.parameters["iter_num"],
             caly_evo_step_steps.inputs.parameters["cnt_num"],
         ),
-        executor=run_executor,
+        executor=wise_executor(expl_mode, run_executor),
         **run_config,
     )
     caly_evo_step_steps.add(run_dp_optim)
diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
index 3ba7334e..3afb8b93 100644
--- a/dpgen2/superop/prep_run_calypso.py
+++ b/dpgen2/superop/prep_run_calypso.py
@@ -157,6 +157,7 @@ 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")
 
     # prep caly input files
     prep_caly_input = Step(
@@ -177,6 +178,15 @@ def _prep_run_caly(
     prep_run_caly_steps.add(prep_caly_input)
 
     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
+    else:
+        raise NotImplementedError(f"Unknown expl mode {expl_mode}")
+
     caly_evo_step_merge = Step(
         "caly-evo-step",
         template=PythonOPTemplate(
@@ -219,8 +229,8 @@ def _prep_run_caly(
             "qhull_input": temp_value,
         },
         key=step_keys["caly-evo-step-{{item}}"],
-        executor=prep_executor,
-        **prep_config,
+        executor=caly_evo_step_merge_executor,
+        **caly_evo_step_merge_config,
     )
     prep_run_caly_steps.add(caly_evo_step_merge)
 
diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
index 547e844d..7da80ed8 100644
--- a/tests/test_prep_run_caly.py
+++ b/tests/test_prep_run_caly.py
@@ -148,7 +148,7 @@ def tearDown(self):
         for i in Path().glob("prep-run-caly-step*"):
             shutil.rmtree(i, ignore_errors=True)
 
-    def test(self):
+    def test_caly_evo_step_merge_debug_mode(self):
         caly_evo_step_op = CalyEvoStepMerge(
             mode="debug",
             name="caly-evo-step",