Limit number of BLAS threads during parallelization. Fixes #1

MatthewThe · Jan 2, 2024 · 0411ec4 · 0411ec4
1 parent 3fce61b
commit 0411ec4
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/job_pool/job_pool.py b/job_pool/job_pool.py
@@ -1,3 +1,4 @@
+from threadpoolctl import threadpool_limits
 import time
 import signal
 from typing import Optional
@@ -58,6 +59,7 @@ class AbnormalWorkerTerminationError(Exception):
 
 
 class JobPool:
+    @threadpool_limits.wrap(limits=1, user_api='blas')
     def __init__(
         self,
         processes: int = 1,

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -10,6 +10,7 @@ packages = [{include = "job_pool"}]
 [tool.poetry.dependencies]
 python = "^3.8"
 tqdm = "^4.66.1"
+threadpoolctl = "^3.2.0"
 
 
 [tool.poetry.group.dev.dependencies]