From f8c2337449d59ddccf14601d26ba58009167b00b Mon Sep 17 00:00:00 2001 From: Eric Snow Date: Wed, 18 Jun 2025 17:57:14 -0600 Subject: [PATCH] gh-133485: Use interpreters.Interpreter in InterpreterPoolExecutor (gh-133957) Most importantly, this resolves the issues with functions and types defined in __main__. It also expands the number of supported objects and simplifies the implementation. (cherry picked from commit 725da50520c5cc3002f6f55bf67c56853507c604) Co-authored-by: Eric Snow --- Doc/library/concurrent.futures.rst | 19 +- Lib/concurrent/futures/interpreter.py | 181 +++--------- Lib/test/test_concurrent_futures/test_init.py | 4 + .../test_interpreter_pool.py | 263 ++++++++++++++---- Python/crossinterp.c | 4 +- 5 files changed, 266 insertions(+), 205 deletions(-) diff --git a/Doc/library/concurrent.futures.rst b/Doc/library/concurrent.futures.rst index 3c8d9ab111e09e..dd92765038c4f7 100644 --- a/Doc/library/concurrent.futures.rst +++ b/Doc/library/concurrent.futures.rst @@ -265,7 +265,7 @@ Each worker's interpreter is isolated from all the other interpreters. "Isolated" means each interpreter has its own runtime state and operates completely independently. For example, if you redirect :data:`sys.stdout` in one interpreter, it will not be automatically -redirected any other interpreter. If you import a module in one +redirected to any other interpreter. If you import a module in one interpreter, it is not automatically imported in any other. You would need to import the module separately in interpreter where you need it. In fact, each module imported in an interpreter is @@ -287,7 +287,7 @@ efficient alternative is to serialize with :mod:`pickle` and then send the bytes over a shared :mod:`socket ` or :func:`pipe `. -.. class:: InterpreterPoolExecutor(max_workers=None, thread_name_prefix='', initializer=None, initargs=(), shared=None) +.. class:: InterpreterPoolExecutor(max_workers=None, thread_name_prefix='', initializer=None, initargs=()) A :class:`ThreadPoolExecutor` subclass that executes calls asynchronously using a pool of at most *max_workers* threads. Each thread runs @@ -304,21 +304,10 @@ the bytes over a shared :mod:`socket ` or and *initargs* using :mod:`pickle` when sending them to the worker's interpreter. - .. note:: - Functions defined in the ``__main__`` module cannot be pickled - and thus cannot be used. - .. note:: The executor may replace uncaught exceptions from *initializer* with :class:`~concurrent.futures.interpreter.ExecutionFailed`. - The optional *shared* argument is a :class:`dict` of objects that all - interpreters in the pool share. The *shared* items are added to each - interpreter's ``__main__`` module. Not all objects are shareable. - Shareable objects include the builtin singletons, :class:`str` - and :class:`bytes`, and :class:`memoryview`. See :pep:`734` - for more info. - Other caveats from parent :class:`ThreadPoolExecutor` apply here. :meth:`~Executor.submit` and :meth:`~Executor.map` work like normal, @@ -326,10 +315,6 @@ except the worker serializes the callable and arguments using :mod:`pickle` when sending them to its interpreter. The worker likewise serializes the return value when sending it back. -.. note:: - Functions defined in the ``__main__`` module cannot be pickled - and thus cannot be used. - When a worker's current task raises an uncaught exception, the worker always tries to preserve the exception as-is. If that is successful then it also sets the ``__cause__`` to a corresponding diff --git a/Lib/concurrent/futures/interpreter.py b/Lib/concurrent/futures/interpreter.py index f12b4ac33cda27..cbb60ce80c1813 100644 --- a/Lib/concurrent/futures/interpreter.py +++ b/Lib/concurrent/futures/interpreter.py @@ -1,56 +1,39 @@ """Implements InterpreterPoolExecutor.""" -import contextlib -import pickle +from concurrent import interpreters +import sys import textwrap from . import thread as _thread -import _interpreters -import _interpqueues +import traceback -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): +def do_call(results, func, args, kwargs): + try: + return func(*args, **kwargs) + except BaseException as exc: + # Send the captured exception out on the results queue, + # but still leave it unhandled for the interpreter to handle. try: - formatted = self.excinfo.errdisplay - except Exception: - return super().__str__() - else: - return textwrap.dedent(f""" -{super().__str__()} - -Uncaught in the interpreter: - -{formatted} - """.strip()) + results.put(exc) + except interpreters.NotShareableError: + # The exception is not shareable. + print('exception is not shareable:', file=sys.stderr) + traceback.print_exception(exc) + results.put(None) + raise # re-raise class WorkerContext(_thread.WorkerContext): @classmethod - def prepare(cls, initializer, initargs, shared): + def prepare(cls, initializer, initargs): def resolve_task(fn, args, kwargs): if isinstance(fn, str): # XXX Circle back to this later. raise TypeError('scripts not supported') 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. task = (fn, args, kwargs) - data = pickle.dumps(task) - return data + return task if initializer is not None: try: @@ -62,68 +45,24 @@ def resolve_task(fn, args, kwargs): else: initdata = None def create_context(): - return cls(initdata, shared) + return cls(initdata) 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. - _interpqueues.put(resultsid, (None, exc)) - raise # re-raise - - @classmethod - def _send_script_result(cls, resultsid): - _interpqueues.put(resultsid, (None, None)) - - @classmethod - def _call(cls, func, args, kwargs, resultsid): - with cls._capture_exc(resultsid): - res = func(*args or (), **kwargs or {}) - # Send the result back. - with cls._capture_exc(resultsid): - _interpqueues.put(resultsid, (res, None)) - - @classmethod - def _call_pickled(cls, pickled, resultsid): - with cls._capture_exc(resultsid): - fn, args, kwargs = pickle.loads(pickled) - cls._call(fn, args, kwargs, resultsid) - - def __init__(self, initdata, shared=None): + def __init__(self, initdata): self.initdata = initdata - self.shared = dict(shared) if shared else None - self.interpid = None - self.resultsid = None + self.interp = None + self.results = None def __del__(self): - if self.interpid is not None: + if self.interp 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) + assert self.interp is None, self.interp + self.interp = interpreters.create() try: - _interpreters.incref(self.interpid) - maxsize = 0 - self.resultsid = _interpqueues.create(maxsize) - - self._exec(f'from {__name__} import WorkerContext') - - if self.shared: - _interpreters.set___main___attrs( - self.interpid, self.shared, restrict=True) + self.results = interpreters.create_queue(maxsize) if self.initdata: self.run(self.initdata) @@ -132,53 +71,25 @@ def initialize(self): 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 + interp = self.interp + results = self.results + self.results = None + self.interp = None + if results is not None: + del results + if interp is not None: + interp.close() def run(self, task): - data = task - script = f'WorkerContext._call_pickled({data!r}, {self.resultsid})' - 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: + return self.interp.call(do_call, self.results, *task) + except interpreters.ExecutionFailed as wrapper: + # Wait for the exception data to show up. + exc = self.results.get() + if exc is None: + # The exception must have been not shareable. 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, exc), unboundop = obj - assert unboundop is None, unboundop - if exc is not None: - assert res is None, res - assert exc_wrapper is not None - raise exc from exc_wrapper - return res + raise exc from wrapper class BrokenInterpreterPool(_thread.BrokenThreadPool): @@ -192,11 +103,11 @@ class InterpreterPoolExecutor(_thread.ThreadPoolExecutor): BROKEN = BrokenInterpreterPool @classmethod - def prepare_context(cls, initializer, initargs, shared): - return WorkerContext.prepare(initializer, initargs, shared) + def prepare_context(cls, initializer, initargs): + return WorkerContext.prepare(initializer, initargs) def __init__(self, max_workers=None, thread_name_prefix='', - initializer=None, initargs=(), shared=None): + initializer=None, initargs=()): """Initializes a new InterpreterPoolExecutor instance. Args: @@ -206,8 +117,6 @@ def __init__(self, max_workers=None, thread_name_prefix='', 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) + initializer, initargs) diff --git a/Lib/test/test_concurrent_futures/test_init.py b/Lib/test/test_concurrent_futures/test_init.py index df640929309318..6b8484c0d5f197 100644 --- a/Lib/test/test_concurrent_futures/test_init.py +++ b/Lib/test/test_concurrent_futures/test_init.py @@ -20,6 +20,10 @@ def init(x): global INITIALIZER_STATUS INITIALIZER_STATUS = x + # InterpreterPoolInitializerTest.test_initializer fails + # if we don't have a LOAD_GLOBAL. (It could be any global.) + # We will address this separately. + INITIALIZER_STATUS def get_init_status(): return INITIALIZER_STATUS diff --git a/Lib/test/test_concurrent_futures/test_interpreter_pool.py b/Lib/test/test_concurrent_futures/test_interpreter_pool.py index 5fd5684e1035e9..844dfdd6fc901c 100644 --- a/Lib/test/test_concurrent_futures/test_interpreter_pool.py +++ b/Lib/test/test_concurrent_futures/test_interpreter_pool.py @@ -2,35 +2,78 @@ import contextlib import io import os -import pickle +import sys import time import unittest -from concurrent.futures.interpreter import ( - ExecutionFailed, BrokenInterpreterPool, -) +from concurrent.futures.interpreter import BrokenInterpreterPool +from concurrent import interpreters from concurrent.interpreters import _queues as queues import _interpreters from test import support +from test.support import os_helper +from test.support import script_helper import test.test_asyncio.utils as testasyncio_utils from .executor import ExecutorTest, mul from .util import BaseTestCase, InterpreterPoolMixin, setup_module +WINDOWS = sys.platform.startswith('win') + + +@contextlib.contextmanager +def nonblocking(fd): + blocking = os.get_blocking(fd) + if blocking: + os.set_blocking(fd, False) + try: + yield + finally: + if blocking: + os.set_blocking(fd, blocking) + + +def read_file_with_timeout(fd, nbytes, timeout): + with nonblocking(fd): + end = time.time() + timeout + try: + return os.read(fd, nbytes) + except BlockingIOError: + pass + while time.time() < end: + try: + return os.read(fd, nbytes) + except BlockingIOError: + continue + else: + raise TimeoutError('nothing to read') + + +if not WINDOWS: + import select + def read_file_with_timeout(fd, nbytes, timeout): + r, _, _ = select.select([fd], [], [], timeout) + if fd not in r: + raise TimeoutError('nothing to read') + return os.read(fd, nbytes) + + def noop(): pass def write_msg(fd, msg): + import os os.write(fd, msg + b'\0') -def read_msg(fd): +def read_msg(fd, timeout=10.0): msg = b'' - while ch := os.read(fd, 1): - if ch == b'\0': - return msg + ch = read_file_with_timeout(fd, 1, timeout) + while ch != b'\0': msg += ch + ch = os.read(fd, 1) + return msg def get_current_name(): @@ -113,6 +156,38 @@ def test_init_func(self): self.assertEqual(before, b'\0') self.assertEqual(after, msg) + def test_init_with___main___global(self): + # See https://github.com/python/cpython/pull/133957#issuecomment-2927415311. + text = """if True: + from concurrent.futures import InterpreterPoolExecutor + + INITIALIZER_STATUS = 'uninitialized' + + def init(x): + global INITIALIZER_STATUS + INITIALIZER_STATUS = x + INITIALIZER_STATUS + + def get_init_status(): + return INITIALIZER_STATUS + + if __name__ == "__main__": + exe = InterpreterPoolExecutor(initializer=init, + initargs=('initialized',)) + fut = exe.submit(get_init_status) + print(fut.result()) # 'initialized' + exe.shutdown(wait=True) + print(INITIALIZER_STATUS) # 'uninitialized' + """ + with os_helper.temp_dir() as tempdir: + filename = script_helper.make_script(tempdir, 'my-script', text) + res = script_helper.assert_python_ok(filename) + stdout = res.out.decode('utf-8').strip() + self.assertEqual(stdout.splitlines(), [ + 'initialized', + 'uninitialized', + ]) + def test_init_closure(self): count = 0 def init1(): @@ -121,10 +196,19 @@ def init2(): nonlocal count count += 1 - with self.assertRaises(pickle.PicklingError): - self.executor_type(initializer=init1) - with self.assertRaises(pickle.PicklingError): - self.executor_type(initializer=init2) + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=init1) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() + + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=init2) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() def test_init_instance_method(self): class Spam: @@ -132,26 +216,12 @@ def initializer(self): raise NotImplementedError spam = Spam() - with self.assertRaises(pickle.PicklingError): - self.executor_type(initializer=spam.initializer) - - def test_init_shared(self): - msg = b'eggs' - r, w = self.pipe() - script = f"""if True: - import os - if __name__ != '__main__': - import __main__ - spam = __main__.spam - os.write({w}, spam + b'\\0') - """ - - executor = self.executor_type(shared={'spam': msg}) - fut = executor.submit(exec, script) - fut.result() - after = read_msg(r) - - self.assertEqual(after, msg) + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=spam.initializer) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() @unittest.expectedFailure def test_init_exception_in_script(self): @@ -178,8 +248,6 @@ def test_init_exception_in_func(self): stderr = stderr.getvalue() self.assertIn('ExecutionFailed: Exception: spam', stderr) self.assertIn('Uncaught in the interpreter:', stderr) - self.assertIn('The above exception was the direct cause of the following exception:', - stderr) @unittest.expectedFailure def test_submit_script(self): @@ -208,10 +276,14 @@ def task2(): return spam executor = self.executor_type() - with self.assertRaises(pickle.PicklingError): - executor.submit(task1) - with self.assertRaises(pickle.PicklingError): - executor.submit(task2) + + fut = executor.submit(task1) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() + + fut = executor.submit(task2) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() def test_submit_local_instance(self): class Spam: @@ -219,8 +291,9 @@ def __init__(self): self.value = True executor = self.executor_type() - with self.assertRaises(pickle.PicklingError): - executor.submit(Spam) + fut = executor.submit(Spam) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() def test_submit_instance_method(self): class Spam: @@ -229,8 +302,9 @@ def run(self): spam = Spam() executor = self.executor_type() - with self.assertRaises(pickle.PicklingError): - executor.submit(spam.run) + fut = executor.submit(spam.run) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() def test_submit_func_globals(self): executor = self.executor_type() @@ -242,13 +316,14 @@ def test_submit_func_globals(self): @unittest.expectedFailure def test_submit_exception_in_script(self): + # Scripts are not supported currently. fut = self.executor.submit('raise Exception("spam")') with self.assertRaises(Exception) as captured: fut.result() self.assertIs(type(captured.exception), Exception) self.assertEqual(str(captured.exception), 'spam') cause = captured.exception.__cause__ - self.assertIs(type(cause), ExecutionFailed) + self.assertIs(type(cause), interpreters.ExecutionFailed) for attr in ('__name__', '__qualname__', '__module__'): self.assertEqual(getattr(cause.excinfo.type, attr), getattr(Exception, attr)) @@ -261,7 +336,7 @@ def test_submit_exception_in_func(self): self.assertIs(type(captured.exception), Exception) self.assertEqual(str(captured.exception), 'spam') cause = captured.exception.__cause__ - self.assertIs(type(cause), ExecutionFailed) + self.assertIs(type(cause), interpreters.ExecutionFailed) for attr in ('__name__', '__qualname__', '__module__'): self.assertEqual(getattr(cause.excinfo.type, attr), getattr(Exception, attr)) @@ -269,16 +344,93 @@ def test_submit_exception_in_func(self): def test_saturation(self): blocker = queues.create() - executor = self.executor_type(4, shared=dict(blocker=blocker)) + executor = self.executor_type(4) for i in range(15 * executor._max_workers): - executor.submit(exec, 'import __main__; __main__.blocker.get()') - #executor.submit('blocker.get()') + executor.submit(blocker.get) self.assertEqual(len(executor._threads), executor._max_workers) for i in range(15 * executor._max_workers): blocker.put_nowait(None) executor.shutdown(wait=True) + def test_blocking(self): + # There is no guarantee that a worker will be created for every + # submitted task. That's because there's a race between: + # + # * a new worker thread, created when task A was just submitted, + # becoming non-idle when it picks up task A + # * after task B is added to the queue, a new worker thread + # is started only if there are no idle workers + # (the check in ThreadPoolExecutor._adjust_thread_count()) + # + # That means we must not block waiting for *all* tasks to report + # "ready" before we unblock the known-ready workers. + ready = queues.create() + blocker = queues.create() + + def run(taskid, ready, blocker): + # There can't be any globals here. + ready.put_nowait(taskid) + blocker.get() # blocking + + numtasks = 10 + futures = [] + with self.executor_type() as executor: + # Request the jobs. + for i in range(numtasks): + fut = executor.submit(run, i, ready, blocker) + futures.append(fut) + pending = numtasks + while pending > 0: + # Wait for any to be ready. + done = 0 + for _ in range(pending): + try: + ready.get(timeout=1) # blocking + except interpreters.QueueEmpty: + pass + else: + done += 1 + pending -= done + # Unblock the workers. + for _ in range(done): + blocker.put_nowait(None) + + def test_blocking_with_limited_workers(self): + # This is essentially the same as test_blocking, + # but we explicitly force a limited number of workers, + # instead of it happening implicitly sometimes due to a race. + ready = queues.create() + blocker = queues.create() + + def run(taskid, ready, blocker): + # There can't be any globals here. + ready.put_nowait(taskid) + blocker.get() # blocking + + numtasks = 10 + futures = [] + with self.executor_type(4) as executor: + # Request the jobs. + for i in range(numtasks): + fut = executor.submit(run, i, ready, blocker) + futures.append(fut) + pending = numtasks + while pending > 0: + # Wait for any to be ready. + done = 0 + for _ in range(pending): + try: + ready.get(timeout=1) # blocking + except interpreters.QueueEmpty: + pass + else: + done += 1 + pending -= done + # Unblock the workers. + for _ in range(done): + blocker.put_nowait(None) + @support.requires_gil_enabled("gh-117344: test is flaky without the GIL") def test_idle_thread_reuse(self): executor = self.executor_type() @@ -289,12 +441,21 @@ def test_idle_thread_reuse(self): executor.shutdown(wait=True) def test_pickle_errors_propagate(self): - # GH-125864: Pickle errors happen before the script tries to execute, so the - # queue used to wait infinitely. - + # GH-125864: Pickle errors happen before the script tries to execute, + # so the queue used to wait infinitely. fut = self.executor.submit(PickleShenanigans(0)) - with self.assertRaisesRegex(RuntimeError, "gotcha"): + expected = interpreters.NotShareableError + with self.assertRaisesRegex(expected, 'args not shareable') as cm: fut.result() + self.assertRegex(str(cm.exception.__cause__), 'unpickled') + + def test_no_stale_references(self): + # Weak references don't cross between interpreters. + raise unittest.SkipTest('not applicable') + + def test_free_reference(self): + # Weak references don't cross between interpreters. + raise unittest.SkipTest('not applicable') class AsyncioTest(InterpretersMixin, testasyncio_utils.TestCase): diff --git a/Python/crossinterp.c b/Python/crossinterp.c index aa958fe2f5cbd7..c44bcd559000b3 100644 --- a/Python/crossinterp.c +++ b/Python/crossinterp.c @@ -674,7 +674,9 @@ _PyPickle_Loads(struct _unpickle_context *ctx, PyObject *pickled) finally: if (exc != NULL) { - sync_module_capture_exc(tstate, &ctx->main); + if (_PyErr_Occurred(tstate)) { + sync_module_capture_exc(tstate, &ctx->main); + } // We restore the original exception. // It might make sense to chain it (__context__). _PyErr_SetRaisedException(tstate, exc);