secondary-school-grade-archive/deploy/ocr-common.sh

#!/usr/bin/env bash
# OCR Worker 共用函数：同机 RapidOCR + systemd 开机自启
OCR_SCREEN_NAME="${OCR_SCREEN_NAME:-ocr-worker}"
OCR_PORT="${OCR_PORT:-23567}"
OCR_SERVICE_NAME="${OCR_SERVICE_NAME:-ocr-worker}"

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：有 nvidia-smi 且能看到 GPU 才开（无卡时驱动在也会失败）
  if command -v nvidia-smi >/dev/null \
    && nvidia-smi >/dev/null 2>&1 \
    && 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:]]"
}

stop_ocr_screen() {
  if ocr_screen_running; then
    log_info "停止旧版 OCR screen 会话…"
    screen -S "${OCR_SCREEN_NAME}" -X quit 2>/dev/null || true
    sleep 1
  fi
}

# 注册并启用 systemd（一键部署 / 更新时调用）
setup_ocr_systemd() {
  local worker_dir use_gpu
  worker_dir="$(ocr_worker_dir)"
  if [[ ! -x "${worker_dir}/.venv/bin/uvicorn" ]]; then
    log_warn "OCR Worker 未安装，跳过 systemd 注册"
    return 1
  fi

  use_gpu="$(detect_ocr_use_gpu)"
  stop_ocr_screen

  log_info "注册 systemd 服务 ${OCR_SERVICE_NAME}（OCR_USE_GPU=${use_gpu}, 端口 ${OCR_PORT}）…"
  cat > "/etc/systemd/system/${OCR_SERVICE_NAME}.service" <<EOF
[Unit]
Description=Grade Archive OCR Worker (RapidOCR)
After=network.target

[Service]
Type=simple
WorkingDirectory=${worker_dir}
Environment=OCR_USE_GPU=${use_gpu}
Environment=OCR_PORT=${OCR_PORT}
Environment=OCR_HOST=0.0.0.0
ExecStart=${worker_dir}/.venv/bin/uvicorn app:app --host 0.0.0.0 --port ${OCR_PORT} --log-level info
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
EOF

  systemctl daemon-reload
  systemctl enable "${OCR_SERVICE_NAME}"
}

start_ocr_service() {
  setup_ocr_systemd || return 1
  log_info "启动 OCR systemd 服务…"
  systemctl restart "${OCR_SERVICE_NAME}"
}

stop_ocr_service() {
  systemctl stop "${OCR_SERVICE_NAME}" 2>/dev/null || true
  stop_ocr_screen
}

# 兼容旧脚本名
start_ocr_screen() {
  start_ocr_service
}

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}（journalctl -u ${OCR_SERVICE_NAME} -f）"
    fi
    sleep 2
  done
  echo ""
  log_warn "OCR 尚未响应（可能仍在下载/加载模型，主程序可先继续）"
  log_warn "稍后执行: systemctl status ${OCR_SERVICE_NAME}"
  return 1
}

show_ocr_status() {
  echo ""
  echo "--- OCR Worker (systemd: ${OCR_SERVICE_NAME}) ---"
  if systemctl is-enabled "${OCR_SERVICE_NAME}" &>/dev/null; then
    echo "开机自启: $(systemctl is-enabled "${OCR_SERVICE_NAME}" 2>/dev/null)"
  else
    echo "开机自启: 未注册"
  fi
  if systemctl is-active "${OCR_SERVICE_NAME}" &>/dev/null; then
    echo "服务状态: active"
  else
    echo "服务状态: $(systemctl is-active "${OCR_SERVICE_NAME}" 2>/dev/null || echo inactive)"
  fi
  if ocr_screen_running; then
    echo "screen(旧): 仍在运行（建议交由 systemd 接管）"
    screen -list 2>/dev/null | grep "${OCR_SCREEN_NAME}" || true
  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
}