prain: stats planner and TOML writer

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)