prain: web UI with sessions and drafts

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 = """<!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

@@ -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)