275 lines
12 KiB
Python
275 lines
12 KiB
Python
"""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()
|