V1.2

apex/__init__.py

@@ -1,5 +1,5 @@
 import os
-__version__ = '1.1.1'
+__version__ = '1.2.0'
 LOCAL_PATH = os.getcwd()
 
 

apex/archive.py

@@ -0,0 +1,271 @@
+import glob
+import logging
+import json
+import os
+from typing import List
+from pathlib import Path
+from monty.json import MontyEncoder
+from monty.serialization import loadfn, dumpfn
+from apex.utils import (
+    judge_flow,
+    json2dict,
+    update_dict,
+    return_prop_list,
+    load_config_file
+)
+from apex.database.DatabaseFactory import DatabaseFactory
+from apex.config import Config
+
+
+class ResultStorage:
+    def __init__(self, work_dir):
+        self.work_dir = Path(work_dir).absolute()
+        self.result_dict = {'work_path': str(self.work_dir)}
+
+    @property
+    def result_data(self):
+        return self.result_dict
+
+    @property
+    def work_dir_path(self):
+        return str(self.work_dir)
+
+    @json2dict
+    def sync_relax(self, relax_param: dict):
+        # sync results from relaxation task
+        confs = relax_param["structures"]
+        interaction = relax_param["interaction"]
+        conf_dirs = []
+        for conf in confs:
+            conf_dirs.extend(glob.glob(str(self.work_dir / conf)))
+        conf_dirs.sort()
+        for ii in conf_dirs:
+            relax_task = os.path.join(ii, 'relaxation/relax_task')
+            inter = os.path.join(relax_task, 'inter.json')
+            task = os.path.join(relax_task, 'task.json')
+            structure = os.path.join(relax_task, 'structure.json')
+            result = os.path.join(relax_task, 'result.json')
+            if not (
+                os.path.isfile(inter)
+                and os.path.isfile(task)
+                and os.path.isfile(result)
+            ):
+                logging.warning(
+                    f"relaxation result path is not complete, will skip result extraction from {relax_task}"
+                )
+                continue
+            logging.info(msg=f"extract results from {relax_task}")
+            conf_key = os.path.relpath(ii, self.work_dir)
+            conf_dict = {"interaction": loadfn(inter),
+                         "parameter": loadfn(task),
+                         "structure_info": loadfn(structure),
+                         "result": loadfn(result)}
+            new_dict = {conf_key: {"relaxation": conf_dict}}
+            update_dict(self.result_dict, new_dict)
+
+    @json2dict
+    def sync_props(self, props_param: dict, archive_tasks: bool = False):
+        # sync results from property test
+        confs = props_param["structures"]
+        interaction = props_param["interaction"]
+        properties = props_param["properties"]
+        prop_list = return_prop_list(properties)
+        conf_dirs = []
+        for conf in confs:
+            conf_dirs.extend(glob.glob(str(self.work_dir / conf)))
+        conf_dirs.sort()
+        for ii in conf_dirs:
+            for jj in prop_list:
+                prop_dir = os.path.join(ii, jj)
+                result = os.path.join(prop_dir, 'result.json')
+                param = os.path.join(prop_dir, 'param.json')
+                task_list_path = os.path.join(prop_dir, 'task_list.json')
+                if not os.path.isfile(result):
+                    logging.warning(
+                        f"Property post-process is not complete, will skip result extraction from {prop_dir}"
+                    )
+                    continue
+                logging.info(msg=f"extract results from {prop_dir}")
+                conf_key = os.path.relpath(ii, self.work_dir)
+                result_dict = loadfn(result)
+                try:
+                    param_dict = loadfn(param)
+                except FileNotFoundError:
+                    logging.warning(f'{param} file not found')
+                    param_dict = None
+                prop_dict = {"parameter": param_dict, "result": result_dict}
+                # extract running details of each task
+                if archive_tasks:
+                    logging.debug(msg='Archive running details of tasks...')
+                    logging.warning(
+                        msg='You are trying to archive detailed running log of each task into database,'
+                            'which may exceed the limitation of database allowance.'
+                            'Please consider split the data or only archive details of the most important property.'
+                    )
+                    try:
+                        task_list = loadfn(task_list_path)
+                        result_task_path = [os.path.join(prop_dir, task, 'result_task.json') for task in task_list]
+                    except FileNotFoundError:
+                        logging.warning(f'{task_list_path} file not found, will track all tasks listed {prop_dir}')
+                        result_task_path = glob.glob(os.path.join(prop_dir, 'task.*', 'result_task.json'))
+                    result_task_path.sort()
+                    task_result_list = [loadfn(kk) for kk in result_task_path]
+                    prop_dict["tasks"] = task_result_list
+
+                new_dict = {conf_key: {jj: prop_dict}}
+                update_dict(self.result_dict, new_dict)
+
+
+def connect_database(config):
+    # connect to database
+    if config.database_type == 'mongodb':
+        logging.info(msg='Use database type: MongoDB')
+        database = DatabaseFactory.create_database(
+            'mongodb',
+            'mongodb',
+            config.mongodb_database,
+            config.mongodb_collection,
+            **config.mongodb_config_dict
+        )
+    elif config.database_type == 'dynamodb':
+        logging.info(msg='Use database type: DynamoDB')
+        database = DatabaseFactory.create_database(
+            'dynamodb',
+            'dynamodb',
+            config.dynamodb_table_name,
+            **config.dynamodb_config_dict
+        )
+    else:
+        raise RuntimeError(f'unsupported database type: {config.database_type}')
+    return database
+
+
+def archive2db(config, data: dict, data_id: str):
+    database = connect_database(config)
+    # archive results database
+    if config.archive_method == 'sync':
+        logging.debug(msg='Archive method: sync')
+        database.sync(data, data_id, depth=2)
+    elif config.archive_method == 'record':
+        logging.debug(msg='Archive method: record')
+        database.record(data, data_id)
+    else:
+        raise TypeError(
+            f'Unrecognized archive method: {config.archive_method}. '
+            f"Should select from 'sync' and 'record'."
+        )
+
+
+def archive_workdir(relax_param, props_param, config, work_dir, flow_type):
+    print(f'=> Begin archiving {work_dir}')
+    # extract results json
+    store = ResultStorage(work_dir)
+    if relax_param and flow_type != 'props':
+        store.sync_relax(relax_param)
+    if props_param and flow_type != 'relax':
+        store.sync_props(props_param, config.archive_tasks)
+
+    # define archive key
+    if config.archive_key:
+        data_id = config.archive_key
+    else:
+        data_id = str(store.work_dir_path)
+
+    dump_file = os.path.join(store.work_dir_path, 'all_result.json')
+    if os.path.isfile(dump_file):
+        logging.info(msg='all_result.json exists, and will be updated.')
+        orig_data = loadfn(dump_file)
+        update_dict(orig_data, store.result_data, depth=2)
+        dumpfn(orig_data, dump_file, indent=4)
+    else:
+        dumpfn(store.result_data, dump_file, indent=4)
+
+    if config.database_type != 'local':
+        data_json_str = json.dumps(store.result_data, cls=MontyEncoder, indent=4)
+        data_dict = json.loads(data_json_str)
+        data_dict['_id'] = data_id
+
+        archive2db(config, data_dict, data_id)
+
+
+def archive2db_from_json(config, json_file):
+    logging.info(msg=f'Archive from local json file: {json_file}')
+    data_dict = loadfn(json_file)
+    data_json_str = json.dumps(data_dict, cls=MontyEncoder, indent=4)
+    data_dict = json.loads(data_json_str)
+    # define archive key
+    if config.archive_key:
+        data_id = config.archive_key
+    else:
+        data_id = str(data_dict["work_path"])
+    data_dict['_id'] = data_id
+
+    archive2db(config, data_dict, data_id)
+
+
+def archive_result(
+        parameters: List[os.PathLike],
+        config_dict: dict,
+        work_dir: List[os.PathLike],
+        indicated_flow_type: str,
+        database_type,
+        method,
+        archive_tasks,
+        is_result
+):
+    global_config = Config(**config_dict)
+    # re-specify args
+    if database_type:
+        global_config.database_type = database_type
+    if method:
+        global_config.archive_method = method
+    if archive_tasks:
+        global_config.archive_tasks = archive_tasks
+
+    if is_result:
+        # archive local results json file
+        json_file_list = []
+        # Parameter here stands for json files that store test results and be archived directly
+        for ii in parameters:
+            glob_list = glob.glob(os.path.abspath(ii))
+            json_file_list.extend(glob_list)
+            json_file_list.sort()
+        for ii in json_file_list:
+            archive2db_from_json(global_config, ii)
+    else:
+        _, _, flow_type, relax_param, props_param = judge_flow(
+            [loadfn(jj) for jj in parameters],
+            indicated_flow_type
+        )
+        # archive work directories
+        work_dir_list = []
+        for ii in work_dir:
+            glob_list = glob.glob(os.path.abspath(ii))
+            work_dir_list.extend(glob_list)
+            work_dir_list.sort()
+        for ii in work_dir_list:
+            archive_workdir(relax_param, props_param, global_config, ii, flow_type)
+
+
+def archive_from_args(
+        parameters: List[os.PathLike],
+        config_file: os.PathLike,
+        work_dirs: List[os.PathLike],
+        indicated_flow_type: str,
+        database_type,
+        method,
+        archive_tasks,
+        is_result
+):
+    print('-------Archive result Mode-------')
+    archive_result(
+        parameters=parameters,
+        config_dict=load_config_file(config_file),
+        work_dir=work_dirs,
+        indicated_flow_type=indicated_flow_type,
+        database_type=database_type,
+        method=method,
+        archive_tasks=archive_tasks,
+        is_result=is_result
+    )
+    print('Complete!')

apex/config.py

@@ -18,7 +18,7 @@ class Config:
     mode: str = 'default'
     debug_copy_method: str = "copy"
     debug_pool_workers: int = 1
-    debug_workdir: str = '.'
+    debug_workdir: str = None
 
     # dflow s3 config
     dflow_s3_config: dict = None
@@ -57,18 +57,30 @@ class Config:
     group_size: int = None
     pool_size: int = None
     upload_python_packages: list = field(default_factory=list)
-    # lammps
     lammps_image_name: str = None
     lammps_run_command: str = None
-    # vasp
     vasp_image_name: str = None
     vasp_run_command: str = None
-    # abacus
     abacus_image_name: str = None
     abacus_run_command: str = None
-
     is_bohrium_dflow: bool = False
 
+    database_type: str = 'local'
+    archive_method: str = 'sync'
+    archive_key: str = None
+    archive_tasks: bool = False
+
+    # MongoDB config
+    mongodb_config: dict = None
+    mongodb_host: str = "localhost"
+    mongodb_port: int = 27017
+    mongodb_database: str = "apex_results"
+    mongodb_collection: str = "default_collection"
+
+    # DynamoDB config
+    dynamodb_config: dict = None
+    dynamodb_table_name: str = "apex_results"
+
     def __post_init__(self):
         # judge if running dflow on the Bohrium
         try:
@@ -83,7 +95,7 @@ def __post_init__(self):
         if not (self.context_type or self.machine):
             self.machine_dict = None
         elif self.context_type in ["Bohrium", "bohrium",
-                                    "BohriumContext", "boriumcontext"]:
+                                   "BohriumContext", "boriumcontext"]:
             self.machine_dict = {
                 "batch_type": self.batch_type,
                 "context_type": self.context_type,
@@ -101,7 +113,7 @@ def __post_init__(self):
             if self.machine:
                 update_dict(self.machine_dict, self.machine)
         elif self.context_type in ["SSHContext", "sshcontext",
-                                    "SSH", "ssh"]:
+                                   "SSH", "ssh"]:
             self.machine_dict = {
                 "batch_type": self.batch_type,
                 "context_type": self.context_type,
@@ -119,7 +131,7 @@ def __post_init__(self):
             if self.machine:
                 update_dict(self.machine_dict, self.machine)
         elif self.context_type in ["LocalContext", "localcontext"
-                                    "Local", "local"]:
+                                                   "Local", "local"]:
             self.machine_dict = {
                 "batch_type": self.batch_type,
                 "context_type": self.context_type,
@@ -130,7 +142,7 @@ def __post_init__(self):
             if self.machine:
                 update_dict(self.machine_dict, self.machine)
         elif self.context_type in ["LazyLocalContext", "lazylocalcontext"
-                                    "LazyLocal", "lazylocal"]:
+                                                       "LazyLocal", "lazylocal"]:
             self.machine_dict = {
                 "batch_type": self.batch_type,
                 "context_type": self.context_type,
@@ -204,6 +216,23 @@ def dispatcher_config_dict(self):
             update_dict(dispatcher_config, self.dispatcher_config)
         return dispatcher_config
 
+    @property
+    def mongodb_config_dict(self):
+        mongodb_config = {
+            "host": self.mongodb_host,
+            "port": self.mongodb_port
+        }
+        if self.mongodb_config:
+            update_dict(mongodb_config, self.mongodb_config)
+        return mongodb_config
+
+    @property
+    def dynamodb_config_dict(self):
+        dynamodb_config = {}
+        if self.mongodb_config:
+            update_dict(dynamodb_config, self.dynamodb_config)
+        return dynamodb_config
+
     @property
     def basic_config_dict(self):
         basic_config = {