Coverage for python/lsst/utils/threads.py: 88%

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# 

12"""Support for threading and multi-processing.""" 

13 

14from __future__ import annotations 

15 

16__all__ = ["disable_implicit_threading", "set_thread_envvars"] 

17 

18import logging 

19import os 

20import sys 

21 

22try: 

23 from threadpoolctl import threadpool_limits 

24except ImportError: 

25 threadpool_limits = None 

26 

27_LOG = logging.getLogger(__name__) 

28 

29 

30def set_thread_envvars(num_threads: int = 1, override: bool = False) -> None: 

31 """Set common threading environment variables to the given value. 

32 

33 Parameters 

34 ---------- 

35 num_threads : `int`, optional 

36 Number of threads to use when setting the environment variable values. 

37 Default to 1 (disable threading). 

38 override : `bool`, optional 

39 Controls whether a previously set value should be over-ridden. Defaults 

40 to `False`. 

41 """ 

42 envvars = ( 

43 "OPENBLAS_NUM_THREADS", 

44 "GOTO_NUM_THREADS", 

45 "OMP_NUM_THREADS", 

46 "MKL_NUM_THREADS", 

47 "MKL_DOMAIN_NUM_THREADS", 

48 "MPI_NUM_THREADS", 

49 "NUMEXPR_NUM_THREADS", 

50 "NUMEXPR_MAX_THREADS", 

51 "NUMBA_NUM_THREADS", 

52 "ARROW_IO_THREADS", 

53 "VECLIB_MAXIMUM_THREADS", 

54 "RAYON_NUM_THREADS", 

55 "POLARS_MAX_THREADS", 

56 "BLIS_NUM_THREADS", 

57 ) 

58 

59 for var in envvars: 

60 if override or var not in os.environ: 

61 os.environ[var] = str(num_threads) 

62 

63 # Also specify an explicit value for OMP_PROC_BIND to tell OpenMP not to 

64 # set CPU affinity. 

65 var = "OMP_PROC_BIND" 

66 if override or var not in os.environ: 

67 os.environ[var] = "false" 

68 

69 

70def disable_implicit_threading() -> None: 

71 """Do whatever is necessary to try to prevent implicit threading. 

72 

73 Notes 

74 ----- 

75 Explicitly limits the number of threads allowed to be used by ``numexpr`` 

76 and ``pyarrow`` (if already imported) and attempts to limit the number of 

77 threads in all APIs supported by the ``threadpoolctl`` package. 

78 """ 

79 # Force one thread and force override. 

80 set_thread_envvars(1, True) 

81 

82 try: 

83 # This must be a deferred import since importing it immediately 

84 # triggers the environment variable examination. 

85 # Catch this in case numexpr is not installed. 

86 import numexpr.utils 

87 except ImportError: 

88 pass 

89 else: 

90 numexpr.utils.set_num_threads(1) 

91 

92 # pyarrow sizes its thread pools from the environment only when each pool 

93 # is first used, so pools that may already exist must be resized 

94 # explicitly. If pyarrow has not been imported the environment variables 

95 # set above are sufficient and there is no need to pay for an import here. 

96 if (pyarrow := sys.modules.get("pyarrow")) is not None: 96 ↛ 108line 96 didn't jump to line 108 because the condition on line 96 was always true

97 pyarrow.set_cpu_count(1) 

98 pyarrow.set_io_thread_count(1) 

99 

100 # numba records the value of NUMBA_NUM_THREADS as the maximum size of its 

101 # thread pool, and numba.config.reload_config() (which numba runs on every 

102 # jit compilation) raises if the environment variable no longer matches 

103 # that recorded value once the pool has been launched. The number of 

104 # threads actually used defaults to that recorded maximum too, so limiting 

105 # an already-imported numba requires masking execution with 

106 # set_num_threads in addition to keeping the recorded value and the 

107 # environment variable consistent. 

108 if (numba := sys.modules.get("numba")) is not None: 108 ↛ 123line 108 didn't jump to line 123 because the condition on line 108 was always true

109 try: 

110 # The pool has not been launched yet, so make numba adopt the 

111 # NUMBA_NUM_THREADS=1 set above before it launches. 

112 numba.config.reload_config() 

113 except RuntimeError: 

114 # The pool is already running at a size that cannot be lowered, so 

115 # numba refuses the new value. Realign NUMBA_NUM_THREADS with the 

116 # launched value instead so later reloads stay consistent. 

117 os.environ["NUMBA_NUM_THREADS"] = str(numba.config.NUMBA_NUM_THREADS) 

118 numba.config.reload_config() 

119 # Mask execution down to a single thread. 

120 numba.set_num_threads(1) 

121 

122 # Try to set threads for openblas and openmp 

123 if threadpool_limits is not None: 

124 threadpool_limits(limits=1) 

125 else: 

126 _LOG.warning( 

127 "threadpoolctl is not installed: thread pools of already-loaded libraries cannot be " 

128 "limited and implicit threading may remain enabled." 

129 )