From 653894bf329f333d1182088ad0e7db997ee1fccb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mihai=20Capot=C4=83?= <mihai@mihaic.ro>
Date: Tue, 7 Nov 2017 14:22:30 -0800
Subject: [PATCH] Restrict multiprocessing to available CPUs

See issue #291.
---
 brainiak/fcma/voxelselector.py      | 13 +++++++++----
 brainiak/searchlight/searchlight.py | 18 +++++++++++++++---
 docs/newsfragments/293.feature      |  1 +
 3 files changed, 25 insertions(+), 7 deletions(-)
 create mode 100644 docs/newsfragments/293.feature

diff --git a/brainiak/fcma/voxelselector.py b/brainiak/fcma/voxelselector.py
index 2c4b25768..95631a66c 100644
--- a/brainiak/fcma/voxelselector.py
+++ b/brainiak/fcma/voxelselector.py
@@ -20,6 +20,7 @@
 # (Intel Labs), 2016
 
 import numpy as np
+import os
 import time
 from mpi4py import MPI
 from scipy.stats.mstats import zscore
@@ -89,10 +90,10 @@ class VoxelSelector:
         The number of voxels assigned to a worker each time
 
     process_num: Optional[int]
-        The number of processes used in cross validation.
+        The maximum number of processes used in cross validation.
         If None, the number of processes will equal
-        the number of hardware threads,
-        i.e. the number returned by cpu_count().
+        the number of available hardware threads,
+        i.e. len(os.sched_getaffinity(0)), falling back to os.cpu_count().
         If 0, cross validation will not use python multiprocessing.
 
     master_rank: int, default 0
@@ -116,7 +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
-        self.process_num = np.min((process_num, multiprocessing.cpu_count()))
+        try:
+            usable_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            usable_cpus = os.cpu_count()
+        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 11a156215..0c5fbe80d 100644
--- a/brainiak/searchlight/searchlight.py
+++ b/brainiak/searchlight/searchlight.py
@@ -12,6 +12,7 @@
 #    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
@@ -387,14 +388,25 @@ def run_block_function(self, block_fn, extra_block_fn_params=None,
         extra_block_fn_params: tuple
             Extra parameters to pass to the block function
 
-        pool_size:    Number of parallel processes in shared memory
-                      process pool
+        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()).
 
         """
         rank = self.comm.rank
 
         results = []
-        with Pool(pool_size) as pool:
+        try:
+            usable_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            usable_cpus = os.cpu_count()
+        if pool_size is None:
+            processes = usable_cpus
+        else:
+            processes = min(pool_size, usable_cpus)
+        with Pool(processes) as pool:
             for idx, block in enumerate(self.blocks):
                 result = pool.apply_async(
                     block_fn,
diff --git a/docs/newsfragments/293.feature b/docs/newsfragments/293.feature
new file mode 100644
index 000000000..6cdc6483b
--- /dev/null
+++ b/docs/newsfragments/293.feature
@@ -0,0 +1 @@
+Restricted multiprocessing to the available number of CPUs.