diff --git a/docs/quickcli.md b/docs/quickcli.md index e9003715..05eaffb4 100644 --- a/docs/quickcli.md +++ b/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`. diff --git a/dpgen2/entrypoint/common.py b/dpgen2/entrypoint/common.py index 0d0af9e8..56493cd3 100644 --- a/dpgen2/entrypoint/common.py +++ b/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, ): diff --git a/dpgen2/entrypoint/main.py b/dpgen2/entrypoint/main.py index 00ee55e3..af1485ca 100644 --- a/dpgen2/entrypoint/main.py +++ b/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,8 +397,7 @@ 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, @@ -406,8 +405,7 @@ def main(): 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: diff --git a/dpgen2/entrypoint/workflow.py b/dpgen2/entrypoint/workflow.py index 6e7e363a..c61e530f 100644 --- a/dpgen2/entrypoint/workflow.py +++ b/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.") diff --git a/pyproject.toml b/pyproject.toml index ccf004f9..abf9b439 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ 'scipy', 'lbg', 'packaging', + 'pyyaml', 'fpop', 'dpgui', 'cp2kdata', diff --git a/tests/entrypoint/test_common.py b/tests/entrypoint/test_common.py new file mode 100644 index 00000000..889f5d4b --- /dev/null +++ b/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)