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.
208 lines
10 KiB
Python
208 lines
10 KiB
Python
import json
|
|
import os
|
|
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, TmpCase, 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):
|
|
with chat.post(self.url, "/warm", {"session": session, "messages": [],
|
|
"draft": draft}) as response:
|
|
return response.status
|
|
|
|
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:
|
|
page = response.read()
|
|
self.assertIn(b'<textarea id="draft"', page)
|
|
self.assertIn(b'"/status"', page) # the monitor polls it
|
|
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.assertEqual(self.warm("bob", "typing along"), 204) # 204: nothing warms up
|
|
self.server.streams[0] = 0
|
|
self.assertEqual(self.warm("bob", "now"), 202) # 202: a warm-up starts
|
|
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_a_persona_leads_every_warm_up_and_answer(self):
|
|
self.start(system="You answer in German.", persona="german")
|
|
self.warm("s1", "Wie spät")
|
|
wait_for(lambda: self.stub.bodies)
|
|
with self.ask("s1", "Wie spät ist es?") as r:
|
|
r.read()
|
|
system = {"role": "system", "content": "You answer in German."}
|
|
self.assertEqual([b["messages"][0] for b in self.stub.bodies], [system, system])
|
|
self.assertEqual(self.stub.bodies[1]["messages"][1:],
|
|
[{"role": "user", "content": "Wie spät ist es?"}])
|
|
with urllib.request.urlopen(self.url + "/status") as r:
|
|
self.assertEqual(json.load(r)["persona"], "german")
|
|
|
|
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)
|
|
|
|
|
|
class Status(TmpCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.stub = StubVllm()
|
|
self.addCleanup(self.stub.close)
|
|
|
|
def status(self, **options):
|
|
server = make_server(self.stub.url, 0, **options)
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
self.addCleanup(server.server_close)
|
|
self.addCleanup(server.shutdown)
|
|
with urllib.request.urlopen(f"http://127.0.0.1:{server.server_port}/status") as r:
|
|
return json.load(r)
|
|
|
|
def test_status_reports_vllm_the_gpu_and_the_experts(self):
|
|
self.stub.metrics = (
|
|
"# HELP vllm:num_requests_running Number of requests running.\n"
|
|
'vllm:num_requests_running{engine="0",model_name="m"} 1.0\n'
|
|
'vllm:num_requests_waiting{engine="0",model_name="m"} 0.0\n'
|
|
'vllm:num_requests_waiting_by_reason{engine="0",reason="capacity"} 5.0\n'
|
|
'vllm:kv_cache_usage_perc{engine="0",model_name="m"} 0.25\n'
|
|
'vllm:prefix_cache_queries_total{engine="0",model_name="m"} 40.0\n'
|
|
'vllm:prefix_cache_hits_total{engine="0",model_name="m"} 30.0\n'
|
|
'vllm:prompt_tokens_total{engine="0",model_name="m"} 100.0\n'
|
|
'vllm:generation_tokens_total{engine="0",model_name="m"} 7.0\n')
|
|
os.makedirs(os.path.join(self.dir, "card0", "device")) # an iGPU: no VRAM files
|
|
gpu = os.path.join(self.dir, "card1", "device")
|
|
os.makedirs(os.path.join(gpu, "hwmon", "hwmon4"))
|
|
for name, value in {"mem_info_vram_total": 8573157376, "mem_info_vram_used": 8000000000,
|
|
"gpu_busy_percent": 96, "hwmon/hwmon4/temp1_input": 61000,
|
|
"hwmon/hwmon4/power1_average": 120500000,
|
|
"hwmon/hwmon4/fan1_input": 1800}.items():
|
|
self.write(os.path.join("card1", "device", name), f"{value}\n")
|
|
stats = {"layers": {"0": {"uses": [3, 0], "vram": [0], "hits": 2, "misses": 1}}}
|
|
stats_path = self.write("stats.json", json.dumps(stats))
|
|
got = self.status(stats_path=stats_path, sysfs=self.dir)
|
|
self.assertEqual(got["model"], "stub-model")
|
|
self.assertEqual(got["vllm"], {"running": 1.0, "waiting": 0.0, "kv_usage": 0.25,
|
|
"prefix_queries": 40.0, "prefix_hits": 30.0,
|
|
"prompt_tokens": 100.0, "generation_tokens": 7.0})
|
|
self.assertEqual(got["gpu"], {"vram_used": 8000000000, "vram_total": 8573157376,
|
|
"busy": 96, "temp": 61.0, "power": 120.5, "fan": 1800.0})
|
|
self.assertEqual(got["experts"], stats)
|
|
self.assertEqual(got["streams"], 0)
|
|
|
|
def test_status_without_metrics_gpu_or_stats(self):
|
|
got = self.status(sysfs=self.dir) # the stub answers /metrics with JSON, not metrics
|
|
self.assertEqual((got["gpu"], got["experts"]), (None, None))
|
|
self.assertEqual(got["vllm"]["running"], 0.0)
|