258 lines
10 KiB
Python
258 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
bge 嵌入服务故障切换代理 (2026-09-05)
|
||
======================================
|
||
zhiyid 只认 localhost:8000 — 本代理常驻 8000：
|
||
* 远端优先: 192.168.5.104:8000 (ZSB GPU / DML bge-m3)
|
||
* 远端故障 → 自动切本地 127.0.0.1:8001 (fallback bge)
|
||
* 远端恢复 → 自动切回（每次请求前探测 /health）
|
||
|
||
启动: python3 bge-failover-proxy.py
|
||
端口: 8000 (env PROXY_PORT)
|
||
远端: 192.168.5.104:8000 (env REMOTE_EMBED_URL)
|
||
本地: 127.0.0.1:8001 (env LOCAL_EMBED_URL)
|
||
"""
|
||
import json
|
||
import logging
|
||
import os
|
||
import socket
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from http.server import HTTPServer, BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from urllib import request as urlrequest
|
||
from urllib.error import URLError
|
||
|
||
PROXY_PORT = int(os.environ.get("PROXY_PORT", "8000"))
|
||
REMOTE_EMBED_URL = os.environ.get("REMOTE_EMBED_URL", "http://192.168.5.104:8000")
|
||
LOCAL_EMBED_URL = os.environ.get("LOCAL_EMBED_URL", "http://127.0.0.1:8001")
|
||
HEALTH_PATH = "/health"
|
||
EMBED_PATH = "/v1/embeddings"
|
||
HEALTH_TIMEOUT = float(os.environ.get("HEALTH_TIMEOUT", "1.5"))
|
||
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "60"))
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="[bge-proxy] %(asctime)s %(message)s",
|
||
datefmt="%H:%M:%S",
|
||
)
|
||
log = logging.getLogger("bge-proxy")
|
||
|
||
_state_lock = threading.Lock()
|
||
# 标记文件：proxy 拉起的 local 需在 proxy 重启后仍记得（避免重启后 local 永不自动停）
|
||
LOCAL_STAMP = os.path.expanduser("~/.hermes/run/bge_proxy_local.stamp")
|
||
_state = {"remote_ok": None, "last_check": 0.0, "using": None,
|
||
"local_started_by_proxy": os.path.exists(LOCAL_STAMP), # 重启后从标记恢复
|
||
"remote_ok_since": None,
|
||
"remote_retry_after": 0.0} # embedding 实发失败后的冷却截止（秒，300s）
|
||
|
||
# 按需拉起本地 fallback 的参数（2026-09-05: bge-embed 不再常驻，104 故障时才拉起）
|
||
LOCAL_START_TIMEOUT = float(os.environ.get("LOCAL_START_TIMEOUT", "30")) # 拉起后等待就绪上限
|
||
LOCAL_STOP_GRACE = float(os.environ.get("LOCAL_STOP_GRACE", "60")) # 远端稳定多久后停本地
|
||
BGE_EMBED_SERVICE = os.environ.get("BGE_EMBED_SERVICE", "bge-embed.service")
|
||
|
||
|
||
def _probe_remote() -> bool:
|
||
"""探测远端 /health，1.5s 超时。"""
|
||
try:
|
||
req = urlrequest.Request(REMOTE_EMBED_URL + HEALTH_PATH)
|
||
with urlrequest.urlopen(req, timeout=HEALTH_TIMEOUT) as r:
|
||
if r.status != 200:
|
||
return False
|
||
body = r.read(200).decode("utf-8", "ignore")
|
||
return "ok" in body
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def _local_healthy() -> bool:
|
||
"""探测本地 fallback (8001) /health。"""
|
||
try:
|
||
req = urlrequest.Request(LOCAL_EMBED_URL + HEALTH_PATH)
|
||
with urlrequest.urlopen(req, timeout=1.5) as r:
|
||
if r.status != 200:
|
||
return False
|
||
return "ok" in r.read(200).decode("utf-8", "ignore")
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def _ensure_local() -> bool:
|
||
"""远端不可用时确保本地可用：未启动则 systemctl --user 拉起 bge-embed。
|
||
|
||
返回 True = 本地可用；False = 本地拉起失败/超时（转发将 502）。
|
||
"""
|
||
if _local_healthy():
|
||
return True
|
||
with _state_lock:
|
||
if not _state["local_started_by_proxy"]:
|
||
log.info("远端不可用 → 拉起 %s (按需)", BGE_EMBED_SERVICE)
|
||
try:
|
||
subprocess.Popen(
|
||
["systemctl", "--user", "start", BGE_EMBED_SERVICE],
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||
)
|
||
_state["local_started_by_proxy"] = True
|
||
try:
|
||
os.makedirs(os.path.dirname(LOCAL_STAMP), exist_ok=True)
|
||
open(LOCAL_STAMP, "w").close()
|
||
except Exception:
|
||
pass
|
||
except Exception as e:
|
||
log.error("拉起 %s 失败: %s", BGE_EMBED_SERVICE, e)
|
||
# 轮询等待模型加载（int8 onnx ~10-20s）
|
||
deadline = time.time() + LOCAL_START_TIMEOUT
|
||
while time.time() < deadline:
|
||
if _local_healthy():
|
||
return True
|
||
time.sleep(1)
|
||
log.error("本地 fallback %s 拉起超时 (%ss)", BGE_EMBED_SERVICE, LOCAL_START_TIMEOUT)
|
||
return False
|
||
|
||
|
||
def _maybe_stop_local(remote_ok: bool):
|
||
"""远端已恢复且稳定 → 停掉 proxy 拉起的本地 fallback，回到省电模式。"""
|
||
now = time.time()
|
||
with _state_lock:
|
||
if remote_ok and _state["local_started_by_proxy"]:
|
||
if _state["remote_ok_since"] is None:
|
||
_state["remote_ok_since"] = now
|
||
elif now - _state["remote_ok_since"] >= LOCAL_STOP_GRACE:
|
||
log.info("远端已稳定 %ss → 停 %s (释放内存)", LOCAL_STOP_GRACE, BGE_EMBED_SERVICE)
|
||
try:
|
||
subprocess.Popen(
|
||
["systemctl", "--user", "stop", BGE_EMBED_SERVICE],
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||
)
|
||
_state["local_started_by_proxy"] = False
|
||
_state["remote_ok_since"] = None
|
||
try:
|
||
os.remove(LOCAL_STAMP)
|
||
except OSError:
|
||
pass
|
||
except Exception as e:
|
||
log.error("停止 %s 失败: %s", BGE_EMBED_SERVICE, e)
|
||
elif not remote_ok:
|
||
_state["remote_ok_since"] = None
|
||
|
||
|
||
def choose_target() -> tuple[str, bool]:
|
||
"""返回 (目标URL, 是否远端)。带 3s 冷却 + 状态缓存 + 按需拉起本地 fallback。"""
|
||
global _state
|
||
now = time.time()
|
||
with _state_lock:
|
||
cache_fresh = _state["remote_ok"] is not None and now - _state["last_check"] < 3.0
|
||
in_cooldown = _state["remote_ok"] is False and now < _state["remote_retry_after"]
|
||
if cache_fresh or in_cooldown:
|
||
remote_ok = _state["remote_ok"]
|
||
else:
|
||
remote_ok = _probe_remote()
|
||
_state["remote_ok"] = remote_ok
|
||
_state["last_check"] = now
|
||
_maybe_stop_local(remote_ok)
|
||
if remote_ok:
|
||
using = "remote(104-GPU)"
|
||
if using != _state["using"]:
|
||
log.info("切换 → %s", using)
|
||
with _state_lock:
|
||
_state["using"] = using
|
||
return REMOTE_EMBED_URL, True
|
||
# 远端不可用 → 确保本地可用
|
||
local_ok = _ensure_local()
|
||
using = "local(106-CPU)" if local_ok else "local(106-CPU, 未就绪)"
|
||
if using != _state["using"]:
|
||
log.info("切换 → %s", using)
|
||
with _state_lock:
|
||
_state["using"] = using
|
||
return LOCAL_EMBED_URL, False
|
||
|
||
|
||
def _forward(url: str, method: str, body: bytes, headers: dict) -> tuple[int, dict, bytes]:
|
||
"""把请求转发到目标，返回 (状态码, 响应头, 响应体)。"""
|
||
req = urlrequest.Request(url, data=body if method == "POST" else None, method=method)
|
||
for k, v in headers.items():
|
||
if k.lower() not in ("host", "content-length", "connection", "accept-encoding"):
|
||
req.add_header(k, v)
|
||
try:
|
||
with urlrequest.urlopen(req, timeout=REQUEST_TIMEOUT) as r:
|
||
resp_body = r.read()
|
||
resp_headers = dict(r.headers.items())
|
||
return r.status, resp_headers, resp_body
|
||
except URLError as e:
|
||
return 502, {"Content-Type": "application/json"}, json.dumps(
|
||
{"error": f"proxy forward failed: {e}"}
|
||
).encode()
|
||
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
protocol_version = "HTTP/1.1"
|
||
|
||
def log_message(self, fmt, *args):
|
||
log.info("%s %s" % (self.address_string(), fmt % args))
|
||
|
||
def _handle(self):
|
||
length = int(self.headers.get("Content-Length", 0) or 0)
|
||
body = self.rfile.read(length) if length else b""
|
||
target, _remote_ok = choose_target()
|
||
status, resp_headers, resp_body = _forward(
|
||
target + self.path, self.command, body, dict(self.headers)
|
||
)
|
||
# 远端 embedding 实发失败（500/502/503/504，health 探测不出的故障）
|
||
# → 标故障 + 30s 冷却 + 立即降级 local 重试一次
|
||
if status >= 500 and target.startswith(REMOTE_EMBED_URL):
|
||
with _state_lock:
|
||
_state["remote_ok"] = False
|
||
# 300s 冷却: 104 可能出现 health OK 但 embedding 500（实发故障），
|
||
# 短冷却会导致反复 remote→500→local 抖动并可能反复启停 bge-embed
|
||
_state["remote_retry_after"] = time.time() + 300
|
||
_state["using"] = None # 强制下次 log 切换
|
||
log.info("远端 embedding 实发失败(%s) → 降级 local 重试", status)
|
||
if _ensure_local():
|
||
target = LOCAL_EMBED_URL
|
||
status, resp_headers, resp_body = _forward(
|
||
target + self.path, self.command, body, dict(self.headers)
|
||
)
|
||
self.send_response(status)
|
||
# 只转发安全响应头；Content-Length 必须用实际长度
|
||
for k, v in resp_headers.items():
|
||
if k.lower() in ("content-type",):
|
||
self.send_header(k, v)
|
||
self.send_header("Content-Length", str(len(resp_body)))
|
||
self.send_header("X-BGE-Proxy", "remote-104" if target.startswith(REMOTE_EMBED_URL) else "local-106")
|
||
self.end_headers()
|
||
if self.command != "HEAD":
|
||
try:
|
||
self.wfile.write(resp_body)
|
||
except (BrokenPipeError, ConnectionResetError):
|
||
# 客户端提前断开（探测超时/取消）——不影响其他请求（多线程）
|
||
pass
|
||
|
||
do_GET = _handle
|
||
do_POST = _handle
|
||
do_PUT = _handle
|
||
do_DELETE = _handle
|
||
|
||
|
||
def _startup_probe():
|
||
"""启动时打一条日志说明当前选路。"""
|
||
remote_ok = _probe_remote()
|
||
with _state_lock:
|
||
_state["remote_ok"] = remote_ok
|
||
_state["last_check"] = time.time()
|
||
_state["using"] = "remote(104-GPU)" if remote_ok else "local(106-CPU)"
|
||
log.info("启动选路: %s (remote=%s)", _state["using"], remote_ok)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
log.info("bge 故障切换代理启动 0.0.0.0:%s (Threading)", PROXY_PORT)
|
||
log.info(" 远端: %s", REMOTE_EMBED_URL)
|
||
log.info(" 本地 fallback: %s", LOCAL_EMBED_URL)
|
||
_startup_probe()
|
||
# ThreadingHTTPServer: 单请求卡住不阻塞 /health 探测和其他请求
|
||
server = ThreadingHTTPServer(("0.0.0.0", PROXY_PORT), Handler)
|
||
server.daemon_threads = True
|
||
try:
|
||
server.serve_forever()
|
||
except KeyboardInterrupt:
|
||
log.info("停止")
|