feat(fabric): free VRAM + GPU telemetry + capacity-aware placement

- Agent /gpu-status + heartbeat: per-GPU name, mem_free_mb, util_pct, temp_c,
  power_w; node-level vram_free_mb + driver_version (extended nvidia-smi query).
- Schema: consolidate capacity/telemetry into Heartbeat/NodeInfo/GpuMemInfo
  (retire the control-plane _node_capacity side-map; restore response_model).
- CP /registry/models: add total_gpu_free_mb (alongside total_gpu_capacity_mb);
  /registry/nodes + /registry/hypervisors carry vram_total_mb / vram_free_mb.
- Scheduler: registry.place() tie-breaks equal-state candidates by most free
  VRAM; /place anti-OOM guard (503 when the chosen node lacks the model's known
  footprint).
- tests/test_vram_capacity.py: parser, N/A+failure, heartbeat round-trip,
  registry aggregation, place tie-break (5/5 pass; runnable via python or pytest).

Additive and backward-compatible. Deployed fleet-wide (12 nodes).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

agent/main.py

@@ -9,7 +9,6 @@ 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
@@ -26,22 +25,6 @@ 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=["*"])
 
@@ -62,38 +45,65 @@ def _detect_ip() -> str:
     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 []."""
+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,memory.total,memory.used",
+            ["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 capacity query failed: %s", e)
-        return []
+        log.warning("nvidia-smi query failed: %s", e)
+        return empty
     if r.returncode != 0:
-        return []
+        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) < 2:
+        if len(parts) < 4:
             continue
-        try:
-            idx = int(parts[0])
-            total = int(float(parts[1]))
-        except ValueError:
+        idx = _to_int(parts[0])
+        if idx is None:
             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
+        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():
@@ -117,20 +127,22 @@ def _get_pid_vram_map() -> dict[int, int]:
     return m
 
 
-def _build_heartbeat() -> HeartbeatExt:
+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())
-    gpus = _get_gpu_mem_info()
-    return HeartbeatExt(
+    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=gpus,
-        vram_total_mb=sum(g["mem_total_mb"] for g in gpus),
+        gpus=info["gpus"],
+        vram_total_mb=info["vram_total_mb"],
+        vram_free_mb=info["vram_free_mb"],
+        driver_version=info["driver_version"],
     )
 
 
@@ -147,9 +159,11 @@ async def list_modules():
 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)
+    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}")

control/main.py

@@ -6,7 +6,6 @@ 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,
@@ -28,25 +27,6 @@ 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():
@@ -60,40 +40,24 @@ async def health():
     }
 
 @app.post("/registry/heartbeat")
-async def heartbeat(hb: HeartbeatExt):
+async def heartbeat(hb: Heartbeat):
     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")
+@app.get("/registry/models", response_model=RegistryModelsResponse)
 async def list_models():
     modules = registry.get_all_modules()
-    alive = registry.get_alive_nodes()
-    base = RegistryModelsResponse(
+    return RegistryModelsResponse(
         models=modules,
-        nodes_count=len(alive),
+        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(),
     )
-    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")
+@app.get("/registry/nodes", response_model=RegistryNodesResponse)
 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}
+    return RegistryNodesResponse(nodes=registry.get_alive_nodes())
 
 @app.get("/registry/hypervisors")
 async def list_hypervisors():
@@ -107,6 +71,8 @@ async def list_hypervisors():
             "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)}
 
@@ -116,6 +82,19 @@ async def place(req: PlaceRequest):
     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
 

control/registry.py

@@ -0,0 +1,118 @@
+"""In-memory registry for nodes and modules."""
+from __future__ import annotations
+import time
+import logging
+
+from fabric.shared.schemas import (
+    Heartbeat, NodeInfo, ModuleStatus, ModuleState,
+    PlaceRequest, PlaceResponse, PriorityClass,
+)
+
+log = logging.getLogger("fabric.registry")
+
+STALE_NODE_TIMEOUT_S = 60
+
+
+class Registry:
+    def __init__(self):
+        self.nodes: dict[str, NodeInfo] = {}
+        self._version = 0
+
+    @property
+    def version(self) -> int:
+        return self._version
+
+    def handle_heartbeat(self, hb: Heartbeat) -> None:
+        node = self.nodes.get(hb.node_id)
+        if node is None:
+            log.info("New node registered: %s (%s)", hb.node_id, hb.node_ip)
+        self.nodes[hb.node_id] = NodeInfo(
+            node_id=hb.node_id,
+            node_ip=hb.node_ip,
+            agent_port=hb.agent_port,
+            hypervisor=hb.hypervisor,
+            gpu_status=hb.gpu_status,
+            modules=hb.modules,
+            last_heartbeat=time.time(),
+            alive=True,
+            gpus=hb.gpus,
+            vram_total_mb=hb.vram_total_mb,
+            vram_free_mb=hb.vram_free_mb,
+            driver_version=hb.driver_version,
+        )
+        self._version += 1
+
+    def prune_stale_nodes(self) -> list[str]:
+        now = time.time()
+        stale = []
+        for nid, node in list(self.nodes.items()):
+            if now - node.last_heartbeat > STALE_NODE_TIMEOUT_S:
+                node.alive = False
+                stale.append(nid)
+        return stale
+
+    def get_all_modules(self) -> list[ModuleStatus]:
+        modules = []
+        for node in self.nodes.values():
+            if not node.alive:
+                continue
+            for m in node.modules:
+                modules.append(m)
+        return modules
+
+    def get_alive_nodes(self) -> list[NodeInfo]:
+        return [n for n in self.nodes.values() if n.alive]
+
+    def find_module(self, model_name: str) -> list[tuple[NodeInfo, ModuleStatus]]:
+        results = []
+        for node in self.nodes.values():
+            if not node.alive:
+                continue
+            for m in node.modules:
+                if m.name == model_name or model_name in m.aliases:
+                    results.append((node, m))
+        return results
+
+    def place(self, req: PlaceRequest) -> PlaceResponse | None:
+        candidates = self.find_module(req.model)
+        if not candidates:
+            return None
+
+        state_priority = {
+            ModuleState.running: 0,
+            ModuleState.loaded: 1,
+            ModuleState.idle: 2,
+            ModuleState.sleeping: 3,
+            ModuleState.warming: 4,
+            ModuleState.unloaded: 5,
+            ModuleState.error: 6,
+        }
+        # Capacity-aware tie-break: among equally-good states, prefer the node
+        # with the most free physical VRAM (more headroom => safer to (re)load).
+        candidates.sort(key=lambda c: (state_priority.get(c[1].state, 99), -(c[0].vram_free_mb or 0)))
+
+        node, module = candidates[0]
+        return PlaceResponse(
+            node_id=node.node_id,
+            node_ip=node.node_ip,
+            infer_url=module.infer_url,
+            module_name=module.name,
+            state=module.state,
+        )
+
+    def total_gpu_vram_mb(self) -> int:
+        total = 0
+        for node in self.nodes.values():
+            if not node.alive:
+                continue
+            for g in node.gpu_status:
+                total += g.vram_mb
+        return total
+
+    def total_gpu_capacity_mb(self) -> int:
+        """Sum of physical-GPU total VRAM across alive nodes (hardware capacity)."""
+        return sum(n.vram_total_mb for n in self.nodes.values() if n.alive)
+
+    def total_gpu_free_mb(self) -> int:
+        """Sum of physical-GPU free VRAM across alive nodes (placement headroom)."""
+        return sum(n.vram_free_mb for n in self.nodes.values() if n.alive)

shared/schemas.py

@@ -0,0 +1,175 @@
+"""Shared Pydantic schemas for Fabric control plane and node-agent."""
+from __future__ import annotations
+from enum import Enum
+from pydantic import BaseModel, Field
+from typing import Optional
+import time
+
+
+class PriorityClass(str, Enum):
+    low = "low"
+    normal = "normal"
+    high = "high"
+    critical = "critical"
+
+
+class ReclaimStrategy(str, Enum):
+    process_kill = "process-kill"
+    vllm_sleep = "vllm-sleep"
+    cuda_checkpoint = "cuda-checkpoint"
+    none = "none"
+
+
+class ModuleState(str, Enum):
+    loaded = "loaded"
+    running = "running"
+    idle = "idle"
+    sleeping = "sleeping"
+    unloaded = "unloaded"
+    error = "error"
+    warming = "warming"
+
+
+class RuntimeType(str, Enum):
+    systemd_wrap = "systemd-wrap"
+    subprocess = "subprocess"
+    vllm = "vllm"
+    python_class = "python-class"
+    comfy_graph = "comfy-graph"
+    proc_exec = "proc-exec"
+
+
+class GpuResources(BaseModel):
+    gpus: int = 1
+    vram_mb: int | str = "auto"
+
+
+class LifecycleConfig(BaseModel):
+    reclaim: ReclaimStrategy = ReclaimStrategy.process_kill
+    idle_evict_s: int = 300
+    warm_default: bool = False
+    idle_evict_s: int = 300
+    warm_default: bool = False
+    warmup_s: int = 10
+
+
+class SchedulingConfig(BaseModel):
+    priority_class: PriorityClass = PriorityClass.normal
+    min_replicas: int = 0
+    max_replicas: int = 8
+    batch: bool = False
+
+
+class SystemdWrapConfig(BaseModel):
+    systemd_unit: str
+    health_url: str | None = None
+    health_timeout_s: int = 120
+    infer_url: str
+    infer_method: str = "POST"
+
+
+class ModuleManifest(BaseModel):
+    name: str
+    description: str = ""
+    version: str = "1.0"
+    modality: str = "general"
+    aliases: list[str] = Field(default_factory=list)
+    runtime: RuntimeType = RuntimeType.systemd_wrap
+    resources: GpuResources = Field(default_factory=GpuResources)
+    lifecycle: LifecycleConfig = Field(default_factory=LifecycleConfig)
+    scheduling: SchedulingConfig = Field(default_factory=SchedulingConfig)
+    systemd: SystemdWrapConfig | None = None
+    capabilities: list[str] = Field(default_factory=list)
+
+
+class GpuPidUsage(BaseModel):
+    gpu_index: int
+    pid: int
+    vram_mb: int
+    process_name: str = ""
+
+
+class GpuMemInfo(BaseModel):
+    """Per-physical-GPU memory capacity + telemetry (additive). One entry per
+    physical GPU (counted once, NOT per process). Absent fields degrade to
+    0 / None so old and new components interoperate."""
+    gpu_index: int = 0
+    name: str = ""
+    mem_total_mb: int = 0
+    mem_used_mb: int = 0
+    mem_free_mb: int = 0
+    util_pct: Optional[int] = None
+    temp_c: Optional[int] = None
+    power_w: Optional[float] = None
+
+
+class ModuleStatus(BaseModel):
+    name: str
+    state: ModuleState
+    vram_mb: int = 0
+    last_infer_at: float | None = None
+    health_ok: bool = False
+    infer_url: str = ""
+    modality: str = "general"
+    aliases: list[str] = Field(default_factory=list)
+    priority_class: PriorityClass = PriorityClass.normal
+    reclaim: ReclaimStrategy = ReclaimStrategy.process_kill
+    idle_evict_s: int = 300
+    warm_default: bool = False
+
+
+class Heartbeat(BaseModel):
+    node_id: str
+    node_ip: str
+    agent_port: int = 7701
+    hypervisor: str = "unknown"
+    gpu_status: list[GpuPidUsage] = Field(default_factory=list)
+    modules: list[ModuleStatus] = Field(default_factory=list)
+    uptime_s: int = 0
+    timestamp: float = Field(default_factory=time.time)
+    # --- additive VRAM-capacity + telemetry (per physical GPU) ---
+    gpus: list[GpuMemInfo] = Field(default_factory=list)
+    vram_total_mb: int = 0
+    vram_free_mb: int = 0
+    driver_version: str = ""
+
+
+class NodeInfo(BaseModel):
+    node_id: str
+    node_ip: str
+    agent_port: int = 7701
+    hypervisor: str = "unknown"
+    gpu_status: list[GpuPidUsage] = Field(default_factory=list)
+    modules: list[ModuleStatus] = Field(default_factory=list)
+    last_heartbeat: float = 0
+    alive: bool = True
+    # --- additive VRAM-capacity + telemetry (per physical GPU) ---
+    gpus: list[GpuMemInfo] = Field(default_factory=list)
+    vram_total_mb: int = 0
+    vram_free_mb: int = 0
+    driver_version: str = ""
+
+
+class PlaceRequest(BaseModel):
+    model: str
+    priority: PriorityClass = PriorityClass.normal
+
+
+class PlaceResponse(BaseModel):
+    node_id: str
+    node_ip: str
+    infer_url: str
+    module_name: str
+    state: ModuleState
+
+
+class RegistryModelsResponse(BaseModel):
+    models: list[ModuleStatus]
+    nodes_count: int
+    total_gpu_vram_mb: int = 0
+    total_gpu_capacity_mb: int = 0
+    total_gpu_free_mb: int = 0
+
+
+class RegistryNodesResponse(BaseModel):
+    nodes: list[NodeInfo]

tests/test_vram_capacity.py

@@ -0,0 +1,112 @@
+"""Tests for additive GPU VRAM capacity + telemetry reporting.
+
+Covers: nvidia-smi parser (capacity/free/telemetry, N/A + failure robustness),
+heartbeat round-trip, registry persistence + aggregation, and the capacity-aware
+placement tie-break. Runs under pytest OR standalone: `python tests/test_vram_capacity.py`.
+"""
+import sys
+from unittest import mock
+
+sys.path.insert(0, "/srv")  # make `import fabric.*` resolve when run directly
+
+from fabric.shared.schemas import (
+    Heartbeat, GpuMemInfo, ModuleStatus, ModuleState, PlaceRequest,
+)
+from fabric.control.registry import Registry
+import fabric.agent.main as agent
+
+
+FAKE_SMI = (
+    "0, NVIDIA L40S, 49140, 29894, 45, 62, 210.50, 550.54.15\n"
+    "1, NVIDIA L40S, 49140, 1024, 0, 38, 90.00, 550.54.15\n"
+)
+
+
+def test_get_gpu_info_parses_capacity_and_telemetry():
+    with mock.patch.object(agent.subprocess, "run",
+                           lambda *a, **k: mock.Mock(returncode=0, stdout=FAKE_SMI)):
+        info = agent._get_gpu_info()
+    assert len(info["gpus"]) == 2
+    g0 = info["gpus"][0]
+    assert g0["gpu_index"] == 0
+    assert g0["name"] == "NVIDIA L40S"
+    assert g0["mem_total_mb"] == 49140
+    assert g0["mem_used_mb"] == 29894
+    assert g0["mem_free_mb"] == 49140 - 29894
+    assert g0["util_pct"] == 45
+    assert g0["temp_c"] == 62
+    assert g0["power_w"] == 210.50
+    assert info["driver_version"] == "550.54.15"
+    assert info["vram_total_mb"] == 49140 * 2
+    assert info["vram_free_mb"] == (49140 - 29894) + (49140 - 1024)
+
+
+def test_get_gpu_info_handles_na_and_failure():
+    na = "0, GPU, 24576, 100, [N/A], [N/A], [N/A], 550.0\n"
+    with mock.patch.object(agent.subprocess, "run",
+                           lambda *a, **k: mock.Mock(returncode=0, stdout=na)):
+        info = agent._get_gpu_info()
+    g = info["gpus"][0]
+    assert g["mem_free_mb"] == 24576 - 100
+    assert g["util_pct"] is None and g["temp_c"] is None and g["power_w"] is None
+    # non-zero return code -> zeros, no crash
+    with mock.patch.object(agent.subprocess, "run",
+                           lambda *a, **k: mock.Mock(returncode=9, stdout="")):
+        empty = agent._get_gpu_info()
+    assert empty["gpus"] == [] and empty["vram_total_mb"] == 0
+    # nvidia-smi missing entirely -> zeros, no crash
+    with mock.patch.object(agent.subprocess, "run",
+                           mock.Mock(side_effect=FileNotFoundError())):
+        empty2 = agent._get_gpu_info()
+    assert empty2["vram_total_mb"] == 0 and empty2["vram_free_mb"] == 0
+
+
+def test_heartbeat_roundtrip_carries_capacity():
+    hb = Heartbeat(
+        node_id="n1", node_ip="1.2.3.4",
+        gpus=[{"gpu_index": 0, "mem_total_mb": 49140, "mem_used_mb": 140, "mem_free_mb": 49000}],
+        vram_total_mb=49140, vram_free_mb=49000, driver_version="550.0",
+    )
+    d = hb.model_dump()
+    assert d["vram_total_mb"] == 49140 and d["vram_free_mb"] == 49000
+    assert d["gpus"][0]["mem_free_mb"] == 49000
+    hb2 = Heartbeat(**d)  # what the CP does when parsing the request body
+    assert hb2.vram_total_mb == 49140
+    assert isinstance(hb2.gpus[0], GpuMemInfo)
+
+
+def test_registry_carries_and_aggregates_capacity():
+    reg = Registry()
+    reg.handle_heartbeat(Heartbeat(node_id="a", node_ip="10.0.0.1",
+                                   vram_total_mb=49140, vram_free_mb=20000))
+    reg.handle_heartbeat(Heartbeat(node_id="b", node_ip="10.0.0.2",
+                                   vram_total_mb=24576, vram_free_mb=10000))
+    nodes = {n.node_id: n for n in reg.get_alive_nodes()}
+    assert nodes["a"].vram_total_mb == 49140 and nodes["a"].vram_free_mb == 20000
+    assert reg.total_gpu_capacity_mb() == 49140 + 24576
+    assert reg.total_gpu_free_mb() == 20000 + 10000
+
+
+def test_place_prefers_node_with_more_free_vram():
+    reg = Registry()
+    mod = ModuleStatus(name="m", state=ModuleState.loaded, infer_url="http://x")
+    reg.handle_heartbeat(Heartbeat(node_id="low", node_ip="10.0.0.1",
+                                   modules=[mod], vram_total_mb=49140, vram_free_mb=5000))
+    reg.handle_heartbeat(Heartbeat(node_id="high", node_ip="10.0.0.2",
+                                   modules=[mod], vram_total_mb=49140, vram_free_mb=40000))
+    res = reg.place(PlaceRequest(model="m"))
+    assert res is not None and res.node_id == "high"  # equal state -> most-free wins
+
+
+if __name__ == "__main__":
+    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]
+    failed = 0
+    for fn in fns:
+        try:
+            fn()
+            print("PASS", fn.__name__)
+        except Exception as e:
+            failed += 1
+            print("FAIL", fn.__name__, "->", repr(e))
+    print("\n%d/%d passed" % (len(fns) - failed, len(fns)))
+    sys.exit(1 if failed else 0)