Simplify API version validation

dpgen/data/gen.py

@@ -10,7 +10,6 @@
 
 import dpdata
 import numpy as np
-from packaging.version import Version
 from pymatgen.core import Structure
 from pymatgen.io.vasp import Incar
 
@@ -28,7 +27,7 @@
     make_abacus_scf_stru,
     make_supercell_abacus,
 )
-from dpgen.generator.lib.utils import symlink_user_forward_files
+from dpgen.generator.lib.utils import check_api_version, symlink_user_forward_files
 from dpgen.generator.lib.vasp import incar_upper
 from dpgen.remote.decide_machine import convert_mdata
 from dpgen.util import load_file
@@ -1158,27 +1157,23 @@
     #        relax_run_tasks.append(ii)
     run_tasks = [os.path.basename(ii) for ii in relax_run_tasks]
 
-    api_version = mdata.get("api_version", "1.0")
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
-
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["fp_machine"],
-            mdata["fp_resources"],
-            commands=[fp_command],
-            work_path=work_dir,
-            run_tasks=run_tasks,
-            group_size=fp_group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="fp.log",
-            errlog="fp.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["fp_machine"],
+        mdata["fp_resources"],
+        commands=[fp_command],
+        work_path=work_dir,
+        run_tasks=run_tasks,
+        group_size=fp_group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="fp.log",
+        errlog="fp.log",
+    )
+    submission.run_submission()
 
 
 def coll_abacus_md(jdata):
@@ -1298,27 +1293,23 @@
     #        relax_run_tasks.append(ii)
     run_tasks = [os.path.basename(ii) for ii in relax_run_tasks]
 
-    api_version = mdata.get("api_version", "1.0")
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
-
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["fp_machine"],
-            mdata["fp_resources"],
-            commands=[fp_command],
-            work_path=work_dir,
-            run_tasks=run_tasks,
-            group_size=fp_group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="fp.log",
-            errlog="fp.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["fp_machine"],
+        mdata["fp_resources"],
+        commands=[fp_command],
+        work_path=work_dir,
+        run_tasks=run_tasks,
+        group_size=fp_group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="fp.log",
+        errlog="fp.log",
+    )
+    submission.run_submission()
 
 
 def run_vasp_md(jdata, mdata):
@@ -1359,27 +1350,24 @@
     run_tasks = [ii.replace(work_dir + "/", "") for ii in md_run_tasks]
     # dlog.info("md_work_dir", work_dir)
     # dlog.info("run_tasks",run_tasks)
-    api_version = mdata.get("api_version", "1.0")
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
 
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["fp_machine"],
-            mdata["fp_resources"],
-            commands=[fp_command],
-            work_path=work_dir,
-            run_tasks=run_tasks,
-            group_size=fp_group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="fp.log",
-            errlog="fp.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["fp_machine"],
+        mdata["fp_resources"],
+        commands=[fp_command],
+        work_path=work_dir,
+        run_tasks=run_tasks,
+        group_size=fp_group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="fp.log",
+        errlog="fp.log",
+    )
+    submission.run_submission()
 
 
 def run_abacus_md(jdata, mdata):
@@ -1435,27 +1423,24 @@
     run_tasks = [ii.replace(work_dir + "/", "") for ii in md_run_tasks]
     # dlog.info("md_work_dir", work_dir)
     # dlog.info("run_tasks",run_tasks)
-    api_version = mdata.get("api_version", "1.0")
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
 
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["fp_machine"],
-            mdata["fp_resources"],
-            commands=[fp_command],
-            work_path=work_dir,
-            run_tasks=run_tasks,
-            group_size=fp_group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="fp.log",
-            errlog="fp.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["fp_machine"],
+        mdata["fp_resources"],
+        commands=[fp_command],
+        work_path=work_dir,
+        run_tasks=run_tasks,
+        group_size=fp_group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="fp.log",
+        errlog="fp.log",
+    )
+    submission.run_submission()
 
 
 def gen_init_bulk(args):

dpgen/generator/lib/utils.py

@@ -6,6 +6,8 @@
 import re
 import shutil
 
+from packaging.version import Version
+
 iter_format = "%06d"
 task_format = "%02d"
 log_iter_head = "iter " + iter_format + " task " + task_format + ": "
@@ -110,3 +112,11 @@
             abs_file = os.path.abspath(file)
             os.symlink(abs_file, os.path.join(task, os.path.basename(file)))
     return
+
+
+def check_api_version(mdata):
+    if Version(mdata.get("api_version", "1.0")) < Version("1.0"):
+        raise RuntimeError(
+            "API version below 1.0 is no longer supported. Please upgrade to version 1.0 or newer."
+        )
+    return

dpgen/generator/run.py

@@ -78,6 +78,7 @@
 )
 from dpgen.generator.lib.siesta import make_siesta_input
 from dpgen.generator.lib.utils import (
+    check_api_version,
     create_path,
     log_iter,
     log_task,
@@ -874,31 +875,27 @@
     except Exception:
         train_group_size = 1
 
-    api_version = mdata.get("api_version", "1.0")
-
     user_forward_files = mdata.get("train" + "_user_forward_files", [])
     forward_files += [os.path.basename(file) for file in user_forward_files]
     backward_files += mdata.get("train" + "_user_backward_files", [])
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
 
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["train_machine"],
-            mdata["train_resources"],
-            commands=commands,
-            work_path=work_path,
-            run_tasks=run_tasks,
-            group_size=train_group_size,
-            forward_common_files=trans_comm_data,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="train.log",
-            errlog="train.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["train_machine"],
+        mdata["train_resources"],
+        commands=commands,
+        work_path=work_path,
+        run_tasks=run_tasks,
+        group_size=train_group_size,
+        forward_common_files=trans_comm_data,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="train.log",
+        errlog="train.log",
+    )
+    submission.run_submission()
 
 
 def post_train(iter_index, jdata, mdata):
@@ -2090,31 +2087,28 @@
     user_forward_files = mdata.get("model_devi" + "_user_forward_files", [])
     forward_files += [os.path.basename(file) for file in user_forward_files]
     backward_files += mdata.get("model_devi" + "_user_backward_files", [])
-    api_version = mdata.get("api_version", "1.0")
     if len(run_tasks) == 0:
         raise RuntimeError(
             "run_tasks for model_devi should not be empty! Please check your files."
         )
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
 
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["model_devi_machine"],
-            mdata["model_devi_resources"],
-            commands=commands,
-            work_path=work_path,
-            run_tasks=run_tasks,
-            group_size=model_devi_group_size,
-            forward_common_files=model_names,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog="model_devi.log",
-            errlog="model_devi.log",
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["model_devi_machine"],
+        mdata["model_devi_resources"],
+        commands=commands,
+        work_path=work_path,
+        run_tasks=run_tasks,
+        group_size=model_devi_group_size,
+        forward_common_files=model_names,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog="model_devi.log",
+        errlog="model_devi.log",
+    )
+    submission.run_submission()
 
 
 def run_model_devi(iter_index, jdata, mdata):
@@ -3964,27 +3958,23 @@
     forward_files += [os.path.basename(file) for file in user_forward_files]
     backward_files += mdata.get("fp" + "_user_backward_files", [])
 
-    api_version = mdata.get("api_version", "1.0")
-    if Version(api_version) < Version("1.0"):
-        raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
-        )
-
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            mdata["fp_machine"],
-            mdata["fp_resources"],
-            commands=[fp_command],
-            work_path=work_path,
-            run_tasks=run_tasks,
-            group_size=fp_group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog=log_file,
-            errlog=log_file,
-        )
-        submission.run_submission()
+    ### Submit the jobs
+    check_api_version(mdata)
+
+    submission = make_submission(
+        mdata["fp_machine"],
+        mdata["fp_resources"],
+        commands=[fp_command],
+        work_path=work_path,
+        run_tasks=run_tasks,
+        group_size=fp_group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog=log_file,
+        errlog=log_file,
+    )
+    submission.run_submission()
 
 
 def run_fp(iter_index, jdata, mdata):