fix: merge op output error

wangzyphysics · May 28, 2024 · 29e00d3 · 29e00d3
1 parent b12b8c0
commit 29e00d3
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 16 deletions.
diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
@@ -53,7 +53,7 @@ def __init__(
         self,
         name: str,
         prep_caly_input_op: Type[OP],
-        caly_evo_step_op: OPTemplate,
+        caly_evo_step_op: Type[OP],
         prep_caly_model_devi_op: Type[OP],
         run_caly_model_devi_op: Type[OP],
         prep_config: dict = normalize_step_dict({}),
@@ -143,7 +143,7 @@ def _prep_run_caly(
     prep_run_caly_steps: Steps,
     step_keys: Dict[str, Any],
     prep_caly_input_op: Type[OP],
-    caly_evo_step_op: OPTemplate,
+    caly_evo_step_op: Type[OP],
     prep_caly_model_devi_op: Type[OP],
     run_caly_model_devi_op: Type[OP],
     prep_config: dict = normalize_step_dict({}),

diff --git a/tests/test_merge_caly_evo_step.py b/tests/test_merge_caly_evo_step.py
@@ -270,12 +270,11 @@ def setUp(self):
         )
 
     def tearDown(self):
-        pass
-        # shutil.rmtree(self.work_dir, ignore_errors=True)
-        # for i in Path().glob("caly-evo-step-*"):
-        #     shutil.rmtree(i, ignore_errors=True)
-        # for i in Path().glob("caly_task*"):
-        #     shutil.rmtree(i, ignore_errors=True)
+        shutil.rmtree(self.work_dir, ignore_errors=True)
+        for i in Path().glob("caly-evo-step-*"):
+            shutil.rmtree(i, ignore_errors=True)
+        for i in Path().glob("caly_task*"):
+            shutil.rmtree(i, ignore_errors=True)
 
     def test_caly_evo_step(self):
         steps = CalyEvoStepMerge(

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -69,6 +69,9 @@
     BaseExplorationTaskGroup,
     ExplorationTask,
 )
+from dpgen2.op.caly_evo_step_merge import (
+    CalyEvoStepMerge,
+)
 from dpgen2.op.prep_caly_dp_optim import (
     PrepCalyDPOptim,
 )
@@ -81,9 +84,6 @@
 from dpgen2.op.run_caly_model_devi import (
     RunCalyModelDevi,
 )
-from dpgen2.superop.caly_evo_step import (
-    CalyEvoStep,
-)
 from dpgen2.superop.prep_run_calypso import (
     PrepRunCaly,
 )
@@ -149,11 +149,12 @@ def tearDown(self):
             shutil.rmtree(i, ignore_errors=True)
 
     def test(self):
-        caly_evo_step_op = CalyEvoStep(
-            "caly-evo-run",
-            MockedCollRunCaly,
-            PrepCalyDPOptim,
-            MockedRunCalyDPOptim,
+        caly_evo_step_op = CalyEvoStepMerge(
+            mode="debug",
+            name="caly-evo-step",
+            collect_run_caly=MockedCollRunCaly,
+            prep_dp_optim=PrepCalyDPOptim,
+            run_dp_optim=MockedRunCalyDPOptim,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,