From cf9a3477920beedb42655ec24156975da643ed4b Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 13:52:25 +0100 Subject: [PATCH] prain: web UI with sessions and drafts --- prain/web.py | 274 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_web.py | 138 +++++++++++++++++++++++ 2 files changed, 412 insertions(+) create mode 100644 prain/web.py create mode 100644 tests/test_web.py diff --git a/prain/web.py b/prain/web.py new file mode 100644 index 0000000..4e3040e --- /dev/null +++ b/prain/web.py @@ -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 = """ + +prain + +
+ +
+ +""" + + +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() diff --git a/tests/test_web.py b/tests/test_web.py new file mode 100644 index 0000000..24aa835 --- /dev/null +++ b/tests/test_web.py @@ -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 ""-style markers.""" + out = [] + for raw in response: + line = raw.decode().strip() + if line == "data: [DONE]": + out.append("") + 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'