prain: traces, simulator and tuner

2026-10-27 13:52:09 +01:00
parent 582601101b
commit 168251e81c
2 changed files with 394 additions and 0 deletions
@@ -0,0 +1,83 @@
+import os
+import unittest
+
+from prain.config import LayerPlacement, Miss, Policy, Prefetch
+from prain.sim import (Cost, read_forwards, simulate, synthetic_trace, tune, write_trace)
+from tests.helpers import TmpCase
+
+SMALL = dict(layers=4, experts=32, top_k=4, forwards=120, prompt_every=60, prompt_tokens=32)
+
+
+def forwards_of(records):
+    forwards, last = [], None
+    for r in records:
+        if last is None or r["l"] <= last:
+            forwards.append({})
+        forwards[-1][r["l"]] = r
+        last = r["l"]
+    return forwards
+
+
+def all_ram(slots, layers=4, experts=32):
+    return {layer: LayerPlacement((), tuple(range(experts)), (), slots) for layer in range(layers)}
+
+
+class Traces(TmpCase):
+    def test_round_trip_and_determinism(self):
+        records = synthetic_trace(**SMALL)
+        self.assertEqual(records, synthetic_trace(**SMALL))
+        path = os.path.join(self.dir, "t.jsonl.gz")
+        write_trace(path, records[:8])
+        write_trace(path, records[8:], mode="at")  # runs append to the same trace
+        forwards = read_forwards(path)
+        self.assertEqual(len(forwards), 120)
+        self.assertEqual(forwards[0][0], records[0])
+        self.assertEqual(forwards[0][0]["t"], 32)
+        self.assertEqual(sorted(forwards[1][1]["p"]), ["1", "2"])
+        self.assertNotIn("p", forwards[1][3])  # nothing after the last layer
+
+    def test_copy_cost(self):
+        self.assertAlmostEqual(Cost().copy_ms(False), 0.991, places=3)
+        self.assertAlmostEqual(Cost().copy_ms(True), 4.955, places=3)
+
+
+class Simulate(unittest.TestCase):
+    def setUp(self):
+        self.forwards = forwards_of(synthetic_trace(**SMALL))
+
+    def run_sim(self, policy=Policy(), prefetch=Prefetch(), miss=Miss(), slots=3, **kw):
+        return simulate(self.forwards, all_ram(slots), policy, prefetch, miss, Cost(), **kw)
+
+    def test_belady_bounds_every_policy(self):
+        bound = self.run_sim(belady=True).hit_rate
+        for eviction in ("lru", "lfu", "gdsf"):
+            self.assertLessEqual(self.run_sim(Policy(eviction=eviction)).hit_rate, bound)
+
+    def test_cpu_path_beats_copies_when_the_cpu_is_faster_than_pcie(self):
+        copy = self.run_sim().ms_per_decode
+        self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)
+
+    def test_prefetch_helps_when_predictions_are_right(self):
+        none = self.run_sim().hit_rate
+        good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))
+        self.assertGreater(good.hit_rate, none)
+        self.assertGreater(good.counters["prefetch_hits"], good.counters["prefetch_wasted"])
+
+    def test_pregate_needs_predictions(self):
+        stripped = [{layer: {k: v for k, v in r.items() if k != "p"} for layer, r in f.items()}
+                    for f in self.forwards]
+        with self.assertRaisesRegex(ValueError, "needs a trace with predictions"):
+            simulate(stripped, all_ram(3), Policy(), Prefetch(mode="pregate"), Miss(), Cost())
+
+
+class Tune(unittest.TestCase):
+    def test_tune_beats_the_defaults_within_the_budget(self):
+        forwards = forwards_of(synthetic_trace(**SMALL))
+        cost = Cost(expert_mib=64.0)  # 1 GiB = 16 slots = 4 per layer
+        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost)
+        self.assertLessEqual(sum(per.values()), 16)
+        self.assertEqual(sum(len(r) for r in ram.values()), 16)  # 1 GiB of RAM tier
+        default = next(r for v, r in rows if v == (Policy(), Prefetch(), Miss()))
+        self.assertLessEqual(final.total_ms, default.total_ms)
+        self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
+        self.assertGreaterEqual(bound.hit_rate, 0.0)