Spaces:
Running on Zero
Running on Zero
File size: 6,784 Bytes
d4620ae e9842fb d4620ae e9842fb d4620ae e9842fb b210716 d4620ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """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
|