Max 500 simulations simultaneously on HPC

Author: KriFos1

ensemble/ensemble.py

@@ -569,21 +569,30 @@ def calc_prediction(self, input_state=None, save_prediction=None):
                 en_pred = [self.sim.run_fwd_sim(state, member_index) for state, member_index in
                            tqdm(zip(list_state, list_member_index), total=len(list_state))]
             elif self.sim.input_dict.get('hpc', False): # Run prediction in parallel on hpc
-                _ = [self.sim.run_fwd_sim(state, member_index, nosim=True) for state, member_index in
-                           zip(list_state, list_member_index)]
-                # Run call_sim on the hpc
-                job_id=self.sim.SLURM_HPC_run(self.ne, filename=self.sim.input_dict['runfile'])
-                # Wait for the simulations to finish
-                if job_id:
-                    self.sim.wait_for_jobs(job_id)
-                else:
-                    print("Job submission failed. Exiting.")
-                # Extract the results
+                batch_size = 500 # If more than 500 ensemble members, we limit the runs to batches of 500
+                # Split the ensemble into batches of 500
                 en_pred = []
-                for member_i in list_member_index:
-                    self.sim.extract_data(member_i)
-                    en_pred.append(deepcopy(self.sim.pred_data))
-                    self.sim.remove_folder(member_i)
+                batch_en = [np.arange(start, start + batch_size) for start in
+                            np.arange(0, self.ne - batch_size, batch_size)]
+                if len(batch_en): # if self.ne is less than batch_size
+                    batch_en.append(np.arange(batch_en[-1][-1]+1, self.ne))
+                else:
+                    batch_en.append(np.arange(0, self.ne))
+                for n_e in batch_en:
+                    _ = [self.sim.run_fwd_sim(state, member_index, nosim=True) for state, member_index in
+                               zip(list_state[n_e], list_member_index[n_e])]
+                    # Run call_sim on the hpc
+                    job_id=self.sim.SLURM_HPC_run(n_e, venv=sys.executable, filename=self.sim.input_dict['runfile'])
+                    # Wait for the simulations to finish
+                    if job_id:
+                        self.sim.wait_for_jobs(job_id)
+                    else:
+                        print("Job submission failed. Exiting.")
+                    # Extract the results
+                    for member_i in list_member_index[n_e]:
+                        self.sim.extract_data(member_i)
+                        en_pred.append(deepcopy(self.sim.pred_data))
+                        self.sim.remove_folder(member_i)
 
             else: # Run prediction in parallel using p_map
                 en_pred = p_map(self.sim.run_fwd_sim, list_state,

simulator/opm.py

@@ -95,7 +95,7 @@ def check_sim_end(self, finished_member=None):
         return finished_member
 
     @staticmethod
-    def SLURM_HPC_run(num_runs, filename=None):
+    def SLURM_HPC_run(n_e, venv, filename=None):
         """
         HPC run manager for SLURM.
 
@@ -106,7 +106,7 @@ def SLURM_HPC_run(num_runs, filename=None):
         slurm_script = f"""#!/bin/bash                                                                                               
 #SBATCH --partition=comp                                                                                  
 #SBATCH --job-name=EnDA                                                                               
-#SBATCH --array=0-{num_runs - 1}                                                                            
+#SBATCH --array={n_e[0]}-{n_e[-1]}                                                                            
 #SBATCH --time=01:00:00                                                                                   
 #SBATCH --mem=4G                                                                                          
 #SBATCH --cpus-per-task=1                                                                                 
@@ -118,7 +118,7 @@ def SLURM_HPC_run(num_runs, filename=None):
 export LMOD_DISABLE_SAME_NAME_AUTOSWAP=no                                                                 
 module load opm-simulators                                                                                
 
-source ../../../code/venv/bin/activate                                                                    
+source {venv}                                                                    
 
 # Set folder based on SLURM_ARRAY_TASK_ID
 folder="En_${{SLURM_ARRAY_TASK_ID}}/"