Coverage for src / monte_neo / utils / parallel.py: 75%

1"""Parallel processing utilities.""" 

2 

3from __future__ import annotations 

4 

5import os 

6from collections.abc import Callable, Iterable 

7from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor 

8from typing import Any 

9 

10from monte_neo.utils.logger import get_logger 

11 

12logger = get_logger(__name__) 

13 

14 

15class ParallelExecutor: 

16 """Parallel execution helper.""" 

17 

18 def __init__( 

19 self, 

20 n_workers: int | None = None, 

21 use_processes: bool = True, 

22 initializer: Callable | None = None, 

23 initargs: tuple = (), 

24 ) -> None: 

25 """Initialize executor. 

26 

27 Args: 

28 n_workers: Number of workers (None = CPU count). 

29 use_processes: Use processes (True) or threads (False). 

30 initializer: Function to initialize each worker. 

31 initargs: Arguments for initializer. 

32 """ 

33 self.n_workers = n_workers or os.cpu_count() or 4 

34 self.use_processes = use_processes 

35 self.initializer = initializer 

36 self.initargs = initargs 

37 self._pool: ProcessPoolExecutor | ThreadPoolExecutor | None = None 

38 

39 def __enter__(self) -> ParallelExecutor: 

40 """Context manager entry.""" 

41 if self._pool is None: 

42 executor_cls = ( 

43 ProcessPoolExecutor if self.use_processes else ThreadPoolExecutor 

44 ) 

45 kwargs: dict[str, Any] = {"max_workers": self.n_workers} 

46 if self.initializer: 

47 kwargs["initializer"] = self.initializer 

48 kwargs["initargs"] = self.initargs 

49 

50 self._pool = executor_cls(**kwargs) 

51 return self 

52 

53 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 

54 """Context manager exit.""" 

55 if self._pool: 

56 self._shutdown_requested = True 

57 # Cancel all pending futures if possible 

58 if hasattr(self._pool, "_pending_work_items"): # ProcessPoolExecutor internal 

59 try: 

60 for future in list(self._pool._pending_work_items.values()): 

61 future.cancel() 

62 except Exception: 

63 pass 

64 

65 # Use wait=False to avoid hanging on exit 

66 # cancel_futures=True is supported in Python 3.9+ 

67 try: 

68 self._pool.shutdown(wait=False, cancel_futures=True) 

69 except TypeError: 

70 # Fallback for older Python versions 

71 self._pool.shutdown(wait=False) 

72 

73 self._pool = None 

74 

75 def map( 

76 self, 

77 func: Callable, 

78 items: Iterable, 

79 show_progress: bool = False, 

80 ) -> list[Any]: 

81 """Map function over items in parallel. 

82 

83 Args: 

84 func: Function to apply. 

85 items: Items to process. 

86 show_progress: Show progress (requires tqdm). 

87 

88 Returns: 

89 List of results. 

90 """ 

91 items = list(items) 

92 

93 if len(items) == 0: 

94 return [] 

95 

96 # Reduce overhead for small batches or single worker 

97 if len(items) == 1 or self.n_workers == 1: 

98 return [func(item) for item in items] 

99 

100 executor_cls = ProcessPoolExecutor if self.use_processes else ThreadPoolExecutor 

101 results = [] 

102 

103 executor = self._pool 

104 is_temp_pool = False 

105 if executor is None: 

106 executor = executor_cls(max_workers=self.n_workers) 

107 is_temp_pool = True 

108 

109 try: 

110 # Submit tasks and store futures 

111 futures_map = {executor.submit(func, item): i for i, item in enumerate(items)} 

112 pending = set(futures_map.keys()) 

113 

114 while pending: 

115 if getattr(self, "_shutdown_requested", False): 

116 # Cancel all remaining if shutdown requested 

117 for f in pending: 

118 f.cancel() 

119 break 

120 

121 # Wait for some futures to complete with a small timeout to allow checking _shutdown_requested 

122 done, pending = as_completed_with_timeout(pending, timeout=0.1) 

123 

124 for future in done: 

125 idx = futures_map[future] 

126 try: 

127 result = future.result() 

128 results.append((idx, result)) 

129 except Exception as e: 

130 # Only log if not a cancellation/shutdown error 

131 if not getattr(self, "_shutdown_requested", False): 

132 logger.error(f"Error processing item {idx}: {e}") 

133 results.append((idx, None)) 

134 

135 except (KeyboardInterrupt, SystemExit): 

136 self._shutdown_requested = True 

137 # Kill workers immediately 

138 if executor: 

139 try: 

140 executor.shutdown(wait=False, cancel_futures=True) 

141 except (TypeError, Exception): 

142 executor.shutdown(wait=False) 

143 raise 

144 finally: 

145 # Clean up properly if it was a temp pool 

146 if is_temp_pool and executor: 

147 try: 

148 executor.shutdown(wait=False, cancel_futures=True) 

149 except (TypeError, Exception): 

150 executor.shutdown(wait=False) 

151 

152 # Sort by original order 

153 results.sort(key=lambda x: x[0]) 

154 return [r[1] for r in results] 

155 

156 def starmap( 

157 self, 

158 func: Callable, 

159 args_list: Iterable[tuple], 

160 ) -> list[Any]: 

161 """Starmap function over argument tuples. 

162 

163 Args: 

164 func: Function to apply. 

165 args_list: List of argument tuples. 

166 

167 Returns: 

168 List of results. 

169 """ 

170 

171 def wrapper(args: Any) -> Any: 

172 return func(*args) 

173 

174 return self.map(wrapper, args_list) 

175 

176 def map_reduce( 

177 self, 

178 map_func: Callable, 

179 reduce_func: Callable, 

180 items: Iterable, 

181 initial: Any = None, 

182 ) -> Any: 

183 """Map-reduce pattern. 

184 

185 Args: 

186 map_func: Function to apply to each item. 

187 reduce_func: Function to combine results. 

188 items: Items to process. 

189 initial: Initial value for reduction. 

190 

191 Returns: 

192 Reduced result. 

193 """ 

194 mapped = self.map(map_func, items) 

195 

196 result = initial 

197 for item in mapped: 

198 if item is not None: 

199 if result is None: 

200 result = item 

201 else: 

202 result = reduce_func(result, item) 

203 

204 return result 

205 

206 

207def as_completed_with_timeout(fs, timeout=None): 

208 """Wait for some futures to complete with a timeout.""" 

209 from concurrent.futures import FIRST_COMPLETED, wait 

210 

211 done_set = wait(fs, timeout=timeout, return_when=FIRST_COMPLETED).done 

212 return done_set, fs - done_set