fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: web UI with sessions and drafts

Browse Source
This commit is contained in:
fabolous005
2026-10-27 13:52:25 +01:00
parent 0531aaed44
commit cf9a347792
2 changed files with 412 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
prain/web.py
+274
View File
@@ -0,0 +1,274 @@
"""prain web: a one-page browser chat that warms experts up while you type.
A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes
{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
machines.
Several browsers at once ([serve] and [draft] tables of the placement file):
- At most draft.full_requests answers run with all their experts. An answer that arrives
while they are taken runs as a draft right away, on the experts already in VRAM (see
prain.drafts), instead of waiting or evicting the other answers' experts.
- With draft.refine = "regenerate", the oldest waiting draft takes the next free full slot:
its draft request is cancelled and the answer is generated again with all experts. The
page shows the draft in grey and the full answer over it.
- Every session has its own Warmer. warmup="idle": no warm-ups while any answer streams;
"always": warm-ups while vLLM has a free request slot (serve.max_requests). Sending cancels
that session's warm-up. A browser that goes away closes the upstream request, and vLLM
cancels it.
"""
import json
import threading
import urllib.error
from collections import OrderedDict, deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from prain.chat import Warmer, chat_body, draft_body, first_model, post, send_request, warm_body
from prain.config import Draft, Serve
MAX_SESSIONS = 64
PAGE = """<!doctype html>
<meta charset="utf-8">
<title>prain</title>
<style>
body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
.user { color: #666; margin-top: 1rem; }
.assistant { white-space: pre-wrap; }
.draft { color: #999; }
textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
#status { color: #888; font-size: 13px; }
</style>
<div id="log"></div>
<textarea id="draft" autofocus
placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
<div id="status"></div>
<script>
const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
const messages = [];
const log = document.getElementById("log");
const draft = document.getElementById("draft");
const status = document.getElementById("status");
let streaming = false;
function show(role, text) {
const div = document.createElement("div");
div.className = role;
div.textContent = text;
log.appendChild(div);
return div;
}
draft.addEventListener("input", () => {
if (!streaming && draft.value.trim())
fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
});
draft.addEventListener("keydown", async (event) => {
if (event.key !== "Enter" || event.shiftKey || streaming || !draft.value.trim()) return;
event.preventDefault();
const text = draft.value.trim();
messages.push({role: "user", content: text});
show("user", "you> " + text);
draft.value = "";
streaming = true;
const out = show("assistant", "");
const finalSpan = out.appendChild(document.createElement("span"));
const draftSpan = out.appendChild(document.createElement("span"));
draftSpan.className = "draft";
let mode = "full", draftText = "", finalText = "", refined = false;
const render = () => { // the full text, then the grey rest of the draft while they agree
finalSpan.textContent = finalText;
draftSpan.textContent = mode === "draft" ? draftText
: mode === "refine" && draftText.startsWith(finalText) ? draftText.slice(finalText.length) : "";
};
let start = performance.now(), first = 0, pieces = 0, buffer = "";
const response = await fetch("/chat", {method: "POST", body: JSON.stringify({session, messages})});
if (!response.ok) {
out.textContent = "error " + response.status + ": " + await response.text();
messages.pop();
streaming = false;
return;
}
const reader = response.body.getReader();
const decoder = new TextDecoder();
for (;;) {
const {done, value} = await reader.read();
if (done) break;
buffer += decoder.decode(value, {stream: true});
const lines = buffer.split("\\n");
buffer = lines.pop();
for (const line of lines) {
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
const data = JSON.parse(line.slice(6));
if (data.prain === "draft") {
mode = "draft";
status.textContent = "GPU busy: drafting from the experts in VRAM, refined when it is free";
} else if (data.prain === "refine") {
mode = "refine";
refined = true;
start = performance.now(), first = 0, pieces = 0;
status.textContent = "refining with all experts";
} else if (data.prain === "error") {
mode = "draft";
refined = false;
status.textContent = "refining failed (HTTP " + data.status + "), keeping the draft";
} else {
const piece = data.choices[0].delta.content;
if (!piece) continue;
if (pieces++ === 0) first = performance.now();
if (mode === "draft") draftText += piece; else finalText += piece;
}
render();
}
}
messages.push({role: "assistant", content: mode === "draft" ? draftText : finalText});
let same = 0;
while (same < draftText.length && draftText[same] === finalText[same]) same++;
const seconds = (performance.now() - first) / 1000;
status.textContent = (refined ? `draft kept ${same} of ${draftText.length} chars; refined: `
: mode === "draft" ? "draft, not refined: " : "") +
`first token ${((first - start) / 1000).toFixed(1)} s, ` +
`${pieces > 1 ? ((pieces - 1) / seconds).toFixed(1) : "0.0"} tok/s`;
streaming = false;
});
</script>
"""
def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
draft=Draft()):
model = first_model(base_url)
lock = threading.Condition()
sessions = OrderedDict() # session -> Warmer, least recently used first
streams = [0] # answers streaming right now, all sessions, full or draft
full = [0] # answers holding one of the draft.full_requests full slots
waiting = deque() # drafts waiting for a full slot to refine, oldest first
refine = draft.refine == "regenerate"
def warmer(session):
with lock:
if session not in sessions:
sessions[session] = Warmer(lambda job: send_request(
base_url, "/v1/chat/completions", warm_body(model, *job, priority)))
if len(sessions) > MAX_SESSIONS:
sessions.popitem(last=False)[1].cancel()
sessions.move_to_end(session)
return sessions[session]
def my_turn(ticket): # call with the lock held
return full[0] < draft.full_requests and bool(waiting) and waiting[0] is ticket
def turn_now(ticket):
with lock:
return my_turn(ticket)
class Handler(BaseHTTPRequestHandler):
def do_GET(self):
self.reply(200, "text/html; charset=utf-8", PAGE.encode())
def do_POST(self):
body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
session = str(body.get("session", ""))
if self.path == "/warm":
with lock:
busy = streams[0]
if body["draft"].strip() and ((warmup == "idle" and not busy) or
(warmup == "always" and busy < serve.max_requests)):
warmer(session).update((body["messages"], body["draft"]))
return self.reply(204, "text/plain", b"")
if self.path != "/chat":
return self.reply(404, "text/plain", b"not found")
warmer(session).cancel()
messages, ticket = body["messages"], object()
with lock:
streams[0] += 1
holds_full = full[0] < draft.full_requests and not waiting
if holds_full:
full[0] += 1
elif refine:
waiting.append(ticket)
try:
if holds_full:
if not self.relay(chat_body(model, messages), opens=True):
return
else:
if not self.relay(draft_body(model, messages), opens=True, event="draft",
stop=(lambda: turn_now(ticket)) if refine else None):
return
if refine:
with lock:
lock.wait_for(lambda: my_turn(ticket))
waiting.popleft()
full[0] += 1
holds_full = True
self.relay(chat_body(model, messages), opens=False, event="refine")
self.send_event("[DONE]")
except (BrokenPipeError, ConnectionResetError):
pass # the browser left; closing upstream makes vLLM cancel the request
finally:
with lock:
streams[0] -= 1
if holds_full:
full[0] -= 1
if ticket in waiting:
waiting.remove(ticket)
lock.notify_all()
def relay(self, body, opens, event=None, stop=None):
"""Stream one vLLM answer to the browser, without its [DONE]. opens: this starts
the response, so a vLLM error becomes the HTTP status (returns False). stop():
checked after every chunk; True closes the upstream request early."""
try:
upstream = post(base_url, "/v1/chat/completions", body)
except urllib.error.HTTPError as e: # vLLM rejected it: pass its answer on
with e:
data = e.read()
if opens:
self.reply(e.code, "application/json", data)
return False
self.send_event({"prain": "error", "status": e.code})
return True
with upstream:
if opens:
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.end_headers()
if event:
self.send_event({"prain": event})
for line in upstream:
if line.strip() in (b"", b"data: [DONE]"):
continue
self.wfile.write(line.rstrip(b"\n") + b"\n\n")
self.wfile.flush()
if stop is not None and stop():
break
return True
def send_event(self, data):
text = data if isinstance(data, str) else json.dumps(data)
self.wfile.write(f"data: {text}\n\n".encode())
self.wfile.flush()
def reply(self, code, kind, data):
self.send_response(code)
self.send_header("Content-Type", kind)
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
def log_message(self, *args):
pass
server = ThreadingHTTPServer((host, port), Handler)
server.sessions, server.streams, server.full = sessions, streams, full # for tests
return server
def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()):
server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
print(f"prain web: http://127.0.0.1:{server.server_port} (vLLM at {base_url}; "
f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
server.serve_forever()
tests/test_web.py
+138
View File
@@ -0,0 +1,138 @@
import json
import threading
import unittest
import urllib.error
import urllib.request
import prain.chat as chat
from prain.config import Draft, Serve
from prain.web import make_server
from tests.helpers import StubVllm, wait_for
def events(response):
"""The proxy's stream as a list: content pieces, and "<draft>"-style markers."""
out = []
for raw in response:
line = raw.decode().strip()
if line == "data: [DONE]":
out.append("<done>")
elif line.startswith("data: "):
data = json.loads(line[6:])
out.append(f"<{data['prain']}>" if "prain" in data
else data["choices"][0]["delta"]["content"])
return out
class Web(unittest.TestCase):
def setUp(self):
self.stub = StubVllm()
self.addCleanup(self.stub.close)
self.start()
def start(self, **options):
self.server = make_server(self.stub.url, 0, **options)
threading.Thread(target=self.server.serve_forever, daemon=True).start()
self.addCleanup(self.server.server_close)
self.addCleanup(self.server.shutdown)
self.url = f"http://127.0.0.1:{self.server.server_port}"
def warm(self, session, draft):
chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read()
def ask(self, session, text):
return chat.post(self.url, "/chat", {"session": session,
"messages": [{"role": "user", "content": text}]})
def test_page_warm_up_and_chat(self):
with urllib.request.urlopen(self.url) as response:
self.assertIn(b'<textarea id="draft"', response.read())
self.warm("s1", "Why is ")
wait_for(lambda: self.stub.bodies)
self.assertEqual(self.stub.bodies[0]["messages"], [{"role": "user", "content": "Why is"}])
self.assertEqual(self.stub.bodies[0]["priority"], 1)
with self.ask("s1", "Why is") as r:
self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
self.assertEqual(self.stub.bodies[1]["model"], "stub-model")
self.assertNotIn("vllm_xargs", self.stub.bodies[1])
def test_sessions_warm_up_independently(self):
self.stub.delay = 0.3
self.warm("alice", "first")
self.warm("bob", "second") # alice's warm-up is in flight; bob's is not dropped
wait_for(lambda: len(self.stub.bodies) == 2)
self.assertEqual({b["messages"][-1]["content"] for b in self.stub.bodies},
{"first", "second"})
self.assertEqual(set(self.server.sessions), {"alice", "bob"})
def test_no_warm_up_while_another_answer_streams(self):
self.server.streams[0] = 1 # someone else's answer is streaming
self.warm("bob", "typing along")
self.server.streams[0] = 0
self.warm("bob", "now")
wait_for(lambda: self.stub.bodies)
self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["now"])
def test_always_mode_warms_up_while_a_request_slot_is_free(self):
self.start(warmup="always", serve=Serve(max_requests=3))
self.server.streams[0] = 3 # every vLLM request slot is busy
self.warm("bob", "no slot")
self.server.streams[0] = 2
self.warm("bob", "one slot")
wait_for(lambda: self.stub.bodies)
self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["one slot"])
def test_vllm_errors_reach_the_browser(self):
self.stub.fail = 400
with self.assertRaises(urllib.error.HTTPError) as caught:
chat.post(self.url, "/chat", {"session": "s", "messages": []})
self.assertEqual(caught.exception.code, 400)
with caught.exception as error:
self.assertIn(b"rejected", error.read())
wait_for(lambda: self.server.streams[0] == 0)
def test_second_answer_drafts_now_and_refines_when_the_first_one_ends(self):
self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5 # full answers are slow
alice = []
thread = threading.Thread(target=lambda: alice.extend(events(self.ask("alice", "a"))))
thread.start()
wait_for(lambda: self.server.full[0] == 1)
with self.ask("bob", "b") as r:
bob = events(r)
thread.join()
self.assertEqual(alice, ["Hel", "lo", "!", "<done>"])
self.assertEqual(bob, ["<draft>", "Hel", "lo", "!", "<refine>", "Hel", "lo", "!", "<done>"])
self.assertEqual([(b["messages"][0]["content"], b.get("vllm_xargs"), b.get("cache_salt"))
for b in self.stub.bodies],
[("a", None, None), ("b", {"prain": "draft"}, "prain-draft"),
("b", None, None)])
wait_for(lambda: self.server.full[0] == 0 and self.server.streams[0] == 0)
def test_a_running_draft_is_cut_short_once_a_full_slot_frees(self):
self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.3
self.stub.pause = lambda body: 0.6 if "vllm_xargs" in body else 0.0 # drafts trickle
thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
thread.start()
wait_for(lambda: self.server.full[0] == 1)
with self.ask("bob", "b") as r: # alice ends at 0.3 s, between bob's "Hel" and "lo"
self.assertEqual(events(r), ["<draft>", "Hel", "lo", "<refine>", "Hel", "lo", "!",
"<done>"])
thread.join()
def test_refine_none_keeps_the_draft_and_more_full_slots_mean_no_draft(self):
self.start(draft=Draft(refine="none"))
self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5
thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
thread.start()
wait_for(lambda: self.server.full[0] == 1)
with self.ask("bob", "b") as r:
self.assertEqual(events(r), ["<draft>", "Hel", "lo", "!", "<done>"])
thread.join()
self.start(serve=Serve(max_requests=2), draft=Draft(full_requests=2))
thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
thread.start()
wait_for(lambda: self.server.full[0] == 1)
with self.ask("bob", "b") as r:
self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
thread.join()
self.assertEqual(sum("vllm_xargs" in b for b in self.stub.bodies), 1)