fix: ut

wangzyphysics · May 28, 2024 · 0628717 · 0628717
1 parent fa3558c
commit 0628717
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
@@ -179,11 +179,11 @@ def _prep_run_caly(
 
     temp_value = None
     if expl_mode == "default":
-        caly_evo_step_executor = prep_executor
         caly_evo_step_config = prep_config
+        caly_evo_step_executor = prep_executor
     elif expl_mode == "merge":
-        caly_evo_step_executor = run_executor
         caly_evo_step_config = run_config
+        caly_evo_step_executor = run_executor
     else:
         raise NotImplementedError(f"Unknown expl mode `{expl_mode}`")
 
@@ -230,7 +230,8 @@ def _prep_run_caly(
                 "qhull_input": temp_value,
             },
             key=step_keys["caly-evo-step-{{item}}"],
-            executor=caly_evo_step_executor,
+            with_param=argo_range(prep_caly_input.outputs.parameters["ntasks"]),
+            # executor=caly_evo_step_executor,
             **caly_evo_step_config,
         )
         prep_run_caly_steps.add(caly_evo_step)
@@ -273,7 +274,8 @@ def _prep_run_caly(
                 "qhull_input": temp_value,
             },
             key=step_keys["caly-evo-step-{{item}}"],
-            executor=caly_evo_step_executor,
+            with_param=argo_range(prep_caly_input.outputs.parameters["ntasks"]),
+            # executor=caly_evo_step_executor,
             **caly_evo_step_config,
         )
         prep_run_caly_steps.add(caly_evo_step)

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -160,7 +160,7 @@ def test_caly_evo_step_merge_merge_mode(self):
             run_dp_optim=MockedRunCalyDPOptim,
             prep_config=prep_default_config,
             run_config=run_default_config,
-            upload_python_packages=upload_python_packages,
+            upload_python_packages=None,
         )
         prep_run_caly_op = PrepRunCaly(
             "prep-run-calypso",