fabric/agent/main.py

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

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_info() -> dict:
    """Per-physical-GPU memory + telemetry via nvidia-smi (each GPU counted once,
    NOT per process). Returns {gpus, vram_total_mb, vram_free_mb, driver_version}.
    Robust if nvidia-smi is missing or fails -> zeros / empty list."""
    empty = {"gpus": [], "vram_total_mb": 0, "vram_free_mb": 0, "driver_version": ""}
    try:
        r = subprocess.run(
            ["nvidia-smi",
             "--query-gpu=index,name,memory.total,memory.used,utilization.gpu,temperature.gpu,power.draw,driver_version",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5,
        )
    except Exception as e:
        log.warning("nvidia-smi query failed: %s", e)
        return empty
    if r.returncode != 0:
        return empty

    def _to_int(s):
        try:
            return int(float(s))
        except (ValueError, TypeError):
            return None

    def _to_float(s):
        try:
            return float(s)
        except (ValueError, TypeError):
            return None

    gpus: list[dict] = []
    driver = ""
    for line in r.stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 4:
            continue
        idx = _to_int(parts[0])
        if idx is None:
            continue
        total = _to_int(parts[2]) or 0
        used = _to_int(parts[3]) or 0
        if len(parts) > 7 and parts[7]:
            driver = parts[7]
        gpus.append({
            "gpu_index": idx,
            "name": parts[1],
            "mem_total_mb": total,
            "mem_used_mb": used,
            "mem_free_mb": max(total - used, 0),
            "util_pct": _to_int(parts[4]) if len(parts) > 4 else None,
            "temp_c": _to_int(parts[5]) if len(parts) > 5 else None,
            "power_w": _to_float(parts[6]) if len(parts) > 6 else None,
        })
    return {
        "gpus": gpus,
        "vram_total_mb": sum(g["mem_total_mb"] for g in gpus),
        "vram_free_mb": sum(g["mem_free_mb"] for g in gpus),
        "driver_version": driver,
    }


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() -> Heartbeat:
    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())
    info = _get_gpu_info()
    return Heartbeat(
        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=info["gpus"],
        vram_total_mb=info["vram_total_mb"],
        vram_free_mb=info["vram_free_mb"],
        driver_version=info["driver_version"],
    )


@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()
    info = _get_gpu_info()
    summary["gpus"] = info["gpus"]
    summary["vram_total_mb"] = info["vram_total_mb"]
    summary["vram_free_mb"] = info["vram_free_mb"]
    summary["driver_version"] = info["driver_version"]
    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))