|
@@ -0,0 +1,215 @@
|
|
|
|
|
+"""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 pydantic import BaseModel
|
|
|
|
|
+
|
|
|
|
|
+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()
|
|
|
|
|
+
|
|
|
|
|
+# Additive VRAM-capacity reporting: per-node physical-GPU capacity captured from
|
|
|
|
|
+# heartbeats, keyed by node_id. Kept outside the strict registry schemas so the
|
|
|
|
|
+# change stays backward-compatible (absent fields degrade to 0 / []).
|
|
|
|
|
+_node_capacity: dict[str, dict] = {}
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class GpuMemInfo(BaseModel):
|
|
|
|
|
+ """Per-physical-GPU memory capacity entry (mirrors the agent payload)."""
|
|
|
|
|
+ gpu_index: int = 0
|
|
|
|
|
+ mem_total_mb: int = 0
|
|
|
|
|
+ mem_used_mb: int = 0
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class HeartbeatExt(Heartbeat):
|
|
|
|
|
+ """Heartbeat + additive VRAM-capacity fields. Old agents omit these and the
|
|
|
|
|
+ defaults apply; updated agents carry per-GPU capacity through to the registry."""
|
|
|
|
|
+ gpus: list[GpuMemInfo] = []
|
|
|
|
|
+ vram_total_mb: int = 0
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@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: HeartbeatExt):
|
|
|
|
|
+ registry.handle_heartbeat(hb)
|
|
|
|
|
+ _node_capacity[hb.node_id] = {
|
|
|
|
|
+ "vram_total_mb": int(hb.vram_total_mb or 0),
|
|
|
|
|
+ "gpus": [g.model_dump(mode="json") for g in hb.gpus],
|
|
|
|
|
+ }
|
|
|
|
|
+ return {"ok": True, "version": registry.version}
|
|
|
|
|
+
|
|
|
|
|
+@app.get("/registry/models")
|
|
|
|
|
+async def list_models():
|
|
|
|
|
+ modules = registry.get_all_modules()
|
|
|
|
|
+ alive = registry.get_alive_nodes()
|
|
|
|
|
+ base = RegistryModelsResponse(
|
|
|
|
|
+ models=modules,
|
|
|
|
|
+ nodes_count=len(alive),
|
|
|
|
|
+ total_gpu_vram_mb=registry.total_gpu_vram_mb(),
|
|
|
|
|
+ )
|
|
|
|
|
+ data = base.model_dump(mode="json")
|
|
|
|
|
+ data["total_gpu_capacity_mb"] = sum(
|
|
|
|
|
+ int((_node_capacity.get(n.node_id) or {}).get("vram_total_mb", 0) or 0)
|
|
|
|
|
+ for n in alive
|
|
|
|
|
+ )
|
|
|
|
|
+ return data
|
|
|
|
|
+
|
|
|
|
|
+@app.get("/registry/nodes")
|
|
|
|
|
+async def list_nodes():
|
|
|
|
|
+ nodes_out = []
|
|
|
|
|
+ for n in registry.get_alive_nodes():
|
|
|
|
|
+ d = n.model_dump(mode="json")
|
|
|
|
|
+ cap = _node_capacity.get(n.node_id) or {}
|
|
|
|
|
+ d["vram_total_mb"] = int(cap.get("vram_total_mb", 0) or 0)
|
|
|
|
|
+ d["gpus"] = cap.get("gpus", [])
|
|
|
|
|
+ nodes_out.append(d)
|
|
|
|
|
+ return {"nodes": nodes_out}
|
|
|
|
|
+
|
|
|
|
|
+@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),
|
|
|
|
|
+ })
|
|
|
|
|
+ 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"):
|
|
|
|
|
+ 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")
|