start_embedding_service.sh 5.21 KB
#!/bin/bash
#
# Start Embedding Service (port 6005).
#
# Design:
# - Run in isolated venv `.venv-embedding` (do not pollute main `.venv`)
# - Text backend default: TEI (Qwen3-Embedding-0.6B)
# - Image backend: clip-as-service or local CN-CLIP (from embeddings/config.py)
#
set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "${PROJECT_ROOT}"

EMBEDDING_VENV="${EMBEDDING_VENV:-${PROJECT_ROOT}/.venv-embedding}"
PYTHON_BIN="${EMBEDDING_VENV}/bin/python"

if [[ ! -x "${PYTHON_BIN}" ]]; then
  echo "ERROR: embedding venv not found: ${EMBEDDING_VENV}" >&2
  echo "Please run: ./scripts/setup_embedding_venv.sh" >&2
  exit 1
fi

# Load .env if present (same behavior as activate.sh, without activating main venv)
ENV_FILE="${PROJECT_ROOT}/.env"
if [ -f "${ENV_FILE}" ]; then
  while IFS= read -r line || [ -n "${line}" ]; do
    line="${line%$'\r'}"
    [[ -z "${line//[[:space:]]/}" ]] && continue
    [[ "${line}" =~ ^[[:space:]]*# ]] && continue
    [[ "${line}" != *=* ]] && continue

    key="${line%%=*}"
    value="${line#*=}"
    key="${key#"${key%%[![:space:]]*}"}"
    key="${key%"${key##*[![:space:]]}"}"
    value="${value#"${value%%[![:space:]]*}"}"

    if [[ ${#value} -ge 2 ]]; then
      first="${value:0:1}"
      last="${value: -1}"
      if [[ ("${first}" == '"' && "${last}" == '"') || ("${first}" == "'" && "${last}" == "'") ]]; then
        value="${value:1:${#value}-2}"
      fi
    fi

    export "${key}=${value}"
  done < "${ENV_FILE}"
fi

DEFAULT_EMBEDDING_SERVICE_HOST=$("${PYTHON_BIN}" -c "from embeddings.config import CONFIG; print(CONFIG.HOST)")
DEFAULT_EMBEDDING_SERVICE_PORT=$("${PYTHON_BIN}" -c "from embeddings.config import CONFIG; print(CONFIG.PORT)")
USE_CLIP_AS_SERVICE=$("${PYTHON_BIN}" -c "from embeddings.config import CONFIG; print('1' if CONFIG.USE_CLIP_AS_SERVICE else '0')")
CLIP_AS_SERVICE_SERVER=$("${PYTHON_BIN}" -c "from embeddings.config import CONFIG; print(CONFIG.CLIP_AS_SERVICE_SERVER)")
TEXT_BACKEND=$("${PYTHON_BIN}" -c "from config.services_config import get_embedding_backend_config; print(get_embedding_backend_config()[0])")
TEI_BASE_URL=$("${PYTHON_BIN}" -c "import os; from config.services_config import get_embedding_backend_config; from embeddings.config import CONFIG; _, cfg = get_embedding_backend_config(); print(os.getenv('TEI_BASE_URL') or cfg.get('base_url') or CONFIG.TEI_BASE_URL)")
ENABLE_IMAGE_MODEL="${EMBEDDING_ENABLE_IMAGE_MODEL:-true}"
ENABLE_IMAGE_MODEL="$(echo "${ENABLE_IMAGE_MODEL}" | tr '[:upper:]' '[:lower:]')"
if [[ "${ENABLE_IMAGE_MODEL}" == "1" || "${ENABLE_IMAGE_MODEL}" == "true" || "${ENABLE_IMAGE_MODEL}" == "yes" ]]; then
  IMAGE_MODEL_ENABLED=1
else
  IMAGE_MODEL_ENABLED=0
fi

EMBEDDING_SERVICE_HOST="${EMBEDDING_HOST:-${DEFAULT_EMBEDDING_SERVICE_HOST}}"
EMBEDDING_SERVICE_PORT="${EMBEDDING_PORT:-${DEFAULT_EMBEDDING_SERVICE_PORT}}"

if [[ "${TEXT_BACKEND}" == "tei" ]]; then
  if ! curl -sf "${TEI_BASE_URL%/}/health" >/dev/null 2>&1; then
    echo "ERROR: TEI backend is selected but TEI is not reachable: ${TEI_BASE_URL}/health" >&2
    echo "Please start TEI first: ./scripts/start_tei_service.sh" >&2
    exit 1
  fi
fi

if [[ "${IMAGE_MODEL_ENABLED}" == "1" && "${USE_CLIP_AS_SERVICE}" == "1" ]]; then
  CLIP_SERVER="${CLIP_AS_SERVICE_SERVER#*://}"
  CLIP_HOST="${CLIP_SERVER%:*}"
  CLIP_PORT="${CLIP_SERVER##*:}"
  if [[ -z "${CLIP_HOST}" || -z "${CLIP_PORT}" ]]; then
    echo "ERROR: invalid CLIP_AS_SERVICE_SERVER: ${CLIP_AS_SERVICE_SERVER}" >&2
    exit 1
  fi
  if ! timeout 2 bash -c "cat < /dev/null > /dev/tcp/${CLIP_HOST}/${CLIP_PORT}" >/dev/null 2>&1; then
    echo "ERROR: clip-as-service is not reachable at ${CLIP_AS_SERVICE_SERVER}." >&2
    echo "Please start CN-CLIP service first: ./scripts/service_ctl.sh start cnclip" >&2
    exit 1
  fi
  if ! "${PYTHON_BIN}" - <<'PY'
try:
    import pkg_resources  # noqa: F401
except Exception:
    raise SystemExit(1)
PY
  then
    echo "ERROR: clip-as-service image embedding requires pkg_resources in .venv-embedding." >&2
    echo "Please run: ./scripts/setup_embedding_venv.sh" >&2
    exit 1
  fi
fi

echo "========================================"
echo "Starting Embedding Service"
echo "========================================"
echo "Python: ${PYTHON_BIN}"
echo "Host: ${EMBEDDING_SERVICE_HOST}"
echo "Port: ${EMBEDDING_SERVICE_PORT}"
echo "Text backend: ${TEXT_BACKEND}"
if [[ "${TEXT_BACKEND}" == "tei" ]]; then
  echo "TEI URL: ${TEI_BASE_URL}"
fi
if [[ "${IMAGE_MODEL_ENABLED}" == "0" ]]; then
  echo "Image backend: disabled"
elif [[ "${USE_CLIP_AS_SERVICE}" == "1" ]]; then
  echo "Image backend: clip-as-service (${CLIP_AS_SERVICE_SERVER})"
fi
echo
echo "Tips:"
echo "  - Use a single worker (GPU models cannot be safely duplicated across workers)."
echo "  - Clients can set EMBEDDING_SERVICE_URL=http://localhost:${EMBEDDING_SERVICE_PORT}"
echo

UVICORN_LOG_LEVEL="${EMBEDDING_UVICORN_LOG_LEVEL:-info}"
UVICORN_ACCESS_LOG="${EMBEDDING_UVICORN_ACCESS_LOG:-true}"
UVICORN_ARGS=(
  --host "${EMBEDDING_SERVICE_HOST}"
  --port "${EMBEDDING_SERVICE_PORT}"
  --workers 1
  --log-level "${UVICORN_LOG_LEVEL}"
)
if [[ "${UVICORN_ACCESS_LOG}" == "0" || "${UVICORN_ACCESS_LOG}" == "false" || "${UVICORN_ACCESS_LOG}" == "no" ]]; then
  UVICORN_ARGS+=(--no-access-log)
fi

exec "${PYTHON_BIN}" -m uvicorn embeddings.server:app "${UVICORN_ARGS[@]}"