"""Device validation and lazy PaddleOCR-VL pipeline creation.""" from __future__ import annotations import logging import os import platform import time from dataclasses import dataclass from typing import Any @dataclass class RuntimeConfig: device: str threads: int | None = None device_id: int = 0 class PipelineProvider: """Create the large OCR pipeline only when a command actually needs it.""" def __init__(self, config: RuntimeConfig, logger: logging.Logger): self.config = config self.logger = logger self._pipeline: Any | None = None self._paddle: Any | None = None self._device_name: str | None = None self.import_seconds = 0.0 self.setup_seconds = 0.0 self.model_init_seconds = 0.0 @property def resolved_device(self) -> str: return "cpu" if self.config.device == "cpu" else f"gpu:{self.config.device_id}" def prepare(self) -> None: """Validate and configure Paddle without loading the OCR model.""" if self._paddle is not None: return started = time.perf_counter() try: import paddle except ImportError as exc: package = "paddlepaddle" if self.config.device == "cpu" else "paddlepaddle-gpu" raise RuntimeError(f"当前子项目未安装 {package}") from exc self.import_seconds = time.perf_counter() - started self._paddle = paddle setup_started = time.perf_counter() if self.config.device == "cpu": from paddle import core total_cores = os.cpu_count() or 4 threads = self.config.threads or max(1, total_cores - 2) if threads < 1: raise ValueError("CPU 线程数必须大于等于 1") self.config.threads = threads core.set_num_threads(threads) self._device_name = platform.processor() or "CPU" paddle.set_device("cpu") self.logger.info( "CPU_CONFIGURED threads=%d total_cores=%d reserved_cores=%d", threads, total_cores, max(0, total_cores - threads), ) else: if not paddle.is_compiled_with_cuda(): raise RuntimeError("当前 PaddlePaddle 不是 CUDA 构建;不会回退到 CPU") try: device_count = paddle.device.cuda.device_count() except Exception as exc: raise RuntimeError(f"无法查询 CUDA 设备: {exc}") from exc if device_count < 1: raise RuntimeError("未检测到 NVIDIA CUDA GPU;不会回退到 CPU") if self.config.device_id < 0 or self.config.device_id >= device_count: raise RuntimeError( f"GPU {self.config.device_id} 不存在,当前检测到 {device_count} 个设备" ) paddle.set_device(self.resolved_device) paddle.device.cuda.synchronize(self.config.device_id) try: self._device_name = paddle.device.cuda.get_device_name(self.config.device_id) except Exception: self._device_name = "unknown" self.logger.info( "GPU_CONFIGURED device=%s device_name=%s device_count=%d", self.resolved_device, self._device_name, device_count, ) self.setup_seconds = time.perf_counter() - setup_started self.logger.info( "RUNTIME_PREPARED device=%s paddle_version=%s import_seconds=%.3f setup_seconds=%.3f", self.resolved_device, paddle.__version__, self.import_seconds, self.setup_seconds, ) def get(self): self.prepare() if self._pipeline is None: self.logger.info( "MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", self.resolved_device, ) started = time.perf_counter() from paddleocr import PaddleOCRVL self._pipeline = PaddleOCRVL( pipeline_version="v1.6", device=self.resolved_device, ) self.synchronize() self.model_init_seconds = time.perf_counter() - started self.logger.info( "MODEL_INITIALIZED seconds=%.3f device=%s", self.model_init_seconds, self.resolved_device, ) return self._pipeline def synchronize(self) -> None: if self.config.device == "gpu" and self._paddle is not None: self._paddle.device.cuda.synchronize(self.config.device_id) def gpu_memory(self) -> dict[str, float | None]: stats: dict[str, float | None] = { "allocated_mb": None, "reserved_mb": None, "max_allocated_mb": None, "max_reserved_mb": None, } if self.config.device != "gpu" or self._paddle is None: return stats functions = { "allocated_mb": "memory_allocated", "reserved_mb": "memory_reserved", "max_allocated_mb": "max_memory_allocated", "max_reserved_mb": "max_memory_reserved", } for key, name in functions.items(): function = getattr(self._paddle.device.cuda, name, None) if function is None: continue try: stats[key] = round( float(function(self.config.device_id)) / (1024**2), 2 ) except Exception: pass return stats def metadata(self) -> dict[str, Any]: paddle_version = self._paddle.__version__ if self._paddle is not None else None return { "device": self.resolved_device, "device_name": self._device_name, "cpu_threads": self.config.threads if self.config.device == "cpu" else None, "python_version": platform.python_version(), "platform": platform.platform(), "paddle_version": paddle_version, "pipeline_version": "v1.6", "runtime_import_seconds": round(self.import_seconds, 3), "runtime_setup_seconds": round(self.setup_seconds, 3), "model_init_seconds": round(self.model_init_seconds, 3), }