Coverage for python/lsst/utils/threads.py: 29%
22 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-13 09:31 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-04-13 09:31 +0000
1# This file is part of utils.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
11#
12from __future__ import annotations
14"""Support for threading and multi-processing."""
16__all__ = ["set_thread_envvars", "disable_implicit_threading"]
18import os
20try:
21 from threadpoolctl import threadpool_limits
22except ImportError:
23 threadpool_limits = None
26def set_thread_envvars(num_threads: int = 1, override: bool = False) -> None:
27 """Set common threading environment variables to the given value.
29 Parameters
30 ----------
31 num_threads : `int`, optional
32 Number of threads to use when setting the environment variable values.
33 Default to 1 (disable threading).
34 override : `bool`, optional
35 Controls whether a previously set value should be over-ridden. Defaults
36 to `False`.
37 """
38 envvars = (
39 "OPENBLAS_NUM_THREADS",
40 "GOTO_NUM_THREADS",
41 "OMP_NUM_THREADS",
42 "MKL_NUM_THREADS",
43 "MKL_DOMAIN_NUM_THREADS",
44 "MPI_NUM_THREADS",
45 "NUMEXPR_NUM_THREADS",
46 "NUMEXPR_MAX_THREADS",
47 )
49 for var in envvars:
50 if override or var not in os.environ:
51 os.environ[var] = str(num_threads)
54def disable_implicit_threading() -> None:
55 """Do whatever is necessary to try to prevent implicit threading.
57 Notes
58 -----
59 Explicitly limits the number of threads allowed to be used by ``numexpr``
60 and attempts to limit the number of threads in all APIs supported by
61 the ``threadpoolctl`` package.
62 """
63 # Force one thread and force override.
64 set_thread_envvars(1, True)
66 try:
67 # This must be a deferred import since importing it immediately
68 # triggers the environment variable examination.
69 # Catch this in case numexpr is not installed.
70 import numexpr.utils
71 except ImportError:
72 pass
73 else:
74 numexpr.utils.set_num_threads(1)
76 # Try to set threads for openblas and openmp
77 if threadpool_limits is not None:
78 threadpool_limits(limits=1)