From a5f011dcb8966fe9c3a7ba62554f282b0c0c369c Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 13:52:06 +0100 Subject: [PATCH] prain: placement and option tables --- prain/config.py | 184 +++++++++++++++++++++++++++++++++++++++++++ tests/test_config.py | 96 ++++++++++++++++++++++ 2 files changed, 280 insertions(+) diff --git a/prain/config.py b/prain/config.py index a840238..d746288 100644 --- a/prain/config.py +++ b/prain/config.py @@ -57,3 +57,187 @@ def layer_index(name): """'language_model.model.layers.12.experts' -> 12, or None unless exactly one number.""" numbers = [int(part) for part in name.split(".") if part.isdigit()] return numbers[0] if len(numbers) == 1 else None + + +@dataclass(frozen=True) +class LayerPlacement: + vram: tuple # pinned experts: slot i holds vram[i] for the whole run + ram: tuple # pinned host memory, copied into a cache slot on demand + disk: tuple # read from the pack file on demand + cache_slots: int # VRAM slots for ram/disk experts, managed by [policy] + + @property + def num_slots(self): + return len(self.vram) + self.cache_slots + + +def _opt(default, choices=None, minimum=0, maximum=None): + return field(default=default, metadata={"choices": choices, "min": minimum, "max": maximum}) + + +@dataclass(frozen=True) +class Policy: + """[policy]: how the VRAM cache slots of every layer are managed.""" + eviction: str = _opt("lfu", ("lru", "lfu", "gdsf")) + heat: str = _opt("weight", ("count", "weight", "recency")) + lfu_half_life: float = _opt(64.0, minimum=1) # forwards until an LFU score halves + disk_cost: float = _opt(4.0, minimum=1) # gdsf: refetching a disk expert vs. a RAM one + recency_half_life: float = _opt(32.0, minimum=1) # heat "recency": tokens per halving + scan_tokens: int = _opt(64, minimum=1) # forwards with this many tokens are scans + scan_keep: int = _opt(0) # cache slots a scan may not evict + promote: int = _opt(2) # experts a scan admits at full priority + + +@dataclass(frozen=True) +class Prefetch: + """[prefetch]: copy experts before their layer runs.""" + mode: str = _opt("none", ("none", "pregate")) + distance: int = _opt(1, minimum=1, maximum=2) # layers ahead + budget: int = _opt(2) # experts per layer and forward + min_prob: float = _opt(0.05, maximum=1) # ignore predictions below this router probability + + +@dataclass(frozen=True) +class Miss: + """[miss]: what happens to an expert that is not in VRAM.""" + mode: str = _opt("copy", ("copy", "cpu", "auto")) + cpu_max_tokens: int = _opt(1, minimum=1) # auto: CPU runs experts with <= this many tokens + cpu_threads: int = _opt(4, minimum=1) # OMP threads of the vLLM worker + admit: int = _opt(1) # CPU-run experts copied into VRAM afterwards, hottest first + + +@dataclass(frozen=True) +class Serve: + """[serve]: the vLLM settings prain and its clients both need to know.""" + max_requests: int = _opt(2, minimum=1) # requests vLLM runs at once (--max-num-seqs) + kv_cache_gib: float = _opt(0.75, minimum=0.01) # KV cache, shared by all requests + + +@dataclass(frozen=True) +class Draft: + """[draft]: answers that arrive while other answers already use the GPU.""" + full_requests: int = _opt(1, minimum=1) # answers that run with all their experts at once + min_weight: float = _opt(0.5, maximum=1) # router weight a draft token must still compute + refine: str = _opt("regenerate", ("regenerate", "none")) # what happens to a draft later + + +def _section(cls, raw, where): + if not isinstance(raw, dict): + raise ValueError(f"{where} must be a table") + spec = {f.name: f for f in fields(cls)} + _check_keys(raw, set(spec), where) + for name, value in raw.items(): + meta, default = spec[name].metadata, spec[name].default + if meta["choices"]: + if value not in meta["choices"]: + raise ValueError(f"{where} {name} must be one of {meta['choices']}, got {value!r}") + continue + number = type(value) is int or (type(default) is float and type(value) is float) + top = meta["max"] + if not number or value < meta["min"] or (top is not None and value > top): + raise ValueError(f"{where} {name} must be a number in " + f"[{meta['min']}, {top if top is not None else 'inf'}], got {value!r}") + return cls(**raw) + + +@dataclass(frozen=True) +class Config: + pack_dir: str + default_tier: str = "ram" + cache_slots: int = 16 + stats_path: str | None = None + stats_every: int = 256 + trace_path: str | None = None + policy: Policy = Policy() + prefetch: Prefetch = Prefetch() + miss: Miss = Miss() + serve: Serve = Serve() + draft: Draft = Draft() + layers: tuple = () # ((selector, section), ...) in file order; later sections win + + def placement(self, layer, num_experts): + tier = [self.default_tier] * num_experts + cache_slots = self.cache_slots + for selector, section in self.layers: + if selector != "*" and layer not in parse_ids(selector): + continue + seen = {} + for name in TIERS: + for expert in parse_ids(section.get(name, []), num_experts): + if expert in seen: + raise ValueError( + f"[layers.\"{selector}\"] puts expert {expert} in both " + f"{seen[expert]} and {name}" + ) + seen[expert] = tier[expert] = name + cache_slots = section.get("cache_slots", cache_slots) + vram, ram, disk = ( + tuple(e for e in range(num_experts) if tier[e] == name) for name in TIERS + ) + if (ram or disk) and cache_slots < 1: + raise ValueError( + f"layer {layer}: cache_slots must be >= 1 while experts live in ram or disk" + ) + if (ram or disk) and self.policy.scan_keep >= cache_slots: + raise ValueError(f"layer {layer}: policy.scan_keep must be < cache_slots ({cache_slots})") + return LayerPlacement(vram, ram, disk, cache_slots) + + +def _check_keys(table, allowed, where): + unknown = set(table) - allowed + if unknown: + raise ValueError(f"{where}: unknown keys {sorted(unknown)}, allowed: {sorted(allowed)}") + + +def _check_count(value, minimum, where): + if type(value) is not int or value < minimum: + raise ValueError(f"{where} must be an integer >= {minimum}, got {value!r}") + + +def _path(config_path, value): + """A path from the file at config_path: '~' expands, relative paths start at that file's + directory, so a config can point at '../models/...' from anywhere.""" + if value is None: + return None + return os.path.join(os.path.dirname(os.path.abspath(config_path)), os.path.expanduser(value)) + + +def load_config(path): + with open(path, "rb") as f: + raw = tomllib.load(f) + _check_keys(raw, _TOP_KEYS, path) + if "pack_dir" not in raw: + raise ValueError(f"{path}: pack_dir is required") + layers = [] + for selector, section in raw.get("layers", {}).items(): + where = f'{path} [layers."{selector}"]' + if selector != "*": + parse_ids(selector) + if not isinstance(section, dict): + raise ValueError(f"{where} must be a table") + _check_keys(section, _LAYER_KEYS, where) + if "cache_slots" in section: + _check_count(section["cache_slots"], 0, f"{where} cache_slots") + layers.append((selector, section)) + config = Config( + pack_dir=_path(path, raw["pack_dir"]), + default_tier=raw.get("default_tier", "ram"), + cache_slots=raw.get("cache_slots", 16), + stats_path=_path(path, raw.get("stats_path")), + stats_every=raw.get("stats_every", 256), + trace_path=_path(path, raw.get("trace_path")), + policy=_section(Policy, raw.get("policy", {}), f"{path} [policy]"), + prefetch=_section(Prefetch, raw.get("prefetch", {}), f"{path} [prefetch]"), + miss=_section(Miss, raw.get("miss", {}), f"{path} [miss]"), + serve=_section(Serve, raw.get("serve", {}), f"{path} [serve]"), + draft=_section(Draft, raw.get("draft", {}), f"{path} [draft]"), + layers=tuple(layers), + ) + if config.draft.full_requests > config.serve.max_requests: + raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests " + f"({config.serve.max_requests})") + if config.default_tier not in TIERS: + raise ValueError(f"{path}: default_tier must be one of {TIERS}") + _check_count(config.cache_slots, 0, f"{path}: cache_slots") + _check_count(config.stats_every, 1, f"{path}: stats_every") + return config diff --git a/tests/test_config.py b/tests/test_config.py index 3d37887..8099645 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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'))