fabric/control/main.py

"""Fabric Control Plane — registry + placement + scheduler + installer."""
from __future__ import annotations
import asyncio
import logging
import time

from fastapi import FastAPI, HTTPException, Query, Body
from fastapi.middleware.cors import CORSMiddleware

from fabric.shared.schemas import (
    Heartbeat, PlaceRequest, PlaceResponse,
    RegistryModelsResponse, RegistryNodesResponse,
)
from fabric.control.registry import Registry
from fabric.control.scheduler import Scheduler
from fabric.control.installer import Installer, RECIPES
from fabric.control import cuda_checkpoint

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
log = logging.getLogger("fabric.control")

app = FastAPI(title="Fabric Control Plane", version="0.2.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

registry = Registry()
scheduler = Scheduler(registry)
installer = Installer()
_start_time = time.time()


@app.get("/health")
async def health():
    return {
        "ok": True,
        "version": registry.version,
        "nodes": len(registry.get_alive_nodes()),
        "modules": len(registry.get_all_modules()),
        "uptime_s": int(time.time() - _start_time),
        "cuda_checkpoint_available": cuda_checkpoint.is_available(),
    }

@app.post("/registry/heartbeat")
async def heartbeat(hb: Heartbeat):
    registry.handle_heartbeat(hb)
    return {"ok": True, "version": registry.version}

@app.get("/registry/models", response_model=RegistryModelsResponse)
async def list_models():
    modules = registry.get_all_modules()
    return RegistryModelsResponse(
        models=modules,
        nodes_count=len(registry.get_alive_nodes()),
        total_gpu_vram_mb=registry.total_gpu_vram_mb(),
        total_gpu_capacity_mb=registry.total_gpu_capacity_mb(),
        total_gpu_free_mb=registry.total_gpu_free_mb(),
    )

@app.get("/registry/nodes", response_model=RegistryNodesResponse)
async def list_nodes():
    return RegistryNodesResponse(nodes=registry.get_alive_nodes())

@app.get("/registry/hypervisors")
async def list_hypervisors():
    """Group nodes by hypervisor — multi-hypervisor visibility."""
    by_hyp: dict[str, list[dict]] = {}
    for n in registry.get_alive_nodes():
        hyp = n.hypervisor or "unknown"
        by_hyp.setdefault(hyp, []).append({
            "node_id": n.node_id,
            "node_ip": n.node_ip,
            "modules": len(n.modules),
            "gpu_processes": len(n.gpu_status),
            "vram_used_mb": sum(g.vram_mb for g in n.gpu_status),
            "vram_total_mb": n.vram_total_mb,
            "vram_free_mb": n.vram_free_mb,
        })
    return {"hypervisors": by_hyp, "total_hypervisors": len(by_hyp)}

@app.post("/place", response_model=PlaceResponse)
async def place(req: PlaceRequest):
    result = registry.place(req)
    if result is None:
        raise HTTPException(status_code=404, detail=f"Model '{req.model}' not found in any node")
    if result.state in ("unloaded", "error"):
        # Anti-OOM guard: refuse to (re)load onto a node that lacks the headroom
        # the model is known to need. Footprint estimated from the model's last
        # observed VRAM usage; skipped when unknown (0) or capacity data absent.
        candidates = registry.find_module(req.model)
        footprint = max((m.vram_mb for (_n, m) in candidates), default=0)
        chosen = registry.nodes.get(result.node_id)
        free = (chosen.vram_free_mb if chosen else 0) or 0
        if footprint > 0 and free > 0 and free < footprint:
            raise HTTPException(
                status_code=503,
                detail=f"Insufficient free VRAM on '{result.node_id}' for '{req.model}': "
                       f"needs ~{footprint}mb, free {free}mb",
            )
        await scheduler.ensure_loaded(result.node_id, result.module_name)
    return result

@app.post("/admin/evict")
async def admin_evict(node_id: str, module_name: str, reason: str = "admin"):
    ok = await scheduler.evict(node_id, module_name, reason)
    if not ok:
        raise HTTPException(400, "Eviction failed")
    return {"ok": True}

@app.post("/admin/load")
async def admin_load(node_id: str, module_name: str):
    ok = await scheduler.ensure_loaded(node_id, module_name)
    if not ok:
        raise HTTPException(400, "Load failed")
    return {"ok": True}

@app.get("/scheduler/utilization")
async def scheduler_utilization():
    return scheduler.get_utilization()

@app.post("/scheduler/preempt")
async def scheduler_preempt(model: str, reason: str = "placement"):
    result = await scheduler.preempt_for(model, reason)
    if result is None:
        raise HTTPException(503, "No resources available for preemption")
    return result

@app.post("/batch")
async def batch_submit(model: str, items: list[dict] = Body(...), priority: str = "normal"):
    return await scheduler.submit_batch(model, items, priority)

@app.get("/batch/{job_id}")
async def batch_status(job_id: str):
    status = scheduler.get_batch_status(job_id)
    if status is None:
        raise HTTPException(404, "Batch job not found")
    return status

# --- v2: Installer endpoints ---

@app.get("/install/recipes")
async def list_recipes():
    return {"recipes": installer.list_recipes()}

@app.post("/install")
async def install_model(hf_repo: str, recipe: str = "vllm", name: str | None = None):
    try:
        return await installer.install(hf_repo, recipe, name)
    except ValueError as e:
        raise HTTPException(400, str(e))

@app.get("/install/{job_id}")
async def install_status(job_id: str):
    s = installer.get_status(job_id)
    if s is None:
        raise HTTPException(404, "Install job not found")
    return s

@app.get("/install")
async def list_install_jobs():
    return {"jobs": installer.list_jobs()}

# --- v2: cuda-checkpoint endpoints ---

@app.get("/checkpoint/availability")
async def checkpoint_availability():
    return {"available": cuda_checkpoint.is_available(), "binary": cuda_checkpoint.CUDA_CHECKPOINT_BIN}

@app.post("/checkpoint/{pid}/toggle")
async def checkpoint_toggle(pid: int, timeout_s: int = 30):
    ok, msg = await cuda_checkpoint.checkpoint(pid, timeout_s)
    return {"ok": ok, "pid": pid, "message": msg}

@app.get("/checkpoint/{pid}/state")
async def checkpoint_state(pid: int):
    state = await cuda_checkpoint.get_state(pid)
    return {"pid": pid, "state": state}


async def _prune_loop():
    while True:
        await asyncio.sleep(15)
        stale = registry.prune_stale_nodes()
        if stale:
            log.warning("Pruned stale nodes: %s", stale)

async def _idle_sweep_loop():
    while True:
        await asyncio.sleep(30)
        await scheduler.idle_sweep()

@app.on_event("startup")
async def startup():
    asyncio.create_task(_prune_loop())
    asyncio.create_task(_idle_sweep_loop())
    log.info("Fabric Control Plane v0.2.0 started on port 7700")