assign cpu and gpu

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 文件的路径