support extra FP outputs; add post_command for Gaussian

dpgen2/entrypoint/args.py

@@ -459,6 +459,7 @@ def fp_args(inputs, run):
     doc_inputs_config = "Configuration for preparing vasp inputs"
     doc_run_config = "Configuration for running vasp tasks"
     doc_task_max = "Maximum number of vasp tasks for each iteration"
+    doc_extra_output_files = "Extra output file names, support wildcards"
 
     return [
         Argument(
@@ -476,6 +477,7 @@ def fp_args(inputs, run):
             doc=doc_run_config,
         ),
         Argument("task_max", int, optional=True, default=10, doc=doc_task_max),
+        Argument("extra_output_files", list, optional=True, default=[], doc=doc_extra_output_files),
     ]
 
 

dpgen2/entrypoint/submit.py

@@ -564,6 +564,7 @@ def workflow_concurrent_learning(
 
     fp_config["inputs"] = fp_inputs
     fp_config["run"] = config["fp"]["run_config"]
+    fp_config["extra_output_files"] = config["fp"]["extra_output_files"]
     if fp_style == "deepmd":
         assert (
             "teacher_model_path" in fp_config["run"]

dpgen2/fp/abacus.py

@@ -166,6 +166,7 @@ def get_output_sign(cls):
             {
                 "log": Artifact(Path),
                 "labeled_data": Artifact(Path),
+                "extra_outputs": Artifact(List[Path]),
             }
         )
 
@@ -202,10 +203,15 @@ def execute(
         out_name = fp_default_out_data_name
         sys.to("deepmd/npy", workdir / out_name)
 
+        extra_outputs = []
+        for fname in ip["config"]["extra_output_files"]:
+            extra_outputs += list(workdir.glob(fname))
+
         return OPIO(
             {
                 "log": workdir / "log",
                 "labeled_data": workdir / out_name,
+                "extra_outputs": extra_outputs,
             }
         )
 

dpgen2/fp/cp2k.py

@@ -126,6 +126,7 @@ def get_output_sign(cls):
             {
                 "log": Artifact(Path),
                 "labeled_data": Artifact(Path),
+                "extra_outputs": Artifact(List[Path]),
             }
         )
 
@@ -170,10 +171,15 @@ def execute(
         out_name = fp_default_out_data_name
         sys.to("deepmd/npy", workdir / out_name)
 
+        extra_outputs = []
+        for fname in ip["config"]["extra_output_files"]:
+            extra_outputs += list(workdir.glob(fname))
+
         return OPIO(
             {
                 "log": workdir / "output.log",
                 "labeled_data": workdir / out_name,
+                "extra_outputs": extra_outputs,
             }
         )
 

dpgen2/fp/gaussian.py

@@ -134,6 +134,7 @@ def run_task(
         self,
         command: str,
         out: str,
+        post_command: str,
     ) -> Tuple[str, str]:
         r"""Defines how one FP task runs
 
@@ -170,6 +171,23 @@ def run_task(
                 )
             )
             raise TransientError("gaussian failed")
+        if post_command is not None:
+            ret, out, err = run_command(post_command, shell=True)
+            if ret != 0:
+                logging.error(
+                    "".join(
+                        (
+                            "gaussian postprocessing failed\n",
+                            "out msg: ",
+                            out,
+                            "\n",
+                            "err msg: ",
+                            err,
+                            "\n",
+                        )
+                    )
+                )
+                raise TransientError("gaussian postprocessing failed")
         # convert the output to deepmd/npy format
         sys = dpdata.LabeledSystem(gaussian_output_name, fmt="gaussian/log")
         sys.to("deepmd/npy", out_name)
@@ -187,6 +205,7 @@ def args() -> List[dargs.Argument]:
 
         doc_gaussian_cmd = "The command of Gaussian"
         doc_gaussian_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`."
+        doc_post_command = "The command after Gaussian"
         return [
             Argument(
                 "command", str, optional=True, default="g16", doc=doc_gaussian_cmd
@@ -198,4 +217,7 @@ def args() -> List[dargs.Argument]:
                 default=fp_default_out_data_name,
                 doc=doc_gaussian_out,
             ),
+            Argument(
+                "post_command", str, optional=True, default=None, doc=doc_post_command
+            )
         ]

dpgen2/fp/run_fp.py

@@ -58,6 +58,7 @@ def get_output_sign(cls):
             {
                 "log": Artifact(Path),
                 "labeled_data": Artifact(Path),
+                "extra_outputs": Artifact(List[Path]),
             }
         )
 
@@ -196,9 +197,14 @@ def execute(
                     Path(iname).symlink_to(ii)
             out_name, log_name = self.run_task(**config)
 
+        extra_outputs = []
+        for fname in ip["config"]["extra_output_files"]:
+            extra_outputs += list(work_dir.glob(fname))
+
         return OPIO(
             {
                 "log": work_dir / log_name,
                 "labeled_data": work_dir / out_name,
+                "extra_outputs": extra_outputs,
             }
         )

dpgen2/superop/prep_run_fp.py

@@ -170,7 +170,7 @@ def _prep_run_fp(
                 "int('{{item}}')",
                 input_parameter=["task_name"],
                 input_artifact=["task_path"],
-                output_artifact=["log", "labeled_data"],
+                output_artifact=["log", "labeled_data", "extra_outputs"],
                 **template_slice_config,
             ),
             python_packages=upload_python_packages,