|
|
@@ -115,6 +115,33 @@ def test_registry_module_carries_required_vram():
|
|
|
assert cands and cands[0][1].required_vram_mb == 20480
|
|
|
|
|
|
|
|
|
+def test_capacity_deduplicates_per_hypervisor():
|
|
|
+ """Nodes on the SAME hypervisor see the same passthrough GPUs -> counted once."""
|
|
|
+ reg = Registry()
|
|
|
+ # Two nodes on hypervisor "a" each see all 8 GPUs (same physical hardware).
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="a1", node_ip="10.0.0.1", hypervisor="a",
|
|
|
+ vram_total_mb=393120, vram_free_mb=200000))
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="a2", node_ip="10.0.0.2", hypervisor="a",
|
|
|
+ vram_total_mb=393120, vram_free_mb=190000))
|
|
|
+ # One node on a different hypervisor with its own single GPU.
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="g1", node_ip="10.0.1.1", hypervisor="gfx2",
|
|
|
+ vram_total_mb=24564, vram_free_mb=12000))
|
|
|
+ # Capacity: max per hypervisor, summed -> 393120 (a) + 24564 (gfx2), NOT x2.
|
|
|
+ assert reg.total_gpu_capacity_mb() == 393120 + 24564
|
|
|
+ # Free: same dedup (max per hypervisor).
|
|
|
+ assert reg.total_gpu_free_mb() == 200000 + 12000
|
|
|
+
|
|
|
+
|
|
|
+def test_capacity_counts_unknown_hypervisor_nodes_individually():
|
|
|
+ """Unknown-hypervisor nodes can't be proven to share GPUs -> counted each."""
|
|
|
+ reg = Registry()
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="u1", node_ip="10.0.0.1",
|
|
|
+ vram_total_mb=49140, vram_free_mb=20000))
|
|
|
+ reg.handle_heartbeat(Heartbeat(node_id="u2", node_ip="10.0.0.2",
|
|
|
+ vram_total_mb=24576, vram_free_mb=10000))
|
|
|
+ assert reg.total_gpu_capacity_mb() == 49140 + 24576
|
|
|
+
|
|
|
+
|
|
|
if __name__ == "__main__":
|
|
|
fns = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)]
|
|
|
failed = 0
|