|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import logging
|
| import os
|
| from pathlib import Path
|
| from typing import Any
|
|
|
| import torch
|
| import torch.distributed as dist
|
| from natsort import natsorted
|
| from protenix.config import parse_configs, parse_sys_args
|
| from protenix.config.extend_types import RequiredValue
|
| from protenix.utils.distributed import DIST_WRAPPER
|
|
|
| from pxdbench.pxd_configs.eval import eval_configs
|
| from pxdbench.run import find_files_with_ext, run_task
|
| from pxdbench.utils import convert_cifs_to_pdbs
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
|
|
| def scan_tasks(data_dir, is_mmcif=False):
|
| data_dir = Path(data_dir).resolve()
|
| paths = []
|
|
|
| pattern = "*.cif" if is_mmcif else "*.pdb"
|
| for pdb_file in data_dir.rglob(pattern):
|
| paths.append(os.path.dirname(pdb_file))
|
|
|
| return sorted(list(set(paths)))
|
|
|
|
|
| class EvalRunner(object):
|
| def __init__(self, configs: Any) -> None:
|
| self.configs = configs
|
| self.root_dir = self.configs.data_dir
|
| self.dump_dir = self.configs.dump_dir
|
| self.init_env()
|
|
|
| def init_env(self) -> None:
|
| self.print(
|
| f"Distributed environment: world size: {DIST_WRAPPER.world_size}, "
|
| + f"global rank: {DIST_WRAPPER.rank}, local rank: {DIST_WRAPPER.local_rank}"
|
| )
|
| self.use_cuda = torch.cuda.device_count() > 0
|
| if self.use_cuda:
|
| self.device = torch.device("cuda:{}".format(DIST_WRAPPER.local_rank))
|
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
| all_gpu_ids = ",".join(str(x) for x in range(torch.cuda.device_count()))
|
| devices = os.getenv("CUDA_VISIBLE_DEVICES", all_gpu_ids)
|
| logging.info(
|
| f"LOCAL_RANK: {DIST_WRAPPER.local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]"
|
| )
|
| torch.cuda.set_device(self.device)
|
| else:
|
| self.device = torch.device("cpu")
|
| if DIST_WRAPPER.world_size > 1:
|
| dist.init_process_group(backend="nccl")
|
| logging.info("Finished init ENV.")
|
|
|
| def print(self, msg: str):
|
| if DIST_WRAPPER.rank == 0:
|
| logger.info(msg)
|
|
|
| def run(self):
|
| input_dirs = scan_tasks(self.root_dir, self.configs.is_mmcif)
|
| valid_input_dirs = []
|
| for data_dir in input_dirs:
|
| exp_name = os.path.relpath(data_dir, self.root_dir)
|
| if os.path.exists(
|
| os.path.join(self.dump_dir, exp_name, "summary_output.json")
|
| ):
|
| self.print(f"Found summary file for {exp_name} - Skip!")
|
| continue
|
| valid_input_dirs.append(data_dir)
|
|
|
| valid_input_dirs = natsorted(valid_input_dirs)
|
| logging.info(f"There are {len(valid_input_dirs)} tasks to evaluate")
|
| print(valid_input_dirs)
|
|
|
| sub_input_dirs = valid_input_dirs[DIST_WRAPPER.rank :: DIST_WRAPPER.world_size]
|
| for i, data_dir in enumerate(sub_input_dirs):
|
| logging.info(
|
| f"Begin to evaluate [{i + 1}/{len(sub_input_dirs)}]: {data_dir}"
|
| )
|
| exp_name = os.path.relpath(data_dir, self.root_dir)
|
| dump_dir = os.path.join(self.dump_dir, exp_name)
|
| os.makedirs(dump_dir, exist_ok=True)
|
| if self.configs.is_mmcif:
|
| pdb_dir, pdb_names, _, _ = convert_cifs_to_pdbs(
|
| data_dir,
|
| out_pdb_dir=os.path.join(data_dir, "converted_pdbs"),
|
| )
|
| else:
|
| pdb_dir = data_dir
|
| pdb_names = find_files_with_ext(data_dir, "pdb")
|
| logging.info(f"There are {len(pdb_names)} pdbs in this task.")
|
|
|
| input_data = {
|
| "task": "monomer",
|
| "name": exp_name,
|
| "pdb_dir": pdb_dir,
|
| "pdb_names": pdb_names,
|
| "out_dir": dump_dir,
|
| }
|
| run_task(
|
| input_data,
|
| self.configs,
|
| device_id=DIST_WRAPPER.local_rank,
|
| seed=self.configs.seed,
|
| )
|
| logging.info("Eval done!")
|
|
|
|
|
| def main():
|
| config_dict = {
|
| "data_dir": RequiredValue(str),
|
| "dump_dir": RequiredValue(str),
|
| "is_mmcif": False,
|
| "seed": 2025,
|
| **eval_configs,
|
| }
|
| configs = parse_configs(config_dict, arg_str=parse_sys_args())
|
| runner = EvalRunner(configs)
|
| runner.run()
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
