prain: placement and option tables

2026-10-27 13:52:06 +01:00
parent db1736c42d
commit a5f011dcb8
2 changed files with 280 additions and 0 deletions
@@ -23,3 +23,99 @@ class ParseIds(unittest.TestCase):
    def test_layer_index(self):
        self.assertEqual(mc.layer_index("language_model.model.layers.12.experts"), 12)
        self.assertIsNone(mc.layer_index("model.experts"))
+
+
+class Placement(TmpCase):
+    def test_sections_apply_in_file_order(self):
+        config = mc.load_config(self.write("p.toml", """
+pack_dir = "/packs"
+default_tier = "ram"
+cache_slots = 4
+
+[layers."*"]
+vram = "0-1"
+
+[layers."2-3"]
+disk = "6-7"
+vram = "1-2"
+cache_slots = 2
+"""))
+        p0 = config.placement(0, 8)
+        self.assertEqual((p0.vram, p0.ram, p0.disk, p0.cache_slots),
+                         ((0, 1), (2, 3, 4, 5, 6, 7), (), 4))
+        p3 = config.placement(3, 8)
+        self.assertEqual((p3.vram, p3.ram, p3.disk, p3.cache_slots),
+                         ((0, 1, 2), (3, 4, 5), (6, 7), 2))
+        self.assertEqual(p3.num_slots, 5)
+
+    def test_rejects_typos_conflicts_and_missing_slots(self):
+        with self.assertRaisesRegex(ValueError, "unknown keys"):
+            mc.load_config(self.write("a.toml", 'pack_dir = "/p"\ncache_slot = 3\n'))
+        with self.assertRaisesRegex(ValueError, "default_tier"):
+            mc.load_config(self.write("b.toml", 'pack_dir = "/p"\ndefault_tier = "gpu"\n'))
+        both = mc.load_config(
+            self.write("c.toml", 'pack_dir = "/p"\n[layers."*"]\nvram = "1"\ndisk = "1"\n'))
+        with self.assertRaisesRegex(ValueError, "both vram and disk"):
+            both.placement(0, 4)
+        no_slots = mc.load_config(self.write("d.toml", 'pack_dir = "/p"\ncache_slots = 0\n'))
+        with self.assertRaisesRegex(ValueError, "cache_slots must be >= 1"):
+            no_slots.placement(0, 4)
+
+    def test_all_in_vram_needs_no_cache_slots(self):
+        config = mc.load_config(
+            self.write("v.toml", 'pack_dir = "/p"\ndefault_tier = "vram"\ncache_slots = 0\n'))
+        self.assertEqual(config.placement(5, 4).num_slots, 4)
+
+    def test_policy_prefetch_and_miss_tables(self):
+        config = mc.load_config(self.write("t.toml", """
+pack_dir = "/p"
+cache_slots = 3
+[policy]
+eviction = "gdsf"
+lfu_half_life = 16
+scan_keep = 1
+[prefetch]
+mode = "pregate"
+min_prob = 0.1
+[miss]
+mode = "auto"
+"""))
+        self.assertEqual((config.policy.eviction, config.policy.lfu_half_life, config.policy.heat),
+                         ("gdsf", 16, "weight"))
+        self.assertEqual((config.prefetch.mode, config.prefetch.distance, config.miss.mode),
+                         ("pregate", 1, "auto"))
+        self.assertEqual(mc.load_config(self.write("u.toml", 'pack_dir = "/p"\n')).policy, mc.Policy())
+
+    def test_rejects_bad_table_values(self):
+        for body, error in (('[policy]\neviction = "fifo"', "eviction must be one of"),
+                            ('[policy]\nscan_keep = -1', "scan_keep must be a number"),
+                            ('[prefetch]\ndistance = 3', "distance must be a number in \\[1, 2\\]"),
+                            ('[prefetch]\nbudget = 1.5', "budget must be a number"),
+                            ('[miss]\nmode = "gpu"', "mode must be one of"),
+                            ('[miss]\nthreads = 4', "unknown keys")):
+            with self.assertRaisesRegex(ValueError, error, msg=body):
+                mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
+        keep_all = mc.load_config(self.write("k.toml", 'pack_dir = "/p"\ncache_slots = 2\n'
+                                                       '[policy]\nscan_keep = 2\n'))
+        with self.assertRaisesRegex(ValueError, "scan_keep must be < cache_slots"):
+            keep_all.placement(0, 4)
+
+    def test_paths_are_relative_to_the_config_file(self):
+        config = mc.load_config(self.write("r.toml", 'pack_dir = "packs/x"\nstats_path = "~/s.json"\n'
+                                                     'trace_path = "/abs/t.gz"\n'))
+        self.assertEqual(config.pack_dir, os.path.join(self.dir, "packs/x"))
+        self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
+        self.assertEqual(config.trace_path, "/abs/t.gz")
+
+    def test_serve_and_draft_tables(self):
+        config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
+                                                     'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'
+                                                     'refine = "none"\n'))
+        self.assertEqual(config.serve, mc.Serve(max_requests=4, kv_cache_gib=1.5))
+        self.assertEqual(config.draft, mc.Draft(full_requests=2, min_weight=0.5, refine="none"))
+        self.assertEqual(mc.load_config(self.write("d.toml", 'pack_dir = "/p"\n')).serve.max_requests, 2)
+        for body, error in (('[serve]\nmax_requests = 0', "max_requests must be a number"),
+                            ('[draft]\nmin_weight = 1.5', "min_weight must be a number in \\[0, 1\\]"),
+                            ('[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'))