From 4de8d3ee8846539b490c8f142c910f3351eb226e Mon Sep 17 00:00:00 2001 From: root Date: Fri, 5 Jan 2024 14:23:02 +0800 Subject: [PATCH] assign cpu and gpu --- runner.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/runner.py b/runner.py index 4f2ac8c..c425134 100755 --- a/runner.py +++ b/runner.py @@ -104,16 +104,20 @@ class SimulationRunner: def run_simulation(self): start_time = time.time() result = None + # 保留1个核心不使用 + num_cores = max(1, multiprocessing.cpu_count() - 1) # 避免核心数小于 1 + os.environ['OMP_NUM_THREADS'] = str(num_cores) try: env_vars = { "NAME": self.pdb_file.stem, "NSTEPS": str(self.nsteps), "DT": str(self.dt), - "GMXRC_PATH": str(self.gmxrc_path), + "GMXRC_PATH": str(self.gmxrc_path) if self.gmxrc_path else "", "PATH": os.environ.get("PATH", ""), "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), "HOME": os.environ["HOME"], - "CUDA_VISIBLE_DEVICES": str(self.gpu_id) if self.gpu_id is not None else "" + "CUDA_VISIBLE_DEVICES": str(self.gpu_id) if self.gpu_id is not None else "", + 'OMP_NUM_THREADS': str(num_cores) } logger.info(f"pdb_file: {self.pdb_file.name}") logger.info(f"Executing script at: {self.bash_script}") @@ -175,7 +179,7 @@ def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_p return # 创建一个进程池，每个GPU运行两个任务 - with multiprocessing.Pool(processes=num_gpus * 2) as pool: + with multiprocessing.Pool(processes=num_gpus * 2) 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)) @@ -186,7 +190,7 @@ def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_p if __name__ == "__main__": NSTEPS = 50000000 # Example: 50000000 steps DT = 0.002 # Example: 2 fs time step - PDB_FOLDER_PATH = Path("./pdb_test") # Assuming the PDB files are in a folder named 'pdb_files' in the current directory + PDB_FOLDER_PATH = Path("./pdb_testb") # Assuming the PDB files are in a folder named 'pdb_files' in the current directory # 传入自定义的bash脚本路径 CUSTOM_BASH_SCRIPT_PATH = Path('md_gromacs.sh') # 传入 GMXRC 文件的路径