93 lines
4.9 KiB
Python
93 lines
4.9 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_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))
|