Synchronize Athena operating modes with live deployment

platform/docker/profile-controller/profile_controller.py

@@ -13,6 +13,7 @@ import logging
 import os
 import socket
 import threading
+import time
 import urllib.parse
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
@@ -25,6 +26,22 @@ ALLOWED = tuple(x.strip() for x in os.environ.get(
 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()
+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()
 LOCK = threading.Lock()
 log = logging.getLogger("profile-controller")
 
@@ -68,15 +85,152 @@ def labelled_containers(label: str) -> list[dict]:
     return json.loads(body)
 
 
-def image_container() -> dict:
+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) == IMAGE_WORKER]
+               if item.get("Labels", {}).get(IMAGE_LABEL_KEY) == kind]
     if len(matches) != 1:
         raise RuntimeError(
-            f"expected exactly one image worker {IMAGE_WORKER!r}, found {len(matches)}")
+            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 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 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
@@ -85,6 +239,31 @@ def stop_container(item: dict, timeout: int = 120) -> None:
         raise RuntimeError(f"failed to stop container: HTTP {status}")
 
 
+def start_container(item: dict) -> None:
+    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()
@@ -94,22 +273,186 @@ def stop_inference() -> dict:
         return {"active_profile": None, "previous_profile": previous}
 
 
-def set_image_worker(running: bool) -> dict:
+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()
+        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)
-            if item.get("State") != "running":
-                status, _ = docker_request("POST", f"/containers/{item['Id']}/start")
-                if status not in (204, 304):
-                    raise RuntimeError(f"failed to start image worker: HTTP {status}")
+            # 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()
+            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)
-        return {"image_worker": "running" if running else "stopped"}
+            # 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()
+            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()
+            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()
+            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()
+            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()
+            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()
+            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 active_profile(items: dict[str, dict] | None = None) -> str | None:
@@ -125,7 +468,14 @@ def activate(profile: str) -> dict:
         raise ValueError("profile is not allowlisted")
     with LOCK:
         # Defensive mutual exclusion even if a caller bypasses the router.
-        stop_container(image_container())
+        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()
+        start_container(tts_container())
         items = containers()
         missing = [name for name in ALLOWED if name not in items]
         if missing:
@@ -189,7 +539,70 @@ class Handler(BaseHTTPRequestHandler):
             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")
+            yue2 = yue2_container() if YUE2_WORKER else {}
+            yue2_status = yue2.get("Status", "")
+            yue2_health = ("disabled" if not YUE2_WORKER else
+                           "healthy" if "(healthy)" in yue2_status else
+                           "unhealthy" if "(unhealthy)" in yue2_status else
+                           "starting" if yue2.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")
+            trellis = trellis_container() if TRELLIS_WORKER else {}
+            trellis_status = trellis.get("Status", "")
+            trellis_health = ("disabled" if not TRELLIS_WORKER else
+                              "healthy" if "(healthy)" in trellis_status else
+                              "unhealthy" if "(unhealthy)" in trellis_status else
+                              "starting" if trellis.get("State") == "running" else
+                              "stopped")
             self.reply(200, {"active_profile": active_profile(items),
+                             "music_worker": music.get("State", "disabled"),
+                             "music_health": music_health,
+                             "yue2_worker": yue2.get("State", "disabled"),
+                             "yue2_health": yue2_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,
+                             "trellis_worker": trellis.get("State", "disabled"),
+                             "trellis_health": trellis_health,
                              "profiles": {name: items.get(name, {}).get(
                                  "State", "missing") for name in ALLOWED}})
         except Exception as exc:
@@ -207,9 +620,65 @@ class Handler(BaseHTTPRequestHandler):
                 log.exception("stopping inference failed")
                 self.reply(503, {"error": str(exc)})
             return
-        if self.path in ("/workers/image/start", "/workers/image/stop"):
+        if self.path in {"/workers/music/start", "/workers/music/stop"}:
             try:
-                self.reply(200, set_image_worker(self.path.endswith("/start")))
+                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
+        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)})