ENH: disable OpenBLAS affinity settings

OpenBLAS starts as many threads as it can and assigns them to cpus
This includes the main thread Python is running in.
This means any threads or processes Python starts must all share this
cpu, nullifying any parallization efforts.
Luckily this behaviour can be disabled via an environment variable which
we need to set before loading dotblas which may be linked to OpenBLAS.
Libraries loading BLAS before numpy are stills screwed.

1 files changed, 13 insertions, 0 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 778eed8c3..8c569ea15 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -1,5 +1,6 @@
 from __future__ import division, absolute_import, print_function
 
+import os
 import sys
 import warnings
 import collections
@@ -1074,9 +1075,17 @@ def outer(a, b, out=None):
     return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis,:], out)
 
 # try to import blas optimized dot if available
+envbak = os.environ.copy()
 try:
     # importing this changes the dot function for basic 4 types
     # to blas-optimized versions.
+
+    # disables openblas affinity setting of the main thread that limits
+    # python threads or processes to one core
+    if 'OPENBLAS_MAIN_FREE' not in os.environ:
+        os.environ['OPENBLAS_MAIN_FREE'] = '1'
+    if 'GOTOBLAS_MAIN_FREE' not in os.environ:
+        os.environ['GOTOBLAS_MAIN_FREE'] = '1'
     from ._dotblas import dot, vdot, inner, alterdot, restoredot
 except ImportError:
     # docstrings are in add_newdocs.py
@@ -1088,6 +1097,10 @@ except ImportError:
         pass
     def restoredot():
         pass
+finally:
+    os.environ.clear()
+    os.environ.update(envbak)
+    del envbak
 
 def tensordot(a, b, axes=2):
     """