160 lines
7.2 KiB
Python
160 lines
7.2 KiB
Python
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)
|