llama-cpp/llama-swap/controller.py

#!/usr/bin/env python3
"""llama-swap controller — hot-swap models via API"""
import os, time
from typing import Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
import docker, httpx
from pydantic import BaseModel

app = FastAPI(title="llama-swap", version="1.1.0")
client = docker.from_env()

MODELS = {
    "ornith-9b": {"profile": "ornith-9b", "file": "ornith-9b-Q8_0.gguf", "name": "Ornith 9B Q8 (coding, MIT, 32K)", "env_file": "/workspace/envs/.env.ornith-9b"},
    "qwen35-9b": {"profile": "qwen35-9b", "file": "Qwen3.5-9B.Q8_0.gguf", "name": "Qwen3.5 9B Q8 (reasoning, 32K)", "env_file": "/workspace/envs/.env.qwen35-9b"},
    "qwen3-4b": {"profile": "qwen3-4b", "file": "Qwen3-4B-Q4_K_M.gguf", "name": "Qwen3 4B Q4 (thinking, 16K)", "env_file": "/workspace/envs/.env.qwen3-4b"},
    "gemma4-e2b": {"profile": "gemma4-e2b", "file": "google_gemma-4-E2B-it-Q4_K_M.gguf", "name": "Gemma 4 E2B (multimodal, 24K)", "env_file": "/workspace/envs/.env.gemma4-e2b"},
    "gemma4-e4b": {"profile": "gemma4-e4b", "file": "google_gemma-4-E4B-it-Q4_K_M.gguf", "name": "Gemma 4 E4B (multimodal, 24K, CPU-split)", "env_file": "/workspace/envs/.env.gemma4-e4b"},
    "smollm3-3b": {"profile": "smollm3-3b", "file": "HuggingFaceTB_SmolLM3-3B-Q4_K_M.gguf", "name": "SmolLM3 3B (thinking+tools, 24K)", "env_file": "/workspace/envs/.env.smollm3-3b"}
}
CONTAINER_NAME, LLAMA_API = "llama_server", "http://llama_server:8080"
IMAGE = "local/llama-cpp-turboquant:server-cuda-sm75-mmq"
NETWORK = "llama-cpp_llama-net"

def load_env_file(path: str) -> dict:
    """Parse .env file into dict"""
    env = {}
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, val = line.split('=', 1)
                    env[key.strip()] = val.strip()
    except: pass
    return env

def get_active_model() -> Optional[str]:
    try:
        c = client.containers.get(CONTAINER_NAME)
        if c.status != "running": return None
        model = c.labels.get("llama-swap.model")
        return model if model in MODELS else None
    except: pass
    return None

def stop_active():
    try:
        c = client.containers.get(CONTAINER_NAME)
        c.stop(timeout=10); c.remove()
    except: pass

def start_model(model_id: str) -> dict:
    if model_id not in MODELS: raise ValueError(f"Unknown model: {model_id}")
    stop_active()
    
    meta = MODELS[model_id]
    env = load_env_file(meta["env_file"])
    
    # Build command args from env vars
    cmd_args = [
        "/app/llama-server",
        "--model", f"/models/{env.get('MODEL_FILE', meta['file'])}",
        "--host", "0.0.0.0",
        "--port", "8080",
        "--ctx-size", env.get('CTX_SIZE', '32768'),
        "--n-gpu-layers", env.get('N_GPU_LAYERS', '99'),
        "--threads", env.get('THREADS', '6'),
        "--threads-batch", env.get('THREADS_BATCH', '6'),
        "--batch-size", env.get('BATCH_SIZE', '512'),
        "--ubatch-size", env.get('UBATCH_SIZE', '128'),
        "--cache-type-k", env.get('CACHE_TYPE_K', 'f16'),
        "--cache-type-v", env.get('CACHE_TYPE_V', 'f16'),
        "--parallel", env.get('PARALLEL', '1')
    ]
    
    # Add EXTRA_ARGS if present (strip quotes if present)
    extra = env.get('EXTRA_ARGS', '').strip().strip('"').strip("'")
    if extra:
        cmd_args.extend(extra.split())
    
    # Create container with GPU device request
    try:
        container = client.containers.create(
            image=IMAGE,
            name=CONTAINER_NAME,
            entrypoint=[],  # Override image entrypoint
            command=cmd_args,
            detach=True,
            runtime="nvidia",
            device_requests=[
                docker.types.DeviceRequest(
                    count=-1,  # all GPUs
                    capabilities=[['gpu', 'compute', 'utility']]
                )
            ],
            environment={
                "NVIDIA_VISIBLE_DEVICES": "all",
                "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
                **env
            },
            volumes={
                "/home/moze/Sources/llama-cpp/models": {"bind": "/models", "mode": "ro"}
            },
            ports={"8080/tcp": 8080},
            network=NETWORK,
            shm_size="1g",
            ulimits=[docker.types.Ulimit(name='memlock', soft=-1, hard=-1)],
            restart_policy={"Name": "unless-stopped"},
            labels={
                "com.docker.compose.project": "llama-cpp",
                "com.docker.compose.service": f"llama-{model_id}",
                "llama-swap.model": model_id
            },
            healthcheck={
                "test": ["CMD-SHELL", "curl -sf http://localhost:8080/health | grep -q '\"status\":\"ok\"'"],
                "interval": 20000000000,  # 20s in nanoseconds
                "timeout": 10000000000,
                "retries": 10,
                "start_period": 180000000000  # 180s for 9B models with mlock
            }
        )
        container.start()
        
        # Wait for health check
        for _ in range(90):
            try:
                container.reload()
                health = container.attrs.get("State", {}).get("Health", {}).get("Status")
                if health == "healthy":
                    return {"model": model_id, "status": "healthy", "container": CONTAINER_NAME}
            except: pass
            time.sleep(2)
        
        return {"model": model_id, "status": "starting", "container": CONTAINER_NAME}
    except Exception as e:
        raise RuntimeError(f"Failed to start {model_id}: {str(e)}")

class SwitchRequest(BaseModel):
    model: str

@app.get("/")
async def root(): return {"service": "llama-swap", "version": "1.1.0"}

@app.get("/models")
async def list_models():
    active = get_active_model()
    return {"models": [{"id": mid, "name": m["name"], "file": m["file"], "active": mid==active} for mid, m in MODELS.items()], "active": active}

@app.post("/models/switch")
async def switch_model(req: SwitchRequest):
    if req.model not in MODELS: raise HTTPException(404, f"Unknown: {req.model}")
    if get_active_model() == req.model: return {"model": req.model, "status": "already_active"}
    try: return start_model(req.model)
    except Exception as e: raise HTTPException(500, str(e))

@app.get("/status")
async def status():
    active = get_active_model()
    if not active: return {"status": "no_model", "model": None}
    try:
        c = client.containers.get(CONTAINER_NAME)
        return {"status": "running", "model": active, "health": c.attrs.get("State",{}).get("Health",{}).get("Status","unknown")}
    except: return {"status": "no_model", "model": None}

@app.api_route("/v1/{path:path}", methods=["GET","POST","PUT","DELETE","PATCH"])
async def proxy_llama(path: str, request: Request):
    if not get_active_model(): raise HTTPException(503, "No model active")
    url = f"{LLAMA_API}/v1/{path}"
    headers = dict(request.headers); headers.pop("host", None)
    async with httpx.AsyncClient(timeout=300.0) as http:
        if request.method == "GET":
            resp = await http.get(url, headers=headers, params=request.query_params)
        else:
            resp = await http.request(request.method, url, headers=headers, params=request.query_params, content=await request.body())
        if "text/event-stream" in resp.headers.get("content-type",""):
            return StreamingResponse(resp.aiter_bytes(), media_type=resp.headers.get("content-type"), headers=dict(resp.headers))
        return JSONResponse(content=resp.json() if "json" in resp.headers.get("content-type","") else resp.text, status_code=resp.status_code, headers=dict(resp.headers))