fix for nomp. only use if you are anyway in a child process

This commit is contained in:
Jan Kieseler
2024-09-13 10:06:28 +02:00
parent e3c048f2ed
commit 854402320c
+19 -5
View File
@@ -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)