加强 log

2024-01-11 11:04:38 +08:00
parent 660f741274
commit e23905ce84
2 changed files with 83 additions and 95 deletions
--- a/runner.py
+++ b/runner.py
@@ -18,12 +18,7 @@ import multiprocessing
 import shutil
 import os

-if Path('simulation_log.log').exists():
-    Path('simulation_log.log').unlink()
-# 设置日志记录
-logging.basicConfig(level=logging.INFO, filename='simulation_log.log', filemode='a',
-                    format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger()
+

@dataclass
 class SimulationRunner:
@@ -31,13 +26,14 @@ class SimulationRunner:
    nsteps: int
    dt: float
    base_folder: Path
-    bash_script: Path = None  # Bash脚本路径作为可选参数
-    gmxrc_path: Path = None  # GMXRC文件路径作为可选参数
+    bash_script: Path = None
+    gmxrc_path: Path = None
    gpu_id: int = None
-    temp_folder: Path = field(init=False)  # 用于存储临时数据和结果数据的路径(轨迹数据处理中间数据等)
+    temp_folder: Path = field(init=False)
    runner_folder: Path = field(init=False)
    tpr_file: Path = field(init=False)
    xtc_file: Path = field(init=False)
+    logger: logging.Logger = field(init=False)

    def __post_init__(self):
        # 初始化文件夹和文件路径
@@ -50,6 +46,18 @@ class SimulationRunner:
        self.bash_script = self.bash_script.absolute() if self.bash_script else Path(__file__).resolve().parent / "md_gromacs.sh"
        # 设置 GMXRC_PATH 环境变量
        self.gmxrc_path = self.gmxrc_path
+        pdb_id = self.pdb_file.stem
+        self.logger = self.setup_logging(pdb_id, self.base_folder)
+
+    @staticmethod
+    def setup_logging(pdb_id, log_folder):
+        log_file = log_folder / f"{pdb_id}_simulation_log.log"
+        if log_file.exists():
+            log_file.unlink()
+
+        logging.basicConfig(level=logging.INFO, filename=log_file, filemode='a',
+                            format='%(asctime)s - %(levelname)s - %(message)s')
+        return logging.getLogger(pdb_id)

    def copy_pdb(self):
        shutil.copy(self.pdb_file, self.runner_folder / self.pdb_file.name)
@@ -120,39 +128,29 @@ class SimulationRunner:
                "HOME": os.environ["HOME"],
                "CUDA_VISIBLE_DEVICES": str(self.gpu_id) if self.gpu_id is not None else ""
            }
-            logger.info(f"pdb_file: {self.pdb_file.name}")
-            logger.info(f"Executing script at: {self.bash_script}")
+            self.logger.info(f"pdb_file: {self.pdb_file.name}")
+            self.logger.info(f"Executing script at: {self.bash_script}")
            result = subprocess.run(["bash", str(self.bash_script)], env=env_vars, cwd=self.runner_folder,
                                    capture_output=True, text=True, check=True)
        except subprocess.CalledProcessError as e:
-            logger.error(f"Error in simulation for {self.pdb_file.name}: {e}")
+            self.logger.error(f"Error in simulation for {self.pdb_file.name}: {e}")
            if e.stdout:
-                logger.error(f"Standard Output:\n{e.stdout}")
+                self.logger.error(f"Standard Output:\n{e.stdout}")
            if e.stderr:
-                logger.error(f"Standard Error:\n{e.stderr}")
+                self.logger.error(f"Standard Error:\n{e.stderr}")

        end_time = time.time()
        duration = end_time - start_time

        if result:
-            logger.info(f"Simulation for {self.pdb_file.name} completed successfully in {duration:.2f} seconds.")
+            self.logger.info(f"Simulation for {self.pdb_file.name} completed successfully in {duration:.2f} seconds.")
            if result.stdout:
-                logger.info(f"Shell Script Output:\n{result.stdout}")
+                self.logger.info(f"Shell Script Output:\n{result.stdout}")
            if result.stderr:
-                logger.error(f"Shell Script Error Output:\n{result.stderr}")
+                self.logger.error(f"Shell Script Error Output:\n{result.stderr}")
        else:
-            logger.error(f"Simulation for {self.pdb_file.name} failed in {duration:.2f} seconds.")
+            self.logger.error(f"Simulation for {self.pdb_file.name} failed in {duration:.2f} seconds.")

-# def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_path):
-#     pdb_folder = Path(pdb_folder_path).resolve()
-#     for pdb_file in pdb_folder.glob("*.pdb"):
-#         runner = SimulationRunner(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path)
-#         runner.copy_pdb()
-#         logger.info(f"Running simulation for {pdb_file.name} in {runner.runner_folder}...")
-#         runner.run_simulation()
-#         logger.info(f"Finished simulation for {pdb_file.name}.")
-#         runner.process_trajectory(extract_interval=100)  # 例如，每100ps抽取一次轨迹
-#         logger.info(f"Finished processing trajectory for {pdb_file.name}. per 100ps")
 def detect_gpus():
    """检测系统上的GPU数量。"""
    try:
@@ -161,37 +159,39 @@ def detect_gpus():
    except subprocess.CalledProcessError:
        return 0

+def setup_global_logging(log_folder):
+    log_file = log_folder / "simulation_log.log"
+    if log_file.exists():
+        log_file.unlink()
+
+    logging.basicConfig(level=logging.INFO, filename=log_file, filemode='a',
+                        format='%(asctime)s - %(levelname)s - %(message)s')
+    return logging.getLogger()
+
+
 def run_simulation_task(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path, gpu_id):
-    runner = SimulationRunner(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path)
+    runner = SimulationRunner(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path, gpu_id)
    runner.set_gpu(gpu_id)  # 设置要使用的GPU
    runner.copy_pdb()
-    logger.info(f"Running simulation for {pdb_file.name} on GPU {gpu_id}...")
    runner.run_simulation()
    runner.process_trajectory(extract_interval=100)
-    logger.info(f"Finished processing trajectory for {pdb_file.name} on GPU {gpu_id}.")

 def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_path):
-    pdb_folder = Path(pdb_folder_path).resolve()
+    pdb_folder = Path(pdb_folder_path)
+    setup_global_logging(pdb_folder)  # 设置全局日志记录器
+
    pdb_files = list(pdb_folder.glob("*.pdb"))
    num_gpus = detect_gpus()
-    
+
    if num_gpus == 0:
-        logger.error("No GPUs detected, exiting.")
+        logging.error("No GPUs detected, exiting.")
        return

-    # 创建一个进程池，每个GPU运行1个任务
-    # with multiprocessing.Pool(processes=num_gpus * 1) as pool: # num_gpus * 1
-    #     for i, pdb_file in enumerate(pdb_files):
-    #         gpu_id = i % num_gpus  # 分配GPU
-    #         pool.apply_async(run_simulation_task, (pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path, gpu_id))
-        
-    #     pool.close()
-    #     pool.join()
-    with multiprocessing.Pool(processes=1) as pool: # num_gpus * 1
+    with multiprocessing.Pool(processes=1) as pool:
        for i, pdb_file in enumerate(pdb_files):
            gpu_id = i % num_gpus  # 分配GPU
            pool.apply_async(run_simulation_task, (pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path, '0'))
-        
+
        pool.close()
        pool.join()