PaddleOCR-VL-1.6_Demo/ocr_app/runtime.py

"""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),
        }