Coverage for cosmolayer/store/parallel.py: 91%

1"""Threaded execution over molecule-sized chunks of segment- or atom-level 

2arrays. 

3 

4Each pass splits *molecules* (not segments or atoms) into contiguous 

5ranges so threads own disjoint molecules and need no locking. 

6""" 

7 

8import os 

9from collections.abc import Callable, Iterator 

10from concurrent.futures import ThreadPoolExecutor, as_completed 

11from contextlib import nullcontext 

12 

13import threadpoolctl 

14 

15 

16def resolve_num_threads(num_threads: int | None) -> int: 

17 """Resolve a possibly-unset thread count to a concrete positive value. 

18 

19 Parameters 

20 ---------- 

21 num_threads : int | None 

22 Requested number of threads, or None to mean "every available CPU 

23 core". 

24 

25 Returns 

26 ------- 

27 int 

28 ``num_threads`` if given, else ``os.cpu_count()`` (falling back to 

29 1 if that returns None, e.g. in a constrained container). 

30 """ 

31 if num_threads is not None: 

32 return num_threads 

33 return os.cpu_count() or 1 

34 

35 

36def molecule_chunks(num_items: int, num_threads: int) -> Iterator[tuple[int, int]]: 

37 """Split ``range(num_items)`` into up to ``num_threads`` contiguous, 

38 roughly equal chunks. 

39 

40 Parameters 

41 ---------- 

42 num_items : int 

43 Total number of molecules (or other items) to split. 

44 num_threads : int 

45 Number of chunks to aim for. Fewer are yielded if 

46 ``num_items < num_threads``. 

47 

48 Yields 

49 ------ 

50 tuple[int, int] 

51 ``(start, stop)`` bounds of each chunk, usable as a Python slice. 

52 """ 

53 if num_items == 0: 

54 return 

55 chunk_size = (num_items + num_threads - 1) // num_threads 

56 for start in range(0, num_items, chunk_size): 

57 yield start, min(start + chunk_size, num_items) 

58 

59 

60def run_in_threads( 

61 fn: Callable[[int, int], None], 

62 num_items: int, 

63 *, 

64 num_threads: int | None = None, 

65 limit_blas: bool = False, 

66) -> None: 

67 """Run ``fn(start, stop)`` once per molecule chunk, across threads. 

68 

69 Chunks cover disjoint molecules, so ``fn`` may write only to its own 

70 ``[start, stop)`` range. 

71 

72 Parameters 

73 ---------- 

74 fn : Callable[[int, int], None] 

75 Called once per chunk with that chunk's molecule bounds. 

76 Exceptions in a thread are re-raised in the caller. 

77 num_items : int 

78 Number of molecules to split across threads. 

79 num_threads : int | None, optional 

80 Thread count. ``None`` (default) uses every CPU core. 

81 limit_blas : bool, optional 

82 If True, cap BLAS to 1 thread for the duration of the call, so 

83 numpy matrix products do not oversubscribe the CPU. 

84 """ 

85 num_threads = resolve_num_threads(num_threads) 

86 limiter = threadpoolctl.threadpool_limits(limits=1) if limit_blas else nullcontext() 

87 with limiter, ThreadPoolExecutor(max_workers=num_threads) as executor: 

88 futures = [ 

89 executor.submit(fn, start, stop) 

90 for start, stop in molecule_chunks(num_items, num_threads) 

91 ] 

92 for future in as_completed(futures): 

93 future.result()