332 lines
14 KiB
Python
332 lines
14 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()
|
||
# ⭐ 统一使用 GlobalMemoryStatusEx（与 monitor.py / Task Manager 一致）
|
||
# psutil.virtual_memory().available 在 Windows 上可比 Task Manager 少报 ~1.7GB，
|
||
# 导致 4.2GB 可用时误判为 2.5GB 触发降级。
|
||
if _vm:
|
||
try:
|
||
_avail_phys_gb = _vm.get("avail_phys_gb", 0)
|
||
_total_phys_gb = _vm.get("total_phys_gb", 0)
|
||
_used_phys_gb = _total_phys_gb - _avail_phys_gb
|
||
_pagefile_pct = _vm.get("commit_pct", 0)
|
||
_pagefile_used_gb = _vm.get("total_pagefile_gb", 0) - _vm.get("avail_pagefile_gb", 0)
|
||
_commit_gb = _used_phys_gb + _pagefile_used_gb
|
||
_commit_total_gb = _total_phys_gb + _vm.get("total_pagefile_gb", 0)
|
||
_commit_pct = (_commit_gb / max(_commit_total_gb, 1)) * 100
|
||
# ⭐ 放宽阈值：提交量>95%才强制CPU，物理内存<2GB才强制CPU
|
||
# 避免内存充足时误判（用户报告提交量81.8%被误判为不足）
|
||
if _commit_pct > 95.0:
|
||
_msg = (
|
||
f"系统内存压力过高: 物理剩余{_avail_phys_gb:.1f}GB，"
|
||
f"提交量{_commit_pct:.1f}% ({_commit_gb:.1f}GB/{_commit_total_gb:.1f}GB)"
|
||
)
|
||
_write_crash_log(f"{datetime.now()} [easyocr_worker] 系统内存压力预检失败: {_msg}")
|
||
os.environ["EASYOCR_FORCE_CPU"] = "1"
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "" # 禁用CUDA，防止加载CUDA库导致崩溃
|
||
sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
|
||
sys.stderr.flush()
|
||
elif _avail_phys_gb < 2.0:
|
||
_msg = (
|
||
f"物理内存不足: 剩余{_avail_phys_gb:.1f}GB"
|
||
f"(总计{_total_phys_gb:.1f}GB)"
|
||
f"，CUDA加载需要≥2GB物理内存"
|
||
)
|
||
_write_crash_log(f"{datetime.now()} [easyocr_worker] 物理内存预检失败: {_msg}")
|
||
os.environ["EASYOCR_FORCE_CPU"] = "1"
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "" # 禁用CUDA，防止加载CUDA库导致崩溃
|
||
sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
|
||
sys.stderr.flush()
|
||
except Exception:
|
||
pass # GlobalMemoryStatusEx 失败时不阻断
|
||
# ⭐ 放宽退出阈值：百分比>98% 或 剩余<1GB 时退出（防冻结）
|
||
# 避免内存充足时误退出（用户报告提交量81.8%被误判为不足）
|
||
if _vm and (_vm.get("commit_pct", 0) > 98.0 or _vm.get("avail_pagefile_gb", 99) < 1.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
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "" # 禁用CUDA，防止加载CUDA库导致崩溃
|
||
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
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "" # 禁用CUDA，防止后续崩溃
|
||
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()
|
||
|
||
|
||
def _cleanup_gpu():
|
||
"""退出时清理 CUDA/GPU 内存，防止子进程被 kill 后 GPU 内存泄漏到系统。"""
|
||
try:
|
||
import gc
|
||
gc.collect()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.ipc_collect()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _signal_handler(signum, frame):
|
||
"""捕获 SIGTERM/SIGINT，先清理 GPU 再退出（proc.terminate() 触发此处理）。"""
|
||
_cleanup_gpu()
|
||
sys.exit(0)
|
||
|
||
|
||
# 注册清理钩子：atexit（正常退出）+ signal（SIGTERM/SIGINT）
|
||
import atexit
|
||
atexit.register(_cleanup_gpu)
|
||
try:
|
||
import signal
|
||
signal.signal(signal.SIGTERM, _signal_handler)
|
||
signal.signal(signal.SIGINT, _signal_handler)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
_main()
|
||
except Exception as e:
|
||
_flush_exit(f"_main 崩溃: {e}\n{traceback.format_exc()}")
|
||
finally:
|
||
_cleanup_gpu()
|