diff options
Diffstat (limited to 'Lib/concurrent/futures/interpreter.py')
| -rw-r--r-- | Lib/concurrent/futures/interpreter.py | 241 |
1 files changed, 241 insertions, 0 deletions
diff --git a/Lib/concurrent/futures/interpreter.py b/Lib/concurrent/futures/interpreter.py new file mode 100644 index 0000000..fd7941a --- /dev/null +++ b/Lib/concurrent/futures/interpreter.py @@ -0,0 +1,241 @@ +"""Implements InterpreterPoolExecutor.""" + +import contextlib +import pickle +import textwrap +from . import thread as _thread +import _interpreters +import _interpqueues + + +class ExecutionFailed(_interpreters.InterpreterError): + """An unhandled exception happened during execution.""" + + def __init__(self, excinfo): + msg = excinfo.formatted + if not msg: + if excinfo.type and excinfo.msg: + msg = f'{excinfo.type.__name__}: {excinfo.msg}' + else: + msg = excinfo.type.__name__ or excinfo.msg + super().__init__(msg) + self.excinfo = excinfo + + def __str__(self): + try: + formatted = self.excinfo.errdisplay + except Exception: + return super().__str__() + else: + return textwrap.dedent(f""" +{super().__str__()} + +Uncaught in the interpreter: + +{formatted} + """.strip()) + + +UNBOUND = 2 # error; this should not happen. + + +class WorkerContext(_thread.WorkerContext): + + @classmethod + def prepare(cls, initializer, initargs, shared): + def resolve_task(fn, args, kwargs): + if isinstance(fn, str): + # XXX Circle back to this later. + raise TypeError('scripts not supported') + if args or kwargs: + raise ValueError(f'a script does not take args or kwargs, got {args!r} and {kwargs!r}') + data = textwrap.dedent(fn) + kind = 'script' + # Make sure the script compiles. + # Ideally we wouldn't throw away the resulting code + # object. However, there isn't much to be done until + # code objects are shareable and/or we do a better job + # of supporting code objects in _interpreters.exec(). + compile(data, '<string>', 'exec') + else: + # Functions defined in the __main__ module can't be pickled, + # so they can't be used here. In the future, we could possibly + # borrow from multiprocessing to work around this. + data = pickle.dumps((fn, args, kwargs)) + kind = 'function' + return (data, kind) + + if initializer is not None: + try: + initdata = resolve_task(initializer, initargs, {}) + except ValueError: + if isinstance(initializer, str) and initargs: + raise ValueError(f'an initializer script does not take args, got {initargs!r}') + raise # re-raise + else: + initdata = None + def create_context(): + return cls(initdata, shared) + return create_context, resolve_task + + @classmethod + @contextlib.contextmanager + def _capture_exc(cls, resultsid): + try: + yield + except BaseException as exc: + # Send the captured exception out on the results queue, + # but still leave it unhandled for the interpreter to handle. + err = pickle.dumps(exc) + _interpqueues.put(resultsid, (None, err), 1, UNBOUND) + raise # re-raise + + @classmethod + def _send_script_result(cls, resultsid): + _interpqueues.put(resultsid, (None, None), 0, UNBOUND) + + @classmethod + def _call(cls, func, args, kwargs, resultsid): + with cls._capture_exc(resultsid): + res = func(*args or (), **kwargs or {}) + # Send the result back. + try: + _interpqueues.put(resultsid, (res, None), 0, UNBOUND) + except _interpreters.NotShareableError: + res = pickle.dumps(res) + _interpqueues.put(resultsid, (res, None), 1, UNBOUND) + + @classmethod + def _call_pickled(cls, pickled, resultsid): + fn, args, kwargs = pickle.loads(pickled) + cls._call(fn, args, kwargs, resultsid) + + def __init__(self, initdata, shared=None): + self.initdata = initdata + self.shared = dict(shared) if shared else None + self.interpid = None + self.resultsid = None + + def __del__(self): + if self.interpid is not None: + self.finalize() + + def _exec(self, script): + assert self.interpid is not None + excinfo = _interpreters.exec(self.interpid, script, restrict=True) + if excinfo is not None: + raise ExecutionFailed(excinfo) + + def initialize(self): + assert self.interpid is None, self.interpid + self.interpid = _interpreters.create(reqrefs=True) + try: + _interpreters.incref(self.interpid) + + maxsize = 0 + fmt = 0 + self.resultsid = _interpqueues.create(maxsize, fmt, UNBOUND) + + self._exec(f'from {__name__} import WorkerContext') + + if self.shared: + _interpreters.set___main___attrs( + self.interpid, self.shared, restrict=True) + + if self.initdata: + self.run(self.initdata) + except BaseException: + self.finalize() + raise # re-raise + + def finalize(self): + interpid = self.interpid + resultsid = self.resultsid + self.resultsid = None + self.interpid = None + if resultsid is not None: + try: + _interpqueues.destroy(resultsid) + except _interpqueues.QueueNotFoundError: + pass + if interpid is not None: + try: + _interpreters.decref(interpid) + except _interpreters.InterpreterNotFoundError: + pass + + def run(self, task): + data, kind = task + if kind == 'script': + raise NotImplementedError('script kind disabled') + script = f""" +with WorkerContext._capture_exc({self.resultsid}): +{textwrap.indent(data, ' ')} +WorkerContext._send_script_result({self.resultsid})""" + elif kind == 'function': + script = f'WorkerContext._call_pickled({data!r}, {self.resultsid})' + else: + raise NotImplementedError(kind) + + try: + self._exec(script) + except ExecutionFailed as exc: + exc_wrapper = exc + else: + exc_wrapper = None + + # Return the result, or raise the exception. + while True: + try: + obj = _interpqueues.get(self.resultsid) + except _interpqueues.QueueNotFoundError: + raise # re-raise + except _interpqueues.QueueError: + continue + except ModuleNotFoundError: + # interpreters.queues doesn't exist, which means + # QueueEmpty doesn't. Act as though it does. + continue + else: + break + (res, excdata), pickled, unboundop = obj + assert unboundop is None, unboundop + if excdata is not None: + assert res is None, res + assert pickled + assert exc_wrapper is not None + exc = pickle.loads(excdata) + raise exc from exc_wrapper + return pickle.loads(res) if pickled else res + + +class BrokenInterpreterPool(_thread.BrokenThreadPool): + """ + Raised when a worker thread in an InterpreterPoolExecutor failed initializing. + """ + + +class InterpreterPoolExecutor(_thread.ThreadPoolExecutor): + + BROKEN = BrokenInterpreterPool + + @classmethod + def prepare_context(cls, initializer, initargs, shared): + return WorkerContext.prepare(initializer, initargs, shared) + + def __init__(self, max_workers=None, thread_name_prefix='', + initializer=None, initargs=(), shared=None): + """Initializes a new InterpreterPoolExecutor instance. + + Args: + max_workers: The maximum number of interpreters that can be used to + execute the given calls. + thread_name_prefix: An optional name prefix to give our threads. + initializer: A callable or script used to initialize + each worker interpreter. + initargs: A tuple of arguments to pass to the initializer. + shared: A mapping of shareabled objects to be inserted into + each worker interpreter. + """ + super().__init__(max_workers, thread_name_prefix, + initializer, initargs, shared=shared) |