prain: command line

2026-10-27 13:52:10 +01:00
parent 168251e81c
commit 7a78cdee79
2 changed files with 276 additions and 0 deletions
@@ -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)