139 lines
6.3 KiB
Python
139 lines
6.3 KiB
Python
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)
|