mihaic · Nov 11, 2017 · Nov 10, 2017 · Nov 10, 2017
diff --git a/brainiak/fcma/voxelselector.py b/brainiak/fcma/voxelselector.py
@@ -20,14 +20,14 @@
 # (Intel Labs), 2016
 
 import numpy as np
-import os
 import time
 from mpi4py import MPI
 from scipy.stats.mstats import zscore
 from sklearn import model_selection
 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
@@ -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
@@ -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
@@ -130,6 +130,7 @@ def finalize_options(self):
         'pymanopt',
         'theano',
         'pybind11>=1.7',
+        'psutil',
         'nibabel',
         'typing',
     ],