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

prain: command line

Browse Source
This commit is contained in:
fabolous005
2026-10-27 13:52:10 +01:00
parent 168251e81c
commit 7a78cdee79
2 changed files with 276 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
prain/cli.py
+198
View File
@@ -0,0 +1,198 @@
"""prain command line: check, plan or tune a placement, simulate it, serve with vLLM, and
chat in a TUI or the browser."""
import argparse
import json
import os
import sys
from dataclasses import asdict
from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
moe_dims, plan, render_toml)
def _model_dims(hf_config_path):
with open(hf_config_path) as f:
num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f))
return num_layers, num_experts, top_k, expert_bytes(hidden, inter)
def check(config_path, hf_config_path):
config = load_config(config_path)
num_layers, num_experts, top_k, size = _model_dims(hf_config_path)
print("layer vram ram disk cache_slots")
vram_rows = ram_rows = few_slots = 0
for layer in range(num_layers):
p = config.placement(layer, num_experts)
vram_rows += p.num_slots
ram_rows += len(p.ram)
few_slots += bool((p.ram or p.disk) and p.cache_slots < top_k)
print(f"{layer:5} {len(p.vram):4} {len(p.ram):4} {len(p.disk):4} {p.cache_slots:11}")
if few_slots:
print(f"warning: {few_slots} layers have fewer cache slots than top_k={top_k}; "
"their decode steps run in several rounds", file=sys.stderr)
print(f"expert size {size / 2**20:.2f} MiB")
print(f"VRAM {vram_rows * size / GIB:.2f} GiB (pinned experts + cache slots)")
print(f"KV {config.serve.kv_cache_gib:.2f} GiB (shared by up to "
f"{config.serve.max_requests} parallel requests)")
print(f"RAM {ram_rows * size / GIB:.2f} GiB (pinned host memory)")
print(f"disk {num_layers * num_experts * size / GIB:.2f} GiB (pack in {config.pack_dir})")
def _variant_text(variant):
policy, prefetch, miss = variant
fetch = "none" if prefetch.mode == "none" else f"pregate d{prefetch.distance} b{prefetch.budget}"
cpu = miss.mode if miss.mode == "copy" else f"{miss.mode} admit{miss.admit}"
return (f"{policy.eviction:5} {policy.heat:8} keep{policy.scan_keep} "
f"promote{policy.promote} {fetch:15} {cpu}")
def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
from prain.sim import describe, load_cost, read_forwards, simulate
config = load_config(config_path)
_, num_experts, _, size = _model_dims(hf_config_path)
forwards = read_forwards(trace)
cost = load_cost(cost_path, expert_mib or size / 2**20)
layers = sorted({layer for fwd in forwards for layer in fwd})
placements = {layer: config.placement(layer, num_experts) for layer in layers}
variant = (config.policy, config.prefetch, config.miss)
result = simulate(forwards, placements, *variant, cost)
bound = simulate(forwards, placements, Policy(scan_tokens=config.policy.scan_tokens),
Prefetch(), Miss(), cost, belady=True)
print("\n".join(describe(result, bound, cost)))
def tune_command(args):
from prain.sim import describe, load_cost, read_forwards, tune
_, num_experts, _, size = _model_dims(args.hf_config)
cost = load_cost(args.cost, args.expert_mib or size / 2**20)
forwards = read_forwards(args.trace)
best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
args.ram_gib, cost)
print(" rank total_s decode_ms hit% eviction heat keep promote prefetch miss",
file=sys.stderr)
for rank, (variant, r) in enumerate(rows[: args.top], 1):
print(f"{rank:5} {r.total_ms / 1000:7.2f} {r.ms_per_decode:9.1f} {100 * r.hit_rate:4.1f}"
f" {_variant_text(variant)}", file=sys.stderr)
print("best, with per-layer cache slots:", file=sys.stderr)
print("\n".join(" " + line for line in describe(final, bound, cost)), file=sys.stderr)
policy, prefetch, miss = best
layers = {layer: {k: v for k, v in (("ram", format_ids(ram[layer])),
("cache_slots", per[layer])) if v != ""}
for layer in per}
print(render_toml(
f"generated by `prain tune` from {os.path.basename(args.trace)}",
{"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk",
"cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)},
{"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)},
layers,
), end="")
SERVE_OWNED = ("--enforce-eager", "--moe-backend", "--scheduling-policy", "--max-num-seqs",
"--kv-cache-memory-bytes") # set by `prain serve` from the placement file
def serve(args):
config = load_config(args.config) # fail now, not after vLLM spent minutes loading
extra = args.vllm_args[1:] if args.vllm_args[:1] == ["--"] else args.vllm_args
owned = [a for a in extra if a.split("=")[0] in SERVE_OWNED]
if owned:
raise SystemExit(f"prain serve: {' '.join(owned)} comes from the placement file "
"([serve] max_requests, kv_cache_gib); remove it after --")
os.environ["PRAIN_CONFIG"] = os.path.abspath(args.config)
if config.miss.mode != "copy": # vLLM pins torch to 1 CPU thread unless this is set
os.environ.setdefault("OMP_NUM_THREADS", str(config.miss.cpu_threads))
os.execvp("vllm", ["vllm", "serve", args.model, "--enforce-eager", "--moe-backend", "triton",
"--scheduling-policy", "priority",
"--max-num-seqs", str(config.serve.max_requests),
"--kv-cache-memory-bytes", str(int(config.serve.kv_cache_gib * GIB)),
*extra])
def main(argv=None):
parser = argparse.ArgumentParser(prog="prain")
sub = parser.add_subparsers(dest="command", required=True)
c = sub.add_parser("check", help="validate a placement file and print what it costs")
c.add_argument("config")
c.add_argument("hf_config", help="the model's config.json")
p = sub.add_parser("plan", help="print a placement file built from routing stats")
p.add_argument("stats", help="stats_path written by a prain run")
p.add_argument("hf_config", help="the model's config.json")
p.add_argument("--expert-vram-gib", type=float, required=True,
help="VRAM for experts: pinned ones plus cache slots")
p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts")
p.add_argument("--cache-slots", type=int, default=16)
p.add_argument("--pack-dir", required=True)
for name, text in (("sim", "simulate a placement file on a routing trace"),
("tune", "search policies and per-layer slots on a trace; prints a placement file")):
s = sub.add_parser(name, help=text)
s.add_argument("trace", help="trace_path written by a prain run, or `prain synth`")
if name == "sim":
s.add_argument("config")
s.add_argument("hf_config", help="the model's config.json")
s.add_argument("--cost", help="JSON from scripts/measure.py (default: RX 6600 XT guesses)")
s.add_argument("--expert-mib", type=float, help="what-if: expert size, e.g. 2.9 for int4")
if name == "tune":
s.add_argument("--expert-vram-gib", type=float, required=True)
s.add_argument("--ram-gib", type=float, required=True)
s.add_argument("--pack-dir", required=True)
s.add_argument("--top", type=int, default=10, help="variants to list")
y = sub.add_parser("synth", help="write a synthetic routing trace, to try sim and tune")
y.add_argument("out")
for flag, default in (("forwards", 300), ("layers", 30), ("experts", 128), ("top-k", 8),
("seed", 0)):
y.add_argument(f"--{flag}", type=int, default=default)
s = sub.add_parser("serve", help="exec `vllm serve` with the placement active")
s.add_argument("config")
s.add_argument("model")
s.add_argument("vllm_args", nargs=argparse.REMAINDER, help="extra vllm serve args after --")
for name, what in (("tui", "the terminal"), ("web", "the browser")):
ui = sub.add_parser(name, help=f"chat in {what}; experts warm up while you type")
ui.add_argument("--url", default="http://127.0.0.1:8000", help="the vLLM server")
ui.add_argument("--warmup", choices=("idle", "always", "off"), default="idle",
help="idle: not while an answer streams (default)")
ui.add_argument("--warmup-priority", type=int, default=1,
help="vLLM priority of warm-ups; 0 for a server without "
"--scheduling-policy priority")
if name == "web":
ui.add_argument("--port", type=int, default=8080)
ui.add_argument("--config", help="placement file: [serve] max_requests and the "
"[draft] table (default: their defaults)")
args = parser.parse_args(argv)
if args.command == "check":
check(args.config, args.hf_config)
elif args.command == "plan":
_, num_experts, _, size = _model_dims(args.hf_config)
with open(args.stats) as f:
stats = json.load(f)
print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib,
args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats)), end="")
elif args.command == "sim":
sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib)
elif args.command == "tune":
tune_command(args)
elif args.command == "synth":
from prain.sim import synthetic_trace, write_trace
write_trace(args.out, synthetic_trace(args.layers, args.experts, args.top_k,
args.forwards, seed=args.seed))
elif args.command == "serve":
serve(args)
elif args.command == "tui":
from prain.tui import main as tui # curses only loads when needed
tui(args.url, args.warmup, args.warmup_priority)
else:
from prain.config import Config
from prain.web import serve as web
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)
if __name__ == "__main__":
main()
tests/test_cli.py
+78
View File
@@ -0,0 +1,78 @@
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_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)