|
@@ -1,8 +1,9 @@
|
|
|
-"""Tests for additive GPU VRAM capacity + telemetry reporting.
|
|
|
|
|
|
|
+"""Tests for additive GPU VRAM capacity + telemetry + capacity-aware placement.
|
|
|
|
|
|
|
|
Covers: nvidia-smi parser (capacity/free/telemetry, N/A + failure robustness),
|
|
Covers: nvidia-smi parser (capacity/free/telemetry, N/A + failure robustness),
|
|
|
-heartbeat round-trip, registry persistence + aggregation, and the capacity-aware
|
|
|
|
|
-placement tie-break. Runs under pytest OR standalone: `python tests/test_vram_capacity.py`.
|
|
|
|
|
|
|
+heartbeat round-trip, registry persistence + aggregation, capacity-aware placement
|
|
|
|
|
+tie-break, and the declared-requirement (required_vram_mb) plumbing for anti-OOM.
|
|
|
|
|
+Runs under pytest OR standalone: `python tests/test_vram_capacity.py`.
|
|
|
"""
|
|
"""
|
|
|
import sys
|
|
import sys
|
|
|
from unittest import mock
|
|
from unittest import mock
|
|
@@ -11,8 +12,10 @@ sys.path.insert(0, "/srv") # make `import fabric.*` resolve when run directly
|
|
|
|
|
|
|
|
from fabric.shared.schemas import (
|
|
from fabric.shared.schemas import (
|
|
|
Heartbeat, GpuMemInfo, ModuleStatus, ModuleState, PlaceRequest,
|
|
Heartbeat, GpuMemInfo, ModuleStatus, ModuleState, PlaceRequest,
|
|
|
|
|
+ ModuleManifest, GpuResources,
|
|
|
)
|
|
)
|
|
|
from fabric.control.registry import Registry
|
|
from fabric.control.registry import Registry
|
|
|
|
|
+from fabric.agent.systemd_wrap import SystemdWrapRuntime
|
|
|
import fabric.agent.main as agent
|
|
import fabric.agent.main as agent
|
|
|
|
|
|
|
|
|
|
|
|
@@ -49,12 +52,10 @@ def test_get_gpu_info_handles_na_and_failure():
|
|
|
g = info["gpus"][0]
|
|
g = info["gpus"][0]
|
|
|
assert g["mem_free_mb"] == 24576 - 100
|
|
assert g["mem_free_mb"] == 24576 - 100
|
|
|
assert g["util_pct"] is None and g["temp_c"] is None and g["power_w"] is None
|
|
assert g["util_pct"] is None and g["temp_c"] is None and g["power_w"] is None
|
|
|
- # non-zero return code -> zeros, no crash
|
|
|
|
|
with mock.patch.object(agent.subprocess, "run",
|
|
with mock.patch.object(agent.subprocess, "run",
|
|
|
lambda *a, **k: mock.Mock(returncode=9, stdout="")):
|
|
lambda *a, **k: mock.Mock(returncode=9, stdout="")):
|
|
|
empty = agent._get_gpu_info()
|
|
empty = agent._get_gpu_info()
|
|
|
assert empty["gpus"] == [] and empty["vram_total_mb"] == 0
|
|
assert empty["gpus"] == [] and empty["vram_total_mb"] == 0
|
|
|
- # nvidia-smi missing entirely -> zeros, no crash
|
|
|
|
|
with mock.patch.object(agent.subprocess, "run",
|
|
with mock.patch.object(agent.subprocess, "run",
|
|
|
mock.Mock(side_effect=FileNotFoundError())):
|
|
mock.Mock(side_effect=FileNotFoundError())):
|
|
|
empty2 = agent._get_gpu_info()
|
|
empty2 = agent._get_gpu_info()
|
|
@@ -70,7 +71,7 @@ def test_heartbeat_roundtrip_carries_capacity():
|
|
|
d = hb.model_dump()
|
|
d = hb.model_dump()
|
|
|
assert d["vram_total_mb"] == 49140 and d["vram_free_mb"] == 49000
|
|
assert d["vram_total_mb"] == 49140 and d["vram_free_mb"] == 49000
|
|
|
assert d["gpus"][0]["mem_free_mb"] == 49000
|
|
assert d["gpus"][0]["mem_free_mb"] == 49000
|
|
|
- hb2 = Heartbeat(**d) # what the CP does when parsing the request body
|
|
|
|
|
|
|
+ hb2 = Heartbeat(**d)
|
|
|
assert hb2.vram_total_mb == 49140
|
|
assert hb2.vram_total_mb == 49140
|
|
|
assert isinstance(hb2.gpus[0], GpuMemInfo)
|
|
assert isinstance(hb2.gpus[0], GpuMemInfo)
|
|
|
|
|
|
|
@@ -95,7 +96,23 @@ def test_place_prefers_node_with_more_free_vram():
|
|
|
reg.handle_heartbeat(Heartbeat(node_id="high", node_ip="10.0.0.2",
|
|
reg.handle_heartbeat(Heartbeat(node_id="high", node_ip="10.0.0.2",
|
|
|
modules=[mod], vram_total_mb=49140, vram_free_mb=40000))
|
|
modules=[mod], vram_total_mb=49140, vram_free_mb=40000))
|
|
|
res = reg.place(PlaceRequest(model="m"))
|
|
res = reg.place(PlaceRequest(model="m"))
|
|
|
- assert res is not None and res.node_id == "high" # equal state -> most-free wins
|
|
|
|
|
|
|
+ assert res is not None and res.node_id == "high"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_systemd_wrap_reports_required_vram_from_manifest():
|
|
|
|
|
+ man = ModuleManifest(name="big", resources=GpuResources(gpus=1, vram_mb=20480))
|
|
|
|
|
+ assert SystemdWrapRuntime(man).get_status().required_vram_mb == 20480
|
|
|
|
|
+ # default resources -> vram_mb="auto" -> reported as 0 (unknown)
|
|
|
|
|
+ assert SystemdWrapRuntime(ModuleManifest(name="auto")).get_status().required_vram_mb == 0
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_registry_module_carries_required_vram():
|
|
|
|
|
+ reg = Registry()
|
|
|
|
|
+ mod = ModuleStatus(name="m", state=ModuleState.unloaded, required_vram_mb=20480)
|
|
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="n", node_ip="1.1.1.1",
|
|
|
|
|
+ modules=[mod], vram_total_mb=49140, vram_free_mb=5000))
|
|
|
|
|
+ cands = reg.find_module("m")
|
|
|
|
|
+ assert cands and cands[0][1].required_vram_mb == 20480
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if __name__ == "__main__":
|