diff --git a/nbsite/_parallel.py b/nbsite/_parallel.py new file mode 100644 index 0000000..0da0ece --- /dev/null +++ b/nbsite/_parallel.py @@ -0,0 +1,87 @@ +"""Make Sphinx parallel builds survive a worker dying mid-task. + +Sphinx runs the reading and writing phases in forked workers and collects +each result over a pipe. If a worker dies without sending anything (killed by +the OS, a segfault in a C extension, ...) the pipe reaches EOF, ``pipe.recv()`` +raises ``EOFError`` and the build is aborted, which happens regularly on CI. + +``ParallelTasks`` is patched here so the chunk of a worker that died is re-run +in the main process instead of aborting the build. +""" +import traceback + +from sphinx.errors import SphinxParallelError +from sphinx.util import logging, parallel as _parallel + +logger = logging.getLogger(__name__) + +_orig_init = _parallel.ParallelTasks.__init__ +_orig_add_task = _parallel.ParallelTasks.add_task + + +def _init(self, nproc): + _orig_init(self, nproc) + # Sphinx keeps no reference to the task functions, and ``Process.start`` + # deletes the target and arguments the process was created with, so they + # have to be recorded to be able to re-run the task of a dead worker. + self._task_funcs = {} + + +def _add_task(self, task_func, arg=None, result_func=None): + self._task_funcs[self._taskid] = task_func + _orig_add_task(self, task_func, arg, result_func) + + +def _run_in_main_process(self, tid): + """Run the task of the worker that died, mimicking ``_process``.""" + proc = self._procs[tid] + proc.join() + logger.warning( + 'parallel worker (exitcode %s) died without sending a result, ' + 'running its task in the main process', proc.exitcode + ) + func, arg = self._task_funcs[tid], self._args[tid] + try: + result = func() if arg is None else func(arg) + except BaseException as err: + errmsg = traceback.format_exception_only(err.__class__, err)[0].strip() + return True, [], (errmsg, traceback.format_exc()) + return False, [], result + + +def _join_one(self): + joined_any = False + for tid, pipe in self._precvs.items(): + if pipe.poll(): + try: + exc, logs, result = pipe.recv() + except EOFError: + exc, logs, result = _run_in_main_process(self, tid) + if exc: + raise SphinxParallelError(*result) + for log in logs: + logger.handle(log) + self._task_funcs.pop(tid) + self._result_funcs.pop(tid)(self._args.pop(tid), result) + self._procs[tid].join() + self._precvs.pop(tid) + self._pworking -= 1 + joined_any = True + break + + while self._precvs_waiting and self._pworking < self.nproc: + newtid, newprecv = self._precvs_waiting.popitem() + self._precvs[newtid] = newprecv + self._procs[newtid].start() + self._pworking += 1 + + return joined_any + + +def patch_parallel_tasks(): + """Patch ``sphinx.util.parallel.ParallelTasks``, idempotently.""" + if _parallel.ParallelTasks._join_one is _join_one: + return + _parallel.ParallelTasks.__init__ = _init + _parallel.ParallelTasks.add_task = _add_task + _parallel.ParallelTasks._join_one = _join_one diff --git a/nbsite/cmd.py b/nbsite/cmd.py index 63133e9..a3a3e87 100644 --- a/nbsite/cmd.py +++ b/nbsite/cmd.py @@ -7,42 +7,12 @@ from os.path import dirname from sphinx.application import Sphinx -from sphinx.util import parallel as _sphinx_parallel +from ._parallel import patch_parallel_tasks from .scripts import clean_dist_html, fix_links from .util import copy_files -_orig_join_one = _sphinx_parallel.ParallelTasks._join_one - - -def _safe_join_one(self): - try: - return _orig_join_one(self) - except EOFError: - pass - for tid, proc in self._procs.copy().items(): - if proc.exitcode is None: - continue - # Worker died without sending; finish chunk in main proc so build completes. - print( - f"sphinx parallel worker (tid={tid}, exitcode={proc.exitcode}) died; " - f"running chunk serial in main process", - flush=True, - file=sys.stderr, - ) - _, func, arg = proc._args - proc.join() - self._procs.pop(tid) - self._precvs.pop(tid) - self._args.pop(tid) - result_func = self._result_funcs.pop(tid) - self._pworking -= 1 - result_func(arg, func(arg) if arg is not None else func()) - break - return True - - -_sphinx_parallel.ParallelTasks._join_one = _safe_join_one +patch_parallel_tasks() DEFAULT_SITE_ORDERING = [ "Introduction", diff --git a/nbsite/tests/test_parallel.py b/nbsite/tests/test_parallel.py new file mode 100644 index 0000000..55d092f --- /dev/null +++ b/nbsite/tests/test_parallel.py @@ -0,0 +1,67 @@ +import multiprocessing +import os + +import pytest + +from sphinx.errors import SphinxParallelError +from sphinx.util.parallel import ParallelTasks, parallel_available + +from nbsite._parallel import patch_parallel_tasks + +pytestmark = pytest.mark.skipif( + not parallel_available, reason='Sphinx parallel builds need forking' +) + + +@pytest.fixture(autouse=True) +def patched(): + patch_parallel_tasks() + + +def _die_in_worker(arg): + if multiprocessing.parent_process() is not None: + # Emulate a worker killed before it could send back a result + os._exit(1) + return f'{arg}-main' + + +def _echo(arg): + return arg + + +def _raise(arg): + raise ValueError('boom') + + +def test_task_of_dead_worker_is_run_in_main_process(): + results = [] + tasks = ParallelTasks(2) + tasks.add_task(_die_in_worker, 'a', lambda arg, result: results.append((arg, result))) + tasks.join() + assert results == [('a', 'a-main')] + + +def test_dead_worker_does_not_stop_the_other_tasks(): + results = [] + tasks = ParallelTasks(2) + tasks.add_task(_die_in_worker, 'a', lambda arg, result: results.append((arg, result))) + for arg in 'bc': + tasks.add_task(_echo, arg, lambda arg, result: results.append((arg, result))) + tasks.join() + assert sorted(results) == [('a', 'a-main'), ('b', 'b'), ('c', 'c')] + + +def test_results_are_collected_normally(): + results = [] + tasks = ParallelTasks(2) + for arg in 'abc': + tasks.add_task(_echo, arg, lambda arg, result: results.append((arg, result))) + tasks.join() + assert sorted(results) == [('a', 'a'), ('b', 'b'), ('c', 'c')] + + +def test_failing_task_still_raises(): + tasks = ParallelTasks(2) + tasks.add_task(_raise, 'a', lambda arg, result: None) + with pytest.raises(SphinxParallelError): + tasks.join()