Files
ComfyUI/tests-unit/comfy_test/test_system_monitor.py

142 lines
3.7 KiB
Python

from types import SimpleNamespace
from comfy import system_monitor
class _Psutil:
cpu_calls = 0
@classmethod
def cpu_percent(cls):
cls.cpu_calls += 1
return 37.5
@staticmethod
def virtual_memory():
return SimpleNamespace(total=1000, available=400)
@staticmethod
def disk_partitions(all=False):
assert all is False
return [
SimpleNamespace(mountpoint="/"),
SimpleNamespace(mountpoint="/Volumes/Private\x00 Data"),
]
@staticmethod
def disk_usage(mountpoint):
return SimpleNamespace(total=2000, free=500)
class _Device:
type = "cuda"
index = 0
class _Models:
device = _Device()
@classmethod
def get_torch_device(cls):
return cls.device
@classmethod
def get_all_torch_devices(cls):
return [cls.device]
@staticmethod
def get_total_memory(device, torch_total_too=False):
assert torch_total_too
return 3000, 2500
@staticmethod
def get_free_memory(device, torch_free_too=False):
assert torch_free_too
return 1000, 900
@staticmethod
def get_torch_device_name(device):
return "Fallback GPU"
class _Nvml:
NVML_TEMPERATURE_GPU = 0
@staticmethod
def nvmlDeviceGetHandleByIndex(index):
return index
@staticmethod
def nvmlDeviceGetMemoryInfo(handle):
return SimpleNamespace(total=4000, free=1500)
@staticmethod
def nvmlDeviceGetUtilizationRates(handle):
return SimpleNamespace(gpu=61)
@staticmethod
def nvmlDeviceGetTemperature(handle, sensor):
return 72
@staticmethod
def nvmlDeviceGetName(handle):
return b"Example GPU"
def test_projection_contains_bounded_metrics_without_mount_paths():
snapshot = system_monitor._collect_snapshot(_Psutil, _Models, _Nvml)
assert snapshot == {
"cpu": {"utilization_percent": 37.5},
"memory": {"total": 1000, "available": 400},
"volumes": [
{"id": "volume-0", "label": "Root", "total": 2000, "available": 500},
{"id": "volume-1", "label": "Private Data", "total": 2000, "available": 500},
],
"accelerators": [{
"id": "accelerator-0",
"name": "Example GPU",
"memory_total": 4000,
"memory_available": 1500,
"utilization_percent": 61.0,
"temperature_c": 72.0,
}],
}
assert "/Volumes" not in repr(snapshot)
def test_snapshot_is_cached_and_returned_as_an_independent_value(monkeypatch):
times = iter([1.0, 1.1, 1.3])
monkeypatch.setattr(system_monitor.time, "monotonic", lambda: next(times))
monkeypatch.setattr(
system_monitor,
"_collect_snapshot",
lambda: {"sample": _Psutil.cpu_percent()},
)
monkeypatch.setattr(system_monitor, "_cached_at", -1.0)
monkeypatch.setattr(system_monitor, "_cached_snapshot", None)
_Psutil.cpu_calls = 0
first = system_monitor.get_system_monitor_snapshot()
first["sample"] = -1
second = system_monitor.get_system_monitor_snapshot()
third = system_monitor.get_system_monitor_snapshot()
assert second == {"sample": 37.5}
assert third == {"sample": 37.5}
assert _Psutil.cpu_calls == 2
def test_missing_optional_sensors_are_null():
snapshot = system_monitor._collect_snapshot(_Psutil, _Models, None)
accelerator = snapshot["accelerators"][0]
assert accelerator["name"] == "Fallback GPU"
assert accelerator["utilization_percent"] is None
assert accelerator["temperature_c"] is None
def test_byte_totals_fit_the_cross_language_safe_integer_contract():
assert system_monitor._bytes(2**63) == 2**53 - 1
assert system_monitor._bytes(-1) == 0