msgbot/easyocr_worker.py

#!/usr/bin/env python3
"""EasyOCR 子进程 worker。

通过 stdin 读 base64 图片、stdout 写 JSON 结果。GPU crash 只杀此进程，不冻主进程。
协议：
  stdin:  {"img":"<base64_png_bytes>"}
  stdout: {"ok":true,"text":"识别结果"}
  stdout: {"ok":false,"error":"错误描述"}

初始化慢（首次 3-10s），初始化完成后常驻处理。
"""

import json
import sys
import base64
import io
import os
import traceback
from datetime import datetime

# ── 崩溃日志文件：写入源码目录根（easyocr_crash.log），统一管理、便于父进程读取 ──
def _get_crash_log_path() -> str:
    try:
        base = os.path.dirname(os.path.abspath(__file__))
    except Exception:
        base = "."
    return os.path.join(base, "easyocr_crash.log")

CRASH_LOG = _get_crash_log_path()

def _write_crash_log(msg: str):
    try:
        with open(CRASH_LOG, "a", encoding="utf-8") as f:
            f.write(msg + "\n")
            f.flush()
    except Exception:
        pass

# ── 顶层异常捕获：确保崩溃原因写入 stderr + 落盘（否则父进程看到空 stderr） ──
def _flush_exit(msg, code=1):
    full = f"[easyocr_worker] FATAL: {msg}"
    sys.stderr.write(full + "\n")
    sys.stderr.flush()
    _write_crash_log(f"{datetime.now()} {full}")
    sys.exit(code)

# ── 全局未捕获异常钩子：任何漏网的 Python 异常都落盘 ──
def _global_excepthook(exc_type, exc_val, exc_tb):
    msg = "".join(traceback.format_exception(exc_type, exc_val, exc_tb))
    sys.stderr.write(f"[easyocr_worker] 未捕获异常:\n{msg}")
    sys.stderr.flush()
    _write_crash_log(f"{datetime.now()} [easyocr_worker] 未捕获异常:\n{msg}")

sys.excepthook = _global_excepthook

# ── faulthandler：捕获原生层崩溃（segfault / CUDA 初始化竞态），Python try/except 抓不到 ──
#    必须在重型 import 之前启用，连 import 阶段的原生崩溃也能 dump 出来。
try:
    import faulthandler
    _fh = open(CRASH_LOG, "a", encoding="utf-8")
    faulthandler.enable(file=_fh, all_threads=True)
except Exception:
    pass

try:
    # ── 线程/进程限制：避免 MKL/OpenMP 在子进程多线程冲突 ──
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")

    from PIL import Image
    import numpy as np
    import easyocr
except Exception as e:
    _flush_exit(f"import 失败: {e}\n{traceback.format_exc()}")


def _init_reader():
    """初始化 EasyOCR reader（GPU 优先，父进程可强制 CPU）。"""
    try:
        import torch
        # 父进程 EASYOCR_FORCE_CPU=1 → 强制 CPU（显存反复 OOM 时降级，防冻结）
        _force_cpu = os.environ.get("EASYOCR_FORCE_CPU") == "1"
        use_gpu = (not _force_cpu) and torch.cuda.is_available()
    except Exception:
        use_gpu = False

    try:
        reader = easyocr.Reader(["ch_sim", "en"], gpu=use_gpu)
        mode = "GPU" if use_gpu else "CPU"
        sys.stderr.write(f"[easyocr_worker] Reader 就绪（{mode}模式）\n")
        sys.stderr.flush()
        return reader
    except Exception as e:
        # GPU 不可用 → CPU 兜底
        sys.stderr.write(f"[easyocr_worker] GPU 初始化失败({e})，尝试 CPU\n")
        sys.stderr.flush()
        reader = easyocr.Reader(["ch_sim", "en"], gpu=False)
        sys.stderr.write("[easyocr_worker] Reader 就绪（CPU 模式）\n")
        sys.stderr.flush()
        return reader


def _process_image(img_b64: str, reader) -> str:
    """解码 base64 图片 → EasyOCR 识别 → 返回空格拼接文本。"""
    if reader is None:
        return ""
    raw = base64.b64decode(img_b64)
    img = Image.open(io.BytesIO(raw))
    # 保持与 monitor.py 一致：RGBA → numpy
    if img.mode != "RGBA":
        img = img.convert("RGBA")
    img_np = np.array(img)
    results = reader.readtext(img_np)
    return "\n".join([r[1] for r in results]).strip()


def _main():
    _write_crash_log(f"{datetime.now()} [easyocr_worker] 启动 (pid={os.getpid()})")
    reader = None
    try:
        reader = _init_reader()
    except Exception as e:
        sys.stderr.write(f"[easyocr_worker] Reader 初始化失败: {e}\n")
        sys.stderr.flush()
        _write_crash_log(f"{datetime.now()} [easyocr_worker] Reader 初始化失败: {e}")
    # ⭐ 就绪握手：即使 Reader 失败也要发 ready，否则父进程永远等握手
    sys.stdout.write(json.dumps({"ready": True}, ensure_ascii=False) + "\n")
    sys.stdout.flush()
    sys.stderr.write("[easyocr_worker] 就绪，等待 stdin 输入\n")
    sys.stderr.flush()

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            payload = json.loads(line)
            img_b64 = payload.get("img", "")
            if not img_b64:
                raise ValueError("缺少 img 字段")
            sys.stderr.write(f"[easyocr_worker] 收到请求，处理中... (img长度={len(img_b64)})\n")
            sys.stderr.flush()
            text = _process_image(img_b64, reader)
            sys.stderr.write(f"[easyocr_worker] 处理完成，结果长度={len(text)}\n")
            sys.stderr.flush()
            out = json.dumps({"ok": True, "text": text}, ensure_ascii=False)
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
        except Exception as e:
            estr = str(e)
            # ⭐ CUDA OOM / kernel error：GPU context 已损坏，继续跑只会反复卡死占显存
            #   → 立即退出进程释放显存，父进程检测到 stdout 关闭后自动重启/降级 CPU
            if "cuda" in estr.lower() or "out of memory" in estr.lower():
                _write_crash_log(f"{datetime.now()} [easyocr_worker] CUDA 错误，主动退出释放显存: {estr[:200]}")
                _flush_exit(f"CUDA 错误，主动退出释放显存: {estr[:200]}")
            out = json.dumps({"ok": False, "error": str(e)}, ensure_ascii=False)
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
            sys.stderr.write(f"[easyocr_worker] 处理失败: {e}\n")
            sys.stderr.flush()


if __name__ == "__main__":
    try:
        _main()
    except Exception as e:
        _flush_exit(f"_main 崩溃: {e}\n{traceback.format_exc()}")