clossers/didi-lot1-ai
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
didi-lot1-ai/ai_platform/modules/embeddings/INDEX.md

12 KiB
Raw Blame History

Embeddings Module — INDEX

Embedding generation API for DIDI knowledge graph + semantic search. Serves BAAI/bge-m3 (1024-dim, multilingual, 8K context) via vLLM (or llama.cpp) behind an OpenAI-compatible FastAPI wrapper, used by didi-brain for atom storage in pgvector and for /v1/gather query embedding.

  • Stack: Python 3.10+, FastAPI, uvicorn, httpx, Pydantic v2; backends via vllm/vllm-openai:v0.8.5 or ghcr.io/ggml-org/llama.cpp:server
  • URLs:
    • API wrapper: http://10.11.10.12:14100 (Prod) / http://10.11.10.12:54100 (Dev)
    • vLLM internal: :14101 (Prod) / :54101 (Dev)
    • llama.cpp internal: :14110 (Prod) / :54110 (Dev)
  • Container layout (2-tier):
    • didiAI-embeddings-api (FastAPI wrapper + auth + rate limiting)
    • didiAI-embeddings-vllm (GPU embedding server) OR didiAI-embeddings-llamacpp (CPU/lightweight server)
  • Default model: BAAI/bge-m3 — 1024 dim, multilingual, 8K context

Ce face

Generate text embeddings for downstream DIDI consumers:

  • didi-brain semantic atom storage — text claims/atoms get embedded once at ingestion, stored as pgvector columns.
  • didi-brain /v1/gather query embedding — incoming claim text gets embedded, used as the kNN probe vector.
  • Cross-encoder reranking — separate rerank module (port 14200) handles re-scoring; this module only does dense embeddings.
  • Catalog / retrieval consumers — anything that needs a vector representation hits this single endpoint.

The module is OpenAI-compatible: drop-in for openai SDK code that uses client.embeddings.create().

API endpoints

All routes are mounted by src/embeddings/api/app.py via create_app(). Health is unauthenticated and excluded from rate-limiting; everything else passes through verify_bearer_token if EMB_API_TOKENS is set.

Method Path Description
POST /v1/embeddings OpenAI-compatible. Body: {input, model, encoding_format?, dimensions?, backend?}. Returns {object, data[], model, usage, backend}.
GET /v1/models List available embedding models. Optional ?backend=vllm|llamacpp filter.
GET /v1/backends List enabled backends as {backends: ["vllm", ...]}.
GET /health Detailed health: per-backend status, overall healthy/degraded/unhealthy.
GET /ready Simple readiness probe ({ready: true}) — for k8s.

Source files:

  • /home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/embeddings.py — POST /v1/embeddings
  • /home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/models.py — /v1/models, /v1/backends
  • /home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/health.py — /health, /ready

Errors: 400 (invalid params / backend disabled), 401 (auth), 429 (rate limit, sets Retry-After), 503 (backend connection), 504 (backend timeout), 500 (other backend failures).

Backends

Two interchangeable engines, selected per-request via the backend field or globally via EMB_DEFAULT_BACKEND:

  • vLLM (production / GPU) — vllm/vllm-openai:v0.8.5 launched with --task embed, high-throughput batching, NVIDIA CUDA 12.x required. Listens internally on 14101 (Prod) / 54101 (Dev).
  • llama.cpp (dev / lightweight) — ghcr.io/ggml-org/llama.cpp:server-b4769 launched with --embedding, runs on CPU or GPU with GGUF-quantized models. Listens on 14110 / 54110.
  • API wrapper — own FastAPI app on 14100 (Prod) / 54100 (Dev) that proxies to whichever backend(s) are enabled.

Both backends speak OpenAI-compatible HTTP, so the wrapper just routes via httpx/openai SDK. Backend abstraction lives in src/embeddings/backends/{base,vllm_backend,llamacpp_backend,registry}.py.

Structura fisiere

embeddings/
├── INDEX.md                   (this file)
├── README.md                  quick-start + config table
├── API.md                     full HTTP API reference + SDK examples
├── pyproject.toml             package def, optional extras: vllm, llamacpp, all, dev
├── .env.example               all EMB_* env vars documented
├── deploy/
│   ├── Dockerfile             multi-stage, python 3.11-slim + uv
│   ├── docker-compose.yml     3 services + 3 profiles (api/vllm/llamacpp)
│   └── deploy.sh              wrapper around `docker compose --profile`
├── src/embeddings/
│   ├── __init__.py            re-exports EmbeddingClient
│   ├── cli.py                 `python -m embeddings.cli` entry point
│   ├── client.py              EmbeddingClient (Python library API)
│   ├── config.py              EmbeddingSettings (pydantic-settings, EMB_ prefix)
│   ├── schemas.py             EmbeddingRequest/Response/Data, base64 encoder
│   ├── types.py               BackendType enum, ModelInfo
│   ├── exceptions.py          BackendNotAvailable/Enabled, RateLimit, Timeout, etc.
│   ├── logging.py             structured logging (text or JSON)
│   ├── api/
│   │   ├── app.py             create_app() factory + lifespan
│   │   ├── dependencies.py    FastAPI deps (auth, concurrency limiter, client)
│   │   ├── middleware.py      RateLimitMiddleware, RequestIdMiddleware
│   │   └── routes/
│   │       ├── embeddings.py  POST /v1/embeddings
│   │       ├── models.py      /v1/models, /v1/backends
│   │       └── health.py      /health, /ready
│   └── backends/
│       ├── base.py            abstract Backend interface
│       ├── vllm_backend.py    vLLM HTTP client
│       ├── llamacpp_backend.py llama.cpp HTTP client
│       └── registry.py        BackendRegistry (selects + caches backends)
└── tests/
    ├── conftest.py
    ├── test_config.py
    ├── test_schemas.py
    └── test_types.py

How didi-brain uses it

didi-brain consumes this service through shared/embedding_client.py (a thin wrapper around the openai SDK pointed at this API):

  • Ingestion path — when atoms/claims are written, brain calls EmbeddingClient.embed(text) to get a 1024-dim vector and stores it in pgvector alongside the row. One round-trip per batch.
  • Query path — for /v1/gather, the incoming claim text is embedded the same way, then the resulting vector is used as the probe in a pgvector <=> (cosine) kNN search to retrieve candidate atoms.
  • Reranking — top-k candidates from the dense search are forwarded to the separate rerank module (cross-encoder, port 14200) for fine-grained scoring. That module does not call this one — they're parallel concerns.
  • Catalog API — also retrieves via the same embedding pipeline (text → vector → kNN), reusing this single endpoint.

Because the wrapper is OpenAI-compatible, EmbeddingClient can be a vanilla openai.OpenAI(base_url=..., api_key=...) instance — no DIDI-specific client code needed in brain.

Configuration

All env vars use the EMB_ prefix. Required vars have no defaults — the app refuses to start without them.

Required

Variable Description
EMB_DEFAULT_BACKEND vllm or llamacpp
EMB_ENABLE_VLLM true / false
EMB_ENABLE_LLAMACPP true / false
EMB_EXTERNAL_URL Public URL exposed in the OpenAPI spec

Common optional

Variable Default Description
EMB_PORT 14100 (Prod) / 54100 (Dev) API server port
EMB_HOST 0.0.0.0 API bind address
EMB_API_TOKENS unset Comma-separated bearer tokens; auth disabled if unset
EMB_VLLM_BASE_URL http://localhost:54101 Where the API wrapper reaches vLLM
EMB_VLLM_API_KEY unset Optional API key for the vLLM server
EMB_LLAMACPP_BASE_URL http://localhost:54110 Where the API wrapper reaches llama.cpp
EMB_REQUEST_TIMEOUT 120.0 Per-request backend timeout (s)
EMB_CONNECT_TIMEOUT 10.0 TCP connect timeout (s)
EMB_RATE_LIMIT_RPS 20.0 Token-bucket rate (req/s)
EMB_RATE_LIMIT_BURST 40 Token-bucket burst capacity
EMB_MAX_CONCURRENT_REQUESTS 20 In-flight cap
EMB_LOG_LEVEL INFO DEBUG / INFO / WARNING / ERROR
EMB_LOG_JSON false Emit JSON-formatted log lines

The API app reads only the variables above (fields on EmbeddingSettings in config.py). EMB_VLLM_BASE_URL / EMB_LLAMACPP_BASE_URL just tell the wrapper where to reach the backend servers.

Backend service variables (docker-compose only)

These are not fields on EmbeddingSettings — they are consumed by deploy/docker-compose.yml to launch the vLLM / llama.cpp containers (model id, GPU pinning, memory, context). They configure the backend server, not the API app.

Variable Default Description
EMB_VLLM_MODEL BAAI/bge-m3 HF model id loaded by the vLLM container
EMB_VLLM_GPU 0 CUDA_VISIBLE_DEVICES for the vLLM container
EMB_VLLM_GPU_UTIL 0.50 vLLM --gpu-memory-utilization
EMB_VLLM_MAX_LEN 8192 vLLM --max-model-len
EMB_LLAMACPP_MODEL bge-m3-q4_k_m.gguf GGUF filename inside MODELS_DIR
EMB_LLAMACPP_CTX 8192 llama.cpp context size
EMB_LLAMACPP_THREADS 4 llama.cpp threads
EMB_LLAMACPP_PARALLEL 4 llama.cpp parallel slots
HF_CACHE_DIR /cai2_ds_storage/hf_cache Mounted into the vLLM container
HF_TOKEN unset Forwarded as HUGGING_FACE_HUB_TOKEN
MODELS_DIR /cai2_ds_storage/models GGUF model directory for llama.cpp

See /home/admin365/didi_mono/ai_platform/modules/embeddings/.env.example for the canonical, fully commented list.

Deployment

GPU host required for the production profile (e.g. 10.11.10.17). All three services live on the shared didi-network Docker network.

cd /home/admin365/didi_mono/ai_platform/modules/embeddings/deploy

# Configure
cp ../.env.example .env
$EDITOR .env

# Production: API + vLLM (GPU)
./deploy.sh --profile vllm -d

# Lightweight: API + llama.cpp
./deploy.sh --profile llamacpp -d

# API only (use external embedding servers via EMB_*_BASE_URL)
./deploy.sh --profile api -d

# Logs / shutdown
./deploy.sh --profile vllm --logs
./deploy.sh --profile vllm --down

Compose profiles:

  • api — only the FastAPI wrapper
  • vllm — wrapper + vLLM (GPU)
  • llamacpp — wrapper + llama.cpp (CPU/GPU)

Performance / model

  • Model: BAAI/bge-m3 — 1024 dim, multilingual (100+ languages), 8K token context, supports dense + sparse + multi-vector (this module uses dense only).
  • Throughput: ~100–200 embeddings/s on H100 with vLLM batching (depends on input length).
  • Latency P95: ~50–100 ms per request (single text, warm GPU). Larger batches amortize well — keep request batches at 16–64 inputs for best throughput.
  • Concurrency knobs: EMB_MAX_CONCURRENT_REQUESTS (in-flight at the wrapper) and EMB_RATE_LIMIT_RPS (token-bucket) bound the load reaching the GPU.
  • GPU memory: bge-m3 fits comfortably in <4 GB; EMB_VLLM_GPU_UTIL=0.50 is intentionally low to allow GPU sharing with reranker / other workloads.

Related

  • didi-brain consumes via shared/embedding_client.py for both ingestion (text → pgvector storage) and /v1/gather query embedding.
  • rerank module (separate, port 14200) handles cross-encoder scoring on top of dense kNN candidates from this module — they're complementary, not chained inside this service.
  • Catalog API uses the same endpoint for retrieval-side embeddings.
  • pgvector in the brain Postgres stores the resulting 1024-dim vectors (cosine distance index).
  • Local Python use: from embeddings import EmbeddingClient (works without the HTTP wrapper if you want in-process inference and have the vllm/llamacpp extras installed).