lingyuzeng/analysis_pdb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

assign cpu and gpu

Browse Source
This commit is contained in:
root
2024-01-05 14:23:02 +08:00
parent b44ea7119a
commit 4de8d3ee88
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
runner.py
View File

@@ -104,16 +104,20 @@ class SimulationRunner:
def run_simulation(self): def run_simulation(self):
start_time = time.time() start_time = time.time()
result = None result = None
# 保留1个核心不使用
num_cores = max(1, multiprocessing.cpu_count() - 1) # 避免核心数小于 1
os.environ['OMP_NUM_THREADS'] = str(num_cores)
try: try:
env_vars = { env_vars = {
"NAME": self.pdb_file.stem, "NAME": self.pdb_file.stem,
"NSTEPS": str(self.nsteps), "NSTEPS": str(self.nsteps),
"DT": str(self.dt), "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", ""), "PATH": os.environ.get("PATH", ""),
"LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""),
"HOME": os.environ["HOME"], "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"pdb_file: {self.pdb_file.name}")
logger.info(f"Executing script at: {self.bash_script}") 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 return
# 创建一个进程池，每个GPU运行两个任务 # 创建一个进程池，每个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): for i, pdb_file in enumerate(pdb_files):
gpu_id = i % num_gpus # 分配GPU 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, gpu_id))
@@ -186,7 +190,7 @@ def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_p
if __name__ == "__main__": if __name__ == "__main__":
NSTEPS = 50000000 # Example: 50000000 steps NSTEPS = 50000000 # Example: 50000000 steps
DT = 0.002 # Example: 2 fs time step 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脚本路径 # 传入自定义的bash脚本路径
CUSTOM_BASH_SCRIPT_PATH = Path('md_gromacs.sh') CUSTOM_BASH_SCRIPT_PATH = Path('md_gromacs.sh')
# 传入 GMXRC 文件的路径 # 传入 GMXRC 文件的路径