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:
parent
91aae61b65
commit
a4475abfb2
4 changed files with 48 additions and 16 deletions
|
|
@ -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}
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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]
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue