prain: chat client with cancellable warm-up
This commit is contained in:
+119
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user