333 lines
13 KiB
Python
333 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Strict Docker profile switcher for the local AI stack.
|
|
|
|
Only status and activation of a fixed set of labelled llama.cpp containers are
|
|
exposed. Callers cannot provide images, commands, mounts or Docker API paths.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import http.client
|
|
import json
|
|
import logging
|
|
import os
|
|
import socket
|
|
import threading
|
|
import urllib.parse
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
HOST = os.environ.get("CONTROLLER_HOST", "0.0.0.0")
|
|
PORT = int(os.environ.get("CONTROLLER_PORT", "8090"))
|
|
SOCKET_PATH = os.environ.get("DOCKER_SOCKET", "/var/run/docker.sock")
|
|
TOKEN_FILE = os.environ.get("CONTROLLER_TOKEN_FILE", "/run/secrets/controller-token")
|
|
ALLOWED = tuple(x.strip() for x in os.environ.get(
|
|
"ALLOWED_PROFILES", "fast,medium,large,ultra,uncensored,experimental").split(",") if x.strip())
|
|
LABEL_KEY = "com.mike-ai.llama-profile"
|
|
IMAGE_LABEL_KEY = "com.mike-ai.image-worker"
|
|
IMAGE_WORKER = os.environ.get("IMAGE_WORKER", "image")
|
|
RESTORE_WORKER = os.environ.get("RESTORE_WORKER", "restore")
|
|
TTS_LABEL_KEY = "com.mike-ai.tts-worker"
|
|
TTS_WORKER = os.environ.get("TTS_WORKER", "qwen3")
|
|
MUSIC_LABEL_KEY = "com.mike-ai.music-worker"
|
|
MUSIC_WORKER = os.environ.get("MUSIC_WORKER", "").strip()
|
|
LOCK = threading.Lock()
|
|
log = logging.getLogger("profile-controller")
|
|
|
|
|
|
class UnixConnection(http.client.HTTPConnection):
|
|
def connect(self) -> None:
|
|
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
self.sock.connect(SOCKET_PATH)
|
|
|
|
|
|
def docker_request(method: str, path: str) -> tuple[int, bytes]:
|
|
conn = UnixConnection("localhost", timeout=30)
|
|
try:
|
|
conn.request(method, path, headers={"Content-Type": "application/json"})
|
|
response = conn.getresponse()
|
|
return response.status, response.read()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def containers() -> dict[str, dict]:
|
|
filters = urllib.parse.quote(json.dumps({"label": [LABEL_KEY]}))
|
|
status, body = docker_request("GET", f"/containers/json?all=1&filters={filters}")
|
|
if status != 200:
|
|
raise RuntimeError(f"Docker list failed with HTTP {status}")
|
|
result: dict[str, dict] = {}
|
|
for item in json.loads(body):
|
|
profile = item.get("Labels", {}).get(LABEL_KEY)
|
|
if profile in ALLOWED:
|
|
if profile in result:
|
|
raise RuntimeError(f"duplicate container for profile {profile}")
|
|
result[profile] = item
|
|
return result
|
|
|
|
|
|
def labelled_containers(label: str) -> list[dict]:
|
|
filters = urllib.parse.quote(json.dumps({"label": [label]}))
|
|
status, body = docker_request("GET", f"/containers/json?all=1&filters={filters}")
|
|
if status != 200:
|
|
raise RuntimeError(f"Docker list failed with HTTP {status}")
|
|
return json.loads(body)
|
|
|
|
|
|
def image_container(kind: str = IMAGE_WORKER) -> dict:
|
|
matches = [item for item in labelled_containers(IMAGE_LABEL_KEY)
|
|
if item.get("Labels", {}).get(IMAGE_LABEL_KEY) == kind]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one image worker {kind!r}, found {len(matches)}")
|
|
return matches[0]
|
|
|
|
|
|
def image_containers() -> list[dict]:
|
|
"""All allowlisted GPU workers that must never overlap an LLM."""
|
|
allowed = {IMAGE_WORKER, RESTORE_WORKER}
|
|
return [item for item in labelled_containers(IMAGE_LABEL_KEY)
|
|
if item.get("Labels", {}).get(IMAGE_LABEL_KEY) in allowed]
|
|
|
|
|
|
def tts_container() -> dict:
|
|
matches = [item for item in labelled_containers(TTS_LABEL_KEY)
|
|
if item.get("Labels", {}).get(TTS_LABEL_KEY) == TTS_WORKER]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one TTS worker {TTS_WORKER!r}, found {len(matches)}")
|
|
return matches[0]
|
|
|
|
|
|
def music_container() -> dict:
|
|
if not MUSIC_WORKER:
|
|
raise RuntimeError("music worker is not configured")
|
|
matches = [item for item in labelled_containers(MUSIC_LABEL_KEY)
|
|
if item.get("Labels", {}).get(MUSIC_LABEL_KEY) == MUSIC_WORKER]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one music worker {MUSIC_WORKER!r}, found {len(matches)}")
|
|
return matches[0]
|
|
|
|
|
|
def stop_music_if_configured() -> None:
|
|
if MUSIC_WORKER:
|
|
stop_container(music_container(), timeout=30)
|
|
|
|
|
|
def stop_container(item: dict, timeout: int = 120) -> None:
|
|
if item.get("State") != "running":
|
|
return
|
|
status, _ = docker_request("POST", f"/containers/{item['Id']}/stop?t={timeout}")
|
|
if status not in (204, 304):
|
|
raise RuntimeError(f"failed to stop container: HTTP {status}")
|
|
|
|
|
|
def start_container(item: dict) -> None:
|
|
if item.get("State") == "running":
|
|
return
|
|
status, _ = docker_request("POST", f"/containers/{item['Id']}/start")
|
|
if status not in (204, 304):
|
|
raise RuntimeError(f"failed to start container: HTTP {status}")
|
|
|
|
|
|
def stop_inference() -> dict:
|
|
with LOCK:
|
|
items = containers()
|
|
previous = active_profile(items)
|
|
for item in items.values():
|
|
stop_container(item)
|
|
return {"active_profile": None, "previous_profile": previous}
|
|
|
|
|
|
def set_image_worker(running: bool, kind: str = IMAGE_WORKER) -> dict:
|
|
if kind not in {IMAGE_WORKER, RESTORE_WORKER}:
|
|
raise ValueError("worker is not allowlisted")
|
|
with LOCK:
|
|
item = image_container(kind)
|
|
if running:
|
|
# The image worker may never overlap a llama profile on the 5080.
|
|
for profile_item in containers().values():
|
|
stop_container(profile_item)
|
|
# The 9B beta text encoder temporarily borrows the RTX 3060 from
|
|
# Qwen3-TTS. The gateway retains Piper as a fallback meanwhile.
|
|
stop_container(tts_container(), timeout=30)
|
|
stop_music_if_configured()
|
|
for other in image_containers():
|
|
if other["Id"] != item["Id"]:
|
|
stop_container(other, timeout=20)
|
|
start_container(item)
|
|
else:
|
|
# CUDA/PyTorch may not react promptly to SIGTERM after an OOM.
|
|
# Bound recovery time and let Docker issue SIGKILL afterwards.
|
|
stop_container(item, timeout=20)
|
|
# TTS is restored by the following profile activation. Keeping it
|
|
# stopped here lets the router verify that both GPUs really
|
|
# released the image model before Qwen and TTS are reloaded.
|
|
return {"image_worker": kind,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def set_music_worker(running: bool) -> dict:
|
|
"""Start ACE-Step exclusively, or stop it before LLM restoration."""
|
|
with LOCK:
|
|
item = music_container()
|
|
if running:
|
|
for profile_item in containers().values():
|
|
stop_container(profile_item)
|
|
for worker in image_containers():
|
|
stop_container(worker, timeout=20)
|
|
stop_container(tts_container(), timeout=30)
|
|
start_container(item)
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
return {"music_worker": MUSIC_WORKER,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def active_profile(items: dict[str, dict] | None = None) -> str | None:
|
|
items = items or containers()
|
|
active = [name for name, item in items.items() if item.get("State") == "running"]
|
|
if len(active) > 1:
|
|
raise RuntimeError(f"multiple llama profiles active: {', '.join(active)}")
|
|
return active[0] if active else None
|
|
|
|
|
|
def activate(profile: str) -> dict:
|
|
if profile not in ALLOWED:
|
|
raise ValueError("profile is not allowlisted")
|
|
with LOCK:
|
|
# Defensive mutual exclusion even if a caller bypasses the router.
|
|
for worker in image_containers():
|
|
stop_container(worker)
|
|
stop_music_if_configured()
|
|
start_container(tts_container())
|
|
items = containers()
|
|
missing = [name for name in ALLOWED if name not in items]
|
|
if missing:
|
|
raise RuntimeError("profile containers missing: " + ", ".join(missing))
|
|
current = active_profile(items)
|
|
if current == profile:
|
|
return {"active_profile": current, "changed": False}
|
|
for name, item in items.items():
|
|
if name == profile or item.get("State") != "running":
|
|
continue
|
|
status, _ = docker_request("POST", f"/containers/{item['Id']}/stop?t=120")
|
|
if status not in (204, 304):
|
|
raise RuntimeError(f"failed to stop profile {name}: HTTP {status}")
|
|
target = items[profile]
|
|
status, _ = docker_request("POST", f"/containers/{target['Id']}/start")
|
|
if status not in (204, 304):
|
|
raise RuntimeError(f"failed to start profile {profile}: HTTP {status}")
|
|
log.info("activated profile %s (previous=%s)", profile, current)
|
|
return {"active_profile": profile, "changed": True}
|
|
|
|
|
|
def load_token() -> str:
|
|
token = os.environ.get("CONTROLLER_TOKEN", "").strip()
|
|
if not token:
|
|
with open(TOKEN_FILE, encoding="utf-8") as handle:
|
|
token = handle.read().strip()
|
|
if len(token) < 32:
|
|
raise RuntimeError("controller token is missing or too short")
|
|
return token
|
|
|
|
|
|
TOKEN = load_token()
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
server_version = "mike-ai-profile-controller/1"
|
|
|
|
def log_message(self, fmt: str, *args: object) -> None:
|
|
log.info("%s - %s", self.client_address[0], fmt % args)
|
|
|
|
def reply(self, status: int, payload: dict) -> None:
|
|
body = json.dumps(payload, separators=(",", ":")).encode()
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def authenticated(self) -> bool:
|
|
return self.headers.get("Authorization", "") == f"Bearer {TOKEN}"
|
|
|
|
def do_GET(self) -> None: # noqa: N802
|
|
if self.path == "/health":
|
|
self.reply(200, {"status": "ok"})
|
|
return
|
|
if self.path != "/status":
|
|
self.reply(404, {"error": "not found"})
|
|
return
|
|
if not self.authenticated():
|
|
self.reply(401, {"error": "unauthorized"})
|
|
return
|
|
try:
|
|
items = containers()
|
|
music = music_container() if MUSIC_WORKER else {}
|
|
music_status = music.get("Status", "")
|
|
music_health = ("disabled" if not MUSIC_WORKER else
|
|
"healthy" if "(healthy)" in music_status else
|
|
"unhealthy" if "(unhealthy)" in music_status else
|
|
"starting" if music.get("State") == "running" else
|
|
"stopped")
|
|
self.reply(200, {"active_profile": active_profile(items),
|
|
"music_worker": music.get("State", "disabled"),
|
|
"music_health": music_health,
|
|
"profiles": {name: items.get(name, {}).get(
|
|
"State", "missing") for name in ALLOWED}})
|
|
except Exception as exc:
|
|
log.exception("status failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
if not self.authenticated():
|
|
self.reply(401, {"error": "unauthorized"})
|
|
return
|
|
if self.path == "/inference/stop":
|
|
try:
|
|
self.reply(200, stop_inference())
|
|
except Exception as exc:
|
|
log.exception("stopping inference failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
return
|
|
if self.path in {"/workers/music/start", "/workers/music/stop"}:
|
|
try:
|
|
self.reply(200, set_music_worker(self.path.endswith("/start")))
|
|
except Exception as exc:
|
|
log.exception("music worker transition failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
return
|
|
worker_paths = {
|
|
"/workers/image/start": (IMAGE_WORKER, True),
|
|
"/workers/image/stop": (IMAGE_WORKER, False),
|
|
"/workers/restore/start": (RESTORE_WORKER, True),
|
|
"/workers/restore/stop": (RESTORE_WORKER, False),
|
|
}
|
|
if self.path in worker_paths:
|
|
try:
|
|
kind, running = worker_paths[self.path]
|
|
self.reply(200, set_image_worker(running, kind))
|
|
except Exception as exc:
|
|
log.exception("image worker transition failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
return
|
|
prefix, suffix = "/profiles/", "/activate"
|
|
if not self.path.startswith(prefix) or not self.path.endswith(suffix):
|
|
self.reply(404, {"error": "not found"})
|
|
return
|
|
profile = self.path[len(prefix):-len(suffix)]
|
|
if profile not in ALLOWED:
|
|
self.reply(400, {"error": "profile is not allowlisted"})
|
|
return
|
|
try:
|
|
self.reply(200, activate(profile))
|
|
except Exception as exc:
|
|
log.exception("activation failed for %s", profile)
|
|
self.reply(503, {"error": str(exc)})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()
|