219 lines
6.4 KiB
Python
219 lines
6.4 KiB
Python
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)
|
|
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)
|
|
|
|
res: Any
|
|
is_exception: bool
|
|
if exc_info := interpreters.exec(self._interpreter_id, self._run_func):
|
|
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
|