Coverage for src/asketmc_bot/rerank.py: 59%
248 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1#!/usr/bin/env python3.10
2"""CrossEncoder reranker for NodeWithScore candidates.
4Public API:
5 * await init_reranker(force: bool = False) -> None
6 * await shutdown_reranker() -> None
7 * await rerank(query: str, nodes: List[NodeWithScore]) -> List[NodeWithScore]
9Behavior:
10 - Configuration is read from `config` with defaults.
11 - Lazy initialization guarded by an asyncio.Lock.
12 - Inference is executed in a ThreadPoolExecutor via run_in_executor.
13 - A timeout is applied to waiting for the inference result; the underlying thread work cannot be forcibly stopped.
14 - Optional retries on timeout/exception.
15 - Query validation: strip, non-empty, max length, allowlist regex.
17Logging:
18 - Uses stdlib logging under logger name "asketmc.rerank".
19"""
21from __future__ import annotations
23import asyncio
24import logging
25import re
26import sys
27import time
28from concurrent.futures import ThreadPoolExecutor
29from typing import List, Optional, Tuple
31import torch
32from llama_index.core.schema import NodeWithScore
33from sentence_transformers import CrossEncoder
35from asketmc_bot import config as cfg
37# -----------------------------------------------------------------------------
38# Logging
39# -----------------------------------------------------------------------------
41_LOG = logging.getLogger("asketmc.rerank")
42_LOG.setLevel(logging.DEBUG if getattr(cfg, "DEBUG", False) else logging.INFO)
43if not _LOG.handlers:
44 _handler = logging.StreamHandler(sys.stdout)
45 _handler.setFormatter(
46 logging.Formatter("[%(asctime)s] [%(levelname)s] %(name)s: %(message)s")
47 )
48 _LOG.addHandler(_handler)
50# -----------------------------------------------------------------------------
51# Configuration defaults
52# -----------------------------------------------------------------------------
54_DEFAULT_ALLOWED_RE = (
55 r"^[\w\s\.\,\:\;\!\?\(\)\-\[\]\{\}'\"/@#%&\*\+\=\<\>\$€£…©®™°|\\`~№"
56 r"А-Яа-яЁё]+$"
57)
58ALLOWED_CHARS_REGEX = getattr(cfg, "ALLOWED_CHARS_REGEX", _DEFAULT_ALLOWED_RE)
59_QUERY_MAX_CHARS = int(getattr(cfg, "QUERY_MAX_CHARS", 2048))
61_RERANKER_MODEL_NAME = getattr(cfg, "RERANKER_MODEL_NAME", "BAAI/bge-reranker-large")
62_RERANK_INPUT_K = int(getattr(cfg, "RERANK_INPUT_K", 20))
63_RERANK_OUTPUT_K = int(getattr(cfg, "RERANK_OUTPUT_K", 5))
64_MAX_LEN = int(getattr(cfg, "MAX_LEN", 512))
65_BATCH_SIZE = int(getattr(cfg, "BATCH_SIZE", 16))
66_RERANKER_DEVICE = str(getattr(cfg, "RERANKER_DEVICE", "cpu")).lower()
67_EXECUTOR_WORKERS = int(getattr(cfg, "EXECUTOR_WORKERS", 4))
68_PREDICT_TIMEOUT_SEC = float(getattr(cfg, "PREDICT_TIMEOUT_SEC", 120.0))
69_PREDICT_RETRIES = int(getattr(cfg, "PREDICT_RETRIES", 1)) # 0 means no retries
71_SHUTDOWN_WAIT_SEC = float(getattr(cfg, "RERANK_SHUTDOWN_WAIT_SEC", 10.0))
73# -----------------------------------------------------------------------------
74# Globals (lazy-initialized; protected by _INIT_LOCK)
75# -----------------------------------------------------------------------------
77_RERANKER: Optional[CrossEncoder] = None
78_EXECUTOR: Optional[ThreadPoolExecutor] = None
79_INIT_LOCK = asyncio.Lock()
81_INFLIGHT: Optional[asyncio.Future[List[float]]] = None
83# -----------------------------------------------------------------------------
84# Helpers
85# -----------------------------------------------------------------------------
87_ALLOWED_RE = re.compile(ALLOWED_CHARS_REGEX, re.UNICODE)
90def _sanitize_input(text: str, *, max_len: int = _QUERY_MAX_CHARS) -> str:
91 """Validate and normalize user input."""
92 if not isinstance(text, str): 92 ↛ 93line 92 didn't jump to line 93, because the condition on line 92 was never true
93 raise ValueError("query must be a string")
94 s = text.strip()
95 if not s: 95 ↛ 96line 95 didn't jump to line 96, because the condition on line 95 was never true
96 raise ValueError("query is empty")
97 if len(s) > max_len: 97 ↛ 98line 97 didn't jump to line 98, because the condition on line 97 was never true
98 raise ValueError(f"query exceeds {max_len} characters")
99 if not _ALLOWED_RE.fullmatch(s): 99 ↛ 100line 99 didn't jump to line 100, because the condition on line 99 was never true
100 raise ValueError("query contains disallowed characters")
101 return s
104def _choose_device() -> str:
105 """Return device string based on config and CUDA availability."""
106 if _RERANKER_DEVICE == "cpu": 106 ↛ 108line 106 didn't jump to line 108, because the condition on line 106 was always true
107 return "cpu"
108 if _RERANKER_DEVICE == "cuda":
109 if torch.cuda.is_available():
110 return "cuda"
111 raise RuntimeError("RERANKER_DEVICE='cuda' but CUDA is not available")
112 raise ValueError("RERANKER_DEVICE must be 'cpu' or 'cuda'")
115def _node_text(n: NodeWithScore) -> str:
116 """Extract textual content from a NodeWithScore safely."""
117 node = getattr(n, "node", n)
118 if hasattr(node, "get_content"): 118 ↛ 123line 118 didn't jump to line 123, because the condition on line 118 was always true
119 try:
120 return node.get_content() or ""
121 except Exception:
122 return ""
123 if hasattr(node, "text"):
124 return getattr(node, "text", "") or ""
125 return ""
128def _node_id(n: NodeWithScore) -> str:
129 """Extract stable node identifier for logs."""
130 node = getattr(n, "node", n)
131 for attr in ("node_id", "id_", "id", "ref_doc_id"): 131 ↛ 139line 131 didn't jump to line 139, because the loop on line 131 didn't complete
132 if hasattr(node, attr):
133 try:
134 v = getattr(node, attr)
135 if v: 135 ↛ 131line 135 didn't jump to line 131, because the condition on line 135 was always true
136 return str(v)
137 except Exception:
138 continue
139 return "n/a"
142def _filter_pairs(
143 query: str, nodes: List[NodeWithScore]
144) -> Tuple[List[List[str]], List[NodeWithScore]]:
145 """Return (query, doc) pairs and corresponding nodes, skipping empty docs."""
146 cand_nodes: List[NodeWithScore] = nodes[:_RERANK_INPUT_K]
147 pairs: List[List[str]] = []
148 filtered_nodes: List[NodeWithScore] = []
150 for idx, n in enumerate(cand_nodes):
151 content = _node_text(n)
152 if content.strip(): 152 ↛ 156line 152 didn't jump to line 156, because the condition on line 152 was always true
153 pairs.append([query, content])
154 filtered_nodes.append(n)
155 else:
156 _LOG.debug(
157 "[_filter_pairs] skip empty doc idx=%d node_id=%s",
158 idx,
159 _node_id(n),
160 )
162 _LOG.debug(
163 "[_filter_pairs] selected %d/%d candidates",
164 len(filtered_nodes),
165 len(cand_nodes),
166 )
167 return pairs, filtered_nodes
170def _clear_inflight_cb(fut: asyncio.Future) -> None:
171 global _INFLIGHT
172 if _INFLIGHT is fut: 172 ↛ exitline 172 didn't return from function '_clear_inflight_cb', because the condition on line 172 was always true
173 _INFLIGHT = None
176# -----------------------------------------------------------------------------
177# Lifecycle
178# -----------------------------------------------------------------------------
180async def init_reranker(force: bool = False) -> None:
181 """Initialize CrossEncoder and thread pool safely; re-init on force=True."""
182 global _RERANKER, _EXECUTOR, _INFLIGHT
184 async with _INIT_LOCK:
185 if _INFLIGHT is not None and not _INFLIGHT.done(): 185 ↛ 186line 185 didn't jump to line 186, because the condition on line 185 was never true
186 _LOG.info("[init_reranker] inference in-flight; skip init")
187 return
189 if _RERANKER is not None and _EXECUTOR is not None and not force:
190 _LOG.info("[init_reranker] already initialized; skip (force=False)")
191 return
193 if _EXECUTOR is not None: 193 ↛ 194line 193 didn't jump to line 194, because the condition on line 193 was never true
194 _LOG.info("[init_reranker] shutting down previous executor")
195 try:
196 _EXECUTOR.shutdown(wait=True, cancel_futures=True)
197 except TypeError:
198 _EXECUTOR.shutdown(wait=True)
199 _EXECUTOR = None
201 if _RERANKER is not None: 201 ↛ 202line 201 didn't jump to line 202, because the condition on line 201 was never true
202 try:
203 if torch.cuda.is_available():
204 mem_before = torch.cuda.memory_allocated()
205 del _RERANKER
206 _RERANKER = None
207 if torch.cuda.is_available():
208 torch.cuda.empty_cache()
209 mem_after = torch.cuda.memory_allocated()
210 _LOG.info(
211 "[init_reranker] CUDA cache cleared: %.2f -> %.2f MB",
212 mem_before / (1024**2),
213 mem_after / (1024**2),
214 )
215 except Exception as exc: # pragma: no cover
216 _LOG.warning("[init_reranker] error during previous model cleanup: %s", exc)
218 device = _choose_device()
219 _LOG.info(
220 "[init_reranker] loading CrossEncoder model=%s device=%s max_length=%d workers=%d",
221 _RERANKER_MODEL_NAME,
222 device,
223 _MAX_LEN,
224 _EXECUTOR_WORKERS,
225 )
226 _RERANKER = CrossEncoder(
227 _RERANKER_MODEL_NAME,
228 device=device,
229 max_length=_MAX_LEN,
230 )
231 _EXECUTOR = ThreadPoolExecutor(max_workers=_EXECUTOR_WORKERS)
232 _LOG.info(
233 "[init_reranker] CrossEncoder ready (id=%r) on %s; executor workers=%d",
234 id(_RERANKER),
235 device,
236 _EXECUTOR_WORKERS,
237 )
240async def shutdown_reranker() -> None:
241 """Release thread pool and GPU memory (if any)."""
242 global _RERANKER, _EXECUTOR, _INFLIGHT
244 async with _INIT_LOCK:
245 _LOG.info("[shutdown_reranker] called")
247 inflight = _INFLIGHT
248 if inflight is not None and not inflight.done(): 248 ↛ 249line 248 didn't jump to line 249, because the condition on line 248 was never true
249 try:
250 await asyncio.wait_for(asyncio.shield(inflight), timeout=_SHUTDOWN_WAIT_SEC)
251 except asyncio.TimeoutError:
252 _LOG.warning(
253 "[shutdown_reranker] in-flight inference did not finish within %.1fs; "
254 "skipping model/CUDA cleanup",
255 _SHUTDOWN_WAIT_SEC,
256 )
257 if _EXECUTOR is not None:
258 try:
259 _EXECUTOR.shutdown(wait=False, cancel_futures=True)
260 except TypeError:
261 _EXECUTOR.shutdown(wait=False)
262 _EXECUTOR = None
263 return
265 try:
266 if _EXECUTOR is not None: 266 ↛ 273line 266 didn't jump to line 273, because the condition on line 266 was always true
267 try:
268 _EXECUTOR.shutdown(wait=True, cancel_futures=True)
269 except TypeError:
270 _EXECUTOR.shutdown(wait=True)
271 _EXECUTOR = None
273 if _RERANKER is not None: 273 ↛ 290line 273 didn't jump to line 290, because the condition on line 273 was always true
274 try:
275 if torch.cuda.is_available(): 275 ↛ 276line 275 didn't jump to line 276, because the condition on line 275 was never true
276 mem_before = torch.cuda.memory_allocated()
277 del _RERANKER
278 _RERANKER = None
279 if torch.cuda.is_available(): 279 ↛ 280line 279 didn't jump to line 280, because the condition on line 279 was never true
280 torch.cuda.empty_cache()
281 mem_after = torch.cuda.memory_allocated()
282 _LOG.info(
283 "[shutdown_reranker] CUDA cache cleared: %.2f -> %.2f MB",
284 mem_before / (1024**2),
285 mem_after / (1024**2),
286 )
287 except Exception as exc: # pragma: no cover
288 _LOG.warning("[shutdown_reranker] error during model cleanup: %s", exc)
289 finally:
290 _INFLIGHT = None
291 _LOG.info("[shutdown_reranker] resources released")
294# -----------------------------------------------------------------------------
295# Rerank
296# -----------------------------------------------------------------------------
298async def rerank(query: str, nodes: List[NodeWithScore]) -> List[NodeWithScore]:
299 """Score candidate nodes with a CrossEncoder and return top-K results."""
300 global _RERANKER, _EXECUTOR, _INFLIGHT
302 query = _sanitize_input(query, max_len=_QUERY_MAX_CHARS)
304 if not nodes:
305 _LOG.debug("[rerank] empty input nodes -> []")
306 return []
308 pairs, valid_nodes = _filter_pairs(query, nodes)
309 if not pairs: 309 ↛ 310line 309 didn't jump to line 310, because the condition on line 309 was never true
310 _LOG.warning("[rerank] no valid documents after filtering -> []")
311 return []
313 loop = asyncio.get_running_loop()
315 attempt = 0
316 last_err: Optional[Exception] = None
318 while attempt <= _PREDICT_RETRIES: 318 ↛ 455line 318 didn't jump to line 455, because the condition on line 318 was always true
319 attempt += 1
321 # Ensure initialized WITHOUT holding _INIT_LOCK here (avoid deadlock).
322 if _RERANKER is None or _EXECUTOR is None: 322 ↛ 323line 322 didn't jump to line 323, because the condition on line 322 was never true
323 _LOG.info("[rerank] reranker not initialized; attempting init")
324 try:
325 await init_reranker(force=False)
326 except Exception as exc:
327 _LOG.exception("[rerank] init error -> []: %s", exc)
328 return []
329 if _RERANKER is None or _EXECUTOR is None:
330 _LOG.error("[rerank] init failed -> []")
331 return []
333 # If another inference is running, wait for it (do not return [] silently).
334 inflight_to_wait: Optional[asyncio.Future[List[float]]] = None
335 async with _INIT_LOCK:
336 if _INFLIGHT is not None and not _INFLIGHT.done(): 336 ↛ 335line 336 didn't jump to line 335
337 inflight_to_wait = _INFLIGHT
339 if inflight_to_wait is not None: 339 ↛ 340line 339 didn't jump to line 340, because the condition on line 339 was never true
340 _LOG.info("[rerank] another inference in-flight; awaiting completion")
341 try:
342 await asyncio.wait_for(
343 asyncio.shield(inflight_to_wait),
344 timeout=_PREDICT_TIMEOUT_SEC,
345 )
346 except asyncio.TimeoutError as exc:
347 _LOG.warning(
348 "[rerank] prior in-flight inference did not finish within %.1fs -> []",
349 _PREDICT_TIMEOUT_SEC,
350 )
351 last_err = exc
352 return []
353 except Exception as exc:
354 _LOG.warning(
355 "[rerank] prior in-flight inference failed; continuing: %s",
356 exc,
357 exc_info=True,
358 )
359 # Loop again: either inflight cleared or will be handled again.
360 continue
362 model = _RERANKER
363 executor = _EXECUTOR
364 if model is None or executor is None: 364 ↛ 365line 364 didn't jump to line 365, because the condition on line 364 was never true
365 _LOG.error("[rerank] model/executor unexpectedly missing -> []")
366 return []
368 def _predict_once() -> List[float]:
369 t0 = time.perf_counter()
370 scores = model.predict(
371 pairs,
372 batch_size=_BATCH_SIZE,
373 convert_to_numpy=True,
374 show_progress_bar=False,
375 )
376 dt = time.perf_counter() - t0
377 _LOG.info(
378 "[_predict_once] model=%s items=%d batch=%d took=%.3fs",
379 _RERANKER_MODEL_NAME,
380 len(pairs),
381 _BATCH_SIZE,
382 dt,
383 )
384 try:
385 return scores.tolist() # type: ignore[attr-defined]
386 except Exception:
387 return list(scores)
389 fut: Optional[asyncio.Future[List[float]]] = None
390 async with _INIT_LOCK: 390 ↛ 318line 390 didn't jump to line 318, because the continue on line 394 wasn't executed
391 if _INFLIGHT is not None and not _INFLIGHT.done(): 391 ↛ 393line 391 didn't jump to line 393, because the condition on line 391 was never true
392 # A concurrent caller started work between checks; retry loop to await it.
393 _LOG.info("[rerank] inference became in-flight concurrently; retry")
394 continue
396 fut = loop.run_in_executor(executor, _predict_once)
397 _INFLIGHT = fut
398 fut.add_done_callback(_clear_inflight_cb)
400 try:
401 assert fut is not None
402 scores = await asyncio.wait_for(
403 asyncio.shield(fut), timeout=_PREDICT_TIMEOUT_SEC
404 )
406 ranked = sorted(
407 zip(valid_nodes, scores),
408 key=lambda x: x[1],
409 reverse=True,
410 )
411 if not ranked: 411 ↛ 412line 411 didn't jump to line 412, because the condition on line 411 was never true
412 _LOG.warning("[rerank] empty result after scoring -> []")
413 return []
415 # Persist rerank score into NodeWithScore.score for downstream consumers.
416 for n, s in ranked:
417 try:
418 n.score = float(s)
419 except Exception:
420 pass
422 top_k = _RERANK_OUTPUT_K
423 out = [n for n, _ in ranked[:top_k]]
424 _LOG.debug(
425 "[rerank] top%d node_ids=%s",
426 top_k,
427 [_node_id(n) for n in out],
428 )
429 return out
431 except asyncio.TimeoutError as exc:
432 last_err = exc
433 _LOG.warning(
434 "[rerank] inference wait timeout after %.1fs (attempt %d/%d)",
435 _PREDICT_TIMEOUT_SEC,
436 attempt,
437 _PREDICT_RETRIES + 1,
438 )
439 if attempt > _PREDICT_RETRIES:
440 return []
441 await asyncio.sleep(min(0.25 * attempt, 1.0))
443 except Exception as exc:
444 last_err = exc
445 _LOG.exception(
446 "[rerank] inference error on attempt %d/%d: %s",
447 attempt,
448 _PREDICT_RETRIES + 1,
449 exc,
450 )
451 if attempt > _PREDICT_RETRIES:
452 return []
453 await asyncio.sleep(min(0.25 * attempt, 1.0))
455 _LOG.error(
456 "[rerank] giving up after %d attempts -> [] (%s)",
457 _PREDICT_RETRIES + 1,
458 last_err,
459 )
460 return []