import os import unittest import prain.config as mc from tests.helpers import TmpCase class ParseIds(unittest.TestCase): def test_ranges_lists_and_star(self): self.assertEqual(mc.parse_ids("0-3,7", 8), (0, 1, 2, 3, 7)) self.assertEqual(mc.parse_ids([5, 1, 5], 8), (1, 5)) self.assertEqual(mc.parse_ids("*", 3), (0, 1, 2)) def test_rejects_bad_specs(self): for spec in ("3-1", "a", "1,", "0-8", [1.5]): with self.assertRaises(ValueError, msg=spec): mc.parse_ids(spec, 8) def test_format_is_the_inverse(self): self.assertEqual(mc.format_ids([7, 0, 1, 2, 3, 9, 10]), "0-3,7,9-10") self.assertEqual(mc.parse_ids(mc.format_ids([4, 2, 3]), 8), (2, 3, 4)) 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')) 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)