diff --git a/prain/cli.py b/prain/cli.py new file mode 100644 index 0000000..3bc586e --- /dev/null +++ b/prain/cli.py @@ -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() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..4acc410 --- /dev/null +++ b/tests/test_cli.py @@ -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)