Further refactoring of concurrency concerns

This commit is contained in:
James R. Barlow
2020-04-26 03:49:20 -07:00
parent db3e75e33e
commit af3c3c6466
3 changed files with 69 additions and 110 deletions
+28 -2
View File
@@ -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
View File
@@ -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(),
+37 -59
View File
@@ -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