Coverage for python/lsst/utils/threads.py: 24%
25 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-01 02:29 -0700
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-01 02:29 -0700
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)
53 # Also specify an explicit value for OMP_PROC_BIND to tell OpenMP not to
54 # set CPU affinity.
55 var = "OMP_PROC_BIND"
56 if override or var not in os.environ:
57 os.environ[var] = "false"
60def disable_implicit_threading() -> None:
61 """Do whatever is necessary to try to prevent implicit threading.
63 Notes
64 -----
65 Explicitly limits the number of threads allowed to be used by ``numexpr``
66 and attempts to limit the number of threads in all APIs supported by
67 the ``threadpoolctl`` package.
68 """
69 # Force one thread and force override.
70 set_thread_envvars(1, True)
72 try:
73 # This must be a deferred import since importing it immediately
74 # triggers the environment variable examination.
75 # Catch this in case numexpr is not installed.
76 import numexpr.utils
77 except ImportError:
78 pass
79 else:
80 numexpr.utils.set_num_threads(1)
82 # Try to set threads for openblas and openmp
83 if threadpool_limits is not None:
84 threadpool_limits(limits=1)