Limiting the number of active CPU cores with Numba

user_16234612 0
Viewed 1132

I'm trying to limit the number of active CPU cores when running a funtion with the decorator @njit(parallel=True) from the Numba library.

So far I've done something like this (example):

from numba import njit, prange, set_num_threads

import numpy as np


@njit(parallel=True)
def func():

    for i in prange(int(1e10)):

        print(np.cos(i))


if __name__ == '__main__':
    
    num_threads = input(f'Enter max number of threads to use: ')

    set_num_threads(int(num_threads))

    func()

but even when setting the number of threads to 1 a quick check with Task Manager shows that the workload is still being shared by all the cores (4 in my case).

In the following screenshot it is possible to see this. The red circles show the time of starting and interruption of the program.

Task Manager Screenshot

Is there a better way of limiting the number of active cores?

numba parallel-processing python
1 Answers

You are stressing your cpu cores by using the print function and displaying them in the console, not because of the calculation. In this example I accumulate the values before printing only once. In this way, you can see more easily the real cpu usage of the code.

from numba import njit, prange, set_num_threads

import numpy as np


@njit(parallel=True)
def func():

    k = 0
    for i in prange(int(1e10)):
        k = k + np.cos(i)

    print(k)



if __name__ == '__main__':

    num_threads = input(f'Enter max number of threads to use: ')

    set_num_threads(int(num_threads))
user_20253190
Related