470 lines
17 KiB
Python
470 lines
17 KiB
Python
# SPDX-FileCopyrightText: 2026 James R. Barlow
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Batch job tracking and execution for the OCRmyPDF web interface.
|
|
|
|
A *batch* is one submission from the browser: several input files plus a
|
|
single set of OCR options. Each file becomes a *job* that runs ``ocrmypdf``
|
|
in a subprocess. Jobs from all batches share one bounded thread pool, so a
|
|
large batch cannot monopolize the machine.
|
|
|
|
Running OCRmyPDF out-of-process (rather than calling ``ocrmypdf.ocr()``) keeps
|
|
a crash, a hang, or a runaway memory allocation contained in a child process
|
|
that we can time out and kill.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import zipfile
|
|
from collections import deque
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
|
|
from webui.config import Settings
|
|
from webui.options import OcrOptions
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
#: ``ocrmypdf.exceptions.ExitCode`` values, phrased for someone who is looking
|
|
#: at a web page rather than a terminal. Kept as literals so the web
|
|
#: interface does not have to import ocrmypdf at startup.
|
|
EXIT_CODE_MESSAGES = {
|
|
1: "A bad or missing argument was supplied.", # bad_args
|
|
2: "The input file is not a valid PDF or image.", # input_file
|
|
3: "A required external program is missing.", # missing_dependency
|
|
4: "The output PDF was created but is not valid.", # invalid_output_pdf
|
|
5: "The file could not be read or written.", # file_access_error
|
|
6: ( # already_done_ocr
|
|
"This file already contains text. Choose “Skip them” or "
|
|
"“Redo the existing OCR” and try again."
|
|
),
|
|
7: "An external program failed while processing this file.", # child_process
|
|
8: "This PDF is encrypted. Remove the password and try again.", # encrypted_pdf
|
|
9: "The requested combination of options is not valid.", # invalid_config
|
|
10: ( # pdfa_conversion_failed
|
|
"The PDF was produced but could not be converted to PDF/A. "
|
|
"Try the standard PDF output format."
|
|
),
|
|
15: "An unexpected error occurred while processing this file.", # other_error
|
|
130: "Processing was interrupted.", # ctrl_c
|
|
}
|
|
|
|
#: Strip ANSI escapes that rich may emit even when not attached to a terminal.
|
|
_ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[a-zA-Z]")
|
|
|
|
|
|
class JobStatus(StrEnum):
|
|
"""Lifecycle of a single file."""
|
|
|
|
pending = "pending"
|
|
running = "running"
|
|
succeeded = "succeeded"
|
|
failed = "failed"
|
|
cancelled = "cancelled"
|
|
|
|
|
|
TERMINAL_STATUSES = frozenset(
|
|
{JobStatus.succeeded, JobStatus.failed, JobStatus.cancelled}
|
|
)
|
|
|
|
|
|
def safe_display_name(raw: str | None, fallback: str) -> str:
|
|
"""Reduce a browser-supplied filename to something safe to echo back.
|
|
|
|
The result is only ever used as display text and as the ``filename`` of a
|
|
download; it never touches the filesystem or a command line.
|
|
"""
|
|
if not raw:
|
|
return fallback
|
|
# Browsers may send a full path (notably older Edge and some drag sources).
|
|
name = raw.replace("\\", "/").split("/")[-1]
|
|
name = name.replace("\x00", "").strip().strip(".")
|
|
# Control characters and quotes would corrupt the Content-Disposition
|
|
# header; anything else is left alone so unicode filenames survive.
|
|
name = "".join(ch for ch in name if ch.isprintable() and ch not in '"\r\n')
|
|
name = name[:200].strip()
|
|
return name or fallback
|
|
|
|
|
|
@dataclass
|
|
class FileJob:
|
|
"""One input file and its result."""
|
|
|
|
id: str
|
|
display_name: str
|
|
input_path: Path
|
|
output_path: Path
|
|
size_bytes: int
|
|
status: JobStatus = JobStatus.pending
|
|
error: str | None = None
|
|
started_at: float | None = None
|
|
finished_at: float | None = None
|
|
log_lines: deque[str] = field(default_factory=lambda: deque(maxlen=200))
|
|
_process: subprocess.Popen | None = field(default=None, repr=False)
|
|
|
|
@property
|
|
def is_image(self) -> bool:
|
|
"""True when the input is an image rather than a PDF."""
|
|
return self.input_path.suffix.lower() != ".pdf"
|
|
|
|
@property
|
|
def output_name(self) -> str:
|
|
"""Filename offered to the browser for this job's result."""
|
|
stem = Path(self.display_name).stem or "output"
|
|
return f"{stem}.pdf"
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Serialize this job for the status API."""
|
|
return {
|
|
"id": self.id,
|
|
"name": self.display_name,
|
|
"output_name": self.output_name,
|
|
"size_bytes": self.size_bytes,
|
|
"status": self.status.value,
|
|
"error": self.error,
|
|
"output_size_bytes": (
|
|
self.output_path.stat().st_size
|
|
if self.status is JobStatus.succeeded and self.output_path.exists()
|
|
else None
|
|
),
|
|
"duration_seconds": (
|
|
round(self.finished_at - self.started_at, 1)
|
|
if self.started_at and self.finished_at
|
|
else None
|
|
),
|
|
"log_tail": list(self.log_lines)[-8:],
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class Batch:
|
|
"""One submission: several files sharing a set of options."""
|
|
|
|
id: str
|
|
directory: Path
|
|
options: OcrOptions
|
|
jobs: list[FileJob]
|
|
created_at: float = field(default_factory=time.time)
|
|
cancelled: bool = False
|
|
_futures: list[Future] = field(default_factory=list, repr=False)
|
|
|
|
@property
|
|
def finished(self) -> bool:
|
|
"""True once no job in this batch can still change state."""
|
|
return all(job.status in TERMINAL_STATUSES for job in self.jobs)
|
|
|
|
@property
|
|
def succeeded_jobs(self) -> list[FileJob]:
|
|
"""Jobs whose output is on disk and downloadable."""
|
|
return [job for job in self.jobs if job.status is JobStatus.succeeded]
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Serialize this batch and all of its jobs for the status API."""
|
|
jobs = [job.to_dict() for job in self.jobs]
|
|
counts = {status.value: 0 for status in JobStatus}
|
|
for job in self.jobs:
|
|
counts[job.status.value] += 1
|
|
return {
|
|
"id": self.id,
|
|
"created_at": self.created_at,
|
|
"finished": self.finished,
|
|
"cancelled": self.cancelled,
|
|
"counts": counts,
|
|
"total": len(self.jobs),
|
|
"completed": sum(counts[status.value] for status in TERMINAL_STATUSES),
|
|
"downloadable": len(self.succeeded_jobs),
|
|
"files": jobs,
|
|
}
|
|
|
|
|
|
class JobManager:
|
|
"""Owns the batch registry, the worker pool, and the reaper thread."""
|
|
|
|
def __init__(self, settings: Settings):
|
|
"""Create the pool and start the background reaper."""
|
|
self.settings = settings
|
|
self._batches: dict[str, Batch] = {}
|
|
self._lock = threading.RLock()
|
|
self._pool = ThreadPoolExecutor(
|
|
max_workers=settings.workers, thread_name_prefix="ocr-worker"
|
|
)
|
|
self._stopping = threading.Event()
|
|
self._reaper = threading.Thread(
|
|
target=self._reap_loop, name="ocr-reaper", daemon=True
|
|
)
|
|
settings.work_dir.mkdir(parents=True, exist_ok=True)
|
|
self._reaper.start()
|
|
|
|
# -- lifecycle ---------------------------------------------------------
|
|
|
|
def shutdown(self) -> None:
|
|
"""Stop accepting work, cancel everything, and clean the work dir."""
|
|
self._stopping.set()
|
|
with self._lock:
|
|
batch_ids = list(self._batches)
|
|
for batch_id in batch_ids:
|
|
self.delete_batch(batch_id)
|
|
self._pool.shutdown(wait=False, cancel_futures=True)
|
|
|
|
# -- batch management --------------------------------------------------
|
|
|
|
def create_batch(self, options: OcrOptions) -> Batch:
|
|
"""Register an empty batch and create its scratch directory."""
|
|
batch_id = uuid.uuid4().hex
|
|
directory = self.settings.work_dir / batch_id
|
|
(directory / "in").mkdir(parents=True, exist_ok=True)
|
|
(directory / "out").mkdir(parents=True, exist_ok=True)
|
|
batch = Batch(id=batch_id, directory=directory, options=options, jobs=[])
|
|
with self._lock:
|
|
self._batches[batch_id] = batch
|
|
return batch
|
|
|
|
def add_file(self, batch: Batch, display_name: str, suffix: str) -> FileJob:
|
|
"""Allocate an on-disk slot for one uploaded file.
|
|
|
|
The caller streams bytes into ``job.input_path`` and then calls
|
|
:meth:`start`. Filenames on disk are generated, never derived from
|
|
user input.
|
|
"""
|
|
index = len(batch.jobs)
|
|
job_id = f"{index:04d}"
|
|
job = FileJob(
|
|
id=job_id,
|
|
display_name=display_name,
|
|
input_path=batch.directory / "in" / f"{job_id}{suffix}",
|
|
output_path=batch.directory / "out" / f"{job_id}.pdf",
|
|
size_bytes=0,
|
|
)
|
|
job.log_lines = deque(maxlen=self.settings.max_log_lines)
|
|
batch.jobs.append(job)
|
|
return job
|
|
|
|
def start(self, batch: Batch) -> None:
|
|
"""Queue every file in the batch for processing."""
|
|
with self._lock:
|
|
batch._futures = [
|
|
self._pool.submit(self._run_job, batch, job) for job in batch.jobs
|
|
]
|
|
|
|
def get_batch(self, batch_id: str) -> Batch | None:
|
|
"""Return a batch by id, or None if it is unknown or expired."""
|
|
with self._lock:
|
|
return self._batches.get(batch_id)
|
|
|
|
def get_job(self, batch_id: str, job_id: str) -> tuple[Batch, FileJob] | None:
|
|
"""Return the (batch, job) pair for an id pair, or None."""
|
|
batch = self.get_batch(batch_id)
|
|
if batch is None:
|
|
return None
|
|
for job in batch.jobs:
|
|
if job.id == job_id:
|
|
return batch, job
|
|
return None
|
|
|
|
def delete_batch(self, batch_id: str) -> bool:
|
|
"""Cancel any running work and remove the batch from disk."""
|
|
with self._lock:
|
|
batch = self._batches.pop(batch_id, None)
|
|
if batch is None:
|
|
return False
|
|
batch.cancelled = True
|
|
for future in batch._futures:
|
|
future.cancel()
|
|
for job in batch.jobs:
|
|
proc = job._process
|
|
if proc is not None and proc.poll() is None:
|
|
self._terminate(proc)
|
|
if job.status in (JobStatus.pending, JobStatus.running):
|
|
job.status = JobStatus.cancelled
|
|
shutil.rmtree(batch.directory, ignore_errors=True)
|
|
return True
|
|
|
|
# -- results -----------------------------------------------------------
|
|
|
|
def build_zip(self, batch: Batch) -> Path | None:
|
|
"""Bundle every successful output into one archive.
|
|
|
|
The archive is written once and reused; a batch is immutable after its
|
|
jobs finish.
|
|
"""
|
|
jobs = batch.succeeded_jobs
|
|
if not jobs:
|
|
return None
|
|
zip_path = batch.directory / "ocrmypdf-results.zip"
|
|
if zip_path.exists():
|
|
return zip_path
|
|
|
|
used: dict[str, int] = {}
|
|
tmp_path = zip_path.with_suffix(".zip.part")
|
|
# ZIP_STORED: the members are already-compressed PDFs, so deflating
|
|
# them again costs CPU we'd rather spend on OCR.
|
|
with zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_STORED) as archive:
|
|
for job in jobs:
|
|
name = job.output_name
|
|
if name in used:
|
|
used[name] += 1
|
|
stem = Path(name).stem
|
|
name = f"{stem} ({used[name]}).pdf"
|
|
else:
|
|
used[name] = 0
|
|
archive.write(job.output_path, arcname=name)
|
|
tmp_path.replace(zip_path)
|
|
return zip_path
|
|
|
|
# -- execution ---------------------------------------------------------
|
|
|
|
def _run_job(self, batch: Batch, job: FileJob) -> None:
|
|
if batch.cancelled or self._stopping.is_set():
|
|
job.status = JobStatus.cancelled
|
|
return
|
|
|
|
job.status = JobStatus.running
|
|
job.started_at = time.time()
|
|
args = [
|
|
sys.executable,
|
|
"-u",
|
|
"-m",
|
|
"ocrmypdf",
|
|
*batch.options.to_args(jobs=self.settings.ocr_jobs, is_image=job.is_image),
|
|
"--",
|
|
str(job.input_path),
|
|
str(job.output_path),
|
|
]
|
|
|
|
try:
|
|
self._execute(job, args)
|
|
except Exception: # noqa: BLE001 - a worker must never die silently
|
|
log.exception("Unexpected failure processing %s", job.display_name)
|
|
job.status = JobStatus.failed
|
|
job.error = "An internal error occurred while processing this file."
|
|
finally:
|
|
job.finished_at = time.time()
|
|
job._process = None
|
|
if job.status is JobStatus.running: # defensive
|
|
job.status = JobStatus.failed
|
|
job.error = job.error or "Processing ended unexpectedly."
|
|
if job.status is not JobStatus.succeeded:
|
|
job.output_path.unlink(missing_ok=True)
|
|
# The input is not needed once we have a result; drop it early so
|
|
# a big batch doesn't hold twice its size on disk.
|
|
job.input_path.unlink(missing_ok=True)
|
|
|
|
def _execute(self, job: FileJob, args: list[str]) -> None:
|
|
proc = subprocess.Popen(
|
|
args,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.PIPE,
|
|
# Never inherit the parent's stdin; ocrmypdf must not block on it.
|
|
stdin=subprocess.DEVNULL,
|
|
text=True,
|
|
errors="replace",
|
|
bufsize=1,
|
|
)
|
|
job._process = proc
|
|
deadline = time.monotonic() + self.settings.job_timeout_seconds
|
|
timed_out = False
|
|
|
|
assert proc.stderr is not None
|
|
for line in proc.stderr:
|
|
clean = _ANSI_RE.sub("", line).rstrip()
|
|
if clean:
|
|
job.log_lines.append(clean)
|
|
if time.monotonic() > deadline:
|
|
timed_out = True
|
|
self._terminate(proc)
|
|
break
|
|
|
|
try:
|
|
returncode = proc.wait(timeout=30)
|
|
except subprocess.TimeoutExpired:
|
|
self._terminate(proc)
|
|
returncode = proc.wait(timeout=10)
|
|
|
|
if timed_out:
|
|
job.status = JobStatus.failed
|
|
job.error = (
|
|
f"Timed out after {self.settings.job_timeout_seconds // 60} minutes."
|
|
)
|
|
return
|
|
|
|
if returncode != 0:
|
|
job.status = JobStatus.failed
|
|
job.error = EXIT_CODE_MESSAGES.get(
|
|
returncode, f"ocrmypdf exited with code {returncode}."
|
|
)
|
|
return
|
|
|
|
if not job.output_path.exists() or job.output_path.stat().st_size == 0:
|
|
job.status = JobStatus.failed
|
|
job.error = "ocrmypdf reported success but produced no output."
|
|
return
|
|
|
|
job.status = JobStatus.succeeded
|
|
|
|
@staticmethod
|
|
def _terminate(proc: subprocess.Popen) -> None:
|
|
"""Ask a child to stop, then insist."""
|
|
try:
|
|
proc.terminate()
|
|
proc.wait(timeout=10)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
except OSError:
|
|
pass
|
|
|
|
# -- retention ---------------------------------------------------------
|
|
|
|
def _reap_loop(self) -> None:
|
|
"""Delete batches that have outlived their TTL."""
|
|
while not self._stopping.wait(timeout=60):
|
|
try:
|
|
self.sweep()
|
|
except Exception: # noqa: BLE001 - the reaper must keep running
|
|
log.exception("Error during batch sweep")
|
|
|
|
def sweep(self, now: float | None = None) -> int:
|
|
"""Remove expired batches. Returns the number deleted."""
|
|
now = now if now is not None else time.time()
|
|
ttl = self.settings.batch_ttl_seconds
|
|
with self._lock:
|
|
expired = [
|
|
batch.id
|
|
for batch in self._batches.values()
|
|
# Unfinished batches are kept alive by their age too, so a
|
|
# wedged job can never pin a directory forever.
|
|
if now - batch.created_at > ttl
|
|
]
|
|
for batch_id in expired:
|
|
log.info("Reaping expired batch %s", batch_id)
|
|
self.delete_batch(batch_id)
|
|
self._sweep_orphans(now)
|
|
return len(expired)
|
|
|
|
def _sweep_orphans(self, now: float) -> None:
|
|
"""Remove directories left behind by a previous process."""
|
|
with self._lock:
|
|
known = set(self._batches)
|
|
try:
|
|
entries = list(self.settings.work_dir.iterdir())
|
|
except OSError:
|
|
return
|
|
for entry in entries:
|
|
if entry.name in known or not entry.is_dir():
|
|
continue
|
|
try:
|
|
age = now - entry.stat().st_mtime
|
|
except OSError:
|
|
continue
|
|
if age > self.settings.batch_ttl_seconds:
|
|
shutil.rmtree(entry, ignore_errors=True)
|