Coverage for local_installation_linux/mumott/core/numba_setup.py: 80%

1import logging 

2import warnings 

3import os 

4from multiprocessing import cpu_count 

5import numba 

6from numba.core.errors import NumbaPerformanceWarning 

7 

8logger = logging.getLogger(__name__) 

9 

10 

11def numba_setup() -> None: 

12 """This function sets defaults for environmental variables and numba 

13 parameters to achieve a balanced performance for most situations. 

14 Advanced users may want to customize these parameters. 

15 """ 

16 

17 # This is a relatively useless and very verbose and frequent warning. 

18 warnings.filterwarnings('ignore', category=NumbaPerformanceWarning) 

19 

20 try: 

21 number_of_cpus = len(os.sched_getaffinity(os.getpid())) 

22 except AttributeError: 

23 logger.info('Using multiprocessing.cpu_count() to determine the number of availible CPUs.') 

24 number_of_cpus = cpu_count() 

25 if number_of_cpus >= 8: 25 ↛ 28line 25 didn't jump to line 28, because the condition on line 25 was never false

26 number_of_threads = 8 

27 else: 

28 number_of_threads = number_of_cpus 

29 logger.info(f'Setting the number of threads to {number_of_threads}. ' 

30 'If your physical cores are fewer than this number, ' 

31 'you may want to use numba.set_num_threads(n), ' 

32 'and os.environ["OPENBLAS_NUM_THREADS"] = f"{n}" ' 

33 'to set the number of threads to the number of ' 

34 'physical cores n.') 

35 numba.set_num_threads(number_of_threads) 

36 os.environ['OPENBLAS_NUM_THREADS'] = f'{number_of_threads}' 

37 

38 # numba.cuda has a very high output of `INFO`-level messages. 

39 numba_logger = logging.getLogger('numba') 

40 numba_logger.setLevel(logging.WARNING) 

41 logger.info('Setting numba log level to WARNING.')