prain: chat client with cancellable warm-up

prain/chat.py

@@ -0,0 +1,119 @@
+"""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 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()

tests/test_chat.py

@@ -0,0 +1,60 @@
+import threading
+import time
+import unittest
+
+import prain.chat as chat
+from tests.helpers import StubVllm, wait_for
+
+
+class Chat(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def warmer(self):
+        return chat.Warmer(lambda draft: chat.send_request(
+            self.stub.url, "/v1/chat/completions", chat.warm_body("m", [], draft, priority=1)))
+
+    def test_warm_body_continues_the_half_typed_message(self):
+        body = chat.warm_body("m", [{"role": "assistant", "content": "hi"}], "  Explain mix ", 1)
+        self.assertEqual(body["messages"][-1], {"role": "user", "content": "Explain mix"})
+        self.assertEqual(
+            (body["max_tokens"], body["add_generation_prompt"], body["continue_final_message"],
+             body["stream"], body["priority"]), (1, False, True, True, 1))
+        self.assertNotIn("priority", chat.warm_body("m", [], "x"))
+
+    def test_draft_body_marks_the_request_and_salts_its_cache(self):
+        body = chat.draft_body("m", [{"role": "user", "content": "hi"}])
+        self.assertEqual(body, {"model": "m", "messages": [{"role": "user", "content": "hi"}],
+                                "stream": True, "vllm_xargs": {"prain": "draft"},
+                                "cache_salt": "prain-draft"})
+
+    def test_first_model_and_streamed_text(self):
+        self.assertEqual(chat.first_model(self.stub.url), "stub-model")
+        with chat.post(self.stub.url, "/v1/chat/completions",
+                       chat.chat_body("stub-model", [], max_tokens=5)) as response:
+            self.assertEqual(list(chat.stream_text(response)), ["Hel", "lo", "!"])
+        self.assertEqual((self.stub.bodies[0]["stream"], self.stub.bodies[0]["max_tokens"]),
+                         (True, 5))
+
+    def test_warmer_keeps_one_request_in_flight_and_sends_the_newest_draft(self):
+        self.stub.delay = 0.3
+        warmer = self.warmer()
+        warmer.update("a")
+        wait_for(lambda: len(self.stub.bodies) == 1)
+        warmer.update("ab")  # "a" is still in flight: these two queue up,
+        warmer.update("abc")  # and only the newest is sent
+        wait_for(lambda: not warmer.busy)
+        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["a", "abc"])
+
+    def test_cancel_drops_the_queue_and_aborts_the_request_in_flight(self):
+        self.stub.delay = 3.0
+        warmer = self.warmer()
+        warmer.update("slow")
+        wait_for(lambda: self.stub.bodies)
+        warmer.update("queued")
+        start = time.monotonic()
+        warmer.cancel()
+        wait_for(lambda: not warmer.busy, timeout=2.0)
+        self.assertLess(time.monotonic() - start, 1.0)  # did not wait out the 3 s answer
+        self.assertEqual(len(self.stub.bodies), 1)  # "queued" was never sent