233 lines
11 KiB
Bash
Executable File
233 lines
11 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
#
|
|
# serve_session_export_smoke.sh — regression test for "session export/restore
|
|
# only ever worked on the single-user path".
|
|
#
|
|
# The defect: /admin/session/export looked the session up in the model's
|
|
# pn_prefix_cache. A BATCHED model (QWEN35_BATCH_SLOTS=1 — every model a K80
|
|
# deployment runs) keeps its session state in the batch scheduler's warm decode slot,
|
|
# and nothing ever puts an entry in that model's prefix cache, so the lookup
|
|
# returned NULL and EVERY export against a batched model answered 404
|
|
# session_not_found. /session/restore had the mirror-image defect: it imported
|
|
# into the prefix cache, so it reported 200 restored:true while writing to a
|
|
# store no batched generation ever reads. /session/release cleared only the
|
|
# prefix cache too, leaving a departed session's context bound to its slot.
|
|
#
|
|
# Since the host KV mirror was removed (Phase 10b / INV3), export is the ONLY
|
|
# way a session's state can leave the GPU: a serve restart loses all warm
|
|
# context and the client never receives the blob the design says it owns.
|
|
#
|
|
# Asserts, against a LIVE session on the model under test:
|
|
# 1. export returns HTTP 200 with a non-empty blob, n_tok > 0, a
|
|
# weights_version, and the committed token ids. (This alone is the
|
|
# assertion that would have caught the defect.)
|
|
# 2. weights_version guard: the same blob with a corrupted weights_version
|
|
# is refused 409 weights_version_mismatch, and nothing is restored.
|
|
# 3. release actually drops the session: a second export answers 404.
|
|
# 4. restore of the exported blob answers 200 restored:true.
|
|
# 5. ROUND TRIP IS LOSSLESS: re-exporting after the restore yields a blob
|
|
# byte-identical to the original, with the same n_tok and token ids.
|
|
# This is the discriminating check — it proves the state left the GPU,
|
|
# crossed the wire format, and came back into a live slot intact.
|
|
# 6. The session still generates coherently after the restore (recall of a
|
|
# planted token that was in the context before the export).
|
|
#
|
|
# Usage (server must already be running):
|
|
# BASE_URL=http://127.0.0.1:8201 MODEL=saber-orch \
|
|
# tests/serve_session_export_smoke.sh
|
|
#
|
|
# Env:
|
|
# BASE_URL default http://127.0.0.1:8201
|
|
# MODEL default saber-orch
|
|
# SECRET planted recall token (default ZEPHYR-7-QUARTZ)
|
|
# FILLER words of real context in front of the planted token (default 200)
|
|
# MAX_TOKENS default 300
|
|
# REQ_TIMEOUT per-request curl timeout, seconds (default 600)
|
|
|
|
set -u
|
|
|
|
BASE_URL="${BASE_URL:-http://127.0.0.1:8201}"
|
|
MODEL="${MODEL:-saber-orch}"
|
|
SECRET="${SECRET:-ZEPHYR-7-QUARTZ}"
|
|
FILLER="${FILLER:-200}"
|
|
MAX_TOKENS="${MAX_TOKENS:-300}"
|
|
REQ_TIMEOUT="${REQ_TIMEOUT:-600}"
|
|
|
|
WORK="$(mktemp -d)"
|
|
SID="expsmoke-$$-$(date +%s)"
|
|
PASS=0
|
|
FAIL=0
|
|
|
|
cleanup() {
|
|
curl -s -m 30 -o /dev/null "$BASE_URL/session/release" \
|
|
-H 'Content-Type: application/json' \
|
|
-d "{\"name\":\"$MODEL\",\"session_id\":\"$SID\"}" || true
|
|
rm -rf "$WORK"
|
|
}
|
|
trap cleanup EXIT
|
|
|
|
ok() { PASS=$((PASS+1)); echo " PASS $*"; }
|
|
bad() { FAIL=$((FAIL+1)); echo " FAIL $*"; }
|
|
|
|
# jq is not assumed — pull scalars with python3, which the K80 test machine has.
|
|
jget() { # jget <file> <expr> e.g. jget f 'd["n_tok"]'
|
|
python3 -c '
|
|
import json,sys
|
|
d=json.load(open(sys.argv[1]))
|
|
v=eval(sys.argv[2])
|
|
print("" if v is None else v)
|
|
' "$1" "$2" 2>/dev/null
|
|
}
|
|
|
|
post() { # post <path> <json-file-in> <json-file-out> -> echoes http code
|
|
curl -s -m "$REQ_TIMEOUT" -o "$3" -w '%{http_code}' \
|
|
"$BASE_URL$1" -H 'Content-Type: application/json' --data-binary "@$2"
|
|
}
|
|
|
|
echo "== session export/restore round trip: model=$MODEL sid=$SID =="
|
|
|
|
# ── Turn 1: plant a distinctive token behind real context ──────────────────
|
|
python3 - "$WORK/turn1.json" "$MODEL" "$SID" "$SECRET" "$FILLER" "$MAX_TOKENS" <<'PY'
|
|
import json, sys
|
|
out, model, sid, secret, filler, maxtok = sys.argv[1:7]
|
|
words = ("The maintenance log records routine checks of the cooling loop, the "
|
|
"power rails and the interconnect fabric. ").split()
|
|
body = " ".join(words[i % len(words)] for i in range(int(filler)))
|
|
msg = (f"{body}\n\nRemember this access code exactly: {secret}. "
|
|
"Reply with only the word: stored")
|
|
json.dump({"model": model, "session_id": sid, "max_tokens": int(maxtok),
|
|
"messages": [{"role": "user", "content": msg}]}, open(out, "w"))
|
|
PY
|
|
|
|
code=$(post /v1/chat/completions "$WORK/turn1.json" "$WORK/turn1.out")
|
|
if [ "$code" = "200" ] && [ "$(jget "$WORK/turn1.out" 'd["usage"]["completion_tokens"]')" -gt 0 ] 2>/dev/null; then
|
|
ok "turn 1 generated (${code}, $(jget "$WORK/turn1.out" 'd["usage"]["prompt_tokens"]') prompt tokens)"
|
|
else
|
|
bad "turn 1 did not generate (HTTP $code)"; cat "$WORK/turn1.out"; echo
|
|
echo "== $PASS passed, $FAIL failed =="; exit 1
|
|
fi
|
|
|
|
# ── 1. Export ──────────────────────────────────────────────────────────────
|
|
cat > "$WORK/exp.json" <<EOF
|
|
{"name":"$MODEL","session_id":"$SID"}
|
|
EOF
|
|
code=$(post /admin/session/export "$WORK/exp.json" "$WORK/exp1.out")
|
|
N_TOK=$(jget "$WORK/exp1.out" 'd.get("n_tok")')
|
|
WV=$(jget "$WORK/exp1.out" 'd.get("weights_version")')
|
|
BLOB_BYTES=$(jget "$WORK/exp1.out" 'd.get("blob_bytes")')
|
|
if [ "$code" = "200" ] && [ -n "$N_TOK" ] && [ "$N_TOK" -gt 0 ] 2>/dev/null \
|
|
&& [ -n "$WV" ] && [ "${BLOB_BYTES:-0}" -gt 0 ] 2>/dev/null; then
|
|
ok "export 200: n_tok=$N_TOK blob_bytes=$BLOB_BYTES weights_version=$WV"
|
|
else
|
|
bad "export failed (HTTP $code) — THE REGRESSION"; head -c 300 "$WORK/exp1.out"; echo
|
|
echo "== $PASS passed, $FAIL failed =="; exit 1
|
|
fi
|
|
if [ "$(jget "$WORK/exp1.out" 'len(d.get("tokens") or [])')" = "$N_TOK" ]; then
|
|
ok "export returned the $N_TOK committed token ids"
|
|
else
|
|
bad "export tokens[] length != n_tok"
|
|
fi
|
|
|
|
# Build the restore bodies (good + weights_version-corrupted).
|
|
python3 - "$WORK/exp1.out" "$MODEL" "$SID" "$WORK/res_good.json" "$WORK/res_badwv.json" <<'PY'
|
|
import base64, json, sys, struct
|
|
src, model, sid, good, badwv = sys.argv[1:6]
|
|
d = json.load(open(src))
|
|
raw = base64.b64decode(d["blob_b64"])
|
|
json.dump({"name": model, "session_id": sid, "blob": d["blob_b64"],
|
|
"tokens": d["tokens"]}, open(good, "w"))
|
|
# The blob header is {magic, format_version, weights_version, ...}; flip a bit
|
|
# in the weights_version field so only the identity tag differs.
|
|
b = bytearray(raw)
|
|
wv_off = None
|
|
tag = struct.pack("<Q", int(d["weights_version"]))
|
|
for off in range(0, min(len(b), 64) - 8):
|
|
if bytes(b[off:off + 8]) == tag:
|
|
wv_off = off
|
|
break
|
|
assert wv_off is not None, "weights_version tag not found in the blob header"
|
|
b[wv_off] ^= 0xFF
|
|
json.dump({"name": model, "session_id": sid,
|
|
"blob": base64.b64encode(bytes(b)).decode(),
|
|
"tokens": d["tokens"]}, open(badwv, "w"))
|
|
PY
|
|
|
|
# ── 2. weights_version guard ───────────────────────────────────────────────
|
|
code=$(post /session/restore "$WORK/res_badwv.json" "$WORK/badwv.out")
|
|
if [ "$code" = "409" ] && grep -q weights_version_mismatch "$WORK/badwv.out"; then
|
|
ok "weights_version guard: corrupted tag refused 409 weights_version_mismatch"
|
|
else
|
|
bad "weights_version guard did not fire (HTTP $code): $(head -c 200 "$WORK/badwv.out")"
|
|
fi
|
|
|
|
# ── 3. Release really drops the session ────────────────────────────────────
|
|
post /session/release "$WORK/exp.json" "$WORK/rel.out" >/dev/null
|
|
code=$(post /admin/session/export "$WORK/exp.json" "$WORK/exp_gone.out")
|
|
if [ "$code" = "404" ]; then
|
|
ok "release dropped the session (export now 404)"
|
|
else
|
|
bad "release did not drop the session (export still HTTP $code)"
|
|
fi
|
|
|
|
# ── 4. Restore ─────────────────────────────────────────────────────────────
|
|
code=$(post /session/restore "$WORK/res_good.json" "$WORK/res.out")
|
|
if [ "$code" = "200" ] && [ "$(jget "$WORK/res.out" 'd.get("n_tok")')" = "$N_TOK" ]; then
|
|
ok "restore 200: n_tok=$N_TOK"
|
|
else
|
|
bad "restore failed (HTTP $code): $(head -c 300 "$WORK/res.out")"
|
|
echo "== $PASS passed, $FAIL failed =="; exit 1
|
|
fi
|
|
|
|
# ── 5. Round trip is lossless ──────────────────────────────────────────────
|
|
code=$(post /admin/session/export "$WORK/exp.json" "$WORK/exp2.out")
|
|
if [ "$code" != "200" ]; then
|
|
bad "re-export after restore failed (HTTP $code)"
|
|
else
|
|
if python3 - "$WORK/exp1.out" "$WORK/exp2.out" <<'PY'
|
|
import json, sys
|
|
a = json.load(open(sys.argv[1])); b = json.load(open(sys.argv[2]))
|
|
same = (a["n_tok"] == b["n_tok"] and a["tokens"] == b["tokens"]
|
|
and a["blob_b64"] == b["blob_b64"])
|
|
sys.exit(0 if same else 1)
|
|
PY
|
|
then
|
|
ok "round trip lossless: re-exported blob byte-identical to the original"
|
|
else
|
|
bad "round trip LOSSY: re-exported blob differs from the original"
|
|
python3 - "$WORK/exp1.out" "$WORK/exp2.out" <<'PY'
|
|
import json, sys
|
|
a = json.load(open(sys.argv[1])); b = json.load(open(sys.argv[2]))
|
|
print(f" n_tok {a['n_tok']} -> {b['n_tok']}; "
|
|
f"blob_bytes {a['blob_bytes']} -> {b['blob_bytes']}; "
|
|
f"tokens equal: {a['tokens'] == b['tokens']}")
|
|
PY
|
|
fi
|
|
fi
|
|
|
|
# ── 6. The restored session still answers ──────────────────────────────────
|
|
python3 - "$WORK/turn2.json" "$WORK/turn1.json" "$WORK/turn1.out" "$MAX_TOKENS" <<'PY'
|
|
import json, sys
|
|
out, t1, t1o, maxtok = sys.argv[1:5]
|
|
req = json.load(open(t1)); res = json.load(open(t1o))
|
|
msgs = list(req["messages"])
|
|
msgs.append({"role": "assistant",
|
|
"content": res["choices"][0]["message"].get("content") or "stored"})
|
|
msgs.append({"role": "user",
|
|
"content": "What was the access code? Reply with only the code."})
|
|
req["messages"] = msgs
|
|
req["max_tokens"] = int(maxtok)
|
|
json.dump(req, open(out, "w"))
|
|
PY
|
|
code=$(post /v1/chat/completions "$WORK/turn2.json" "$WORK/turn2.out")
|
|
ANS=$(jget "$WORK/turn2.out" '(d["choices"][0]["message"].get("content") or "") + " " + (d["choices"][0]["message"].get("reasoning_content") or "")')
|
|
if [ "$code" = "200" ] && printf '%s' "$ANS" | grep -q "$SECRET"; then
|
|
ok "post-restore turn recalls the planted token ($SECRET)"
|
|
else
|
|
bad "post-restore turn did not recall the planted token (HTTP $code): $(printf '%s' "$ANS" | head -c 200)"
|
|
fi
|
|
|
|
echo "== $PASS passed, $FAIL failed =="
|
|
[ "$FAIL" -eq 0 ]
|