833 lines
33 KiB
Python
833 lines
33 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 time
|
|
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")
|
|
QWEN_IMAGE_TEST_WORKER = os.environ.get(
|
|
"QWEN_IMAGE_TEST_WORKER", "qwen-image-2.1-test").strip()
|
|
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()
|
|
YUE2_WORKER = os.environ.get("YUE2_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()
|
|
TRELLIS_LABEL_KEY = "com.mike-ai.trellis-worker"
|
|
TRELLIS_WORKER = os.environ.get("TRELLIS_WORKER", "").strip()
|
|
VIDEO_LABEL_KEY = "com.mike-ai.video-worker"
|
|
VIDEO_WORKER = os.environ.get("VIDEO_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, payload: dict | None = None) -> tuple[int, bytes]:
|
|
conn = UnixConnection("localhost", timeout=30)
|
|
try:
|
|
body = json.dumps(payload).encode() if payload is not None else None
|
|
conn.request(method, path, body=body,
|
|
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}
|
|
if QWEN_IMAGE_TEST_WORKER:
|
|
allowed.add(QWEN_IMAGE_TEST_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 yue2_container() -> dict:
|
|
if not YUE2_WORKER:
|
|
raise RuntimeError("YuE2 worker is not configured")
|
|
matches = [item for item in labelled_containers(MUSIC_LABEL_KEY)
|
|
if item.get("Labels", {}).get(MUSIC_LABEL_KEY) == YUE2_WORKER]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one YuE2 worker {YUE2_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 trellis_container() -> dict:
|
|
if not TRELLIS_WORKER:
|
|
raise RuntimeError("TRELLIS worker is not configured")
|
|
matches = [item for item in labelled_containers(TRELLIS_LABEL_KEY)
|
|
if item.get("Labels", {}).get(TRELLIS_LABEL_KEY) == TRELLIS_WORKER]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one TRELLIS worker {TRELLIS_WORKER!r}, found {len(matches)}")
|
|
return matches[0]
|
|
|
|
|
|
def video_container() -> dict:
|
|
if not VIDEO_WORKER:
|
|
raise RuntimeError("video worker is not configured")
|
|
matches = [item for item in labelled_containers(VIDEO_LABEL_KEY)
|
|
if item.get("Labels", {}).get(VIDEO_LABEL_KEY) == VIDEO_WORKER]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(
|
|
f"expected exactly one video worker {VIDEO_WORKER!r}, found {len(matches)}")
|
|
return matches[0]
|
|
|
|
|
|
def docker_exec(item: dict, command: str) -> str:
|
|
status, body = docker_request("POST", f"/containers/{item['Id']}/exec", {
|
|
"AttachStdout": True, "AttachStderr": True, "Tty": True,
|
|
"Cmd": ["/bin/sh", "-c", command],
|
|
})
|
|
if status != 201:
|
|
raise RuntimeError(f"failed to create container exec: HTTP {status}")
|
|
exec_id = json.loads(body)["Id"]
|
|
status, body = docker_request("POST", f"/exec/{exec_id}/start",
|
|
{"Detach": False, "Tty": True})
|
|
if status != 200:
|
|
raise RuntimeError(f"failed to start container exec: HTTP {status}")
|
|
return body.decode(errors="replace").strip()
|
|
|
|
|
|
VIDEO_UI_PROCESSES = r'''for p in /proc/[0-9]*; do
|
|
[ -r "$p/cmdline" ] || continue
|
|
exe=$(readlink "$p/exe" 2>/dev/null) || continue
|
|
c=$(tr '\000' ' ' < "$p/cmdline")
|
|
case "${exe##*/}:$c" in
|
|
ltx-desktop:*--type=*)
|
|
printf '%s\n' "${p##*/}";;
|
|
esac
|
|
done'''
|
|
|
|
VIDEO_UI_STATE: str | None = None
|
|
|
|
|
|
def video_ui_state(item: dict | None = None) -> str:
|
|
item = item or video_container()
|
|
if item.get("State") != "running":
|
|
return "unavailable"
|
|
output = docker_exec(item, VIDEO_UI_PROCESSES + r''' | while read p; do
|
|
state=$(sed -n 's/^State:[[:space:]]*\([A-Z]\).*/\1/p' "/proc/$p/status")
|
|
printf '%s\n' "$state"
|
|
done''')
|
|
states = [line.strip() for line in output.splitlines() if line.strip()]
|
|
if not states:
|
|
return "unavailable"
|
|
return "sleeping" if all(state == "T" for state in states) else "awake"
|
|
|
|
|
|
def cached_video_ui_state(item: dict | None = None) -> str:
|
|
global VIDEO_UI_STATE
|
|
if VIDEO_UI_STATE is None:
|
|
VIDEO_UI_STATE = video_ui_state(item)
|
|
return VIDEO_UI_STATE
|
|
|
|
|
|
def set_video_ui(awake: bool) -> dict:
|
|
global VIDEO_UI_STATE
|
|
with LOCK:
|
|
item = video_container()
|
|
if item.get("State") != "running":
|
|
raise RuntimeError("LTX-2 is not running")
|
|
signal = "CONT" if awake else "STOP"
|
|
docker_exec(item, VIDEO_UI_PROCESSES +
|
|
f''' | while read p; do kill -{signal} "$p"; done''')
|
|
time.sleep(0.25)
|
|
VIDEO_UI_STATE = video_ui_state(item)
|
|
return {"video_ui": VIDEO_UI_STATE}
|
|
|
|
|
|
def stop_video_if_configured() -> None:
|
|
if VIDEO_WORKER:
|
|
stop_container(video_container(), timeout=30)
|
|
|
|
|
|
def stop_music_if_configured() -> None:
|
|
if MUSIC_WORKER:
|
|
stop_container(music_container(), timeout=30)
|
|
|
|
|
|
def stop_yue2_if_configured() -> None:
|
|
if YUE2_WORKER:
|
|
stop_container(yue2_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_trellis_if_configured() -> None:
|
|
if TRELLIS_WORKER:
|
|
stop_container(trellis_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 wait_container_healthy(item: dict, timeout: int) -> None:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
status, body = docker_request("GET", f"/containers/{item['Id']}/json")
|
|
if status != 200:
|
|
raise RuntimeError(f"failed to inspect container: HTTP {status}")
|
|
state = json.loads(body).get("State", {})
|
|
health = state.get("Health", {}).get("Status")
|
|
if state.get("Status") == "running" and health == "healthy":
|
|
return
|
|
if state.get("Status") in {"dead", "exited"} or health == "unhealthy":
|
|
raise RuntimeError(
|
|
f"container {item.get('Names', ['unknown'])[0]} failed: "
|
|
f"state={state.get('Status')} health={health}")
|
|
time.sleep(1)
|
|
raise RuntimeError(
|
|
f"container {item.get('Names', ['unknown'])[0]} did not become healthy")
|
|
|
|
|
|
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, QWEN_IMAGE_TEST_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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools()
|
|
stop_trellis_if_configured()
|
|
stop_video_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)
|
|
stop_separator_if_configured()
|
|
stop_voice_tools()
|
|
stop_yue2_if_configured()
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
start_container(item)
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
return {"music_worker": MUSIC_WORKER,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def set_yue2_worker(running: bool) -> dict:
|
|
"""Start YuE2 exclusively, or stop it before another mode is loaded."""
|
|
with LOCK:
|
|
item = yue2_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()
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
start_container(item)
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
return {"yue2_worker": YUE2_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_yue2_if_configured()
|
|
stop_voice_tools()
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools("voice")
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools("voicechange")
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools("applio")
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
start_container(item)
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
return {"applio_worker": APPLIO_WORKER,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def set_trellis_worker(running: bool) -> dict:
|
|
"""Start TRELLIS.2 exclusively, or stop it before another mode is loaded."""
|
|
with LOCK:
|
|
item = trellis_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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools()
|
|
start_container(item)
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
return {"trellis_worker": TRELLIS_WORKER,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def set_video_worker(running: bool) -> dict:
|
|
"""Start LTX-2 exclusively, or stop it before another mode is loaded."""
|
|
global VIDEO_UI_STATE
|
|
with LOCK:
|
|
item = video_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_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools()
|
|
stop_trellis_if_configured()
|
|
start_container(item)
|
|
VIDEO_UI_STATE = "awake"
|
|
else:
|
|
stop_container(item, timeout=30)
|
|
VIDEO_UI_STATE = "unavailable"
|
|
return {"video_worker": VIDEO_WORKER,
|
|
"state": "running" if running else "stopped"}
|
|
|
|
|
|
def active_profile(items: dict[str, dict] | None = None) -> str | None:
|
|
items = containers() if items is None else items
|
|
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:
|
|
# Validate before stopping any working service.
|
|
items = containers()
|
|
missing = [name for name in ALLOWED if name not in items]
|
|
if missing:
|
|
raise RuntimeError("profile containers missing: " + ", ".join(missing))
|
|
# Defensive mutual exclusion even if a caller bypasses the router.
|
|
for worker in image_containers():
|
|
stop_container(worker)
|
|
stop_music_if_configured()
|
|
stop_yue2_if_configured()
|
|
stop_separator_if_configured()
|
|
stop_voice_tools()
|
|
stop_trellis_if_configured()
|
|
stop_video_if_configured()
|
|
start_container(tts_container())
|
|
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 status_snapshot() -> dict:
|
|
"""One Docker snapshot; missing optional workers cannot hide the LLM."""
|
|
status, body = docker_request("GET", "/containers/json?all=1")
|
|
if status != 200:
|
|
raise RuntimeError(f"Docker list failed with HTTP {status}")
|
|
records = json.loads(body)
|
|
profiles = {}
|
|
for item in records:
|
|
name = item.get("Labels", {}).get(LABEL_KEY)
|
|
if name in ALLOWED:
|
|
if name in profiles:
|
|
raise RuntimeError(f"duplicate container for profile {name}")
|
|
profiles[name] = item
|
|
result = {"active_profile": active_profile(profiles),
|
|
"profiles": {name: profiles.get(name, {}).get("State", "missing")
|
|
for name in ALLOWED}, "worker_errors": {}}
|
|
workers = (
|
|
("music", MUSIC_LABEL_KEY, MUSIC_WORKER),
|
|
("yue2", MUSIC_LABEL_KEY, YUE2_WORKER),
|
|
("separator", SEPARATOR_LABEL_KEY, SEPARATOR_WORKER),
|
|
("voice", VOICE_LABEL_KEY, VOICE_WORKER),
|
|
("voice_change", VOICE_CHANGE_LABEL_KEY, VOICE_CHANGE_WORKER),
|
|
("applio", APPLIO_LABEL_KEY, APPLIO_WORKER),
|
|
("trellis", TRELLIS_LABEL_KEY, TRELLIS_WORKER),
|
|
("video", VIDEO_LABEL_KEY, VIDEO_WORKER),
|
|
)
|
|
result["video_ui"] = "unavailable"
|
|
for name, label, kind in workers:
|
|
matches = [item for item in records
|
|
if kind and item.get("Labels", {}).get(label) == kind]
|
|
if not kind:
|
|
state = health = "disabled"
|
|
elif len(matches) != 1:
|
|
state = health = "missing" if not matches else "error"
|
|
result["worker_errors"][name] = f"expected one container, found {len(matches)}"
|
|
else:
|
|
item = matches[0]
|
|
state = item.get("State", "unknown")
|
|
description = item.get("Status", "")
|
|
health = ("healthy" if "(healthy)" in description else
|
|
"unhealthy" if "(unhealthy)" in description else
|
|
"starting" if state == "running" else "stopped")
|
|
if name == "video":
|
|
try:
|
|
result["video_ui"] = cached_video_ui_state(item)
|
|
except Exception as exc:
|
|
result["worker_errors"]["video_ui"] = str(exc)
|
|
result[f"{name}_worker"] = state
|
|
result[f"{name}_health"] = health
|
|
return result
|
|
|
|
|
|
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 not in {"/status", "/profiles/status"}:
|
|
self.reply(404, {"error": "not found"})
|
|
return
|
|
if not self.authenticated():
|
|
self.reply(401, {"error": "unauthorized"})
|
|
return
|
|
try:
|
|
if self.path == "/profiles/status":
|
|
items = containers()
|
|
self.reply(200, {"active_profile": active_profile(items)})
|
|
else:
|
|
self.reply(200, status_snapshot())
|
|
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/yue2/start", "/workers/yue2/stop"}:
|
|
try:
|
|
self.reply(200, set_yue2_worker(self.path.endswith("/start")))
|
|
except Exception as exc:
|
|
log.exception("YuE2 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
|
|
if self.path in {"/workers/trellis/start", "/workers/trellis/stop"}:
|
|
try:
|
|
self.reply(200, set_trellis_worker(self.path.endswith("/start")))
|
|
except Exception as exc:
|
|
log.exception("TRELLIS worker transition failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
return
|
|
if self.path in {"/workers/video/start", "/workers/video/stop"}:
|
|
try:
|
|
self.reply(200, set_video_worker(self.path.endswith("/start")))
|
|
except Exception as exc:
|
|
log.exception("video worker transition failed")
|
|
self.reply(503, {"error": str(exc)})
|
|
return
|
|
if self.path in {"/workers/video/ui/sleep", "/workers/video/ui/wake"}:
|
|
try:
|
|
self.reply(200, set_video_ui(self.path.endswith("/wake")))
|
|
except Exception as exc:
|
|
log.exception("video UI 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),
|
|
"/workers/qwen-image-test/start": (QWEN_IMAGE_TEST_WORKER, True),
|
|
"/workers/qwen-image-test/stop": (QWEN_IMAGE_TEST_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()
|