74 lines
2.9 KiB
Python
Executable File
74 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# llama-queue-proxy: 排队 + 超时快速失败 + 健康透传（方案 A）
|
||
# 上游: llama-server 127.0.0.1:8080（双槽）；本代理监听 :8081
|
||
# 行为: 上游并行槽位已满时排队，排队超过 QUEUE_TIMEOUT 秒 → 503（调用方退避重试）
|
||
import os, threading, urllib.request, urllib.error, http.server
|
||
|
||
UPSTREAM = os.environ.get("UPSTREAM_URL", "http://127.0.0.1:8080")
|
||
PORT = int(os.environ.get("PROXY_PORT", "8081"))
|
||
SLOTS = int(os.environ.get("SLOTS", "2"))
|
||
QUEUE_TIMEOUT = float(os.environ.get("QUEUE_TIMEOUT", "15"))
|
||
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "180"))
|
||
|
||
sem = threading.Semaphore(SLOTS)
|
||
|
||
class H(http.server.BaseHTTPRequestHandler):
|
||
protocol_version = "HTTP/1.1"
|
||
def log_message(self, *a):
|
||
pass
|
||
|
||
def _send(self, code, body, ct="application/json"):
|
||
b = body.encode() if isinstance(body, str) else body
|
||
self.send_response(code)
|
||
self.send_header("Content-Type", ct)
|
||
self.send_header("Content-Length", str(len(b)))
|
||
self.end_headers()
|
||
try:
|
||
self.wfile.write(b)
|
||
except Exception:
|
||
pass
|
||
|
||
def _forward(self, body, headers):
|
||
req = urllib.request.Request(UPSTREAM + self.path, data=body,
|
||
method=self.command, headers=headers)
|
||
try:
|
||
up = urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT)
|
||
except urllib.error.HTTPError as e:
|
||
up = e
|
||
except Exception as e:
|
||
self._send(502, '{"error":{"message":"upstream unreachable: %s"}}' % str(e)[:120])
|
||
return
|
||
data = up.read()
|
||
ct = up.headers.get("Content-Type", "application/json")
|
||
self.send_response(up.getcode())
|
||
self.send_header("Content-Type", ct)
|
||
self.send_header("Content-Length", str(len(data)))
|
||
self.end_headers()
|
||
try:
|
||
self.wfile.write(data)
|
||
except Exception:
|
||
pass
|
||
|
||
def do_POST(self):
|
||
ln = int(self.headers.get("Content-Length", 0) or 0)
|
||
body = self.rfile.read(ln) if ln else b""
|
||
if not sem.acquire(timeout=QUEUE_TIMEOUT):
|
||
self._send(503, '{"error":{"message":"llama busy: queue timeout after %ds, retry later"}}' % int(QUEUE_TIMEOUT))
|
||
return
|
||
try:
|
||
hdrs = {k: v for k, v in self.headers.items()
|
||
if k.lower() not in ("content-length", "host", "connection", "accept-encoding")}
|
||
hdrs["Content-Length"] = str(len(body))
|
||
self._forward(body, hdrs)
|
||
finally:
|
||
sem.release()
|
||
|
||
def do_GET(self):
|
||
self._forward(None, {})
|
||
|
||
if __name__ == "__main__":
|
||
srv = http.server.ThreadingHTTPServer(("127.0.0.1", PORT), H)
|
||
print("[llama-queue-proxy] listen :%d -> %s slots=%d queue_timeout=%ds request_timeout=%ds"
|
||
% (PORT, UPSTREAM, SLOTS, QUEUE_TIMEOUT, REQUEST_TIMEOUT), flush=True)
|
||
srv.serve_forever()
|