init_reaction: fix compatibility with new dpdispatcher (#755)

fix compatibility as the key was changed in the dpdispatcher
deepmodeling · Jun 14, 2022 · aca0d32 · aca0d32
1 parent 800fb51
commit aca0d32
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/dpgen/data/reaction.py b/dpgen/data/reaction.py
@@ -107,15 +107,18 @@ def link_trj(jdata):
 
 def run_build_dataset(jdata, mdata, log_file="build_log"):
     work_path = build_path
+    # compatible with new dpdispatcher and old dpgen.dispatcher
+    build_ntasks = mdata["build_resources"].get("cpu_per_node", mdata["build_resources"]["task_per_node"])
+    fp_ntasks = mdata["fp_resources"].get("cpu_per_node", mdata["fp_resources"]["task_per_node"])
     build_command = "{cmd} -n {dataset_name} -a {type_map} -d {lammpstrj} -c {cutoff} -s {dataset_size} -k \"{qmkeywords}\" --nprocjob {nprocjob} --nproc {nproc}".format(
         cmd=mdata["build_command"],
         type_map=" ".join(jdata["type_map"]),
         lammpstrj=trj_path,
         cutoff=jdata["cutoff"],
         dataset_size=jdata["dataset_size"],
         qmkeywords=jdata["qmkeywords"],
-        nprocjob=mdata["fp_resources"]["task_per_node"],
-        nproc=mdata["build_resources"]["task_per_node"],
+        nprocjob=fp_ntasks,
+        nproc=build_ntasks,
         dataset_name=dataset_name
     )
     run_tasks = glob.glob(os.path.join(work_path, 'task.*'))