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

25 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-02-27 11:49 +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# 

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

13 

14from __future__ import annotations 

15 

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

17 

18import os 

19 

20try: 

21 from threadpoolctl import threadpool_limits 

22except ImportError: 

23 threadpool_limits = None 

24 

25 

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

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

28 

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 ) 

48 

49 for var in envvars: 

50 if override or var not in os.environ: 

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

52 

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" 

58 

59 

60def disable_implicit_threading() -> None: 

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

62 

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) 

71 

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) 

81 

82 # Try to set threads for openblas and openmp 

83 if threadpool_limits is not None: 

84 threadpool_limits(limits=1)