From bd73f7c13791797abb3650bdb115743529d3f722 Mon Sep 17 00:00:00 2001 From: Lehman Garrison Date: Fri, 26 Jul 2024 10:02:58 -0400 Subject: [PATCH 1/2] Add prepare_sim.Nthread_per_load parameter --- abacusnbody/hod/prepare_sim.py | 13 ++++++++++--- scripts/hod/config/abacus_hod.yaml | 3 ++- scripts/hod/config/lc_hod.yaml | 1 + 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/abacusnbody/hod/prepare_sim.py b/abacusnbody/hod/prepare_sim.py index 3ae7fe02..984fea20 100644 --- a/abacusnbody/hod/prepare_sim.py +++ b/abacusnbody/hod/prepare_sim.py @@ -1094,9 +1094,16 @@ def main( else: shearmark = None # N_dim = config['HOD_params']['Ndim'] - nthread = int( - np.floor(multiprocessing.cpu_count() / config['prepare_sim']['Nparallel_load']) - ) + nthread = config['prepare_sim'].get('Nthread_per_load', 'auto') + if nthread == 'auto': + nthread = int( + np.floor( + len(os.sched_getaffinity(0)) / config['prepare_sim']['Nparallel_load'] + ) + ) + print(f'prepare_sim inferred Nthread_per_load = {nthread}') + else: + nthread = int(nthread) p = multiprocessing.Pool(config['prepare_sim']['Nparallel_load']) p.starmap( diff --git a/scripts/hod/config/abacus_hod.yaml b/scripts/hod/config/abacus_hod.yaml index 68eb9442..d94f1027 100644 --- a/scripts/hod/config/abacus_hod.yaml +++ b/scripts/hod/config/abacus_hod.yaml @@ -13,7 +13,8 @@ sim_params: cleaned_halos: True # load cleaned halos? prepare_sim: - Nparallel_load: 5 # number of thread for organizing simulation outputs (prepare_sim) + Nparallel_load: 5 # number of processes. peak memory usage will increase by this factor. + Nthread_per_load: 'auto' # number of threads per process (auto uses the affinity mask) # HOD parameters HOD_params: diff --git a/scripts/hod/config/lc_hod.yaml b/scripts/hod/config/lc_hod.yaml index d20487ba..7c88dc49 100644 --- a/scripts/hod/config/lc_hod.yaml +++ b/scripts/hod/config/lc_hod.yaml @@ -14,6 +14,7 @@ sim_params: prepare_sim: Nparallel_load: 1 # not sure if this makes a difference since we have a single slab + Nthread_per_load: 'auto' # HOD parameters HOD_params: From 4723bccd45e0d55fc9f6f4291a2a70e537b41076 Mon Sep 17 00:00:00 2001 From: Lehman Garrison Date: Fri, 26 Jul 2024 10:06:10 -0400 Subject: [PATCH 2/2] use integer division --- abacusnbody/hod/prepare_sim.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/abacusnbody/hod/prepare_sim.py b/abacusnbody/hod/prepare_sim.py index 984fea20..51171338 100644 --- a/abacusnbody/hod/prepare_sim.py +++ b/abacusnbody/hod/prepare_sim.py @@ -1096,10 +1096,8 @@ def main( # N_dim = config['HOD_params']['Ndim'] nthread = config['prepare_sim'].get('Nthread_per_load', 'auto') if nthread == 'auto': - nthread = int( - np.floor( - len(os.sched_getaffinity(0)) / config['prepare_sim']['Nparallel_load'] - ) + nthread = ( + len(os.sched_getaffinity(0)) // config['prepare_sim']['Nparallel_load'] ) print(f'prepare_sim inferred Nthread_per_load = {nthread}') else: