diff --git a/.github/workflows/deepinterpolation.yml b/.github/workflows/deepinterpolation.yml index 913022ac6e..1d8d246575 100644 --- a/.github/workflows/deepinterpolation.yml +++ b/.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: | diff --git a/src/spikeinterface/comparison/tests/test_multisortingcomparison.py b/src/spikeinterface/comparison/tests/test_multisortingcomparison.py index 90c24898e2..6fad43bde0 100644 --- a/src/spikeinterface/comparison/tests/test_multisortingcomparison.py +++ b/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(): diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 962d5aa040..737fc1c668 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/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) diff --git a/src/spikeinterface/core/tests/test_globals.py b/src/spikeinterface/core/tests/test_globals.py index ba31f6f08a..b8f54dfe0e 100644 --- a/src/spikeinterface/core/tests/test_globals.py +++ b/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() diff --git a/src/spikeinterface/core/tests/test_job_tools.py b/src/spikeinterface/core/tests/test_job_tools.py index 596f0da932..1f0f7a3891 100644 --- a/src/spikeinterface/core/tests/test_job_tools.py +++ b/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,20 +223,20 @@ 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): @@ -214,17 +244,16 @@ def test_fix_job_kwargs(): 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(): diff --git a/src/spikeinterface/preprocessing/deepinterpolation/train.py b/src/spikeinterface/preprocessing/deepinterpolation/train.py index da698df5fd..144ac7ff71 100644 --- a/src/spikeinterface/preprocessing/deepinterpolation/train.py +++ b/src/spikeinterface/preprocessing/deepinterpolation/train.py @@ -1,13 +1,13 @@ -import os +import multiprocessing as mp import warnings +from collections.abc import Callable +from concurrent.futures import ProcessPoolExecutor from pathlib import Path -from typing import Callable, Optional -from concurrent.futures import ProcessPoolExecutor -import multiprocessing as mp +from spikeinterface.core import BaseRecording +from spikeinterface.core.job_tools import get_usable_cpu_count from .tf_utils import import_tf -from spikeinterface.core import BaseRecording global train_func @@ -17,17 +17,17 @@ def train_deepinterpolation( model_folder: str | Path, model_name: str, desired_shape: tuple[int, int], - train_start_s: Optional[float] = None, - train_end_s: Optional[float] = None, - train_duration_s: Optional[float] = None, - test_start_s: Optional[float] = None, - test_end_s: Optional[float] = None, - test_duration_s: Optional[float] = None, - test_recordings: Optional[BaseRecording | list[BaseRecording]] = None, + train_start_s: float | None = None, + train_end_s: float | None = None, + train_duration_s: float | None = None, + test_start_s: float | None = None, + test_end_s: float | None = None, + test_duration_s: float | None = None, + test_recordings: BaseRecording | list[BaseRecording] | None = None, pre_frame: int = 30, post_frame: int = 30, pre_post_omission: int = 1, - existing_model_path: Optional[str | Path] = None, + existing_model_path: str | Path | None = None, verbose: bool = True, nb_gpus: int = 1, steps_per_epoch: int = 10, @@ -42,7 +42,7 @@ def train_deepinterpolation( network: Callable | None = None, use_gpu: bool = True, disable_tf_logger: bool = True, - memory_gpu: Optional[int] = None, + memory_gpu: int | None = None, ): """ Train a deepinterpolation model from a recording extractor. @@ -118,7 +118,7 @@ def train_deepinterpolation( """ if nb_workers == -1: - nb_workers = os.cpu_count() + nb_workers = get_usable_cpu_count() args = ( recordings, @@ -171,11 +171,11 @@ def train_deepinterpolation_process( test_start_s: float, test_end_s: float, test_duration_s: float | None, - test_recordings: Optional[BaseRecording | list[BaseRecording]] = None, + test_recordings: BaseRecording | list[BaseRecording] | None = None, pre_frame: int = 30, post_frame: int = 30, pre_post_omission: int = 1, - existing_model_path: Optional[str | Path] = None, + existing_model_path: str | Path | None = None, verbose: bool = True, nb_gpus: int = 1, steps_per_epoch: int = 10, @@ -190,9 +190,10 @@ def train_deepinterpolation_process( network: Callable | None = None, use_gpu: bool = True, disable_tf_logger: bool = True, - memory_gpu: Optional[int] = None, + memory_gpu: int | None = None, ): from deepinterpolation.trainor_collection import core_trainer + from .generators import SpikeInterfaceRecordingGenerator # initialize TF diff --git a/src/spikeinterface/sorters/external/spyking_circus.py b/src/spikeinterface/sorters/external/spyking_circus.py index 274836a707..ef9187e6dd 100644 --- a/src/spikeinterface/sorters/external/spyking_circus.py +++ b/src/spikeinterface/sorters/external/spyking_circus.py @@ -1,19 +1,17 @@ -from pathlib import Path -import os import importlib.util -from importlib.metadata import version import sys +from importlib.metadata import version +from pathlib import Path import numpy as np from numpy.lib.format import open_memmap +from probeinterface import write_prb - +from spikeinterface.core.job_tools import get_usable_cpu_count from spikeinterface.extractors.extractor_classes import SpykingCircusSortingExtractor from spikeinterface.sorters.basesorter import BaseSorter from spikeinterface.sorters.utils import ShellScript -from probeinterface import write_prb - class SpykingcircusSorter(BaseSorter): """SpykingCircus Sorter object.""" @@ -82,7 +80,7 @@ def _check_params(cls, recording, sorter_output_folder, params): # check and re dump params p = params if p["num_workers"] is None: - p["num_workers"] = np.maximum(1, int(os.cpu_count() / 2)) + p["num_workers"] = np.maximum(1, int(get_usable_cpu_count() / 2)) return p @classmethod