prain: stats planner and TOML writer

prain/config.py

@@ -241,3 +241,64 @@ def load_config(path):
     _check_count(config.cache_slots, 0, f"{path}: cache_slots")
     _check_count(config.stats_every, 1, f"{path}: stats_every")
     return config
+
+
+def moe_dims(hf_config):
+    """(num_layers, num_experts, top_k, hidden, moe_intermediate) of a HF config.json dict."""
+    c = hf_config.get("text_config") or hf_config
+
+    def first(*names):
+        for name in names:
+            if c.get(name) is not None:
+                return c[name]
+        raise KeyError(f"config.json has none of {names}")
+
+    return (
+        first("num_hidden_layers"),
+        first("num_experts", "num_local_experts", "n_routed_experts"),
+        first("top_k_experts", "num_experts_per_tok"),
+        first("hidden_size"),
+        first("moe_intermediate_size", "intermediate_size"),
+    )
+
+
+def expert_bytes(hidden, moe_intermediate, dtype_bytes=2):
+    """One gated expert: gate+up [2I, H] plus down [H, I]."""
+    return 3 * hidden * moe_intermediate * dtype_bytes
+
+
+def render_toml(header, top, tables=None, layers=None):
+    """A placement file: top-level keys, [tables], then [layers."N"] sections (empty ones skipped)."""
+    lines = [f"# {header}"] + [f"{k} = {json.dumps(v)}" for k, v in top.items() if v is not None]
+    for name, values in (tables or {}).items():
+        lines += ["", f"[{name}]"] + [f"{k} = {json.dumps(v)}" for k, v in values.items()]
+    for layer, values in sorted((layers or {}).items()):
+        if values:
+            lines += ["", f'[layers."{layer}"]'] + [f"{k} = {json.dumps(v)}" for k, v in values.items()]
+    return "\n".join(lines) + "\n"
+
+
+def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path):
+    """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk."""
+    uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()}
+    vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses)
+    if vram_rows < 0:
+        raise ValueError("expert VRAM budget is smaller than the cache slots alone")
+    ram_rows = int(ram_gib * GIB // expert_size)
+    ranked = sorted(
+        ((layer, e) for layer in uses for e in range(num_experts)),
+        key=lambda le: (-uses[le[0]][le[1]], le),
+    )
+    tier = {le: "vram" if i < vram_rows else "ram" if i < vram_rows + ram_rows else "disk"
+            for i, le in enumerate(ranked)}
+    layers = {
+        layer: {name: format_ids(ids) for name in ("vram", "ram")
+                if (ids := [e for e in range(num_experts) if tier[(layer, e)] == name])}
+        for layer in uses
+    }
+    return render_toml(
+        "generated by `prain plan`; experts not listed live on disk",
+        {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots,
+         "stats_path": stats_path},
+        layers=layers,
+    )

tests/test_config.py

@@ -119,3 +119,41 @@ mode = "auto"
                             ('[draft]\nfull_requests = 3', "full_requests must be <= \\[serve\\] max_requests \\(2\\)")):
             with self.assertRaisesRegex(ValueError, error, msg=body):
                 mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
+
+
+class Plan(TmpCase):
+    def test_hottest_to_vram_then_ram_then_disk(self):
+        stats = {"layers": {"0": {"uses": [5, 0, 9, 1]}, "1": {"uses": [2, 7, 0, 0]}}}
+        # 1 GiB experts: 4 GiB VRAM = 2 cache slots + 2 pinned, 3 GiB RAM = 3 experts
+        text = mc.plan(stats, 4, mc.GIB, 4, 3, 1, "/packs", "/s.json")
+        self.assertEqual(text, """\
+# generated by `prain plan`; experts not listed live on disk
+pack_dir = "/packs"
+default_tier = "disk"
+cache_slots = 1
+stats_path = "/s.json"
+
+[layers."0"]
+vram = "2"
+ram = "0,3"
+
+[layers."1"]
+vram = "1"
+ram = "0"
+""")
+        p0 = mc.load_config(self.write("plan.toml", text)).placement(0, 4)
+        self.assertEqual((p0.vram, p0.ram, p0.disk), ((2,), (0, 3), (1,)))
+
+    def test_budget_smaller_than_cache_slots(self):
+        with self.assertRaisesRegex(ValueError, "smaller than the cache slots"):
+            mc.plan({"layers": {"0": {"uses": [1]}}}, 1, mc.GIB, 1, 0, 2, "/p", "/s")
+
+    def test_render_toml_round_trips_tables(self):
+        text = mc.render_toml("h", {"pack_dir": "/p", "stats_path": None},
+                              {"policy": {"eviction": "lru", "lfu_half_life": 8.0}},
+                              {3: {"cache_slots": 2}, 4: {}})
+        config = mc.load_config(self.write("r.toml", text))
+        self.assertEqual((config.policy.eviction, config.policy.lfu_half_life), ("lru", 8.0))
+        self.assertEqual(config.placement(3, 4).cache_slots, 2)
+        self.assertNotIn("stats_path", text)
+        self.assertNotIn('"4"', text)