fix(gpu): admit benchmark loads through lifecycle broker

This commit is contained in:
Joseph Magly
2026-08-26 22:48:59 -04:00
parent 5fb64fd7e8
commit 9ab258359e
3 changed files with 136 additions and 23 deletions
+68 -23
View File
@@ -937,6 +937,26 @@ def _clear_gpu(*, release_lifecycle: bool = True):
_gpu_lifecycle.release(reason="unload")
def _release_benchmark_pipeline(pipeline_ref, *, reason: str) -> None:
"""Free a benchmark-local model before releasing supervisor ownership."""
pipeline = pipeline_ref[0]
if pipeline is not None:
if getattr(pipeline, "handle", None):
pipeline.handle.model = None
pipeline.handle.tokenizer = None
gc.collect()
if torch.cuda.is_available():
torch.cuda.synchronize()
dev.empty_cache()
memory = measure_torch_memory(torch)
if memory.allocated_bytes or memory.reserved_bytes:
_gpu_lifecycle.resize(memory)
raise RuntimeError(
"benchmark CUDA allocations remain after cleanup; retaining GPU lease"
)
_gpu_lifecycle.release(reason=reason)
def _checkpoint_is_available(checkpoint: str | None) -> bool:
"""Return whether *checkpoint* is a recoverable model directory."""
if not checkpoint:
@@ -1425,6 +1445,10 @@ def benchmark(
stage_key = result.stage
if result.status == "running":
run_logs.append(f"{stage_key.upper()}{result.message}")
elif stage_key == "summon" and result.status == "done":
memory = measure_torch_memory(torch)
_gpu_lifecycle.resize(memory)
_gpu_lifecycle.ready(memory)
quantization = load_settings.quantization
@@ -1471,9 +1495,25 @@ def benchmark(
except Exception as e:
nonlocal run_error
run_error = e
finally:
try:
_release_benchmark_pipeline(
pipeline_ref,
reason="benchmark_complete" if run_error is None else "benchmark_failed",
)
except Exception as cleanup_error:
if run_error is None:
run_error = cleanup_error
try:
_gpu_lifecycle.loading(model_id)
except Exception as error:
run_error = error
_gpu_lifecycle.release(reason="benchmark_admission_failed")
worker = threading.Thread(target=run_pipeline, daemon=True)
worker.start()
worker_started = run_error is None
if worker_started:
worker.start()
# Stream log updates while pipeline runs
last_count = len(all_logs)
@@ -1488,7 +1528,8 @@ def benchmark(
)
time.sleep(0.5)
worker.join()
if worker_started:
worker.join()
elapsed = time.time() - t_start
# Collect results
@@ -1584,15 +1625,7 @@ def benchmark(
# before the next benchmark iteration. _clear_gpu() only clears
# _state["model"], not the benchmark-local pipeline object.
if pipeline_ref[0] is not None:
try:
if hasattr(pipeline_ref[0], "handle") and pipeline_ref[0].handle:
pipeline_ref[0].handle.model = None
pipeline_ref[0].handle.tokenizer = None
except Exception:
pass
pipeline_ref[0] = None
gc.collect()
dev.empty_cache()
_release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete")
yield (
f"**{method_key} complete** ({mi + 1}/{len(methods_to_test)}) \u2014 {_bench_elapsed()}",
@@ -1789,7 +1822,10 @@ def benchmark_multi_model(
all_logs.append(f" [{_mid.split('/')[-1]}] {msg}")
def on_stage(result):
pass
if result.stage == "summon" and result.status == "done":
memory = measure_torch_memory(torch)
_gpu_lifecycle.resize(memory)
_gpu_lifecycle.ready(memory)
try:
load_settings = _resolve_ui_load_settings(
@@ -1842,9 +1878,25 @@ def benchmark_multi_model(
except Exception as e:
nonlocal run_error
run_error = e
finally:
try:
_release_benchmark_pipeline(
pipeline_ref,
reason="benchmark_complete" if run_error is None else "benchmark_failed",
)
except Exception as cleanup_error:
if run_error is None:
run_error = cleanup_error
try:
_gpu_lifecycle.loading(model_id)
except Exception as error:
run_error = error
_gpu_lifecycle.release(reason="benchmark_admission_failed")
worker = threading.Thread(target=run_pipeline, daemon=True)
worker.start()
worker_started = run_error is None
if worker_started:
worker.start()
last_count = len(all_logs)
while worker.is_alive():
@@ -1858,7 +1910,8 @@ def benchmark_multi_model(
)
time.sleep(0.5)
worker.join()
if worker_started:
worker.join()
elapsed = time.time() - t_start
entry = {
@@ -1950,15 +2003,7 @@ def benchmark_multi_model(
# Explicitly free pipeline and model before next iteration
if pipeline_ref[0] is not None:
try:
if hasattr(pipeline_ref[0], "handle") and pipeline_ref[0].handle:
pipeline_ref[0].handle.model = None
pipeline_ref[0].handle.tokenizer = None
except Exception:
pass
pipeline_ref[0] = None
gc.collect()
dev.empty_cache()
_release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete")
yield (
f"**{model_id} complete** ({mi + 1}/{len(model_choices)}) \u2014 {_mm_elapsed()}",