From 9bac2b7aaebb0d29f26b557b406efd05057c5203 Mon Sep 17 00:00:00 2001
From: Adam Tyson <code@adamltyson.com>
Date: Fri, 13 Oct 2023 14:40:47 +0100
Subject: [PATCH] Limit n processes on Windows to 61

---
 brainglobe_utils/general/system.py | 59 +++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 17 deletions(-)

diff --git a/brainglobe_utils/general/system.py b/brainglobe_utils/general/system.py
index 17a3a0c..d137cb5 100644
--- a/brainglobe_utils/general/system.py
+++ b/brainglobe_utils/general/system.py
@@ -14,6 +14,10 @@
 
 from brainglobe_utils.general.string import get_text_lines
 
+# On Windows, max_workers must be less than or equal to 61
+# https://docs.python.org/3/library/concurrent.futures.html#processpoolexecutor
+MAX_PROCESSES_WINDOWS = 61
+
 
 def replace_extension(file, new_extension, check_leading_period=True):
     """
@@ -132,35 +136,29 @@ def get_num_processes(
     :return: Number of processes to
     """
     logging.debug("Determining the maximum number of CPU cores to use")
-    try:
-        os.environ["SLURM_JOB_ID"]
-        n_cpu_cores = (
-            slurmio.SlurmJobParameters().allocated_cores - min_free_cpu_cores
-        )
-    except KeyError:
-        n_cpu_cores = psutil.cpu_count() - min_free_cpu_cores
+
+    n_cpu_cores = get_cores_available()
+    n_cpu_cores = n_cpu_cores - min_free_cpu_cores
 
     logging.debug(f"Number of CPU cores available is: {n_cpu_cores}")
 
     if ram_needed_per_process is not None:
-        cores_w_sufficient_ram = how_many_cores_with_sufficient_ram(
+        n_processes = limit_cores_based_on_memory(
+            n_cpu_cores,
             ram_needed_per_process,
-            fraction_free_ram=fraction_free_ram,
-            max_ram_usage=max_ram_usage,
-        )
-        n_processes = min(n_cpu_cores, cores_w_sufficient_ram)
-        logging.debug(
-            f"Based on memory requirements, up to {cores_w_sufficient_ram} "
-            f"cores could be used. Therefore setting the number of "
-            f"processes to {n_processes}."
+            fraction_free_ram,
+            max_ram_usage,
         )
     else:
         n_processes = n_cpu_cores
 
+    if platform.system() == "Windows":
+        n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS)
+
     if n_max_processes is not None:
         if n_max_processes < n_processes:
             logging.debug(
-                f"Forcing the number of processes to {n_max_processes} based"
+                f"Limiting the number of processes to {n_max_processes} based"
                 f" on other considerations."
             )
         n_processes = min(n_processes, n_max_processes)
@@ -174,6 +172,33 @@ def get_num_processes(
     return int(n_processes)
 
 
+def get_cores_available():
+    try:
+        os.environ["SLURM_JOB_ID"]
+        n_cpu_cores = slurmio.SlurmJobParameters().allocated_cores
+    except KeyError:
+        n_cpu_cores = psutil.cpu_count()
+
+    return n_cpu_cores
+
+
+def limit_cores_based_on_memory(
+    n_cpu_cores, ram_needed_per_process, fraction_free_ram, max_ram_usage
+):
+    cores_w_sufficient_ram = how_many_cores_with_sufficient_ram(
+        ram_needed_per_process,
+        fraction_free_ram=fraction_free_ram,
+        max_ram_usage=max_ram_usage,
+    )
+    n_processes = min(n_cpu_cores, cores_w_sufficient_ram)
+    logging.debug(
+        f"Based on memory requirements, up to {cores_w_sufficient_ram} "
+        f"cores could be used. Therefore setting the number of "
+        f"processes to {n_processes}."
+    )
+    return n_processes
+
+
 def how_many_cores_with_sufficient_ram(
     ram_needed_per_cpu, fraction_free_ram=0.1, max_ram_usage=None
 ):