"""System-friendly multiprocessing batch OCR with structured timing logs.""" from __future__ import annotations import argparse import logging import os import random import sys import time from logging.handlers import QueueHandler, QueueListener from multiprocessing import Manager, Pool, cpu_count from pathlib import Path from ocr_logging import default_log_path, setup_run_logger PROJECT_ROOT = Path(__file__).resolve().parent _WORKER_LOG_QUEUE = None _WORKER_INIT_METRICS: dict = {} def _worker_logger() -> logging.Logger: logger = logging.getLogger(f"ocr.batch.worker.{os.getpid()}") if logger.handlers: return logger logger.setLevel(logging.INFO) logger.propagate = False if _WORKER_LOG_QUEUE is not None: logger.addHandler(QueueHandler(_WORKER_LOG_QUEUE)) return logger def _init_worker(threads: int, stagger_max: float, log_queue) -> None: """Stagger startup, lower process priority, and load one model per worker.""" global _pipeline, _WORKER_LOG_QUEUE, _WORKER_INIT_METRICS _WORKER_LOG_QUEUE = log_queue logger = _worker_logger() worker_started = time.perf_counter() delay = random.uniform(0, stagger_max) logger.info("WORKER_START threads=%d stagger_delay_seconds=%.3f", threads, delay) time.sleep(delay) import_started = time.perf_counter() from paddle import core core.set_num_threads(threads) import_seconds = time.perf_counter() - import_started try: import psutil process = psutil.Process() if sys.platform == "win32": process.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS) else: process.nice(10) priority_status = "lowered" except Exception as exc: priority_status = f"unchanged:{type(exc).__name__}" model_started = time.perf_counter() from paddleocr import PaddleOCRVL _pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu") model_seconds = time.perf_counter() - model_started startup_total = time.perf_counter() - worker_started _WORKER_INIT_METRICS = { "pid": os.getpid(), "threads": threads, "stagger_delay_seconds": round(delay, 3), "import_seconds": round(import_seconds, 3), "model_init_seconds": round(model_seconds, 3), "startup_total_seconds": round(startup_total, 3), "priority": priority_status, } logger.info( "WORKER_READY threads=%d import_seconds=%.3f model_init_seconds=%.3f startup_total_seconds=%.3f priority=%s", threads, import_seconds, model_seconds, startup_total, priority_status, ) def _ocr_task(image_path: str) -> dict: global _pipeline, _WORKER_INIT_METRICS logger = _worker_logger() started = time.perf_counter() logger.info("IMAGE_START path=%s", image_path) try: result = _pipeline.predict(image_path) elapsed = time.perf_counter() - started first = result[0] blocks = [ {"label": block.label, "bbox": block.bbox, "content": block.content} for block in first["parsing_res_list"] if block.content.strip() ] response = { "path": image_path, "status": "completed", "elapsed": round(elapsed, 3), "width": first.get("width"), "height": first.get("height"), "layout_boxes": len(first["layout_det_res"]["boxes"]), "parsed_blocks": len(first["parsing_res_list"]), "blocks": blocks, "worker_pid": os.getpid(), "worker_init": _WORKER_INIT_METRICS, } logger.info( "IMAGE_COMPLETED path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d", image_path, elapsed, response["width"], response["height"], response["layout_boxes"], response["parsed_blocks"], len(blocks), ) return response except Exception as exc: elapsed = time.perf_counter() - started logger.exception("IMAGE_FAILED path=%s seconds=%.3f error=%s", image_path, elapsed, exc) return { "path": image_path, "status": "failed", "elapsed": round(elapsed, 3), "error": f"{type(exc).__name__}: {exc}", "blocks": [], "worker_pid": os.getpid(), "worker_init": _WORKER_INIT_METRICS, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="批量 OCR — 多进程并行（系统友好版）", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("dir", type=Path, help="图片目录") parser.add_argument("--workers", type=int, default=2, help="并行进程数") parser.add_argument("--threads", type=int, default=None, help="每进程线程数") parser.add_argument("--stagger", type=float, default=15.0, help="Worker 启动错峰窗口秒数") parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") parser.add_argument("--verbose", action="store_true", help="输出详细日志") parser.add_argument("--no-result", action="store_true", help="不记录 OCR 文本块") return parser.parse_args() def main() -> int: program_started = time.perf_counter() args = parse_args() image_dir = args.dir.expanduser().resolve() log_file = args.log_file or default_log_path(PROJECT_ROOT, "batch", image_dir.name, device="cpu") logger = setup_run_logger("ocr.batch.main", log_file, verbose=args.verbose) if not image_dir.is_dir(): logger.error("INPUT_DIRECTORY_NOT_FOUND path=%s", image_dir) return 1 if args.workers < 1 or args.stagger < 0: logger.error("INVALID_ARGUMENT workers=%d stagger=%.3f", args.workers, args.stagger) return 2 scan_started = time.perf_counter() extensions = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.tif", "*.webp") images = sorted(path for extension in extensions for path in image_dir.glob(extension)) scan_seconds = time.perf_counter() - scan_started if not images: logger.error("NO_IMAGES_FOUND path=%s scan_seconds=%.3f", image_dir, scan_seconds) return 1 total_cores = cpu_count() workers = min(args.workers, len(images)) threads = args.threads or max(1, (total_cores - 1) // workers) if threads < 1: logger.error("INVALID_ARGUMENT threads=%d", threads) return 2 total_cpu_used = workers * threads estimated_mem = workers * 2.0 + 2 try: import psutil available_gb = psutil.virtual_memory().available / (1024**3) except ImportError: available_gb = None logger.info( "PROGRAM_STARTED directory=%s image_count=%d scan_seconds=%.3f workers=%d threads_per_worker=%d total_cores=%d planned_threads=%d reserved_cores=%d stagger_seconds=%.3f estimated_memory_gb=%.1f available_memory_gb=%s", image_dir, len(images), scan_seconds, workers, threads, total_cores, total_cpu_used, max(0, total_cores - total_cpu_used), args.stagger, estimated_mem, f"{available_gb:.1f}" if available_gb is not None else "unknown", ) if available_gb is not None and available_gb <= estimated_mem: logger.warning( "MEMORY_PRESSURE estimated_memory_gb=%.1f available_memory_gb=%.1f recommendation=reduce_workers", estimated_mem, available_gb, ) response = input("可用内存可能不足，是否继续？[y/N] ").strip().lower() if response != "y": logger.warning("PROGRAM_CANCELLED_BY_USER") return 0 pool_started = time.perf_counter() results: list[dict] = [] try: with Manager() as manager: log_queue = manager.Queue() listener_handlers = tuple(logger.handlers) listener = QueueListener(log_queue, *listener_handlers, respect_handler_level=True) listener.start() try: with Pool( processes=workers, initializer=_init_worker, initargs=(threads, args.stagger, log_queue), ) as pool: for completed, result in enumerate( pool.imap_unordered(_ocr_task, [str(path) for path in images], chunksize=1), start=1, ): results.append(result) logger.info( "BATCH_PROGRESS completed=%d total=%d path=%s status=%s image_seconds=%.3f worker_pid=%s", completed, len(images), result["path"], result["status"], result["elapsed"], result.get("worker_pid"), ) finally: listener.stop() except KeyboardInterrupt: logger.warning("PROGRAM_INTERRUPTED elapsed_seconds=%.3f", time.perf_counter() - program_started) return 130 except Exception as exc: logger.exception("POOL_FAILED error=%s elapsed_seconds=%.3f", exc, time.perf_counter() - program_started) return 1 pool_seconds = time.perf_counter() - pool_started completed_results = [result for result in results if result["status"] == "completed"] failed_results = [result for result in results if result["status"] == "failed"] worker_metrics = { result["worker_pid"]: result.get("worker_init", {}) for result in results if result.get("worker_pid") is not None } worker_model_init_total = sum( metrics.get("model_init_seconds", 0.0) for metrics in worker_metrics.values() ) worker_model_init_average = ( worker_model_init_total / len(worker_metrics) if worker_metrics else 0.0 ) serial_estimate = sum(result["elapsed"] for result in results) average = serial_estimate / len(results) if results else 0.0 speedup = serial_estimate / pool_seconds if pool_seconds else 0.0 program_total = time.perf_counter() - program_started for worker_pid, metrics in sorted(worker_metrics.items()): logger.info( "WORKER_SUMMARY pid=%s threads=%s stagger_delay_seconds=%s import_seconds=%s model_init_seconds=%s startup_total_seconds=%s priority=%s", worker_pid, metrics.get("threads"), metrics.get("stagger_delay_seconds"), metrics.get("import_seconds"), metrics.get("model_init_seconds"), metrics.get("startup_total_seconds"), metrics.get("priority"), ) for result in sorted(results, key=lambda item: item["path"]): if result["status"] == "completed": logger.info( "IMAGE_SUMMARY path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d", result["path"], result["elapsed"], result["width"], result["height"], result["layout_boxes"], result["parsed_blocks"], ) if not args.no_result: for index, block in enumerate(result["blocks"], start=1): logger.info( "OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s", result["path"], index, block["label"], block["bbox"], block["content"].replace("\r", "").replace("\n", "\\n"), ) else: logger.error("IMAGE_SUMMARY path=%s status=failed error=%s", result["path"], result["error"]) logger.info( "BATCH_SUMMARY image_count=%d completed=%d failed=%d scan_seconds=%.3f pool_seconds=%.3f worker_count=%d worker_model_init_total_seconds=%.3f worker_model_init_average_seconds=%.3f serial_estimate_seconds=%.3f average_image_seconds=%.3f speedup=%.3f program_total_seconds=%.3f workers=%d threads_per_worker=%d log=%s", len(images), len(completed_results), len(failed_results), scan_seconds, pool_seconds, len(worker_metrics), worker_model_init_total, worker_model_init_average, serial_estimate, average, speedup, program_total, workers, threads, log_file.resolve(), ) return 0 if not failed_results else 3 if __name__ == "__main__": raise SystemExit(main())