Coverage for tests/test_threads.py: 83%

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# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22import os 

23import unittest 

24import unittest.mock 

25 

26from lsst.utils.threads import disable_implicit_threading, set_thread_envvars 

27 

28try: 

29 import numexpr 

30except ImportError: 

31 numexpr = None 

32try: 

33 import threadpoolctl 

34except ImportError: 

35 threadpoolctl = None 

36try: 

37 import pyarrow 

38except ImportError: 

39 pyarrow = None 

40try: 

41 import numba 

42except ImportError: 

43 numba = None 

44 

45 

46class ThreadsTestCase(unittest.TestCase): 

47 """Tests for threads.""" 

48 

49 def testDisable(self): 

50 set_thread_envvars(2, override=True) 

51 self.assertEqual(os.environ["OMP_NUM_THREADS"], "2") 

52 set_thread_envvars(3, override=False) 

53 self.assertEqual(os.environ["OMP_NUM_THREADS"], "2") 

54 

55 disable_implicit_threading() 

56 self.assertEqual(os.environ["OMP_NUM_THREADS"], "1") 

57 self.assertEqual(os.environ["OMP_PROC_BIND"], "false") 

58 

59 # Check that we have only one thread. 

60 if numexpr: 60 ↛ 62line 60 didn't jump to line 62 because the condition on line 60 was always true

61 self.assertEqual(numexpr.utils.get_num_threads(), 1) 

62 if threadpoolctl: 62 ↛ exitline 62 didn't return from function 'testDisable' because the condition on line 62 was always true

63 info = threadpoolctl.threadpool_info() 

64 for api in info: 

65 self.assertEqual(api["num_threads"], 1, f"API: {api}") 

66 

67 def testEnvVarsForEnvOnlyLibraries(self): 

68 """Libraries that are only controllable via environment variables 

69 must be included in the list of variables that are set. 

70 """ 

71 set_thread_envvars(2, override=True) 

72 for var in ( 

73 "NUMBA_NUM_THREADS", 

74 "ARROW_IO_THREADS", 

75 "VECLIB_MAXIMUM_THREADS", 

76 "RAYON_NUM_THREADS", 

77 "POLARS_MAX_THREADS", 

78 "BLIS_NUM_THREADS", 

79 ): 

80 self.assertEqual(os.environ.get(var), "2", f"Variable: {var}") 

81 

82 @unittest.skipIf(pyarrow is None, "pyarrow is not available") 

83 def testDisablePyarrow(self): 

84 """Already-created pyarrow thread pools must be resized since they 

85 do not react to environment variables after creation. 

86 """ 

87 pyarrow.set_cpu_count(4) 

88 pyarrow.set_io_thread_count(4) 

89 disable_implicit_threading() 

90 self.assertEqual(pyarrow.cpu_count(), 1) 

91 self.assertEqual(pyarrow.io_thread_count(), 1) 

92 

93 @unittest.skipIf(numba is None, "numba is not available") 

94 def testDisableNumba(self): 

95 """An already-imported numba must be limited to a single thread, and 

96 jit compilation afterwards must still work. 

97 """ 

98 if numba.config.NUMBA_NUM_THREADS > 1: 98 ↛ 100line 98 didn't jump to line 100 because the condition on line 98 was always true

99 numba.set_num_threads(2) 

100 disable_implicit_threading() 

101 self.assertEqual(numba.get_num_threads(), 1) 

102 

103 # numba re-reads NUMBA_NUM_THREADS every time it reloads its 

104 # configuration, which happens on each jit compilation. Compiling a 

105 # function must not raise because the environment variable and numba's 

106 # recorded thread count disagree. 

107 @numba.njit("float64(float64)") 

108 def add_one(x: float) -> float: 

109 return x + 1.0 

110 

111 self.assertEqual(add_one(1.0), 2.0) 

112 self.assertEqual(numba.get_num_threads(), 1) 

113 

114 @unittest.skipIf(threadpoolctl is None, "threadpoolctl is not available") 

115 def testMissingThreadpoolctlWarning(self): 

116 """Absence of threadpoolctl silently weakens the thread disabling 

117 so it must be reported. 

118 """ 

119 with ( 

120 unittest.mock.patch("lsst.utils.threads.threadpool_limits", None), 

121 self.assertLogs("lsst.utils.threads", "WARNING") as cm, 

122 ): 

123 disable_implicit_threading() 

124 self.assertIn("threadpoolctl", "\n".join(cm.output)) 

125 

126 

127if __name__ == "__main__": 

128 unittest.main()