From 854402320c0e71ba2bde0a6178b96b0273d1d0d6 Mon Sep 17 00:00:00 2001
From: Jan Kieseler <jkiesele@cern.ch>
Date: Fri, 13 Sep 2024 10:06:28 +0200
Subject: [PATCH] fix for nomp. only use if you are anyway in a child process

---
 bind/G4Calo.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/bind/G4Calo.py b/bind/G4Calo.py
index d8cd757..a54ce6a 100644
--- a/bind/G4Calo.py
+++ b/bind/G4Calo.py
@@ -163,6 +163,12 @@ def _run_mini_batch_silent(*args, **kwargs):
     #I tried
     return _run_mini_batch(*args, **kwargs)
 
+_g4calo_threading_lock = None
+
+def init_g4calo_threading_lock():
+    global _g4calo_threading_lock
+    _g4calo_threading_lock = multiprocessing.Lock()
+
 
 def run_batch(
         gd : GeometryDescriptor,
@@ -171,7 +177,7 @@ def run_batch(
         minEnergy_GeV: float,
         maxEnergy_GeV: float = -1.0,
         filename: str = "",
-        threading_lock = None):
+        no_mp: bool = False):
     '''
     splits the batch in jobs depending on how many cores are available and runs mini batches in parallel
     '''
@@ -181,6 +187,9 @@ def run_batch(
     #make sure to adjust cores such that at least 80 events are run per core
     nCores = min(nCores, nEvents // 80 + 1)
 
+    if no_mp:
+        nCores = 1
+
     print(f"Running on {nCores} cores")
     nEventsPerCore = nEvents // nCores
     print(f"Running {nEventsPerCore} events per core")
@@ -189,7 +198,7 @@ def run_batch(
     
     nevents = [nEventsPerCore if i < nCores - 1 else nEventsLastCore for i in range(nCores)]
 
-    with threading_lock if threading_lock is not None else contextlib.nullcontext():
+    with _g4calo_threading_lock if _g4calo_threading_lock is not None else contextlib.nullcontext():
         # check if file exists, if so, read last seed. If not create it
         if os.path.exists(os.path.expanduser("~/.g4calo_seeds.txt")):
             with open(os.path.expanduser("~/.g4calo_seeds.txt"), "r") as f:
@@ -201,9 +210,14 @@ def run_batch(
 
     seed += 1
     print(f"Batch seed: {seed}")
-    #use a multiprocessing pool to run the mini batches in parallel
-    with multiprocessing.Pool(nCores) as pool:
-        dfs = pool.starmap(_run_mini_batch_silent, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i) for i in range(nCores)])
+    
+    if no_mp:
+        dfs = []
+        for i in range(nCores):
+            dfs.append(_run_mini_batch_silent(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i))
+    else:
+        with multiprocessing.Pool(nCores) as pool:
+            dfs = pool.starmap(_run_mini_batch_silent, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, seed, i) for i in range(nCores)])
 
     alldf = pd.concat(dfs)
     alldf.reset_index(drop=True, inplace=True)