prain: warm-up and draft measurement scripts

scripts/draft_effect.py

@@ -0,0 +1,54 @@
+"""A second answer that arrives while a first one streams, through `prain web`: when does
+its first token come, and how much of its draft survives the refine?
+
+Run it against a proxy with drafts and one without ([draft] full_requests = max_requests)
+and compare. usage: python3 scripts/draft_effect.py http://127.0.0.1:8080
+"""
+
+import json
+import os
+import sys
+import threading
+import time
+
+from prain.chat import post
+
+LONG = [{"role": "user", "content": "Write a detailed 400-word history of the printing press."}]
+SHORT = [{"role": "user", "content": "Name three uses of copper, one line each."}]
+
+
+def answer(url, session, messages, out):
+    start, mode = time.monotonic(), "full"
+    out.update(first=None, draft="", final="")
+    with post(url, "/chat", {"session": session, "messages": messages}) as response:
+        for raw in response:
+            line = raw.decode().strip()
+            if not line.startswith("data: ") or line == "data: [DONE]":
+                continue
+            data = json.loads(line[6:])
+            if "prain" in data:
+                mode = data["prain"]
+                continue
+            piece = data["choices"][0]["delta"].get("content")
+            if piece:
+                out["first"] = out["first"] or time.monotonic() - start
+                out["draft" if mode == "draft" else "final"] += piece
+    out["total"] = time.monotonic() - start
+
+
+def main():
+    url = sys.argv[1]
+    first, second = {}, {}
+    thread = threading.Thread(target=answer, args=(url, "first", LONG, first))
+    thread.start()
+    time.sleep(5)  # the first answer is decoding now
+    answer(url, "second", SHORT, second)
+    thread.join()
+    kept = len(os.path.commonprefix([second["draft"], second["final"]]))
+    print(f"second answer: first token {second['first']:.1f} s, done after {second['total']:.1f} s; "
+          + (f"draft {len(second['draft'])} chars, {kept} kept by the refine"
+             if second["draft"] else "no draft"))
+
+
+if __name__ == "__main__":
+    main()

scripts/warm_effect.py

@@ -0,0 +1,38 @@
+"""Time to first token for a fresh prompt: cold, and warmed up as if it had been typed.
+
+usage (prain installed): python3 scripts/warm_effect.py http://127.0.0.1:8000
+"""
+
+import sys
+import time
+import uuid
+
+from prain.chat import chat_body, first_model, post, stream_text, warm_body
+
+PROMPT = ("You are reviewing a design for a home server that runs a mixture-of-experts language "
+          "model on a single consumer graphics card with eight gigabytes of memory. The model has "
+          "thirty layers with one hundred and twenty-eight experts each, and only a few experts per "
+          "layer fit on the card at once. The rest wait in system memory or on an NVMe drive and "
+          "are copied over PCI Express when a token needs them. List the three biggest risks of "
+          "this design and, for each, one concrete way to measure it.")
+
+
+def first_token_seconds(base_url, model, text):
+    start = time.monotonic()
+    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=1)
+    with post(base_url, "/v1/chat/completions", body) as response:
+        next(stream_text(response), None)
+    return time.monotonic() - start
+
+
+def main(base_url):
+    model = first_model(base_url)
+    cold = first_token_seconds(base_url, model, f"[{uuid.uuid4().hex[:8]}] {PROMPT}")
+    text = f"[{uuid.uuid4().hex[:8]}] {PROMPT}"  # a new tag: no prefix-cache hit from above
+    post(base_url, "/v1/chat/completions", warm_body(model, [], text[:-40])).read()
+    warm = first_token_seconds(base_url, model, text)
+    print(f"first token: cold {cold:.2f} s, warmed {warm:.2f} s")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])