fabric/tests/test_vram_capacity.py

"""Tests for additive GPU VRAM capacity + telemetry + capacity-aware placement.

Covers: nvidia-smi parser (capacity/free/telemetry, N/A + failure robustness),
heartbeat round-trip, registry persistence + aggregation, capacity-aware placement
tie-break, and the declared-requirement (required_vram_mb) plumbing for anti-OOM.
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,
    ModuleManifest, GpuResources,
)
from fabric.control.registry import Registry
from fabric.agent.systemd_wrap import SystemdWrapRuntime
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
    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
    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)
    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"


def test_systemd_wrap_reports_required_vram_from_manifest():
    man = ModuleManifest(name="big", resources=GpuResources(gpus=1, vram_mb=20480))
    assert SystemdWrapRuntime(man).get_status().required_vram_mb == 20480
    # default resources -> vram_mb="auto" -> reported as 0 (unknown)
    assert SystemdWrapRuntime(ModuleManifest(name="auto")).get_status().required_vram_mb == 0


def test_registry_module_carries_required_vram():
    reg = Registry()
    mod = ModuleStatus(name="m", state=ModuleState.unloaded, required_vram_mb=20480)
    reg.handle_heartbeat(Heartbeat(node_id="n", node_ip="1.1.1.1",
                                   modules=[mod], vram_total_mb=49140, vram_free_mb=5000))
    cands = reg.find_module("m")
    assert cands and cands[0][1].required_vram_mb == 20480


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)