From d274d88929d7d87b0e41313325ca56b1e4739b87 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Sat, 30 Jan 2021 17:36:30 -0800 Subject: [PATCH] Refactor to eliminate global state in _concurrent --- src/ocrmypdf/__init__.py | 1 + src/ocrmypdf/_concurrent.py | 123 ++++++++++---- src/ocrmypdf/_pipeline.py | 8 +- src/ocrmypdf/_sync.py | 19 +-- src/ocrmypdf/builtin_plugins/concurrency.py | 149 +++++++---------- src/ocrmypdf/extra_plugins/awslambda.py | 174 ++++++++++---------- src/ocrmypdf/optimize.py | 29 +++- src/ocrmypdf/pdfinfo/info.py | 23 ++- src/ocrmypdf/pluginspec.py | 37 +---- tests/test_validation.py | 2 +- 10 files changed, 301 insertions(+), 264 deletions(-) diff --git a/src/ocrmypdf/__init__.py b/src/ocrmypdf/__init__.py index 081324e9..80586658 100644 --- a/src/ocrmypdf/__init__.py +++ b/src/ocrmypdf/__init__.py @@ -8,6 +8,7 @@ from pluggy import HookimplMarker as _HookimplMarker from ocrmypdf import helpers, hocrtransform, leptonica, pdfa, pdfinfo +from ocrmypdf._concurrent import Executor from ocrmypdf._jobcontext import PageContext, PdfContext from ocrmypdf._version import PROGRAM_NAME, __version__ from ocrmypdf.api import Verbosity, configure_logging, ocr diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py index e3fb5b91..48c13e08 100644 --- a/src/ocrmypdf/_concurrent.py +++ b/src/ocrmypdf/_concurrent.py @@ -4,47 +4,110 @@ # 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 sys import threading +from abc import ABC, abstractmethod +from functools import partial from typing import Callable, Iterable, Optional -pool_lock = threading.Lock() +from tqdm import tqdm def _task_noop(*_args, **_kwargs): return -def _model(**_kwargs) -> None: - raise RuntimeError("Parallel executor not set up") +class Executor(ABC): + pool_lock = threading.Lock() + + def __call__( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Optional[Callable] = None, + task: Optional[Callable] = None, + task_arguments: Optional[Iterable] = None, + task_finished: Optional[Callable] = None, + ) -> 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. + """ + + if not task_arguments: + return # Nothing to do! + if not worker_initializer: + worker_initializer = _task_noop + if not task_finished: + task_finished = _task_noop + if not task: + task = _task_noop + + with self.pool_lock: + self._execute( + 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, + ) + + @abstractmethod + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + """Custom executors should override this method.""" -def set_execution_model(model): - global _model - _model = model +def setup_executor(plugin_manager) -> Executor: + return plugin_manager.hook.get_executor() -def exec_progress_pool( - *, - use_threads: bool, - max_workers: int, - tqdm_kwargs: dict, - 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 +class SerialExecutor(Executor): + """Implements a purely sequential executor using the parallel protocol. - with pool_lock: - _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, - ) + The current process/thread will be the worker that executes all tasks + in order. As such, ``worker_initializer`` will never be called. + """ + + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + with tqdm(**tqdm_kwargs) as pbar: + for args in task_arguments: + result = task(args) + task_finished(result, pbar) diff --git a/src/ocrmypdf/_pipeline.py b/src/ocrmypdf/_pipeline.py index 6de1f2e9..d4c171d8 100644 --- a/src/ocrmypdf/_pipeline.py +++ b/src/ocrmypdf/_pipeline.py @@ -22,6 +22,7 @@ from PIL import Image, ImageColor, ImageDraw from tqdm import tqdm from ocrmypdf import leptonica +from ocrmypdf._concurrent import Executor from ocrmypdf._exec import unpaper from ocrmypdf._jobcontext import PageContext, PdfContext from ocrmypdf._version import PROGRAM_NAME @@ -146,6 +147,8 @@ def triage(original_filename, input_file, output_file, options): def get_pdfinfo( input_file, + *, + executor: Executor, detailed_analysis=False, progbar=False, max_workers=None, @@ -158,6 +161,7 @@ def get_pdfinfo( progbar=progbar, max_workers=max_workers, check_pages=check_pages, + executor=executor, ) except pikepdf.PasswordError: raise EncryptedPdfError() @@ -792,7 +796,7 @@ def metadata_fixup(working_file: Path, context: PdfContext): return output_file -def optimize_pdf(input_file: Path, context: PdfContext): +def optimize_pdf(input_file: Path, context: PdfContext, executor: Executor): output_file = context.get_path('optimize.pdf') save_settings = dict( compress_streams=True, @@ -800,7 +804,7 @@ def optimize_pdf(input_file: Path, context: PdfContext): object_stream_mode=pikepdf.ObjectStreamMode.generate, linearize=should_linearize(input_file, context), ) - optimize(input_file, output_file, context, save_settings) + optimize(input_file, output_file, context, save_settings, executor) return output_file diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index 1d78f1d9..68577a5d 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -15,10 +15,9 @@ from pathlib import Path from tempfile import mkdtemp from typing import List, NamedTuple, Optional, Tuple -import pikepdf import PIL -from ocrmypdf._concurrent import exec_progress_pool, set_execution_model +from ocrmypdf._concurrent import Executor, setup_executor from ocrmypdf._graft import OcrGrafter from ocrmypdf._jobcontext import PageContext, PdfContext, cleanup_working_files from ocrmypdf._logging import PageNumberFilter @@ -224,14 +223,14 @@ def exec_page_sync(page_context: PageContext): ) -def post_process(pdf_file, context: PdfContext): +def post_process(pdf_file, context: PdfContext, executor: Executor): pdf_out = pdf_file if context.options.output_type.startswith('pdfa'): ps_stub_out = generate_postscript_stub(context) pdf_out = convert_to_pdfa(pdf_out, ps_stub_out, context) pdf_out = metadata_fixup(pdf_out, context) - return optimize_pdf(pdf_out, context) + return optimize_pdf(pdf_out, context, executor) def worker_init(max_pixels: int): @@ -242,7 +241,7 @@ def worker_init(max_pixels: int): pikepdf_enable_mmap() -def exec_concurrent(context: PdfContext): +def exec_concurrent(context: PdfContext, executor: Executor): """Execute the pipeline concurrently""" # Run exec_page_sync on every page context @@ -269,7 +268,7 @@ def exec_concurrent(context: PdfContext): finally: tls.pageno = None - exec_progress_pool( + executor( use_threads=options.use_threads, max_workers=max_workers, tqdm_kwargs=dict( @@ -296,7 +295,7 @@ def exec_concurrent(context: PdfContext): # PDF/A and metadata log.info("Postprocessing...") - pdf = post_process(pdf, context) + pdf = post_process(pdf, context, executor) # Copy PDF file to destination copy_final(pdf, options.output_file, context) @@ -346,8 +345,7 @@ def run_pipeline(options, *, plugin_manager, api=False): pikepdf_enable_mmap() - set_execution_model(plugin_manager.hook.get_parallel_executor()) - + executor = setup_executor(plugin_manager) try: check_requested_output_file(options) start_input_file, original_filename = create_input_file(options, work_folder) @@ -360,6 +358,7 @@ def run_pipeline(options, *, plugin_manager, api=False): # Gather pdfinfo and create context pdfinfo = get_pdfinfo( origin_pdf, + executor=executor, detailed_analysis=options.redo_ocr, progbar=options.progress_bar, max_workers=options.jobs if not options.use_threads else 1, # To help debug @@ -372,7 +371,7 @@ def run_pipeline(options, *, plugin_manager, api=False): validate_pdfinfo_options(context) # Execute the pipeline - exec_concurrent(context) + exec_concurrent(context, executor) if options.output_file == '-': log.info("Output sent to stdout") diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py index 3b3e0d74..b61654dc 100644 --- a/src/ocrmypdf/builtin_plugins/concurrency.py +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -26,7 +26,7 @@ from typing import Callable, Iterable, Optional, Union from tqdm import tqdm -from ocrmypdf import hookimpl +from ocrmypdf import Executor, hookimpl from ocrmypdf.exceptions import InputFileError Queue = Union[multiprocessing.Queue, queue.Queue] @@ -91,99 +91,68 @@ def thread_init(_queue: Queue, user_init: Callable[[], None], _loglevel): return -def exec_progress_pool( - *, - use_threads: bool, - max_workers: int, - tqdm_kwargs: dict, - worker_initializer: Optional[Callable], - task: Callable, - task_arguments: Optional[Iterable] = None, - task_finished: Callable, -): +class StandardExecutor(Executor): + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + 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 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 + # 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() - if not worker_initializer: - - def _noop(): - return - - worker_initializer = _noop - - _exec_progress_pool( - max_workers=max_workers, - tqdm_kwargs=tqdm_kwargs, - worker_initializer=worker_initializer, - task=task, - task_arguments=task_arguments, - task_finished=task_finished, - log_queue=log_queue, - pool_class=pool_class, - initializer=initializer, - ) - - -def _exec_progress_pool( - *, - max_workers: int, - tqdm_kwargs: dict, - worker_initializer: Callable, - task: Callable, - task_arguments: Optional[Iterable] = None, - task_finished: Callable, - log_queue: Queue, - pool_class: Callable, - initializer: Callable, -): - # 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. + 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() - raise - finally: - # Terminate log listener - log_queue.put_nowait(None) - pool.close() - pool.join() + # 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() + listener.join() @hookimpl -def get_parallel_executor(): - return exec_progress_pool +def get_executor(): + return StandardExecutor() diff --git a/src/ocrmypdf/extra_plugins/awslambda.py b/src/ocrmypdf/extra_plugins/awslambda.py index 4f7a7af7..8579683d 100644 --- a/src/ocrmypdf/extra_plugins/awslambda.py +++ b/src/ocrmypdf/extra_plugins/awslambda.py @@ -19,7 +19,7 @@ from multiprocessing.connection import Connection, wait from typing import Callable, Iterable, Optional from unittest.mock import Mock -from ocrmypdf import hookimpl +from ocrmypdf import Executor, hookimpl from ocrmypdf.exceptions import InputFileError @@ -80,97 +80,101 @@ def process_loop( return -def lambda_pool_impl( - *, - use_threads: bool, - max_workers: int, - tqdm_kwargs: dict, - worker_initializer: Callable, - task: Callable, - task_arguments: Optional[Iterable] = None, - task_finished: Callable, -): - pbar = Mock() +class LambdaExecutor(Executor): + def _execute( + self, + *, + use_threads: bool, + max_workers: int, + tqdm_kwargs: dict, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + ): + pbar = Mock() - if use_threads and max_workers == 1: - for args in task_arguments: - result = task(args) - task_finished(result, pbar) - return + if use_threads and max_workers == 1: + for args in task_arguments: + result = task(args) + task_finished(result, pbar) + return - _lambda_pool_impl( - max_workers=max_workers, - worker_initializer=worker_initializer, - task=task, - task_arguments=task_arguments, - task_finished=task_finished, - pbar=pbar, - ) - - -def _lambda_pool_impl( - *, - max_workers: int, - worker_initializer: Callable, - task: Callable, - task_arguments: Optional[Iterable] = None, - task_finished: Callable, - pbar, -): - task_arguments = list(task_arguments) - grouped_args = list(zip_longest(*list(split_every(max_workers, task_arguments)))) - if not grouped_args: - return - - processes = [] - connections = [] - for chunk in grouped_args: - parent_conn, child_conn = Pipe() - - worker_args = [args for args in chunk if args is not None] - process = Process( - target=process_loop, - args=( - child_conn, - worker_initializer, - logging.getLogger("").level, - task, - worker_args, - ), + self._lambda_pool_impl( + max_workers=max_workers, + worker_initializer=worker_initializer, + task=task, + task_arguments=task_arguments, + task_finished=task_finished, + pbar=pbar, ) - process.daemon = True - processes.append(process) - connections.append(parent_conn) - for process in processes: - process.start() + def _lambda_pool_impl( + self, + *, + max_workers: int, + worker_initializer: Callable, + task: Callable, + task_arguments: Iterable, + task_finished: Callable, + pbar, + ): + task_arguments = list(task_arguments) + grouped_args = list( + zip_longest(*list(split_every(max_workers, task_arguments))) + ) + if not grouped_args: + return - while connections: - for r in wait(connections): - try: - msg_type, msg = r.recv() - except EOFError: - connections.remove(r) - continue + processes = [] + connections = [] + for chunk in grouped_args: + parent_conn, child_conn = Pipe() - if msg_type == MessageType.result: - if task_finished: - task_finished(msg, pbar) - elif msg_type == 'log': - record = msg - logger = logging.getLogger(record.name) - logger.handle(record) - elif msg_type == MessageType.complete: - connections.remove(r) - elif msg_type == MessageType.exception: - for process in processes: - process.terminate() - raise msg + worker_args = [args for args in chunk if args is not None] + process = Process( + target=process_loop, + args=( + child_conn, + worker_initializer, + logging.getLogger("").level, + task, + worker_args, + ), + ) + process.daemon = True + processes.append(process) + connections.append(parent_conn) - for process in processes: - process.join() + for process in processes: + process.start() + + while connections: + for r in wait(connections): + try: + msg_type, msg = r.recv() + except EOFError: + connections.remove(r) + continue + + if msg_type == MessageType.result: + if task_finished: + task_finished(msg, pbar) + elif msg_type == 'log': + record = msg + logger = logging.getLogger(record.name) + logger.handle(record) + elif msg_type == MessageType.complete: + connections.remove(r) + elif msg_type == MessageType.exception: + for process in processes: + process.terminate() + raise msg + + for process in processes: + process.join() @hookimpl -def get_parallel_executor(): - return lambda_pool_impl +def get_executor(): + return LambdaExecutor() diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py index 3724f731..5993014d 100644 --- a/src/ocrmypdf/optimize.py +++ b/src/ocrmypdf/optimize.py @@ -34,7 +34,7 @@ from PIL import Image from tqdm import tqdm from ocrmypdf import leptonica -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import Executor, SerialExecutor from ocrmypdf._exec import jbig2enc, pngquant from ocrmypdf._jobcontext import PdfContext from ocrmypdf.exceptions import OutputFileAccessError @@ -301,7 +301,7 @@ def extract_images_jbig2(pike: Pdf, root: Path, options) -> Dict[int, List[XrefE def _produce_jbig2_images( - jbig2_groups: Dict[int, List[XrefExt]], root: Path, options + jbig2_groups: Dict[int, List[XrefExt]], root: Path, options, executor: Executor ) -> None: """Produce JBIG2 images from their groups""" @@ -333,7 +333,7 @@ def _produce_jbig2_images( jbig2_args = jbig2_single_args jbig2_convert = jbig2enc.convert_single_mp - exec_progress_pool( + executor( use_threads=True, max_workers=options.jobs, tqdm_kwargs=dict( @@ -348,7 +348,11 @@ def _produce_jbig2_images( def convert_to_jbig2( - pike: Pdf, jbig2_groups: Dict[int, List[XrefExt]], root: Path, options + pike: Pdf, + jbig2_groups: Dict[int, List[XrefExt]], + root: Path, + options, + executor: Executor, ) -> None: """Convert images to JBIG2 and insert into PDF. @@ -363,7 +367,7 @@ def convert_to_jbig2( and needs no dictionary. Currently this must be lossless JBIG2. """ - _produce_jbig2_images(jbig2_groups, root, options) + _produce_jbig2_images(jbig2_groups, root, options, executor) for group, xref_exts in jbig2_groups.items(): prefix = f'group{group:08d}' @@ -457,6 +461,7 @@ def transcode_pngs( image_name_fn: Callable[[Path, Xref], Path], root: Path, options, + executor, ) -> None: modified: MutableSet[Xref] = set() if options.optimize >= 2: @@ -476,7 +481,7 @@ def transcode_pngs( ) modified.add(xref) - exec_progress_pool( + executor( use_threads=True, max_workers=options.jobs, tqdm_kwargs=dict( @@ -569,7 +574,13 @@ def rewrite_png(pike: Pdf, im_obj: Object, compdata) -> None: # pragma: no cove im_obj.write(compdata.read(), filter=Name.FlateDecode, decode_parms=dparms) -def optimize(input_file: Path, output_file: Path, context, save_settings) -> None: +def optimize( + input_file: Path, + output_file: Path, + context, + save_settings, + executor: Executor = SerialExecutor(), +) -> None: options = context.options if options.optimize == 0: safe_symlink(input_file, output_file) @@ -591,10 +602,10 @@ def optimize(input_file: Path, output_file: Path, context, save_settings) -> Non # if options.optimize >= 2: # Try pngifying the jpegs # transcode_pngs(pike, jpegs, jpg_name, root, options) - transcode_pngs(pike, pngs, png_name, root, options) + transcode_pngs(pike, pngs, png_name, root, options, executor) jbig2_groups = extract_images_jbig2(pike, root, options) - convert_to_jbig2(pike, jbig2_groups, root, options) + convert_to_jbig2(pike, jbig2_groups, root, options, executor) target_file = output_file.with_suffix('.opt.pdf') pike.remove_unreferenced_resources() diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index 1ee82cf9..3733feaf 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -10,6 +10,7 @@ import atexit import logging import re from collections import defaultdict, namedtuple +from contextlib import ExitStack from decimal import Decimal from enum import Enum from functools import partial @@ -22,7 +23,7 @@ from warnings import warn import pikepdf from pikepdf import Object, Pdf, PdfMatrix -from ocrmypdf._concurrent import exec_progress_pool +from ocrmypdf._concurrent import Executor, SerialExecutor from ocrmypdf.exceptions import EncryptedPdfError, InputFileError from ocrmypdf.helpers import Resolution, available_cpu_count, pikepdf_enable_mmap from ocrmypdf.pdfinfo.layout import get_page_analysis, get_text_boxes @@ -592,12 +593,21 @@ def _pdf_pageinfo_sync_init(pdf: Pdf, infile: Path, pdfminer_loglevel): def _pdf_pageinfo_sync(args): pageno, thread_pdf, infile, check_pages, detailed_analysis = args pdf = thread_pdf if thread_pdf is not None else worker_pdf - page = PageInfo(pdf, pageno, infile, check_pages, detailed_analysis) - return page + with ExitStack() as stack: + if not pdf: # When called with SerialExecutor + pdf = stack.enter_context(pikepdf.open(infile)) + page = PageInfo(pdf, pageno, infile, check_pages, detailed_analysis) + return page def _pdf_pageinfo_concurrent( - pdf, infile, progbar, max_workers, check_pages, detailed_analysis=False + pdf, + executor: Executor, + infile, + progbar, + max_workers, + check_pages, + detailed_analysis=False, ): pages = [None] * len(pdf.pages) @@ -629,7 +639,7 @@ def _pdf_pageinfo_concurrent( (n, initial_pdf, infile, check_pages, detailed_analysis) for n in range(total) ) assert n_workers == 1 if use_threads else n_workers >= 1, "Not multithreadable" - exec_progress_pool( + executor( use_threads=use_threads, max_workers=n_workers, tqdm_kwargs=dict( @@ -829,10 +839,12 @@ class PdfInfo: def __init__( self, infile, + *, detailed_analysis: bool = False, progbar: bool = False, max_workers: int = None, check_pages=None, + executor: Executor = SerialExecutor(), ): self._infile = infile if check_pages is None: @@ -843,6 +855,7 @@ class PdfInfo: raise EncryptedPdfError() # Triggered by encryption with empty passwd self._pages = _pdf_pageinfo_concurrent( pdf, + executor, infile, progbar, max_workers, diff --git a/src/ocrmypdf/pluginspec.py b/src/ocrmypdf/pluginspec.py index 538be7a1..f3d5d6fa 100644 --- a/src/ocrmypdf/pluginspec.py +++ b/src/ocrmypdf/pluginspec.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, AbstractSet, Callable, Iterable, List, Optiona import pluggy +from ocrmypdf._concurrent import Executor from ocrmypdf.helpers import Resolution if TYPE_CHECKING: @@ -62,44 +63,16 @@ 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: +def get_executor() -> Executor: """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. + + OCRmyPDF's executors are analogous to the standard Python executors in + conconcurrent.futures, but they do not work the same way. """ diff --git a/tests/test_validation.py b/tests/test_validation.py index deed9769..f7ce5286 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -173,7 +173,7 @@ def test_false_action_store_true(): def test_no_progress_bar(progress_bar, resources): opts = make_opts(progress_bar=progress_bar, input_file=(resources / 'trivial.pdf')) plugin_manager = get_plugin_manager(opts.plugins) - with patch('ocrmypdf.builtin_plugins.concurrency.tqdm', autospec=True) as tqdmpatch: + with patch('ocrmypdf._concurrent.tqdm', autospec=True) as tqdmpatch: vd._check_options(opts, plugin_manager, set()) pdfinfo = PdfInfo(opts.input_file, progbar=opts.progress_bar) assert pdfinfo is not None