Mikei386
fdaaef98cc
- prompt, seed, width, height, size, steps, guidance, quality, seconds, model, created - /images zeigt die Metadaten (falls vorhanden)
1036 lines
41 KiB
Python
Executable File
1036 lines
41 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""AI Profile Router – OpenAI-kompatibler Proxy vor llama.cpp.
|
||
|
||
Leitet OpenAI-kompatible Requests transparent an den lokalen llama.cpp-Server
|
||
weiter (Streaming, Tool Calls, JSON) und schaltet zwischen drei festen
|
||
Profilen um:
|
||
|
||
Profil Kontext
|
||
------ --------
|
||
fast 73728
|
||
medium 94208
|
||
long 131072
|
||
|
||
Virtuelle Modelle: qwen-fast, qwen-medium, qwen-long
|
||
Kommandos: POST /fast, /medium, /long (Profilwechsel)
|
||
GET /status (Zustand)
|
||
|
||
Bildgenerierung (FLUX.2 [klein] 4B Base):
|
||
POST /v1/images/generations (OpenAI-kompatibel)
|
||
GET /images (Liste)
|
||
GET /images/<datei> (PNG-Download)
|
||
|
||
Der Router agiert als Modell-Orchestrator: vor der Generierung wird
|
||
llama.cpp gestoppt, der Bild-Worker lädt FLUX, generiert und entlädt
|
||
das Modell wieder; danach wird das vorherige Qwen-Profil wiederher-
|
||
gestellt und erst dann geantwortet (try/finally – Qwen wird auch bei
|
||
Fehlgeschlagener Generierung wiederhergestellt).
|
||
|
||
Nur Python-Standardbibliothek. Logging nach stdout (journald).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import queue
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
import http.client
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Konfiguration (über Umgebungsvariablen, vgl. systemd-Unit)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
HOST = os.environ.get("ROUTER_HOST", "0.0.0.0")
|
||
PORT = int(os.environ.get("ROUTER_PORT", "8081"))
|
||
UPSTREAM_URL = os.environ.get("UPSTREAM_URL", "http://127.0.0.1:8080").rstrip("/")
|
||
PROFILE_SCRIPT = os.environ.get("PROFILE_SCRIPT", "/usr/local/bin/llama-profile")
|
||
PROFILE_DIR = os.environ.get(
|
||
"PROFILE_DIR", "/etc/systemd/system/mike-ai-llama-ui.service.d")
|
||
|
||
SWITCH_TIMEOUT = float(os.environ.get("SWITCH_TIMEOUT", "600")) # s, Warten auf llama.cpp
|
||
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "600")) # s, Read-Timeout Upstream
|
||
CONNECT_TIMEOUT = float(os.environ.get("CONNECT_TIMEOUT", "10")) # s, Connect-Timeout
|
||
POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "2")) # s, Polling-Intervall
|
||
|
||
# --- Bildgenerierung (FLUX.2 [klein] 4B Base) ---
|
||
LLAMA_SERVICE = os.environ.get("LLAMA_SERVICE", "mike-ai-llama-ui.service")
|
||
SYSTEMCTL_BIN = os.environ.get("SYSTEMCTL_BIN", "systemctl")
|
||
IMAGE_WORKER = os.environ.get(
|
||
"IMAGE_WORKER", "/opt/mike-ai/ai-profile-router/image_worker.py")
|
||
IMAGE_PYTHON = os.environ.get(
|
||
"IMAGE_PYTHON", "/opt/mike-ai/ai-profile-router/venv/bin/python")
|
||
IMAGE_DIR = os.environ.get(
|
||
"IMAGE_DIR", "/opt/mike-ai/ai-profile-router/images")
|
||
IMAGE_WORKER_LOG = os.environ.get(
|
||
"IMAGE_WORKER_LOG", "/opt/mike-ai/ai-profile-router/image_worker.log")
|
||
IMAGE_START_TIMEOUT = float(os.environ.get("IMAGE_START_TIMEOUT", "120")) # s, Worker-Start
|
||
IMAGE_GEN_TIMEOUT = float(os.environ.get("IMAGE_GEN_TIMEOUT", "1800")) # s, pro Bild
|
||
IMAGE_VRAM_FREE_TIMEOUT = float(os.environ.get("IMAGE_VRAM_FREE_TIMEOUT", "90")) # s, VRAM-Abgabe
|
||
|
||
# Erlaubte Auflösungen (Breite x Höhe). FLUX.2 klein ist für 1 MP
|
||
# ausgelegt; 1920x1088 (≈2 MP) wird zusätzlich unterstützt.
|
||
IMAGE_SIZES = {
|
||
"1024x1024": (1024, 1024),
|
||
"1536x1024": (1536, 1024),
|
||
"1024x1536": (1024, 1536),
|
||
"1920x1088": (1920, 1088),
|
||
"1088x1920": (1088, 1920),
|
||
}
|
||
# Qualitätsstufen → Inference-Schritte (guidance bleibt offiziell 4.0).
|
||
# Auf der RTX 5080 gemessen: 30 vs. 50 Steps liefern praktisch dieselbe
|
||
# Qualität (1024x1024: 31,3 s vs. 45,3 s). Default ist daher "standard".
|
||
IMAGE_QUALITY = {"standard": 30, "high": 50}
|
||
IMAGE_DEFAULT_QUALITY = "standard"
|
||
IMAGE_MAX_N = 4
|
||
|
||
# Chat-Waiting: Während eines Image-Jobs oder Profilwechsels ist Qwen
|
||
# down. Chat-Requests warten (statt 502) bis Qwen wieder bereit ist.
|
||
CHAT_WAIT_TIMEOUT = float(os.environ.get("CHAT_WAIT_TIMEOUT", "300")) # s, max. Warten
|
||
CHAT_DRAIN_TIMEOUT = float(os.environ.get("CHAT_DRAIN_TIMEOUT", "60")) # s, max. Warten auf aktive Chats
|
||
|
||
PROFILES = {"fast": 73728, "medium": 94208, "long": 131072}
|
||
VIRTUAL_MODELS = {f"qwen-{name}": name for name in PROFILES}
|
||
|
||
log = logging.getLogger("ai-profile-router")
|
||
|
||
# Hop-by-hop-Header, die nicht an Upstream/Client weitergereicht werden.
|
||
HOP_BY_HOP = {
|
||
"host", "connection", "keep-alive", "proxy-authenticate",
|
||
"proxy-authorization", "te", "trailer", "transfer-encoding",
|
||
"upgrade", "content-length",
|
||
}
|
||
|
||
|
||
def _parse_upstream(url: str) -> tuple[str, int]:
|
||
"""'http://127.0.0.1:8080' -> ('127.0.0.1', 8080)"""
|
||
hostport = url.split("://", 1)[-1]
|
||
host, _, port = hostport.partition(":")
|
||
return host, int(port) if port else 80
|
||
|
||
|
||
UPSTREAM_HOST, UPSTREAM_PORT = _parse_upstream(UPSTREAM_URL)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Zustand
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _ImageState:
|
||
"""Zustand der Bildgenerierung (nur für Status-Reporting)."""
|
||
def __init__(self) -> None:
|
||
self.phase = "idle" # siehe PHASES unten
|
||
self.worker: "_Worker | None" = None
|
||
self.last_error: str | None = None
|
||
self.last_image: str | None = None
|
||
self.last_seconds: float | None = None
|
||
|
||
|
||
IMAGE_PHASES = (
|
||
"idle", "stopping-qwen", "loading-image", "generating",
|
||
"unloading-image", "restoring-qwen",
|
||
)
|
||
|
||
|
||
class _State:
|
||
"""Gemeinsamer, thread-sicherer Zustand.
|
||
|
||
lock : zentraler GPU-/Model-Lock. Wird von Profilwechsel UND
|
||
Image-Generation gehalten → gegenseitiger Ausschluss,
|
||
kein Race zwischen beiden.
|
||
avail_lock : schützt qwen_unavailable + active_chats (Chat-Waiting).
|
||
"""
|
||
lock = threading.Lock() # GPU-/Model-Lock (Profilwechsel + Image)
|
||
switching: str | None = None # Profil, das gerade gewechselt wird
|
||
started = time.time()
|
||
image = _ImageState()
|
||
# Qwen-Verfügbarkeit für das Chat-Waiting:
|
||
qwen_unavailable = False # True, wenn Qwen down/neu geladen wird
|
||
active_chats = 0 # Anzahl laufender Chat-Requests
|
||
avail_lock = threading.Lock() # schützt die beiden Felder oben
|
||
|
||
|
||
STATE = _State()
|
||
|
||
|
||
def _wait_chats_drained(timeout: float | None = None) -> None:
|
||
"""Wartet, bis keine aktiven Chat-Requests mehr laufen.
|
||
|
||
Wird von Profilwechsel/Image-Job aufgerufen, BEVOR Qwen gestoppt wird.
|
||
Verhindert, dass ein laufender Chat auf ein gestopptes Qwen trifft (502).
|
||
"""
|
||
timeout = CHAT_DRAIN_TIMEOUT if timeout is None else timeout
|
||
deadline = time.monotonic() + timeout
|
||
while True:
|
||
with STATE.avail_lock:
|
||
if STATE.active_chats == 0:
|
||
return
|
||
n = STATE.active_chats
|
||
if time.monotonic() > deadline:
|
||
log.warning("Chat-Drain-Timeout nach %.0f s (%d aktive Chats) – "
|
||
"fahre trotzdem fort", timeout, n)
|
||
return
|
||
time.sleep(0.5)
|
||
|
||
|
||
def _set_qwen_unavailable(unavailable: bool) -> None:
|
||
with STATE.avail_lock:
|
||
STATE.qwen_unavailable = unavailable
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Upstream (llama.cpp)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def upstream_status() -> dict:
|
||
"""Prüft llama.cpp: erreichbar? welches Modell? welcher Kontext?"""
|
||
try:
|
||
conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT,
|
||
timeout=CONNECT_TIMEOUT)
|
||
conn.request("GET", "/v1/models")
|
||
resp = conn.getresponse()
|
||
data = json.loads(resp.read())
|
||
conn.close()
|
||
except (OSError, ValueError) as e:
|
||
return {"reachable": False, "error": str(e)}
|
||
models = data.get("data") or []
|
||
if not models:
|
||
return {"reachable": True, "model": None, "ctx": None}
|
||
m = models[0]
|
||
return {"reachable": True,
|
||
"model": m.get("id"),
|
||
"ctx": (m.get("meta") or {}).get("n_ctx")}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Profile
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _read(path: str) -> str:
|
||
with open(path, encoding="utf-8") as f:
|
||
return f.read().strip()
|
||
|
||
|
||
def current_profile() -> str | None:
|
||
"""Aktives Profil, ermittelt durch Vergleich der override.conf."""
|
||
try:
|
||
override = _read(os.path.join(PROFILE_DIR, "override.conf"))
|
||
except OSError:
|
||
return None
|
||
for name in PROFILES:
|
||
try:
|
||
ref = _read(os.path.join(PROFILE_DIR, f"profile-{name}.conf.disabled"))
|
||
except OSError:
|
||
continue
|
||
if override == ref:
|
||
return name
|
||
return None
|
||
|
||
|
||
def _wait_ready(profile: str, deadline: float) -> None:
|
||
"""Wartet, bis llama.cpp das Profil geladen hat (Modell + ctx)."""
|
||
expected_ctx = PROFILES[profile]
|
||
while True:
|
||
status = upstream_status()
|
||
if (status["reachable"] and status.get("model")
|
||
and status.get("ctx") == expected_ctx):
|
||
log.info("llama.cpp bereit: Profil=%s Modell=%s ctx=%s",
|
||
profile, status.get("model"), status.get("ctx"))
|
||
return
|
||
if time.monotonic() > deadline:
|
||
raise RuntimeError(
|
||
f"llama.cpp nach {SWITCH_TIMEOUT:.0f} s nicht bereit "
|
||
f"(erwartet ctx {expected_ctx}, aktuell: {status.get('ctx')})")
|
||
time.sleep(POLL_INTERVAL)
|
||
|
||
|
||
def switch_profile(profile: str, implicit: bool = False) -> None:
|
||
"""Stellt sicher, dass das Profil aktiv ist, und wartet bis es geladen ist.
|
||
|
||
Wirft RuntimeError, wenn das Profil nicht aktiviert werden konnte.
|
||
|
||
implicit=True (ausgelöst durch ein virtuelles Modell in einem Chat-Request):
|
||
Wenn das Profil bereits aktiv ist, aber llama.cpp down ist, wird sofort
|
||
eine RuntimeError geworfen (kein stiller Neustart). Der Nutzer kann den
|
||
Neustart explizit über /<profil> anstoßen.
|
||
"""
|
||
if profile not in PROFILES:
|
||
raise ValueError(f"unbekanntes Profil: {profile!r} "
|
||
f"(erlaubt: {', '.join(PROFILES)})")
|
||
# Kein Fast-Fail: Wenn ein Image-Job läuft (hält den GPU-Lock), wartet
|
||
# der Profilwechsel auf den GPU-Lock (blockiert), bis der Image-Job
|
||
# fertig ist. So bekommen Chat-Requests kein 502, sondern warten.
|
||
with STATE.lock:
|
||
STATE.switching = profile
|
||
try:
|
||
cur = current_profile()
|
||
up = upstream_status()
|
||
ready = (up["reachable"] and up.get("model")
|
||
and up.get("ctx") == PROFILES[profile])
|
||
if cur == profile and ready:
|
||
log.info("Profil %s ist bereits aktiv", profile)
|
||
return
|
||
# Qwen wird neu geladen/gewechselt → für Chats nicht verfügbar.
|
||
_set_qwen_unavailable(True)
|
||
try:
|
||
_wait_chats_drained()
|
||
if cur == profile and up["reachable"] and not ready:
|
||
# Modell wird gerade geladen (z.B. nach einem Wechsel)
|
||
log.info("Warte, bis Profil %s geladen ist ...", profile)
|
||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||
return
|
||
if cur == profile and not up["reachable"] and implicit:
|
||
raise RuntimeError(
|
||
f"llama.cpp nicht erreichbar (Profil {profile} ist "
|
||
f"bereits aktiv; Neustart über /{profile})")
|
||
log.info("Profilwechsel: %s -> %s", cur, profile)
|
||
try:
|
||
proc = subprocess.run(
|
||
[PROFILE_SCRIPT, profile],
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
timeout=120,
|
||
)
|
||
out = proc.stdout.decode(errors="replace").strip()
|
||
if out:
|
||
log.info("llama-profile: %s", out[-500:])
|
||
if proc.returncode != 0:
|
||
# whiptail bricht das Skript ohne TTY ab – der Wechsel
|
||
# selbst (cp + systemctl restart) ist dann erledigt.
|
||
log.warning("llama-profile Exit-Code %d (ohne TTY "
|
||
"erwartet)", proc.returncode)
|
||
except subprocess.TimeoutExpired:
|
||
log.error("llama-profile hat 120 s überschritten")
|
||
if current_profile() != profile:
|
||
raise RuntimeError(
|
||
f"Profildatei wurde nicht gesetzt (erwartet: {profile})")
|
||
log.info("Warte, bis llama.cpp das Profil geladen hat ...")
|
||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||
finally:
|
||
_set_qwen_unavailable(False)
|
||
finally:
|
||
STATE.switching = None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Bildgenerierung (FLUX.2 [klein] 4B Base)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _Worker:
|
||
"""Verwaltet den Bild-Worker-Prozess (stdin/stdout-JSON-Protokoll)."""
|
||
|
||
def __init__(self) -> None:
|
||
self.proc: subprocess.Popen | None = None
|
||
self.model_loaded = False
|
||
self._queue: queue.Queue[dict] = queue.Queue()
|
||
self._reader: threading.Thread | None = None
|
||
|
||
def alive(self) -> bool:
|
||
return self.proc is not None and self.proc.poll() is None
|
||
|
||
def start(self) -> None:
|
||
if self.alive():
|
||
return
|
||
log.info("starte Bild-Worker: %s %s", IMAGE_PYTHON, IMAGE_WORKER)
|
||
logf = open(IMAGE_WORKER_LOG, "ab")
|
||
self.proc = subprocess.Popen(
|
||
[IMAGE_PYTHON, IMAGE_WORKER],
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=logf,
|
||
text=True,
|
||
bufsize=1,
|
||
)
|
||
self._reader = threading.Thread(target=self._read_loop, daemon=True)
|
||
self._reader.start()
|
||
try:
|
||
msg = self._queue.get(timeout=IMAGE_START_TIMEOUT)
|
||
except queue.Empty:
|
||
self.stop()
|
||
raise RuntimeError("Bild-Worker hat nicht gestartet")
|
||
if msg.get("status") != "ready":
|
||
self.stop()
|
||
raise RuntimeError(f"Bild-Worker-Startfehler: {msg}")
|
||
log.info("Bild-Worker bereit")
|
||
|
||
def _read_loop(self) -> None:
|
||
assert self.proc is not None and self.proc.stdout is not None
|
||
for line in self.proc.stdout:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
self._queue.put(json.loads(line))
|
||
except ValueError:
|
||
log.warning("Worker-Zeile (kein JSON): %s", line[:200])
|
||
|
||
def request(self, payload: dict, timeout: float) -> dict:
|
||
if not self.alive():
|
||
raise RuntimeError("Bild-Worker ist nicht aktiv")
|
||
assert self.proc is not None and self.proc.stdin is not None
|
||
self.proc.stdin.write(json.dumps(payload) + "\n")
|
||
self.proc.stdin.flush()
|
||
try:
|
||
return self._queue.get(timeout=timeout)
|
||
except queue.Empty:
|
||
raise RuntimeError(
|
||
f"Bild-Worker hat nach {timeout:.0f} s nicht geantwortet "
|
||
f"(cmd={payload.get('cmd')})")
|
||
|
||
def stop(self) -> None:
|
||
if self.proc is not None and self.proc.poll() is None:
|
||
self.proc.terminate()
|
||
try:
|
||
self.proc.wait(timeout=10)
|
||
except subprocess.TimeoutExpired:
|
||
self.proc.kill()
|
||
self.proc = None
|
||
self.model_loaded = False
|
||
|
||
|
||
def _worker() -> _Worker:
|
||
"""Worker-Instanz liefern (startet bei Bedarf)."""
|
||
img = STATE.image
|
||
if not img.worker or not img.worker.alive():
|
||
if img.worker:
|
||
img.worker.stop()
|
||
img.worker = _Worker()
|
||
img.worker.start()
|
||
return img.worker
|
||
|
||
|
||
def _wait_upstream_down(deadline: float) -> None:
|
||
"""Wartet, bis llama.cpp den Port freigegeben hat (VRAM frei)."""
|
||
while time.monotonic() < deadline:
|
||
if not upstream_status()["reachable"]:
|
||
return
|
||
time.sleep(1)
|
||
raise RuntimeError("llama.cpp gibt Port/VRAM nicht frei")
|
||
|
||
|
||
def _vram_used_mib() -> int | None:
|
||
"""Aktuelle VRAM-Belegung in MiB (via nvidia-smi), None bei Fehler."""
|
||
try:
|
||
out = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=memory.used",
|
||
"--format=csv,noheader,nounits"],
|
||
stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
|
||
stderr=subprocess.DEVNULL, timeout=10,
|
||
).stdout.decode().strip()
|
||
return int(out.splitlines()[0].split()[0])
|
||
except (OSError, ValueError, IndexError):
|
||
return None
|
||
|
||
|
||
def _wait_vram_free(threshold_mib: int = 1000,
|
||
timeout: float | None = None) -> None:
|
||
"""Wartet, bis der VRAM unter threshold_mib fällt (FLUX entladen).
|
||
|
||
Wird nach dem Beenden des Bild-Workers aufgerufen, um sicherzustellen,
|
||
dass der VRAM (inkl. CUDA-Kontext) frei ist, bevor Qwen neu startet.
|
||
Wenn nvidia-smi nicht verfügbar ist (z.B. lokale Tests), wird der
|
||
Check übersprungen.
|
||
"""
|
||
timeout = IMAGE_VRAM_FREE_TIMEOUT if timeout is None else timeout
|
||
deadline = time.monotonic() + timeout
|
||
last = _vram_used_mib()
|
||
if last is None:
|
||
log.info("VRAM-Check übersprungen (nvidia-smi nicht verfügbar)")
|
||
return
|
||
while time.monotonic() < deadline:
|
||
if last <= threshold_mib:
|
||
log.info("VRAM frei: %d MiB", last)
|
||
return
|
||
time.sleep(1)
|
||
last = _vram_used_mib()
|
||
if last is None:
|
||
log.info("VRAM-Check übersprungen (nvidia-smi nicht verfügbar)")
|
||
return
|
||
raise RuntimeError(
|
||
f"VRAM nach {timeout:.0f} s nicht frei (letzte Messung: "
|
||
f"{last} MiB, erwartet <= {threshold_mib} MiB)")
|
||
|
||
|
||
def _restore_qwen(profile: str) -> None:
|
||
"""Startet llama.cpp mit dem gemerkten Profil und wartet auf Readiness."""
|
||
log.info("stelle Qwen-Profil %s wieder her ...", profile)
|
||
try:
|
||
subprocess.run([SYSTEMCTL_BIN, "start", LLAMA_SERVICE],
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||
timeout=120)
|
||
except subprocess.TimeoutExpired:
|
||
log.error("systemctl start hat 120 s überschritten")
|
||
_wait_ready(profile, time.monotonic() + SWITCH_TIMEOUT)
|
||
|
||
|
||
def generate_image(prompt: str, width: int, height: int, steps: int,
|
||
guidance: float, seed: int | None, n: int,
|
||
quality: str = "standard"
|
||
) -> tuple[list[str], str | None]:
|
||
"""Orchestriert die Bildgenerierung inkl. Qwen-Hotswap.
|
||
|
||
Hält den zentralen GPU-Lock (gegenseitiger Ausschluss mit Profilwechsel).
|
||
Ablauf: Qwen stoppen → Worker laden → generieren → Worker beenden
|
||
(VRAM + CUDA-Kontext frei) → Qwen wiederherstellen. Qwen wird auch bei
|
||
Fehlern wiederhergestellt (try/finally).
|
||
"""
|
||
img = STATE.image
|
||
with STATE.lock:
|
||
if img.phase != "idle":
|
||
raise RuntimeError(f"Bildgenerierung läuft ({img.phase})")
|
||
profile = current_profile()
|
||
if profile is None:
|
||
raise RuntimeError("kein aktives Qwen-Profil (override.conf?)")
|
||
os.makedirs(IMAGE_DIR, exist_ok=True)
|
||
results: list[str] = []
|
||
warning: str | None = None
|
||
# Qwen wird gestoppt → für Chats nicht verfügbar (die warten).
|
||
_set_qwen_unavailable(True)
|
||
try:
|
||
_wait_chats_drained()
|
||
|
||
# 1) Qwen stoppen (VRAM freigeben).
|
||
img.phase = "stopping-qwen"
|
||
subprocess.run([SYSTEMCTL_BIN, "stop", LLAMA_SERVICE],
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||
timeout=120)
|
||
_wait_upstream_down(time.monotonic() + 60)
|
||
|
||
# 2) Worker starten (Modell wird beim ersten generate geladen).
|
||
img.phase = "loading-image"
|
||
worker = _worker()
|
||
|
||
# 3) Generieren.
|
||
for i in range(n):
|
||
img.phase = "generating"
|
||
filename = time.strftime("%Y%m%d-%H%M%S") + \
|
||
f"-{os.urandom(2).hex()}.png"
|
||
output = os.path.join(IMAGE_DIR, filename)
|
||
resp = worker.request({
|
||
"cmd": "generate",
|
||
"prompt": prompt,
|
||
"width": width,
|
||
"height": height,
|
||
"steps": steps,
|
||
"guidance": guidance,
|
||
"seed": seed,
|
||
"output": output,
|
||
}, timeout=IMAGE_GEN_TIMEOUT)
|
||
if resp.get("status") != "ok":
|
||
raise RuntimeError(
|
||
resp.get("message", "Bildgenerierung fehlgeschlagen"))
|
||
worker.model_loaded = True
|
||
results.append(filename)
|
||
img.last_image = filename
|
||
img.last_seconds = resp.get("seconds")
|
||
# Metadaten speichern (Sidecar-JSON).
|
||
meta = {
|
||
"prompt": prompt,
|
||
"seed": seed,
|
||
"width": width,
|
||
"height": height,
|
||
"size": f"{width}x{height}",
|
||
"steps": steps,
|
||
"guidance": guidance,
|
||
"quality": quality,
|
||
"seconds": resp.get("seconds"),
|
||
"model": "FLUX.2-klein-base-4B",
|
||
"created": time.strftime("%Y-%m-%dT%H:%M:%S"),
|
||
}
|
||
meta_path = os.path.join(IMAGE_DIR, filename[:-4] + ".json")
|
||
try:
|
||
with open(meta_path, "w", encoding="utf-8") as f:
|
||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||
except OSError as e:
|
||
log.warning("Metadaten-Speicherung fehlgeschlagen: %s", e)
|
||
log.info("Bild %d/%d: %s (%.1f s)", i + 1, n, filename,
|
||
resp.get("seconds", 0))
|
||
|
||
# 4) Worker vollständig beenden (VRAM + CUDA-Kontext freigeben).
|
||
img.phase = "unloading-image"
|
||
worker.stop()
|
||
img.worker = None
|
||
try:
|
||
_wait_vram_free()
|
||
except RuntimeError as e:
|
||
log.warning("VRAM-Check: %s (fahre mit Qwen-Restore fort)", e)
|
||
except Exception as e:
|
||
img.last_error = str(e)
|
||
log.error("Bildgenerierung fehlgeschlagen: %s", e)
|
||
# Worker sicher beenden (falls noch aktiv), VRAM freigeben.
|
||
if img.worker is not None:
|
||
img.worker.stop()
|
||
img.worker = None
|
||
raise
|
||
finally:
|
||
# 5) Qwen immer wiederherstellen.
|
||
img.phase = "restoring-qwen"
|
||
try:
|
||
_restore_qwen(profile)
|
||
_set_qwen_unavailable(False)
|
||
except Exception as e:
|
||
warning = f"Qwen-Wiederherstellung fehlgeschlagen: {e}"
|
||
img.last_error = warning
|
||
log.error(warning)
|
||
# Qwen ist down → qwen_unavailable bleibt True.
|
||
img.phase = "idle"
|
||
return results, warning
|
||
|
||
|
||
def _image_filename_ok(name: str) -> bool:
|
||
return bool(re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*\.png", name))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# HTTP-Handler
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
server_version = "AIProfileRouter/1.0"
|
||
timeout = 60 # Socket-Timeout für Client-Requests (s)
|
||
|
||
# ---------- Routing ----------
|
||
|
||
def do_GET(self):
|
||
self._route()
|
||
|
||
def do_POST(self):
|
||
self._route()
|
||
|
||
def _route(self):
|
||
path = self.path.split("?", 1)[0]
|
||
started = time.monotonic()
|
||
try:
|
||
if path == "/v1/models" and self.command == "GET":
|
||
self._send_json(200, self._models_payload())
|
||
elif path == "/status":
|
||
self._send_json(200, self._status_payload())
|
||
elif path == "/v1/images/generations" and self.command == "POST":
|
||
self._image_generate()
|
||
elif path == "/images" and self.command == "GET":
|
||
self._images_list()
|
||
elif path.startswith("/images/") and self.command == "GET":
|
||
self._image_serve(path[len("/images/"):])
|
||
elif path in ("/fast", "/medium", "/long"):
|
||
self._switch(path[1:])
|
||
elif (self.command == "POST" and path.startswith("/")
|
||
and path.count("/") == 1):
|
||
# Kommandonamensraum: unbekanntes Profil
|
||
self._send_error(400, f"unbekanntes Profil: {path[1:]}",
|
||
"invalid_request_error", "invalid_profile")
|
||
else:
|
||
self._forward()
|
||
except BrokenPipeError:
|
||
log.warning("Client getrennt: %s %s", self.command, path)
|
||
except Exception:
|
||
log.exception("Fehler bei %s %s", self.command, path)
|
||
self._safe_error(500, "interner Router-Fehler")
|
||
finally:
|
||
log.info("%s %s -> %s in %.3f s", self.command, path,
|
||
getattr(self, "_last_code", "-"), time.monotonic() - started)
|
||
|
||
# ---------- Router-eigene Endpunkte ----------
|
||
|
||
@staticmethod
|
||
def _models_payload() -> dict:
|
||
return {
|
||
"object": "list",
|
||
"data": [
|
||
{
|
||
"id": f"qwen-{name}",
|
||
"object": "model",
|
||
"created": 0,
|
||
"owned_by": "ai-profile-router",
|
||
"context_length": ctx,
|
||
"context_window": ctx,
|
||
}
|
||
for name, ctx in PROFILES.items()
|
||
],
|
||
}
|
||
|
||
def _status_payload(self) -> dict:
|
||
up = upstream_status()
|
||
img = STATE.image
|
||
with STATE.avail_lock:
|
||
qwen_unavailable = STATE.qwen_unavailable
|
||
active_chats = STATE.active_chats
|
||
return {
|
||
"router": "ai-profile-router",
|
||
"uptime_seconds": round(time.time() - STATE.started, 1),
|
||
"current_profile": current_profile(),
|
||
"switching": STATE.switching,
|
||
"profiles": PROFILES,
|
||
"upstream": {
|
||
"url": UPSTREAM_URL,
|
||
"reachable": up["reachable"],
|
||
"model": up.get("model"),
|
||
"ctx": up.get("ctx"),
|
||
},
|
||
"qwen": {
|
||
"available": not qwen_unavailable,
|
||
"active_chats": active_chats,
|
||
},
|
||
"image": {
|
||
"phase": img.phase,
|
||
"worker": "running" if (img.worker and img.worker.alive())
|
||
else "stopped",
|
||
"model_loaded": bool(img.worker and img.worker.model_loaded),
|
||
"last_image": img.last_image,
|
||
"last_seconds": img.last_seconds,
|
||
"last_error": img.last_error,
|
||
},
|
||
}
|
||
|
||
# ---------- Bildgenerierung ----------
|
||
|
||
def _image_generate(self) -> None:
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
try:
|
||
data = json.loads(self.rfile.read(length))
|
||
except ValueError:
|
||
self._send_error(400, "ungültiges JSON",
|
||
"invalid_request_error", "invalid_json")
|
||
return
|
||
if not isinstance(data, dict):
|
||
self._send_error(400, "Request muss ein JSON-Objekt sein",
|
||
"invalid_request_error", "invalid_request")
|
||
return
|
||
|
||
prompt = data.get("prompt")
|
||
if not isinstance(prompt, str) or not prompt.strip():
|
||
self._send_error(400, "'prompt' fehlt oder ist leer",
|
||
"invalid_request_error", "missing_prompt")
|
||
return
|
||
if len(prompt) > 8000:
|
||
self._send_error(400, "'prompt' zu lang (max 8000 Zeichen)",
|
||
"invalid_request_error", "prompt_too_long")
|
||
return
|
||
|
||
# Größe
|
||
size = data.get("size", "1024x1024")
|
||
if size not in IMAGE_SIZES:
|
||
self._send_error(
|
||
400, f"ungültige Größe: {size!r} "
|
||
f"(erlaubt: {', '.join(IMAGE_SIZES)})",
|
||
"invalid_request_error", "invalid_size")
|
||
return
|
||
width, height = IMAGE_SIZES[size]
|
||
|
||
# Anzahl
|
||
n = data.get("n", 1)
|
||
if not isinstance(n, int) or isinstance(n, bool) or not 1 <= n <= IMAGE_MAX_N:
|
||
self._send_error(400, f"'n' muss eine Ganzzahl 1..{IMAGE_MAX_N} sein",
|
||
"invalid_request_error", "invalid_n")
|
||
return
|
||
|
||
# Qualität / Schritte / Guidance
|
||
quality = data.get("quality", IMAGE_DEFAULT_QUALITY)
|
||
if quality not in IMAGE_QUALITY:
|
||
self._send_error(400, f"ungültige Qualität: {quality!r} "
|
||
f"(erlaubt: {', '.join(IMAGE_QUALITY)})",
|
||
"invalid_request_error", "invalid_quality")
|
||
return
|
||
steps = data.get("steps", IMAGE_QUALITY[quality])
|
||
if not isinstance(steps, int) or isinstance(steps, bool) or not 4 <= steps <= 150:
|
||
self._send_error(400, "'steps' muss eine Ganzzahl 4..150 sein",
|
||
"invalid_request_error", "invalid_steps")
|
||
return
|
||
guidance = data.get("guidance", 4.0)
|
||
try:
|
||
guidance = float(guidance)
|
||
except (TypeError, ValueError):
|
||
self._send_error(400, "'guidance' muss eine Zahl sein",
|
||
"invalid_request_error", "invalid_guidance")
|
||
return
|
||
if not 1.0 <= guidance <= 10.0:
|
||
self._send_error(400, "'guidance' muss zwischen 1.0 und 10.0 sein",
|
||
"invalid_request_error", "invalid_guidance")
|
||
return
|
||
|
||
seed = data.get("seed")
|
||
if seed is not None:
|
||
try:
|
||
seed = int(seed)
|
||
except (TypeError, ValueError):
|
||
self._send_error(400, "'seed' muss eine Ganzzahl sein",
|
||
"invalid_request_error", "invalid_seed")
|
||
return
|
||
if not 0 <= seed <= 2**32 - 1:
|
||
self._send_error(400, "'seed' muss zwischen 0 und 4294967295 sein",
|
||
"invalid_request_error", "invalid_seed")
|
||
return
|
||
|
||
response_format = data.get("response_format", "url")
|
||
if response_format not in ("url", "b64_json"):
|
||
self._send_error(400, "'response_format' muss 'url' oder 'b64_json' sein",
|
||
"invalid_request_error", "invalid_response_format")
|
||
return
|
||
|
||
# Generierung (blockt mehrere Minuten – eigener Thread-Timeout).
|
||
self.timeout = None
|
||
try:
|
||
results, warning = generate_image(
|
||
prompt.strip(), width, height, steps, guidance, seed, n,
|
||
quality)
|
||
except (ValueError, RuntimeError) as e:
|
||
self._send_error(503, str(e), "server_error", "image_generation_failed")
|
||
return
|
||
|
||
# Antwort bauen
|
||
host = self.headers.get("Host") or f"{HOST}:{PORT}"
|
||
if not host.startswith(("http://", "https://")):
|
||
host = f"http://{host}"
|
||
items = []
|
||
for filename in results:
|
||
path = os.path.join(IMAGE_DIR, filename)
|
||
item: dict = {"url": f"{host}/images/{filename}", "b64_json": None}
|
||
if response_format == "b64_json":
|
||
with open(path, "rb") as f:
|
||
item["b64_json"] = base64.b64encode(f.read()).decode()
|
||
item["url"] = None
|
||
items.append(item)
|
||
payload: dict = {"created": int(time.time()), "data": items}
|
||
if warning:
|
||
payload["router_warning"] = warning
|
||
self._send_json(200, payload)
|
||
|
||
def _images_list(self) -> None:
|
||
if not os.path.isdir(IMAGE_DIR):
|
||
self._send_json(200, {"images": []})
|
||
return
|
||
entries = []
|
||
for name in sorted(os.listdir(IMAGE_DIR), reverse=True):
|
||
if not _image_filename_ok(name):
|
||
continue
|
||
path = os.path.join(IMAGE_DIR, name)
|
||
try:
|
||
st = os.stat(path)
|
||
except OSError:
|
||
continue
|
||
entry = {
|
||
"name": name,
|
||
"url": f"/images/{name}",
|
||
"bytes": st.st_size,
|
||
"modified": int(st.st_mtime),
|
||
}
|
||
# Metadaten laden (Sidecar-JSON, falls vorhanden).
|
||
meta_path = os.path.join(IMAGE_DIR, name[:-4] + ".json")
|
||
if os.path.isfile(meta_path):
|
||
try:
|
||
with open(meta_path, encoding="utf-8") as f:
|
||
entry["meta"] = json.load(f)
|
||
except (OSError, ValueError):
|
||
pass
|
||
entries.append(entry)
|
||
self._send_json(200, {"images": entries[:200]})
|
||
|
||
def _image_serve(self, name: str) -> None:
|
||
if not _image_filename_ok(name):
|
||
self._send_error(400, "ungültiger Dateiname",
|
||
"invalid_request_error", "invalid_filename")
|
||
return
|
||
path = os.path.join(IMAGE_DIR, name)
|
||
if not os.path.isfile(path):
|
||
self._send_error(404, "Bild nicht gefunden",
|
||
"invalid_request_error", "not_found")
|
||
return
|
||
data = open(path, "rb").read()
|
||
self._last_code = 200
|
||
self.send_response(200)
|
||
self.send_header("Content-Type", "image/png")
|
||
self.send_header("Content-Length", str(len(data)))
|
||
self.send_header("Cache-Control", "public, max-age=86400")
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
self.wfile.write(data)
|
||
|
||
def _switch(self, profile: str) -> None:
|
||
if profile not in PROFILES:
|
||
self._send_error(400, f"unbekanntes Profil: {profile}",
|
||
"invalid_request_error", "invalid_profile")
|
||
return
|
||
try:
|
||
switch_profile(profile)
|
||
except (ValueError, RuntimeError) as e:
|
||
self._send_error(503, str(e), "server_error", "profile_switch_failed")
|
||
return
|
||
up = upstream_status()
|
||
self._send_json(200, {
|
||
"status": "ok",
|
||
"profile": profile,
|
||
"context_length": PROFILES[profile],
|
||
"model": up.get("model"),
|
||
})
|
||
|
||
# ---------- Transparentes Forwarding ----------
|
||
|
||
def _forward(self) -> None:
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length > 0 else None
|
||
|
||
# Virtuelles Modell? -> Profil sicherstellen, dann Modell ersetzen.
|
||
if body is not None and self.path.startswith("/v1/"):
|
||
try:
|
||
data = json.loads(body)
|
||
except ValueError:
|
||
data = None
|
||
model = data.get("model") if isinstance(data, dict) else None
|
||
if isinstance(model, str) and model in VIRTUAL_MODELS:
|
||
profile = VIRTUAL_MODELS[model]
|
||
try:
|
||
switch_profile(profile, implicit=True)
|
||
except (ValueError, RuntimeError) as e:
|
||
self._send_error(502, str(e), "server_error",
|
||
"upstream_unavailable")
|
||
return
|
||
up = upstream_status()
|
||
if not up["reachable"] or not up.get("model"):
|
||
self._send_error(502, "llama.cpp nicht erreichbar",
|
||
"server_error", "upstream_unavailable")
|
||
return
|
||
data["model"] = up["model"]
|
||
body = json.dumps(data).encode()
|
||
elif isinstance(model, str) and model.startswith("qwen-"):
|
||
# qwen-* ist der Namensraum des Routers
|
||
self._send_error(400, f"unbekanntes virtuelles Modell: {model}",
|
||
"invalid_request_error", "unknown_model")
|
||
return
|
||
|
||
# An llama.cpp weiterleiten (mit Chat-Waiting, Streaming bleibt erhalten).
|
||
self._proxy_with_wait(body)
|
||
|
||
def _proxy_with_wait(self, body: bytes | None) -> None:
|
||
"""Leitet an llama.cpp weiter, wartet aber erst, bis Qwen verfügbar ist.
|
||
|
||
Während eines Image-Jobs oder Profilwechsels ist Qwen down. Statt
|
||
502 zu liefern, wartet der Request (mit Timeout), bis Qwen wieder
|
||
bereit ist. Mehrere Chats können parallel laufen (active_chats).
|
||
|
||
Race-frei: Der Check auf qwen_unavailable und das Inkrement von
|
||
active_chats sind atomar (avail_lock). Ein Image-Job/Profilwechsel
|
||
setzt qwen_unavailable=True und wartet auf active_chats==0, BEVOR
|
||
er Qwen stoppt – ein laufender Chat wird daher nie unterbrochen.
|
||
"""
|
||
deadline = time.monotonic() + CHAT_WAIT_TIMEOUT
|
||
while True:
|
||
with STATE.avail_lock:
|
||
if not STATE.qwen_unavailable:
|
||
STATE.active_chats += 1
|
||
break
|
||
if time.monotonic() > deadline:
|
||
self._send_error(
|
||
503,
|
||
"Qwen wird neu geladen (Image-Job oder Profilwechsel), "
|
||
"bitte später erneut",
|
||
"server_error", "qwen_reloading")
|
||
return
|
||
time.sleep(0.5)
|
||
try:
|
||
self._proxy(body)
|
||
finally:
|
||
with STATE.avail_lock:
|
||
STATE.active_chats -= 1
|
||
|
||
def _proxy(self, body: bytes | None) -> None:
|
||
# An llama.cpp weiterleiten (Streaming bleibt erhalten).
|
||
try:
|
||
conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT,
|
||
timeout=CONNECT_TIMEOUT)
|
||
conn.connect()
|
||
conn.sock.settimeout(REQUEST_TIMEOUT)
|
||
headers = {k: v for k, v in self.headers.items()
|
||
if k.lower() not in HOP_BY_HOP}
|
||
conn.request(self.command, self.path, body=body, headers=headers)
|
||
resp = conn.getresponse()
|
||
except (OSError, http.client.HTTPException) as e:
|
||
self._send_error(502, f"llama.cpp nicht erreichbar: {e}",
|
||
"server_error", "upstream_unavailable")
|
||
return
|
||
|
||
self._last_code = resp.status
|
||
self.send_response(resp.status)
|
||
for k, v in resp.getheaders():
|
||
if k.lower() not in HOP_BY_HOP:
|
||
self.send_header(k, v)
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
try:
|
||
while True:
|
||
chunk = resp.read(16384)
|
||
if not chunk:
|
||
break
|
||
self.wfile.write(chunk)
|
||
self.wfile.flush()
|
||
except (OSError, http.client.HTTPException) as e:
|
||
log.warning("Upstream-Stream abgebrochen: %s", e)
|
||
finally:
|
||
conn.close()
|
||
|
||
# ---------- Antworten ----------
|
||
|
||
def _send_json(self, code: int, payload: dict) -> None:
|
||
body = json.dumps(payload).encode()
|
||
self._last_code = code
|
||
self.send_response(code)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def _send_error(self, code: int, message: str, etype: str, ecode: str) -> None:
|
||
# OpenAI-kompatibles Fehlerformat
|
||
self._send_json(code, {"error": {"message": message,
|
||
"type": etype,
|
||
"code": ecode}})
|
||
|
||
def _safe_error(self, code: int, message: str) -> None:
|
||
try:
|
||
self._send_error(code, message, "server_error", "internal_error")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _FlushHandler(logging.StreamHandler):
|
||
"""StreamHandler, der nach jedem Record flusht (journald)."""
|
||
|
||
def emit(self, record):
|
||
super().emit(record)
|
||
self.flush()
|
||
|
||
|
||
def main() -> None:
|
||
handler = _FlushHandler(sys.stdout)
|
||
handler.setFormatter(logging.Formatter(
|
||
"%(asctime)s %(levelname)s %(message)s"))
|
||
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"),
|
||
handlers=[handler])
|
||
log.info("AI Profile Router startet: %s:%s -> %s (Profile: %s)",
|
||
HOST, PORT, UPSTREAM_URL, ", ".join(PROFILES))
|
||
server = ThreadingHTTPServer((HOST, PORT), Handler)
|
||
server.daemon_threads = True
|
||
try:
|
||
server.serve_forever()
|
||
except KeyboardInterrupt:
|
||
pass
|
||
finally:
|
||
server.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|