New format of submit input. add dargs check. (#78)

* fix bugs of import functions

* fix bug of context

* add example of new input

* new format of submit input. add dargs check.

* command line option for old-style compatibility

* fix s3 config of commands

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

dpgen2/entrypoint/download.py

@@ -3,7 +3,7 @@
     Workflow,
 )
 from dpgen2.utils import (
-    dflow_config,
+    workflow_config_from_dict,
 )
 from dpgen2.utils.dflow_query import (
     matched_step_key,
@@ -21,8 +21,7 @@ def download(
         wf_keys : Optional[List] = None,
         prefix : Optional[str] = None,
 ):
-    dflow_config_data = wf_config.get('dflow_config', None)
-    dflow_config(dflow_config_data)
+    workflow_config_from_dict(wf_config)
 
     wf = Workflow(id=workflow_id)
 

dpgen2/entrypoint/main.py

@@ -66,6 +66,9 @@ def main_parser() -> argparse.ArgumentParser:
     parser_run.add_argument(
         "CONFIG", help="the config file in json format defining the workflow."
     )
+    parser_run.add_argument(
+        "-o", "--old-compatible", action='store_true', help="compatible with old-style input script used in dpgen2 < 0.0.6."
+    )
 
     ##########################################
     # resubmit
@@ -81,11 +84,14 @@ def main_parser() -> argparse.ArgumentParser:
         "ID", help="the ID of the existing workflow."
     )
     parser_resubmit.add_argument(
-        "--list", action='store_true', help="list the Steps of the existing workflow."
+        "-l", "--list", action='store_true', help="list the Steps of the existing workflow."
     )
     parser_resubmit.add_argument(
         "--reuse", type=str, nargs='+', default=None, help="specify which Steps to reuse."
     )
+    parser_resubmit.add_argument(
+        "-o", "--old-compatible", action='store_true', help="compatible with old-style input script used in dpgen2 < 0.0.6."
+    )
 
     ##########################################
     # status
@@ -185,13 +191,20 @@ def main():
     if args.command == "submit":
         with open(args.CONFIG) as fp:
             config = json.load(fp)
-        submit_concurrent_learning(config)
+        submit_concurrent_learning(
+            config,
+            old_style=args.old_compatible,
+        )
     elif args.command == "resubmit":
         with open(args.CONFIG) as fp:
             config = json.load(fp)
         wfid = args.ID
         resubmit_concurrent_learning(
-            config, wfid, list_steps=args.list, reuse=args.reuse,
+            config, 
+            wfid, 
+            list_steps=args.list, 
+            reuse=args.reuse,
+            old_style=args.old_compatible,
         )
     elif args.command == "status":
         with open(args.CONFIG) as fp:

dpgen2/entrypoint/status.py

@@ -3,7 +3,7 @@
     Workflow,
 )
 from dpgen2.utils import (
-    dflow_config,
+    workflow_config_from_dict,
 )
 from dpgen2.utils.dflow_query import (
     get_last_scheduler,
@@ -16,8 +16,7 @@ def status(
         workflow_id,
         wf_config : Optional[Dict] = {}, 
 ):
-    dflow_config_data = wf_config.get('dflow_config', None)
-    dflow_config(dflow_config_data)
+    workflow_config_from_dict(wf_config)
 
     wf = Workflow(id=workflow_id)
 

dpgen2/entrypoint/submit.py

@@ -70,13 +70,14 @@
     load_object_from_file,
     normalize_alloy_conf_dict,
     generate_alloy_conf_file_content,
-    dflow_config,
     sort_slice_ops,
     print_keys_in_nice_format,
+    workflow_config_from_dict,
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
+from dpgen2.entrypoint.submit_args import normalize as normalize_submit_args
 from typing import (
-    Union, List,
+    Union, List, Dict, Optional,
 )
 
 default_config = normalize_step_dict(
@@ -196,24 +197,30 @@ def make_conf_list(
 
 def make_naive_exploration_scheduler(
         config,
+        old_style = False,
 ):
     # use npt task group
-    model_devi_jobs = config['model_devi_jobs']
-    sys_configs = config['sys_configs']
-    mass_map = config['mass_map']
-    type_map = config['type_map']
-    numb_models = config['numb_models']
-    fp_task_max = config['fp_task_max']
-    conv_accuracy = config['conv_accuracy']
-    max_numb_iter = config['max_numb_iter']
-    fatal_at_max = config.get('fatal_at_max', True)
+    model_devi_jobs = config['model_devi_jobs'] if old_style else config['explore']['stages']
+    sys_configs = config['sys_configs'] if old_style else config['explore']['configurations']
+    sys_prefix = config.get('sys_prefix')
+    if sys_prefix is not None:
+        for ii in range(len(sys_configs)):
+            if isinstance(sys_configs[ii], list):
+                sys_configs[ii] = [os.path.join(sys_prefix, jj) for jj in sys_prefix[ii]]
+    mass_map = config['mass_map'] if old_style else config['inputs']['mass_map']
+    type_map = config['type_map'] if old_style else config['inputs']['type_map']
+    numb_models = config['numb_models'] if old_style else config['train']['numb_models']
+    fp_task_max = config['fp_task_max'] if old_style else config['fp']['task_max']
+    conv_accuracy = config['conv_accuracy'] if old_style else config['explore']['conv_accuracy']
+    max_numb_iter = config['max_numb_iter'] if old_style else config['explore']['max_numb_iter']
+    fatal_at_max = config.get('fatal_at_max', True) if old_style else config['explore']['fatal_at_max']
     scheduler = ExplorationScheduler()
 
     for job in model_devi_jobs:
         # task group
         tgroup = NPTTaskGroup()
         ##  ignore the expansion of sys_idx
-        # get all file names of md initial configuraitons
+        # get all file names of md initial configurations
         sys_idx = job['sys_idx']
         conf_list = []        
         for ii in sys_idx:
@@ -244,8 +251,10 @@ def make_naive_exploration_scheduler(
         stage.add_task_group(tasks)
         # trust level
         trust_level = TrustLevel(
-            config['model_devi_f_trust_lo'],
-            config['model_devi_f_trust_hi'],
+            config['model_devi_f_trust_lo'] if old_style else config['explore']['f_trust_lo'],
+            config['model_devi_f_trust_hi'] if old_style else config['explore']['f_trust_hi'],
+            level_v_lo=config.get('model_devi_v_trust_lo') if old_style else config['explore']['v_trust_lo'],
+            level_v_hi=config.get('model_devi_v_trust_hi') if old_style else config['explore']['v_trust_hi'],
         )
         # selector
         selector = ConfSelectorLammpsFrames(
@@ -285,22 +294,23 @@ def get_kspacing_kgamma_from_incar(
 
 
 def workflow_concurrent_learning(
-        config,
+        config : Dict,
+        old_style : Optional[bool] = False,
 ):
-    default_config = normalize_step_dict(config.get('default_config', {}))
-
-    train_style = config['train_style']
-    explore_style = config['explore_style']
-    fp_style = config['fp_style']
-    prep_train_config = normalize_step_dict(config.get('prep_train_config', default_config))
-    run_train_config = normalize_step_dict(config.get('run_train_config', default_config))
-    prep_explore_config = normalize_step_dict(config.get('prep_explore_config', default_config))
-    run_explore_config = normalize_step_dict(config.get('run_explore_config', default_config))
-    prep_fp_config = normalize_step_dict(config.get('prep_fp_config', default_config))
-    run_fp_config = normalize_step_dict(config.get('run_fp_config', default_config))
-    select_confs_config = normalize_step_dict(config.get('select_confs_config', default_config))
-    collect_data_config = normalize_step_dict(config.get('collect_data_config', default_config))
-    cl_step_config = normalize_step_dict(config.get('cl_step_config', default_config))
+    default_config = normalize_step_dict(config.get('default_config', {})) if old_style else config['default_step_config']
+
+    train_style = config.get('train_style', 'dp') if old_style else config['train']['type']
+    explore_style = config.get('explore_style', 'lmp') if old_style else config['explore']['type']
+    fp_style = config.get('fp_style', 'vasp') if old_style else config['fp']['type']
+    prep_train_config = normalize_step_dict(config.get('prep_train_config', default_config)) if old_style else config['step_configs']['prep_train_config']
+    run_train_config = normalize_step_dict(config.get('run_train_config', default_config)) if old_style else config['step_configs']['run_train_config']
+    prep_explore_config = normalize_step_dict(config.get('prep_explore_config', default_config)) if old_style else config['step_configs']['prep_explore_config']
+    run_explore_config = normalize_step_dict(config.get('run_explore_config', default_config)) if old_style else config['step_configs']['run_explore_config']
+    prep_fp_config = normalize_step_dict(config.get('prep_fp_config', default_config)) if old_style else config['step_configs']['prep_fp_config']
+    run_fp_config = normalize_step_dict(config.get('run_fp_config', default_config)) if old_style else config['step_configs']['run_fp_config']
+    select_confs_config = normalize_step_dict(config.get('select_confs_config', default_config)) if old_style else config['step_configs']['select_confs_config']
+    collect_data_config = normalize_step_dict(config.get('collect_data_config', default_config)) if old_style else config['step_configs']['collect_data_config']
+    cl_step_config = normalize_step_dict(config.get('cl_step_config', default_config)) if old_style else config['step_configs']['cl_step_config']
     upload_python_package = config.get('upload_python_package', None)
     init_models_paths = config.get('training_iter0_model_path')
 
@@ -319,24 +329,27 @@ def workflow_concurrent_learning(
         cl_step_config = cl_step_config,
         upload_python_package = upload_python_package,
     )
-    scheduler = make_naive_exploration_scheduler(config)
-
-    type_map = config['type_map']
-    numb_models = config['numb_models']
-    template_script = config['default_training_param']
-    train_config = {}
-    lmp_config = config.get('lmp_config', {})
-    fp_config = config.get('fp_config', {})
-    kspacing, kgamma = get_kspacing_kgamma_from_incar(config['fp_incar'])
-    fp_pp_files = config['fp_pp_files']
-    incar_file = config['fp_incar']    
+    scheduler = make_naive_exploration_scheduler(config, old_style=old_style)
+
+    type_map = config['type_map'] if old_style else config['inputs']['type_map']
+    numb_models = config['numb_models'] if old_style else config['train']['numb_models']
+    template_script = config['default_training_param'] if old_style else config['train']['template_script']
+    train_config = {} if old_style else config['train']['config']
+    lmp_config = config.get('lmp_config', {}) if old_style else config['explore']['config']
+    fp_config = config.get('fp_config', {}) if old_style else config['fp']['config']
+    kspacing, kgamma = get_kspacing_kgamma_from_incar(config['fp_incar'] if old_style else config['fp']['incar'])
+    fp_pp_files = config['fp_pp_files'] if old_style else config['fp']['pp_files']
+    incar_file = config['fp_incar'] if old_style else config['fp']['incar']
     fp_inputs = VaspInputs(
         kspacing = kspacing,
         kgamma = kgamma,
         incar_template_name = incar_file,
         potcar_names = fp_pp_files,
     )
-    init_data = config['init_data_sys']
+    init_data_prefix = config.get('init_data_prefix') if old_style else config['inputs']['init_data_prefix']
+    init_data = config['init_data_sys'] if old_style else config['inputs']['init_data_sys']
+    if init_data_prefix is not None:
+        init_data = [os.path.join(init_data_prefix, ii) for ii in init_data_sys]
     if isinstance(init_data,str):
         init_data = expand_sys_str(init_data)
     init_data = upload_artifact(init_data)
@@ -372,8 +385,7 @@ def workflow_concurrent_learning(
 def wf_global_workflow(
         wf_config,
 ):
-    dflow_config_data = wf_config.get('dflow_config', None)
-    dflow_config(dflow_config_data)
+    workflow_config_from_dict(wf_config)
 
     # lebesgue context
     from dflow.plugins.lebesgue import LebesgueContext
@@ -391,10 +403,13 @@ def wf_global_workflow(
 def submit_concurrent_learning(
         wf_config,
         reuse_step = None,
+        old_style = False,
 ):
+    wf_config = normalize_submit_args(wf_config)
+
     context = wf_global_workflow(wf_config)
 
-    dpgen_step = workflow_concurrent_learning(wf_config)
+    dpgen_step = workflow_concurrent_learning(wf_config, old_style=old_style)
 
     wf = Workflow(name="dpgen", context=context)
     wf.add(dpgen_step)
@@ -450,7 +465,10 @@ def resubmit_concurrent_learning(
         wfid,
         list_steps = False,
         reuse = None,
+        old_style = False,
 ):
+    wf_config = normalize_submit_args(wf_config)
+
     context = wf_global_workflow(wf_config)
 
     old_wf = Workflow(id=wfid)
@@ -474,6 +492,7 @@ def resubmit_concurrent_learning(
     wf = submit_concurrent_learning(
         wf_config, 
         reuse_step=reuse_step,
+        old_style=old_style,
     )
 
     return wf