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

22 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-02 14:10 +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 

13 

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

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 

54def disable_implicit_threading() -> None: 

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

56 

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) 

65 

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) 

75 

76 # Try to set threads for openblas and openmp 

77 if threadpool_limits is not None: 

78 threadpool_limits(limits=1)