dekun
123 lines
3.4 KiB
Bash
123 lines
3.4 KiB
Bash
#!/usr/bin/env bash
|
||
# OCR Worker 共用函数：同机 GPU 识别 + screen 常驻
|
||
OCR_SCREEN_NAME="${OCR_SCREEN_NAME:-ocr-worker}"
|
||
OCR_PORT="${OCR_PORT:-23567}"
|
||
|
||
ocr_worker_dir() {
|
||
echo "${INSTALL_DIR}/deploy/ocr-worker"
|
||
}
|
||
|
||
detect_ocr_use_gpu() {
|
||
if [[ "${OCR_USE_GPU:-auto}" == "false" ]]; then
|
||
echo "false"
|
||
return
|
||
fi
|
||
if [[ "${OCR_USE_GPU:-auto}" == "true" ]]; then
|
||
if ldconfig -p 2>/dev/null | grep -q libcudnn; then
|
||
echo "true"
|
||
else
|
||
log_warn "未检测到 cuDNN 库，GPU OCR 不可用，改用 CPU"
|
||
echo "false"
|
||
fi
|
||
return
|
||
fi
|
||
# auto
|
||
if command -v nvidia-smi >/dev/null && ldconfig -p 2>/dev/null | grep -q libcudnn; then
|
||
echo "true"
|
||
else
|
||
echo "false"
|
||
fi
|
||
}
|
||
|
||
install_ocr_worker() {
|
||
local worker_dir
|
||
worker_dir="$(ocr_worker_dir)"
|
||
if [[ ! -d "${worker_dir}" ]]; then
|
||
log_error "未找到 ${worker_dir}"
|
||
return 1
|
||
fi
|
||
log_info "安装/更新 OCR Worker (RapidOCR/ONNX)…"
|
||
chmod +x "${worker_dir}"/*.sh 2>/dev/null || true
|
||
OCR_PORT="${OCR_PORT}" bash "${worker_dir}/install.sh"
|
||
}
|
||
|
||
ocr_screen_running() {
|
||
screen -list 2>/dev/null | grep -q "\.${OCR_SCREEN_NAME}[[:space:]]"
|
||
}
|
||
|
||
start_ocr_screen() {
|
||
local worker_dir log_file
|
||
worker_dir="$(ocr_worker_dir)"
|
||
if [[ ! -x "${worker_dir}/.venv/bin/uvicorn" ]]; then
|
||
log_warn "OCR Worker 未安装，跳过 screen 启动"
|
||
return 1
|
||
fi
|
||
if ! command -v screen >/dev/null; then
|
||
log_error "未安装 screen，请 apt install screen"
|
||
return 1
|
||
fi
|
||
|
||
log_info "启动 OCR Worker → screen「${OCR_SCREEN_NAME}」(RapidOCR, 端口 ${OCR_PORT})"
|
||
if ocr_screen_running; then
|
||
screen -S "${OCR_SCREEN_NAME}" -X quit 2>/dev/null || true
|
||
sleep 1
|
||
fi
|
||
|
||
mkdir -p "${INSTALL_DIR}/logs" 2>/dev/null || true
|
||
log_file="${INSTALL_DIR}/logs/ocr-worker.log"
|
||
|
||
screen -dmS "${OCR_SCREEN_NAME}" bash -c "
|
||
cd '${worker_dir}' &&
|
||
export OCR_PORT='${OCR_PORT}' OCR_HOST=0.0.0.0 &&
|
||
exec bash run.sh >> '${log_file}' 2>&1
|
||
"
|
||
sleep 2
|
||
log_info "OCR 日志: ${log_file}"
|
||
}
|
||
|
||
stop_ocr_screen() {
|
||
if ocr_screen_running; then
|
||
log_info "停止 OCR screen 会话…"
|
||
screen -S "${OCR_SCREEN_NAME}" -X quit 2>/dev/null || true
|
||
fi
|
||
}
|
||
|
||
wait_ocr_healthy() {
|
||
local max="${1:-30}"
|
||
local i
|
||
log_info "等待 OCR 就绪（最多 ${max}×2 秒，首次加载模型较慢）…"
|
||
for i in $(seq 1 "${max}"); do
|
||
if curl -sf "http://127.0.0.1:${OCR_PORT}/health" >/dev/null 2>&1; then
|
||
log_info "OCR 健康检查通过 — http://127.0.0.1:${OCR_PORT}/health"
|
||
return 0
|
||
fi
|
||
if (( i % 5 == 0 )); then
|
||
echo -ne "\r${YELLOW}[INFO]${NC} 仍在等待 OCR… ${i}/${max}（可另开终端: tail -f ${INSTALL_DIR}/logs/ocr-worker.log）"
|
||
fi
|
||
sleep 2
|
||
done
|
||
echo ""
|
||
log_warn "OCR 尚未响应（可能仍在下载/加载模型，主程序可先继续）"
|
||
log_warn "稍后执行: bash ${INSTALL_DIR}/deploy/ocr-screen.sh status"
|
||
return 1
|
||
}
|
||
|
||
show_ocr_status() {
|
||
echo ""
|
||
echo "--- OCR Worker (screen: ${OCR_SCREEN_NAME}) ---"
|
||
if ocr_screen_running; then
|
||
echo "screen: 运行中"
|
||
screen -list 2>/dev/null | grep "${OCR_SCREEN_NAME}" || true
|
||
else
|
||
echo "screen: 未运行"
|
||
fi
|
||
if curl -sf "http://127.0.0.1:${OCR_PORT}/health" 2>/dev/null; then
|
||
echo ""
|
||
else
|
||
echo "health: 无响应 (http://127.0.0.1:${OCR_PORT}/health)"
|
||
fi
|
||
if command -v nvidia-smi >/dev/null; then
|
||
nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu --format=csv,noheader 2>/dev/null || true
|
||
fi
|
||
}
|