diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py new file mode 100644 index 00000000..d6e49e91 --- /dev/null +++ b/src/ocrmypdf/_concurrent.py @@ -0,0 +1,110 @@ +# © 2020 James R. Barlow: github.com/jbarlow83 +# +# This file is part of OCRmyPDF. +# +# OCRmyPDF is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# OCRmyPDF is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with OCRmyPDF. If not, see . + +import logging +import logging.handlers +import multiprocessing +import os +import signal +import sys +import threading +from multiprocessing import Pool as ProcessPool +from multiprocessing.dummy import Pool as ThreadPool +from pathlib import Path + +from tqdm import tqdm + + +def log_listener(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 = queue.get() + if record is None: + break + logger = logging.getLogger(record.name) + logger.handle(record) + except Exception: + import traceback + + print("Logging problem", file=sys.stderr) + traceback.print_exc(file=sys.stderr) + + +def exec_progress_pool( + *, + use_threads, + max_workers, + tqdm_kwargs, + task_initializer=None, + task_initargs=None, + task=None, + task_arguments=None, + task_finished=None, +): + log_queue = multiprocessing.Queue(-1) + listener = threading.Thread(target=log_listener, args=(log_queue,)) + + if use_threads: + pool_class = ThreadPool + else: + pool_class = ProcessPool + listener.start() + + with tqdm(**tqdm_kwargs) as pbar: + pool = pool_class( + processes=max_workers, + initializer=task_initializer, + initargs=(log_queue, *task_initargs), + ) + try: + results = pool.imap_unordered(task, task_arguments) + while True: + try: + result = results.next() + task_finished(result, pbar) + except StopIteration: + break + 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() diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index 07c355ab..b20ccea8 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -21,14 +21,13 @@ import multiprocessing import os import signal import sys -import threading from collections import namedtuple from pathlib import Path from tempfile import mkdtemp import PIL -from tqdm import tqdm +from ._concurrent import exec_progress_pool from ._graft import OcrGrafter from ._jobcontext import PDFContext, cleanup_working_files, make_logger from ._pipeline import ( @@ -253,63 +252,35 @@ def exec_concurrent(context): context.log.info("Using Tesseract OpenMP thread limit %d", tess_threads) if context.options.use_threads: - from multiprocessing.dummy import Pool - initializer = worker_thread_init else: - Pool = multiprocessing.Pool initializer = worker_init sidecars = [None] * len(context.pdfinfo) ocrgraft = OcrGrafter(context) - log_queue = multiprocessing.Queue(-1) - listener = threading.Thread(target=log_listener, args=(log_queue,)) - listener.start() - with tqdm( - total=(2 * len(context.pdfinfo)), - desc='OCR', - unit='page', - unit_scale=0.5, - disable=not context.options.progress_bar, - ) as pbar: - pool = Pool( - processes=max_workers, - initializer=initializer, - initargs=(log_queue, PIL.Image.MAX_IMAGE_PIXELS), - ) - try: - results = pool.imap_unordered(exec_page_sync, context.get_page_contexts()) - while True: - try: - page_result = results.next() - sidecars[page_result.pageno] = page_result.text - pbar.update() - ocrgraft.graft_page(page_result) - pbar.update() - except StopIteration: - break - 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() + def update_page(result, pbar): + sidecars[result.pageno] = result.text + pbar.update() + ocrgraft.graft_page(result) + pbar.update() - listener.join() + exec_progress_pool( + use_threads=context.options.use_threads, + max_workers=max_workers, + tqdm_kwargs=dict( + total=(2 * len(context.pdfinfo)), + desc='OCR', + unit='page', + unit_scale=0.5, + disable=not context.options.progress_bar, + ), + task_initializer=initializer, + task_initargs=(PIL.Image.MAX_IMAGE_PIXELS,), + task=exec_page_sync, + task_arguments=context.get_page_contexts(), + task_finished=update_page, + ) # Output sidecar text if context.options.sidecar: diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index b50b34c7..85259c2f 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -641,7 +641,7 @@ def _pdf_pageinfo_concurrent(pdf, infile, pages_xml, detailed_analysis, progbar) global worker_pdf worker_pdf = pdf pool = Pool( - processes=4, # max_workers, + processes=1, # max_workers, initializer=_pdf_pageinfo_sync_init, initargs=tuple(), )