vh
30f36f2478
The vendor shim silently trims state at 32,000 tokens; the admission proxy is what turns the seat contract into a loud refusal, so it must front the shim in any deployment. Committing it with the Amendment 1 confirmation driver so executing the swap does not depend on a bench scratch dir surviving.
93 lines
4.9 KiB
Python
93 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Thin admission wrapper for the h2o confirmation run (Amendment 1, brokkr 63d7a1b).
|
|
|
|
Sits IN FRONT of the unmodified shipped h2o_lightning_shim.py and refuses over-cap
|
|
input loudly, mirroring INC's refusal. Counts with h2o's OWN tokenizer: renders the
|
|
state exactly as the shim does (imported verbatim from the shipped file) and asks
|
|
vLLM's /tokenize for the count. Refuses > CAP tokens with an explicit 422; forwards
|
|
everything else untouched. No trim. The shim is not modified.
|
|
"""
|
|
import importlib.util, json, os, sys, time, urllib.request
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
CAP = int(os.environ.get("ADMIT_CAP", "32000"))
|
|
UP = os.environ.get("ADMIT_UP", "http://127.0.0.1:18091") # the shim
|
|
VLLM = os.environ.get("ADMIT_VLLM", "http://127.0.0.1:18099")
|
|
PORT = int(os.environ.get("ADMIT_PORT", "18092"))
|
|
SHIM_PATH = os.environ.get("ADMIT_SHIM", "/tank/bench/pkg-h2o/h2o_lightning_shim.py")
|
|
|
|
spec = importlib.util.spec_from_file_location("h2o_shim", SHIM_PATH)
|
|
shim_mod = importlib.util.module_from_spec(spec)
|
|
sys.modules["h2o_shim"] = shim_mod
|
|
spec.loader.exec_module(shim_mod)
|
|
|
|
CFG = shim_mod.Contract(json.load(open("/tank/bench/pkg-h2o/serve_config.json")))
|
|
CFG.max_state_tokens = 10**9 # wrapper must count the UNTRUNCATED render; the shim's own trim is what we police
|
|
|
|
def sent_tokens(body):
|
|
"""Tokens per prompt as vLLM will actually receive them, counted by vLLM's own tokenizer via
|
|
the shim's own renderers (chat path: tokenize over messages, matching the image-free prompt
|
|
the shim sends per question). Returns (max_per_prompt, shim_accounted_total).
|
|
The shim accounts usage as shared_prefix + sum(question tails); we mirror that exactly."""
|
|
import concurrent.futures
|
|
state = body.get("state") or ""
|
|
if isinstance(state, dict):
|
|
state = json.dumps(state)
|
|
state_txt = shim_mod.render_state(state, CFG.state_compact) # UNTRUNCATED: what admission guards
|
|
all_toks = []
|
|
for qid, q in body.get("questions", {}).items():
|
|
prompt_text = CFG.question_prompt(state_txt, q, qid)[0] # EXACTLY what the shim sends per question
|
|
req = urllib.request.Request(VLLM + "/tokenize",
|
|
data=json.dumps({"model": CFG.model, "prompt": prompt_text, "add_special_tokens": False}).encode(),
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=30) as r:
|
|
all_toks.append(json.loads(r.read())["tokens"])
|
|
mx = max(len(t) for t in all_toks)
|
|
shared = shim_mod.common_prefix_len(all_toks)
|
|
total = shared + sum(len(t) - shared for t in all_toks) # the shim's own accounting: head + tails
|
|
return mx, total
|
|
|
|
class H(BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
def log_message(self, *a): pass
|
|
def _send(self, code, obj):
|
|
b = json.dumps(obj).encode()
|
|
self.send_response(code); self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b)
|
|
def do_GET(self):
|
|
try:
|
|
req = urllib.request.Request(UP + self.path)
|
|
with urllib.request.urlopen(req, timeout=30) as r:
|
|
self._send(r.status, json.loads(r.read()))
|
|
except urllib.error.HTTPError as e:
|
|
self._send(e.code, json.loads(e.read() or b"{}"))
|
|
def do_POST(self):
|
|
n = int(self.headers.get("Content-Length", 0))
|
|
raw = self.rfile.read(n)
|
|
mx = total = None
|
|
if self.path.split("?")[0] in ("/v1/systemone", "/api/alpha/decisions"):
|
|
try:
|
|
body = json.loads(raw)
|
|
except ValueError:
|
|
return self._send(422, {"error": "admission: body is not JSON"})
|
|
try:
|
|
mx, total = sent_tokens(body)
|
|
except Exception as e:
|
|
return self._send(502, {"error": f"admission: tokenizer count failed: {e}"})
|
|
if mx > CAP or total > CAP:
|
|
return self._send(422, {"error": f"admission: {mx} max-prompt / {total} total tokens against the "
|
|
f"{CAP}-token seat limit; refused, not trimmed",
|
|
"max_prompt_tokens": mx, "shim_accounted_total": total, "cap": CAP})
|
|
req = urllib.request.Request(UP + self.path, data=raw,
|
|
headers={"Content-Type": "application/json"}, method="POST")
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=300) as r:
|
|
obj = json.loads(r.read())
|
|
obj.setdefault("_admission", {})["max_prompt_tokens"] = mx if "mx" in dir() else None
|
|
obj["_admission"]["shim_accounted_total"] = total if "total" in dir() else None
|
|
self._send(r.status, obj)
|
|
except urllib.error.HTTPError as e:
|
|
self._send(e.code, json.loads(e.read() or b"{}"))
|
|
|
|
ThreadingHTTPServer(("127.0.0.1", PORT), H).serve_forever()
|