57 lines
2.4 KiB
Python
57 lines
2.4 KiB
Python
from pathlib import Path
|
|
|
|
from hardware_fakes import FakeNvml
|
|
|
|
from modelforge_api.domain.enums import Availability
|
|
from modelforge_api.hardware.collectors import (
|
|
NodeIdentityProvider,
|
|
NvidiaNvmlCollector,
|
|
SystemHostCollector,
|
|
)
|
|
|
|
|
|
def test_nvml_unavailable_is_graceful_and_does_not_shutdown_uninitialized() -> None:
|
|
api = FakeNvml(init_error=True)
|
|
result = NvidiaNvmlCollector(api).collect()
|
|
assert result.availability is Availability.UNAVAILABLE
|
|
assert result.inventory == []
|
|
assert api.shutdown_calls == 0
|
|
|
|
|
|
def test_zero_one_and_multiple_gpus_and_cleanup() -> None:
|
|
zero_api = FakeNvml(count=0)
|
|
assert NvidiaNvmlCollector(zero_api).collect().inventory == []
|
|
assert zero_api.shutdown_calls == 1
|
|
one = NvidiaNvmlCollector(FakeNvml(count=1)).collect()
|
|
assert one.inventory[0].device_uuid == "GPU-0"
|
|
assert one.inventory[0].total_vram_bytes.value == 16 * 1024**3
|
|
multiple = NvidiaNvmlCollector(FakeNvml(count=2)).collect()
|
|
assert [item.device_uuid for item in multiple.inventory] == ["GPU-0", "GPU-1"]
|
|
|
|
|
|
def test_optional_metric_unsupported_and_individual_device_failure_are_isolated() -> None:
|
|
result = NvidiaNvmlCollector(FakeNvml(count=1, unsupported_power=True)).collect()
|
|
assert result.telemetry[0].power_draw_w.availability is Availability.UNSUPPORTED
|
|
degraded = NvidiaNvmlCollector(FakeNvml(count=2, device_error=1)).collect()
|
|
assert degraded.availability is Availability.TEMPORARILY_FAILED
|
|
assert len(degraded.inventory) == 1
|
|
|
|
|
|
def test_host_inventory_and_persisted_node_identity(tmp_path: Path) -> None:
|
|
identity_file = tmp_path / "node-id"
|
|
provider = NodeIdentityProvider(identity_file, explicit_identity="stable-node")
|
|
collector = SystemHostCollector(provider, {"artifacts": tmp_path})
|
|
result = collector.collect()
|
|
assert result.identity_key == "stable-node"
|
|
assert result.total_ram_bytes.value and result.total_ram_bytes.value > 0
|
|
assert result.logical_cpu_count.value and result.logical_cpu_count.value > 0
|
|
assert result.storage[0].free_bytes.value is not None
|
|
|
|
|
|
def test_generated_node_identity_is_stable(monkeypatch, tmp_path: Path) -> None:
|
|
monkeypatch.setattr(NodeIdentityProvider, "_system_machine_id", staticmethod(lambda: None))
|
|
provider = NodeIdentityProvider(tmp_path / "node-id")
|
|
first = provider.resolve()
|
|
second = NodeIdentityProvider(tmp_path / "node-id").resolve()
|
|
assert first == second
|