youpaishidifu
4d7017be6a
- knowledge_base.py: _split_body_by_list_items检测表格区域，不切分表格行 - knowledge_base.py: 添加独立表格section，确保表格内容可被搜索匹配 - knowledge_base.py: 表格兜底策略：最佳文档含操作/支持表格时强制加入并置顶 - knowledge_base.py: reranking后无条件优先返回含操作/支持的表格section - .gitignore: 添加msg_history.db和诊断脚本到忽略列表
289 lines
12 KiB
Python
289 lines
12 KiB
Python
#!/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_virtual_memory_info() -> dict:
|
||
"""获取虚拟内存（页面文件）状态，用于诊断 WinError 1455。"""
|
||
try:
|
||
import ctypes
|
||
class MEMORYSTATUSEX(ctypes.Structure):
|
||
_fields_ = [
|
||
("dwLength", ctypes.c_ulong),
|
||
("dwMemoryLoad", ctypes.c_ulong),
|
||
("ullTotalPhys", ctypes.c_ulonglong),
|
||
("ullAvailPhys", ctypes.c_ulonglong),
|
||
("ullTotalPageFile", ctypes.c_ulonglong),
|
||
("ullAvailPageFile", ctypes.c_ulonglong),
|
||
("ullTotalVirtual", ctypes.c_ulonglong),
|
||
("ullAvailVirtual", ctypes.c_ulonglong),
|
||
("ullAvailExtendedVirtual", ctypes.c_ulonglong),
|
||
]
|
||
mem = MEMORYSTATUSEX()
|
||
mem.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
|
||
ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(mem))
|
||
return {
|
||
"total_phys_gb": round(mem.ullTotalPhys / (1024**3), 1),
|
||
"avail_phys_gb": round(mem.ullAvailPhys / (1024**3), 1),
|
||
"total_pagefile_gb": round(mem.ullTotalPageFile / (1024**3), 1),
|
||
"avail_pagefile_gb": round(mem.ullAvailPageFile / (1024**3), 1),
|
||
"commit_pct": round((1 - mem.ullAvailPageFile / max(mem.ullTotalPageFile, 1)) * 100, 1),
|
||
}
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
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")
|
||
|
||
# ── 预检：物理内存+虚拟内存是否足够加载 CUDA 库 ──
|
||
_vm = _get_virtual_memory_info()
|
||
# 检查物理内存：CUDA加载需要足够物理内存，否则会使用页面文件导致冻结
|
||
try:
|
||
import psutil
|
||
_mem = psutil.virtual_memory()
|
||
_avail_phys_gb = _mem.available / (1024**3)
|
||
_used_phys_gb = _mem.used / (1024**3)
|
||
_total_phys_gb = _mem.total / (1024**3)
|
||
# ⭐ 页面文件百分比 + 总提交量检查：防页面文件94%+时仍尝试CUDA导致冻结
|
||
_pagefile_pct = _vm.get("commit_pct", 0) if _vm else 0
|
||
_pagefile_used_gb = _total_phys_gb + (_vm.get("total_pagefile_gb", 0) - _vm.get("avail_pagefile_gb", 0)) - _avail_phys_gb if _vm else 0
|
||
_commit_gb = _used_phys_gb + (_vm.get("total_pagefile_gb", 0) - _vm.get("avail_pagefile_gb", 0)) if _vm else _used_phys_gb
|
||
if _pagefile_pct > 85.0 or _commit_gb > _total_phys_gb * 1.5:
|
||
_msg = (
|
||
f"系统内存压力过高: 物理剩余{_avail_phys_gb:.1f}GB({_mem.percent}%)，"
|
||
f"页面文件已用{_pagefile_pct}%，总提交{_commit_gb:.1f}GB>{_total_phys_gb*1.5:.1f}GB阈值"
|
||
)
|
||
_write_crash_log(f"{datetime.now()} [easyocr_worker] 系统内存压力预检失败: {_msg}")
|
||
os.environ["EASYOCR_FORCE_CPU"] = "1"
|
||
sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
|
||
sys.stderr.flush()
|
||
elif _avail_phys_gb < 3.0:
|
||
_msg = (
|
||
f"物理内存不足: 剩余{_avail_phys_gb:.1f}GB"
|
||
f"(总计{_total_phys_gb:.1f}GB，使用{_mem.percent}%)"
|
||
f"，CUDA加载需要≥3GB物理内存"
|
||
)
|
||
_write_crash_log(f"{datetime.now()} [easyocr_worker] 物理内存预检失败: {_msg}")
|
||
# 设置环境变量强制CPU模式，而不是直接退出
|
||
os.environ["EASYOCR_FORCE_CPU"] = "1"
|
||
sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
|
||
sys.stderr.flush()
|
||
except Exception:
|
||
pass # psutil不可用时不阻断
|
||
# ⭐ 检查页面文件：百分比>85% 或 剩余<2GB 时退出（防冻结）
|
||
if _vm and (_vm.get("commit_pct", 0) > 85.0 or _vm.get("avail_pagefile_gb", 99) < 2.0):
|
||
_msg = (
|
||
f"页面文件不足: 已用{_vm.get('commit_pct', '?')}%"
|
||
f"(剩余{_vm.get('avail_pagefile_gb', '?')}GB/{_vm.get('total_pagefile_gb', '?')}GB)"
|
||
f"，CUDA库需要充足的虚拟地址空间"
|
||
)
|
||
_write_crash_log(f"{datetime.now()} [easyocr_worker] 页面文件预检失败: {_msg}")
|
||
_flush_exit(f"页面文件不足: {_msg}", code=2)
|
||
|
||
from PIL import Image
|
||
import numpy as np
|
||
import easyocr
|
||
except Exception as e:
|
||
_vm = _get_virtual_memory_info()
|
||
_vm_hint = ""
|
||
if _vm:
|
||
_vm_hint = (
|
||
f"\n[virtual memory] 物理内存: {_vm.get('avail_phys_gb',0)}/{_vm.get('total_phys_gb',0)} GB"
|
||
f" | 页面文件: {_vm.get('avail_pagefile_gb',0)}/{_vm.get('total_pagefile_gb',0)} GB"
|
||
f" | 提交使用率: {_vm.get('commit_pct',0)}%"
|
||
)
|
||
# WinError 1455 = 页面文件太小 → 标记为永久失败，不重试
|
||
_is_pagefile_error = "1455" in str(e) or "页面文件" in str(e)
|
||
_exit_code = 2 if _is_pagefile_error else 1 # exit(2) = 虚拟内存不足
|
||
_write_crash_log(
|
||
f"{datetime.now()} [easyocr_worker] import {_exit_code}: "
|
||
f"{'页面文件不足' if _is_pagefile_error else 'import失败'}: {e}{_vm_hint}"
|
||
)
|
||
_flush_exit(f"import 失败{'(页面文件不足)' if _is_pagefile_error else ''}: {e}{_vm_hint}", code=_exit_code)
|
||
|
||
|
||
def _init_reader():
|
||
"""初始化 EasyOCR reader（GPU 优先，父进程可强制 CPU）。"""
|
||
# 先检查环境变量是否强制CPU
|
||
_force_cpu = os.environ.get("EASYOCR_FORCE_CPU") == "1"
|
||
|
||
# 再次检查物理内存：加载torch+CUDA需要大量内存，内存不足时强制CPU
|
||
if not _force_cpu:
|
||
try:
|
||
import psutil
|
||
_mem = psutil.virtual_memory()
|
||
_avail_phys_gb = _mem.available / (1024**3)
|
||
if _avail_phys_gb < 3.0:
|
||
_force_cpu = True
|
||
sys.stderr.write(
|
||
f"[easyocr_worker] ⚠️ 物理内存不足(剩余{_avail_phys_gb:.1f}GB)，"
|
||
f"强制CPU模式（避免CUDA加载导致冻结）\n"
|
||
)
|
||
sys.stderr.flush()
|
||
except Exception:
|
||
pass # psutil不可用时不阻断
|
||
|
||
try:
|
||
import torch
|
||
# 父进程 EASYOCR_FORCE_CPU=1 → 强制 CPU（显存反复 OOM 时降级，防冻结）
|
||
use_gpu = (not _force_cpu) and torch.cuda.is_available()
|
||
if use_gpu:
|
||
# CUDA 健康检查：尝试创建一个小张量，避免驱动/显存问题导致后续崩溃
|
||
try:
|
||
_test_tensor = torch.zeros(1, device="cuda")
|
||
del _test_tensor
|
||
torch.cuda.empty_cache()
|
||
sys.stderr.write(f"[easyocr_worker] CUDA 健康检查通过 (设备: {torch.cuda.get_device_name(0)})\n")
|
||
sys.stderr.flush()
|
||
except Exception as cuda_err:
|
||
sys.stderr.write(f"[easyocr_worker] CUDA 健康检查失败({cuda_err})，降级 CPU\n")
|
||
sys.stderr.flush()
|
||
use_gpu = False
|
||
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()}")
|