Repository: ruoyuwang1995nya/pfd-kit

Issue:
finetune流程中conf_generation无法读取指定prefix里的文件
json内容如下：
```
"conf_generation": {
        "init_configurations":
            {
            "type": "file",
            "prefix": "./poscar/",
            "fmt": "vasp/poscar",
            "files": [
                      "B1C20H42S01POSCAR.vasp",
            .....
           ]},
```
报错结果如下：
```
Traceback (most recent call last):
  File "/opt/mamba/bin/dp-dist", line 8, in <module>
    sys.exit(main())
  File "/opt/mamba/lib/python3.10/site-packages/distill/entrypoint/main.py", line 210, in main
    submit_ft(
  File "/opt/mamba/lib/python3.10/site-packages/distill/entrypoint/submit.py", line 563, in submit_ft
    ft_step=workflow_finetune(wf_config)
  File "/opt/mamba/lib/python3.10/site-packages/distill/entrypoint/submit.py", line 444, in workflow_finetune
    init_confs=upload_artifact(config["conf_generation"]["init_configurations"]["files"])
  File "/opt/mamba/lib/python3.10/site-packages/dflow/utils.py", line 228, in upload_artifact
    raise RuntimeError("File or directory %s not found" % p)
RuntimeError: File or directory B1C20H42S01POSCAR.vasp not found
```
该问题可以通过把prefix改成./，并把结构文件放在json当前目录解决，不过这样对于多个结构文件来说不太方便。

Hints from issue comments:
@ruoyuwang1995nya: issue resolved

Solution patch:
diff --git a/examples/distillation/example.json b/examples/distillation/example.json
index 6aa2724..7d6dbbf 100644
--- a/examples/distillation/example.json
+++ b/examples/distillation/example.json
@@ -82,6 +82,9 @@
         "/opt/mamba/envs/openlam/lib/python3.12/site-packages/fpop"
 
     ],
+    "task":{
+        "type":"dist"
+    },
     "inputs": {
         "type_map": [
             "H",
@@ -203,7 +206,6 @@
             "Ts",
             "Og"
         ],
-        "teacher_model":"./lgps_ft.ckpt.pt",
         "_comment": "all"
     },
     "train": {
diff --git a/examples/finetune/example.json b/examples/finetune/example.json
index c423467..55f2537 100644
--- a/examples/finetune/example.json
+++ b/examples/finetune/example.json
@@ -138,6 +138,10 @@
     "upload_python_packages": [
         "/opt/mamba/envs/openlam/lib/python3.12/site-packages/ase"
     ],
+    "task":{
+        "type":"finetune"
+    },
+    "init_training": false,
     "inputs": {
         "type_map": [
             "H",
@@ -280,35 +284,32 @@
             }
             ]},
     "exploration": {
-        "init_training":true,
-        "skip_aimd":false,
         "max_iter":2,
-        "converge_config":{
+        "convergence":{
             "type":"energy_rmse",
             "RMSE":0.01},
         "filter":[{
             "type":"distance"
         }],
-        "md":{
-            "type": "lmp",
-            "config": {
-                "command": "lmp -var restart 0",
-                "shuffle_models": false, 
-                "head": null},
-            "stages":[[
-                {   "_comment": "group 1 stage 1 of finetune-exploration",
-                    "conf_idx": [0],
-                    "n_sample":3,
-                    "exploration":{
-                        "type": "lmp-md",
-                        "ensemble": "npt",
-                        "dt":0.005,
-                        "nsteps": 2000,
-                        "temps": [500],
-                        "press":[1],
-                        "trj_freq": 200},
-                    "max_sample": 10000
-                    }]]}},
+        "type": "lmp",
+        "config": {
+            "command": "lmp -var restart 0",
+            "shuffle_models": false, 
+            "head": null},
+        "stages":[[
+            {  "_comment": "group 1 stage 1 of finetune-exploration",
+                "conf_idx": [0],
+                "n_sample":3,
+                "exploration":{
+                "type": "lmp-md",
+                "ensemble": "npt",
+                "dt":0.005,
+                "nsteps": 2000,
+                "temps": [500],
+                "press":[1],
+                "trj_freq": 200},
+                "max_sample": 10000
+                    }]]},
     "aimd":{
             "type": "fpop_abacus",
             "task_max": 5,
diff --git a/pfd/entrypoint/args.py b/pfd/entrypoint/args.py
index eff5bd4..43a4305 100644
--- a/pfd/entrypoint/args.py
+++ b/pfd/entrypoint/args.py
@@ -9,22 +9,13 @@
     Variant,
 )
 
-import dpgen2
-from dpgen2.conf import (
-    conf_styles,
-)
-from dpgen2.constants import (
-    default_image,
-)
-from dpgen2.exploration.report import (
-    conv_styles,
-)
+from pfd.exploration.converge import ConvTypes as conv_styles
+
 from dpgen2.fp import (
     fp_styles,
 )
-from dpgen2.op.run_dp_train import (
-    RunDPTrain,
-)
+from pfd.train import train_styles
+
 from dpgen2.op.run_lmp import (
     RunLmp,
 )
@@ -41,44 +32,68 @@ def make_link(content, ref_key):
     )
 
 
-def dp_dist_train_args():
-    doc_config = "Configuration of training"
-    doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
-    dock_student_model_path = "The path of student model"
-    doc_student_model_uri = "The URI of student model"
-    doc_optional_files = "Optional files for training"
+def task_dist():
+    return []
 
+
+def task_finetune():
+    doc_init_train = "Training before exploration"
+    doc_skip_aimd = "Skip aimd exploration"
     return [
         Argument(
-            "config",
-            dict,
-            RunDPTrain.training_args(),
-            optional=True,
-            default=RunDPTrain.normalize_config({}),
-            doc=doc_config,
-        ),
-        Argument(
-            "template_script", [List[str], str], optional=False, doc=doc_template_script
+            "init_training", bool, optional=True, default=False, doc=doc_init_train
         ),
-        Argument("student_model_path", str, optional=True, doc=dock_student_model_path),
+        Argument("skip_aimd", bool, optional=True, default=True, doc=doc_skip_aimd),
+    ]
+
+
+def variant_task():
+    doc = "Select task type"
+    return Variant(
+        "type",
+        [Argument("finetune", dict, task_finetune()), Argument("dist", dict, [])],
+    )
+
+
+def conf_args():
+    doc_fmt = "Format of input structure files"
+    return [
+        Argument("type", str, optional=True, default="file"),
+        Argument("prefix", str, optional=True, default="./"),
+        Argument("fmt", str, optional=True, default="vasp/poscar", doc=doc_fmt),
+        Argument("files", [str, List[str]], optional=True),
+        Argument("confs_uri", [str, List[str]], optional=True, default=None),
+    ]
+
+
+def pert_gen():
+    doc_atom_pert_distance = "Perturb distance for atoms, in Angstrom"
+    doc_pert_num = "Number of perturbed structures"
+    return [
+        Argument("conf_idx", [str, List[int]], optional=True, default="default"),
         Argument(
-            "student_model_uri",
-            str,
+            "atom_pert_distance",
+            float,
             optional=True,
-            default=None,
-            doc=doc_student_model_uri,
+            default=0.1,
+            doc=doc_atom_pert_distance,
         ),
+        Argument("atom_pert_fraction", float, optional=True, default=0.03),
+        Argument("pert_num", int, optional=True, default=1, doc=doc_atom_pert_distance),
+    ]
+
+
+def conf_gen_args():
+    doc = "Generation of atomic configuration for pfd exploration"
+    return [
         Argument(
-            "optional_files",
-            list,
-            optional=True,
-            default=None,
-            doc=doc_optional_files,
+            "init_configurations", dict, conf_args(), alias=["confs", "init_confs"]
         ),
+        Argument("pert_generation", List[dict], pert_gen()),
     ]
 
 
-def dp_train_args():
+def train_args(run_train):
     doc_numb_models = "Number of models trained for evaluating the model deviation"
     doc_config = "Configuration of training"
     doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. "
@@ -90,10 +105,10 @@ def dp_train_args():
         Argument(
             "config",
             dict,
-            RunDPTrain.training_args(),
+            run_train.training_args(),
             optional=True,
-            default=RunDPTrain.normalize_config({}),
-            doc=doc_numb_models,
+            default=run_train.normalize_config({}),
+            doc=doc_config,
         ),
         Argument("numb_models", int, optional=True, default=1, doc=doc_numb_models),
         Argument(
@@ -125,13 +140,14 @@ def dp_train_args():
 
 
 def variant_train():
-    doc = "the type of the training"
+    doc = "the type of the training model"
+    train_list = []
+    for kk in train_styles.keys():
+        train_list.append(Argument(kk, dict, train_args(train_styles[kk]["run"])))
+
     return Variant(
         "type",
-        [
-            Argument("dp", dict, dp_train_args()),
-            Argument("dp-dist", dict, dp_dist_train_args()),
-        ],
+        train_list,
         doc=doc,
     )
 
@@ -155,25 +171,6 @@ def variant_conv():
     )
 
 
-def variant_conf():
-    doc = "the type of the initial configuration generator."
-    var_list = []
-    for kk in conf_styles.keys():
-        var_list.append(
-            Argument(
-                kk,
-                dict,
-                conf_styles[kk].args(),
-                doc=conf_styles[kk].doc(),
-            )
-        )
-    return Variant(
-        "type",
-        var_list,
-        doc=doc,
-    )
-
-
 def lmp_args():
     doc_config = "Configuration of lmp exploration"
     doc_max_numb_iter = "Maximum number of iterations per stage"
@@ -182,13 +179,13 @@ def lmp_args():
     )
     doc_output_nopbc = "Remove pbc of the output configurations"
     doc_convergence = "The method of convergence check."
-    doc_configuration = "A list of initial configurations."
     doc_stages = (
         "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. "
         "The outer list provides the enumeration of the exploration stages. "
         "Then each stage is defined by a list of exploration task groups. "
         "Each task group is described in :ref:`the task group definition<task_group_sec>` "
     )
+    doc_filter = "Filter configuration for DFT calculation"
 
     return [
         Argument(
@@ -200,7 +197,12 @@ def lmp_args():
             doc=doc_config,
         ),
         Argument(
-            "max_numb_iter", int, optional=True, default=10, doc=doc_max_numb_iter
+            "max_numb_iter",
+            int,
+            optional=True,
+            default=10,
+            doc=doc_max_numb_iter,
+            alias=["max_iter"],
         ),
         Argument(
             "fatal_at_max", bool, optional=True, default=True, doc=doc_fatal_at_max
@@ -215,16 +217,14 @@ def lmp_args():
             [variant_conv()],
             optional=False,
             doc=doc_convergence,
+            alias=["converge_config"],
         ),
         Argument(
-            "configurations",
-            list,
-            [],
-            [variant_conf()],
-            optional=False,
-            repeat=True,
-            doc=doc_configuration,
-            alias=["configuration"],
+            "filter",
+            List[dict],
+            optional=True,
+            default=[{"type": "distance"}],
+            doc=doc_filter,
         ),
         Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
     ]
@@ -299,16 +299,6 @@ def caly_args():
             optional=False,
             doc=doc_convergence,
         ),
-        Argument(
-            "configurations",
-            list,
-            [],
-            [variant_conf()],
-            optional=False,
-            repeat=True,
-            doc=doc_configuration,
-            alias=["configuration"],
-        ),
         Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
     ]
 
@@ -400,7 +390,7 @@ def input_args():
 
     return [
         Argument("type_map", List[str], optional=False, doc=doc_type_map),
-        Argument("mass_map", List[float], optional=False, doc=doc_mass_map),
+        Argument("mass_map", List[float], optional=True, doc=doc_mass_map),
         Argument(
             "init_data_prefix",
             str,
@@ -426,6 +416,20 @@ def input_args():
             default=None,
             doc=doc_init_data_uri,
         ),
+        Argument(
+            "teacher_models_paths",
+            [List[str], str],
+            optional=True,
+            default=None,
+            doc=doc_init_sys,
+        ),
+        Argument(
+            "teacher_models_uri",
+            str,
+            optional=True,
+            default=None,
+            doc=doc_init_data_uri,
+        ),
         Argument(
             "multitask",
             bool,
@@ -553,7 +557,7 @@ def default_step_config_args():
     ]
 
 
-def dpgen_step_config_args(default_config):
+def pfd_step_config_args(default_config):
     doc_prep_train_config = "Configuration for prepare train"
     doc_run_train_config = "Configuration for run train"
     doc_prep_explore_config = "Configuration for prepare exploration"
@@ -562,7 +566,6 @@ def dpgen_step_config_args(default_config):
     doc_run_fp_config = "Configuration for run fp"
     doc_select_confs_config = "Configuration for the select confs"
     doc_collect_data_config = "Configuration for the collect data"
-    doc_cl_step_config = "Configuration for the concurrent learning step"
 
     return [
         Argument(
@@ -629,14 +632,6 @@ def dpgen_step_config_args(default_config):
             default=default_config,
             doc=doc_collect_data_config,
         ),
-        Argument(
-            "cl_step_config",
-            dict,
-            step_conf_args(),
-            optional=True,
-            default=default_config,
-            doc=doc_cl_step_config,
-        ),
     ]
 
 
@@ -644,11 +639,13 @@ def submit_args(default_step_config=normalize_step_dict({})):
     doc_bohrium_config = "Configurations for the Bohrium platform."
     doc_step_configs = "Configurations for executing dflow steps"
     doc_upload_python_packages = "Upload python package, for debug purpose"
-    doc_inputs = "The input parameter and artifacts for dpgen2"
+    doc_task = "Task type, `finetune` or `dist`"
+    doc_conf_gen = "The inputparameter and artifacts for confs generation"
+    doc_inputs = "The input parameter and artifacts for pfd"
     doc_train = "The configuration for training"
     doc_explore = "The configuration for exploration"
     doc_fp = "The configuration for FP"
-    doc_name = "The workflow name, 'dpgen' for default"
+    doc_name = "The workflow name, 'pfd' for default"
     doc_parallelism = "The parallelism for the workflow. Accept an int that stands for the maximum number of running pods for the workflow. None for default"
 
     return (
@@ -666,7 +663,7 @@ def submit_args(default_step_config=normalize_step_dict({})):
             Argument(
                 "step_configs",
                 dict,
-                dpgen_step_config_args(default_step_config),
+                pfd_step_config_args(default_step_config),
                 optional=True,
                 default={},
                 doc=doc_step_configs,
@@ -679,20 +676,31 @@ def submit_args(default_step_config=normalize_step_dict({})):
                 doc=doc_upload_python_packages,
                 alias=["upload_python_package"],
             ),
+            Argument(
+                "conf_generation",
+                dict,
+                conf_gen_args(),
+                optional=False,
+                doc=doc_conf_gen,
+            ),
+            # task formatter
+            Argument("task", dict, [], [variant_task()], optional=False, doc=doc_task),
+            # formatter for `input` section
             Argument("inputs", dict, input_args(), optional=False, doc=doc_inputs),
             Argument(
                 "train", dict, [], [variant_train()], optional=False, doc=doc_train
             ),
             Argument(
-                "explore",
+                "exploration",
                 dict,
                 [],
                 [variant_explore()],
                 optional=False,
                 doc=doc_explore,
+                alias=["explore"],
             ),
-            Argument("fp", dict, [], [variant_fp()], optional=False, doc=doc_fp),
-            Argument("name", str, optional=True, default="dpgen", doc=doc_name),
+            Argument("fp", dict, [], [variant_fp()], optional=True, doc=doc_fp),
+            Argument("name", str, optional=True, default="pfd", doc=doc_name),
             Argument(
                 "parallelism", int, optional=True, default=None, doc=doc_parallelism
             ),
@@ -702,12 +710,12 @@ def submit_args(default_step_config=normalize_step_dict({})):
 
 def normalize(data):
     default_step_config = normalize_step_dict(data.get("default_step_config", {}))
+    # nested arguments formatter
     defs = submit_args(default_step_config)
-
     base = Argument("base", dict, defs)
     data = base.normalize_value(data, trim_pattern="_*")
     # not possible to strictly check arguments, dirty hack!
-    # base.check_value(data, strict=False)
+    base.check_value(data, strict=False)
     return data
 
 
diff --git a/pfd/entrypoint/submit.py b/pfd/entrypoint/submit.py
index 9566350..536e0f6 100644
--- a/pfd/entrypoint/submit.py
+++ b/pfd/entrypoint/submit.py
@@ -1,5 +1,6 @@
 from copy import deepcopy
 import os
+import copy
 from typing import Dict, List, Optional, Tuple
 import json
 import dpdata
@@ -271,8 +272,8 @@ def workflow_dist(
 
     # uploaded python packages
     upload_python_packages = []
-    custom_packages = config.get("upload_python_packages", [])
-    upload_python_packages.extend(custom_packages)
+    if custom_packages := config.get("upload_python_packages"):
+        upload_python_packages.extend(custom_packages)
     upload_python_packages.extend(list(dpdata.__path__))
     upload_python_packages.extend(list(dflow.__path__))
     upload_python_packages.extend(list(pfd.__path__))
@@ -287,8 +288,8 @@ def workflow_dist(
     # exploration
     explore_config = config["exploration"]["config"]
     expl_stages = config["exploration"]["stages"]
-    converge_config = config["exploration"]["converge_config"]
-    max_iter = config["exploration"].get("max_iter", 1)
+    converge_config = config["exploration"]["convergence"]
+    max_iter = config["exploration"]["max_numb_iter"]
     # train
     train_config = config["train"]["config"]
     # type_map_train=config["train"]["type_map"]
@@ -310,7 +311,18 @@ def workflow_dist(
     init_confs = upload_artifact(
         config["conf_generation"]["init_configurations"]["files"]
     )
-    teacher_model = upload_artifact([config["inputs"]["teacher_model"]])
+    # teacher models
+    teacher_models_paths = config["inputs"]["teacher_models_paths"]
+    print(teacher_models_paths)
+    if config["inputs"]["teacher_models_uri"] is not None:
+        print("Using uploaded model at: ", config["inputs"]["teacher_models_uri"])
+        teacher_models = get_artifact_from_uri(config["train"]["init_models_uri"])
+    elif teacher_models_paths is not None:
+        teacher_models = upload_artifact_and_print_uri(
+            teacher_models_paths, "teacher_models"
+        )
+    else:
+        raise FileNotFoundError("Pre-trained model must exist!")
 
     # init_data
     if config["inputs"]["init_data_uri"] is not None:
@@ -358,7 +370,7 @@ def workflow_dist(
         },
         artifacts={
             "init_confs": init_confs,
-            "teacher_model": teacher_model,
+            "teacher_model": teacher_models,
             "init_data": init_data,
             "iter_data": upload_artifact([]),
         },
@@ -390,8 +402,8 @@ def workflow_finetune(config: Dict) -> Step:
     run_lmp_config = config["step_configs"].get("run_explore_config", default_config)
     # uploaded python packages
     upload_python_packages = []
-    custom_packages = config.get("upload_python_packages", [])
-    upload_python_packages.extend(custom_packages)
+    if custom_packages := config.get("upload_python_packages"):
+        upload_python_packages.extend(custom_packages)
     upload_python_packages.extend(list(dpdata.__path__))
     upload_python_packages.extend(list(dflow.__path__))
     upload_python_packages.extend(list(pfd.__path__))
@@ -406,19 +418,23 @@ def workflow_finetune(config: Dict) -> Step:
         mass_map = [getattr(elements, ii).mass for ii in type_map]
     train_style = config["train"]["type"]
     train_config = config["train"]["config"]
-    numb_models = config["train"].get("numb_models", 1)
+    numb_models = config["train"]["numb_models"]
     pert_config = config["conf_generation"]
-    explore_config = config["exploration"]["md"]["config"]
-    max_iter = config["exploration"]["md"].get("max_iter", 1)
-    converge_config = config["exploration"]["converge_config"]
+
+    explore_config = config["exploration"]["config"]
+    max_iter = config["exploration"]["max_numb_iter"]
+    converge_config = config["exploration"]["convergence"]
     # conf selectors
     conf_filters = get_conf_filters(config["exploration"]["filter"])
     render = TrajRenderLammps(nopbc=False)
     conf_selector = ConfSelectorFrames(render, config["fp"]["task_max"], conf_filters)
-    explore_style = config["exploration"]["md"]["type"]
-    expl_stages = config["exploration"]["md"]["stages"]
-    init_training = config["exploration"].get("init_training", False)
-    skip_aimd = config["exploration"].get("skip_aimd", True)
+    explore_style = config["exploration"]["type"]
+    expl_stages = config["exploration"]["stages"]
+    # task
+    init_training = config["task"]["init_training"]
+    if init_training is False:
+        print("No initial training before exploration")
+    skip_aimd = config["task"]["skip_aimd"]
     if skip_aimd is True:
         print("AIMD is exploration skipped!")
 
@@ -433,10 +449,18 @@ def workflow_finetune(config: Dict) -> Step:
     # read training template
     with open(config["train"]["template_script"], "r") as fp:
         template_script = json.load(fp)
-    init_confs = upload_artifact(
-        config["conf_generation"]["init_configurations"]["files"]
-    )
-
+    # init_confs
+    if config["conf_generation"]["init_configurations"]["confs_uri"] is not None:
+        init_confs = get_artifact_from_uri(
+            config["conf_generation"]["init_configurations"]["confs_uri"]
+        )
+    elif config["conf_generation"]["init_configurations"]["files"] is not None:
+        init_confs_prefix = config["conf_generation"]["init_configurations"]["prefix"]
+        init_confs = config["conf_generation"]["init_configurations"]["files"]
+        init_confs = get_systems_from_data(init_confs, init_confs_prefix)
+        init_confs = upload_artifact_and_print_uri(init_confs, "init_data")
+    else:
+        raise RuntimeError("init_confs must be provided")
     # init_data
     if config["inputs"]["init_data_uri"] is not None:
         init_data = get_artifact_from_uri(config["inputs"]["init_data_uri"])
@@ -448,10 +472,10 @@ def workflow_finetune(config: Dict) -> Step:
     else:
         init_data = upload_artifact([])
     iter_data = upload_artifact([])
-    init_models_paths = config["train"].get("init_models_paths", None)
-    if config["train"].get("init_models_url") is not None:
-        print("Using uploaded model at: ", config["train"].get("init_models_url"))
-        init_models = get_artifact_from_uri(config["train"]["init_models_url"])
+    init_models_paths = config["train"]["init_models_paths"]
+    if config["train"]["init_models_uri"] is not None:
+        print("Using uploaded model at: ", config["train"]["init_models_uri"])
+        init_models = get_artifact_from_uri(config["train"]["init_models_uri"])
     elif init_models_paths is not None:
         init_models = upload_artifact_and_print_uri(init_models_paths, "init_models")
     else:
@@ -465,13 +489,16 @@ def workflow_finetune(config: Dict) -> Step:
     fp_config["run"] = config["fp"]["run_config"]
     fp_config["extra_output_files"] = config["fp"].get("extra_output_files", [])
 
+    # aimd exploration
     aimd_config = {}
-    aimd_inputs_config = config["aimd"]["inputs_config"]
-    aimd_type = config["aimd"]["type"]
-    aimd_inputs = fp_styles[aimd_type]["inputs"](**aimd_inputs_config)
-    aimd_config["inputs"] = aimd_inputs
-    aimd_config["run"] = config["aimd"]["run_config"]
-    aimd_config["extra_output_files"] = config["aimd"].get("extra_output_files", [])
+    if skip_aimd is not True:
+        aimd_config.update(fp_config)
+        aimd_inputs_config = copy.deepcopy(fp_inputs_config)
+        aimd_inputs_config.update(config["aimd"]["inputs_config"])
+        aimd_inputs = fp_styles[fp_type]["inputs"](**aimd_inputs_config)
+        aimd_config["inputs"] = aimd_inputs
+        # aimd_config["run"] = config["aimd"]["run_config"]
+        # aimd_config["extra_output_files"] = config["aimd"].get("extra_output_files", [])
 
     # make distillation op
     ft_op = make_ft_op(
diff --git a/pfd/exploration/converge/check_conv.py b/pfd/exploration/converge/check_conv.py
index 8fc2b45..0264333 100644
--- a/pfd/exploration/converge/check_conv.py
+++ b/pfd/exploration/converge/check_conv.py
@@ -6,3 +6,16 @@ class CheckConv(ABC):
     @abstractmethod
     def check_conv(self):
         pass
+
+    @classmethod
+    @abstractmethod
+    def doc(cls):
+        return "The default method doc"
+
+    @classmethod
+    @abstractmethod
+    def args(cls):
+        """
+        The default arguments for the method
+        """
+        return []
diff --git a/pfd/exploration/converge/energy_conv.py b/pfd/exploration/converge/energy_conv.py
index bb2b072..ceb024d 100644
--- a/pfd/exploration/converge/energy_conv.py
+++ b/pfd/exploration/converge/energy_conv.py
@@ -1,4 +1,5 @@
 from .check_conv import CheckConv
+from dargs import Argument
 
 
 class EnerConvRMSE(CheckConv):
@@ -16,3 +17,11 @@ def check_conv(self, test_res_ls, conv_config: dict):
         if weighted_rmse < conv_rmse:
             converged = True
         return converged
+
+    @classmethod
+    def args(cls):
+        return [Argument("RMSE", float, optional=True, default=0.01)]
+
+    @classmethod
+    def doc(cls):
+        return "Converge by RMSE of energy per atom"
diff --git a/pfd/exploration/converge/force_conv.py b/pfd/exploration/converge/force_conv.py
index 6b7df1a..310118f 100644
--- a/pfd/exploration/converge/force_conv.py
+++ b/pfd/exploration/converge/force_conv.py
@@ -1,4 +1,5 @@
 from .check_conv import CheckConv
+from dargs import Argument
 
 
 class ForceConvRMSE(CheckConv):
@@ -16,3 +17,11 @@ def check_conv(self, test_res_ls, conv_config: dict):
         if weighted_rmse < conv_rmse:
             converged = True
         return converged
+
+    @classmethod
+    def args(cls):
+        return [Argument("RMSE", float, optional=True, default=0.01)]
+
+    @classmethod
+    def doc(cls):
+        return "Converge by RMSE of atomic forces"


Test patch:
diff --git a/tests/test_dist.py b/tests/test_dist.py
new file mode 100644
index 0000000..9b828c2
--- /dev/null
+++ b/tests/test_dist.py
@@ -0,0 +1,101 @@
+import unittest
+import os
+import json
+import sys
+
+sys.path.append("../")
+import tempfile
+from pfd.entrypoint.submit import submit_dist
+from dflow.python import OPIO
+
+
+class TestWorkDist(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+        # poscar
+        self.poscar = os.path.join(self.test_dir, "foo.vasp")
+        with open(self.poscar, "w") as f:
+            f.write("foo")
+        # model file
+        self.model = os.path.join(self.test_dir, "foo.pt")
+        with open(self.model, "w") as f:
+            f.write("foo")
+        # model train script file
+        self.train_script = os.path.join(self.test_dir, "foo.json")
+        with open(self.train_script, "w") as f:
+            json.dump({"foo": "foo"}, f, indent=4)
+        self.config_param = {
+            "default_step_config": {
+                "template_config": {
+                    "image": "registry.dp.tech/dptech/deepmd-kit:2024Q1-d23cf3e"
+                },
+            },
+            "task": {"type": "finetune"},
+            "inputs": {
+                "type_map": ["foo"],
+                "mass_map": [1.0],
+                "teacher_models_paths": [self.model],
+            },
+            "conf_generation": {
+                "init_configurations": {
+                    "type": "file",
+                    "prefix": "./",
+                    "fmt": "vasp/poscar",
+                    "files": [self.poscar],
+                },
+                "pert_generation": [
+                    {
+                        "conf_idx": [0],
+                        "atom_pert_distance": 0.1,
+                        "cell_pert_fraction": 0.03,
+                        "pert_num": 5,
+                    }
+                ],
+            },
+            "exploration": {
+                "init_training": True,
+                "skip_aimd": True,
+                "max_iter": 1,
+                "converge_config": {"type": "energy_rmse", "RMSE": 0.01},
+                "filter": [{"type": "distance"}],
+                "type": "lmp",
+                "config": {
+                    "command": "lmp -var restart 0",
+                    "shuffle_models": False,
+                    "head": None,
+                },
+                "stages": [
+                    [
+                        {
+                            "conf_idx": [0],
+                            "n_sample": 1,
+                            "exploration": {
+                                "type": "lmp-md",
+                                "ensemble": "npt",
+                                "dt": 0.005,
+                                "nsteps": 100,
+                                "temps": [300],
+                                "press": [1],
+                                "trj_freq": 1,
+                            },
+                            "max_sample": 100,
+                        }
+                    ]
+                ],
+            },
+            "train": {
+                "comment": "Training script for downstream DeePMD model",
+                "type": "dp",
+                "config": {"impl": "pytorch", "init_model_policy": "no"},
+                "template_script": self.train_script,
+            },
+        }
+
+    os.environ["DFLOW_DEBUG"] = "1"
+
+    def test_wf_io(self):
+        assert submit_dist(self.config_param, no_submission=True)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_ft.py b/tests/test_ft.py
new file mode 100644
index 0000000..c488a8a
--- /dev/null
+++ b/tests/test_ft.py
@@ -0,0 +1,133 @@
+import unittest
+import os
+import json
+import sys
+
+sys.path.append("../")
+import tempfile
+from pfd.entrypoint.submit import submit_ft
+from dflow.python import OPIO
+
+
+class TestWorkflowFinetune(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = tempfile.mkdtemp()
+        # poscar
+        self.poscar = os.path.join(self.test_dir, "foo.vasp")
+        with open(self.poscar, "w") as f:
+            f.write("foo")
+        # INPUT file
+        self.input = os.path.join(self.test_dir, "foo.scf")
+        with open(self.input, "w") as f:
+            f.write("foo")
+        # pp file
+        self.pp = os.path.join(self.test_dir, "foo.upf")
+        with open(self.pp, "w") as f:
+            f.write("foo")
+        # orb file
+        self.orb = os.path.join(self.test_dir, "foo.orb")
+        with open(self.orb, "w") as f:
+            f.write("foo")
+        # model file
+        self.model = os.path.join(self.test_dir, "foo.pt")
+        with open(self.model, "w") as f:
+            f.write("foo")
+        # model train script file
+        self.train_script = os.path.join(self.test_dir, "foo.json")
+        with open(self.train_script, "w") as f:
+            json.dump({"foo": "foo"}, f, indent=4)
+        self.config_param = {
+            "default_step_config": {
+                "template_config": {
+                    "image": "registry.dp.tech/dptech/deepmd-kit:2024Q1-d23cf3e"
+                },
+            },
+            "task": {"type": "finetune", "init_training": True, "skip_aimd": False},
+            "inputs": {"type_map": ["foo"], "mass_map": [1.0]},
+            "conf_generation": {
+                "init_configurations": {
+                    "type": "file",
+                    "prefix": "./",
+                    "fmt": "vasp/poscar",
+                    "files": [self.poscar],
+                },
+                "pert_generation": [
+                    {
+                        "conf_idx": [0],
+                        "atom_pert_distance": 0.1,
+                        "cell_pert_fraction": 0.03,
+                        "pert_num": 5,
+                    }
+                ],
+            },
+            "aimd": {"inputs_config": {"input_file": self.input}},
+            "exploration": {
+                "init_training": True,
+                "skip_aimd": True,
+                "max_iter": 1,
+                "converge_config": {"type": "energy_rmse", "RMSE": 0.01},
+                "filter": [{"type": "distance"}],
+                "type": "lmp",
+                "config": {
+                    "command": "lmp -var restart 0",
+                    "shuffle_models": False,
+                    "head": None,
+                },
+                "stages": [
+                    [
+                        {
+                            "conf_idx": [0],
+                            "n_sample": 1,
+                            "exploration": {
+                                "type": "lmp-md",
+                                "ensemble": "npt",
+                                "dt": 0.005,
+                                "nsteps": 100,
+                                "temps": [300],
+                                "press": [1],
+                                "trj_freq": 1,
+                            },
+                            "max_sample": 100,
+                        }
+                    ]
+                ],
+            },
+            "fp": {
+                "type": "fpop_abacus",
+                "task_max": 50,
+                "extra_output_files:": [],
+                "run_config": {
+                    "command": "OMP_NUM_THREADS=4 mpirun -np 8 abacus | tee log"
+                },
+                "inputs_config": {
+                    "input_file": self.input,
+                    "pp_files": {
+                        "foo": self.pp,
+                    },
+                    "orb_files": {
+                        "foo": self.orb,
+                    },
+                },
+            },
+            "train": {
+                "comment": "Training script for downstream DeePMD model",
+                "type": "dp",
+                "config": {
+                    "impl": "pytorch",
+                    "init_model_policy": "no",
+                    "init_model_with_finetune": True,
+                },
+                "template_script": self.train_script,
+                "init_models_paths": [self.model],
+                "numb_models": 1,
+            },
+        }
+
+    os.environ["DFLOW_DEBUG"] = "1"
+
+    def test_wf_io(self):
+        assert submit_ft(self.config_param, no_submission=True)
+
+
+if __name__ == "__main__":
+    unittest.main()