mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 14:40:38 +02:00
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timezone
|
|
import json
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from obliteratus.gpu_lifecycle import (
|
|
GpuLifecyclePublisher,
|
|
MemoryUsage,
|
|
measure_torch_memory,
|
|
)
|
|
|
|
|
|
def test_fake_supervisor_observes_order_identity_and_recoverable_state(tmp_path):
|
|
publisher = GpuLifecyclePublisher(
|
|
tmp_path,
|
|
heartbeat_seconds=0.01,
|
|
clock=lambda: datetime(2026, 1, 2, tzinfo=timezone.utc),
|
|
run_id="run-1",
|
|
)
|
|
memory = MemoryUsage(allocated_bytes=10, reserved_bytes=12, device_count=1)
|
|
publisher.loading("org/model")
|
|
publisher.resize(memory)
|
|
publisher.ready(memory)
|
|
time.sleep(1.05)
|
|
publisher.release(reason="test_complete")
|
|
assert publisher.release(reason="duplicate") is None
|
|
|
|
events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()]
|
|
assert [event["event"] for event in events] == [
|
|
"loading", "resize", "ready", "heartbeat", "release",
|
|
]
|
|
assert [event["sequence"] for event in events] == list(range(1, 6))
|
|
assert len({event["event_id"] for event in events}) == 5
|
|
assert all(event["run_id"] == "run-1" for event in events)
|
|
assert events[1]["reserved_vram_bytes"] == 12
|
|
current = json.loads((tmp_path / "current.json").read_text())
|
|
assert current["event"] == "release"
|
|
assert current["reason"] == "test_complete"
|
|
|
|
|
|
def test_disabled_publisher_is_noop():
|
|
publisher = GpuLifecyclePublisher(None)
|
|
assert publisher.loading("model") is None
|
|
assert publisher.resize(MemoryUsage()) is None
|
|
assert publisher.ready() is None
|
|
assert publisher.heartbeat() is None
|
|
assert publisher.release() is None
|
|
|
|
|
|
def test_runtime_directory_must_exist(tmp_path):
|
|
with pytest.raises(ValueError, match="must already exist"):
|
|
GpuLifecyclePublisher(tmp_path / "missing")
|
|
|
|
|
|
def test_runtime_directory_rejects_symlink(tmp_path):
|
|
real = tmp_path / "real"
|
|
real.mkdir()
|
|
link = tmp_path / "link"
|
|
link.symlink_to(real, target_is_directory=True)
|
|
with pytest.raises(ValueError, match="must not be a symlink"):
|
|
GpuLifecyclePublisher(link)
|
|
|
|
|
|
def test_measure_torch_memory_aggregates_devices():
|
|
class FakeCuda:
|
|
is_available = staticmethod(lambda: True)
|
|
device_count = staticmethod(lambda: 2)
|
|
memory_allocated = staticmethod(lambda index: (index + 1) * 10)
|
|
memory_reserved = staticmethod(lambda index: (index + 1) * 20)
|
|
|
|
usage = measure_torch_memory(type("Torch", (), {"cuda": FakeCuda})())
|
|
assert usage == MemoryUsage(allocated_bytes=30, reserved_bytes=60, device_count=2)
|