menekis/OCRmyPDF
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor concurrency so that it is pluggable

Browse Source 
However, this may not be the best idea because it involves global
state that could be overridden by a parallel call to ocrmypdf.ocr.
This commit is contained in:
James R. Barlow
2021-01-24 23:46:29 -08:00
parent 34e564cd7d
commit 26b4d9bb4b
6 changed files with 242 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ocrmypdf/_concurrent.py
+23 -129
View File
@@ -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,
)
src/ocrmypdf/_plugin_manager.py
+8 -2
View File
@@ -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
src/ocrmypdf/_sync.py
+4 -2
View File
@@ -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)
src/ocrmypdf/builtin_plugins/concurrency.py
+164
View File
@@ -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
src/ocrmypdf/pdfinfo/info.py
+1 -1
View File
@@ -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,
src/ocrmypdf/pluginspec.py
+42 -1
View File
@@ -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*.