Coverage for src/image_utils/process_utils.py: 56%

1 from functools import partial

2 from multiprocessing import Pool

3 import random

4 import time

5 from pathlib import Path

6 from joblib import dump, load

7 import time

8 import numpy as np

9 from joblib import Parallel, delayed

10 from joblib import Memory

11

12

13 def timer_func(func):

14 def wrap_func(*args, **kwargs):

15 t1 = time.time()

16 result = func(*args, **kwargs)

17 t2 = time.time()

18 print(f'Function {func.__name__!r} executed in {(t2-t1):.4f}s')

19 return result

20 return wrap_func

21

22

23 def split_range(a, n):

24 k, m = divmod(len(a), n)

25 return (a[i * k + min(i, m): (i + 1) * k + min(i + 1, m)] for i in range(n))

26

27

28 def parallelize(func=None, num_processes=1, shuffle=False, custom_chunking=False, use_joblib=False):

29 def wrapper(iterable, **kwargs):

30 if use_joblib:

31 executor = Parallel(n_jobs=num_processes, max_nbytes=1e6)

32 tasks = (delayed(func)(iterable, idx, **kwargs) for idx in range(len(iterable)))

33 return executor(tasks)

34 else:

35 with Pool(num_processes) as p:

36 num_arr = list(range(len(iterable)))

37 if shuffle:

38 random.shuffle(num_arr)

39 for _ in p.imap_unordered(

40 partial(func, **kwargs),

41 split_range(num_arr, num_processes) if custom_chunking else iterable

42 ):

43 pass

44

45 return wrapper