Fix n_jobs=-1 oversubscribing cores under CPU affinity restrictions

.github/workflows/deepinterpolation.yml

@@ -45,8 +45,9 @@ jobs:
           # install deepinteprolation
           uv pip install --system tensorflow==2.8.4
           uv pip install --system deepinterpolation@git+https://github.com/AllenInstitute/deepinterpolation.git
-          uv pip install --system numpy==1.26.4
           uv pip install --system -e .[full] --group test-core
+          # spikeinterface requires numpy>=2.0.0; reinstall older numpy for deepinterpolation
+          uv pip install --system numpy==1.26.4
       - name: Pip list
         if: ${{ steps.modules-changed.outputs.DEEPINTERPOLATION_CHANGED == 'true' }}
         run: |

src/spikeinterface/comparison/tests/test_multisortingcomparison.py

@@ -6,6 +6,7 @@
 import time
 
 from spikeinterface.core import generate_sorting
+from spikeinterface.core.job_tools import get_usable_cpu_count
 from spikeinterface.extractors import NumpySorting
 from spikeinterface.comparison import compare_multiple_sorters, MultiSortingComparison
 
@@ -111,8 +112,8 @@ def test_parallel():
     elapsed_N_jobs = t_stop - t_start
     print(f"Elapsed N jobs: {elapsed_N_jobs}")
 
-    # there is no guarantee there are more than 1 CPU on GH actions. Let's comment it out
-    if not ON_GITHUB and os.cpu_count() > 1:
+    # there is no guarantee there are more than 1 usable CPU on GH actions. Let's comment it out
+    if not ON_GITHUB and get_usable_cpu_count() > 1:
         assert elapsed_N_jobs < elapsed_1_job
     # check if the results are the same
     for k, cmp in msc_1_job.comparisons.items():

src/spikeinterface/core/job_tools.py

@@ -2,20 +2,23 @@
 Some utils to handle parallel jobs on top of job and/or loky
 """
 
-import numpy as np
-import platform
+import multiprocessing
 import os
-import warnings
-
+import platform
 import sys
-from tqdm.auto import tqdm
-
-from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
-import multiprocessing
 import threading
+import warnings
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+
+import numpy as np
 from threadpoolctl import threadpool_limits
+from tqdm.auto import tqdm
 
-from spikeinterface.core.core_tools import convert_string_to_bytes, convert_bytes_to_str, convert_seconds_to_str
+from spikeinterface.core.core_tools import (
+    convert_bytes_to_str,
+    convert_seconds_to_str,
+    convert_string_to_bytes,
+)
 
 _shared_job_kwargs_doc = """**job_kwargs : keyword arguments for parallel processing:
     * chunk_duration or chunk_size or chunk_memory or total_memory
@@ -30,8 +33,9 @@
     * n_jobs : int | float
         Number of workers that will be requested during multiprocessing. Note that
         the OS determines how this is distributed, but for convenience one can use
-        * -1 the number of workers is the same as number of cores (from os.cpu_count())
-        * float between 0 and 1 uses fraction of total cores (from os.cpu_count())
+        * -1 the number of workers is the same as the number of cores available to this
+          process, respecting CPU affinity restrictions where possible
+        * float between 0 and 1 uses a fraction of that core count
     * progress_bar : bool
         If True, a progress bar is printed
     * mp_context : "fork" | "spawn" | None, default: None
@@ -67,13 +71,26 @@
 )
 
 
+def get_usable_cpu_count():
+    """
+    Number of CPUs usable by this process.
+    """
+    if hasattr(os, "process_cpu_count"):
+        # Python >= 3.13; respects both Windows job object and Linux cgroup/cpuset restrictions
+        return os.process_cpu_count()
+    if hasattr(os, "sched_getaffinity"):
+        # Respects Linux cgroup/cpuset restrictions
+        return len(os.sched_getaffinity(0))
+    return os.cpu_count()
+
+
 def get_best_job_kwargs():
     """
     Gives best possible job_kwargs for the platform.
     Currently this function  is from developer experience, but may be adapted in the future.
     """
 
-    n_cpu = os.cpu_count()
+    n_cpu = get_usable_cpu_count()
 
     if platform.system() == "Linux":
         pool_engine = "process"
@@ -112,7 +129,7 @@ def get_best_job_kwargs():
 
 
 def fix_job_kwargs(runtime_job_kwargs):
-    from .globals import get_global_job_kwargs, is_set_global_job_kwargs_set
+    from .globals import get_global_job_kwargs
 
     job_kwargs = get_global_job_kwargs()
 
@@ -139,22 +156,22 @@ def fix_job_kwargs(runtime_job_kwargs):
             del runtime_job_kwargs_exclude_none[job_key]
     job_kwargs.update(runtime_job_kwargs_exclude_none)
 
-    # if n_jobs is -1, set to os.cpu_count() (n_jobs is always in global job_kwargs)
+    # if n_jobs is -1, set to get_usable_cpu_count() (n_jobs is always in global job_kwargs)
     n_jobs = job_kwargs["n_jobs"]
     assert isinstance(n_jobs, (float, np.integer, int)) and n_jobs != 0, "n_jobs must be a non-zero int or float"
 
     # for a fraction we do fraction of total cores
     if isinstance(n_jobs, float) and 0 < n_jobs <= 1:
-        n_jobs = int(n_jobs * os.cpu_count())
+        n_jobs = int(n_jobs * get_usable_cpu_count())
     # for negative numbers we count down from total cores (with -1 being all)
     elif n_jobs < 0:
-        n_jobs = int(os.cpu_count() + 1 + n_jobs)
+        n_jobs = int(get_usable_cpu_count() + 1 + n_jobs)
     # otherwise we just take the value given
     else:
         n_jobs = int(n_jobs)
 
     n_jobs = max(n_jobs, 1)
-    job_kwargs["n_jobs"] = min(n_jobs, os.cpu_count())
+    job_kwargs["n_jobs"] = min(n_jobs, get_usable_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(
@@ -186,9 +203,7 @@ def split_job_kwargs(mixed_kwargs):
 
 
 def divide_segment_into_chunks(num_frames, chunk_size):
-    if chunk_size is None:
-        chunks = [(0, num_frames)]
-    elif chunk_size > num_frames:
+    if chunk_size is None or chunk_size > num_frames:
         chunks = [(0, num_frames)]
     else:
         n = num_frames // chunk_size
@@ -216,10 +231,8 @@ def divide_time_series_into_chunks(recording, chunk_size):
 
 def ensure_n_jobs(extractor, n_jobs=1):
     if n_jobs == -1:
-        n_jobs = os.cpu_count()
-    elif n_jobs == 0:
-        n_jobs = 1
-    elif n_jobs is None:
+        n_jobs = get_usable_cpu_count()
+    elif n_jobs == 0 or n_jobs is None:
         n_jobs = 1
 
     # ProcessPoolExecutor has a hard limit of 61 for Windows
@@ -474,10 +487,22 @@ def get_chunk_memory(self):
         return self.chunk_size * self.time_series.get_sample_size_in_bytes()
 
     def ensure_chunk_size(
-        self, total_memory=None, chunk_size=None, chunk_memory=None, chunk_duration=None, n_jobs=1, **other_kwargs
+        self,
+        total_memory=None,
+        chunk_size=None,
+        chunk_memory=None,
+        chunk_duration=None,
+        n_jobs=1,
+        **other_kwargs,
     ):
         return ensure_chunk_size(
-            self.time_series, total_memory, chunk_size, chunk_memory, chunk_duration, n_jobs, **other_kwargs
+            self.time_series,
+            total_memory,
+            chunk_size,
+            chunk_memory,
+            chunk_duration,
+            n_jobs,
+            **other_kwargs,
         )
 
     def run(self, slices=None):
@@ -518,9 +543,7 @@ def run(self, slices=None):
             n_jobs = min(self.n_jobs, len(slices))
 
             if self.pool_engine == "process":
-
                 if self.need_worker_index:
-
                     multiprocessing.set_start_method(self.mp_context, force=True)
                     lock = multiprocessing.Lock()
                     array_pid = multiprocessing.Array("i", n_jobs)
@@ -590,7 +613,6 @@ def run(self, slices=None):
                         lock,
                     ),
                 ) as executor:
-
                     slices2 = [(thread_local_data,) + tuple(args) for args in slices]
                     results = executor.map(thread_function_wrapper, slices2)
 

src/spikeinterface/core/tests/test_globals.py

@@ -1,18 +1,13 @@
-import pytest
-import warnings
-from pathlib import Path
-from os import cpu_count
-
 from spikeinterface import (
-    set_global_dataset_folder,
     get_global_dataset_folder,
-    set_global_tmp_folder,
-    get_global_tmp_folder,
-    set_global_job_kwargs,
     get_global_job_kwargs,
+    get_global_tmp_folder,
     reset_global_job_kwargs,
+    set_global_dataset_folder,
+    set_global_job_kwargs,
+    set_global_tmp_folder,
 )
-from spikeinterface.core.job_tools import fix_job_kwargs
+from spikeinterface.core.job_tools import get_usable_cpu_count, fix_job_kwargs
 
 
 def test_global_dataset_folder(create_cache_folder):
@@ -70,7 +65,7 @@ def test_global_job_kwargs():
         max_threads_per_worker=1,
     )
     # test that fix_job_kwargs grabs global kwargs
-    new_job_kwargs = dict(n_jobs=cpu_count())
+    new_job_kwargs = dict(n_jobs=get_usable_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"]
@@ -81,9 +76,9 @@ def test_global_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)
+    excessive_n_jobs = dict(n_jobs=get_usable_cpu_count() + 2)
     job_kwargs_split = fix_job_kwargs(excessive_n_jobs)
-    assert job_kwargs_split["n_jobs"] == cpu_count()
+    assert job_kwargs_split["n_jobs"] == get_usable_cpu_count()
     reset_global_job_kwargs()
 
 

src/spikeinterface/core/tests/test_job_tools.py

@@ -1,18 +1,23 @@
-import pytest
 import os
-
 import time
 
-from spikeinterface.core import generate_recording, set_global_job_kwargs, get_global_job_kwargs, get_best_job_kwargs
+import pytest
 
+from spikeinterface.core import (
+    generate_recording,
+    get_best_job_kwargs,
+    reset_global_job_kwargs,
+    set_global_job_kwargs,
+)
 from spikeinterface.core.job_tools import (
+    TimeSeriesChunkExecutor,
     divide_segment_into_chunks,
-    ensure_n_jobs,
+    divide_time_series_into_chunks,
     ensure_chunk_size,
-    TimeSeriesChunkExecutor,
+    ensure_n_jobs,
     fix_job_kwargs,
+    get_usable_cpu_count,
     split_job_kwargs,
-    divide_time_series_into_chunks,
 )
 
 
@@ -26,6 +31,33 @@ def test_divide_segment_into_chunks():
     assert chunks[2] == (10, 11)
 
 
+def test_get_usable_cpu_count():
+    # smoke test aginst the actual platform
+    n_cpu = get_usable_cpu_count()
+    assert isinstance(n_cpu, int)
+    assert n_cpu >= 1
+
+
+@pytest.mark.parametrize(
+    "os_func_name, return_value, expected",
+    [
+        ("process_cpu_count", 7, 7),
+        ("sched_getaffinity", {0, 1, 2}, 3),
+        ("cpu_count", 5, 5),
+    ],
+)
+def test_get_usable_cpu_count_fallback(monkeypatch, os_func_name, return_value, expected):
+    # verifies the fallback order: process_cpu_count() > sched_getaffinity() > cpu_count()
+    fallback_funcs = ["process_cpu_count", "sched_getaffinity", "cpu_count"]
+    for name in fallback_funcs:
+        if name != os_func_name:
+            monkeypatch.delattr(os, name, raising=False)
+    stub = (lambda pid: return_value) if os_func_name == "sched_getaffinity" else (lambda: return_value)
+    monkeypatch.setattr(os, os_func_name, stub, raising=False)
+
+    assert get_usable_cpu_count() == expected
+
+
 def test_ensure_n_jobs():
     recording = generate_recording()
 
@@ -80,7 +112,6 @@ def test_ensure_chunk_size():
 
 
 def func(segment_index, start_frame, end_frame, worker_dict):
-    import os
 
     #  print('func', segment_index, start_frame, end_frame, worker_dict, os.getpid())
     time.sleep(0.010)
@@ -134,7 +165,6 @@ def __init__(self):
         def __call__(self, res):
             self.pos += 1
             # print(self.pos, res)
-            pass
 
     gathering_func2 = GatherClass()
 
@@ -193,38 +223,37 @@ def test_fix_job_kwargs():
     # test negative n_jobs
     job_kwargs = dict(n_jobs=-1, progress_bar=False, chunk_duration="1s")
     fixed_job_kwargs = fix_job_kwargs(job_kwargs)
-    assert fixed_job_kwargs["n_jobs"] == os.cpu_count()
+    assert fixed_job_kwargs["n_jobs"] == get_usable_cpu_count()
 
     # test float n_jobs
     job_kwargs = dict(n_jobs=0.5, progress_bar=False, chunk_duration="1s")
     fixed_job_kwargs = fix_job_kwargs(job_kwargs)
-    if int(0.5 * os.cpu_count()) > 1:
-        assert fixed_job_kwargs["n_jobs"] == int(0.5 * os.cpu_count())
+    if int(0.5 * get_usable_cpu_count()) > 1:
+        assert fixed_job_kwargs["n_jobs"] == int(0.5 * get_usable_cpu_count())
     else:
         assert fixed_job_kwargs["n_jobs"] == 1
 
     # test float value > 1 is cast to correct int
-    job_kwargs = dict(n_jobs=float(os.cpu_count()), progress_bar=False, chunk_duration="1s")
+    job_kwargs = dict(n_jobs=float(get_usable_cpu_count()), progress_bar=False, chunk_duration="1s")
     fixed_job_kwargs = fix_job_kwargs(job_kwargs)
-    assert fixed_job_kwargs["n_jobs"] == os.cpu_count()
+    assert fixed_job_kwargs["n_jobs"] == get_usable_cpu_count()
 
     # test wrong keys
     with pytest.raises(AssertionError):
         job_kwargs = dict(n_jobs=0, progress_bar=False, chunk_duration="1s", other_param="other")
         fixed_job_kwargs = fix_job_kwargs(job_kwargs)
 
     # test mutually exclusive
-    _old_global = get_global_job_kwargs().copy()
     set_global_job_kwargs(chunk_memory="50M")
     job_kwargs = dict()
-    fixed_job_kwargs = fixed_job_kwargs = fix_job_kwargs(job_kwargs)
+    fixed_job_kwargs = fix_job_kwargs(job_kwargs)
     assert "chunk_memory" in fixed_job_kwargs
 
     job_kwargs = dict(chunk_duration="300ms")
-    fixed_job_kwargs = fixed_job_kwargs = fix_job_kwargs(job_kwargs)
+    fixed_job_kwargs = fix_job_kwargs(job_kwargs)
     assert "chunk_memory" not in fixed_job_kwargs
     assert fixed_job_kwargs["chunk_duration"] == "300ms"
-    set_global_job_kwargs(**_old_global)
+    reset_global_job_kwargs()
 
 
 def test_split_job_kwargs():