feat(server): add workbench API, DSH planner worker and RAG pipeline
This commit is contained in:
yi
@@ -0,0 +1,288 @@
|
||||
"""RAG 管道：Embedding（本地模型，ModelScope 下载，离线可用）"""
|
||||
import sys
|
||||
import threading
|
||||
import math
|
||||
import logging
|
||||
import time
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from ..config import MODELS_DIR, EMBEDDING_MODEL
|
||||
|
||||
_lock = threading.Lock()
|
||||
_encoder = None
|
||||
_encoder_name = None
|
||||
_encoder_load_ms = 0.0
|
||||
_metric_lock = threading.Lock()
|
||||
_last_embedding_metric: dict = {}
|
||||
_embedding_dimension = ""
|
||||
logger = logging.getLogger("yawen.rag.embed")
|
||||
|
||||
# This is deliberately a separate value from the model name. Operators can
|
||||
# bump it when replacing files in the model directory without changing the
|
||||
# configured model identifier (which also invalidates query-vector caches).
|
||||
EMBEDDING_MODEL_VERSION = os.environ.get("EMBEDDING_MODEL_VERSION", "1")
|
||||
|
||||
|
||||
def _set_metric(row: dict) -> None:
|
||||
"""Keep the most recent bounded embedding sample for the RAG envelope.
|
||||
|
||||
The RAG worker is single threaded today, but protecting this assignment
|
||||
makes the helper safe when a worker is configured with concurrent requests.
|
||||
No source text is retained in telemetry.
|
||||
"""
|
||||
global _last_embedding_metric
|
||||
clean = {
|
||||
"process_start_ms": 0.0,
|
||||
"ipc_ms": 0.0,
|
||||
"model_load_ms": 0.0,
|
||||
"tokenizer_ms": 0.0,
|
||||
"encode_ms": 0.0,
|
||||
"postprocess_ms": 0.0,
|
||||
"embedding_source": "sentence-transformers",
|
||||
"embedding_count": 0,
|
||||
"model": str(EMBEDDING_MODEL),
|
||||
"model_version": str(EMBEDDING_MODEL_VERSION),
|
||||
"pid": os.getpid(),
|
||||
"cache_hit": False,
|
||||
}
|
||||
clean.update({key: value for key, value in row.items() if value is not None})
|
||||
for key in ("process_start_ms", "ipc_ms", "model_load_ms", "tokenizer_ms",
|
||||
"encode_ms", "postprocess_ms"):
|
||||
try:
|
||||
clean[key] = round(float(clean[key]), 2)
|
||||
except (TypeError, ValueError):
|
||||
clean[key] = 0.0
|
||||
clean["embedding_count"] = int(clean.get("embedding_count") or 0)
|
||||
clean["pid"] = int(clean.get("pid") or os.getpid())
|
||||
with _metric_lock:
|
||||
_last_embedding_metric = clean
|
||||
|
||||
|
||||
def last_embedding_metric() -> dict:
|
||||
"""Return phase timings for the latest encode without exposing input text."""
|
||||
with _metric_lock:
|
||||
return dict(_last_embedding_metric)
|
||||
|
||||
|
||||
def query_cache_snapshot() -> dict:
|
||||
"""Expose the independent query-vector LRU counters for diagnostics."""
|
||||
info = _cached_query.cache_info()
|
||||
total = info.hits + info.misses
|
||||
return {
|
||||
"size": info.currsize, "max_size": info.maxsize,
|
||||
"hits": info.hits, "misses": info.misses,
|
||||
"hit_rate": round(info.hits / total, 4) if total else 0.0,
|
||||
"model": str(EMBEDDING_MODEL), "model_version": str(EMBEDDING_MODEL_VERSION),
|
||||
"dimension": int(_embedding_dimension or 0),
|
||||
"pid": os.getpid(),
|
||||
}
|
||||
|
||||
|
||||
def clear_query_cache() -> None:
|
||||
"""Clear only query vectors; retrieval candidate pools remain untouched."""
|
||||
_cached_query.cache_clear()
|
||||
|
||||
|
||||
def set_transport_timings(*, process_start_ms: float = 0.0, ipc_ms: float = 0.0) -> None:
|
||||
"""Attach process/IPC timings measured by the parent worker transport.
|
||||
|
||||
``embed.py`` cannot observe subprocess creation or pipe transfer itself;
|
||||
the worker pool can call this after a request and the next envelope then
|
||||
contains the complete set of phases. Keeping this optional preserves the
|
||||
existing embedding API for in-process callers and tests.
|
||||
"""
|
||||
metric = last_embedding_metric()
|
||||
metric.update(process_start_ms=process_start_ms, ipc_ms=ipc_ms)
|
||||
_set_metric(metric)
|
||||
|
||||
|
||||
def ensure_model_downloaded(model_name: str) -> Path:
|
||||
"""模型目录存在即用；否则从 ModelScope 下载（HF 被墙）"""
|
||||
safe = model_name.replace("/", "--")
|
||||
local = MODELS_DIR / safe
|
||||
if (local / "config.json").exists():
|
||||
return local
|
||||
try:
|
||||
from modelscope import snapshot_download
|
||||
print(f"[embed] 从 ModelScope 下载 {model_name} ...", file=sys.stderr)
|
||||
snapshot_download(model_name, local_dir=str(local))
|
||||
return local
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"模型下载失败（ModelScope）: {e}")
|
||||
|
||||
|
||||
def get_encoder():
|
||||
"""单例加载 SentenceTransformer（本地路径）"""
|
||||
global _encoder, _encoder_name, _encoder_load_ms
|
||||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||||
return _encoder
|
||||
with _lock:
|
||||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||||
return _encoder
|
||||
started = time.perf_counter()
|
||||
from sentence_transformers import SentenceTransformer
|
||||
local = ensure_model_downloaded(EMBEDDING_MODEL)
|
||||
print(f"[embed] 加载 {local} ...", file=sys.stderr)
|
||||
_encoder = SentenceTransformer(str(local))
|
||||
_encoder_name = EMBEDDING_MODEL
|
||||
_encoder_load_ms = (time.perf_counter() - started) * 1000
|
||||
return _encoder
|
||||
|
||||
|
||||
def prewarm() -> bool:
|
||||
"""Load the embedding model and run one tiny encode as a best-effort warmup.
|
||||
|
||||
Model loading is intentionally kept out of the API startup path. The
|
||||
caller should run this function from a daemon thread so a slow model
|
||||
download/load never delays health checks or request handling. Returning a
|
||||
boolean keeps this useful for startup diagnostics without making warmup
|
||||
failures fatal to the service.
|
||||
"""
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
encoder = get_encoder()
|
||||
# A real encode initializes tokenizer/runtime kernels as well as the
|
||||
# SentenceTransformer object. The result is discarded; query vectors
|
||||
# continue to use the normal cached path.
|
||||
encoder.encode(["系统启动预热"], normalize_embeddings=True,
|
||||
show_progress_bar=False, batch_size=1)
|
||||
logger.info("Embedding 模型预热完成：%s，耗时 %.3f 秒",
|
||||
EMBEDDING_MODEL, time.perf_counter() - started)
|
||||
return True
|
||||
except Exception:
|
||||
logger.exception("Embedding 模型预热失败：%s，服务将按需加载", EMBEDDING_MODEL)
|
||||
return False
|
||||
|
||||
|
||||
EMBEDDING_POLICY = "full-token-windows-v1"
|
||||
|
||||
|
||||
def text_windows(text: str, encoder=None) -> list[str]:
|
||||
"""Cover every source character without sending over-limit windows to BGE."""
|
||||
enc = encoder or get_encoder()
|
||||
tokenizer = enc.tokenizer
|
||||
maximum = int(enc.max_seq_length)
|
||||
def fits(value):
|
||||
return len(tokenizer.encode(value, add_special_tokens=True, truncation=False)) <= maximum
|
||||
if fits(text):
|
||||
return [text]
|
||||
# Character boundaries preserve exact original text; a binary search uses the
|
||||
# actual model tokenizer rather than assuming that one character is one token.
|
||||
windows, offset = [], 0
|
||||
while offset < len(text):
|
||||
remaining = text[offset:]
|
||||
if fits(remaining):
|
||||
windows.append(remaining)
|
||||
break
|
||||
low, high = 1, min(len(remaining), maximum * 4)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if fits(remaining[:middle]):
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if not fits(remaining[:low]):
|
||||
raise ValueError("单个字符超过向量模型输入上限")
|
||||
windows.append(remaining[:low])
|
||||
offset += low
|
||||
return windows
|
||||
|
||||
|
||||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
# ``get_encoder`` includes model download and construction only on the
|
||||
# first call in each resident worker. Subsequent calls report zero load
|
||||
# time, which makes accidental worker restarts visible in telemetry.
|
||||
enc = get_encoder()
|
||||
global _encoder_load_ms
|
||||
model_load_ms = _encoder_load_ms
|
||||
# Consume the one-shot load measurement. A resident encoder must report
|
||||
# zero on later requests; a non-zero value then reliably indicates a worker
|
||||
# restart or an explicit model replacement.
|
||||
_encoder_load_ms = 0.0
|
||||
# At most 64 originals and 64 encoded windows in one model batch. Long
|
||||
# paragraphs are pooled in the existing vector space, keeping chunk IDs and
|
||||
# consumption records stable while representing their tails as well.
|
||||
output = []
|
||||
tokenizer_ms = encode_ms = postprocess_ms = 0.0
|
||||
flat_count = 0
|
||||
for start in range(0, len(texts), 64):
|
||||
tokenizer_started = time.perf_counter()
|
||||
groups = [text_windows(str(text), enc) for text in texts[start:start + 64]]
|
||||
flat = [window for group in groups for window in group]
|
||||
tokenizer_ms += (time.perf_counter() - tokenizer_started) * 1000
|
||||
flat_count += len(flat)
|
||||
encode_started = time.perf_counter()
|
||||
vectors = enc.encode(flat, normalize_embeddings=True, show_progress_bar=False, batch_size=64)
|
||||
encode_ms += (time.perf_counter() - encode_started) * 1000
|
||||
postprocess_started = time.perf_counter()
|
||||
cursor = 0
|
||||
for group in groups:
|
||||
selected = vectors[cursor:cursor + len(group)]
|
||||
cursor += len(group)
|
||||
if len(group) == 1:
|
||||
output.append([float(v) for v in selected[0]])
|
||||
continue
|
||||
weights = [max(1, len(enc.tokenizer.encode(window, add_special_tokens=False, truncation=False)))
|
||||
for window in group]
|
||||
combined = [sum(float(vector[i]) * weight for vector, weight in zip(selected, weights))
|
||||
for i in range(len(selected[0]))]
|
||||
norm = math.sqrt(sum(value * value for value in combined)) or 1.0
|
||||
output.append([value / norm for value in combined])
|
||||
postprocess_ms += (time.perf_counter() - postprocess_started) * 1000
|
||||
_set_metric({
|
||||
"model_load_ms": model_load_ms,
|
||||
"tokenizer_ms": tokenizer_ms,
|
||||
"encode_ms": encode_ms,
|
||||
"postprocess_ms": postprocess_ms,
|
||||
"embedding_count": len(texts),
|
||||
# ``flat_count`` is useful when long documents are split into windows;
|
||||
# retain it as a diagnostic without changing the public return value.
|
||||
"window_count": flat_count,
|
||||
"embedding_dimension": len(output[0]) if output else 0,
|
||||
"cache_hit": False,
|
||||
})
|
||||
return output
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _cached_query(model: str, policy: str, version: str, dimension: str, text: str) -> tuple[float, ...]:
|
||||
"""Cache only query vectors, keyed by model identity and policy version."""
|
||||
return tuple(embed_texts([text])[0])
|
||||
|
||||
|
||||
def embed_query(text: str) -> list[float]:
|
||||
# lru_cache does not expose whether this invocation hit without comparing
|
||||
# its counters. The comparison is process local and avoids retaining the
|
||||
# query itself in telemetry.
|
||||
before = _cached_query.cache_info()
|
||||
global _embedding_dimension
|
||||
# Keep the cache identity stable across the first call: the actual vector
|
||||
# dimension is discovered only after encoding. Model/version already
|
||||
# identify the encoder contract; ``auto`` prevents the first "unknown"
|
||||
# dimension from creating a second key for the same query.
|
||||
result = list(_cached_query(EMBEDDING_MODEL, EMBEDDING_POLICY,
|
||||
EMBEDDING_MODEL_VERSION, "auto", str(text)))
|
||||
if not _embedding_dimension and result:
|
||||
_embedding_dimension = str(len(result))
|
||||
after = _cached_query.cache_info()
|
||||
cache_hit = after.hits > before.hits
|
||||
metric = last_embedding_metric()
|
||||
metric.update(
|
||||
embedding_count=1,
|
||||
cache_hit=cache_hit,
|
||||
embedding_source="query-vector-lru",
|
||||
model=EMBEDDING_MODEL,
|
||||
model_version=EMBEDDING_MODEL_VERSION,
|
||||
pid=os.getpid(),
|
||||
embedding_dimension=len(result) if result else 0,
|
||||
)
|
||||
if cache_hit:
|
||||
# A cache hit performs no tokenizer, model inference, or serialization;
|
||||
# do not leak timings from the previous (possibly indexing) call.
|
||||
metric.update(model_load_ms=0.0, tokenizer_ms=0.0,
|
||||
encode_ms=0.0, postprocess_ms=0.0, window_count=0)
|
||||
_set_metric(metric)
|
||||
return result
|
||||
Reference in New Issue
Block a user