diff --git a/prain/config.py b/prain/config.py new file mode 100644 index 0000000..a840238 --- /dev/null +++ b/prain/config.py @@ -0,0 +1,59 @@ +"""Placement config: which experts of each MoE layer live in VRAM, RAM or on disk, how the +VRAM slots are managed ([policy], [prefetch] and [miss] tables), and how parallel requests +share the GPU ([serve] and [draft] tables).""" + +import json +import os +import tomllib +from dataclasses import dataclass, field, fields + +TIERS = ("vram", "ram", "disk") +GIB = 1024**3 +_TOP_KEYS = {"pack_dir", "default_tier", "cache_slots", "stats_path", "stats_every", + "trace_path", "policy", "prefetch", "miss", "serve", "draft", "layers"} +_LAYER_KEYS = {"cache_slots", *TIERS} + + +def parse_ids(spec, limit=None): + """'0-3,7' or [0, 1, 2, 3, 7] or '*' -> (0, 1, 2, 3, 7). '*' needs a limit.""" + if spec == "*": + if limit is None: + raise ValueError("'*' is only allowed where the count is known") + ids = set(range(limit)) + elif isinstance(spec, list): + if not all(type(i) is int for i in spec): + raise ValueError(f"expected a list of ints, got {spec!r}") + ids = set(spec) + elif isinstance(spec, str): + ids = set() + for part in spec.replace(" ", "").split(","): + lo, dash, hi = part.partition("-") + if not lo.isdigit() or (dash and not hi.isdigit()): + raise ValueError(f"bad id range {part!r} in {spec!r}") + first, last = int(lo), int(hi) if dash else int(lo) + if last < first: + raise ValueError(f"empty id range {part!r} in {spec!r}") + ids.update(range(first, last + 1)) + else: + raise ValueError(f"expected '0-3,7', a list of ints or '*', got {spec!r}") + if limit is not None and any(not 0 <= i < limit for i in ids): + raise ValueError(f"ids in {spec!r} must be in [0, {limit})") + return tuple(sorted(ids)) + + +def format_ids(ids): + """Inverse of parse_ids: [0, 1, 2, 3, 7] -> '0-3,7'.""" + ids, parts, i = sorted(ids), [], 0 + while i < len(ids): + j = i + while j + 1 < len(ids) and ids[j + 1] == ids[j] + 1: + j += 1 + parts.append(str(ids[i]) if i == j else f"{ids[i]}-{ids[j]}") + i = j + 1 + return ",".join(parts) + + +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 diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..3d37887 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,25 @@ +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"))