feat: support opt-in $ref loading in workflows

README.md

@@ -49,6 +49,14 @@ DP-GEN contains the following workflows:
 
 For detailed usage and parameters, read [DP-GEN documentation](https://docs.deepmodeling.com/projects/dpgen/).
 
+### `$ref` support (secure by default)
+
+DP-GEN now supports loading external JSON/YAML snippets through `"$ref"` in relevant workflows (via `dargs>=0.5.0`).
+For security reasons, it is **disabled by default** and must be enabled explicitly.
+
+- CLI: add `--allow-ref` to `dpgen run`, `dpgen simplify`, and `dpgen init_reaction`
+- Python API path (`dpgen.util.normalize`): pass `allow_ref=True`
+
 ## Tutorials and examples
 
 * [Tutorials](https://tutorials.deepmodeling.com/en/latest/Tutorials/DP-GEN/): basic tutorials for DP-GEN.

doc/init/init-reaction.md

@@ -18,4 +18,10 @@ An example of `reaction.json` is given below:
 
 For detailed parameters, see [parametes](init-reaction-jdata.rst) and [machine parameters](init-reaction-mdata.rst).
 
+To enable loading external JSON/YAML snippets via `$ref` in the reaction parameter file, add `--allow-ref` explicitly (disabled by default for security):
+
+```bash
+dpgen init_reaction reaction.json machine.json --allow-ref
+```
+
 The genereated data can be used to continue DP-GEN concurrent learning workflow. Read [Energy & Fuels, 2021, 35 (1), 762–769](https://10.1021/acs.energyfuels.0c03211) for details.

doc/run/overview-of-the-run-process.md

@@ -16,6 +16,12 @@ Here, we give a general description of the run process. We can execute the run p
 dpgen run param.json machine.json
 ```
 
+To enable loading external JSON/YAML snippets via `$ref` in `param.json`, add `--allow-ref` explicitly (disabled by default for security):
+
+```sh
+dpgen run param.json machine.json --allow-ref
+```
+
 The following files or folders will be created and upgraded by codes：
 
 - iter.00000x contains the main results that DP-GEN generates in the first iteration.

doc/simplify/simplify.md

@@ -6,6 +6,12 @@ Use the following script to start the workflow:
 dpgen simplify param.json machine.json
 ```
 
+To enable loading external JSON/YAML snippets via `$ref`, add `--allow-ref` explicitly (disabled by default for security):
+
+```bash
+dpgen simplify param.json machine.json --allow-ref
+```
+
 Here is an example of `param.json` for QM7 dataset:
 ```json
 {

dpgen/data/reaction.py

@@ -212,7 +212,7 @@ def gen_init_reaction(args):
     mdata = load_file(args.MACHINE)
 
     jdata_arginfo = init_reaction_jdata_arginfo()
-    jdata = normalize(jdata_arginfo, jdata)
+    jdata = normalize(jdata_arginfo, jdata, allow_ref=args.allow_ref)
 
     mdata = convert_mdata(mdata, ["reaxff", "build", "fp"])
     record = "record.reaction"

dpgen/generator/run.py

@@ -4972,12 +4972,12 @@ def set_version(mdata):
     return mdata
 
 
-def run_iter(param_file, machine_file):
+def run_iter(param_file, machine_file, allow_ref: bool = False):
     jdata = load_file(param_file)
     mdata = load_file(machine_file)
 
     jdata_arginfo = run_jdata_arginfo()
-    jdata = normalize(jdata_arginfo, jdata, strict_check=False)
+    jdata = normalize(jdata_arginfo, jdata, strict_check=False, allow_ref=allow_ref)
 
     update_mass_map(jdata)
 
@@ -5718,7 +5718,7 @@ def gen_run(args):
         if args.debug:
             dlog.setLevel(logging.DEBUG)
         dlog.info("start running")
-        run_iter(args.PARAM, args.MACHINE)
+        run_iter(args.PARAM, args.MACHINE, allow_ref=args.allow_ref)
         dlog.info("finished")
 
 

dpgen/main.py

@@ -101,6 +101,11 @@ def main_parser() -> argparse.ArgumentParser:
         nargs="?",
         help="machine file, json/yaml format",
     )
+    parser_init_reaction.add_argument(
+        "--allow-ref",
+        action="store_true",
+        help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.",
+    )
     parser_init_reaction.set_defaults(func=gen_init_reaction)
 
     # run
@@ -110,6 +115,11 @@ def main_parser() -> argparse.ArgumentParser:
     parser_run.add_argument("PARAM", type=str, help="parameter file, json/yaml format")
     parser_run.add_argument("MACHINE", type=str, help="machine file, json/yaml format")
     parser_run.add_argument("-d", "--debug", action="store_true", help="log debug info")
+    parser_run.add_argument(
+        "--allow-ref",
+        action="store_true",
+        help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.",
+    )
     parser_run.set_defaults(func=gen_run)
 
     # run/report
@@ -181,6 +191,11 @@ def main_parser() -> argparse.ArgumentParser:
     parser_run.add_argument("PARAM", type=str, help="parameter file, json/yaml format")
     parser_run.add_argument("MACHINE", type=str, help="machine file, json/yaml format")
     parser_run.add_argument("-d", "--debug", action="store_true", help="log debug info")
+    parser_run.add_argument(
+        "--allow-ref",
+        action="store_true",
+        help="Allow loading external JSON/YAML snippets through `$ref`. Disabled by default for security.",
+    )
     parser_run.set_defaults(func=gen_simplify)
 
     # test

dpgen/simplify/simplify.py

@@ -508,7 +508,7 @@ def make_fp(iter_index, jdata, mdata):
         make_fp_calculation(iter_index, jdata, mdata)
 
 
-def run_iter(param_file, machine_file):
+def run_iter(param_file, machine_file, allow_ref: bool = False):
     """Init (iter 0): init_pick.
 
     tasks (iter > 0):
@@ -526,7 +526,7 @@ def run_iter(param_file, machine_file):
     mdata = load_file(machine_file)
 
     jdata_arginfo = simplify_jdata_arginfo()
-    jdata = normalize(jdata_arginfo, jdata)
+    jdata = normalize(jdata_arginfo, jdata, allow_ref=allow_ref)
 
     # set up electron temperature
     use_ele_temp = jdata.get("use_ele_temp", 0)
@@ -626,5 +626,5 @@ def gen_simplify(args):
         if args.debug:
             dlog.setLevel(logging.DEBUG)
         dlog.info("start simplifying")
-        run_iter(args.PARAM, args.MACHINE)
+        run_iter(args.PARAM, args.MACHINE, allow_ref=args.allow_ref)
         dlog.info("finished")

dpgen/util.py

@@ -78,7 +78,12 @@ def expand_sys_str(root_dir: Union[str, Path]) -> list[str]:
     return matches
 
 
-def normalize(arginfo: Argument, data: dict, strict_check: bool = True) -> dict:
+def normalize(
+    arginfo: Argument,
+    data: dict,
+    strict_check: bool = True,
+    allow_ref: bool = False,
+) -> dict:
     """Normalize and check input data.
 
     Parameters
@@ -89,14 +94,17 @@ def normalize(arginfo: Argument, data: dict, strict_check: bool = True) -> dict:
         input data
     strict_check : bool, default=True
         strict check data or not
+    allow_ref : bool, default=False
+        Whether to allow loading external JSON/YAML snippets via ``$ref``.
+        Disabled by default for security.
 
     Returns
     -------
     dict
         normalized data
     """
-    data = arginfo.normalize_value(data, trim_pattern="_*")
-    arginfo.check_value(data, strict=strict_check)
+    data = arginfo.normalize_value(data, trim_pattern="_*", allow_ref=allow_ref)
+    arginfo.check_value(data, strict=strict_check, allow_ref=allow_ref)
     return data
 
 

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
     'GromacsWrapper>=0.9.0; python_version >= "3.12"',
     'dpdispatcher>=0.3.11',
     'netCDF4',
-    'dargs>=0.4.0',
+    'dargs>=0.5.0',
     'h5py',
     'pymatgen-analysis-defects',
     'openbabel-wheel',