Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/to_process.py: 19%

1from __future__ import annotations 

2 

3__all__ = ( 

4 "current_default_process_limiter", 

5 "process_worker", 

6 "run_sync", 

7) 

8 

9import os 

10import pickle 

11import runpy 

12import subprocess 

13import sys 

14from collections import deque 

15from collections.abc import Callable 

16from types import ModuleType 

17from typing import TypeVar, cast 

18 

19from ._core._eventloop import current_time, get_async_backend, get_cancelled_exc_class 

20from ._core._exceptions import BrokenWorkerProcess 

21from ._core._subprocesses import open_process 

22from ._core._synchronization import CapacityLimiter 

23from ._core._tasks import CancelScope, fail_after 

24from .abc import ByteReceiveStream, ByteSendStream, Process 

25from .lowlevel import RunVar, checkpoint_if_cancelled 

26from .streams.buffered import BufferedByteReceiveStream 

27 

28if sys.version_info >= (3, 11): 

29 from typing import TypeVarTuple, Unpack 

30else: 

31 from typing_extensions import TypeVarTuple, Unpack 

32 

33WORKER_MAX_IDLE_TIME = 300 # 5 minutes 

34 

35T_Retval = TypeVar("T_Retval") 

36PosArgsT = TypeVarTuple("PosArgsT") 

37 

38_process_pool_workers: RunVar[set[Process]] = RunVar("_process_pool_workers") 

39_process_pool_idle_workers: RunVar[deque[tuple[Process, float]]] = RunVar( 

40 "_process_pool_idle_workers" 

41) 

42_default_process_limiter: RunVar[CapacityLimiter] = RunVar("_default_process_limiter") 

43 

44 

45async def run_sync( # type: ignore[return] 

46 func: Callable[[Unpack[PosArgsT]], T_Retval], 

47 *args: Unpack[PosArgsT], 

48 cancellable: bool = False, 

49 limiter: CapacityLimiter | None = None, 

50) -> T_Retval: 

51 """ 

52 Call the given function with the given arguments in a worker process. 

53 

54 If the ``cancellable`` option is enabled and the task waiting for its completion is 

55 cancelled, the worker process running it will be abruptly terminated using SIGKILL 

56 (or ``terminateProcess()`` on Windows). 

57 

58 :param func: a callable 

59 :param args: positional arguments for the callable 

60 :param cancellable: ``True`` to allow cancellation of the operation while it's 

61 running 

62 :param limiter: capacity limiter to use to limit the total amount of processes 

63 running (if omitted, the default limiter is used) 

64 :raises NoEventLoopError: if no supported asynchronous event loop is running in the 

65 current thread 

66 :return: an awaitable that yields the return value of the function. 

67 

68 """ 

69 

70 async def send_raw_command(pickled_cmd: bytes) -> object: 

71 try: 

72 await stdin.send(pickled_cmd) 

73 response = await buffered.receive_until(b"\n", 50) 

74 status, length = response.split(b" ") 

75 if status not in (b"RETURN", b"EXCEPTION"): 

76 raise RuntimeError( 

77 f"Worker process returned unexpected response: {response!r}" 

78 ) 

79 

80 pickled_response = await buffered.receive_exactly(int(length)) 

81 except BaseException as exc: 

82 workers.discard(process) 

83 try: 

84 process.kill() 

85 with CancelScope(shield=True): 

86 await process.aclose() 

87 except ProcessLookupError: 

88 pass 

89 

90 if isinstance(exc, get_cancelled_exc_class()): 

91 raise 

92 else: 

93 raise BrokenWorkerProcess from exc 

94 

95 retval = pickle.loads(pickled_response) 

96 if status == b"EXCEPTION": 

97 assert isinstance(retval, BaseException) 

98 raise retval 

99 else: 

100 return retval 

101 

102 # First pickle the request before trying to reserve a worker process 

103 await checkpoint_if_cancelled() 

104 request = pickle.dumps(("run", func, args), protocol=pickle.HIGHEST_PROTOCOL) 

105 

106 # If this is the first run in this event loop thread, set up the necessary variables 

107 try: 

108 workers = _process_pool_workers.get() 

109 idle_workers = _process_pool_idle_workers.get() 

110 except LookupError: 

111 workers = set() 

112 idle_workers = deque() 

113 _process_pool_workers.set(workers) 

114 _process_pool_idle_workers.set(idle_workers) 

115 get_async_backend().setup_process_pool_exit_at_shutdown(workers) 

116 

117 async with limiter or current_default_process_limiter(): 

118 # Pop processes from the pool (starting from the most recently used) until we 

119 # find one that hasn't exited yet 

120 process: Process 

121 while idle_workers: 

122 process, _idle_since = idle_workers.pop() 

123 if process.returncode is None: 

124 stdin = cast(ByteSendStream, process.stdin) 

125 buffered = BufferedByteReceiveStream( 

126 cast(ByteReceiveStream, process.stdout) 

127 ) 

128 

129 # Prune any other workers that have been idle for WORKER_MAX_IDLE_TIME 

130 # seconds or longer 

131 now = current_time() 

132 killed_processes: list[Process] = [] 

133 while idle_workers: 

134 if now - idle_workers[0][1] < WORKER_MAX_IDLE_TIME: 

135 break 

136 

137 process_to_kill, _idle_since = idle_workers.popleft() 

138 process_to_kill.kill() 

139 workers.remove(process_to_kill) 

140 killed_processes.append(process_to_kill) 

141 

142 with CancelScope(shield=True): 

143 for killed_process in killed_processes: 

144 await killed_process.aclose() 

145 

146 break 

147 

148 workers.remove(process) 

149 else: 

150 command = [sys.executable, "-u", "-m", __name__] 

151 process = await open_process( 

152 command, stdin=subprocess.PIPE, stdout=subprocess.PIPE 

153 ) 

154 try: 

155 stdin = cast(ByteSendStream, process.stdin) 

156 buffered = BufferedByteReceiveStream( 

157 cast(ByteReceiveStream, process.stdout) 

158 ) 

159 with fail_after(20): 

160 message = await buffered.receive(6) 

161 

162 if message != b"READY\n": 

163 raise BrokenWorkerProcess( 

164 f"Worker process returned unexpected response: {message!r}" 

165 ) 

166 

167 main_module_path = getattr(sys.modules["__main__"], "__file__", None) 

168 pickled = pickle.dumps( 

169 ("init", sys.path, main_module_path), 

170 protocol=pickle.HIGHEST_PROTOCOL, 

171 ) 

172 await send_raw_command(pickled) 

173 except (BrokenWorkerProcess, get_cancelled_exc_class()): 

174 raise 

175 except BaseException as exc: 

176 process.kill() 

177 raise BrokenWorkerProcess( 

178 "Error during worker process initialization" 

179 ) from exc 

180 

181 workers.add(process) 

182 

183 with CancelScope(shield=not cancellable): 

184 try: 

185 return cast(T_Retval, await send_raw_command(request)) 

186 finally: 

187 if process in workers: 

188 idle_workers.append((process, current_time())) 

189 

190 

191def current_default_process_limiter() -> CapacityLimiter: 

192 """ 

193 Return the capacity limiter that is used by default to limit the number of worker 

194 processes. 

195 

196 :return: a capacity limiter object 

197 

198 """ 

199 try: 

200 return _default_process_limiter.get() 

201 except LookupError: 

202 limiter = CapacityLimiter(os.cpu_count() or 2) 

203 _default_process_limiter.set(limiter) 

204 return limiter 

205 

206 

207def process_worker() -> None: 

208 # Redirect standard streams to os.devnull so that user code won't interfere with the 

209 # parent-worker communication 

210 stdin = sys.stdin 

211 stdout = sys.stdout 

212 sys.stdin = open(os.devnull) 

213 sys.stdout = open(os.devnull, "w") 

214 sys.stderr = open(os.devnull, "w") 

215 

216 stdout.buffer.write(b"READY\n") 

217 while True: 

218 retval = exception = None 

219 try: 

220 command, *args = pickle.load(stdin.buffer) 

221 except EOFError: 

222 return 

223 except BaseException as exc: 

224 exception = exc 

225 else: 

226 if command == "run": 

227 func, args = args 

228 try: 

229 retval = func(*args) 

230 except BaseException as exc: 

231 exception = exc 

232 elif command == "init": 

233 main_module_path: str | None 

234 sys.path, main_module_path = args 

235 del sys.modules["__main__"] 

236 if main_module_path and os.path.isfile(main_module_path): 

237 # Load the parent's main module but as __mp_main__ instead of 

238 # __main__ (like multiprocessing does) to avoid infinite recursion 

239 try: 

240 main = ModuleType("__mp_main__") 

241 main_content = runpy.run_path( 

242 main_module_path, run_name="__mp_main__" 

243 ) 

244 main.__dict__.update(main_content) 

245 sys.modules["__main__"] = sys.modules["__mp_main__"] = main 

246 except BaseException as exc: 

247 exception = exc 

248 try: 

249 if exception is not None: 

250 status = b"EXCEPTION" 

251 pickled = pickle.dumps(exception, pickle.HIGHEST_PROTOCOL) 

252 else: 

253 status = b"RETURN" 

254 pickled = pickle.dumps(retval, pickle.HIGHEST_PROTOCOL) 

255 except BaseException as exc: 

256 exception = exc 

257 status = b"EXCEPTION" 

258 pickled = pickle.dumps(exc, pickle.HIGHEST_PROTOCOL) 

259 

260 stdout.buffer.write(b"%s %d\n" % (status, len(pickled))) 

261 stdout.buffer.write(pickled) 

262 

263 # Respect SIGTERM 

264 if isinstance(exception, SystemExit): 

265 raise exception 

266 

267 

268if __name__ == "__main__": 

269 process_worker()