fix for nomp. only use if you are anyway in a child process
This commit is contained in:
+19
-5
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user