"""Small isolated RAG process pool with bounded admission and wall-clock deadlines.""" from __future__ import annotations import atexit import json import math import os from pathlib import Path import queue import selectors import subprocess import sys import threading import time class RagWorkerError(RuntimeError): pass class RagWorkerBusy(RagWorkerError): pass class RagWorkerTimeout(RagWorkerError): pass def _number(name: str, default: float, minimum: float, maximum: float) -> float: try: value = float(os.environ.get(name, default)) return min(maximum, max(minimum, value)) if math.isfinite(value) else default except (TypeError, ValueError): return default class _Worker: def __init__(self, command: list[str], stderr_path: Path | None): self.command = command self.stderr_path = stderr_path self.process: subprocess.Popen | None = None self.last_start_ms = 0.0 self.last_reused = False def start(self): if self.process is not None and self.process.poll() is None: self.last_start_ms = 0.0 self.last_reused = True return self.process self.close() started = time.perf_counter() stderr = subprocess.DEVNULL if self.stderr_path: self.stderr_path.parent.mkdir(parents=True, exist_ok=True) stderr = self.stderr_path.open('ab') try: self.process = subprocess.Popen( self.command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=stderr, bufsize=0) finally: if hasattr(stderr, 'close'): stderr.close() os.set_blocking(self.process.stdin.fileno(), False) os.set_blocking(self.process.stdout.fileno(), False) self.last_start_ms = round((time.perf_counter() - started) * 1000, 2) self.last_reused = False return self.process def close(self): process, self.process = self.process, None if process is None: return try: if process.poll() is None: process.kill() process.wait(timeout=0.2) except (OSError, subprocess.TimeoutExpired): pass finally: for pipe in (process.stdin, process.stdout): if pipe: pipe.close() def request(self, payload: dict, deadline: float) -> dict: if time.monotonic() >= deadline: raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。') process = self.start() encoded = memoryview((json.dumps(payload, ensure_ascii=False) + '\n').encode('utf-8')) buffer = bytearray() with selectors.DefaultSelector() as selector: selector.register(process.stdin, selectors.EVENT_WRITE) selector.register(process.stdout, selectors.EVENT_READ) while True: remaining = deadline - time.monotonic() if remaining <= 0: raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。') events = selector.select(remaining) if not events: raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。') for key, _ in events: if key.fileobj is process.stdin: try: sent = os.write(key.fd, encoded[:65536]) except BlockingIOError: continue encoded = encoded[sent:] if not encoded: selector.unregister(process.stdin) else: try: part = os.read(key.fd, 65536) except BlockingIOError: continue if not part: raise RagWorkerError('知识库检索进程意外退出，请重试。') buffer.extend(part) if len(buffer) > 32 * 1024 * 1024: raise RagWorkerError('知识库检索返回内容过大，请减少检索片段数。') if b'\n' in buffer: line, trailing = buffer.split(b'\n', 1) if trailing.strip(): raise RagWorkerError('知识库检索进程返回格式异常，请重试。') try: response = json.loads(line) except (ValueError, UnicodeDecodeError) as exc: raise RagWorkerError('知识库检索进程返回格式异常，请重试。') from exc if not isinstance(response, dict): raise RagWorkerError('知识库检索进程返回格式异常，请重试。') return response class RagWorkerPool: """At most three active workers, a small waiting room, no in-process fallback.""" def __init__(self, *, size: int = 2, total_timeout: float = 45, queue_timeout: float = 3, max_pending: int | None = None, command: list[str] | None = None, stderr_path: Path | None = None): self.size = min(3, max(1, int(size))) self.total_timeout = max(0.01, float(total_timeout)) self.queue_timeout = max(0.0, min(float(queue_timeout), self.total_timeout)) self._closed = False self._workers = [_Worker(command or [sys.executable, str(Path(__file__).with_name('rag_worker.py'))], stderr_path) for _ in range(self.size)] self._available = queue.Queue(maxsize=self.size) for worker in self._workers: self._available.put_nowait(worker) self._admission = threading.BoundedSemaphore(self.size + (max_pending if max_pending is not None else self.size * 2)) def request(self, payload: dict) -> dict: started = time.monotonic() deadline = started + self.total_timeout if self._closed: raise RagWorkerError('知识库检索服务正在关闭，请稍后重试。') if not self._admission.acquire(blocking=False): raise RagWorkerBusy('知识库检索任务较多，等待队列已满，请稍后重试。') worker = None try: try: worker = self._available.get(timeout=min(self.queue_timeout, max(0, deadline - time.monotonic()))) except queue.Empty as exc: raise RagWorkerBusy('知识库检索排队超时，请稍后重试。') from exc queue_ms = round((time.monotonic() - started) * 1000, 2) try: response = worker.request(payload, deadline) if response.get('ok') is not True: from ..system_log import redact detail = redact(response.get('error') or '未知错误', 400) error = RagWorkerError(f'知识库检索失败：{detail}') setattr(error, 'retrieval_metric', response.get('retrieval_metric') or {"error": detail}) raise error if payload.get('op') == 'warmup': return response if not isinstance(response.get('results'), list): raise RagWorkerError('知识库检索返回格式异常，请重试。') except Exception as exc: # A timed-out response may never be consumed by a subsequent request. worker.close() if isinstance(exc, RagWorkerError): raise raise RagWorkerError('知识库检索进程通信失败，请重试。') from exc response['transport_timings'] = { 'queue_ms': queue_ms, 'total_ms': round((time.monotonic() - started) * 1000, 2), 'worker_ms': round((time.monotonic() - started) * 1000 - queue_ms, 2), 'process_start_ms': worker.last_start_ms, 'process_reused': worker.last_reused, 'worker_pid': int(worker.process.pid) if worker.process else 0, } # Child search timing is reported in the response envelope. The # residual after spawn and search is the parent/child pipe and JSON # transport overhead, useful for spotting IPC regressions. child_ms = float((response.get('worker_timings') or {}).get('search_ms') or 0.0) response['transport_timings']['ipc_ms'] = round(max( 0.0, response['transport_timings']['worker_ms'] - child_ms - float(response['transport_timings']['process_start_ms'] or 0.0)), 2) return response finally: if worker is not None: self._available.put_nowait(worker) self._admission.release() def prewarm(self) -> dict: """Start one resident worker and warm its model/client in the background.""" return self.request({'op': 'warmup'}) def close(self): self._closed = True for worker in self._workers: worker.close() _pool: RagWorkerPool | None = None _pool_lock = threading.Lock() def get_pool() -> RagWorkerPool: global _pool with _pool_lock: if _pool is None: from ..config import DATA_DIR _pool = RagWorkerPool( # The local embedding model is loaded once per worker. Keep one # resident worker by default on the single-host deployment to # avoid duplicating several hundred MB and triggering macOS # memory pressure; operators can raise RAG_WORKER_POOL_SIZE # explicitly when concurrent retrieval is required. size=int(_number('RAG_WORKER_POOL_SIZE', 1, 1, 3)), total_timeout=_number('RAG_SEARCH_TIMEOUT_SECONDS', 45, 1, 180), # Content tasks commonly fan out several segment requests at # once. Give queued requests enough time to use the single # memory-efficient worker instead of surfacing a spurious # timeout while an earlier segment is still being searched. queue_timeout=_number('RAG_QUEUE_TIMEOUT_SECONDS', 10, 0, 30), max_pending=int(_number('RAG_MAX_PENDING', 8, 0, 32)), stderr_path=DATA_DIR / 'logs' / 'rag-worker.log') atexit.register(_pool.close) return _pool