Optimize for 4 cores: BM25 preload, httpx pool, orjson, fix lambda

- index: UVICORN_WORKERS 8→4, lifespan BM25 preload, explicit ThreadPoolExecutor(4), orjson
- search: lifespan BM25 preload, httpx limits (max_conn=100, keepalive=20, timeout=30s), fix asyncio.to_thread lambda, orjson
- both: ORJSONResponse as default_response_class

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
q 2026-04-18 21:21:28 +03:00
parent 91aae61b65
commit a4475abfb2
4 changed files with 48 additions and 16 deletions

View file

@ -1,17 +1,19 @@
import asyncio import asyncio
import logging import logging
import os import os
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from functools import lru_cache from functools import lru_cache
from typing import Any from typing import Any
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse, ORJSONResponse
from pydantic import BaseModel from pydantic import BaseModel
HOST = os.getenv("HOST", "0.0.0.0") HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "8000")) PORT = int(os.getenv("PORT", "8000"))
UVICORN_WORKERS = 8 UVICORN_WORKERS = 4
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
logger = logging.getLogger("index-service") logger = logging.getLogger("index-service")
@ -78,6 +80,24 @@ OVERLAP_SIZE = 128
SPARSE_MODEL_NAME = "Qdrant/bm25" SPARSE_MODEL_NAME = "Qdrant/bm25"
FASTEMBED_CACHE_PATH = "/models/fastembed" FASTEMBED_CACHE_PATH = "/models/fastembed"
_thread_pool = ThreadPoolExecutor(max_workers=4)
@lru_cache(maxsize=1)
def get_sparse_model():
from fastembed import SparseTextEmbedding
logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH)
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
@asynccontextmanager
async def lifespan(app: FastAPI):
# Preload BM25 model on startup to avoid cold-start latency
await asyncio.to_thread(get_sparse_model)
logger.info("BM25 model preloaded")
yield
def render_message(message: Message) -> str: def render_message(message: Message) -> str:
parts_list: list[str] = [] parts_list: list[str] = []
@ -186,7 +206,12 @@ def build_chunks(
return result return result
app = FastAPI(title="Index Service", version="0.1.0") app = FastAPI(
title="Index Service",
version="0.1.0",
lifespan=lifespan,
default_response_class=ORJSONResponse,
)
@app.get("/health") @app.get("/health")
@ -205,14 +230,6 @@ async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
) )
@lru_cache(maxsize=1)
def get_sparse_model():
from fastembed import SparseTextEmbedding
logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH)
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
def embed_sparse_texts(texts: list[str]) -> list[dict]: def embed_sparse_texts(texts: list[str]) -> list[dict]:
model = get_sparse_model() model = get_sparse_model()
vectors = [] vectors = []
@ -228,7 +245,9 @@ def embed_sparse_texts(texts: list[str]) -> list[dict]:
@app.post("/sparse_embedding") @app.post("/sparse_embedding")
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]: async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts) vectors = await asyncio.get_event_loop().run_in_executor(
_thread_pool, embed_sparse_texts, payload.texts
)
return {"vectors": vectors} return {"vectors": vectors}

View file

@ -2,3 +2,4 @@ fastapi==0.135.1
uvicorn[standard]==0.42.0 uvicorn[standard]==0.42.0
pydantic==2.12.5 pydantic==2.12.5
fastembed==0.7.4 fastembed==0.7.4
orjson==3.10.18

View file

@ -9,7 +9,7 @@ import httpx
from fastembed import SparseTextEmbedding from fastembed import SparseTextEmbedding
from fastapi import FastAPI, HTTPException, Request from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse, ORJSONResponse
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from qdrant_client import AsyncQdrantClient, models from qdrant_client import AsyncQdrantClient, models
@ -148,7 +148,13 @@ def get_sparse_model() -> SparseTextEmbedding:
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
app.state.http = httpx.AsyncClient() # Preload BM25 and set up HTTP client with connection pooling
await asyncio.to_thread(get_sparse_model)
logger.info("BM25 model preloaded")
app.state.http = httpx.AsyncClient(
timeout=30.0,
limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
)
app.state.qdrant = AsyncQdrantClient( app.state.qdrant = AsyncQdrantClient(
url=QDRANT_URL, url=QDRANT_URL,
api_key=API_KEY, api_key=API_KEY,
@ -160,7 +166,12 @@ async def lifespan(app: FastAPI):
await app.state.qdrant.close() await app.state.qdrant.close()
app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan) app = FastAPI(
title="Search Service",
version="0.1.0",
lifespan=lifespan,
default_response_class=ORJSONResponse,
)
DENSE_PREFETCH_K = 80 DENSE_PREFETCH_K = 80
SPARSE_PREFETCH_K = 200 SPARSE_PREFETCH_K = 200
@ -354,7 +365,7 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
sparse_query = build_sparse_query(question) sparse_query = build_sparse_query(question)
dense_task = embed_dense(client, dense_query) dense_task = embed_dense(client, dense_query)
sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query)) sparse_task = asyncio.to_thread(embed_sparse_sync, sparse_query)
dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task) dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
dense_vectors = [dense_vector] dense_vectors = [dense_vector]

View file

@ -4,3 +4,4 @@ pydantic==2.12.5
httpx==0.28.1 httpx==0.28.1
qdrant-client==1.15.1 qdrant-client==1.15.1
fastembed==0.7.4 fastembed==0.7.4
orjson==3.10.18