From 9bac2b7aaebb0d29f26b557b406efd05057c5203 Mon Sep 17 00:00:00 2001 From: Adam Tyson Date: Fri, 13 Oct 2023 14:40:47 +0100 Subject: [PATCH 1/4] 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 ): From cc7bd1b270429fb8fb90d9bf02ef57ce094a394c Mon Sep 17 00:00:00 2001 From: Adam Tyson Date: Fri, 13 Oct 2023 15:02:44 +0100 Subject: [PATCH 2/4] Test windows cpu count --- brainglobe_utils/general/system.py | 9 ++++-- .../test_unit/test_general/test_system.py | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/brainglobe_utils/general/system.py b/brainglobe_utils/general/system.py index d137cb5..5d0b226 100644 --- a/brainglobe_utils/general/system.py +++ b/brainglobe_utils/general/system.py @@ -152,8 +152,7 @@ def get_num_processes( else: n_processes = n_cpu_cores - if platform.system() == "Windows": - n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS) + n_max_processes = limit_cpus_windows(n_max_processes) if n_max_processes is not None: if n_max_processes < n_processes: @@ -172,6 +171,12 @@ def get_num_processes( return int(n_processes) +def limit_cpus_windows(n_max_processes): + if platform.system() == "Windows": + n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS) + return n_max_processes + + def get_cores_available(): try: os.environ["SLURM_JOB_ID"] diff --git a/tests/tests/test_unit/test_general/test_system.py b/tests/tests/test_unit/test_general/test_system.py index f50d831..4d7bf23 100644 --- a/tests/tests/test_unit/test_general/test_system.py +++ b/tests/tests/test_unit/test_general/test_system.py @@ -2,6 +2,7 @@ import random from pathlib import Path from random import shuffle +from unittest.mock import patch import pytest @@ -90,6 +91,7 @@ def test_get_num_processes(): assert os.cpu_count() == system.get_num_processes(min_free_cpu_cores=0) +## orig def test_max_processes(): max_proc = 5 correct_n = min(os.cpu_count(), max_proc) @@ -98,6 +100,33 @@ def test_max_processes(): ) +def test_max_processes_windows_low(): + cpu_count = 10 + with patch( + "brainglobe_utils.general.system.platform.system", + return_value="Windows", + ): + with patch( + "brainglobe_utils.general.system.psutil.cpu_count", + return_value=cpu_count, + ): + assert system.limit_cpus_windows(cpu_count) == cpu_count + + +def test_max_processes_windows_high(): + cpu_count = 128 + with patch( + "brainglobe_utils.general.system.platform.system", + return_value="Windows", + ): + with patch( + "brainglobe_utils.general.system.psutil.cpu_count", + return_value=cpu_count, + ): + # 61 is max on Windows + assert system.limit_cpus_windows(cpu_count) == 61 + + class Paths: def __init__(self, directory): self.one = directory / "one.aaa" From 5e5ceac8c59cb5401b2a6352b02f0d6792d6d1e0 Mon Sep 17 00:00:00 2001 From: Adam Tyson Date: Fri, 13 Oct 2023 15:10:31 +0100 Subject: [PATCH 3/4] update test --- tests/tests/test_unit/test_general/test_system.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/tests/test_unit/test_general/test_system.py b/tests/tests/test_unit/test_general/test_system.py index 4d7bf23..20bbd47 100644 --- a/tests/tests/test_unit/test_general/test_system.py +++ b/tests/tests/test_unit/test_general/test_system.py @@ -88,10 +88,14 @@ def test_check_path_in_dir(): def test_get_num_processes(): - assert os.cpu_count() == system.get_num_processes(min_free_cpu_cores=0) + cpu_count = 10 + with patch( + "brainglobe_utils.general.system.psutil.cpu_count", + return_value=cpu_count, + ): + assert system.get_num_processes(min_free_cpu_cores=0) == cpu_count -## orig def test_max_processes(): max_proc = 5 correct_n = min(os.cpu_count(), max_proc) From 4154739e6183a36fc108ce74b903d720526b1bd2 Mon Sep 17 00:00:00 2001 From: Adam Tyson Date: Fri, 13 Oct 2023 15:16:30 +0100 Subject: [PATCH 4/4] Deal with no max processes set when setting max on windows --- brainglobe_utils/general/system.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/brainglobe_utils/general/system.py b/brainglobe_utils/general/system.py index 5d0b226..f6ee71b 100644 --- a/brainglobe_utils/general/system.py +++ b/brainglobe_utils/general/system.py @@ -173,7 +173,8 @@ def get_num_processes( def limit_cpus_windows(n_max_processes): if platform.system() == "Windows": - n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS) + if n_max_processes is not None: + n_max_processes = min(n_max_processes, MAX_PROCESSES_WINDOWS) return n_max_processes