diff --git a/brainiak/fcma/voxelselector.py b/brainiak/fcma/voxelselector.py index 95631a66c..d09d2d499 100644 --- a/brainiak/fcma/voxelselector.py +++ b/brainiak/fcma/voxelselector.py @@ -20,7 +20,6 @@ # (Intel Labs), 2016 import numpy as np -import os import time from mpi4py import MPI from scipy.stats.mstats import zscore @@ -28,6 +27,7 @@ import sklearn from . import fcma_extension # type: ignore from . import cython_blas as blas # type: ignore +from ..utils.utils import usable_cpu_count import logging import multiprocessing @@ -92,8 +92,8 @@ class VoxelSelector: process_num: Optional[int] The maximum number of processes used in cross validation. If None, the number of processes will equal - the number of available hardware threads, - i.e. len(os.sched_getaffinity(0)), falling back to os.cpu_count(). + the number of available hardware threads, considering cpusets + restrictions. If 0, cross validation will not use python multiprocessing. master_rank: int, default 0 @@ -117,11 +117,11 @@ def __init__(self, self.num_voxels2 = raw_data2[0].shape[1] \ if raw_data2 is not None else self.num_voxels self.voxel_unit = voxel_unit - try: - usable_cpus = len(os.sched_getaffinity(0)) - except AttributeError: - usable_cpus = os.cpu_count() - self.process_num = np.min((process_num, usable_cpus)) + usable_cpus = usable_cpu_count() + if process_num is None: + self.process_num = usable_cpus + else: + self.process_num = np.min((process_num, usable_cpus)) if self.process_num == 0: self.use_multiprocessing = False else: diff --git a/brainiak/searchlight/searchlight.py b/brainiak/searchlight/searchlight.py index 0c5fbe80d..968a33e99 100644 --- a/brainiak/searchlight/searchlight.py +++ b/brainiak/searchlight/searchlight.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from multiprocessing import Pool import numpy as np from mpi4py import MPI from scipy.spatial.distance import cityblock +from ..utils.utils import usable_cpu_count + """Distributed Searchlight """ @@ -390,18 +391,13 @@ def run_block_function(self, block_fn, extra_block_fn_params=None, pool_size: int Maximum number of processes running the block function in parallel. - The actual number of processes created is - min(pool_size, len(os.sched_getaffinity(0))) (falling back to - min(pool_size, os.cpu_count()). - + If None, number of available hardware threads, considering cpusets + restrictions. """ rank = self.comm.rank results = [] - try: - usable_cpus = len(os.sched_getaffinity(0)) - except AttributeError: - usable_cpus = os.cpu_count() + usable_cpus = usable_cpu_count() if pool_size is None: processes = usable_cpus else: diff --git a/brainiak/utils/utils.py b/brainiak/utils/utils.py index 78c030d7b..4c6772293 100644 --- a/brainiak/utils/utils.py +++ b/brainiak/utils/utils.py @@ -15,6 +15,7 @@ import re import warnings import os.path +import psutil from .fmrisim import generate_stimfunction, _double_gamma_hrf, convolve_hrf """ @@ -647,3 +648,22 @@ def center_mass_exp(interval, scale=1.0): np.exp(-interval_left / scale) - np.exp(-interval_right / scale)) else: return interval_left + scale + + +def usable_cpu_count(): + """Get number of CPUs usable by the current process. + + Takes into consideration cpusets restrictions. + + Returns + ------- + int + """ + try: + result = len(os.sched_getaffinity(0)) + except AttributeError: + try: + result = len(psutil.Process().cpu_affinity()) + except AttributeError: + result = os.cpu_count() + return result diff --git a/setup.py b/setup.py index 8e44cc60b..c207b7258 100644 --- a/setup.py +++ b/setup.py @@ -130,6 +130,7 @@ def finalize_options(self): 'pymanopt', 'theano', 'pybind11>=1.7', + 'psutil', 'nibabel', 'typing', ],