Add RTX 5080 FLUX hot-swap worker

platform/docker/profile-controller/profile_controller.py

@@ -23,6 +23,8 @@ TOKEN_FILE = os.environ.get("CONTROLLER_TOKEN_FILE", "/run/secrets/controller-to
 ALLOWED = tuple(x.strip() for x in os.environ.get(
     "ALLOWED_PROFILES", "fast,medium,large,ultra,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", "flux")
 LOCK = threading.Lock()
 log = logging.getLogger("profile-controller")
 
@@ -58,6 +60,56 @@ def containers() -> dict[str, dict]:
     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() -> dict:
+    matches = [item for item in labelled_containers(IMAGE_LABEL_KEY)
+               if item.get("Labels", {}).get(IMAGE_LABEL_KEY) == IMAGE_WORKER]
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"expected exactly one image worker {IMAGE_WORKER!r}, found {len(matches)}")
+    return matches[0]
+
+
+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 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) -> dict:
+    with LOCK:
+        item = image_container()
+        if running:
+            # A FLUX 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}")
+        else:
+            stop_container(item)
+        return {"image_worker": "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"]
@@ -70,6 +122,8 @@ 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.
+        stop_container(image_container())
         items = containers()
         missing = [name for name in ALLOWED if name not in items]
         if missing:
@@ -144,6 +198,20 @@ class Handler(BaseHTTPRequestHandler):
         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/image/start", "/workers/image/stop"):
+            try:
+                self.reply(200, set_image_worker(self.path.endswith("/start")))
+            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"})