fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: web UI with a model monitor

Browse Source 
A clear chat page (markdown, code blocks with copy, stop, new chat, kept in the browser,
drafts shown as such) and a monitor sidebar, a sheet on phones: tokens/s, a live map of the
30 x 128 experts (in VRAM, used now, copied in), VRAM hit rate, copies, context, prompt cache,
GPU VRAM/load/temperature/power/fan. GET /status merges vLLM /metrics, amdgpu sysfs and the
stats file, which now lists the experts in VRAM. /warm answers 202 when a warm-up starts.
README: the int4 kernels, their A/B, 2048-token chunks as a trade, prompt logprobs limit.
This commit is contained in:
fabolous005
2026-10-29 14:20:23 +01:00
parent d1ccc5c294
commit a14e3ade0b
11 changed files with 977 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+47 -8
View File
@@ -34,9 +34,12 @@ Tune vLLM's fused-MoE kernel for the GPU once (~45 min of Triton compiles the fi
The first start writes the bf16 pack (42.5 GiB, ~3 min); with `quant = "int8"` or `"int4"` it
then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files
(`layerNNN.w13`, `.w2`, `.done`) are not read and can be deleted; a later bf16 or int8 start
writes them again. Chat with `prain tui`, or `prain web --config FILE` for several users with
drafts (`--host 0.0.0.0` serves it to other machines; it has no login). Tool calls work through
vLLM's `gemma4` parser, so coding agents and MCP clients can use the OpenAI-compatible API.
writes them again. Chat with `prain tui`, or in the browser with `prain web --config FILE`:
several users with drafts, and a monitor beside the chat (tokens/s, a live map of which experts
are in VRAM and working, VRAM hit rate, context use, GPU load, temperature and power; the map
needs `stats_path` in the placement file). `--host 0.0.0.0` serves it to other machines; it has
no login. Tool calls work through vLLM's `gemma4` parser, so coding agents and MCP clients can
use the OpenAI-compatible API.
## Expert formats
@@ -44,7 +47,7 @@ vLLM's `gemma4` parser, so coding agents and MCP clients can use the OpenAI-comp
|---|---|---|---|---|---|
| bf16 | 11.34 MiB | 42.5 GiB | as trained | vLLM Triton | torch bf16 matmul |
| int8 | 5.69 MiB | 21.3 GiB | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
| int4 | 3.01 MiB | 11.3 GiB | bf16 scale per 64, MSE clip | vLLM Triton (GPTQ/AWQ, w4a16) | `torch._weight_int4pack_mm_for_cpu` |
| int4 | 3.01 MiB | 11.3 GiB | bf16 scale per 64, MSE clip | vLLM Triton (GPTQ/AWQ, w4a16); prain's on RDNA2 | `torch._weight_int4pack_mm_for_cpu` |
Measured on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4, 2.5 GB/s), i9-13950HX (16 CPU
threads for experts), 31 GiB RAM, quiet host. Every format got 1.66 GiB of expert VRAM, its own
@@ -61,7 +64,7 @@ placement from `prain tune` and tuned kernels; 4k context, 0.75 GiB KV:
The smaller formats win twice: more experts fit in VRAM, and every miss moves fewer bytes over
the link and out of the page cache (the int4 pack fits in free RAM, the bf16 one does not). On
the GPU itself a decode round costs about the same in every format (2.3-2.8 ms). These decode
numbers predate the dense-layer fix below, which makes every format faster.
numbers predate the dense-layer fix and the int4 kernels below, which make decode 3x faster.
## Dense layers on RDNA2
@@ -72,6 +75,32 @@ with at most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GE
off). With the agent setup below, decode went from 1.58 to 2.77 tok/s with unchanged answers
(same quality score). Prefill keeps rocBLAS.
## int4 experts on RDNA2
vLLM's int4 MoE kernel pads every token to a 16-row tile and multiplies in bf16, which RDNA2
lacks: a decode round took 2.7 ms for one token. On gfx10 prain runs int4 rounds of up to 4
tokens through its own GEMV-style kernel (`prain/kernels.py`, 0.25 ms) and larger rounds through
vLLM's kernel in fp16, twice as fast as bf16 there (a round that overflows fp16 is redone in
bf16). Every int4 expert copy goes through a pinned staging buffer: one DMA and one unpack kernel
instead of four pageable copies and a ten-kernel repack, asynchronous from every tier.
A/B on the live agent setup (int4, 16k context, 1 GiB KV, 450 slots): decode is the median of
3 x 48 tokens, prefill a fresh 2.3k-token prompt, quality the greedy answers to the six
`quality.py` prompts (identical to the reference answers, mean logprob of their tokens):
| | decode tok/s | first token, short | first token, 2.3k tokens | identical | mean logprob |
|---|---|---|---|---|---|
| before | 3.49 | 4.3 s | 27.5 s | 4 of 6 | -0.104 |
| decode kernel | 4.73 | 4.2 s | 27.3 s | 4 of 6 | -0.093 |
| staged copies | 3.67 | 3.7 s | 24.5 s | 4 of 6 | -0.098 |
| fp16 rounds | 3.84 | 4.0 s | 23.3 s | 4 of 6 | -0.102 |
| all three | 5.24 | 3.3 s | 20.6 s | 5 of 6 | -0.104 |
Repeated runs of one config differ by up to 5% in decode speed and 1% in prefill. Measured and
dropped: routing statistics on the CPU and a Triton GEMV for the fp32 router (both within noise),
and `[prefetch] mode = "pregate"`, which can now prefetch disk-tier experts too (4.9 vs 5.2 tok/s:
the link is already busy with demand copies).
## Trading quality for speed
`[miss] min_weight` (default 1) is the share of each token's router weight that must be
@@ -91,7 +120,7 @@ negative log-likelihood of the min_weight = 1 answers to six chat prompts (`scri
At 0.6 the answers stay correct but take other words; at 0.4 they show slips ("Hereon any of the
following one-liners"). Pinning only the hottest experts and never copying (`prain plan
--cache-slots 1`, min_weight 0) broke the model (NLL 9) for 25% more decode speed (measured before
the dense-layer fix).
the dense-layer fix). This table predates the int4 kernels above.
## Context length and coding agents
@@ -101,12 +130,22 @@ hold the tokens in flight (2 x `--max-num-batched-tokens`). `kv_cache_gib = 1.0`
tokens, 1.11 requests of 16k (measured).
With the example config (int4, 16k, tools) a cold 4.7k-token prompt took 70 s and a 12k-token
one 176 s: every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link.
one 176 s (before the int4 kernels, which cut prefill by a quarter): every 1024-token prefill chunk
needs nearly all experts, copied over the 2.5 GB/s link.
Agent turns reuse the prefix cache: turns 2 and 3 of a conversation found 4,704 and 4,736 of
~4,750 prompt tokens cached and took ~20 s each (24-token answers). vLLM keeps sliding-window
state only where earlier prompts ended, so the second conversation that shares just the system
prompt recomputes it once (176 s); later ones hit (4.5 s).
Requests with `prompt_logprobs` (perplexity tools) materialize [prompt tokens x 262k vocab] fp32
logits that vLLM's memory plan does not cover: on the 16k config even ~300-token prompts crashed
the engine. `scripts/quality.py score --greedy-only` skips them.
Bigger prefill chunks trade decode speed for prompt speed: at 16k, `--max-num-batched-tokens 2048`
needs `kv_cache_gib = 1.3` (the sliding layers hold two chunks), so 0.3 GiB less for experts (348
slots instead of 450). A 2.3k-token prompt then took 17.4 s instead of 20.6 s, and decode ran at
4.5 instead of 5.2 tok/s. Long agent prompts gain more (half the chunks, half the expert copies).
## Tune for a machine
python scripts/measure.py PACK_DIR/layer029.w13 --quant int4 --threads 16 > cost.json # server stopped
@@ -127,7 +166,7 @@ tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s
Serve with `PRAIN_PROFILE=/abs/dir`. Writing a number N into `dir/start` records the next N engine
steps: spans around the engine step, every module of the model and prain's MoE phases (route,
plan, copy: h2d/repack/d2d, round), each synced with the GPU. `python -m prain.profile dir`
plan, copy: memcpy/h2d/unpack for int4, round), each synced with the GPU. `python -m prain.profile dir`
writes `decode.svg`/`prefill.svg` flame graphs and prints self time and calls per step. The
syncs remove the CPU/GPU overlap: a decode step took 395 ms profiled, 307 ms without.
examples/gemma4-26b-a4b-rx6600xt.toml
+2 -2
View File
@@ -24,8 +24,8 @@ quant = "int4" # then quantized once into an int4 pack
# bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
default_tier = "disk" # runs on any RAM size; `prain tune` does better
cache_slots = 15 # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM
stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
stats_every = 64 # MoE forward passes between stats/trace writes
stats_path = "../models/packs/gemma4-26b-a4b/stats.json" # prain plan and the web monitor read it
stats_every = 8 # MoE forward passes between stats/trace writes (~3 ms each)
trace_path = "../models/packs/gemma4-26b-a4b/trace.jsonl.gz" # routing trace for sim/tune
[policy] # which expert a cache slot gives up
prain/cli.py
+4 -3
View File
@@ -169,8 +169,9 @@ def main(argv=None):
ui.add_argument("--port", type=int, default=8080)
ui.add_argument("--host", default="127.0.0.1",
help="address to bind; 0.0.0.0 for every interface (no login!)")
ui.add_argument("--config", help="placement file: [serve] max_requests and the "
"[draft] table (default: their defaults)")
ui.add_argument("--config", help="placement file: [serve] max_requests, the [draft] "
"table and stats_path for the expert monitor "
"(default: their defaults)")
args = parser.parse_args(argv)
if args.command == "check":
@@ -202,7 +203,7 @@ def main(argv=None):
config = load_config(args.config) if args.config else Config(pack_dir="") # defaults
web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft,
host=args.host)
host=args.host, stats_path=config.stats_path)
if __name__ == "__main__":
prain/policy.py
+4 -2
View File
@@ -236,8 +236,10 @@ COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wa
def write_stats(path, caches):
"""Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
"""Atomically write {"layers": {"<layer>": {"uses": [...], "vram": [experts in VRAM now],
"hits": n, ...}}}."""
data = {"layers": {str(layer): {"uses": c.uses, "vram": sorted({*c.pinned, *c.slots}),
**{k: getattr(c, k) for k in COUNTERS}}
for layer, c in sorted(caches.items())}}
tmp = f"{path}.tmp"
with open(tmp, "w") as f:
prain/web.html
+716
View File
@@ -0,0 +1,716 @@
<!doctype html>
<html lang="en">
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1, viewport-fit=cover">
<meta name="color-scheme" content="light dark">
<title>prain</title>
<style>
:root {
--bg: #f1f3f6; --surface: #fff; --ink: #161b22; --muted: #5b6673; --faint: #8a94a0;
--line: #dde2e8; --bubble: #e5e9ef; --code: #f6f7f9; --focus: #7c4dcc; --live: #f26b1d;
--ok: #23946a; --bad: #c2362f;
--ui: system-ui, -apple-system, "Segoe UI", Roboto, "Noto Sans", Cantarell, sans-serif;
--display: ui-rounded, "SF Pro Rounded", Nunito, system-ui, -apple-system, sans-serif;
--mono: ui-monospace, "SF Mono", "JetBrains Mono", "Cascadia Code", "DejaVu Sans Mono", Menlo, Consolas, monospace;
}
@media (prefers-color-scheme: dark) {
:root { --bg: #0f1216; --surface: #161a20; --ink: #e6eaf0; --muted: #9aa4b1; --faint: #6d7784;
--line: #262d36; --bubble: #232a34; --code: #1b2027; --focus: #a58bea; }
}
* { box-sizing: border-box; }
html, body { height: 100%; margin: 0; }
body { background: var(--bg); color: var(--ink); font: 16px/1.6 var(--ui); -webkit-font-smoothing: antialiased; }
button { font: inherit; color: inherit; }
:focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; }
[hidden] { display: none !important; }
.app { display: grid; grid-template-columns: minmax(0, 1fr) 340px; height: 100dvh; }
.chat { display: flex; flex-direction: column; min-width: 0; height: 100dvh; }
/* header */
.bar { display: flex; align-items: center; justify-content: space-between; gap: 12px; padding: 12px 20px; }
.brand { display: flex; align-items: center; gap: 10px; min-width: 0; }
.mark { display: grid; grid-template-columns: repeat(4, 6px); gap: 2px; }
.mark i { width: 6px; height: 6px; border-radius: 2px; background: var(--line); transition: background .3s; }
.mark i.on { background: var(--live); }
.word { font: 700 20px/1 var(--display); letter-spacing: -.01em; }
.model { font: 13px var(--mono); color: var(--faint); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
.bar-actions { display: flex; gap: 8px; align-items: center; }
.ghost { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 6px 14px; font-size: 14px; cursor: pointer; }
.ghost:hover { border-color: var(--faint); }
.pill { display: none; align-items: center; gap: 8px; }
.dot { width: 8px; height: 8px; border-radius: 50%; background: var(--faint); flex: none; }
.dot.ok { background: var(--ok); }
.dot.live { background: var(--live); animation: pulse 1.2s ease-in-out infinite; }
.dot.bad { background: var(--bad); }
@keyframes pulse { 50% { opacity: .35; } }
/* thread */
.log { flex: 1; overflow-y: auto; padding: 8px 20px 24px; overscroll-behavior: contain; }
.thread { max-width: 760px; margin: 0 auto; }
.empty { max-width: 560px; margin: 14vh auto 0; }
.empty h1 { font: 650 clamp(28px, 4.5vw, 38px)/1.15 var(--display); letter-spacing: -.015em; margin: 0 0 10px; }
.empty p { color: var(--muted); margin: 0 0 6px; }
.empty .live-line { font-size: 14px; color: var(--faint); margin-bottom: 22px; font-variant-numeric: tabular-nums; }
.chips { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 22px; }
.chip { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 8px 14px; font-size: 14px; cursor: pointer; text-align: left; }
.chip:hover { border-color: var(--faint); }
.msg { margin: 22px 0; }
.msg.user { display: flex; justify-content: flex-end; }
.bubble { background: var(--bubble); padding: 10px 16px; border-radius: 20px 20px 6px 20px; max-width: min(85%, 600px); white-space: pre-wrap; overflow-wrap: anywhere; }
.body { overflow-wrap: anywhere; }
.body > :first-child { margin-top: 0; }
.body > :last-child { margin-bottom: 0; }
.body p { margin: 0 0 .8em; }
.body h3, .body h4, .body h5 { font: 650 1.08em/1.3 var(--display); margin: 1.3em 0 .4em; }
.body ul, .body ol { padding-left: 1.4em; margin: 0 0 .8em; }
.body li { margin: .25em 0; }
.body blockquote { border-left: 3px solid var(--line); margin: .8em 0; padding-left: 12px; color: var(--muted); }
.body hr { border: 0; border-top: 1px solid var(--line); margin: 1.2em 0; }
.body a { color: var(--focus); }
.body code { font: .88em var(--mono); background: var(--code); border: 1px solid var(--line); padding: .05em .35em; border-radius: 6px; }
.code { border: 1px solid var(--line); background: var(--code); border-radius: 12px; margin: .9em 0; overflow: hidden; }
.code-head { display: flex; justify-content: space-between; align-items: center; padding: 4px 6px 4px 14px; border-bottom: 1px solid var(--line); font: 12px var(--mono); color: var(--faint); }
.code pre { margin: 0; padding: 12px 14px; overflow-x: auto; font: 13.5px/1.55 var(--mono); tab-size: 4; }
.code code { background: none; border: 0; padding: 0; font-size: inherit; }
.table { overflow-x: auto; margin: .9em 0; }
.body table { border-collapse: collapse; font-size: 14px; }
.body th, .body td { border-bottom: 1px solid var(--line); padding: 6px 12px 6px 0; text-align: left; vertical-align: top; }
.small { border: 0; background: none; border-radius: 6px; padding: 3px 8px; font-size: 12.5px; color: var(--faint); cursor: pointer; }
.small:hover { color: var(--ink); background: var(--bubble); }
.meta { display: flex; flex-wrap: wrap; gap: 4px 12px; align-items: center; margin-top: 6px; font-size: 12.5px; color: var(--faint); font-variant-numeric: tabular-nums; }
.note { display: inline-block; font-size: 13px; color: var(--muted); background: var(--bubble); border-radius: 10px; padding: 5px 10px; margin-bottom: 10px; }
.note.bad { color: var(--bad); }
.is-draft .body, .rest { color: var(--faint); }
.typing::after { content: ""; display: inline-block; width: .5em; height: 1.05em; margin-left: 3px; vertical-align: -.18em; border-radius: 2px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
/* composer */
.composer { padding: 0 20px calc(14px + env(safe-area-inset-bottom)); }
.field { max-width: 760px; margin: 0 auto; display: flex; align-items: flex-end; gap: 8px; background: var(--surface); border: 1px solid var(--line); border-radius: 22px; padding: 7px 7px 7px 16px; box-shadow: 0 10px 30px -18px rgb(20 30 45 / .35); }
.field:focus-within { border-color: color-mix(in srgb, var(--focus) 50%, var(--line)); }
textarea { flex: 1; border: 0; outline: 0; resize: none; background: transparent; color: inherit; font: 16px/1.5 var(--ui); max-height: 40vh; padding: 6px 0; }
.send { width: 36px; height: 36px; flex: none; border-radius: 50%; border: 0; background: var(--ink); color: var(--surface); display: grid; place-items: center; cursor: pointer; }
.send:disabled { opacity: .3; cursor: default; }
.send svg { width: 18px; height: 18px; }
.hint { max-width: 760px; margin: 6px auto 0; min-height: 18px; padding: 0 16px; font-size: 12.5px; color: var(--faint); display: flex; align-items: center; gap: 8px; }
.cells { display: inline-grid; grid-template-columns: repeat(8, 5px); gap: 2px; }
.cells i { width: 5px; height: 5px; border-radius: 1.5px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
.cells i:nth-child(2n) { animation-delay: .25s; } .cells i:nth-child(3n) { animation-delay: .5s; }
/* monitor */
.monitor { background: var(--surface); border-left: 1px solid var(--line); overflow-y: auto; padding: 18px 20px 28px; font-size: 14px; }
.m-top { display: flex; justify-content: space-between; align-items: start; }
.m-title { font: 650 17px/1.2 var(--display); }
.m-sub { font: 12.5px var(--mono); color: var(--faint); margin-top: 2px; }
.close { display: none; }
.state { display: flex; align-items: center; gap: 8px; margin: 14px 0 6px; color: var(--muted); }
.speed { display: flex; align-items: baseline; gap: 6px; }
.speed b { font: 650 40px/1 var(--display); letter-spacing: -.02em; font-variant-numeric: tabular-nums; }
.speed.idle b { color: var(--faint); }
.speed span { color: var(--faint); }
.spark { width: 100%; height: 28px; margin: 6px 0 16px; display: block; }
.spark path { fill: none; stroke: var(--live); stroke-width: 1.5; vector-effect: non-scaling-stroke; }
.scan { background: #16121d; color: #c9c1d6; border-radius: 14px; padding: 12px 12px 10px; }
.scan-head { display: flex; justify-content: space-between; font-size: 12.5px; margin-bottom: 8px; font-variant-numeric: tabular-nums; }
.scan-head b { font-weight: 600; color: #f2edf8; }
.scan canvas { display: block; width: 100%; aspect-ratio: 128 / 64; image-rendering: pixelated; border-radius: 4px; background: #2c2538; }
.axis { display: flex; justify-content: space-between; font: 11px var(--mono); color: #8f86a0; margin-top: 4px; }
.legend { display: flex; flex-wrap: wrap; gap: 4px 12px; margin: 8px 0 0; padding: 0; list-style: none; font-size: 12px; }
.legend i { display: inline-block; width: 9px; height: 9px; border-radius: 2px; margin-right: 5px; vertical-align: -1px; }
.scan-empty { font-size: 12.5px; color: #8f86a0; padding: 18px 4px; }
.kv { margin: 16px 0 0; }
.kv div { display: flex; justify-content: space-between; align-items: baseline; gap: 12px; padding: 7px 0; border-bottom: 1px solid var(--line); }
.kv dt { color: var(--muted); }
.kv dd { margin: 0; font-variant-numeric: tabular-nums; text-align: right; }
.kv small { color: var(--faint); }
.meter { display: block; width: 84px; height: 4px; border-radius: 2px; background: var(--bubble); margin-top: 5px; margin-left: auto; overflow: hidden; }
.meter i { display: block; height: 100%; width: 0; background: var(--ink); transition: width .4s; }
.m-h { font: 650 13px/1 var(--display); color: var(--faint); text-transform: uppercase; letter-spacing: .06em; margin: 24px 0 0; }
.scrim { display: none; }
@media (max-width: 999px) {
.app { grid-template-columns: 1fr; }
.pill, .close { display: inline-flex; }
.monitor { position: fixed; top: 0; right: 0; bottom: 0; width: min(360px, 92vw); z-index: 20; transform: translateX(100%); transition: transform .25s ease; box-shadow: -20px 0 40px -20px rgb(0 0 0 / .35); padding-top: calc(18px + env(safe-area-inset-top)); }
.monitor.open { transform: none; }
.scrim.open { display: block; position: fixed; inset: 0; z-index: 10; background: rgb(10 14 20 / .35); }
}
@media (max-width: 560px) {
.bar { padding: 10px 14px; } .log { padding: 6px 14px 18px; } .composer { padding-left: 10px; padding-right: 10px; }
.model, .hint .keys { display: none; } .empty { margin-top: 8vh; }
}
@media (prefers-reduced-motion: reduce) {
*, *::after { animation: none !important; transition: none !important; }
}
</style>
<div class="app">
<main class="chat">
<header class="bar">
<div class="brand">
<span class="mark" id="mark" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span>
<span class="word">prain</span>
<span class="model" id="model"></span>
</div>
<div class="bar-actions">
<button class="ghost pill" id="pill" aria-controls="monitor" aria-expanded="false">
<span class="dot" id="pill-dot"></span><span id="pill-text">Monitor</span></button>
<button class="ghost" id="new-chat">New chat</button>
</div>
</header>
<section class="log" id="log" role="log">
<div class="thread" id="thread">
<div class="empty" id="empty">
<h1>Ready when you are</h1>
<p>prain moves the experts your message needs into VRAM while you type, so the answer starts sooner.</p>
<div class="live-line" id="empty-live"></div>
<div class="chips">
<button class="chip">Explain what git rebase -i HEAD~3 does</button>
<button class="chip">Write a bash one-liner that counts lines of code per file type</button>
<button class="chip">On Gentoo, how do I rebuild the packages affected by a USE flag change?</button>
</div>
</div>
</div>
</section>
<form class="composer" id="composer">
<div class="field">
<textarea id="draft" rows="1" placeholder="Message prain" aria-label="Message"></textarea>
<button class="send" id="send" type="submit" aria-label="Send" disabled>
<svg viewBox="0 0 24 24" aria-hidden="true"><path id="send-icon" fill="currentColor" d="M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z"/></svg>
</button>
</div>
<div class="hint">
<span id="warm" hidden><span class="cells" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span> Loading experts for your message</span>
<span class="keys" id="keys">Enter sends · Shift+Enter adds a line</span>
</div>
</form>
</main>
<div class="scrim" id="scrim"></div>
<aside class="monitor" id="monitor" aria-label="Model monitor">
<div class="m-top">
<div><div class="m-title">Monitor</div><div class="m-sub" id="m-model">connecting…</div></div>
<button class="ghost close" id="close" aria-label="Close the monitor">Close</button>
</div>
<div class="state"><span class="dot" id="dot"></span><span id="state">Connecting…</span></div>
<div class="speed idle" id="speed-box"><b id="speed">–</b><span id="speed-unit">tokens/s</span></div>
<svg class="spark" viewBox="0 0 60 28" preserveAspectRatio="none" aria-hidden="true"><path id="spark"/></svg>
<section class="scan" aria-label="Experts by layer">
<div class="scan-head"><span>Experts</span><span id="scan-count"></span></div>
<div id="scan-body">
<canvas id="scan" width="128" height="30" role="img" aria-label="Expert map"></canvas>
<div class="axis"><span>layer 1 · expert 1</span><span id="axis-end"></span></div>
<ul class="legend">
<li><i style="background:#6a35a0"></i>In VRAM</li>
<li><i style="background:#ff8a2a"></i>Used now</li>
<li><i style="background:#45d3c7"></i>Copied in</li>
</ul>
</div>
<div class="scan-empty" id="scan-empty" hidden>The map appears once prain writes expert stats: set stats_path in the placement file and pass it with --config.</div>
</section>
<dl class="kv">
<div><dt>VRAM hits</dt><dd id="hits">–</dd></div>
<div><dt>Copies</dt><dd id="copies">–</dd></div>
<div><dt>Context in use</dt><dd><span id="kv">–</span><span class="meter"><i id="kv-bar"></i></span></dd></div>
<div><dt>Prompt cache reuse</dt><dd id="prefix">–</dd></div>
<div><dt>Requests</dt><dd id="reqs">–</dd></div>
</dl>
<h3 class="m-h">GPU</h3>
<dl class="kv">
<div><dt>VRAM</dt><dd><span id="vram">–</span><span class="meter"><i id="vram-bar"></i></span></dd></div>
<div><dt>Load</dt><dd id="load">–</dd></div>
<div><dt>Temperature</dt><dd id="temp">–</dd></div>
<div><dt>Power</dt><dd id="power">–</dd></div>
<div><dt>Fan</dt><dd id="fan">–</dd></div>
</dl>
</aside>
</div>
<script>
const $ = (id) => document.getElementById(id);
const log = $("log"), thread = $("thread"), draft = $("draft"), send = $("send");
const coarse = matchMedia("(pointer: coarse)").matches;
const newId = () => Math.random().toString(36).slice(2) + Date.now().toString(36);
let chat = load() || {session: newId(), messages: []};
let streaming = false, controller = null;
function load() {
try { return JSON.parse(localStorage.getItem("prain.chat")); } catch { return null; }
}
function save() {
try { localStorage.setItem("prain.chat", JSON.stringify(chat)); } catch {}
}
const plain = () => chat.messages.map(({role, content}) => ({role, content}));
/* ---------- markdown: escape first, then a small safe subset ---------- */
const esc = (s) => s.replace(/[&<>"']/g, (c) => ({"&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;"}[c]));
function inline(text) {
return text.split(/(`[^`\n]+`)/g).map((part, i) => i % 2
? `<code>${esc(part.slice(1, -1))}</code>`
: esc(part)
.replace(/\[([^\]]+)\]\((https?:\/\/[^\s)]+)\)/g, '<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>')
.replace(/\*\*([^*]+)\*\*/g, "<strong>$1</strong>")
.replace(/(^|[^*\w])\*([^*\s][^*]*)\*(?![*\w])/g, "$1<em>$2</em>")
).join("");
}
const ITEM = /^\s*([-*+]|\d+[.)])\s+/;
function md(src) {
const lines = src.split("\n"), out = [];
let i = 0, para = [];
const flush = () => { if (para.length) out.push(`<p>${para.map(inline).join("<br>")}</p>`); para = []; };
while (i < lines.length) {
const line = lines[i];
let m;
if ((m = line.match(/^\s*```\s*([\w+#.-]*)/))) {
flush();
const code = [];
for (i++; i < lines.length && !/^\s*```\s*$/.test(lines[i]); i++) code.push(lines[i]);
i++;
out.push(`<div class="code"><div class="code-head"><span>${esc(m[1] || "code")}</span>` +
`<button class="small" data-copy>Copy</button></div><pre><code>${esc(code.join("\n"))}</code></pre></div>`);
} else if ((m = line.match(/^(#{1,6})\s+(.*)/))) {
flush();
const level = Math.min(m[1].length + 2, 5);
out.push(`<h${level}>${inline(m[2])}</h${level}>`);
i++;
} else if (/^\s*\|.*\|\s*$/.test(line) && /^\s*\|?\s*:?-{2,}/.test(lines[i + 1] || "")) {
flush();
const rows = [];
while (i < lines.length && /^\s*\|.*\|\s*$/.test(lines[i])) rows.push(lines[i++]);
const cells = (row, tag) => row.trim().replace(/^\||\|$/g, "").split("|")
.map((c) => `<${tag}>${inline(c.trim())}</${tag}>`).join("");
out.push(`<div class="table"><table><thead><tr>${cells(rows[0], "th")}</tr></thead><tbody>` +
rows.slice(2).map((r) => `<tr>${cells(r, "td")}</tr>`).join("") + "</tbody></table></div>");
} else if (ITEM.test(line)) {
flush();
const ordered = /^\s*\d/.test(line), items = [];
const start = ordered ? parseInt(line.trim(), 10) : 1;
while (i < lines.length && (ITEM.test(lines[i]) || (items.length && /^\s{2,}\S/.test(lines[i])))) {
if (ITEM.test(lines[i])) items.push(lines[i].replace(ITEM, ""));
else items[items.length - 1] += " " + lines[i].trim();
i++;
}
const tag = ordered ? "ol" : "ul";
out.push(`<${tag}${ordered && start !== 1 ? ` start="${start}"` : ""}>` +
items.map((t) => `<li>${inline(t)}</li>`).join("") + `</${tag}>`);
} else if (/^\s*>/.test(line)) {
flush();
const quote = [];
while (i < lines.length && /^\s*>/.test(lines[i])) quote.push(lines[i++].replace(/^\s*>\s?/, ""));
out.push(`<blockquote>${md(quote.join("\n"))}</blockquote>`);
} else if (/^\s*([-*_])(\s*\1){2,}\s*$/.test(line)) {
flush();
out.push("<hr>");
i++;
} else if (!line.trim()) {
flush();
i++;
} else {
para.push(line);
i++;
}
}
flush();
return out.join("");
}
/* ---------- thread ---------- */
function nearBottom() { return log.scrollHeight - log.scrollTop - log.clientHeight < 120; }
function toBottom() { log.scrollTop = log.scrollHeight; }
function el(tag, cls, text) {
const node = document.createElement(tag);
if (cls) node.className = cls;
if (text !== undefined) node.textContent = text;
return node;
}
function addUser(text) {
const row = el("div", "msg user");
row.appendChild(el("div", "bubble", text));
thread.appendChild(row);
}
function addBot() {
const row = el("div", "msg bot");
const view = {row, note: el("div", "note"), body: el("div", "body"), meta: el("div", "meta")};
view.note.hidden = true;
view.meta.hidden = true;
row.append(view.note, view.body, view.meta);
thread.appendChild(row);
return view;
}
function paint(view, text, rest, live) {
view.body.innerHTML = md(text);
if (rest) {
const target = view.body.lastElementChild && view.body.lastElementChild.tagName === "P"
? view.body.lastElementChild : view.body;
target.appendChild(el("span", "rest", rest));
}
view.body.classList.toggle("typing", live);
}
function finish(view, message) {
view.meta.hidden = false;
view.meta.replaceChildren(el("span", "", message.meta || ""));
const copy = el("button", "small", "Copy");
copy.addEventListener("click", () => copied(copy, message.content));
view.meta.appendChild(copy);
}
function render() {
thread.querySelectorAll(".msg").forEach((n) => n.remove());
$("empty").hidden = chat.messages.length > 0;
for (const m of chat.messages) {
if (m.role === "user") { addUser(m.content); continue; }
const view = addBot();
view.row.classList.toggle("is-draft", !!m.draft);
paint(view, m.content, "", false);
finish(view, m);
}
toBottom();
}
async function copyText(text) {
try { await navigator.clipboard.writeText(text); return true; } catch {}
const area = el("textarea"); // plain http: no Clipboard API
area.value = text;
area.setAttribute("readonly", "");
area.style.cssText = "position:fixed;top:0;opacity:0";
document.body.appendChild(area);
area.select();
area.setSelectionRange(0, text.length);
let ok = false;
try { ok = document.execCommand("copy"); } catch {}
area.remove();
return ok;
}
async function copied(button, text) {
const ok = await copyText(text);
button.textContent = ok ? "Copied" : "Copy failed";
setTimeout(() => (button.textContent = "Copy"), 1500);
}
thread.addEventListener("click", (event) => {
const button = event.target.closest("[data-copy]");
if (button) copied(button, button.closest(".code").querySelector("pre").textContent);
const chip = event.target.closest(".chip");
if (chip) { draft.value = chip.textContent; draft.focus(); draft.dispatchEvent(new Event("input")); }
});
/* ---------- sending ---------- */
const ICON_SEND = "M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z";
const ICON_STOP = "M7 7h10v10H7z";
function setStreaming(on) {
streaming = on;
log.setAttribute("aria-busy", on);
$("send-icon").setAttribute("d", on ? ICON_STOP : ICON_SEND);
send.setAttribute("aria-label", on ? "Stop" : "Send");
send.disabled = !on && !draft.value.trim();
if (!on) lightMark(0);
}
function lightMark(n) {
const cells = [...$("mark").children];
cells.forEach((c) => c.classList.remove("on"));
for (let k = 0; k < n; k++) cells[Math.floor(Math.random() * cells.length)].classList.add("on");
}
async function submit() {
const text = draft.value.trim();
if (streaming || !text) return;
const conv = chat; // "New chat" mid-answer must not receive this answer
$("empty").hidden = true;
conv.messages.push({role: "user", content: text});
addUser(text);
save();
draft.value = "";
autosize();
hideWarm();
setStreaming(true);
toBottom();
const view = addBot();
const st = {mode: "full", draft: "", final: "", refined: false};
let start = performance.now(), first = 0, pieces = 0, stopped = false, failure = null;
const show = () => {
const follow = nearBottom();
const text = st.mode === "draft" ? st.draft : st.final;
const rest = st.mode === "refine" && st.draft.startsWith(st.final) ? st.draft.slice(st.final.length) : "";
view.row.classList.toggle("is-draft", st.mode === "draft");
paint(view, text, rest, true);
if (follow) toBottom();
};
const note = (text, bad) => { view.note.hidden = !text; view.note.textContent = text || ""; view.note.classList.toggle("bad", !!bad); };
controller = new AbortController();
try {
const response = await fetch("/chat", {method: "POST", signal: controller.signal,
body: JSON.stringify({session: chat.session, messages: plain()})});
if (!response.ok) throw new Error(`HTTP ${response.status}: ${(await response.text()).slice(0, 200)}`);
const reader = response.body.getReader(), decoder = new TextDecoder();
let buffer = "";
show();
for (;;) {
const {done, value} = await reader.read();
if (done) break;
buffer += decoder.decode(value, {stream: true});
const lines = buffer.split("\n");
buffer = lines.pop();
for (const line of lines) {
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
const data = JSON.parse(line.slice(6));
if (data.prain === "draft") {
st.mode = "draft";
note("Another answer is using the GPU, so this is a draft from the experts already in VRAM.");
} else if (data.prain === "refine") {
st.mode = "refine";
st.refined = true;
start = performance.now(), first = 0, pieces = 0;
note("Refining the draft with all experts…");
} else if (data.prain === "error") {
st.mode = "draft";
st.refined = false;
note(`Refining failed (HTTP ${data.status}); this is the draft.`, true);
} else {
const piece = data.choices?.[0]?.delta?.content; // usage chunks have no choices
if (!piece) continue;
if (pieces++ === 0) first = performance.now();
if (st.mode === "draft") st.draft += piece; else st.final += piece;
lightMark(2);
}
show();
}
}
} catch (error) {
if (error.name === "AbortError") stopped = true; else failure = error;
}
controller = null;
const content = st.mode === "draft" ? st.draft : st.final;
if (conv !== chat) { setStreaming(false); return; }
if (failure && !content) {
view.row.remove();
conv.messages.pop();
save();
const box = addBot();
box.note.hidden = false;
box.note.classList.add("bad");
box.note.textContent = `The model didn't answer (${failure.message}). Your message is back in the box: ` +
"send it again once the monitor shows the model as ready.";
draft.value = text;
autosize();
setStreaming(false);
return;
}
const seconds = (performance.now() - first) / 1000;
const parts = [];
if (first) parts.push(`${((first - start) / 1000).toFixed(1)} s to first token`);
if (pieces > 1) parts.push(`${((pieces - 1) / seconds).toFixed(1)} tokens/s`);
if (st.refined && st.draft) {
let same = 0;
while (same < st.draft.length && st.draft[same] === st.final[same]) same++;
parts.push(`refined, kept ${Math.round(100 * same / st.draft.length)}% of the draft`);
} else if (st.mode === "draft") parts.push("draft only");
if (stopped) parts.push("stopped");
if (failure) parts.push("cut off: " + failure.message);
const message = {role: "assistant", content, meta: parts.join(" · "), draft: st.mode === "draft"};
conv.messages.push(message);
save();
if (st.refined || !view.note.classList.contains("bad")) note("");
paint(view, content, "", false);
finish(view, message);
setStreaming(false);
}
$("composer").addEventListener("submit", (event) => {
event.preventDefault();
if (streaming) controller && controller.abort(); else submit();
});
draft.addEventListener("keydown", (event) => {
if (event.key === "Enter" && !event.shiftKey && !event.isComposing && !coarse) {
event.preventDefault();
submit();
}
});
if (coarse) $("keys").textContent = "Return adds a line";
function autosize() {
draft.style.height = "auto";
draft.style.height = Math.min(draft.scrollHeight, innerHeight * 0.4) + "px";
if (!streaming) send.disabled = !draft.value.trim();
}
let warmTimer = 0;
function hideWarm() { clearTimeout(warmTimer); $("warm").hidden = true; $("keys").hidden = false; }
draft.addEventListener("input", () => {
autosize();
if (streaming || !draft.value.trim()) return;
fetch("/warm", {method: "POST", body: JSON.stringify({session: chat.session, messages: plain(), draft: draft.value})})
.then((response) => {
if (response.status !== 202) return;
$("warm").hidden = false;
$("keys").hidden = true;
clearTimeout(warmTimer);
warmTimer = setTimeout(hideWarm, 2500);
}).catch(() => {});
});
$("new-chat").addEventListener("click", () => {
if (controller) controller.abort();
chat = {session: newId(), messages: []};
save();
render();
draft.focus();
});
/* ---------- monitor ---------- */
const monitor = $("monitor"), scrim = $("scrim"), pill = $("pill");
function openMonitor(on) {
monitor.classList.toggle("open", on);
scrim.classList.toggle("open", on);
pill.setAttribute("aria-expanded", on);
}
pill.addEventListener("click", () => openMonitor(!monitor.classList.contains("open")));
$("close").addEventListener("click", () => openMonitor(false));
scrim.addEventListener("click", () => openMonitor(false));
addEventListener("keydown", (event) => { if (event.key === "Escape") openMonitor(false); });
const fmt = (n, digits = 0) => n.toLocaleString(undefined, {maximumFractionDigits: digits, minimumFractionDigits: digits});
const GIB = 2 ** 30;
const history = []; // polls of the last minute
const speeds = [];
let lastSpeed = 0, scanState = null;
function set(id, text) { $(id).textContent = text; }
function meter(id, share) { $(id).style.width = `${Math.max(0, Math.min(1, share)) * 100}%`; }
function setState(kind, text) {
for (const id of ["dot", "pill-dot"]) $(id).className = "dot " + kind;
set("state", text);
}
const SCAN = {off: [44, 37, 56], vram: [106, 53, 160], warm: [255, 107, 26], hot: [255, 212, 71], copy: [69, 211, 199]};
const mix = (a, b, t) => [0, 1, 2].map((k) => a[k] + (b[k] - a[k]) * t);
function drawScan(experts) {
const canvas = $("scan");
const layers = Object.keys(experts.layers).map(Number).sort((a, b) => a - b).map((k) => experts.layers[k]);
const rows = layers.length, cols = layers[0].uses.length, n = rows * cols;
if (!scanState || scanState.heat.length !== n) {
canvas.width = cols;
canvas.height = rows;
scanState = {heat: new Float32Array(n), flash: new Float32Array(n), uses: null, vram: null};
}
const uses = new Float64Array(n), vram = new Uint8Array(n);
layers.forEach((layer, l) => {
layer.uses.forEach((u, e) => (uses[l * cols + e] = u));
(layer.vram || []).forEach((e) => (vram[l * cols + e] = 1));
});
const {heat, flash} = scanState, ctx = canvas.getContext("2d"), image = ctx.createImageData(cols, rows);
let resident = 0, active = 0;
for (let k = 0; k < n; k++) {
const fired = scanState.uses ? uses[k] - scanState.uses[k] : 0;
heat[k] = Math.min(1, heat[k] * 0.5 + (fired > 0 ? 0.55 + 0.1 * Math.min(fired, 4) : 0));
flash[k] = scanState.vram && vram[k] && !scanState.vram[k] ? 1 : flash[k] * 0.4;
resident += vram[k];
active += fired > 0;
let c = vram[k] ? SCAN.vram : SCAN.off;
if (heat[k] > 0.02) c = mix(c, mix(SCAN.warm, SCAN.hot, heat[k]), Math.min(1, heat[k] * 1.5));
if (flash[k] > 0.02) c = mix(c, SCAN.copy, flash[k] * 0.85);
image.data.set([c[0], c[1], c[2], 255], k * 4);
}
ctx.putImageData(image, 0, 0);
scanState.uses = uses;
scanState.vram = vram;
set("scan-count", `${fmt(resident)} of ${fmt(n)} in VRAM`);
set("axis-end", `layer ${rows} · expert ${cols}`);
canvas.setAttribute("aria-label", `${resident} of ${n} experts in VRAM; ${active} used since the last update`);
set("empty-live", `${fmt(resident)} of ${fmt(n)} experts are in VRAM right now.`);
}
function totals(experts) {
const t = {hits: 0, misses: 0, cpu: 0};
for (const layer of Object.values(experts.layers)) for (const k in t) t[k] += layer[k] || 0;
return t;
}
async function poll() {
let s;
try {
const response = await fetch("/status", {cache: "no-store"});
s = await response.json();
} catch {
setState("bad", "Can't reach prain web");
set("pill-text", "Offline");
return;
}
set("model", s.model);
set("m-model", s.model);
const v = s.vllm, now = s.time;
const entry = {t: now, gen: v ? v.generation_tokens : null, ...(s.experts ? totals(s.experts) : {})};
history.push(entry);
while (history.length > 2 && now - history[0].t > 60) history.shift();
if (!v) {
setState("bad", "Model offline: starting or restarting");
set("pill-text", "Offline");
} else {
const busy = v.running > 0;
const prev = history[history.length - 2];
const dgen = prev && prev.gen != null ? entry.gen - prev.gen : 0; // < 0 after a restart
const rate = dgen > 0 ? dgen / Math.max(entry.t - prev.t, 0.1) : 0;
if (busy && rate) lastSpeed = lastSpeed ? lastSpeed * 0.4 + rate * 0.6 : rate;
speeds.push(busy ? rate : 0);
if (speeds.length > 60) speeds.shift();
$("speed-box").classList.toggle("idle", !busy || !rate);
set("speed", lastSpeed ? fmt(lastSpeed, 1) : "–");
set("speed-unit", busy && rate ? "tokens/s" : lastSpeed ? "tokens/s, last answer" : "tokens/s");
const top = Math.max(...speeds, 1);
$("spark").setAttribute("d", speeds.map((y, x) => `${x ? "L" : "M"}${x + 60 - speeds.length} ${27 - (y / top) * 25}`).join(""));
const waiting = v.waiting ? `, ${fmt(v.waiting)} waiting` : "";
if (busy) setState("live", `${rate ? "Generating" : "Reading the prompt"}${waiting}`);
else setState("ok", "Ready");
set("pill-text", !busy ? "Ready" : rate ? `${fmt(lastSpeed, 1)} tok/s` : "Working");
set("reqs", `${fmt(v.running)} running${waiting}`);
set("kv", `${fmt(v.kv_usage * 100)}%`);
meter("kv-bar", v.kv_usage);
set("prefix", v.prefix_queries ? `${fmt(100 * v.prefix_hits / v.prefix_queries)}%` : "–");
}
if (s.experts && Object.keys(s.experts.layers).length) {
$("scan-empty").hidden = true;
$("scan-body").hidden = false;
drawScan(s.experts);
const old = history.find((h) => h.hits !== undefined);
const d = old ? {hits: entry.hits - old.hits, misses: entry.misses - old.misses, cpu: entry.cpu - old.cpu} : null;
const routed = d && d.hits >= 0 && d.misses >= 0 ? d.hits + d.misses + d.cpu : 0; // < 0 after a restart
if (routed > 0) {
set("hits", `${fmt(100 * d.hits / routed)}%`);
set("copies", `${fmt(d.misses / Math.max(entry.t - old.t, 1), 0)}/s`);
} else {
const all = entry.hits + entry.misses + entry.cpu;
set("hits", all ? `${fmt(100 * entry.hits / all)}% overall` : "–");
set("copies", `${fmt(entry.misses)} so far`);
}
} else {
$("scan-empty").hidden = false;
$("scan-body").hidden = true;
set("empty-live", "");
}
const g = s.gpu;
if (g) {
set("vram", `${fmt(g.vram_used / GIB, 1)} / ${fmt(g.vram_total / GIB, 1)} GiB`);
meter("vram-bar", g.vram_used / g.vram_total);
set("load", g.busy == null ? "–" : `${g.busy}%`);
set("temp", g.temp == null ? "–" : `${fmt(g.temp)} °C`);
set("power", g.power == null ? "–" : `${fmt(g.power)} W`);
set("fan", g.fan == null ? "–" : `${fmt(g.fan)} rpm`);
}
}
let timer = 0;
function schedule() {
clearTimeout(timer);
if (document.visibilityState === "visible") timer = setTimeout(() => poll().finally(schedule), 1500);
}
document.addEventListener("visibilitychange", () => { if (document.visibilityState === "visible") poll().finally(schedule); });
render();
autosize();
poll().finally(schedule);
if (!coarse) draft.focus();
</script>
prain/web.py
+84 -113
View File
@@ -1,10 +1,12 @@
"""prain web: a one-page browser chat that warms experts up while you type.
"""prain web: a browser chat that warms experts up while you type, with a model monitor.
A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes
{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
"refine" | "error"} events of its own. Binds to 127.0.0.1 unless --host says otherwise
(0.0.0.0 serves every interface, without a login).
A small stdlib proxy in front of vLLM. GET / serves the page (web.html). GET /status returns
what the monitor shows: vLLM's counters (/metrics), the GPU (amdgpu sysfs) and prain's expert
stats (the placement file's stats_path, with the experts in VRAM). POST /warm takes
{"session", "messages", "draft"} and returns at once: 202 when a warm-up starts, 204 when not.
POST /chat takes {"session", "messages"} and streams server-sent events: vLLM's chunks, plus
{"prain": "draft" | "refine" | "error"} events of its own. Binds to 127.0.0.1 unless --host
says otherwise (0.0.0.0 serves every interface, without a login).
Several browsers at once ([serve] and [draft] tables of the placement file):
- At most draft.full_requests answers run with all their experts. An answer that arrives
@@ -19,127 +21,88 @@ Several browsers at once ([serve] and [draft] tables of the placement file):
cancels it.
"""
import glob
import json
import os
import threading
import time
import urllib.error
import urllib.request
from collections import OrderedDict, deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from prain.chat import Warmer, chat_body, draft_body, first_model, post, send_request, warm_body
from prain.config import Draft, Serve
MAX_SESSIONS = 64
PAGE = """<!doctype html>
<meta charset="utf-8">
<title>prain</title>
<style>
body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
.user { color: #666; margin-top: 1rem; }
.assistant { white-space: pre-wrap; }
.draft { color: #999; }
textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
#status { color: #888; font-size: 13px; }
</style>
<div id="log"></div>
<textarea id="draft" autofocus
placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
<div id="status"></div>
<script>
const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
const messages = [];
const log = document.getElementById("log");
const draft = document.getElementById("draft");
const status = document.getElementById("status");
let streaming = false;
PAGE = (Path(__file__).parent / "web.html").read_bytes()
METRICS = {"vllm:num_requests_running": "running", "vllm:num_requests_waiting": "waiting",
"vllm:kv_cache_usage_perc": "kv_usage", "vllm:prefix_cache_queries_total":
"prefix_queries", "vllm:prefix_cache_hits_total": "prefix_hits",
"vllm:prompt_tokens_total": "prompt_tokens",
"vllm:generation_tokens_total": "generation_tokens"}
function show(role, text) {
const div = document.createElement("div");
div.className = role;
div.textContent = text;
log.appendChild(div);
return div;
}
draft.addEventListener("input", () => {
if (!streaming && draft.value.trim())
fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
});
def vllm_metrics(base_url):
"""The METRICS of vLLM's Prometheus text, summed over label sets; None when unreachable."""
try:
with urllib.request.urlopen(f"{base_url}/metrics", timeout=2) as response:
text = response.read().decode()
except (OSError, ValueError):
return None
out = dict.fromkeys(METRICS.values(), 0.0)
for line in text.splitlines():
key = METRICS.get(line.split("{", 1)[0].split(" ", 1)[0])
if key and not line.startswith("#"):
out[key] += float(line.rsplit(None, 1)[1])
return out
def _number(path):
try:
with open(path) as f:
return int(f.read().strip())
except (OSError, ValueError):
return None
def gpu_stats(sysfs="/sys/class/drm"):
"""VRAM, load, temperature, power and fan of the GPU with the most VRAM (amdgpu sysfs);
None without one."""
cards = [(_number(os.path.join(d, "mem_info_vram_total")), d)
for d in glob.glob(os.path.join(sysfs, "card*", "device"))]
cards = sorted((total, d) for total, d in cards if total)
if not cards:
return None
total, card = cards[-1]
hwmon = sorted(glob.glob(os.path.join(card, "hwmon", "hwmon*")))
def sensor(name, scale):
value = _number(os.path.join(hwmon[0], name)) if hwmon else None
return None if value is None else value / scale
return {"vram_used": _number(os.path.join(card, "mem_info_vram_used")), "vram_total": total,
"busy": _number(os.path.join(card, "gpu_busy_percent")),
"temp": sensor("temp1_input", 1000), "power": sensor("power1_average", 1e6),
"fan": sensor("fan1_input", 1)}
def expert_stats(path):
"""prain's stats file (prain.policy.write_stats), or None."""
if not path:
return None
try:
with open(path) as f:
return json.load(f)
except (OSError, ValueError):
return None
draft.addEventListener("keydown", async (event) => {
if (event.key !== "Enter" || event.shiftKey || streaming || !draft.value.trim()) return;
event.preventDefault();
const text = draft.value.trim();
messages.push({role: "user", content: text});
show("user", "you> " + text);
draft.value = "";
streaming = true;
const out = show("assistant", "");
const finalSpan = out.appendChild(document.createElement("span"));
const draftSpan = out.appendChild(document.createElement("span"));
draftSpan.className = "draft";
let mode = "full", draftText = "", finalText = "", refined = false;
const render = () => { // the full text, then the grey rest of the draft while they agree
finalSpan.textContent = finalText;
draftSpan.textContent = mode === "draft" ? draftText
: mode === "refine" && draftText.startsWith(finalText) ? draftText.slice(finalText.length) : "";
};
let start = performance.now(), first = 0, pieces = 0, buffer = "";
const response = await fetch("/chat", {method: "POST", body: JSON.stringify({session, messages})});
if (!response.ok) {
out.textContent = "error " + response.status + ": " + await response.text();
messages.pop();
streaming = false;
return;
}
const reader = response.body.getReader();
const decoder = new TextDecoder();
for (;;) {
const {done, value} = await reader.read();
if (done) break;
buffer += decoder.decode(value, {stream: true});
const lines = buffer.split("\\n");
buffer = lines.pop();
for (const line of lines) {
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
const data = JSON.parse(line.slice(6));
if (data.prain === "draft") {
mode = "draft";
status.textContent = "GPU busy: drafting from the experts in VRAM, refined when it is free";
} else if (data.prain === "refine") {
mode = "refine";
refined = true;
start = performance.now(), first = 0, pieces = 0;
status.textContent = "refining with all experts";
} else if (data.prain === "error") {
mode = "draft";
refined = false;
status.textContent = "refining failed (HTTP " + data.status + "), keeping the draft";
} else {
const piece = data.choices[0].delta.content;
if (!piece) continue;
if (pieces++ === 0) first = performance.now();
if (mode === "draft") draftText += piece; else finalText += piece;
}
render();
}
}
messages.push({role: "assistant", content: mode === "draft" ? draftText : finalText});
let same = 0;
while (same < draftText.length && draftText[same] === finalText[same]) same++;
const seconds = (performance.now() - first) / 1000;
status.textContent = (refined ? `draft kept ${same} of ${draftText.length} chars; refined: `
: mode === "draft" ? "draft, not refined: " : "") +
`first token ${((first - start) / 1000).toFixed(1)} s, ` +
`${pieces > 1 ? ((pieces - 1) / seconds).toFixed(1) : "0.0"} tok/s`;
streaming = false;
});
</script>
"""
def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
draft=Draft()):
draft=Draft(), stats_path=None, sysfs="/sys/class/drm"):
model = first_model(base_url)
lock = threading.Condition()
sessions = OrderedDict() # session -> Warmer, least recently used first
@@ -167,7 +130,13 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
class Handler(BaseHTTPRequestHandler):
def do_GET(self):
self.reply(200, "text/html; charset=utf-8", PAGE.encode())
if self.path == "/status":
with lock:
busy = streams[0]
data = {"model": model, "vllm": vllm_metrics(base_url), "gpu": gpu_stats(sysfs),
"experts": expert_stats(stats_path), "streams": busy, "time": time.time()}
return self.reply(200, "application/json", json.dumps(data).encode())
self.reply(200, "text/html; charset=utf-8", PAGE)
def do_POST(self):
body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
@@ -178,6 +147,7 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
if body["draft"].strip() and ((warmup == "idle" and not busy) or
(warmup == "always" and busy < serve.max_requests)):
warmer(session).update((body["messages"], body["draft"]))
return self.reply(202, "text/plain", b"")
return self.reply(204, "text/plain", b"")
if self.path != "/chat":
return self.reply(404, "text/plain", b"not found")
@@ -267,9 +237,10 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
return server
def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1"):
def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1",
stats_path=None):
server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve,
draft=draft)
draft=draft, stats_path=stats_path)
print(f"prain web: http://{host}:{server.server_port} (vLLM at {base_url}; "
f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
server.serve_forever()
pyproject.toml
+3
View File
@@ -17,3 +17,6 @@ build-backend = "setuptools.build_meta"
[tool.setuptools]
packages = ["prain"]
[tool.setuptools.package-data]
prain = ["web.html"]
tests/helpers.py
+3
View File
@@ -41,10 +41,13 @@ class StubVllm:
self.delay = 0.0
self.pause = 0.0
self.fail = None
self.metrics = None # Prometheus text served at GET /metrics
stub = self
class Handler(BaseHTTPRequestHandler):
def do_GET(self):
if self.path == "/metrics" and stub.metrics is not None:
return self.send(200, "text/plain", stub.metrics)
self.send(200, "application/json", {"data": [{"id": "stub-model"}]})
def do_POST(self):
tests/test_policy.py
+4
View File
@@ -119,3 +119,7 @@ class Prefetch(TmpCase):
stats = json.load(f)["layers"]["7"]
self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]),
([1, 1, 1, 0], 1, 2))
self.assertEqual(stats["vram"], [1, 2]) # the experts in VRAM now, for the web monitor
write_stats(path, {7: SlotCache([5, 3], 8, 2, Policy(eviction="lfu"))})
with open(path) as f:
self.assertEqual(json.load(f)["layers"]["7"]["vram"], [3, 5]) # pinned ones too
tests/test_web.py
+61 -5
View File
@@ -1,4 +1,5 @@
import json
import os
import threading
import unittest
import urllib.error
@@ -7,7 +8,7 @@ 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
from tests.helpers import StubVllm, TmpCase, wait_for
def events(response):
@@ -38,7 +39,9 @@ class Web(unittest.TestCase):
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()
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,
@@ -46,7 +49,9 @@ class Web(unittest.TestCase):
def test_page_warm_up_and_chat(self):
with urllib.request.urlopen(self.url) as response:
self.assertIn(b'<textarea id="draft"', response.read())
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"}])
@@ -67,9 +72,9 @@ class Web(unittest.TestCase):
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.assertEqual(self.warm("bob", "typing along"), 204) # 204: nothing warms up
self.server.streams[0] = 0
self.warm("bob", "now")
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"])
@@ -136,3 +141,54 @@ class Web(unittest.TestCase):
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)
tests/test_web_page.py
+49
View File
@@ -0,0 +1,49 @@
import json
import shutil
import subprocess
import unittest
from pathlib import Path
PAGE = (Path(__file__).parent.parent / "prain" / "web.html").read_text()
@unittest.skipUnless(shutil.which("node"), "needs node")
class Markdown(unittest.TestCase):
"""The page's md() (model output -> innerHTML) run in node."""
def md(self, text):
script = PAGE[PAGE.index("<script>") + 8:PAGE.index("/* ---------- thread")]
script = script.split("/* ---------- markdown", 1)[1].split("*/", 1)[1]
out = subprocess.run(["node", "-e", script + f"\nprocess.stdout.write(md({json.dumps(text)}))"],
capture_output=True, text=True, check=True)
return out.stdout
def test_html_in_model_output_is_text(self):
html = self.md('<img src=x onerror="alert(1)"> **<b>hi</b>** `<i>` [x](javascript:alert(1))')
self.assertNotIn("<img", html)
self.assertNotIn("<b>", html)
self.assertNotIn('href="javascript', html)
self.assertIn("&lt;img src=x onerror=&quot;alert(1)&quot;&gt;", html)
self.assertIn("<strong>&lt;b&gt;hi&lt;/b&gt;</strong>", html)
def test_blocks(self):
html = self.md("# Title\n\nSome *very* **bold** text with `code` and "
"[a link](https://x.org/?a=1&b=2).\n\n```python\nif a < b:\n pass\n```\n"
"3. three\n4. four\n\n- one\n more\n- two\n\n| a | b |\n|---|---|\n| 1 | 2 |\n"
"\n> quoted\n\n---\nend")
for part in ("<h3>Title</h3>", "<em>very</em>", "<strong>bold</strong>", "<code>code</code>",
'<a href="https://x.org/?a=1&amp;b=2" target="_blank" rel="noopener noreferrer">',
'<span>python</span>', "<pre><code>if a &lt; b:\n pass</code></pre>",
'<ol start="3"><li>three</li><li>four</li></ol>',
"<ul><li>one more</li><li>two</li></ul>",
"<th>a</th><th>b</th></tr></thead><tbody><tr><td>1</td><td>2</td>",
"<blockquote><p>quoted</p></blockquote>", "<hr>", "<p>end</p>"):
self.assertIn(part, html)
def test_an_open_fence_while_streaming_is_code(self):
self.assertIn("<pre><code>x = 1\n# not a heading</code></pre>",
self.md("Here:\n```\nx = 1\n# not a heading"))
if __name__ == "__main__":
unittest.main()