prain: expert id ranges

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

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"))