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

#!/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,experimental").split(",") if x.strip())
LABEL_KEY = "com.mike-ai.llama-profile"
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 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:
        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()
            self.reply(200, {"active_profile": active_profile(items),
                             "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
        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()