Prechádzať zdrojové kódy

feat(fabric): report per-GPU memory.total -> node vram_total_mb -> CP total_gpu_capacity_mb

Agent (agent/main.py): add per-GPU mem_total_mb (nvidia-smi memory.total) and
node-level vram_total_mb (sum, counted once per physical GPU) to /gpu-status and
the heartbeat payload (HeartbeatExt subclass). Control (control/main.py): carry
node vram_total_mb into /registry/nodes node objects and add top-level
total_gpu_capacity_mb to /registry/models. Additive and backward-compatible;
deployed fleet-wide (12 nodes). Existing total_gpu_vram_mb (sum-of-used) kept.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude 1 mesiac pred
commit
4ab5e8db56
2 zmenil súbory, kde vykonal 416 pridanie a 0 odobranie
  1. 201 0
      agent/main.py
  2. 215 0
      control/main.py

+ 201 - 0
agent/main.py

@@ -0,0 +1,201 @@
+"""Fabric Node Agent — module discovery, GPU monitoring, systemd lifecycle, heartbeat."""
+from __future__ import annotations
+import asyncio
+import logging
+import os
+import socket
+import subprocess
+import time
+
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import httpx
+
+from fabric.shared.schemas import Heartbeat, ModuleState
+from fabric.agent.discovery import scan_modules
+from fabric.agent.gpu_monitor import get_per_pid_vram, get_gpu_count
+from fabric.agent.systemd_wrap import SystemdWrapRuntime
+from fabric.agent.heartbeat import heartbeat_loop
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+log = logging.getLogger("fabric.agent")
+
+NODE_ID = os.environ.get("FABRIC_NODE_ID", socket.gethostname())
+NODE_IP = os.environ.get("FABRIC_NODE_IP", "")
+AGENT_PORT = int(os.environ.get("FABRIC_AGENT_PORT", "7701"))
+HYPERVISOR = os.environ.get("FABRIC_HYPERVISOR", "unknown")
+
+
+class GpuMemInfo(BaseModel):
+    """Per-physical-GPU memory capacity entry (additive VRAM-capacity reporting)."""
+    gpu_index: int = 0
+    mem_total_mb: int = 0
+    mem_used_mb: int = 0
+
+
+class HeartbeatExt(Heartbeat):
+    """Heartbeat + additive VRAM-capacity fields. Subclassing keeps the base
+    payload identical for control planes that haven't been updated yet (they
+    ignore the extra fields) while carrying per-GPU capacity to updated ones."""
+    gpus: list[GpuMemInfo] = []
+    vram_total_mb: int = 0
+
+
+app = FastAPI(title=f"Fabric Agent ({NODE_ID})", version="0.1.0")
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+_runtimes: dict[str, SystemdWrapRuntime] = {}
+_start_time = time.time()
+
+
+def _detect_ip() -> str:
+    if NODE_IP:
+        return NODE_IP
+    try:
+        r = subprocess.run(["hostname", "-I"], capture_output=True, text=True, timeout=5)
+        for ip in r.stdout.strip().split():
+            if not ip.startswith("127."):
+                return ip
+    except Exception:
+        pass
+    return "127.0.0.1"
+
+
+def _get_gpu_mem_info() -> list[dict]:
+    """Per-physical-GPU total/used VRAM via nvidia-smi (one entry per GPU, counted
+    once — NOT per process). Robust if nvidia-smi is missing or fails -> returns []."""
+    try:
+        r = subprocess.run(
+            ["nvidia-smi", "--query-gpu=index,memory.total,memory.used",
+             "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=5,
+        )
+    except Exception as e:
+        log.warning("nvidia-smi capacity query failed: %s", e)
+        return []
+    if r.returncode != 0:
+        return []
+    gpus: list[dict] = []
+    for line in r.stdout.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) < 2:
+            continue
+        try:
+            idx = int(parts[0])
+            total = int(float(parts[1]))
+        except ValueError:
+            continue
+        used = 0
+        if len(parts) >= 3:
+            try:
+                used = int(float(parts[2]))
+            except ValueError:
+                used = 0
+        gpus.append({"gpu_index": idx, "mem_total_mb": total, "mem_used_mb": used})
+    return gpus
+
+
+def _load_modules():
+    global _runtimes
+    manifests = scan_modules()
+    for m in manifests:
+        if m.runtime.value == "systemd-wrap" and m.systemd:
+            _runtimes[m.name] = SystemdWrapRuntime(m)
+            log.info("Loaded systemd-wrap: %s -> %s", m.name, m.systemd.systemd_unit)
+        else:
+            log.warning("Unsupported runtime %s for %s (v0: systemd-wrap only)", m.runtime, m.name)
+    pid_vram = _get_pid_vram_map()
+    for name, rt in _runtimes.items():
+        rt.update_state(pid_vram)
+
+
+def _get_pid_vram_map() -> dict[int, int]:
+    m: dict[int, int] = {}
+    for p in get_per_pid_vram():
+        m[p.pid] = m.get(p.pid, 0) + p.vram_mb
+    return m
+
+
+def _build_heartbeat() -> HeartbeatExt:
+    gpu_status = get_per_pid_vram()
+    pid_vram = _get_pid_vram_map()
+    modules = []
+    for name, rt in _runtimes.items():
+        rt.update_state(pid_vram)
+        modules.append(rt.get_status())
+    gpus = _get_gpu_mem_info()
+    return HeartbeatExt(
+        node_id=NODE_ID, node_ip=_detect_ip(), agent_port=AGENT_PORT,
+        hypervisor=HYPERVISOR, gpu_status=gpu_status, modules=modules,
+        uptime_s=int(time.time() - _start_time),
+        gpus=gpus,
+        vram_total_mb=sum(g["mem_total_mb"] for g in gpus),
+    )
+
+
+@app.get("/health")
+async def health():
+    return {"ok": True, "node_id": NODE_ID, "modules": len(_runtimes), "gpu_count": get_gpu_count(), "uptime_s": int(time.time() - _start_time)}
+
+@app.get("/modules")
+async def list_modules():
+    pid_vram = _get_pid_vram_map()
+    return {"modules": [rt.update_state(pid_vram) or rt.get_status().model_dump() for rt in _runtimes.values()]}
+
+@app.get("/gpu-status")
+async def gpu_status():
+    from fabric.agent.gpu_monitor import get_local_gpu_summary
+    summary = get_local_gpu_summary()
+    gpus = _get_gpu_mem_info()
+    summary["gpus"] = gpus
+    summary["vram_total_mb"] = sum(g["mem_total_mb"] for g in gpus)
+    return summary
+
+@app.post("/load/{module_name}")
+async def load_module(module_name: str):
+    rt = _runtimes.get(module_name)
+    if not rt:
+        raise HTTPException(404, f"Module '{module_name}' not found")
+    ok = rt.start()
+    if not ok:
+        raise HTTPException(500, f"Failed to start '{module_name}'")
+    warmup = rt.manifest.lifecycle.warmup_s
+    if warmup > 0:
+        await asyncio.sleep(min(warmup, 5))
+    return {"ok": True, "module": module_name, "state": "loaded"}
+
+@app.post("/unload/{module_name}")
+async def unload_module(module_name: str):
+    rt = _runtimes.get(module_name)
+    if not rt:
+        raise HTTPException(404, f"Module '{module_name}' not found")
+    return {"ok": rt.stop(), "module": module_name, "state": "unloaded"}
+
+@app.post("/infer/{module_name}")
+async def infer_proxy(module_name: str, request: Request):
+    rt = _runtimes.get(module_name)
+    if not rt:
+        raise HTTPException(404, f"Module '{module_name}' not found")
+    if not rt.is_active():
+        raise HTTPException(503, f"Module '{module_name}' not active")
+    rt.record_infer()
+    body = await request.body()
+    headers = {k: v for k, v in request.headers.items() if k.lower() not in ("host", "content-length", "transfer-encoding")}
+    async with httpx.AsyncClient(timeout=120) as client:
+        try:
+            resp = await client.request(
+                method=rt.manifest.systemd.infer_method if rt.manifest.systemd else "POST",
+                url=rt.infer_url, content=body, headers=headers,
+            )
+            return Response(content=resp.content, status_code=resp.status_code, headers=dict(resp.headers))
+        except Exception as e:
+            raise HTTPException(502, f"Proxy failed: {e}")
+
+
+@app.on_event("startup")
+async def startup():
+    log.info("Fabric Agent starting: node_id=%s ip=%s", NODE_ID, _detect_ip())
+    _load_modules()
+    asyncio.create_task(heartbeat_loop(_build_heartbeat))
+    log.info("Agent ready with %d modules", len(_runtimes))

+ 215 - 0
control/main.py

@@ -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")