singler process, singler gpu

This commit is contained in:
root
2024-01-05 15:11:27 +08:00
parent 4de8d3ee88
commit 27084ba28d

View File

@@ -18,6 +18,8 @@ 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')
@@ -116,8 +118,7 @@ class SimulationRunner:
"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 "",
'OMP_NUM_THREADS': str(num_cores)
"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}")
@@ -178,11 +179,18 @@ def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_p
logger.error("No GPUs detected, exiting.")
return
# 创建一个进程池每个GPU运行个任务
with multiprocessing.Pool(processes=num_gpus * 2) as pool: # num_gpus * 1
# 创建一个进程池每个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
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.apply_async(run_simulation_task, (pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path, '0'))
pool.close()
pool.join()
@@ -190,7 +198,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_testb") # Assuming the PDB files are in a folder named 'pdb_files' in the current directory
PDB_FOLDER_PATH = Path("./pdb_test") # 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 文件的路径