File size: 6,624 Bytes
06555b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 |
from __future__ import annotations
import atexit
import os
import pickle
import sys
from collections import deque
from collections.abc import Callable
from textwrap import dedent
from typing import Any, Final, TypeVar
from . import current_time, to_thread
from ._core._exceptions import BrokenWorkerIntepreter
from ._core._synchronization import CapacityLimiter
from .lowlevel import RunVar
if sys.version_info >= (3, 11):
from typing import TypeVarTuple, Unpack
else:
from typing_extensions import TypeVarTuple, Unpack
UNBOUND: Final = 2 # I have no clue how this works, but it was used in the stdlib
FMT_UNPICKLED: Final = 0
FMT_PICKLED: Final = 1
DEFAULT_CPU_COUNT: Final = 8 # this is just an arbitrarily selected value
MAX_WORKER_IDLE_TIME = (
30 # seconds a subinterpreter can be idle before becoming eligible for pruning
)
T_Retval = TypeVar("T_Retval")
PosArgsT = TypeVarTuple("PosArgsT")
_idle_workers = RunVar[deque["Worker"]]("_available_workers")
_default_interpreter_limiter = RunVar[CapacityLimiter]("_default_interpreter_limiter")
class Worker:
_run_func = compile(
dedent("""
import _interpqueues as queues
import _interpreters as interpreters
from pickle import loads, dumps, HIGHEST_PROTOCOL
item = queues.get(queue_id)[0]
try:
func, args = loads(item)
retval = func(*args)
except BaseException as exc:
is_exception = True
retval = exc
else:
is_exception = False
try:
queues.put(queue_id, (retval, is_exception), FMT_UNPICKLED, UNBOUND)
except interpreters.NotShareableError:
retval = dumps(retval, HIGHEST_PROTOCOL)
queues.put(queue_id, (retval, is_exception), FMT_PICKLED, UNBOUND)
"""),
"<string>",
"exec",
)
last_used: float = 0
_initialized: bool = False
_interpreter_id: int
_queue_id: int
def initialize(self) -> None:
import _interpqueues as queues
import _interpreters as interpreters
self._interpreter_id = interpreters.create()
self._queue_id = queues.create(2, FMT_UNPICKLED, UNBOUND) # type: ignore[call-arg]
self._initialized = True
interpreters.set___main___attrs(
self._interpreter_id,
{
"queue_id": self._queue_id,
"FMT_PICKLED": FMT_PICKLED,
"FMT_UNPICKLED": FMT_UNPICKLED,
"UNBOUND": UNBOUND,
},
)
def destroy(self) -> None:
import _interpqueues as queues
import _interpreters as interpreters
if self._initialized:
interpreters.destroy(self._interpreter_id)
queues.destroy(self._queue_id)
def _call(
self,
func: Callable[..., T_Retval],
args: tuple[Any],
) -> tuple[Any, bool]:
import _interpqueues as queues
import _interpreters as interpreters
if not self._initialized:
self.initialize()
payload = pickle.dumps((func, args), pickle.HIGHEST_PROTOCOL)
queues.put(self._queue_id, payload, FMT_PICKLED, UNBOUND) # type: ignore[call-arg]
res: Any
is_exception: bool
if exc_info := interpreters.exec(self._interpreter_id, self._run_func): # type: ignore[func-returns-value,arg-type]
raise BrokenWorkerIntepreter(exc_info)
(res, is_exception), fmt = queues.get(self._queue_id)[:2]
if fmt == FMT_PICKLED:
res = pickle.loads(res)
return res, is_exception
async def call(
self,
func: Callable[..., T_Retval],
args: tuple[Any],
limiter: CapacityLimiter,
) -> T_Retval:
result, is_exception = await to_thread.run_sync(
self._call,
func,
args,
limiter=limiter,
)
if is_exception:
raise result
return result
def _stop_workers(workers: deque[Worker]) -> None:
for worker in workers:
worker.destroy()
workers.clear()
async def run_sync(
func: Callable[[Unpack[PosArgsT]], T_Retval],
*args: Unpack[PosArgsT],
limiter: CapacityLimiter | None = None,
) -> T_Retval:
"""
Call the given function with the given arguments in a subinterpreter.
If the ``cancellable`` option is enabled and the task waiting for its completion is
cancelled, the call will still run its course but its return value (or any raised
exception) will be ignored.
.. warning:: This feature is **experimental**. The upstream interpreter API has not
yet been finalized or thoroughly tested, so don't rely on this for anything
mission critical.
:param func: a callable
:param args: positional arguments for the callable
:param limiter: capacity limiter to use to limit the total amount of subinterpreters
running (if omitted, the default limiter is used)
:return: the result of the call
:raises BrokenWorkerIntepreter: if there's an internal error in a subinterpreter
"""
if sys.version_info <= (3, 13):
raise RuntimeError("subinterpreters require at least Python 3.13")
if limiter is None:
limiter = current_default_interpreter_limiter()
try:
idle_workers = _idle_workers.get()
except LookupError:
idle_workers = deque()
_idle_workers.set(idle_workers)
atexit.register(_stop_workers, idle_workers)
async with limiter:
try:
worker = idle_workers.pop()
except IndexError:
worker = Worker()
try:
return await worker.call(func, args, limiter)
finally:
# Prune workers that have been idle for too long
now = current_time()
while idle_workers:
if now - idle_workers[0].last_used <= MAX_WORKER_IDLE_TIME:
break
await to_thread.run_sync(idle_workers.popleft().destroy, limiter=limiter)
worker.last_used = current_time()
idle_workers.append(worker)
def current_default_interpreter_limiter() -> CapacityLimiter:
"""
Return the capacity limiter that is used by default to limit the number of
concurrently running subinterpreters.
Defaults to the number of CPU cores.
:return: a capacity limiter object
"""
try:
return _default_interpreter_limiter.get()
except LookupError:
limiter = CapacityLimiter(os.cpu_count() or DEFAULT_CPU_COUNT)
_default_interpreter_limiter.set(limiter)
return limiter
|