support bohrium platform (#102)

Minor bug fixings:
- fix bug in vasp_input
- fix bug of dflow_s3_config in input arg parsing

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

.github/workflows/test.yml

@@ -20,7 +20,7 @@ jobs:
     - name: Install dependencies
       run: | 
         pip install -e .
-        pip install mock coverage 
+        pip install mock coverage pytest
     - name: Test
       run: SKIP_UT_WITH_DFLOW=0 DFLOW_DEBUG=1 coverage run --source=./dpgen2 -m unittest && coverage report
     - uses: codecov/codecov-action@v3

dpgen2/entrypoint/submit_args.py → dpgen2/entrypoint/args.py

@@ -115,11 +115,11 @@ def input_args():
 
 def dflow_conf_args():
     doc_dflow_config = "The configuration passed to dflow"
-    doc_s3_config = "The S3 configuration passed to dflow"
+    doc_dflow_s3_config = "The S3 configuration passed to dflow"
 
     return [
         Argument("dflow_config", dict, optional=True, default=None, doc=doc_dflow_config),
-        Argument("s3_config", dict, optional=True, default=None, doc=doc_s3_config),
+        Argument("dflow_s3_config", dict, optional=True, default=None, doc=doc_dflow_s3_config),
     ]
 
 def lebesgue_conf_args():
@@ -129,6 +129,26 @@ def lebesgue_conf_args():
         Argument("lebesgue_context_config", dict, optional=True, default=None, doc=doc_lebesgue_context_config),
     ]
 
+def bohrium_conf_args():
+    doc_username = "The username of the Bohrium platform"
+    doc_password = "The password of the Bohrium platform"
+    doc_project_id = "The project ID of the Bohrium platform"
+    doc_host = "The host name of the Bohrium platform. Will overwrite `dflow_config['host']`"
+    doc_k8s_api_server = "The k8s server of the Bohrium platform. Will overwrite `dflow_config['k8s_api_server']`"
+    doc_repo_key = "The repo key of the Bohrium platform. Will overwrite `dflow_s3_config['repo_key']`"
+    doc_storage_client = "The storage client of the Bohrium platform. Will overwrite `dflow_s3_config['storage_client']`"
+
+    return [
+        Argument("username", str, optional=False, doc=doc_username),
+        Argument("password", str, optional=False, doc=doc_password),
+        Argument("project_id", int, optional=False, doc=doc_project_id),
+        Argument("host", str, optional=True, default="https://workflows.deepmodeling.com", doc=doc_host),
+        Argument("k8s_api_server", str, optional=True, default="https://workflows.deepmodeling.com", doc=doc_k8s_api_server),
+        Argument("repo_key", str, optional=True, default="oss-bohrium", doc=doc_repo_key),
+        Argument("storage_client", str, optional=True, default="dflow.plugins.bohrium.TiefblueClient", doc=doc_storage_client),
+    ]
+
+
 def default_step_config_args():
     doc_default_step_config = "The default step configuration."
 
@@ -162,6 +182,7 @@ def dpgen_step_config_args(default_config):
 
 
 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"
@@ -173,6 +194,7 @@ def submit_args(default_step_config = normalize_step_dict({})):
         dflow_conf_args() + \
         lebesgue_conf_args() + \
         default_step_config_args() + [
+        Argument("bohrium_config", dict, bohrium_conf_args(), optional=True, default=None, doc=doc_bohrium_config),
         Argument("step_configs", dict, dpgen_step_config_args(default_step_config), optional=True, default={}, doc=doc_step_configs),
         Argument("upload_python_packages", [list,str], optional=True, default=None, doc=doc_upload_python_packages, alias=["upload_python_package"]),
         Argument("inputs", dict, input_args(), optional=False, doc=doc_inputs),

dpgen2/entrypoint/common.py

@@ -0,0 +1,75 @@
+import dflow
+from pathlib import Path
+from dpgen2.utils import (
+    dump_object_to_file,
+    load_object_from_file,
+    normalize_alloy_conf_dict,
+    generate_alloy_conf_file_content,
+    sort_slice_ops,
+    print_keys_in_nice_format,
+    workflow_config_from_dict,
+    matched_step_key,
+    bohrium_config_from_dict,
+)
+from dpgen2.utils.step_config import normalize as normalize_step_dict
+from typing import (
+    Union, List, Dict, Optional,
+)
+
+
+def global_config_workflow(
+        wf_config,
+        do_lebesgue : bool=False,
+):
+    # dflow_config, dflow_s3_config
+    workflow_config_from_dict(wf_config)
+
+    # lebesgue context
+    lebesgue_context = None
+    if do_lebesgue:
+        from dflow.plugins.lebesgue import LebesgueContext
+        lb_context_config = wf_config.get("lebesgue_context_config", None)
+        if lb_context_config:
+            lebesgue_context = LebesgueContext(
+                **lb_context_config,
+            )
+
+    # bohrium configuration
+    if wf_config.get("bohrium_config") is not None:
+        assert(lebesgue_context is None), \
+            "cannot use bohrium and lebesgue at the same time"
+        bohrium_config_from_dict(wf_config["bohrium_config"])        
+
+    return lebesgue_context
+
+
+def expand_sys_str(root_dir: Union[str, Path]) -> List[str]:
+    root_dir = Path(root_dir)
+    matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()]
+    if (root_dir / "type.raw").is_file():
+        matches.append(str(root_dir))
+    return matches
+
+
+def expand_idx (in_list) :
+    ret = []
+    for ii in in_list :
+        if isinstance(ii, int) :
+            ret.append(ii)
+        elif isinstance(ii, str):
+            step_str = ii.split(':')
+            if len(step_str) > 1 :
+                step = int(step_str[1])
+            else :
+                step = 1
+            range_str = step_str[0].split('-')
+            if len(range_str) == 2:
+                ret += range(int(range_str[0]), int(range_str[1]), step)
+            elif len(range_str) == 1 :
+                ret += [int(range_str[0])]
+            else:
+                raise RuntimeError('not expected range string', step_str[0])
+    ret = sorted(list(set(ret)))
+    return ret
+
+

dpgen2/entrypoint/download.py

@@ -2,9 +2,6 @@
 from dflow import (
     Workflow,
 )
-from dpgen2.utils import (
-    workflow_config_from_dict,
-)
 from dpgen2.utils.dflow_query import (
     matched_step_key,
 )
@@ -14,14 +11,23 @@
 from typing import (
     Optional, Dict, Union, List,
 )
+from dpgen2.entrypoint.common import (
+    global_config_workflow,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
+
 
 def download(
         workflow_id,
         wf_config : Optional[Dict] = {}, 
         wf_keys : Optional[List] = None,
         prefix : Optional[str] = None,
 ):
-    workflow_config_from_dict(wf_config)
+    wf_config = normalize_args(wf_config)
+
+    global_config_workflow(wf_config)
 
     wf = Workflow(id=workflow_id)
 

dpgen2/entrypoint/showkey.py

@@ -5,16 +5,24 @@
 )
 from dpgen2.entrypoint.submit import get_resubmit_keys
 from dpgen2.utils import (
-    workflow_config_from_dict,
     print_keys_in_nice_format,
 )
+from dpgen2.entrypoint.common import (
+    global_config_workflow,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
 
 
 def showkey(
         wf_id,
         wf_config,
 ):
-    workflow_config_from_dict(wf_config)    
+    wf_config = normalize_args(wf_config)
+
+    global_config_workflow(wf_config)
+
     wf = Workflow(id=wf_id)
     all_step_keys = get_resubmit_keys(wf)
     prt_str = print_keys_in_nice_format(

dpgen2/entrypoint/status.py

@@ -2,21 +2,27 @@
 from dflow import (
     Workflow,
 )
-from dpgen2.utils import (
-    workflow_config_from_dict,
-)
 from dpgen2.utils.dflow_query import (
     get_last_scheduler,
 )
 from typing import (
     Optional, Dict, Union, List,
 )
+from dpgen2.entrypoint.common import (
+    global_config_workflow,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
+
 
 def status(
         workflow_id,
         wf_config : Optional[Dict] = {}, 
 ):
-    workflow_config_from_dict(wf_config)
+    wf_config = normalize_args(wf_config)
+
+    global_config_workflow(wf_config)
 
     wf = Workflow(id=workflow_id)
 

dpgen2/entrypoint/submit.py

@@ -72,9 +72,17 @@
     print_keys_in_nice_format,
     workflow_config_from_dict,
     matched_step_key,
+    bohrium_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 dpgen2.entrypoint.common import (
+    global_config_workflow,
+    expand_sys_str,
+    expand_idx,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
 from typing import (
     Union, List, Dict, Optional,
 )
@@ -87,12 +95,6 @@
     }
 )
 
-def expand_sys_str(root_dir: Union[str, Path]) -> List[str]:
-    root_dir = Path(root_dir)
-    matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()]
-    if (root_dir / "type.raw").is_file():
-        matches.append(str(root_dir))
-    return matches
 
 def make_concurrent_learning_op (
         train_style : str = 'dp',
@@ -340,7 +342,7 @@ def workflow_concurrent_learning(
     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']
+    fp_config = config.get('fp_config', {}) if old_style else {}
     if old_style:        
         potcar_names = config['fp_pp_files']
         incar_template_name = config['fp_incar']
@@ -393,32 +395,17 @@ def workflow_concurrent_learning(
     return dpgen_step
 
 
-def wf_global_workflow(
-        wf_config,
-):
-    workflow_config_from_dict(wf_config)
-
-    # lebesgue context
-    from dflow.plugins.lebesgue import LebesgueContext
-    lb_context_config = wf_config.get("lebesgue_context_config", None)
-    if lb_context_config:
-        lebesgue_context = LebesgueContext(
-            **lb_context_config,
-        )
-    else :
-        lebesgue_context = None
-
-    return lebesgue_context
-
-
 def submit_concurrent_learning(
         wf_config,
         reuse_step = None,
         old_style = False,
 ):
-    wf_config = normalize_submit_args(wf_config)
+    # normalize args
+    wf_config = normalize_args(wf_config)
+
+    do_lebesgue = wf_config.get("lebesgue_context_config", None) is not None
 
-    context = wf_global_workflow(wf_config)
+    context = global_config_workflow(wf_config, do_lebesgue=do_lebesgue)
 
     dpgen_step = workflow_concurrent_learning(wf_config, old_style=old_style)
 
@@ -439,28 +426,6 @@ def print_list_steps(
     return '\n'.join(ret)
 
 
-def expand_idx (in_list) :
-    ret = []
-    for ii in in_list :
-        if isinstance(ii, int) :
-            ret.append(ii)
-        elif isinstance(ii, str):
-            step_str = ii.split(':')
-            if len(step_str) > 1 :
-                step = int(step_str[1])
-            else :
-                step = 1
-            range_str = step_str[0].split('-')
-            if len(range_str) == 2:
-                ret += range(int(range_str[0]), int(range_str[1]), step)
-            elif len(range_str) == 1 :
-                ret += [int(range_str[0])]
-            else:
-                raise RuntimeError('not expected range string', step_str[0])
-    ret = sorted(list(set(ret)))
-    return ret
-
-
 def successful_step_keys(wf):
     all_step_keys_ = wf.query_keys_of_steps()
     wf_info = wf.query()
@@ -492,9 +457,9 @@ def resubmit_concurrent_learning(
         reuse = None,
         old_style = False,
 ):
-    wf_config = normalize_submit_args(wf_config)
+    wf_config = normalize_args(wf_config)
 
-    context = wf_global_workflow(wf_config)
+    context = global_config_workflow(wf_config)
 
     old_wf = Workflow(id=wfid)
     all_step_keys = get_resubmit_keys(old_wf)

dpgen2/entrypoint/watch.py

@@ -2,9 +2,6 @@
 from dflow import (
     Workflow,
 )
-from dpgen2.utils import (
-    workflow_config_from_dict,
-)
 from dpgen2.utils.dflow_query import (
     matched_step_key,
 )
@@ -14,6 +11,12 @@
 from typing import (
     Optional, Dict, Union, List,
 )
+from dpgen2.entrypoint.common import (
+    global_config_workflow,
+)
+from dpgen2.entrypoint.args import (
+    normalize as normalize_args,
+)
 
 default_watching_keys = [
     "prep-run-train",
@@ -55,7 +58,9 @@ def watch(
         download : Optional[bool] = False,
         prefix : Optional[str] = None,
 ):
-    workflow_config_from_dict(wf_config)
+    wf_config = normalize_args(wf_config)
+
+    global_config_workflow(wf_config)
 
     wf = Workflow(id=workflow_id)