fabolous005
a67238f66b
prain web/tui --persona FILE sends the file as the system message ahead of every warm-up and answer; the page shows its name. Three examples in examples/personas (task, how to answer, environment, how to use it). A persona costs 2.0-3.2 s once (241-467 tokens), then at most 0.3 s per conversation from the prefix cache. README rewritten around launching: quick start, the vLLM flags, a quantization table measured today at equal expert VRAM (bf16 0.90 / int8 2.19 / int4 5.19 tok/s; 2.3k-token prompt 69.5 / 32.3 / 20.5 s; 4 of 6 int8 and 0 of 6 int4 answers identical to bf16, one int4 answer wrong), how to switch formats, the VRAM budget, writing personas, the chats and the monitor.
125 lines
4.8 KiB
Python
125 lines
4.8 KiB
Python
"""OpenAI-compatible client pieces shared by the TUI and the WebUI. Stdlib only."""
|
|
|
|
import http.client
|
|
import json
|
|
import socket
|
|
import threading
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
ABORTED = (OSError, ValueError, http.client.HTTPException) # what a cancelled request raises
|
|
|
|
|
|
def post(base_url, path, body, timeout=600):
|
|
request = urllib.request.Request(f"{base_url}{path}", json.dumps(body).encode(),
|
|
{"Content-Type": "application/json"})
|
|
return urllib.request.urlopen(request, timeout=timeout)
|
|
|
|
|
|
def send_request(base_url, path, body, timeout=600):
|
|
"""POST without waiting for the answer. Returns the connection: call getresponse() to
|
|
read it, or shutdown its socket from another thread to abort it (vLLM then cancels)."""
|
|
url = urllib.parse.urlsplit(base_url)
|
|
connection = http.client.HTTPConnection(url.hostname, url.port or 80, timeout=timeout)
|
|
connection.request("POST", path, json.dumps(body), {"Content-Type": "application/json"})
|
|
return connection
|
|
|
|
|
|
def first_model(base_url):
|
|
with urllib.request.urlopen(f"{base_url}/v1/models", timeout=30) as response:
|
|
return json.load(response)["data"][0]["id"]
|
|
|
|
|
|
def chat_body(model, messages, **extra):
|
|
return {"model": model, "messages": messages, "stream": True, **extra}
|
|
|
|
|
|
def with_system(system, messages):
|
|
"""messages behind the persona's system message; unchanged without one."""
|
|
return [{"role": "system", "content": system}, *messages] if system else list(messages)
|
|
|
|
|
|
def draft_body(model, messages):
|
|
"""An answer computed from the experts already in VRAM: prain.drafts marks the request
|
|
by its vllm_xargs. Its own cache_salt keeps its KV blocks, computed with skipped experts,
|
|
out of the prefix-cache hits of full requests."""
|
|
return chat_body(model, messages, vllm_xargs={"prain": "draft"}, cache_salt="prain-draft")
|
|
|
|
|
|
def warm_body(model, messages, draft, priority=0):
|
|
"""Prefill the half-typed message and stop after one token. vLLM keeps its KV blocks
|
|
(prefix caching) and prain leaves the experts it routed to most in VRAM, so the real
|
|
request only prefills what was typed after this. Streamed, so aborting it is cheap.
|
|
priority > 0 queues it behind real requests (needs --scheduling-policy priority)."""
|
|
body = {"model": model, "messages": [*messages, {"role": "user", "content": draft.strip()}],
|
|
"max_tokens": 1, "temperature": 0, "stream": True,
|
|
"add_generation_prompt": False, "continue_final_message": True}
|
|
if priority:
|
|
body["priority"] = priority
|
|
return body
|
|
|
|
|
|
def stream_text(response):
|
|
"""Yield the content pieces of a streamed chat completion (server-sent events)."""
|
|
for raw in response:
|
|
line = raw.decode().strip()
|
|
if not line.startswith("data: ") or line == "data: [DONE]":
|
|
continue
|
|
piece = json.loads(line[6:])["choices"][0].get("delta", {}).get("content")
|
|
if piece:
|
|
yield piece
|
|
|
|
|
|
class Warmer:
|
|
"""Warm-up requests of one chat session. At most one is in flight; while it runs, newer
|
|
drafts replace each other and only the newest is sent next. cancel() drops the queued
|
|
draft and aborts the one in flight, so a sent message never waits behind a warm-up."""
|
|
|
|
def __init__(self, start):
|
|
self.start = start # start(job) -> connection from send_request()
|
|
self.lock = threading.Lock()
|
|
self.pending = None
|
|
self.busy = False
|
|
self.connection = None
|
|
self.sent = 0
|
|
|
|
def update(self, job):
|
|
with self.lock:
|
|
self.pending = job
|
|
if self.busy:
|
|
return
|
|
self.busy = True
|
|
threading.Thread(target=self._run, daemon=True).start()
|
|
|
|
def cancel(self):
|
|
with self.lock:
|
|
self.pending = None
|
|
connection = self.connection
|
|
if connection is not None and connection.sock is not None:
|
|
try:
|
|
connection.sock.shutdown(socket.SHUT_RDWR) # wakes the blocked reader
|
|
except OSError:
|
|
pass
|
|
|
|
def _run(self):
|
|
while True:
|
|
with self.lock:
|
|
job, self.pending = self.pending, None
|
|
if job is None:
|
|
self.busy = False
|
|
return
|
|
try:
|
|
connection = self.start(job)
|
|
with self.lock:
|
|
self.connection = connection
|
|
self.sent += 1
|
|
with connection.getresponse() as response:
|
|
response.read()
|
|
except ABORTED:
|
|
pass # best effort: the real request reports server errors
|
|
finally:
|
|
with self.lock:
|
|
connection, self.connection = self.connection, None
|
|
if connection is not None:
|
|
connection.close()
|