NumPy parallelism

numpy-parallelism.py

#!/usr/bin/env python3

import sys
import timeit
from concurrent.futures import ThreadPoolExecutor, wait

import numpy as np


N_THREADS = int(sys.argv[1])
ITERATIONS = 100


def work(arr: np.array) -> int:
    arr = arr * 2
    result = arr.sum()
    return result


def main() -> None:
    arr = np.arange(1_000_000, dtype=np.int64)
    with ThreadPoolExecutor(max_workers=N_THREADS) as executor:
        futures = [
            executor.submit(work, arr.copy())
            for _ in range(N_THREADS)
        ]
        wait(futures)


if __name__ == '__main__':
    elapsed = timeit.timeit('main()', number=ITERATIONS, globals=globals())
    print(
        'Took', elapsed, 'seconds to run with',
        N_THREADS, 'threads and', ITERATIONS, 'iterations.')

There's an unnecessary copy there, here's the improved code:

#!/usr/bin/env python3

import sys
import timeit
from concurrent.futures import ThreadPoolExecutor, wait

import numpy as np


N_THREADS = int(sys.argv[1])
ITERATIONS = 100


def work(arr: np.array) -> int:
    arr *= 2
    result = arr.sum()
    return result


def main() -> None:
    arr = np.arange(1_000_000, dtype=np.int64)
    with ThreadPoolExecutor(max_workers=N_THREADS) as executor:
        futures = [
            executor.submit(work, arr)
            for _ in range(N_THREADS)
        ]
        wait(futures)


if __name__ == '__main__':
    elapsed = timeit.timeit('main()', number=ITERATIONS, globals=globals())
    print(
        'Took', elapsed, 'seconds to run with',
        N_THREADS, 'threads and', ITERATIONS, 'iterations.')

Now, there's a new problem, I changed arr = arr * 2 for arr *= 2 which is fine with Python objects, but not with NumPy.