tests/test_gpu_monitor.py

"""Tests for the env-gated GPU monitoring thread (llm_backend.gpu_monitor)."""
import pytest


GB = 1024 ** 3


# --------------------------------------------------------------------------- #
# Env gating
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("value", ["True", "true", "TRUE", "TrUe"])
def test_is_gpu_logging_enabled_true_case_insensitive(monkeypatch, value):
    from llm_backend.gpu_monitor import is_gpu_logging_enabled
    monkeypatch.setenv("GPU_LOGGING", value)
    assert is_gpu_logging_enabled() is True


@pytest.mark.parametrize("value", ["False", "false", "1", "0", "yes", "on", "", " anything "])
def test_is_gpu_logging_enabled_false_for_other_values(monkeypatch, value):
    from llm_backend.gpu_monitor import is_gpu_logging_enabled
    monkeypatch.setenv("GPU_LOGGING", value)
    assert is_gpu_logging_enabled() is False


def test_is_gpu_logging_enabled_false_when_unset(monkeypatch):
    from llm_backend.gpu_monitor import is_gpu_logging_enabled
    monkeypatch.delenv("GPU_LOGGING", raising=False)
    assert is_gpu_logging_enabled() is False


# --------------------------------------------------------------------------- #
# _sample()
# --------------------------------------------------------------------------- #
def _patch_cuda_available(monkeypatch, available):
    import torch
    monkeypatch.setattr(torch.cuda, "is_available", lambda: available)


def test_sample_no_cuda_returns_status_line(monkeypatch):
    from llm_backend.gpu_monitor import _sample
    _patch_cuda_available(monkeypatch, False)
    line = _sample()
    assert line.startswith("[GPU_LOG]")
    assert "status=no_cuda" in line


def test_sample_with_cuda_includes_all_fields(monkeypatch):
    import torch
    from llm_backend.gpu_monitor import _sample
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
    monkeypatch.setattr(torch.cuda, "get_device_name", lambda _idx=None: "NVIDIA T4")
    monkeypatch.setattr(torch.cuda, "utilization", lambda _idx=None: 42)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (12 * GB, 16 * GB))
    monkeypatch.setattr(torch.cuda, "memory_allocated", lambda: 2 * GB)
    monkeypatch.setattr(torch.cuda, "memory_reserved", lambda: int(2.5 * GB))
    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: int(3.877 * GB))

    line = _sample()

    assert line.startswith("[GPU_LOG]")
    assert "device=cuda:0" in line
    assert "name=NVIDIA-T4" in line
    assert "util_pct=42" in line
    assert "allocated_gb=2.000" in line
    assert "reserved_gb=2.500" in line
    assert "peak_allocated_gb=3.877" in line
    assert "free_gb=12.000" in line
    assert "total_gb=16.000" in line
    assert "NVIDIA T4" not in line


def test_sample_with_cuda_util_pct_appears_right_after_name(monkeypatch):
    import torch
    from llm_backend.gpu_monitor import _sample
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
    monkeypatch.setattr(torch.cuda, "get_device_name", lambda _idx=None: "T4")
    monkeypatch.setattr(torch.cuda, "utilization", lambda _idx=None: 7)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (12 * GB, 16 * GB))
    monkeypatch.setattr(torch.cuda, "memory_allocated", lambda: 0)
    monkeypatch.setattr(torch.cuda, "memory_reserved", lambda: 0)
    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0)

    line = _sample()
    fields = line.split()
    name_idx = next(i for i, f in enumerate(fields) if f.startswith("name="))
    util_idx = next(i for i, f in enumerate(fields) if f.startswith("util_pct="))
    assert util_idx == name_idx + 1
    assert "util_pct=7" in line


def test_sample_with_cuda_utilization_error_omits_util_pct(monkeypatch):
    """When torch.cuda.utilization raises (e.g. nvidia-ml-py missing),
    the log line still has all memory fields but no util_pct."""
    import torch
    from llm_backend.gpu_monitor import _sample
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
    monkeypatch.setattr(torch.cuda, "get_device_name", lambda _idx=None: "T4")
    monkeypatch.setattr(torch.cuda, "utilization",
                        lambda _idx=None: (_ for _ in ()).throw(RuntimeError("no NVML")))
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (12 * GB, 16 * GB))
    monkeypatch.setattr(torch.cuda, "memory_allocated", lambda: 0)
    monkeypatch.setattr(torch.cuda, "memory_reserved", lambda: 0)
    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 0)

    line = _sample()
    assert "util_pct=" not in line
    assert "free_gb=12.000" in line
    assert "total_gb=16.000" in line


def test_sample_line_has_iso_timestamp(monkeypatch):
    from llm_backend.gpu_monitor import _sample
    _patch_cuda_available(monkeypatch, False)
    line = _sample()
    # Expect a token like 2026-07-06T14:23:01.234Z
    ts = line.split()[1]
    assert ts.endswith("Z")
    assert ts[4] == "-" and ts[7] == "-" and ts[10] == "T" and ts[13] == ":"


# --------------------------------------------------------------------------- #
# start_monitoring() thread behavior
# --------------------------------------------------------------------------- #
def test_start_monitoring_spawns_daemon_thread_that_exits_on_no_cuda(monkeypatch):
    from llm_backend import gpu_monitor
    _patch_cuda_available(monkeypatch, False)
    thread = gpu_monitor.start_monitoring(interval_seconds=0.01)
    assert thread.daemon is True
    # No-CUDA path prints one line and returns, so the thread should die quickly.
    thread.join(timeout=2.0)
    assert not thread.is_alive()


def test_monitor_loop_swallows_sampling_exceptions(monkeypatch, capsys):
    import torch
    from llm_backend import gpu_monitor

    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(gpu_monitor.time, "sleep", lambda s: None)

    attempts = {"n": 0}

    def flaky_sample():
        attempts["n"] += 1
        # Raise twice, then return a no_cuda line to terminate the loop cleanly.
        if attempts["n"] <= 2:
            raise RuntimeError("kaboom")
        return "[GPU_LOG] x status=no_cuda"

    monkeypatch.setattr(gpu_monitor, "_sample", flaky_sample)

    thread = gpu_monitor.start_monitoring(interval_seconds=0.01)
    thread.join(timeout=2.0)
    assert not thread.is_alive()

    captured = capsys.readouterr().out
    # Exceptions were logged as error lines, not allowed to kill the thread silently.
    assert "status=error" in captured
    assert captured.count("status=error") >= 2
    assert attempts["n"] >= 3