to_process.py 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266
  1. from __future__ import annotations
  2. __all__ = (
  3. "current_default_process_limiter",
  4. "process_worker",
  5. "run_sync",
  6. )
  7. import os
  8. import pickle
  9. import subprocess
  10. import sys
  11. from collections import deque
  12. from collections.abc import Callable
  13. from importlib.util import module_from_spec, spec_from_file_location
  14. from typing import TypeVar, cast
  15. from ._core._eventloop import current_time, get_async_backend, get_cancelled_exc_class
  16. from ._core._exceptions import BrokenWorkerProcess
  17. from ._core._subprocesses import open_process
  18. from ._core._synchronization import CapacityLimiter
  19. from ._core._tasks import CancelScope, fail_after
  20. from .abc import ByteReceiveStream, ByteSendStream, Process
  21. from .lowlevel import RunVar, checkpoint_if_cancelled
  22. from .streams.buffered import BufferedByteReceiveStream
  23. if sys.version_info >= (3, 11):
  24. from typing import TypeVarTuple, Unpack
  25. else:
  26. from typing_extensions import TypeVarTuple, Unpack
  27. WORKER_MAX_IDLE_TIME = 300 # 5 minutes
  28. T_Retval = TypeVar("T_Retval")
  29. PosArgsT = TypeVarTuple("PosArgsT")
  30. _process_pool_workers: RunVar[set[Process]] = RunVar("_process_pool_workers")
  31. _process_pool_idle_workers: RunVar[deque[tuple[Process, float]]] = RunVar(
  32. "_process_pool_idle_workers"
  33. )
  34. _default_process_limiter: RunVar[CapacityLimiter] = RunVar("_default_process_limiter")
  35. async def run_sync( # type: ignore[return]
  36. func: Callable[[Unpack[PosArgsT]], T_Retval],
  37. *args: Unpack[PosArgsT],
  38. cancellable: bool = False,
  39. limiter: CapacityLimiter | None = None,
  40. ) -> T_Retval:
  41. """
  42. Call the given function with the given arguments in a worker process.
  43. If the ``cancellable`` option is enabled and the task waiting for its completion is
  44. cancelled, the worker process running it will be abruptly terminated using SIGKILL
  45. (or ``terminateProcess()`` on Windows).
  46. :param func: a callable
  47. :param args: positional arguments for the callable
  48. :param cancellable: ``True`` to allow cancellation of the operation while it's
  49. running
  50. :param limiter: capacity limiter to use to limit the total amount of processes
  51. running (if omitted, the default limiter is used)
  52. :raises NoEventLoopError: if no supported asynchronous event loop is running in the
  53. current thread
  54. :return: an awaitable that yields the return value of the function.
  55. """
  56. async def send_raw_command(pickled_cmd: bytes) -> object:
  57. try:
  58. await stdin.send(pickled_cmd)
  59. response = await buffered.receive_until(b"\n", 50)
  60. status, length = response.split(b" ")
  61. if status not in (b"RETURN", b"EXCEPTION"):
  62. raise RuntimeError(
  63. f"Worker process returned unexpected response: {response!r}"
  64. )
  65. pickled_response = await buffered.receive_exactly(int(length))
  66. except BaseException as exc:
  67. workers.discard(process)
  68. try:
  69. process.kill()
  70. with CancelScope(shield=True):
  71. await process.aclose()
  72. except ProcessLookupError:
  73. pass
  74. if isinstance(exc, get_cancelled_exc_class()):
  75. raise
  76. else:
  77. raise BrokenWorkerProcess from exc
  78. retval = pickle.loads(pickled_response)
  79. if status == b"EXCEPTION":
  80. assert isinstance(retval, BaseException)
  81. raise retval
  82. else:
  83. return retval
  84. # First pickle the request before trying to reserve a worker process
  85. await checkpoint_if_cancelled()
  86. request = pickle.dumps(("run", func, args), protocol=pickle.HIGHEST_PROTOCOL)
  87. # If this is the first run in this event loop thread, set up the necessary variables
  88. try:
  89. workers = _process_pool_workers.get()
  90. idle_workers = _process_pool_idle_workers.get()
  91. except LookupError:
  92. workers = set()
  93. idle_workers = deque()
  94. _process_pool_workers.set(workers)
  95. _process_pool_idle_workers.set(idle_workers)
  96. get_async_backend().setup_process_pool_exit_at_shutdown(workers)
  97. async with limiter or current_default_process_limiter():
  98. # Pop processes from the pool (starting from the most recently used) until we
  99. # find one that hasn't exited yet
  100. process: Process
  101. while idle_workers:
  102. process, idle_since = idle_workers.pop()
  103. if process.returncode is None:
  104. stdin = cast(ByteSendStream, process.stdin)
  105. buffered = BufferedByteReceiveStream(
  106. cast(ByteReceiveStream, process.stdout)
  107. )
  108. # Prune any other workers that have been idle for WORKER_MAX_IDLE_TIME
  109. # seconds or longer
  110. now = current_time()
  111. killed_processes: list[Process] = []
  112. while idle_workers:
  113. if now - idle_workers[0][1] < WORKER_MAX_IDLE_TIME:
  114. break
  115. process_to_kill, idle_since = idle_workers.popleft()
  116. process_to_kill.kill()
  117. workers.remove(process_to_kill)
  118. killed_processes.append(process_to_kill)
  119. with CancelScope(shield=True):
  120. for killed_process in killed_processes:
  121. await killed_process.aclose()
  122. break
  123. workers.remove(process)
  124. else:
  125. command = [sys.executable, "-u", "-m", __name__]
  126. process = await open_process(
  127. command, stdin=subprocess.PIPE, stdout=subprocess.PIPE
  128. )
  129. try:
  130. stdin = cast(ByteSendStream, process.stdin)
  131. buffered = BufferedByteReceiveStream(
  132. cast(ByteReceiveStream, process.stdout)
  133. )
  134. with fail_after(20):
  135. message = await buffered.receive(6)
  136. if message != b"READY\n":
  137. raise BrokenWorkerProcess(
  138. f"Worker process returned unexpected response: {message!r}"
  139. )
  140. main_module_path = getattr(sys.modules["__main__"], "__file__", None)
  141. pickled = pickle.dumps(
  142. ("init", sys.path, main_module_path),
  143. protocol=pickle.HIGHEST_PROTOCOL,
  144. )
  145. await send_raw_command(pickled)
  146. except (BrokenWorkerProcess, get_cancelled_exc_class()):
  147. raise
  148. except BaseException as exc:
  149. process.kill()
  150. raise BrokenWorkerProcess(
  151. "Error during worker process initialization"
  152. ) from exc
  153. workers.add(process)
  154. with CancelScope(shield=not cancellable):
  155. try:
  156. return cast(T_Retval, await send_raw_command(request))
  157. finally:
  158. if process in workers:
  159. idle_workers.append((process, current_time()))
  160. def current_default_process_limiter() -> CapacityLimiter:
  161. """
  162. Return the capacity limiter that is used by default to limit the number of worker
  163. processes.
  164. :return: a capacity limiter object
  165. """
  166. try:
  167. return _default_process_limiter.get()
  168. except LookupError:
  169. limiter = CapacityLimiter(os.cpu_count() or 2)
  170. _default_process_limiter.set(limiter)
  171. return limiter
  172. def process_worker() -> None:
  173. # Redirect standard streams to os.devnull so that user code won't interfere with the
  174. # parent-worker communication
  175. stdin = sys.stdin
  176. stdout = sys.stdout
  177. sys.stdin = open(os.devnull)
  178. sys.stdout = open(os.devnull, "w")
  179. stdout.buffer.write(b"READY\n")
  180. while True:
  181. retval = exception = None
  182. try:
  183. command, *args = pickle.load(stdin.buffer)
  184. except EOFError:
  185. return
  186. except BaseException as exc:
  187. exception = exc
  188. else:
  189. if command == "run":
  190. func, args = args
  191. try:
  192. retval = func(*args)
  193. except BaseException as exc:
  194. exception = exc
  195. elif command == "init":
  196. main_module_path: str | None
  197. sys.path, main_module_path = args
  198. del sys.modules["__main__"]
  199. if main_module_path and os.path.isfile(main_module_path):
  200. # Load the parent's main module but as __mp_main__ instead of
  201. # __main__ (like multiprocessing does) to avoid infinite recursion
  202. try:
  203. spec = spec_from_file_location("__mp_main__", main_module_path)
  204. if spec and spec.loader:
  205. main = module_from_spec(spec)
  206. spec.loader.exec_module(main)
  207. sys.modules["__main__"] = main
  208. except BaseException as exc:
  209. exception = exc
  210. try:
  211. if exception is not None:
  212. status = b"EXCEPTION"
  213. pickled = pickle.dumps(exception, pickle.HIGHEST_PROTOCOL)
  214. else:
  215. status = b"RETURN"
  216. pickled = pickle.dumps(retval, pickle.HIGHEST_PROTOCOL)
  217. except BaseException as exc:
  218. exception = exc
  219. status = b"EXCEPTION"
  220. pickled = pickle.dumps(exc, pickle.HIGHEST_PROTOCOL)
  221. stdout.buffer.write(b"%s %d\n" % (status, len(pickled)))
  222. stdout.buffer.write(pickled)
  223. # Respect SIGTERM
  224. if isinstance(exception, SystemExit):
  225. raise exception
  226. if __name__ == "__main__":
  227. process_worker()