Implement step configuration (#38)

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

Author: wanghan-iapcm

dpgen2/entrypoint/submit.py

@@ -69,28 +69,36 @@
     dump_object_to_file,
     load_object_from_file,
 )
+from dpgen2.utils.step_config import normalize as normalize_step_dict
+default_config = normalize_step_dict(
+    {
+        "template_config" : {
+            "image" : default_image,
+        }
+    }
+)
 
 def make_concurrent_learning_op (
         train_style : str = 'dp',
         explore_style : str = 'lmp',
         fp_style : str = 'vasp',
-        prep_train_image : str = default_image,
-        run_train_image : str = default_image,
-        prep_explore_image : str = default_image,
-        run_explore_image : str = default_image,
-        prep_fp_image : str = default_image,
-        run_fp_image : str = default_image,
-        select_confs_image : str = default_image,
-        collect_data_image : str = default_image,
+        prep_train_config : str = default_config,
+        run_train_config : str = default_config,
+        prep_explore_config : str = default_config,
+        run_explore_config : str = default_config,
+        prep_fp_config : str = default_config,
+        run_fp_config : str = default_config,
+        select_confs_config : str = default_config,
+        collect_data_config : str = default_config,
         upload_python_package : bool = None,
 ):
     if train_style == 'dp':
         prep_run_train_op = PrepRunDPTrain(
             "prep-run-dp-train",
             PrepDPTrain,
             RunDPTrain,
-            prep_image = prep_train_image,
-            run_image = run_train_image,
+            prep_config = prep_train_config,
+            run_config = run_train_config,
             upload_python_package = upload_python_package,
         )
     else:
@@ -100,8 +108,8 @@ def make_concurrent_learning_op (
             "prep-run-lmp",
             PrepLmp,
             RunLmp,
-            prep_image = prep_explore_image,
-            run_image = run_explore_image,
+            prep_config = prep_explore_config,
+            run_config = run_explore_config,
             upload_python_package = upload_python_package,
         )
     else:
@@ -111,8 +119,8 @@ def make_concurrent_learning_op (
             "prep-run-vasp",
             PrepVasp,
             RunVasp,
-            prep_image = prep_fp_image,
-            run_image = run_fp_image,
+            prep_config = prep_fp_config,
+            run_config = run_fp_config,
             upload_python_package = upload_python_package,
         )
     else:
@@ -126,16 +134,16 @@ def make_concurrent_learning_op (
         SelectConfs,
         prep_run_fp_op,
         CollectData,
-        select_confs_image = select_confs_image,
-        collect_data_image = collect_data_image,
+        select_confs_config = select_confs_config,
+        collect_data_config = collect_data_config,
         upload_python_package = upload_python_package,
     )    
     # dpgen
     dpgen_op = ConcurrentLearning(
         "concurrent-learning",
         block_cl_op,
         upload_python_package = upload_python_package,
-        image = default_image,
+        step_config = default_config,
     )
 
     return dpgen_op
@@ -242,25 +250,25 @@ def workflow_concurrent_learning(
     train_style = config['train_style']
     explore_style = config['explore_style']
     fp_style = config['fp_style']
-    prep_train_image = config['prep_train_image']
-    run_train_image = config['run_train_image']
-    prep_explore_image = config['prep_explore_image']
-    run_explore_image = config['run_explore_image']
-    prep_fp_image = config['prep_fp_image']
-    run_fp_image = config['run_fp_image']
+    prep_train_config = normalize_step_dict(config.get('prep_train_config', {}))
+    run_train_config = normalize_step_dict(config.get('run_train_config', {}))
+    prep_explore_config = normalize_step_dict(config.get('prep_explore_config', {}))
+    run_explore_config = normalize_step_dict(config.get('run_explore_config', {}))
+    prep_fp_config = normalize_step_dict(config.get('prep_fp_config', {}))
+    run_fp_config = normalize_step_dict(config.get('run_fp_config', {}))
     upload_python_package = config.get('upload_python_package', None)
     init_models_paths = config.get('training_iter0_model_path')
 
     concurrent_learning_op = make_concurrent_learning_op(
         train_style,
         explore_style,
         fp_style,
-        prep_train_image = prep_train_image,
-        run_train_image = run_train_image,
-        prep_explore_image = prep_explore_image,
-        run_explore_image = run_explore_image,
-        prep_fp_image = prep_fp_image,
-        run_fp_image = run_fp_image,
+        prep_train_config = prep_train_config,
+        run_train_config = run_train_config,
+        prep_explore_config = prep_explore_config,
+        run_explore_config = run_explore_config,
+        prep_fp_config = prep_fp_config,
+        run_fp_config = run_fp_config,
         upload_python_package = upload_python_package,
     )
     scheduler = make_naive_exploration_scheduler(config)
@@ -315,15 +323,39 @@ def workflow_concurrent_learning(
             "iter_data" : iter_data,
         },
     )
-        
-    wf = Workflow(name="dpgen", host=default_host)
-    wf.add(dpgen_step)
-    
-    return wf
+    return dpgen_step
 
 def submit_concurrent_learning(
-        config,
+        wf_config,
 ):
-    wf = workflow_concurrent_learning(config)
+    # set global config
+    from dflow import config, s3_config
+    dflow_config = wf_config.get('dflow_config', None)
+    if dflow_config :
+        config["host"] = dflow_config.get('host', None)
+        s3_config["endpoint"] = dflow_config.get('s3_endpoint', None)
+        config["k8s_api_server"] = dflow_config.get('k8s_api_server', None)
+        config["token"] = dflow_config.get('token', None)    
+
+    # lebesque context
+    from dflow.plugins.lebesgue import LebesgueContext
+    lb_context_config = wf_config.get("lebesque_context_config", None)
+    if lb_context_config:
+        lebesgue_context = LebesgueContext(
+            **lb_context_config,
+        )
+    else :
+        lebesgue_context = None
+
+    # print('config:', config)
+    # print('s3_config:',s3_config)
+    # print('lebsque context:', lb_context_config)
+
+    dpgen_step = workflow_concurrent_learning(wf_config)
+
+    wf = Workflow(name="dpgen", context=lebesgue_context)
+    wf.add(dpgen_step)
+
     wf.submit()
+
     return wf

dpgen2/flow/dpgen_loop.py

@@ -37,6 +37,9 @@
     load_object_from_file,
     dump_object_to_file,
 )
+from dpgen2.utils.step_config import normalize as normalize_step_dict
+from copy import deepcopy
+
 
 class SchedulerWrapper(OP):
 
@@ -117,7 +120,7 @@ def __init__(
             self,
             name : str,
             block_op : Steps,
-            image : str = "dflow:v1.0",
+            step_config : dict = normalize_step_dict({}),
             upload_python_package : str = None,
     ):
         self._input_parameters={
@@ -174,7 +177,7 @@ def __init__(
             self.step_keys,
             name,
             block_op,
-            image = image,
+            step_config = step_config,
             upload_python_package = upload_python_package,
         )
 
@@ -204,13 +207,13 @@ def __init__(
             self,
             name : str,
             block_op : Steps,
-            image : str = "dflow:v1.0",
+            step_config : dict = normalize_step_dict({}),
             upload_python_package : str = None,
     ):
         self.loop = ConcurrentLearningLoop(
             name+'-loop',
             block_op,
-            image = image,
+            step_config = step_config,
             upload_python_package = upload_python_package,
         )
 
@@ -261,7 +264,7 @@ def __init__(
             name, 
             self.loop,
             self.loop_key,
-            image = image,
+            step_config = step_config,
             upload_python_package = upload_python_package,
         )
 
@@ -295,9 +298,12 @@ def _loop (
         step_keys,
         name : str,
         block_op : OP,
-        image : str = "dflow:v1.0",
+        step_config : dict = normalize_step_dict({}),
         upload_python_package : str = None,
 ):    
+    step_config = deepcopy(step_config)
+    step_template_config = step_config.pop('template_config')
+
     block_step = Step(
         name = name + '-block',
         template = block_op,
@@ -326,8 +332,8 @@ def _loop (
         name = name + '-scheduler',
         template=PythonOPTemplate(
             SchedulerWrapper,
-            image=image,
             python_packages = upload_python_package,
+            **step_template_config,
         ),
         parameters={
             "exploration_report": block_step.outputs.parameters['exploration_report'],
@@ -337,22 +343,24 @@ def _loop (
             "trajs" : block_step.outputs.artifacts['trajs'],
         },
         key = step_keys['scheduler'],
+        **step_config,
     )
     steps.add(scheduler_step)
 
     id_step = Step(
         name = name + '-make-block-id',
         template=PythonOPTemplate(
             MakeBlockId,
-            image=image,
             python_packages = upload_python_package,
+            **step_template_config,
         ),
         parameters={
         },
         artifacts={
             "exploration_scheduler": scheduler_step.outputs.artifacts['exploration_scheduler'],
         },
         key = step_keys['id'],
+        **step_config,
     )
     steps.add(id_step)
 
@@ -409,15 +417,18 @@ def _dpgen(
         name,
         loop_op,
         loop_key,
-        image = "dflow:v1.0",
+        step_config : dict = normalize_step_dict({}),
         upload_python_package : str = None
 ):    
+    step_config = deepcopy(step_config)
+    step_template_config = step_config.pop('template_config')
+
     scheduler_step = Step(
         name = name + '-scheduler',
         template=PythonOPTemplate(
             SchedulerWrapper,
-            image=image,
             python_packages = upload_python_package,
+            **step_template_config,
         ),
         parameters={
             "exploration_report": None,
@@ -427,22 +438,24 @@ def _dpgen(
             "trajs" : None,
         },
         key = step_keys['scheduler'],
+        **step_config,
     )
     steps.add(scheduler_step)
 
     id_step = Step(
         name = name + '-make-block-id',
         template=PythonOPTemplate(
             MakeBlockId,
-            image=image,
             python_packages = upload_python_package,
+            **step_template_config,
         ),
         parameters={
         },
         artifacts={
             "exploration_scheduler": scheduler_step.outputs.artifacts['exploration_scheduler'],
         },
         key = step_keys['id'],
+        **step_config,
     )
     steps.add(id_step)
 

dpgen2/op/select_confs.py

@@ -2,7 +2,8 @@
     OP,
     OPIO,
     OPIOSign,
-    Artifact
+    Artifact,
+    BigParameter,
 )
 import os, json
 from typing import Tuple, List, Set