From af3c3c646606993d3fa6f16486844f0ae2a544e0 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Sat, 11 Apr 2020 00:49:36 -0700 Subject: [PATCH] Further refactoring of concurrency concerns --- src/ocrmypdf/_concurrent.py | 30 ++++++++++- src/ocrmypdf/_sync.py | 53 ++------------------ src/ocrmypdf/pdfinfo/info.py | 96 ++++++++++++++---------------------- 3 files changed, 69 insertions(+), 110 deletions(-) diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py index d6e49e91..724e6484 100644 --- a/src/ocrmypdf/_concurrent.py +++ b/src/ocrmypdf/_concurrent.py @@ -53,6 +53,27 @@ def log_listener(queue): traceback.print_exc(file=sys.stderr) +def process_init(queue, userfn, *userargs): + """Initialize a process pool worker""" + + # Ignore SIGINT (our parent process will kill us gracefully) + signal.signal(signal.SIGINT, signal.SIG_IGN) + + # Reconfigure the root logger for this process to send all messages to a queue + h = logging.handlers.QueueHandler(queue) + root = logging.getLogger() + root.handlers = [] + root.addHandler(h) + + if userfn: + userfn(*userargs) + + +def thread_init(_queue, userfn, *userargs): + if userfn: + userfn(*userargs) + + def exec_progress_pool( *, use_threads, @@ -67,17 +88,22 @@ def exec_progress_pool( log_queue = multiprocessing.Queue(-1) listener = threading.Thread(target=log_listener, args=(log_queue,)) + if not task_initargs: + task_initargs = tuple() + if use_threads: pool_class = ThreadPool + initializer = thread_init else: pool_class = ProcessPool + initializer = process_init listener.start() with tqdm(**tqdm_kwargs) as pbar: pool = pool_class( processes=max_workers, - initializer=task_initializer, - initargs=(log_queue, *task_initargs), + initializer=initializer, + initargs=(log_queue, task_initializer, *task_initargs), ) try: results = pool.imap_unordered(task, task_arguments) diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py index 7996f869..998b67ab 100644 --- a/src/ocrmypdf/_sync.py +++ b/src/ocrmypdf/_sync.py @@ -199,53 +199,13 @@ def post_process(pdf_file, context): return optimize_pdf(pdf_out, context) -def worker_init(queue, max_pixels): - """Initialize a process pool worker""" - - # Ignore SIGINT (our parent process will kill us gracefully) - signal.signal(signal.SIGINT, signal.SIG_IGN) - - # Reconfigure the root logger for this process to send all messages to a queue - h = logging.handlers.QueueHandler(queue) - root = logging.getLogger() - root.handlers = [] - root.addHandler(h) - +def worker_init(max_pixels): # In Windows, child process will not inherit our change to this value in - # the parent process, so ensure workers get it set + # the parent process, so ensure workers get it set. Not needed when running + # threaded, but harmless to set again. PIL.Image.MAX_IMAGE_PIXELS = max_pixels -def worker_thread_init(_queue, max_pixels): - # This is probably not needed since threads should all see the same memory, - # but done for consistency. - PIL.Image.MAX_IMAGE_PIXELS = max_pixels - - -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_concurrent(context): """Execute the pipeline concurrently""" @@ -273,11 +233,6 @@ def exec_concurrent(context): if tess_threads > 1: log.info("Using Tesseract OpenMP thread limit %d", tess_threads) - if context.options.use_threads: - initializer = worker_thread_init - else: - initializer = worker_init - sidecars = [None] * len(context.pdfinfo) ocrgraft = OcrGrafter(context) @@ -297,7 +252,7 @@ def exec_concurrent(context): unit_scale=0.5, disable=not context.options.progress_bar, ), - task_initializer=initializer, + task_initializer=worker_init, task_initargs=(PIL.Image.MAX_IMAGE_PIXELS,), task=exec_page_sync, task_arguments=context.get_page_contexts(), diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index b228e769..97ddb1d8 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -23,15 +23,14 @@ from collections import defaultdict, namedtuple from decimal import Decimal from enum import Enum from math import hypot, isclose -from multiprocessing import Pool from os import PathLike, fspath from pathlib import Path from warnings import warn import pikepdf from pikepdf import PdfMatrix -from tqdm import tqdm +from ocrmypdf._concurrent import exec_progress_pool from ocrmypdf.exceptions import EncryptedPdfError from ocrmypdf.exec import ghostscript from ocrmypdf.helpers import Resolution @@ -618,71 +617,50 @@ worker_pdf = None def _pdf_pageinfo_sync(args): - global worker_pdf - pageno, infile, xmltext, detailed_analysis = args page = PageInfo(worker_pdf, pageno, infile, xmltext, detailed_analysis) return page -def _pdf_pageinfo_sync_init(): - pass - - def _pdf_pageinfo_concurrent(pdf, infile, pages_xml, detailed_analysis, progbar): pages = [None] * len(pdf.pages) - with tqdm( - total=len(pdf.pages), desc="Scan", unit='page', disable=not progbar - ) as pbar: - global worker_pdf - worker_pdf = pdf - pool = Pool( - processes=1, # max_workers, - initializer=_pdf_pageinfo_sync_init, - initargs=tuple(), - ) - contexts = ( - (n, infile, pages_xml[n] if pages_xml else None, detailed_analysis) - for n in range(len(pdf.pages)) - ) - try: - results = pool.imap_unordered(_pdf_pageinfo_sync, contexts, chunksize=1) - while True: - try: - # page = results.next() - page = next(results) - pages[page.pageno] = page - pbar.update() - except StopIteration: - break - except KeyboardInterrupt: - pool.terminate() - 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() - # for n, _ in tqdm( - # enumerate(pdf.pages), - # total=len(pdf.pages), - # desc="Scan", - # unit='page', - # disable=not progbar, - # ): - # page_xml = pages_xml[n] if pages_xml else None - # page = PageInfo(pdf, n, infile, page_xml, detailed_analysis) - # pages.append(page) + def update_pageinfo(result, pbar): + page = result + pages[page.pageno] = page + pbar.update() + + contexts = ( + (n, infile, pages_xml[n] if pages_xml else None, detailed_analysis) + for n in range(len(pdf.pages)) + ) + global worker_pdf + worker_pdf = pdf + + if os.name == 'nt': + # We can't parallelize on Windows, because Windows cannot fork. + # We are trying to fork, then take advantage of the preloaded pikepdf.Pdf + # object in memory to save time reloading it, hence the silly global + # variable. Hey, it works. Threads are not helpful here because they + # will all just fight over the lock. So on Windows just run sequentially. + use_threads = True + max_workers = 1 + else: + use_threads = False + max_workers = min(len(pages), 16) + + exec_progress_pool( + use_threads=use_threads, + max_workers=1, + tqdm_kwargs=dict( + total=len(pdf.pages), desc="Scan", unit='page', disable=not progbar + ), + task_initializer=None, + task_initargs=None, + task=_pdf_pageinfo_sync, + task_arguments=contexts, + task_finished=update_pageinfo, + ) return pages