From bb44cc624de52c9edbabb530b0c3aa2f6b54784e Mon Sep 17 00:00:00 2001 From: zm711 <92116279+zm711@users.noreply.github.com> Date: Thu, 30 May 2024 15:52:19 +0100 Subject: [PATCH 1/3] limit to number of virtual cores for n_jobs --- src/spikeinterface/core/job_tools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 6499fb145f..dd65409b64 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -96,7 +96,8 @@ def fix_job_kwargs(runtime_job_kwargs): else: n_jobs = int(n_jobs) - job_kwargs["n_jobs"] = max(n_jobs, 1) + n_jobs = max(n_jobs, 1) + job_kwargs["n_jobs"] = min(n_jobs, os.cpu_count()) if "n_jobs" not in runtime_job_kwargs and job_kwargs["n_jobs"] == 1 and not is_set_global_job_kwargs_set(): warnings.warn( From 46b65345cb8363b887b2e13725341ed68430ad56 Mon Sep 17 00:00:00 2001 From: zm711 <92116279+zm711@users.noreply.github.com> Date: Fri, 31 May 2024 09:44:24 +0100 Subject: [PATCH 2/3] fix jobs for n_jobs limits --- src/spikeinterface/core/tests/test_globals.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/tests/test_globals.py b/src/spikeinterface/core/tests/test_globals.py index a45bb6f49c..668cdb980f 100644 --- a/src/spikeinterface/core/tests/test_globals.py +++ b/src/spikeinterface/core/tests/test_globals.py @@ -1,6 +1,7 @@ import pytest import warnings from pathlib import Path +from os import cpu_count from spikeinterface import ( set_global_dataset_folder, @@ -64,7 +65,7 @@ def test_global_job_kwargs(): n_jobs=2, chunk_duration="1s", progress_bar=True, mp_context=None, max_threads_per_process=1 ) # test that fix_job_kwargs grabs global kwargs - new_job_kwargs = dict(n_jobs=10) + new_job_kwargs = dict(n_jobs=cpu_count()) job_kwargs_split = fix_job_kwargs(new_job_kwargs) assert job_kwargs_split["n_jobs"] == new_job_kwargs["n_jobs"] assert job_kwargs_split["chunk_duration"] == job_kwargs["chunk_duration"] @@ -74,6 +75,10 @@ def test_global_job_kwargs(): job_kwargs_split = fix_job_kwargs(none_job_kwargs) assert job_kwargs_split["chunk_duration"] == job_kwargs["chunk_duration"] assert job_kwargs_split["progress_bar"] == job_kwargs["progress_bar"] + # test that n_jobs are clipped if using more than virtual cores + excessive_n_jobs = dict(n_jobs=cpu_count() + 2) + job_kwargs_split = fix_job_kwargs(excessive_n_jobs) + assert job_kwargs_split["n_jobs"] == cpu_count() reset_global_job_kwargs() From 80ed220a2dd99496652692a3c1b78217c4d78b10 Mon Sep 17 00:00:00 2001 From: zm711 <92116279+zm711@users.noreply.github.com> Date: Fri, 31 May 2024 09:52:29 +0100 Subject: [PATCH 3/3] switch test to check cores of runner --- src/spikeinterface/core/tests/test_job_tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/core/tests/test_job_tools.py b/src/spikeinterface/core/tests/test_job_tools.py index f23a56ac5d..2f3aff0023 100644 --- a/src/spikeinterface/core/tests/test_job_tools.py +++ b/src/spikeinterface/core/tests/test_job_tools.py @@ -188,9 +188,9 @@ def test_fix_job_kwargs(): assert fixed_job_kwargs["n_jobs"] == 1 # test float value > 1 is cast to correct int - job_kwargs = dict(n_jobs=4.0, progress_bar=False, chunk_duration="1s") + job_kwargs = dict(n_jobs=float(os.cpu_count()), progress_bar=False, chunk_duration="1s") fixed_job_kwargs = fix_job_kwargs(job_kwargs) - assert fixed_job_kwargs["n_jobs"] == 4 + assert fixed_job_kwargs["n_jobs"] == os.cpu_count() # test wrong keys with pytest.raises(AssertionError):