mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 06:30:37 +02:00
93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
"""GPU admission and cleanup contracts for application benchmark paths."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from obliteratus.gpu_lifecycle import MemoryUsage
|
|
|
|
|
|
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))
|
|
|
|
|
|
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,
|
|
}
|
|
)
|
|
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,
|
|
}
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="retaining GPU lease"):
|
|
cleanup(
|
|
[SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
|
|
reason="benchmark_complete",
|
|
)
|
|
|
|
assert lifecycle.events == [("resize", memory)]
|