Further refactoring of concurrency concerns
This commit is contained in:
@@ -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)
|
||||
|
||||
+4
-49
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user