From 26b4d9bb4b4bf508ed8d22419eff8a8e64512c67 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Tue, 22 Dec 2020 00:44:59 -0800 Subject: [PATCH] Refactor concurrency so that it is pluggable However, this may not be the best idea because it involves global state that could be overridden by a parallel call to ocrmypdf.ocr. --- src/ocrmypdf/_concurrent.py | 152 +++--------------- src/ocrmypdf/_plugin_manager.py | 10 +- src/ocrmypdf/_sync.py | 6 +- src/ocrmypdf/builtin_plugins/concurrency.py | 164 ++++++++++++++++++++ src/ocrmypdf/pdfinfo/info.py | 2 +- src/ocrmypdf/pluginspec.py | 43 ++++- 6 files changed, 242 insertions(+), 135 deletions(-) create mode 100644 src/ocrmypdf/builtin_plugins/concurrency.py diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py index f38150dd..8dc3a6c3 100644 --- a/src/ocrmypdf/_concurrent.py +++ b/src/ocrmypdf/_concurrent.py @@ -4,84 +4,20 @@ # 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 logging -import logging.handlers -import multiprocessing -import os -import queue -import signal -import sys -import threading -from contextlib import suppress -from multiprocessing import Pool as ProcessPool -from multiprocessing.dummy import Pool as ThreadPool -from typing import Callable, Iterable, Optional, Union - -from tqdm import tqdm - -from ocrmypdf.exceptions import InputFileError - -Queue = Union[multiprocessing.Queue, queue.Queue] +from typing import Callable, Iterable, Optional -def log_listener(q: Queue): - """Listen to the worker processes and forward the messages to logging - - For simplicity this is a thread rather than a process. Only one process - should actually write to sys.stderr or whatever we're using, so if this is - made into a process the main application needs to be directed to it. - - See https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes - """ - - while True: - try: - record = q.get() - if record is None: - break - logger = logging.getLogger(record.name) - logger.handle(record) - except Exception: # pylint: disable=broad-except - import traceback # pylint: disable=import-outside-toplevel - - print("Logging problem", file=sys.stderr) - traceback.print_exc(file=sys.stderr) - - -def process_sigbus(*args): - raise InputFileError("A worker process lost access to an input file") - - -def process_init(q: Queue, user_init: Callable[[], None], loglevel): - """Initialize a process pool worker""" - - # Ignore SIGINT (our parent process will kill us gracefully) - signal.signal(signal.SIGINT, signal.SIG_IGN) - - # Install SIGBUS handler (so our parent process can abort somewhat gracefully) - with suppress(AttributeError): # Windows and Cygwin do not have SIGBUS - signal.signal(signal.SIGBUS, process_sigbus) - - # Reconfigure the root logger for this process to send all messages to a queue - h = logging.handlers.QueueHandler(q) - root = logging.getLogger() - root.setLevel(loglevel) - root.handlers = [] - root.addHandler(h) - - user_init() +def _task_noop(*_args, **_kwargs): return -def thread_init(_queue: Queue, user_init: Callable[[], None], _loglevel): - # As a thread, block SIGBUS so the main thread deals with it... - with suppress(AttributeError): - # Windows and Cygwin do not have pthread_sigmask or SIGBUS - signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS}) +def _model(**_kwargs): + raise RuntimeError("Parallel executor not set up") - user_init() - return + +def set_execution_model(model): + global _model + _model = model def exec_progress_pool( @@ -89,64 +25,22 @@ def exec_progress_pool( use_threads: bool, max_workers: int, tqdm_kwargs: dict, - task_initializer: Optional[Callable] = None, - task: Optional[Callable] = None, + worker_initializer: Optional[Callable] = None, + task: Callable, task_arguments: Optional[Iterable] = None, task_finished: Optional[Callable] = None, ): + if not worker_initializer: + worker_initializer = _task_noop + if not task_finished: + task_finished = _task_noop - if use_threads: - log_queue = queue.Queue(-1) - pool_class = ThreadPool - initializer = thread_init - else: - log_queue = multiprocessing.Queue(-1) - pool_class = ProcessPool - initializer = process_init - - if not task_initializer: - - def _noop(): - return - - task_initializer = _noop - - # Regardless of whether we use_threads for worker processes, the log_listener - # must be a thread - listener = threading.Thread(target=log_listener, args=(log_queue,)) - listener.start() - - with tqdm(**tqdm_kwargs) as pbar: - pool = pool_class( - processes=max_workers, - initializer=initializer, - initargs=(log_queue, task_initializer, logging.getLogger("").level), - ) - try: - results = pool.imap_unordered(task, task_arguments) - for result in results: - if task_finished: - task_finished(result, pbar) - else: - pbar.update() - except KeyboardInterrupt: - # Terminate pool so we exit instantly - pool.terminate() - # Don't try listener.join() here, will deadlock - 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() - raise - finally: - # Terminate log listener - log_queue.put_nowait(None) - pool.close() - pool.join() - - listener.join() + _model( + use_threads=use_threads, + max_workers=max_workers, + tqdm_kwargs=tqdm_kwargs, + worker_initializer=worker_initializer, + task=task, + task_arguments=task_arguments, + task_finished=task_finished, + ) diff --git a/src/ocrmypdf/_plugin_manager.py b/src/ocrmypdf/_plugin_manager.py index 510ae65b..76e12bd7 100644 --- a/src/ocrmypdf/_plugin_manager.py +++ b/src/ocrmypdf/_plugin_manager.py @@ -32,7 +32,11 @@ class OcrmypdfPluginManager(pluggy.PluginManager): """ def __init__( - self, *args, plugins: List[Union[str, Path]], builtins: bool = True, **kwargs, + self, + *args, + plugins: List[Union[str, Path]], + builtins: bool = True, + **kwargs, ): self.__init_args = args self.__init_kwargs = kwargs @@ -88,7 +92,9 @@ class OcrmypdfPluginManager(pluggy.PluginManager): def get_plugin_manager(plugins: List[Union[str, Path]], builtins=True): pm = OcrmypdfPluginManager( - project_name='ocrmypdf', plugins=plugins, builtins=builtins, + project_name='ocrmypdf', + plugins=plugins, + builtins=builtins, ) return pm diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index af32ad8c..1d78f1d9 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -18,7 +18,7 @@ from typing import List, NamedTuple, Optional, Tuple import pikepdf import PIL -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import exec_progress_pool, set_execution_model from ocrmypdf._graft import OcrGrafter from ocrmypdf._jobcontext import PageContext, PdfContext, cleanup_working_files from ocrmypdf._logging import PageNumberFilter @@ -279,7 +279,7 @@ def exec_concurrent(context: PdfContext): unit_scale=0.5, disable=not options.progress_bar, ), - task_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), + worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), task=exec_page_sync, task_arguments=context.get_page_contexts(), task_finished=update_page, @@ -346,6 +346,8 @@ def run_pipeline(options, *, plugin_manager, api=False): pikepdf_enable_mmap() + set_execution_model(plugin_manager.hook.get_parallel_executor()) + try: check_requested_output_file(options) start_input_file, original_filename = create_input_file(options, work_folder) diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py new file mode 100644 index 00000000..6cf56e42 --- /dev/null +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -0,0 +1,164 @@ +# © 2020 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/. + +# © 2020 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 logging +import logging.handlers +import multiprocessing +import os +import queue +import signal +import sys +import threading +from contextlib import suppress +from multiprocessing import Pool as ProcessPool +from multiprocessing.dummy import Pool as ThreadPool +from typing import Callable, Iterable, Optional, Union + +from tqdm import tqdm + +from ocrmypdf import hookimpl +from ocrmypdf.exceptions import InputFileError + +Queue = Union[multiprocessing.Queue, queue.Queue] + + +def log_listener(q: Queue): + """Listen to the worker processes and forward the messages to logging + + For simplicity this is a thread rather than a process. Only one process + should actually write to sys.stderr or whatever we're using, so if this is + made into a process the main application needs to be directed to it. + + See https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes + """ + + while True: + try: + record = q.get() + if record is None: + break + logger = logging.getLogger(record.name) + logger.handle(record) + except Exception: # pylint: disable=broad-except + import traceback # pylint: disable=import-outside-toplevel + + print("Logging problem", file=sys.stderr) + traceback.print_exc(file=sys.stderr) + + +def process_sigbus(*args): + raise InputFileError("A worker process lost access to an input file") + + +def process_init(q: Queue, user_init: Callable[[], None], loglevel): + """Initialize a process pool worker""" + + # Ignore SIGINT (our parent process will kill us gracefully) + signal.signal(signal.SIGINT, signal.SIG_IGN) + + # Install SIGBUS handler (so our parent process can abort somewhat gracefully) + with suppress(AttributeError): # Windows and Cygwin do not have SIGBUS + # Windows and Cygwin do not have pthread_sigmask or SIGBUS + signal.signal(signal.SIGBUS, process_sigbus) + + # Reconfigure the root logger for this process to send all messages to a queue + h = logging.handlers.QueueHandler(q) + root = logging.getLogger() + root.setLevel(loglevel) + root.handlers = [] + root.addHandler(h) + + user_init() + return + + +def thread_init(_queue: Queue, user_init: Callable[[], None], _loglevel): + # As a thread, block SIGBUS so the main thread deals with it... + with suppress(AttributeError): + signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS}) + + user_init() + return + + +def exec_progress_pool( + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Optional[Iterable] = None, + task_finished: Callable, +): + + if use_threads: + log_queue = queue.Queue(-1) + pool_class = ThreadPool + initializer = thread_init + else: + log_queue = multiprocessing.Queue(-1) + pool_class = ProcessPool + initializer = process_init + + if not worker_initializer: + + def _noop(): + return + + worker_initializer = _noop + + # Regardless of whether we use_threads for worker processes, the log_listener + # must be a thread + listener = threading.Thread(target=log_listener, args=(log_queue,)) + listener.start() + + with tqdm(**tqdm_kwargs) as pbar: + pool = pool_class( + processes=max_workers, + initializer=initializer, + initargs=(log_queue, worker_initializer, logging.getLogger("").level), + ) + try: + results = pool.imap_unordered(task, task_arguments) + for result in results: + if task_finished: + task_finished(result, pbar) + else: + pbar.update() + except KeyboardInterrupt: + # Terminate pool so we exit instantly + pool.terminate() + # Don't try listener.join() here, will deadlock + 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() + raise + finally: + # Terminate log listener + log_queue.put_nowait(None) + pool.close() + pool.join() + + listener.join() + + +@hookimpl +def get_parallel_executor(): + return exec_progress_pool diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index ae02fb98..a11e64cf 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -623,7 +623,7 @@ def _pdf_pageinfo_concurrent( tqdm_kwargs=dict( total=total, desc="Scanning contents", unit='page', disable=not progbar ), - task_initializer=partial( + worker_initializer=partial( _pdf_pageinfo_sync_init, infile, logging.getLogger('pdfminer').level ), task=_pdf_pageinfo_sync, diff --git a/src/ocrmypdf/pluginspec.py b/src/ocrmypdf/pluginspec.py index 9aee6596..538be7a1 100644 --- a/src/ocrmypdf/pluginspec.py +++ b/src/ocrmypdf/pluginspec.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod, abstractstaticmethod from argparse import ArgumentParser, Namespace from collections import namedtuple from pathlib import Path -from typing import TYPE_CHECKING, AbstractSet, List, Optional +from typing import TYPE_CHECKING, AbstractSet, Callable, Iterable, List, Optional import pluggy @@ -62,6 +62,47 @@ def check_options(options: Namespace) -> None: """ +class ParallelExecutor(ABC): + @abstractstaticmethod + def __call__( + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_finished: Callable, + task_arguments: Optional[Iterable] = None, + ): + """ + Args: + use_threads: If False, the workload is the sort that will benefit from + running in a multiprocessing context (for example, it uses Python + heavily, and parallelizing it with threads is not expected to be + performant). + max_workers: The maximum number of workers that should be run. + tdqm_kwargs: Arguments to set up the progress bar. + worker_initializer: Called when the worker is initialized, in the worker's + execution context. Must be possible to marshall to the worker. + task: Called when the worker starts a new task, in the worker's execution + context. Must be possible to marshallable to the worker. + task_finished: Called when a worker finishes a task, in the parent's + context. + task_arguments: An iterable that generates a group of parameters for each + task. This runs in the parent's context, but the parameters must be + marshallable to the worker. + """ + + +@hookspec(firstresult=True) +def get_parallel_executor() -> Callable: + """Called to perform parallel execution + + This may be used to replace OCRmyPDF's default parallel execution system + with a third party alternative. For example, you could make OCRmyPDF run in a + distributed environment. + """ + + @hookspec def validate(pdfinfo: 'PdfInfo', options: Namespace) -> None: """Called to give a plugin an opportunity to review *options* and *pdfinfo*.