feat: accept YAML workflow configurations

docs/quickcli.md

@@ -36,6 +36,8 @@ INFO:root:steps iter-000001--prep-run-train----------------------- finished
 INFO:root:steps iter-000001--prep-run-explore--------------------- finished
 ...
 ```
+Workflow configuration files may use JSON (`.json`) or YAML (`.yaml`/`.yml`). YAML is useful when comments and unquoted keys make a long configuration easier to maintain.
+
 The artifacts can be downloaded on-the-fly with `-d` flag. Note that the existing files are automatically skipped if one sets `dflow_config["archive_mode"] = None`.
 
 

dpgen2/entrypoint/common.py

@@ -1,3 +1,4 @@
+import json
 import os
 from pathlib import (
     Path,
@@ -10,6 +11,7 @@
 )
 
 import dflow
+import yaml
 
 from dpgen2.utils import (
     bohrium_config_from_dict,
@@ -23,6 +25,39 @@
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
 
+def load_config(path: Union[str, Path]) -> Dict:
+    r"""Load a DPGEN2 workflow configuration from JSON or YAML.
+
+    YAML is selected for ``.yaml`` and ``.yml`` files. Other suffixes retain
+    the existing strict JSON parser so malformed JSON continues to fail with
+    its familiar diagnostics.
+
+    Parameters
+    ----------
+    path : str or Path
+        Workflow configuration path.
+
+    Returns
+    -------
+    dict
+        Parsed workflow configuration.
+
+    Raises
+    ------
+    ValueError
+        If the document root is not a mapping.
+    """
+    config_path = Path(path)
+    content = config_path.read_text()
+    if config_path.suffix.lower() in {".yaml", ".yml"}:
+        config = yaml.safe_load(content)
+    else:
+        config = json.loads(content)
+    if not isinstance(config, dict):
+        raise ValueError(f"DPGEN2 configuration root must be a mapping: {config_path}")
+    return config
+
+
 def global_config_workflow(
     wf_config,
 ):

dpgen2/entrypoint/main.py

@@ -1,5 +1,4 @@
 import argparse
-import json
 import logging
 import textwrap
 from typing import (
@@ -16,6 +15,7 @@
 
 from .common import (
     expand_idx,
+    load_config,
 )
 from .download import (
     download,
@@ -75,7 +75,7 @@ def main_parser() -> argparse.ArgumentParser:
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     parser_run.add_argument(
-        "CONFIG", help="the config file in json format defining the workflow."
+        "CONFIG", help="the workflow config file in JSON or YAML format."
     )
 
     ##########################################
@@ -86,7 +86,7 @@ def main_parser() -> argparse.ArgumentParser:
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     parser_resubmit.add_argument(
-        "CONFIG", help="the config file in json format defining the workflow."
+        "CONFIG", help="the workflow config file in JSON or YAML format."
     )
     parser_resubmit.add_argument("ID", help="the ID of the existing workflow.")
     parser_resubmit.add_argument(
@@ -123,7 +123,9 @@ def main_parser() -> argparse.ArgumentParser:
         help="Print the keys of the successful DPGEN2 steps",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser_showkey.add_argument("CONFIG", help="the config file in json format.")
+    parser_showkey.add_argument(
+        "CONFIG", help="the config file in JSON or YAML format."
+    )
     parser_showkey.add_argument("ID", help="the ID of the existing workflow.")
 
     ##########################################
@@ -133,7 +135,7 @@ def main_parser() -> argparse.ArgumentParser:
         help="Print the status (stage, iteration, convergence) of the  DPGEN2 workflow",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser_status.add_argument("CONFIG", help="the config file in json format.")
+    parser_status.add_argument("CONFIG", help="the config file in JSON or YAML format.")
     parser_status.add_argument("ID", help="the ID of the existing workflow.")
 
     ##########################################
@@ -162,7 +164,9 @@ def main_parser() -> argparse.ArgumentParser:
         ),
         formatter_class=argparse.RawTextHelpFormatter,
     )
-    parser_download.add_argument("CONFIG", help="the config file in json format.")
+    parser_download.add_argument(
+        "CONFIG", help="the config file in JSON or YAML format."
+    )
     parser_download.add_argument("ID", help="the ID of the existing workflow.")
     parser_download.add_argument(
         "-l",
@@ -211,7 +215,7 @@ def main_parser() -> argparse.ArgumentParser:
         help="Watch a DPGEN2 workflow",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser_watch.add_argument("CONFIG", help="the config file in json format.")
+    parser_watch.add_argument("CONFIG", help="the config file in JSON or YAML format.")
     parser_watch.add_argument("ID", help="the ID of the existing workflow.")
     parser_watch.add_argument(
         "-k",
@@ -279,7 +283,9 @@ def main_parser() -> argparse.ArgumentParser:
         help="restart a DPGEN2 workflow (for debug mode only).",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser_restart.add_argument("CONFIG", help="the config file in json format.")
+    parser_restart.add_argument(
+        "CONFIG", help="the config file in JSON or YAML format."
+    )
     parser_restart.add_argument("ID", help="the ID of the workflow.")
 
     # --version
@@ -320,14 +326,12 @@ def main():
     dict_args = vars(args)
 
     if args.command == "submit":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         submit_concurrent_learning(
             config,
         )
     elif args.command == "resubmit":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         resubmit_concurrent_learning(
             config,
@@ -338,24 +342,21 @@ def main():
             fold=args.fold,
         )
     elif args.command == "status":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         status(
             wfid,
             config,
         )
     elif args.command == "showkey":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         showkey(
             wfid,
             config,
         )
     elif args.command == "download":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         if args.list_supported is not None and args.list_supported:
             print(print_op_download_setting())
@@ -379,8 +380,7 @@ def main():
                 chk_pnt=args.no_check_point,
             )
     elif args.command == "watch":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         watch(
             wfid,
@@ -397,17 +397,15 @@ def main():
             bind_all=args.bind_all,
         )
     elif args.command == "restart":
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wf = submit_concurrent_learning(
             config,
             no_submission=True,
         )
         wf.id = args.ID
         wf.submit()
     elif args.command in workflow_subcommands:
-        with open(args.CONFIG) as fp:
-            config = json.load(fp)
+        config = load_config(args.CONFIG)
         wfid = args.ID
         execute_workflow_subcommand(args.command, wfid, config)
     elif args.command is None:

dpgen2/entrypoint/workflow.py

@@ -28,7 +28,7 @@ def add_subparser_workflow_subcommand(subparsers, command: str):
         help=f"{command.capitalize()} a DPGEN2 workflow.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser_cmd.add_argument("CONFIG", help="the config file in json format.")
+    parser_cmd.add_argument("CONFIG", help="the config file in JSON or YAML format.")
     parser_cmd.add_argument("ID", help="the ID of the workflow.")
 
 

pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
 	     'scipy',
 	     'lbg',
 	     'packaging',
+	     'pyyaml',
 	     'fpop',
 	     'dpgui',
              'cp2kdata',

tests/entrypoint/test_common.py

@@ -0,0 +1,54 @@
+import json
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+# isort: off
+from .context import (
+    dpgen2,
+)
+from dpgen2.entrypoint.common import (
+    load_config,
+)
+
+# isort: on
+
+
+class TestLoadConfig(unittest.TestCase):
+    def test_json(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "input.json"
+            path.write_text(json.dumps({"type_map": ["H", "O"]}))
+
+            self.assertEqual(load_config(path), {"type_map": ["H", "O"]})
+
+    def test_yaml(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "input.yaml"
+            path.write_text(
+                """# YAML permits comments and unquoted mapping keys.
+type_map:
+  - H
+  - O
+explore:
+  fatal_at_max: false
+"""
+            )
+
+            self.assertEqual(
+                load_config(path),
+                {
+                    "type_map": ["H", "O"],
+                    "explore": {"fatal_at_max": False},
+                },
+            )
+
+    def test_document_root_must_be_mapping(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "input.yml"
+            path.write_text("- not\n- a\n- mapping\n")
+
+            with self.assertRaisesRegex(ValueError, "root must be a mapping"):
+                load_config(path)