improve detection of cpu count on Windows system

NRCan · Oct 22, 2024 · e86cd1c · e86cd1c
1 parent f9acbc5
commit e86cd1c
Showing 1 changed file with 6 additions and 4 deletions.
diff --git 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))