From c072c75fa97300d12e4ee295890180b9bdc279b0 Mon Sep 17 00:00:00 2001
From: Frank <frank@dankware>
Date: Tue, 25 Aug 2026 17:11:42 -0700
Subject: [PATCH] Make ubatch benchmark cleanup failure safe

---
 scripts/bench-ubatch-rocmi4.sh | 19 ++++++++++++++++---
 tests/frank_large_config.bats  |  7 +++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/scripts/bench-ubatch-rocmi4.sh b/scripts/bench-ubatch-rocmi4.sh
index 004a5b8..352050a 100755
--- a/scripts/bench-ubatch-rocmi4.sh
+++ b/scripts/bench-ubatch-rocmi4.sh
@@ -70,7 +70,14 @@ if [[ -f "$MMPROJ" ]]; then
 fi
 
 STopped=0
+SCRATCH_PID=""
 cleanup() {
+    if [[ -n "$SCRATCH_PID" ]] && kill -0 "$SCRATCH_PID" 2>/dev/null; then
+        echo "=== stopping scratch llama-server (pid $SCRATCH_PID) ==="
+        kill "$SCRATCH_PID" 2>/dev/null || true
+        wait "$SCRATCH_PID" 2>/dev/null || true
+        SCRATCH_PID=""
+    fi
     if [ "$STopped" = 1 ]; then
         echo "=== restarting frank-default-large ==="
         sudo systemctl restart frank-default-large || true
@@ -134,6 +141,7 @@ run_one() {
         --reasoning "$REASONING" --reasoning-format "$REASONING_FORMAT" --reasoning-budget -1 \
         -lv 4 >>"$log" 2>&1 &
     local pid=$!
+    SCRATCH_PID="$pid"
     wait_health "$pid" "$port" || { tail -30 "$log"; return 2; }
 
     # one python pass per run; reuse the same boot across runs
@@ -177,8 +185,12 @@ body = json.dumps({
 req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions",
                              data=body, headers={"Content-Type": "application/json"})
 t0 = time.time()
-with urllib.request.urlopen(req, timeout=7200) as resp:
-    out = json.load(resp)
+try:
+    with urllib.request.urlopen(req, timeout=7200) as resp:
+        out = json.load(resp)
+except Exception:
+    print("request failed; preserving benchmark cleanup", file=sys.stderr)
+    raise
 dt = time.time() - t0
 pt = out.get("usage", {}).get("prompt_tokens", 0)
 ct = out.get("usage", {}).get("completion_tokens", 0)
@@ -194,7 +206,7 @@ g_tps = ""
 if evals and ct:
     # eval timing line: "eval time = .. ms / N tokens (... tokens per second)"
     import re
-    match = re.search(r"eval time =\\s*([0-9.]+) ms", evals[-1])
+    match = re.search(r"eval time =\s*([0-9.]+) ms", evals[-1])
     if match:
         ms = float(match.group(1))
         g_tps = f"{ct / (ms / 1000):.2f}"
@@ -205,6 +217,7 @@ PY
     done
     kill "$pid" 2>/dev/null || true
     wait "$pid" 2>/dev/null || true
+    SCRATCH_PID=""
 }
 
 main() {
diff --git a/tests/frank_large_config.bats b/tests/frank_large_config.bats
index ac4808d..2d774bd 100644
--- a/tests/frank_large_config.bats
+++ b/tests/frank_large_config.bats
@@ -73,6 +73,13 @@ TEMPLATE_FILE="$REPO_ROOT/scripts/chat-template.jinja"
     grep -q 'REASONING_FORMAT' "$bench"
 }
 
+@test "large ubatch bench cleans up scratch server on request failure" {
+    local bench="$REPO_ROOT/scripts/bench-ubatch-rocmi4.sh"
+    grep -q 'SCRATCH_PID=' "$bench"
+    grep -q 'kill.*SCRATCH_PID' "$bench"
+    grep -qF 'eval time =\s' "$bench"
+}
+
 @test "frank-default-large uses the tracked Qwen-Sharp template on :7770" {
     grep -q '^CHAT_TEMPLATE_FILE=' "$ENV_FILE"
     grep -q -- '--jinja' "$SERVE_FILE"