Files
AI-Profile-Router/platform/docker/profile-controller/profile_controller.py
T

563 lines
23 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()
SEPARATOR_LABEL_KEY = "com.mike-ai.stem-separator"
SEPARATOR_WORKER = os.environ.get("SEPARATOR_WORKER", "").strip()
VOICE_LABEL_KEY = "com.mike-ai.voice-worker"
VOICE_WORKER = os.environ.get("VOICE_WORKER", "").strip()
VOICE_CHANGE_LABEL_KEY = "com.mike-ai.voice-change-worker"
VOICE_CHANGE_WORKER = os.environ.get("VOICE_CHANGE_WORKER", "").strip()
APPLIO_LABEL_KEY = "com.mike-ai.applio-worker"
APPLIO_WORKER = os.environ.get("APPLIO_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 separator_container() -> dict:
if not SEPARATOR_WORKER:
raise RuntimeError("stem separator is not configured")
matches = [item for item in labelled_containers(SEPARATOR_LABEL_KEY)
if item.get("Labels", {}).get(SEPARATOR_LABEL_KEY) == SEPARATOR_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one stem separator {SEPARATOR_WORKER!r}, found {len(matches)}")
return matches[0]
def voice_container() -> dict:
if not VOICE_WORKER:
raise RuntimeError("voice worker is not configured")
matches = [item for item in labelled_containers(VOICE_LABEL_KEY)
if item.get("Labels", {}).get(VOICE_LABEL_KEY) == VOICE_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one voice worker {VOICE_WORKER!r}, found {len(matches)}")
return matches[0]
def voice_change_container() -> dict:
if not VOICE_CHANGE_WORKER:
raise RuntimeError("voice-change worker is not configured")
matches = [item for item in labelled_containers(VOICE_CHANGE_LABEL_KEY)
if item.get("Labels", {}).get(VOICE_CHANGE_LABEL_KEY) == VOICE_CHANGE_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one voice-change worker {VOICE_CHANGE_WORKER!r}, found {len(matches)}")
return matches[0]
def applio_container() -> dict:
if not APPLIO_WORKER:
raise RuntimeError("Applio worker is not configured")
matches = [item for item in labelled_containers(APPLIO_LABEL_KEY)
if item.get("Labels", {}).get(APPLIO_LABEL_KEY) == APPLIO_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one Applio worker {APPLIO_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_separator_if_configured() -> None:
if SEPARATOR_WORKER:
stop_container(separator_container(), timeout=30)
def stop_voice_if_configured() -> None:
if VOICE_WORKER:
stop_container(voice_container(), timeout=30)
def stop_voice_change_if_configured() -> None:
if VOICE_CHANGE_WORKER:
stop_container(voice_change_container(), timeout=30)
def stop_applio_if_configured() -> None:
if APPLIO_WORKER:
stop_container(applio_container(), timeout=30)
def stop_voice_tools(except_kind: str | None = None) -> None:
if except_kind != "voice":
stop_voice_if_configured()
if except_kind != "voicechange":
stop_voice_change_if_configured()
if except_kind != "applio":
stop_applio_if_configured()
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. TTS is unavailable during this exclusive GPU phase.
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_separator_if_configured()
stop_voice_tools()
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)
stop_separator_if_configured()
stop_voice_tools()
start_container(item)
else:
stop_container(item, timeout=30)
return {"music_worker": MUSIC_WORKER,
"state": "running" if running else "stopped"}
def set_separator_worker(running: bool) -> dict:
"""Start vocal separation exclusively, or stop it before LLM restoration."""
with LOCK:
item = separator_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)
stop_music_if_configured()
stop_voice_tools()
start_container(item)
else:
stop_container(item, timeout=30)
return {"separator_worker": SEPARATOR_WORKER,
"state": "running" if running else "stopped"}
def set_voice_worker(running: bool) -> dict:
"""Start OmniVoice exclusively, or stop it before LLM restoration."""
with LOCK:
item = voice_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)
stop_music_if_configured()
stop_separator_if_configured()
stop_voice_tools("voice")
start_container(item)
else:
stop_container(item, timeout=30)
return {"voice_worker": VOICE_WORKER,
"state": "running" if running else "stopped"}
def set_voice_change_worker(running: bool) -> dict:
"""Start X-VC exclusively, or stop it before another mode is loaded."""
with LOCK:
item = voice_change_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)
stop_music_if_configured()
stop_separator_if_configured()
stop_voice_tools("voicechange")
start_container(item)
else:
stop_container(item, timeout=30)
return {"voice_change_worker": VOICE_CHANGE_WORKER,
"state": "running" if running else "stopped"}
def set_applio_worker(running: bool) -> dict:
"""Start Applio exclusively, or stop it before another mode is loaded."""
with LOCK:
item = applio_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)
stop_music_if_configured()
stop_separator_if_configured()
stop_voice_tools("applio")
start_container(item)
else:
stop_container(item, timeout=30)
return {"applio_worker": APPLIO_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()
stop_separator_if_configured()
stop_voice_tools()
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")
separator = separator_container() if SEPARATOR_WORKER else {}
separator_status = separator.get("Status", "")
separator_health = ("disabled" if not SEPARATOR_WORKER else
"healthy" if "(healthy)" in separator_status else
"unhealthy" if "(unhealthy)" in separator_status else
"starting" if separator.get("State") == "running" else
"stopped")
voice = voice_container() if VOICE_WORKER else {}
voice_status = voice.get("Status", "")
voice_health = ("disabled" if not VOICE_WORKER else
"healthy" if "(healthy)" in voice_status else
"unhealthy" if "(unhealthy)" in voice_status else
"starting" if voice.get("State") == "running" else
"stopped")
voice_change = voice_change_container() if VOICE_CHANGE_WORKER else {}
voice_change_status = voice_change.get("Status", "")
voice_change_health = ("disabled" if not VOICE_CHANGE_WORKER else
"healthy" if "(healthy)" in voice_change_status else
"unhealthy" if "(unhealthy)" in voice_change_status else
"starting" if voice_change.get("State") == "running" else
"stopped")
applio = applio_container() if APPLIO_WORKER else {}
applio_status = applio.get("Status", "")
applio_health = ("disabled" if not APPLIO_WORKER else
"healthy" if "(healthy)" in applio_status else
"unhealthy" if "(unhealthy)" in applio_status else
"starting" if applio.get("State") == "running" else
"stopped")
self.reply(200, {"active_profile": active_profile(items),
"music_worker": music.get("State", "disabled"),
"music_health": music_health,
"separator_worker": separator.get("State", "disabled"),
"separator_health": separator_health,
"voice_worker": voice.get("State", "disabled"),
"voice_health": voice_health,
"voice_change_worker": voice_change.get("State", "disabled"),
"voice_change_health": voice_change_health,
"applio_worker": applio.get("State", "disabled"),
"applio_health": applio_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
if self.path in {"/workers/separator/start", "/workers/separator/stop"}:
try:
self.reply(200, set_separator_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("stem separator transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/voice/start", "/workers/voice/stop"}:
try:
self.reply(200, set_voice_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("voice worker transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/voice-change/start", "/workers/voice-change/stop"}:
try:
self.reply(200, set_voice_change_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("voice-change worker transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/applio/start", "/workers/applio/stop"}:
try:
self.reply(200, set_applio_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("Applio 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()