test(gpu): keep benchmark contracts CPU-only

This commit is contained in:
Joseph Magly
2026-08-26 22:51:13 -04:00
parent 9ab258359e
commit 6df26d07e7
+73 -48
View File
@@ -2,66 +2,91 @@
from __future__ import annotations
import subprocess
import sys
def test_benchmark_gpu_lifecycle_contracts():
script = r'''
import ast
from pathlib import Path
from types import SimpleNamespace
import inspect
import app
import pytest
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:
APP_SOURCE = Path("app.py").read_text(encoding="utf-8")
APP_TREE = ast.parse(APP_SOURCE)
def _function_node(name: str) -> ast.FunctionDef:
return next(
node
for node in APP_TREE.body
if isinstance(node, ast.FunctionDef) and node.name == name
)
def _load_cleanup_function(namespace: dict):
node = _function_node("_release_benchmark_pipeline")
module = ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[]))
exec(compile(module, "app.py", "exec"), namespace)
return namespace["_release_benchmark_pipeline"]
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",
def test_benchmark_entrypoints_admit_before_worker_start():
for name in ("benchmark", "benchmark_multi_model"):
source = ast.get_source_segment(APP_SOURCE, _function_node(name))
assert source.index("_gpu_lifecycle.loading(model_id)") < source.index("worker.start()")
assert "result.status == \"done\"" in source
assert "_release_benchmark_pipeline(" in source
def test_benchmark_cleanup_releases_only_after_cuda_is_gone():
lifecycle = _LifecycleRecorder()
calls = []
cleanup = _load_cleanup_function(
{
"gc": SimpleNamespace(collect=lambda: calls.append("gc")),
"torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
"dev": SimpleNamespace(empty_cache=lambda: calls.append("empty_cache")),
"measure_torch_memory": lambda _torch: MemoryUsage(),
"_gpu_lifecycle": lifecycle,
}
)
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,
handle = SimpleNamespace(model=object(), tokenizer=object())
cleanup([SimpleNamespace(handle=handle)], reason="benchmark_complete")
assert handle.model is None and handle.tokenizer is None
assert calls == ["gc", "empty_cache"]
assert lifecycle.events == [("release", "benchmark_complete")]
def test_benchmark_cleanup_retains_lease_when_cuda_remains():
lifecycle = _LifecycleRecorder()
memory = MemoryUsage(allocated_bytes=1, reserved_bytes=2, device_count=1)
cleanup = _load_cleanup_function(
{
"gc": SimpleNamespace(collect=lambda: None),
"torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
"dev": SimpleNamespace(empty_cache=lambda: None),
"measure_torch_memory": lambda _torch: memory,
"_gpu_lifecycle": lifecycle,
}
)
assert result.returncode == 0, result.stdout + result.stderr
with pytest.raises(RuntimeError, match="retaining GPU lease"):
cleanup(
[SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
reason="benchmark_complete",
)
assert lifecycle.events == [("resize", memory)]