Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/unblob/pool.py: 32%

1import abc 

2import atexit 

3import contextlib 

4import fcntl 

5import multiprocessing as mp 

6import multiprocessing.connection 

7import os 

8import queue 

9import signal 

10import sys 

11import threading 

12from collections.abc import Callable 

13from multiprocessing.queues import JoinableQueue, SimpleQueue 

14from typing import Any 

15 

16from .logging import multiprocessing_breakpoint 

17 

18 

19class PoolBase(abc.ABC): 

20 @abc.abstractmethod 

21 def submit(self, args): 

22 pass 

23 

24 @abc.abstractmethod 

25 def process_until_done(self): 

26 pass 

27 

28 def start(self): 

29 pass 

30 

31 def close(self, *, immediate=False): 

32 pass 

33 

34 def __enter__(self): 

35 self.start() 

36 return self 

37 

38 def __exit__(self, exc_type, _exc_value, _tb): 

39 self.close(immediate=exc_type is not None) 

40 

41 

42class Queue(JoinableQueue): 

43 def is_empty(self) -> bool: 

44 """Check if all ``task_done`` has been called for all items. 

45 

46 Based on ``multiprocessing.JoinableQueue.join``. 

47 """ 

48 with self._cond: # type: ignore 

49 return self._unfinished_tasks._semlock._is_zero() # type: ignore # noqa: SLF001 

50 

51 

52class ResultQueue(SimpleQueue): 

53 @property 

54 def reader(self) -> multiprocessing.connection.Connection: 

55 return self._reader # type: ignore 

56 

57 

58class _Sentinel: 

59 pass 

60 

61 

62_SENTINEL = _Sentinel 

63 

64 

65class WorkerDiedError(RuntimeError): 

66 pass 

67 

68 

69def _worker_process(handler, input_, output, lifeline_worker_side, lifeline_host_side): 

70 # Creates a new process group, making sure no signals are 

71 # propagated from the main process to the worker processes. 

72 os.setpgrp() 

73 

74 # Restore default signal handlers, otherwise workers would inherit 

75 # them from main process. When used as a library, the hosting app 

76 # is free to set-up its own signal handlers. 

77 signal.signal(signal.SIGTERM, signal.SIG_DFL) 

78 signal.signal(signal.SIGINT, signal.SIG_DFL) 

79 

80 os.close( # forked processes inherit open files, we don't need the host FD 

81 lifeline_host_side 

82 ) 

83 

84 def _exit_on_parent_death(): 

85 os.read(lifeline_worker_side, 1) 

86 # We cannot really do anything about this, best to reliably 

87 # abort the process 

88 os._exit(1) 

89 

90 parent_liveness_monitor = threading.Thread( 

91 target=_exit_on_parent_death, daemon=True 

92 ) 

93 parent_liveness_monitor.start() 

94 

95 sys.breakpointhook = multiprocessing_breakpoint 

96 while (args := input_.get()) is not _SENTINEL: 

97 result = handler(args) 

98 output.put(result) 

99 output.put(_SENTINEL) 

100 

101 

102class MultiPool(PoolBase): 

103 def __init__( 

104 self, 

105 process_num: int, 

106 handler: Callable[[Any], Any], 

107 *, 

108 result_callback: Callable[["MultiPool", Any], Any], 

109 ): 

110 if process_num <= 0: 

111 raise ValueError("At process_num must be greater than 0") 

112 

113 self._running = False 

114 self._result_callback = result_callback 

115 mp_context = mp.get_context("fork") 

116 self._input = Queue(ctx=mp_context) 

117 self._input.cancel_join_thread() 

118 self._output = ResultQueue(ctx=mp_context) 

119 # see search results for "death-pipe" or "forkfd concept" 

120 (self._lifeline_worker_side, self._lifeline_host_side) = os.pipe() 

121 fcntl.fcntl(self._lifeline_host_side, fcntl.F_SETFD, fcntl.FD_CLOEXEC) 

122 fcntl.fcntl(self._lifeline_worker_side, fcntl.F_SETFD, fcntl.FD_CLOEXEC) 

123 

124 self._procs = [ 

125 mp_context.Process( 

126 target=_worker_process, 

127 args=( 

128 handler, 

129 self._input, 

130 self._output, 

131 self._lifeline_worker_side, 

132 self._lifeline_host_side, 

133 ), 

134 ) 

135 for _ in range(process_num) 

136 ] 

137 self._tid = threading.get_native_id() 

138 

139 def start(self): 

140 self._running = True 

141 for p in self._procs: 

142 p.start() 

143 # We are the host process, we don't need this anymore. 

144 # Had to keep the file alive until inherited by the forked subprocess 

145 os.close(self._lifeline_worker_side) 

146 atexit.register(self._close_immediate) 

147 

148 def _any_worker_exited(self) -> bool: 

149 sentinels = [p.sentinel for p in self._procs] 

150 return bool(multiprocessing.connection.wait(sentinels, timeout=0)) 

151 

152 def close(self, *, immediate=False): 

153 if not self._running: 

154 return 

155 self._running = False 

156 atexit.unregister(self._close_immediate) 

157 immediate = immediate or self._any_worker_exited() 

158 

159 termination_exception = None 

160 if not immediate: 

161 try: 

162 self._clear_input_queue() 

163 self._request_workers_to_quit() 

164 self._clear_output_queue() 

165 except BaseException as exc: 

166 termination_exception = exc 

167 immediate = True 

168 

169 if immediate: 

170 self._terminate_workers() 

171 

172 self._wait_for_workers_to_quit() 

173 

174 # closing this FD any sooner would cause workers to abort 

175 # immediately, should close only after workers quit 

176 os.close(self._lifeline_host_side) 

177 

178 if termination_exception: 

179 raise termination_exception 

180 

181 def _close_immediate(self): 

182 self.close(immediate=True) 

183 

184 def _terminate_workers(self): 

185 for proc in self._procs: 

186 proc.terminate() 

187 

188 self._input.close() 

189 self._output.close() 

190 

191 def _clear_input_queue(self): 

192 with contextlib.suppress(queue.Empty): 

193 while True: 

194 self._input.get_nowait() 

195 

196 def _request_workers_to_quit(self): 

197 for _ in self._procs: 

198 self._input.put(_SENTINEL) 

199 self._input.close() 

200 

201 def _clear_output_queue(self): 

202 alive = {p.sentinel: p for p in self._procs if p.exitcode is None} 

203 while alive: 

204 ready = multiprocessing.connection.wait([self._output.reader, *alive]) 

205 for fd in ready: 

206 alive.pop(fd, None) # type: ignore[arg-type] 

207 if self._output.reader in ready: 

208 self._output.get() 

209 

210 def _wait_for_workers_to_quit(self): 

211 for p in self._procs: 

212 p.join() 

213 

214 def submit(self, args): 

215 if threading.get_native_id() != self._tid: 

216 raise RuntimeError( 

217 "Submit can only be called from the same " 

218 "thread/process where the pool is created" 

219 ) 

220 self._input.put(args) 

221 

222 def _check_worker_deaths(self, sentinels, ready): 

223 for fd in ready: 

224 if fd not in sentinels: 

225 continue 

226 proc = sentinels.pop(fd) 

227 if proc.exitcode != 0: 

228 exitcode = proc.exitcode 

229 if exitcode is not None and exitcode < 0: 

230 reason = f"killed by signal {-exitcode}" 

231 else: 

232 reason = f"exited with code {exitcode}" 

233 raise WorkerDiedError( 

234 f"Worker process {proc.pid} exited unexpectedly ({reason})" 

235 ) 

236 

237 def process_until_done(self): 

238 sentinels = {p.sentinel: p for p in self._procs} 

239 with contextlib.suppress(EOFError): 

240 while not self._input.is_empty(): 

241 ready = multiprocessing.connection.wait( 

242 [self._output.reader, *sentinels] 

243 ) 

244 self._check_worker_deaths(sentinels, ready) 

245 if self._output.reader in ready: 

246 result = self._output.get() 

247 self._result_callback(self, result) 

248 self._input.task_done() 

249 

250 

251class SinglePool(PoolBase): 

252 def __init__(self, handler, *, result_callback): 

253 self._handler = handler 

254 self._result_callback = result_callback 

255 

256 def submit(self, args): 

257 result = self._handler(args) 

258 self._result_callback(self, result) 

259 

260 def process_until_done(self): 

261 pass 

262 

263 

264def make_pool(process_num, handler, result_callback) -> SinglePool | MultiPool: 

265 if process_num == 1: 

266 return SinglePool(handler=handler, result_callback=result_callback) 

267 

268 return MultiPool( 

269 process_num=process_num, 

270 handler=handler, 

271 result_callback=result_callback, 

272 )