modelbeast/tests/test_node_reporting.py

"""Node-placement reporting: gpu_pool must reflect live occupancy.

The bug this guards: gpu_pool() read a `busy` key that load_gpu_nodes() set to
False once and nothing ever updated, so every node reported idle while running
jobs. Occupancy must come from the runner's `inflight` counters.

Run: python3 tests/test_node_reporting.py
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from server import sysinfo


class FakeRunner:
    """Minimal stand-in — gpu_pool only needs gpu_nodes + _node_capacity."""

    def __init__(self, nodes):
        self.gpu_nodes = nodes

    def _node_capacity(self, node, lane):
        if lane == "gpu":
            return 1
        if lane == "cpu":
            return int(node.get("cpu_slots", 3 if node.get("local") else 2))
        return 0


def test_busy_reflects_inflight():
    runner = FakeRunner([
        {"name": "m3ultra", "local": True, "busy": False,
         "inflight": {"gpu": 1, "cpu": 0}},
        {"name": "m1", "ssh": "x@y", "busy": False,
         "inflight": {"gpu": 0, "cpu": 0}, "operators": ["flux_local"]},
    ])
    pool = {n["name"]: n for n in sysinfo.gpu_pool(runner)}

    # the stale `busy: False` on the node dict must NOT win
    assert pool["m3ultra"]["busy"] is True, "node running a gpu job reported idle"
    assert pool["m3ultra"]["gpu"] == {"running": 1, "slots": 1, "free": 0}
    assert pool["m1"]["busy"] is False
    assert pool["m1"]["gpu"] == {"running": 0, "slots": 1, "free": 1}


def test_local_is_not_named_local():
    """load_gpu_nodes names the local node after the host, so the UI can tell
    the primary apart from the placeholder it used to show."""
    from server import remote
    nodes = remote.load_gpu_nodes()
    assert nodes[0]["local"] is True
    assert nodes[0]["name"] != "local", "local node should carry the hostname"


def test_operator_count_reported():
    runner = FakeRunner([
        {"name": "all-ops", "local": True, "inflight": {"gpu": 0, "cpu": 0}},
        {"name": "fenced", "ssh": "x@y", "inflight": {"gpu": 0, "cpu": 0},
         "operators": ["a", "b", "c"]},
    ])
    pool = {n["name"]: n for n in sysinfo.gpu_pool(runner)}
    assert pool["all-ops"]["operators"] == "all"   # no allowlist == allow-all
    assert pool["fenced"]["operators"] == 3


def test_node_column_migrated():
    from server import db
    assert ("jobs", "node", "TEXT") in db.MIGRATIONS, "jobs.node migration missing"




# --- nodes_detail: the /api/nodes contract fluxgod consumes -------------------

class FakeCon:
    def __init__(self, rows): self._rows = rows
    def execute(self, *a, **k): return self
    def fetchall(self): return self._rows


def _runner_with_ops():
    r = FakeRunner([
        {"name": "m3ultra", "local": True, "inflight": {"gpu": 1, "cpu": 0}},
        {"name": "m1", "ssh": "u@h", "inflight": {"gpu": 1, "cpu": 0},
         "operators": ["mflux_image_edit", "flux_local"], "_notes": ["a note"]},
        {"name": "m2max", "ssh": "u@h2", "inflight": {"gpu": 0, "cpu": 0},
         "operators": ["flux_local"]},
    ])
    r.operators = {"mflux_image_edit": {"resources": "gpu"},
                   "flux_local": {"resources": "gpu"},
                   "ffprobe": {"resources": "cpu"}}
    r.lane_of = lambda op: r.operators.get(op, {}).get("resources", "cpu")
    return r


def test_running_attributed_by_node_column():
    """The whole point: a job that ran on m1 must show under m1, not the primary."""
    rows = [
        {"id": "j1", "operator": "mflux_image_edit", "status": "running",
         "node": "m1", "created_at": 1.0, "started_at": 1.1},
        {"id": "j2", "operator": "mflux_image_edit", "status": "running",
         "node": "m3ultra", "created_at": 2.0, "started_at": 2.1},
    ]
    d = sysinfo.nodes_detail(FakeCon(rows), _runner_with_ops())
    byname = {n["name"]: n for n in d["nodes"]}
    assert [j["id"] for j in byname["m1"]["running"]] == ["j1"], "m1's job misattributed"
    assert [j["id"] for j in byname["m3ultra"]["running"]] == ["j2"]
    assert byname["m2max"]["running"] == []


def test_legacy_null_node_credited_to_primary():
    rows = [{"id": "old", "operator": "flux_local", "status": "running",
             "node": None, "created_at": 1.0, "started_at": 1.1}]
    d = sysinfo.nodes_detail(FakeCon(rows), _runner_with_ops())
    byname = {n["name"]: n for n in d["nodes"]}
    assert [j["id"] for j in byname["m3ultra"]["running"]] == ["old"]


def test_eligible_queued_respects_allowlist():
    rows = [{"id": "q1", "operator": "mflux_image_edit", "status": "queued",
             "node": None, "created_at": 3.0, "started_at": None}]
    d = sysinfo.nodes_detail(FakeCon(rows), _runner_with_ops())
    byname = {n["name"]: n for n in d["nodes"]}
    assert len(byname["m1"]["eligible_queued"]) == 1      # allowlists it
    assert byname["m2max"]["eligible_queued"] == []       # does not
    assert len(byname["m3ultra"]["eligible_queued"]) == 1  # no allowlist == all


def test_contract_keys_match_captured_api():
    """Field names fluxgod's server.py reads — captured from the live API."""
    d = sysinfo.nodes_detail(FakeCon([]), _runner_with_ops())
    assert set(d) == {"nodes", "by_operator", "queued_unassigned", "gpu_ops"}
    assert set(d["nodes"][0]) == {
        "name", "remote", "ssh", "gpu_capacity", "cpu_capacity", "gpu_inflight",
        "cpu_inflight", "operators_explicit", "operators", "operator_count",
        "running", "eligible_queued", "notes"}
    assert d["by_operator"]["flux_local"] == ["m3ultra", "m1", "m2max"]
    assert d["gpu_ops"] == ["flux_local", "mflux_image_edit"]


if __name__ == "__main__":
    for name, fn in sorted(globals().items()):
        if name.startswith("test_") and callable(fn):
            fn()
            print(f"  ok  {name}")
    print("all node-reporting checks passed")