From 9de06f62eec88fcb1dae5f6bda88c249e6962746 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Mon, 6 Dec 2021 01:49:29 -0800 Subject: [PATCH] Use Python executors instead of pools ProcessPool/ThreadPool don't have the ability to notice when a child worker was terminated. ProcessPoolExecutor and ThreadPoolExecutor do notice and provide better error messages. Add tests to check. --- setup.cfg | 1 + src/ocrmypdf/_sync.py | 13 +++ src/ocrmypdf/builtin_plugins/concurrency.py | 48 +++++------ tests/conftest.py | 7 +- .../plugins/tesseract_simulate_oom_killer.py | 83 +++++++++++++++++++ tests/test_concurrency.py | 23 +++++ 6 files changed, 149 insertions(+), 26 deletions(-) create mode 100644 tests/plugins/tesseract_simulate_oom_killer.py create mode 100644 tests/test_concurrency.py diff --git a/setup.cfg b/setup.cfg index 9bc31e8c..ef55ea95 100644 --- a/setup.cfg +++ b/setup.cfg @@ -55,6 +55,7 @@ install_requires = tqdm>=4 importlib-metadata>=4;python_version<'3.8' # until Python 3.8 importlib-resources>=5;python_version<'3.9' # until Python 3.9 + typing-extensions;python_version<'3.8' # until Python 3.8 python_requires = >=3.7 include_package_data = True package_dir = diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index 349851ca..7958674c 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -11,6 +11,8 @@ import logging.handlers import os import sys import threading +from concurrent.futures.process import BrokenProcessPool +from concurrent.futures.thread import BrokenThreadPool from functools import partial from pathlib import Path from tempfile import mkdtemp @@ -430,6 +432,17 @@ def run_pipeline( "image pixel limit." ) return ExitCode.other_error + except ( + BrokenProcessPool if not api else NeverRaise, + BrokenThreadPool if not api else NeverRaise, + ) as e: + log.exception( + "A worker process was terminated unexpectedly. This is known to occur if " + "processing your file takes all available swap space and RAM. It may " + "help to try again with a smaller number of jobs, using the --jobs " + "argument." + ) + return ExitCode.child_process_error except (Exception if not api else NeverRaise): # pylint: disable=broad-except log.exception("An exception occurred while executing the pipeline") return ExitCode.other_error diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py index 74e6b504..8465d989 100644 --- a/src/ocrmypdf/builtin_plugins/concurrency.py +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -19,6 +19,7 @@ import queue import signal import sys import threading +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed from contextlib import suppress from multiprocessing.pool import Pool, ThreadPool from typing import Callable, Iterable, Type, Union @@ -30,7 +31,7 @@ from ocrmypdf._logging import TqdmConsole from ocrmypdf.exceptions import InputFileError from ocrmypdf.helpers import remove_all_log_handlers -ProcessPool = Pool +FuturesExecutorClass = Union[Type[ThreadPoolExecutor], Type[ProcessPoolExecutor]] Queue = Union[multiprocessing.Queue, queue.Queue] UserInit = Callable[[], None] WorkerInit = Callable[[Queue, UserInit, int], None] @@ -110,11 +111,11 @@ class StandardExecutor(Executor): ): if use_threads: log_queue: Queue = queue.Queue(-1) - pool_class: Type[Pool] = ThreadPool + executor_class: FuturesExecutorClass = ThreadPoolExecutor initializer: WorkerInit = thread_init else: log_queue = multiprocessing.Queue(-1) - pool_class = ProcessPool + executor_class = ProcessPoolExecutor initializer = process_init # Regardless of whether we use_threads for worker processes, the log_listener @@ -123,39 +124,36 @@ class StandardExecutor(Executor): listener = threading.Thread(target=log_listener, args=(log_queue,)) listener.start() - with self.pbar_class(**tqdm_kwargs) as pbar: - pool = pool_class( - processes=max_workers, - initializer=initializer, - initargs=(log_queue, worker_initializer, logging.getLogger("").level), - ) + with self.pbar_class(**tqdm_kwargs) as pbar, executor_class( + max_workers=max_workers, + initializer=initializer, + initargs=(log_queue, worker_initializer, logging.getLogger("").level), + ) as executor: + futures = [executor.submit(task, args) for args in task_arguments] try: - results = pool.imap_unordered(task, task_arguments) - for result in results: - if task_finished: - task_finished(result, pbar) - else: - pbar.update() + for future in as_completed(futures): + result = future.result() + task_finished(result, pbar) + pbar.update() except KeyboardInterrupt: # Terminate pool so we exit instantly - pool.terminate() - # Don't try listener.join() here, will deadlock + executor.shutdown(wait=False, cancel_futures=True) raise except Exception: if not os.environ.get("PYTEST_CURRENT_TEST", ""): - # Unless inside pytest, exit immediately because no one wants - # to wait for child processes to finalize results that will be - # thrown away. Inside pytest, we want child processes to exit - # cleanly so that they output an error messages or coverage data - # we need from them. - pool.terminate() + # Normally we shutdown without waiting for other child workers + # on error, because there is no point in waiting for them. Their + # results will be discard. But if the condition above is True, + # then we are running in pytest, and we want everything to exit + # as cleanly as possible so that we get good error messages. + executor.shutdown(wait=False, cancel_futures=True) raise finally: # Terminate log listener log_queue.put_nowait(None) - pool.close() - pool.join() + # When the above succeeds, wait for the listener thread to exit. (If + # an exception occurs, we don't try to join, in case it deadlocks.) listener.join() diff --git a/tests/conftest.py b/tests/conftest.py index 56dc7057..4b20d94b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,7 @@ import platform import sys from pathlib import Path from subprocess import PIPE, CompletedProcess, run -from typing import AnyStr, List, Literal, Tuple, overload +from typing import AnyStr, List, Tuple, overload import pytest @@ -19,6 +19,11 @@ from ocrmypdf._exec import unpaper from ocrmypdf._plugin_manager import get_parser_options_plugins from ocrmypdf.exceptions import ExitCode +try: + from typing import Literal +except ImportError: + from typing_extensions import Literal + def is_linux(): return platform.system() == 'Linux' diff --git a/tests/plugins/tesseract_simulate_oom_killer.py b/tests/plugins/tesseract_simulate_oom_killer.py new file mode 100644 index 00000000..d56345ff --- /dev/null +++ b/tests/plugins/tesseract_simulate_oom_killer.py @@ -0,0 +1,83 @@ +# © 2021 James R. Barlow: github.com/jbarlow83 +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the +# "Software"), to deal in the Software without restriction, including +# without limitation the rights to use, copy, modify, merge, publish, +# distribute, sublicense, and/or sell copies of the Software, and to +# permit persons to whom the Software is furnished to do so, subject to +# the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +# type: ignore + +"""Tesseract no-op plugin that simulates the OOM killer on page 4. + +OCRmyPDF can use a lot of memory, even that it might trigger the +OOM killer on Linux or similar features on other platforms. We want to +ensure we fail with an error rather than deadlock in such cases. + +Page 4 was chosen because of this number's association with bad luck +in many East Asian cultures. +""" + +import os +import signal +import sys +from pathlib import Path + +from ocrmypdf import hookimpl + +# Ugly hack that let us use the NoopOcrEngine without setting up packaging for our +# tests. +# This hack also requires us to set type: ignore +parent_file = Path(__file__).with_name('tesseract_noop.py') +parent = compile(parent_file.read_text(), parent_file, mode='exec') +exec(parent) +NoopOcrEngine = locals()['NoopOcrEngine'] + + +class Page4Engine(NoopOcrEngine): + def __str__(self): + return f"NO-OP Page 4 {NoopOcrEngine.version()}" + + @staticmethod + def generate_hocr(input_file: Path, output_hocr, output_text, options): + if input_file.stem.startswith('000004'): + # Suicide + os.kill(os.getpid(), signal.SIGKILL) + else: + return NoopOcrEngine.generate_hocr( + input_file, output_hocr, output_text, options + ) + + @staticmethod + def generate_pdf(input_file, output_pdf, output_text, options): + if input_file.stem.startswith('000004'): + # Suicide + os.kill(os.getpid(), signal.SIGKILL) + else: + return NoopOcrEngine.generate_pdf( + input_file, output_pdf, output_text, options + ) + + +@hookimpl +def check_options(options): + if options.use_threads: + raise ValueError("I'm not compatible with use_threads") + + +@hookimpl +def get_ocr_engine(): + return Page4Engine() diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py new file mode 100644 index 00000000..57185a07 --- /dev/null +++ b/tests/test_concurrency.py @@ -0,0 +1,23 @@ +# © 2021 James R. Barlow: github.com/jbarlow83 +# +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + + +import pytest + +from ocrmypdf import ExitCode + +from .conftest import run_ocrmypdf_api + + +def test_child_page4(resources, no_outpdf): + exitcode = run_ocrmypdf_api( + resources / 'multipage.pdf', + no_outpdf, + '--force-ocr', + '--plugin', + 'tests/plugins/tesseract_simulate_oom_killer.py', + ) + assert exitcode == ExitCode.child_process_error