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.
113 lines
6.0 KiB
Python
113 lines
6.0 KiB
Python
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from prain.cli import main
|
|
from prain.config import load_config
|
|
from tests.helpers import TmpCase
|
|
|
|
GEMMA4 = {"text_config": {"num_hidden_layers": 30, "num_experts": 128, "top_k_experts": 8,
|
|
"hidden_size": 2816, "moe_intermediate_size": 704}}
|
|
|
|
|
|
class Cli(TmpCase):
|
|
def run_cli(self, *argv):
|
|
out, err = io.StringIO(), io.StringIO()
|
|
with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
|
|
main(list(argv))
|
|
return out.getvalue(), err.getvalue()
|
|
|
|
def test_check_prints_tier_sizes(self):
|
|
hf = self.write("config.json", json.dumps(GEMMA4))
|
|
config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 4\n')
|
|
out, err = self.run_cli("check", config, hf)
|
|
self.assertIn(" 0 0 128 0 4", out)
|
|
self.assertTrue(out.endswith(
|
|
"expert size 11.34 MiB\n"
|
|
"VRAM 1.33 GiB (pinned experts + cache slots)\n"
|
|
"KV 0.75 GiB (shared by up to 2 parallel requests)\n"
|
|
"RAM 42.54 GiB (pinned host memory)\n"
|
|
"disk 42.54 GiB (pack in /packs)\n"), out)
|
|
self.assertIn("warning: 30 layers have fewer cache slots than top_k=8", err)
|
|
|
|
def test_plan_output_is_a_valid_config(self):
|
|
hf = self.write("config.json", json.dumps(GEMMA4))
|
|
stats = {"layers": {str(layer): {"uses": list(range(128))} for layer in range(30)}}
|
|
stats_path = self.write("stats.json", json.dumps(stats))
|
|
out, _ = self.run_cli("plan", stats_path, hf, "--expert-vram-gib", "8",
|
|
"--ram-gib", "16", "--pack-dir", "/packs")
|
|
config = self.write("planned.toml", out)
|
|
check_out, _ = self.run_cli("check", config, hf)
|
|
self.assertIn("VRAM 8.00 GiB", check_out) # budgets are filled, never exceeded
|
|
self.assertIn("RAM 16.00 GiB", check_out)
|
|
|
|
def test_serve_execs_vllm_with_the_required_flags(self):
|
|
config = self.write("p.toml", 'pack_dir = "/packs"\n[miss]\nmode = "auto"\ncpu_threads = 6\n'
|
|
'[serve]\nmax_requests = 3\nkv_cache_gib = 1.0\n')
|
|
with mock.patch("os.execvp") as execvp, mock.patch.dict(os.environ, clear=False):
|
|
os.environ.pop("OMP_NUM_THREADS", None)
|
|
main(["serve", config, "/models/m", "--", "--port", "8001"])
|
|
self.assertEqual(os.environ["PRAIN_CONFIG"], config)
|
|
self.assertEqual(os.environ["OMP_NUM_THREADS"], "6")
|
|
execvp.assert_called_once_with("vllm", [
|
|
"vllm", "serve", "/models/m", "--enforce-eager", "--moe-backend", "triton",
|
|
"--scheduling-policy", "priority", "--max-num-seqs", "3",
|
|
"--kv-cache-memory-bytes", "1073741824", "--port", "8001"])
|
|
with mock.patch("os.execvp") as execvp, self.assertRaisesRegex(
|
|
SystemExit, "--max-num-seqs=4 comes from the placement file"):
|
|
main(["serve", config, "/models/m", "--", "--max-num-seqs=4"])
|
|
execvp.assert_not_called()
|
|
|
|
def test_web_binds_the_host_it_is_given(self):
|
|
with mock.patch("prain.web.serve") as web:
|
|
main(["web", "--port", "8081"])
|
|
self.assertEqual(web.call_args.kwargs["host"], "127.0.0.1")
|
|
main(["web", "--host", "0.0.0.0"])
|
|
self.assertEqual(web.call_args.kwargs["host"], "0.0.0.0")
|
|
|
|
def test_a_persona_file_reaches_web_and_tui(self):
|
|
path = self.write("gentoo-admin.md", "You look after a Gentoo box.\n\n")
|
|
with mock.patch("prain.web.serve") as web:
|
|
main(["web", "--persona", path])
|
|
self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
|
|
("You look after a Gentoo box.", "gentoo-admin"))
|
|
main(["web"])
|
|
self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
|
|
(None, None))
|
|
with mock.patch("prain.tui.main") as tui:
|
|
main(["tui", "--persona", path])
|
|
self.assertEqual(tui.call_args.kwargs["system"], "You look after a Gentoo box.")
|
|
|
|
def test_synth_sim_and_tune(self):
|
|
hf = self.write("config.json", json.dumps(GEMMA4))
|
|
trace = os.path.join(self.dir, "trace.jsonl.gz")
|
|
self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
|
|
config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 5\n')
|
|
out, _ = self.run_cli("sim", trace, config, hf)
|
|
self.assertRegex(out, r"VRAM hit rate \d+\.\d% \(eviction alone at best: \d+\.\d%\)")
|
|
self.assertIn("decode ", out)
|
|
out, err = self.run_cli("tune", trace, hf, "--expert-vram-gib", "0.2",
|
|
"--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
|
|
self.assertIn(" 1 ", err)
|
|
tuned = load_config(self.write("tuned.toml", out))
|
|
self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
|
|
self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
|
|
int(0.2 * 1024 / 11.34) // 3 * 3)
|
|
|
|
def test_quantized_experts_check_and_tune(self):
|
|
hf = self.write("config.json", json.dumps(GEMMA4))
|
|
out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
|
|
self.assertIn("expert size 3.01 MiB\n", out)
|
|
trace = os.path.join(self.dir, "trace.jsonl.gz")
|
|
self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
|
|
out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
|
|
"--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
|
|
tuned = load_config(self.write("tuned.toml", out))
|
|
self.assertEqual(tuned.quant, "int4")
|
|
slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
|
|
self.assertGreater(slots, int(0.2 * 1024 / 11.34)) # int4 experts are ~4x smaller
|
|
self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))
|