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

34 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-19 16:15 -0700

1from functools import partial 

2from multiprocessing import Pool 

3import random 

4import time 

5from pathlib import Path 

6from joblib import dump, load 

7import time 

8import numpy as np 

9from joblib import Parallel, delayed 

10from joblib import Memory 

11 

12 

13def 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 

23def 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 

28def 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