start_embedding_service.sh 7.1 KB
#!/bin/bash
#
# Start Embedding Service (combined/text/image mode).
#
# 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 without activating main venv.
# shellcheck source=scripts/lib/load_env.sh
source "${PROJECT_ROOT}/scripts/lib/load_env.sh"
load_env_file "${PROJECT_ROOT}/.env"

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)")
CLIP_AS_SERVICE_MODEL_NAME=$("${PYTHON_BIN}" -c "from embeddings.config import CONFIG; print(CONFIG.CLIP_AS_SERVICE_MODEL_NAME)")
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)")
SERVICE_KIND="${1:-${EMBEDDING_SERVICE_KIND:-all}}"
SERVICE_KIND="$(echo "${SERVICE_KIND}" | tr '[:upper:]' '[:lower:]')"
if [[ "${SERVICE_KIND}" != "all" && "${SERVICE_KIND}" != "text" && "${SERVICE_KIND}" != "image" ]]; then
  echo "ERROR: invalid embedding service kind: ${SERVICE_KIND}. expected all|text|image" >&2
  exit 1
fi

ENABLE_TEXT_MODEL="${EMBEDDING_ENABLE_TEXT_MODEL:-true}"
ENABLE_TEXT_MODEL="$(echo "${ENABLE_TEXT_MODEL}" | tr '[:upper:]' '[:lower:]')"
ENABLE_IMAGE_MODEL="${EMBEDDING_ENABLE_IMAGE_MODEL:-true}"
ENABLE_IMAGE_MODEL="$(echo "${ENABLE_IMAGE_MODEL}" | tr '[:upper:]' '[:lower:]')"

TEXT_MODEL_ENABLED=0
IMAGE_MODEL_ENABLED=0
if [[ "${SERVICE_KIND}" == "all" || "${SERVICE_KIND}" == "text" ]]; then
  if [[ "${ENABLE_TEXT_MODEL}" == "1" || "${ENABLE_TEXT_MODEL}" == "true" || "${ENABLE_TEXT_MODEL}" == "yes" ]]; then
    TEXT_MODEL_ENABLED=1
  fi
fi
if [[ "${SERVICE_KIND}" == "all" || "${SERVICE_KIND}" == "image" ]]; then
  if [[ "${ENABLE_IMAGE_MODEL}" == "1" || "${ENABLE_IMAGE_MODEL}" == "true" || "${ENABLE_IMAGE_MODEL}" == "yes" ]]; then
    IMAGE_MODEL_ENABLED=1
  fi
fi

EMBEDDING_SERVICE_HOST="${EMBEDDING_HOST:-${DEFAULT_EMBEDDING_SERVICE_HOST}}"
if [[ "${SERVICE_KIND}" == "text" ]]; then
  EMBEDDING_SERVICE_PORT="${EMBEDDING_TEXT_PORT:-6005}"
elif [[ "${SERVICE_KIND}" == "image" ]]; then
  EMBEDDING_SERVICE_PORT="${EMBEDDING_IMAGE_PORT:-6008}"
else
  EMBEDDING_SERVICE_PORT="${EMBEDDING_PORT:-${DEFAULT_EMBEDDING_SERVICE_PORT}}"
fi

export EMBEDDING_SERVICE_KIND="${SERVICE_KIND}"
export EMBEDDING_HOST="${EMBEDDING_SERVICE_HOST}"
export EMBEDDING_PORT="${EMBEDDING_SERVICE_PORT}"
if [[ "${TEXT_MODEL_ENABLED}" == "1" ]]; then
  export EMBEDDING_ENABLE_TEXT_MODEL=true
else
  export EMBEDDING_ENABLE_TEXT_MODEL=false
fi
if [[ "${IMAGE_MODEL_ENABLED}" == "1" ]]; then
  export EMBEDDING_ENABLE_IMAGE_MODEL=true
else
  export EMBEDDING_ENABLE_IMAGE_MODEL=false
fi

if [[ "${TEXT_MODEL_ENABLED}" == "1" && "${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 "Kind: ${SERVICE_KIND}"
echo "Python: ${PYTHON_BIN}"
echo "Host: ${EMBEDDING_SERVICE_HOST}"
echo "Port: ${EMBEDDING_SERVICE_PORT}"
echo "Text backend enabled: ${TEXT_MODEL_ENABLED}"
if [[ "${TEXT_MODEL_ENABLED}" == "1" ]]; then
  echo "Text backend: ${TEXT_BACKEND}"
  echo "Text max inflight: ${TEXT_MAX_INFLIGHT:-32}"
fi
if [[ "${TEXT_MODEL_ENABLED}" == "1" && "${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}, model=${CLIP_AS_SERVICE_MODEL_NAME})"
fi
if [[ "${IMAGE_MODEL_ENABLED}" == "1" ]]; then
  echo "Image max inflight: ${IMAGE_MAX_INFLIGHT:-1}"
fi
if [[ "${SERVICE_KIND}" == "image" ]]; then
  echo "Logs: logs/embedding-image.log"
else
  echo "Logs: logs/embedding.log"
fi
echo
echo "Tips:"
echo "  - Use a single worker (GPU models cannot be safely duplicated across workers)."
if [[ "${SERVICE_KIND}" == "text" ]]; then
  echo "  - Clients can set EMBEDDING_TEXT_SERVICE_URL=http://localhost:${EMBEDDING_SERVICE_PORT}"
elif [[ "${SERVICE_KIND}" == "image" ]]; then
  echo "  - Clients can set EMBEDDING_IMAGE_SERVICE_URL=http://localhost:${EMBEDDING_SERVICE_PORT}"
else
  echo "  - All mode serves both /embed/text and /embed/image on port ${EMBEDDING_SERVICE_PORT}"
fi
echo

UVICORN_LOG_LEVEL="${EMBEDDING_UVICORN_LOG_LEVEL:-info}"
UVICORN_ACCESS_LOG="${EMBEDDING_UVICORN_ACCESS_LOG:-true}"
UVICORN_LOG_CONFIG="${EMBEDDING_UVICORN_LOG_CONFIG:-${PROJECT_ROOT}/config/uvicorn_embedding_logging.json}"
UVICORN_ARGS=(
  --host "${EMBEDDING_SERVICE_HOST}"
  --port "${EMBEDDING_SERVICE_PORT}"
  --workers 1
  --log-level "${UVICORN_LOG_LEVEL}"
)
if [[ -f "${UVICORN_LOG_CONFIG}" ]]; then
  UVICORN_ARGS+=(--log-config "${UVICORN_LOG_CONFIG}")
fi
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[@]}"