prain/prain/policy.py

"""Per-layer VRAM slot cache: which slot holds which expert, what to evict, what to copy
first, how scans and prefetched experts are admitted. Pure Python, no torch."""

import json
import math
import os
from bisect import bisect_right

from prain.config import Policy


class Lru:
    """Evict the expert used longest ago."""

    def __init__(self, num_experts, policy, cost):
        self.last = [0.0] * num_experts

    def score(self, e, now):
        return self.last[e]

    def hit(self, e, now, heat):
        self.last[e] = now

    def admit(self, e, now, heat, low):
        self.last[e] = -math.inf if low else now

    def evicted(self, e):
        pass


class Lfu:
    """Evict the expert with the lowest exponentially decayed heat. Scores survive eviction,
    so an expert that comes back keeps its history."""

    def __init__(self, num_experts, policy, cost):
        self.decay = 0.5 ** (1 / policy.lfu_half_life)
        self.value = [0.0] * num_experts
        self.when = [0] * num_experts

    def score(self, e, now):
        return self.value[e] * self.decay ** (now - self.when[e])

    def hit(self, e, now, heat):
        self.value[e] = self.score(e, now) + heat
        self.when[e] = now

    def admit(self, e, now, heat, low):
        if not low:  # a low-priority admit keeps its history: never-seen experts score 0
            self.hit(e, now, heat)

    def evicted(self, e):
        pass


class Gdsf:
    """GreedyDual-Size-Frequency: score = floor + frequency x refetch cost. The floor rises
    to each victim's score, so experts that stopped being used age out. cost[e] is 1 for a
    RAM expert and policy.disk_cost for a disk one, so slow-to-refetch experts stay longer."""

    def __init__(self, num_experts, policy, cost):
        self.cost = cost
        self.floor = 0.0
        self.freq = [0.0] * num_experts
        self.value = [0.0] * num_experts

    def score(self, e, now):
        return self.value[e]

    def hit(self, e, now, heat):
        self.freq[e] += heat
        self.value[e] = self.floor + self.freq[e] * self.cost[e]

    def admit(self, e, now, heat, low):
        self.freq[e] = 0.0 if low else heat
        self.value[e] = self.floor + self.freq[e] * self.cost[e]

    def evicted(self, e):
        self.floor = max(self.floor, self.value[e])


class Belady:
    """The offline optimum, for the simulator only: evict the expert used furthest in the
    future. future[e] lists the forward numbers that route to e, ascending."""

    def __init__(self, future):
        self.future = future

    def score(self, e, now):
        uses = self.future.get(e, ())
        i = bisect_right(uses, now)
        return -uses[i] if i < len(uses) else -math.inf

    def hit(self, e, now, heat):
        pass

    def admit(self, e, now, heat, low):
        pass

    def evicted(self, e):
        pass


POLICIES = {"lru": Lru, "lfu": Lfu, "gdsf": Gdsf}


class SlotCache:
    """Slots [0, len(pinned)) hold the pinned experts for the whole run. The cache_slots
    slots after them hold experts from RAM or disk, managed by an eviction policy."""

    def __init__(self, pinned, num_experts, cache_slots, policy=Policy(), cost=None, evictor=None):
        self.num_experts = num_experts
        self.cache_slots = cache_slots
        self.policy = policy
        self.pinned = {e: slot for slot, e in enumerate(pinned)}
        self.slots = {}  # cached expert -> slot
        first = len(self.pinned)
        self.free = list(range(first + cache_slots - 1, first - 1, -1))  # pop() -> lowest slot
        cost = cost or [1.0] * num_experts
        self.evictor = evictor or POLICIES[policy.eviction](num_experts, policy, cost)
        self.now = 0  # forwards planned so far
        self.recent = set()  # experts routed to in the last forward
        self.prefetched = set()  # prefetched experts that were not used yet
        self.uses = [0] * num_experts  # forwards that routed to each expert
        self.hits = 0  # routed expert already in VRAM
        self.misses = 0  # routed expert copied in
        self.cpu = 0  # routed expert computed on the CPU instead
        self.prefetches = 0  # experts copied in ahead of time (prefetch or admit-behind)
        self.prefetch_hits = 0  # ... that were used before being evicted
        self.prefetch_wasted = 0  # ... that were evicted unused

    def choose_cpu(self, counts, miss):
        """Experts the CPU should compute: counts maps expert -> routed tokens."""
        if miss.mode == "copy":
            return set()
        away = (e for e in counts if e not in self.pinned and e not in self.slots)
        if miss.mode == "cpu":
            return set(away)
        return {e for e in away if counts[e] <= miss.cpu_max_tokens}

    def plan(self, heat, tokens=1, cpu=()):
        """Split one forward pass into rounds.

        heat: {expert: heat > 0} for every expert routed to in this pass.
        tokens: tokens in the pass; a pass with >= policy.scan_tokens tokens is a scan.
        cpu: routed experts the CPU computes; they get no slot.
        Returns [(slot_of, loads), ...]: slot_of maps each expert computed in that round to
        its slot; loads lists the (expert, slot) copies to do before the round. Experts in
        VRAM run in round 0. The others follow coldest first, so the hottest are copied last
        and are the ones left in VRAM afterwards.
        """
        self.now += 1
        now, scan = self.now, tokens >= self.policy.scan_tokens
        norm = {e: h / tokens for e, h in heat.items()}  # per token, so scans don't dominate
        for e in heat:
            self.uses[e] += 1
        self.recent = set(heat)
        pinned = [e for e in heat if e in self.pinned]
        resident = [e for e in heat if e in self.slots]
        missing = sorted((e for e in heat if e not in self.pinned and e not in self.slots
                          and e not in cpu), key=lambda e: (heat[e], e))
        self.hits += len(pinned) + len(resident)
        self.misses += len(missing)
        self.cpu += len(cpu)
        for e in resident:
            self.evictor.hit(e, now, norm[e])
            if e in self.prefetched:
                self.prefetched.discard(e)
                self.prefetch_hits += 1
        if missing and not self.cache_slots:
            raise RuntimeError(f"experts {missing} are not in VRAM and there are no cache slots")
        protected = set()
        if scan and self.policy.scan_keep:
            ranked = sorted(self.slots, key=lambda e: (self.evictor.score(e, now), e))
            protected = set(ranked[-self.policy.scan_keep:])
        busy = set(resident)  # experts computed in the round being built
        slot_of = {e: self.pinned[e] for e in pinned} | {e: self.slots[e] for e in resident}
        rounds, i = [], 0
        while True:
            room = len(self.free) + sum(1 for e in self.slots if e not in busy and e not in protected)
            group = missing[i:i + room]
            if i < len(missing) and not group and not busy:
                raise RuntimeError("no evictable cache slot: lower policy.scan_keep")
            i += len(group)
            loads = []
            for e in group:
                slot = self._take_slot(now, busy | protected)
                self.slots[e] = slot
                self.evictor.admit(e, now, norm[e], low=scan)
                busy.add(e)
                loads.append((e, slot))
            slot_of.update(loads)
            rounds.append((slot_of, loads))
            if i >= len(missing):
                break
            slot_of, busy = {}, set()
        if scan and self.policy.promote:
            for e in missing[-self.policy.promote:]:
                if e in self.slots:
                    self.evictor.hit(e, now, norm[e])
        return rounds

    def prefetch(self, predicted, budget):
        """Copy predicted experts in ahead of time, at low priority: they are evicted first
        unless a forward uses them. Never evicts an expert routed to in the last forward.
        predicted: [(expert, probability)], most likely first. Returns the (expert, slot) loads."""
        loads, keep = [], set(self.recent)
        for e, _ in predicted:
            if len(loads) >= budget:
                break
            if e in self.pinned or e in self.slots:
                continue
            if not self.free and all(r in keep for r in self.slots):
                break
            slot = self._take_slot(self.now, keep)
            self.slots[e] = slot
            self.evictor.admit(e, self.now, 0.0, low=True)
            self.prefetched.add(e)
            keep.add(e)
            loads.append((e, slot))
        self.prefetches += len(loads)
        return loads

    def _take_slot(self, now, keep):
        if self.free:
            return self.free.pop()
        victim = min((e for e in self.slots if e not in keep),
                     key=lambda e: (self.evictor.score(e, now), e))
        self.evictor.evicted(victim)
        if victim in self.prefetched:
            self.prefetched.discard(victim)
            self.prefetch_wasted += 1
        return self.slots.pop(victim)


COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wasted")


def write_stats(path, caches):
    """Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
    data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
                       for layer, c in sorted(caches.items())}}
    tmp = f"{path}.tmp"
    with open(tmp, "w") as f:
        json.dump(data, f)
    os.replace(tmp, path)