mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 06:30:37 +02:00
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
"""GPU admission and cleanup contracts for application benchmark paths."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def test_benchmark_gpu_lifecycle_contracts():
|
|
script = r'''
|
|
from types import SimpleNamespace
|
|
import inspect
|
|
|
|
import app
|
|
from obliteratus.gpu_lifecycle import MemoryUsage
|
|
|
|
for entrypoint in (app.benchmark, app.benchmark_multi_model):
|
|
source = inspect.getsource(entrypoint)
|
|
assert source.index("_gpu_lifecycle.loading(model_id)") < source.index("worker.start()")
|
|
assert "result.status == \"done\"" in source
|
|
assert "_release_benchmark_pipeline(" in source
|
|
|
|
class LifecycleRecorder:
|
|
def __init__(self):
|
|
self.events = []
|
|
def resize(self, memory):
|
|
self.events.append(("resize", memory))
|
|
def release(self, *, reason):
|
|
self.events.append(("release", reason))
|
|
|
|
lifecycle = LifecycleRecorder()
|
|
handle = SimpleNamespace(model=object(), tokenizer=object())
|
|
pipeline_ref = [SimpleNamespace(handle=handle)]
|
|
calls = []
|
|
app._gpu_lifecycle = lifecycle
|
|
app.gc.collect = lambda: calls.append("gc")
|
|
app.torch.cuda.is_available = lambda: False
|
|
app.dev.empty_cache = lambda: calls.append("empty_cache")
|
|
app.measure_torch_memory = lambda _torch: MemoryUsage()
|
|
app._release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete")
|
|
assert handle.model is None and handle.tokenizer is None
|
|
assert calls == ["gc", "empty_cache"]
|
|
assert lifecycle.events == [("release", "benchmark_complete")]
|
|
|
|
lifecycle = LifecycleRecorder()
|
|
app._gpu_lifecycle = lifecycle
|
|
memory = MemoryUsage(allocated_bytes=1, reserved_bytes=2, device_count=1)
|
|
app.measure_torch_memory = lambda _torch: memory
|
|
try:
|
|
app._release_benchmark_pipeline(
|
|
[SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
|
|
reason="benchmark_complete",
|
|
)
|
|
except RuntimeError as error:
|
|
assert "retaining GPU lease" in str(error)
|
|
else:
|
|
raise AssertionError("cleanup must fail closed while CUDA allocations remain")
|
|
assert lifecycle.events == [("resize", memory)]
|
|
'''
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", script],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120,
|
|
check=False,
|
|
)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|