238 lines
10 KiB
Python
238 lines
10 KiB
Python
"""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
|