Improve passing of arguments to workers
The executor system was built around passing only a single argument to workers, which was always PageContext. For other tasks, all actual arguments were packed in a tuple, which meant we needed intermediate functions to unpack the tuple. The situation is now rationlized and resembles how Python handles argument passing to familiar multiprocessing tools.
This commit is contained in:
@@ -133,5 +133,5 @@ class SerialExecutor(Executor):
|
||||
): # pylint: disable=unused-argument
|
||||
with self.pbar_class(**progress_kwargs) as pbar:
|
||||
for args in task_arguments:
|
||||
result = task(args)
|
||||
result = task(*args)
|
||||
task_finished(result, pbar)
|
||||
|
||||
@@ -25,7 +25,7 @@ def available():
|
||||
return True
|
||||
|
||||
|
||||
def convert_group(*, cwd, infiles, out_prefix, threshold):
|
||||
def convert_group(cwd, infiles, out_prefix, threshold):
|
||||
args = [
|
||||
'jbig2',
|
||||
'-b',
|
||||
@@ -43,21 +43,9 @@ def convert_group(*, cwd, infiles, out_prefix, threshold):
|
||||
return proc
|
||||
|
||||
|
||||
def convert_group_mp(args):
|
||||
return convert_group(
|
||||
cwd=args[0], infiles=args[1], out_prefix=args[2], threshold=args[3]
|
||||
)
|
||||
|
||||
|
||||
def convert_single(*, cwd, infile, outfile, threshold):
|
||||
def convert_single(cwd, infile, outfile, threshold):
|
||||
args = ['jbig2', '--pdf', '-t', str(threshold), infile]
|
||||
with open(outfile, 'wb') as fstdout:
|
||||
proc = run(args, cwd=cwd, stdout=fstdout, stderr=PIPE)
|
||||
proc.check_returncode()
|
||||
return proc
|
||||
|
||||
|
||||
def convert_single_mp(args):
|
||||
return convert_single(
|
||||
cwd=args[0], infile=args[1], outfile=args[2], threshold=args[3]
|
||||
)
|
||||
|
||||
@@ -50,7 +50,3 @@ def quantize(input_file: Path, output_file: Path, quality_min: int, quality_max:
|
||||
if result.returncode == 0:
|
||||
# input_file could be the same as output_file, so we defer the write
|
||||
output_file.write_bytes(result.stdout)
|
||||
|
||||
|
||||
def quantize_mp(args):
|
||||
return quantize(*args)
|
||||
|
||||
@@ -53,6 +53,12 @@ class PdfContext:
|
||||
for n in range(npages):
|
||||
yield PageContext(self, n)
|
||||
|
||||
def get_page_context_args(self) -> Iterator[tuple[PageContext]]:
|
||||
"""Get all ``PageContext`` for this PDF packaged in tuple for args-splatting."""
|
||||
npages = len(self.pdfinfo)
|
||||
for n in range(npages):
|
||||
yield (PageContext(self, n),)
|
||||
|
||||
|
||||
class PageContext:
|
||||
"""Holds our context for a page.
|
||||
|
||||
@@ -87,7 +87,7 @@ def exec_hocr_to_ocr_pdf(context: PdfContext, executor: Executor) -> Sequence[st
|
||||
),
|
||||
worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS),
|
||||
task=exec_hocrtransform_sync,
|
||||
task_arguments=context.get_page_contexts(),
|
||||
task_arguments=context.get_page_context_args(),
|
||||
task_finished=graft_page,
|
||||
)
|
||||
|
||||
|
||||
@@ -128,7 +128,7 @@ def exec_concurrent(context: PdfContext, executor: Executor) -> Sequence[str]:
|
||||
),
|
||||
worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS),
|
||||
task=exec_page_sync,
|
||||
task_arguments=context.get_page_contexts(),
|
||||
task_arguments=context.get_page_context_args(),
|
||||
task_finished=update_page,
|
||||
)
|
||||
|
||||
|
||||
@@ -39,7 +39,6 @@ from ocrmypdf._validation import (
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
|
||||
def exec_page_hocr_sync(page_context: PageContext) -> HOCRResult:
|
||||
"""Execute a pipeline for a single page hOCR."""
|
||||
set_thread_pageno(page_context.pageno + 1)
|
||||
@@ -82,7 +81,7 @@ def exec_pdf_to_hocr(context: PdfContext, executor: Executor) -> None:
|
||||
),
|
||||
worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS),
|
||||
task=exec_page_hocr_sync,
|
||||
task_arguments=context.get_page_contexts(),
|
||||
task_arguments=context.get_page_context_args(),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -134,7 +134,7 @@ class StandardExecutor(Executor):
|
||||
initializer=initializer,
|
||||
initargs=(log_queue, worker_initializer, logging.getLogger("").level),
|
||||
) as executor:
|
||||
futures = [executor.submit(task, args) for args in task_arguments]
|
||||
futures = [executor.submit(task, *args) for args in task_arguments]
|
||||
try:
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
|
||||
@@ -94,7 +94,7 @@ def process_loop(
|
||||
|
||||
for args in task_args:
|
||||
try:
|
||||
result = task(args)
|
||||
result = task(*args)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
conn.send((MessageType.exception, e))
|
||||
break
|
||||
@@ -123,7 +123,7 @@ class LambdaExecutor(Executor):
|
||||
if use_threads and max_workers == 1:
|
||||
with self.pbar_class(**progress_kwargs) as pbar:
|
||||
for args in task_arguments:
|
||||
result = task(args)
|
||||
result = task(*args)
|
||||
task_finished(result, pbar)
|
||||
return
|
||||
|
||||
|
||||
@@ -390,10 +390,10 @@ def _produce_jbig2_images(
|
||||
|
||||
if options.jbig2_page_group_size > 1:
|
||||
jbig2_args = jbig2_group_args
|
||||
jbig2_convert = jbig2enc.convert_group_mp
|
||||
jbig2_convert = jbig2enc.convert_group
|
||||
else:
|
||||
jbig2_args = jbig2_single_args
|
||||
jbig2_convert = jbig2enc.convert_single_mp
|
||||
jbig2_convert = jbig2enc.convert_single
|
||||
|
||||
executor(
|
||||
use_threads=True,
|
||||
@@ -454,9 +454,9 @@ def convert_to_jbig2(
|
||||
)
|
||||
|
||||
|
||||
def _optimize_jpeg(args: tuple[Xref, Path, Path, int]) -> tuple[Xref, Path | None]:
|
||||
xref, in_jpg, opt_jpg, jpeg_quality = args
|
||||
|
||||
def _optimize_jpeg(
|
||||
xref: Xref, in_jpg: Path, opt_jpg: Path, jpeg_quality: int
|
||||
) -> tuple[Xref, Path | None]:
|
||||
with Image.open(in_jpg) as im:
|
||||
im.save(opt_jpg, optimize=True, quality=jpeg_quality)
|
||||
|
||||
@@ -515,8 +515,9 @@ def _find_deflatable_jpeg(
|
||||
return None
|
||||
|
||||
|
||||
def _deflate_jpeg(args: tuple[Pdf, threading.Lock, Xref, int]) -> tuple[Xref, bytes]:
|
||||
pdf, lock, xref, complevel = args
|
||||
def _deflate_jpeg(
|
||||
pdf: Pdf, lock: threading.Lock, xref: Xref, complevel: int
|
||||
) -> tuple[Xref, bytes]:
|
||||
with lock:
|
||||
xobj = pdf.get_object(xref, 0)
|
||||
try:
|
||||
@@ -651,7 +652,7 @@ def transcode_pngs(
|
||||
unit='image',
|
||||
disable=not options.progress_bar,
|
||||
),
|
||||
task=pngquant.quantize_mp,
|
||||
task=pngquant.quantize,
|
||||
task_arguments=pngquant_args(),
|
||||
)
|
||||
|
||||
|
||||
@@ -694,8 +694,7 @@ def _pdf_pageinfo_sync_init(pdf: Pdf, infile: Path, pdfminer_loglevel):
|
||||
atexit.register(on_process_close)
|
||||
|
||||
|
||||
def _pdf_pageinfo_sync(args):
|
||||
pageno, thread_pdf, infile, check_pages, detailed_analysis = args
|
||||
def _pdf_pageinfo_sync(pageno, thread_pdf, infile, check_pages, detailed_analysis):
|
||||
pdf = thread_pdf if thread_pdf is not None else worker_pdf
|
||||
with ExitStack() as stack:
|
||||
if not pdf: # When called with SerialExecutor
|
||||
|
||||
Reference in New Issue
Block a user