fabric/tests/test_vram_capacity.py

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