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 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 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 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 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/dispatcher/Dispatcher.py

@@ -138,21 +138,20 @@
     """
     if Version(api_version) < Version("1.0"):
         raise RuntimeError(
-            f"API version {api_version} has been removed. Please upgrade to 1.0."
+            "API version below 1.0 is no longer supported. Please upgrade to version 1.0 or newer."
         )
 
-    elif Version(api_version) >= Version("1.0"):
-        submission = make_submission(
-            machine,
-            resources,
-            commands=commands,
-            work_path=work_path,
-            run_tasks=run_tasks,
-            group_size=group_size,
-            forward_common_files=forward_common_files,
-            forward_files=forward_files,
-            backward_files=backward_files,
-            outlog=outlog,
-            errlog=errlog,
-        )
-        submission.run_submission()
+    submission = make_submission(
+        machine,
+        resources,
+        commands=commands,
+        work_path=work_path,
+        run_tasks=run_tasks,
+        group_size=group_size,
+        forward_common_files=forward_common_files,
+        forward_files=forward_files,
+        backward_files=backward_files,
+        outlog=outlog,
+        errlog=errlog,
+    )
+    submission.run_submission()

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,12 @@
             abs_file = os.path.abspath(file)
             os.symlink(abs_file, os.path.join(task, os.path.basename(file)))
     return
+
+
+def check_api_version(mdata):
+    """Check if the API version in mdata is at least 1.0."""
+    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