From e86cd1c9d5925f10a5e9e523e99b8c57aca9dbcb Mon Sep 17 00:00:00 2001 From: Jean-Francois Bourgon Date: Tue, 22 Oct 2024 11:17:42 +0000 Subject: [PATCH] improve detection of cpu count on Windows system --- geo_inference/geo_inference.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/geo_inference/geo_inference.py b/geo_inference/geo_inference.py index cd3b484..50946e9 100644 --- a/geo_inference/geo_inference.py +++ b/geo_inference/geo_inference.py @@ -183,11 +183,13 @@ async def async_run_inference(self, """ - # configuring dask - if 'linux' in platform.uname().system.lower(): - num_workers = len(os.sched_getaffinity(0)) - 1 if workers == 0 else workers + # configuring dask with proper number of workers, alternatively we could also use os.getenv('SLURM_CPUS_PER_TASK') + if workers != 0: + num_workers = workers + elif 'linux' in platform.uname().system.lower(): + num_workers = len(os.sched_getaffinity(0)) - 1 else: - num_workers = os.cpu_count() - 1 if workers == 0 else workers + num_workers = os.cpu_count() - 1 print(f"running dask with {num_workers} workers") config.set(scheduler='threads', num_workers=num_workers) config.set(pool=ThreadPool(num_workers))