Coverage for src/cvx/quadprog/_threads.py: 100%
124 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-15 05:55 +0000
1"""BLAS thread count: an opt-in cap, and the probes that describe the machine.
3This package pushes its work into BLAS calls, so the BLAS thread count is a
4first-order performance parameter -- and on Linux with OpenBLAS its default is a
5trap. Contributed measurements on a Ryzen 7 5800X (8 physical cores, 16 logical),
6same machine, same package, only ``OPENBLAS_NUM_THREADS`` differing (#41, #66):
8====================== ================== ======== =======
9run unset (16 threads) ``=1`` penalty
10====================== ================== ======== =======
11exact, ``n = 800`` 5666 ms 77.3 ms 73x
12exact, ``n = 1600`` 15697 ms 517 ms 30x
13``fast=True, n = 200`` 73.7 ms 1.53 ms 48x
14====================== ================== ======== =======
16The failure is a cliff at oversubscription rather than a slope: the same sweep
17scales normally out to four threads (1.24x-1.50x on the exact path) and only then
18falls apart. Threaded MKL asked to oversubscribe by the same factor does not
19degrade at all -- 88 ms against OpenBLAS's 7274 ms at 16 threads on ``n = 800``
20exact -- so this is specific to OpenBLAS, not a property of oversubscription in
21general and not a Linux problem. Windows ships the same threading layer and never
22got worse than 0.34x in five reports; Accelerate exposes no thread knob.
24Two things follow, and this module is the smaller of them.
26**The cap.** :func:`limit` exposes a scoped ``threadpoolctl`` context, reached
27through ``solve_qp(..., blas_threads=...)``. It is opt-in, and deliberately not a
28default: the right thread count differs by BLAS *in opposite directions* -- the
29fast path wants 4 on OpenBLAS, where 16 reads 0.05x, and 16 on MKL, where it is
30still improving -- and it differs by path, since every Windows exact-path sweep is
31best at one thread. No single number serves all three, and ``threadpoolctl``
32itself costs ~100 microseconds, which is real against a 0.2 ms solve at
33``n = 10``. So nothing here runs unless the caller asks for it.
35**The machine.** :func:`_is_openblas`, :func:`_intended_threads` and
36:func:`_physical_cores` answer the three questions that decide whether a process
37is in the configuration above: which BLAS NumPy was built against, how many
38threads it will start, and how many physical cores there actually are. They are
39deliberately built on stdlib and NumPy alone -- no ``threadpoolctl`` -- so that
40they work on a plain ``pip install`` of this package, which is exactly the
41installation the trap is set for.
43**The automatic cap.** :func:`auto_cap_threads` is what consumes them. When
44``blas_threads`` is not given, ``solve_qp`` asks it whether this process is in the
45configuration above *and* the problem is large enough for the collapse to be
46reachable, and caps to the physical core count when both hold.
47:class:`~cvx.quadprog.Sweep` asks the same question once per object, since its
48``n`` is fixed for its lifetime, and applies the answer to its factorisation and
49to its misses but not to its hits -- the context costs ~100 microseconds, which is
50real against a hit that is tens of microseconds at the sizes this matters for. The
51size gate comes from the hardware rather than a constant:
52:func:`dynamic_n_thresh` divides the probed per-core L2 cache by the
53bytes each path touches per variable, so the cap engages where the working set
54stops fitting and threading starts to matter. Every probe is cached for the life
55of the process, so the solve path pays a cached lookup rather than a file read.
57The two callers of ``limit`` want opposite failure modes, and the gate is where
58that is decided. ``blas_threads=N`` is a request, so a missing ``threadpoolctl``
59raises -- silently not capping is not what the caller asked for. The automatic cap
60is a defence nobody asked for, so a missing ``threadpoolctl`` makes
61:func:`_auto_cap_target` decline, with one warning per process pointing at
62``OPENBLAS_NUM_THREADS``. Declining leaves the process exactly as it was before
63this gate existed; raising would make an optional dependency mandatory on the one
64installation the trap is set for (#106).
65"""
67# G is the name from Goldfarb & Idnani (1983), as everywhere else in this package.
68# TRY003 goes with the waivers in _setup.py: an exception raised for a mistake in
69# one argument has to say which argument and what was wrong with it, and a
70# dedicated exception class per message would be worse.
71# ruff: noqa: TRY003
73import functools
74import math
75import os
76import pathlib
77import platform
78import warnings
79from contextlib import AbstractContextManager, nullcontext
80from typing import Any
82import numpy as np
84# Linux exposes the CPU topology here. One file per logical CPU, each holding the
85# set of logical CPUs sharing its physical core -- so the number of *distinct*
86# contents is the number of physical cores. Read rather than inferred from
87# `os.cpu_count() // 2`, which is the physical core count only under 2-way SMT
88# and silently wrong on the parts that have none (#66).
89_SYSFS_CPU = pathlib.Path("/sys/devices/system/cpu")
91# Consulted in this order for the thread count OpenBLAS will start with. Both are
92# read rather than only the OpenBLAS-specific one, because OMP_NUM_THREADS is
93# what a user who has already thought about threading is most likely to have set.
94_THREAD_ENV = ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS")
97def _parse_cache_size(size_str: str) -> int:
98 """Parse cache size strings like '512K', '1024K', '1M', '2048' into bytes.
100 Args:
101 size_str: Raw size string from sysfs or configuration.
103 Returns:
104 Integer size in bytes.
105 """
106 s = size_str.strip().upper()
107 if s.endswith("B"):
108 s = s[:-1]
109 if s.endswith("K"):
110 return int(s[:-1]) * 1024
111 if s.endswith("M"):
112 return int(s[:-1]) * 1024 * 1024
113 return int(s)
116def _probe_sysfs_l2() -> int | None:
117 """Return the per-core L2 size from Linux sysfs, or None if it cannot be read.
119 Scans `/sys/devices/system/cpu/cpu0/cache/index*` for the entry whose `level`
120 file reads 2, since the index a level lands on is not fixed across parts. If
121 no entry declares its level, `index2` is read directly as the conventional
122 placement.
124 Returns:
125 L2 cache size in bytes, or None on a host without the tree, without a
126 Level-2 entry in it, or with one that does not read or parse.
127 """
128 try:
129 cache_dir = _SYSFS_CPU / "cpu0" / "cache"
130 if not cache_dir.exists():
131 return None
133 for index_path in cache_dir.glob("index*"):
134 level_path = index_path / "level"
135 size_path = index_path / "size"
136 if level_path.exists() and level_path.read_text().strip() == "2" and size_path.exists():
137 return _parse_cache_size(size_path.read_text())
139 idx2_size = cache_dir / "index2" / "size"
140 if idx2_size.exists():
141 return _parse_cache_size(idx2_size.read_text())
142 except (OSError, ValueError):
143 return None
145 return None
148def _probe_sysconf_l2() -> int | None:
149 """Return the per-core L2 size from POSIX sysconf, or None if unavailable.
151 Absent on Windows, and present but unanswerable on hosts that report zero or
152 a negative size for the name -- both of which are None here rather than a
153 number the caller would have to re-check.
155 Returns:
156 L2 cache size in bytes, or None.
157 """
158 if not hasattr(os, "sysconf"):
159 return None
161 try:
162 val = os.sysconf("SC_LEVEL2_CACHE_SIZE")
163 except (ValueError, OSError):
164 return None
166 return val if isinstance(val, int) and val > 0 else None
169@functools.cache
170def _probe_l2_cache_bytes() -> int:
171 """Return per-core L2 cache size in bytes, defaulting to 512 KB (524288).
173 Consults Linux sysfs first and POSIX `sysconf` second, each of which answers
174 None where it does not apply, and falls back to 512 KB as the baseline floor
175 -- the value the threshold in :func:`dynamic_n_thresh` is calibrated against
176 on a host that reports nothing.
178 Returns:
179 L2 cache size in bytes.
180 """
181 for probe in (_probe_sysfs_l2, _probe_sysconf_l2):
182 size = probe()
183 if size is not None:
184 return size
186 return 512 * 1024
189@functools.cache
190def dynamic_n_thresh(fast: bool = False) -> int:
191 """Return the problem dimension threshold n_thresh for hardware L2 caching.
193 Derived dynamically from the host CPU's probed L2 cache size per core.
195 - For exact Hessian G (size n x n), memory footprint is 8*n^2 bytes.
196 Setting 8*n^2 = L2 yields n_thresh_exact = sqrt(L2 / 8).
197 - For fast path PDAS (KKT system matrix ~ 2n x 2n), peak footprint is 32*n^2 bytes.
198 Setting 32*n^2 = L2 yields n_thresh_fast = sqrt(L2 / 32).
200 Args:
201 fast: True if using fast path PDAS KKT matrix expansion, False for exact walk.
203 Returns:
204 Integer threshold n.
205 """
206 l2_bytes = _probe_l2_cache_bytes()
207 divisor = 32 if fast else 8
208 return max(16, int(math.sqrt(l2_bytes / divisor)))
211def _threadpoolctl_available() -> bool:
212 """Return whether the cap can actually be applied in this process.
214 Spelled as the same import :func:`limit` performs, rather than as a
215 ``find_spec`` probe, so that the two cannot disagree: whatever makes ``limit``
216 raise makes this False. Not cached, because its only caller is
217 :func:`_auto_cap_target`, which is -- so the import is attempted once per
218 process either way and a second cache would only be a second thing for a test
219 to have to clear.
221 Returns:
222 True if ``threadpoolctl`` imports, False if it is not installed.
223 """
224 try:
225 from threadpoolctl import threadpool_limits # noqa: F401
226 except ImportError:
227 return False
228 return True
231@functools.cache
232def _auto_cap_target() -> int | None:
233 """Return physical core count if oversubscribed on Linux with OpenBLAS, else None.
235 Evaluated once per process and cached so zero file reads or system calls
236 occur during solver execution.
238 Returns:
239 The physical core count to cap to, or None when this process is not in the
240 measured pathology -- or is, but has no way to act on it.
241 """
242 if platform.system() != "Linux":
243 return None
245 if not _is_openblas():
246 return None
248 threads, cores = _intended_threads(), _physical_cores()
249 if threads is None or cores is None or threads <= cores:
250 return None
252 if not _threadpoolctl_available():
253 # This gate fires without being asked, so a missing optional dependency
254 # has to mean "no cap" and not "no solve" (#106). Raising here would turn
255 # a performance defence into an availability failure, on precisely the
256 # plain `pip install` the probes above are written in stdlib to support.
257 #
258 # It is still worth saying so once: the caller is in the one configuration
259 # measured at 73x, and the environment variable fixes it with nothing
260 # installed. Once per process comes free with the cache on this function.
261 warnings.warn(
262 f"This process is configured for {threads} BLAS threads on {cores} physical cores, "
263 "which is the OpenBLAS oversubscription that measures up to 73x slower. Capping it "
264 "per solve needs threadpoolctl, which is not installed: `pip install threadpoolctl`, "
265 f"or set OPENBLAS_NUM_THREADS={cores} to cap the whole process instead.",
266 RuntimeWarning,
267 stacklevel=2,
268 )
269 return None
271 return cores
274def auto_cap_threads(n: int, fast: bool = False) -> int | None:
275 """Return recommended physical thread limit if oversubscribed & n >= n_thresh.
277 Evaluates with zero system calls or file reads on the solve path.
279 Args:
280 n: Dimension of the Hessian G (leading dimension).
281 fast: True if fast path PDAS active-set expansion applies.
283 Returns:
284 Physical core count to cap BLAS threads at, or None if no cap is needed.
285 """
286 target = _auto_cap_target()
287 if target is None:
288 return None
290 if n < dynamic_n_thresh(fast=fast):
291 return None
293 return target
296def limit(threads: int) -> AbstractContextManager[Any]:
297 """Return a context manager capping the BLAS thread count for its body.
299 Args:
300 threads: Maximum number of threads the BLAS may use inside the context.
302 Returns:
303 A ``threadpoolctl`` context manager. It restores the previous limits on
304 exit, so nothing about the process outlives the ``with`` block -- which is
305 the whole reason this is a context manager and not a setting.
307 Raises:
308 ValueError: If ``threads`` is not at least 1.
309 ImportError: If ``threadpoolctl`` is not installed. It is an optional
310 dependency (``pip install cvx-quadprog[threads]``) rather than a
311 required one, because it is needed only by callers who use this.
312 """
313 if threads < 1:
314 raise ValueError(f"blas_threads must be a positive integer. Received {threads}")
316 try:
317 from threadpoolctl import threadpool_limits
318 except ImportError as exc:
319 raise ImportError(
320 "blas_threads needs threadpoolctl, which is an optional dependency of this package. "
321 "Install it with `pip install threadpoolctl`, or set OPENBLAS_NUM_THREADS in the "
322 "environment instead -- that caps the whole process rather than one call."
323 ) from exc
325 # Annotated on the way out because threadpoolctl ships no type information, so
326 # what it returns is `Any` and returning it directly is an untyped escape.
327 limiter: AbstractContextManager[Any] = threadpool_limits(limits=threads, user_api="blas")
328 return limiter
331def scoped_limit(threads: int | None) -> AbstractContextManager[Any]:
332 """Return :func:`limit`, or a context that does nothing when there is no cap.
334 Lets a caller wrap a block once instead of writing the call site twice, which
335 is what :class:`~cvx.quadprog.Sweep` needs: its cap is decided once per object
336 and then applies to two separate blocks.
338 Not used by :func:`~cvx.quadprog.solve_qp`, which keeps its three explicit
339 return paths. Entering a ``nullcontext`` is cheap but not free, and that
340 function is on the small-``n`` dispatch path this package measures in
341 percentage points; the blocks this wraps are a Cholesky and a full solve.
343 Args:
344 threads: Maximum number of threads, or None to leave the process alone.
346 Returns:
347 A context manager, restoring the previous limits on exit where it set any.
349 Raises:
350 ValueError: If ``threads`` is given and is not at least 1.
351 ImportError: If ``threads`` is given and ``threadpoolctl`` is not installed.
352 """
353 return nullcontext() if threads is None else limit(threads)
356def _is_openblas() -> bool:
357 """Return whether NumPy was built against OpenBLAS.
359 NumPy's build configuration is a proxy for the library actually loaded, and an
360 imperfect one -- SciPy could in principle have been built against a different
361 BLAS, and this package calls both. It is used anyway because the alternative,
362 ``threadpoolctl``, would have to become a required dependency to make this work
363 on a plain ``pip install``, which is exactly the installation that matters
364 here. The failure mode of getting it wrong is a decision not taken, or one that
365 names the wrong library while still being right about the core counts.
367 Returns:
368 True if NumPy reports an OpenBLAS-family BLAS, False if it reports
369 anything else or nothing intelligible.
370 """
371 try:
372 config = np.show_config(mode="dicts")
373 name = config["Build Dependencies"]["blas"]["name"]
374 except (KeyError, TypeError):
375 # No config, or a shape this does not know how to read. Either way there
376 # is no evidence of OpenBLAS, which is the answer.
377 return False
379 return "openblas" in str(name).lower()
382def _intended_threads() -> int | None:
383 """Return the thread count OpenBLAS will start with, or None if unpredictable.
385 An explicitly set variable is honoured rather than treated as consent: a user
386 who has set ``OPENBLAS_NUM_THREADS=16`` on eight cores has the same problem as
387 one who set nothing, and the value is right there to be compared. Whether a
388 given value is *fine* is for the caller of this function to decide.
390 Returns:
391 The value of the first of :data:`_THREAD_ENV` that is set, or the logical
392 CPU count when neither is; None if the variable holds something that is
393 not an integer, or if the CPU count is unavailable.
394 """
395 for name in _THREAD_ENV:
396 value = os.environ.get(name)
397 if value is not None:
398 try:
399 return int(value)
400 except ValueError:
401 # A malformed value is OpenBLAS's problem to interpret, not ours
402 # to guess at.
403 return None
405 # Uncapped, OpenBLAS threads to the number of CPUs it detects, which counts
406 # SMT siblings. `os.cpu_count()` overreports under a CPU-set restriction, but
407 # so does the sysfs topology this is compared against, so the comparison
408 # survives it.
409 return os.cpu_count()
412@functools.cache
413def _physical_cores() -> int | None:
414 """Return the number of physical cores, or None if the topology is unreadable.
416 Returns:
417 The number of distinct thread-sibling sets under :data:`_SYSFS_CPU`, which
418 is the physical core count; None where that directory does not exist or
419 cannot be read, which includes every non-Linux platform and some
420 containers.
421 """
422 try:
423 siblings = {path.read_text() for path in _SYSFS_CPU.glob("cpu[0-9]*/topology/thread_siblings_list")}
424 except OSError:
425 return None
427 # An empty glob means no topology to read, not a machine with no cores.
428 cores = len(siblings) or None
429 if cores is not None and hasattr(os, "sched_getaffinity"):
430 try:
431 affinity = len(os.sched_getaffinity(0))
432 if affinity > 0:
433 cores = min(cores, affinity)
434 except OSError:
435 pass
437 return cores