122 lines
6.0 KiB
Python
122 lines
6.0 KiB
Python
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import unittest
|
|
|
|
from prain.config import Miss, Policy
|
|
from prain.policy import SlotCache, write_stats
|
|
from tests.helpers import TmpCase
|
|
|
|
|
|
def lru(**kw):
|
|
return Policy(eviction="lru", **kw)
|
|
|
|
|
|
class Eviction(unittest.TestCase):
|
|
def test_pinned_experts_stay_and_lru_evicts_the_oldest(self):
|
|
cache = SlotCache(pinned=[3, 5], num_experts=8, cache_slots=2, policy=lru())
|
|
self.assertEqual(cache.plan({0: 1, 3: 1}), [({0: 2, 3: 0}, [(0, 2)])])
|
|
self.assertEqual(cache.plan({0: 1, 1: 1}), [({0: 2, 1: 3}, [(1, 3)])])
|
|
self.assertEqual(cache.plan({1: 1, 2: 1}), [({1: 3, 2: 2}, [(2, 2)])]) # 0 was oldest
|
|
self.assertEqual((cache.hits, cache.misses), (3, 3))
|
|
self.assertEqual(cache.uses, [2, 2, 1, 1, 0, 0, 0, 0])
|
|
|
|
def test_lfu_keeps_the_frequent_expert_lru_the_recent_one(self):
|
|
for policy, survivor in ((lru(), 1), (Policy(eviction="lfu"), 0)):
|
|
cache = SlotCache([], 3, 2, policy)
|
|
for e in (0, 0, 0, 1, 2):
|
|
cache.plan({e: 1})
|
|
self.assertIn(survivor, cache.slots, policy.eviction)
|
|
|
|
def test_gdsf_keeps_the_expensive_disk_expert(self):
|
|
cost = [1.0, 4.0, 1.0] # expert 1 lives on disk
|
|
for eviction, survivor in (("lfu", 0), ("gdsf", 1)):
|
|
cache = SlotCache([], 3, 2, Policy(eviction=eviction), cost)
|
|
for e in (0, 0, 1, 2):
|
|
cache.plan({e: 1})
|
|
self.assertIn(survivor, cache.slots, eviction)
|
|
|
|
|
|
class Rounds(unittest.TestCase):
|
|
def test_hottest_experts_are_copied_last_and_stay(self):
|
|
cache = SlotCache([], 10, 2, lru())
|
|
rounds = cache.plan({0: 1, 1: 9, 2: 3, 3: 8, 4: 2, 5: 1})
|
|
self.assertEqual([sorted(slot_of) for slot_of, _ in rounds], [[0, 5], [2, 4], [1, 3]])
|
|
self.assertEqual(set(cache.slots), {1, 3})
|
|
|
|
def test_scan_keeps_the_protected_expert_and_promotes_the_hottest(self):
|
|
cache = SlotCache([], 6, 3, lru(scan_tokens=4, scan_keep=1, promote=1))
|
|
cache.plan({0: 1})
|
|
cache.plan({1: 1})
|
|
rounds = cache.plan({2: 1, 3: 5, 4: 2, 5: 3}, tokens=8)
|
|
self.assertEqual(rounds, [({2: 2, 4: 0}, [(2, 2), (4, 0)]),
|
|
({5: 2, 3: 0}, [(5, 2), (3, 0)])])
|
|
self.assertEqual(set(cache.slots), {1, 3, 5}) # 1 was protected
|
|
self.assertEqual(cache.evictor.last[3], cache.now) # promoted
|
|
self.assertEqual(cache.evictor.last[5], -math.inf) # admitted at low priority
|
|
|
|
def test_cpu_experts_get_no_slot(self):
|
|
cache = SlotCache([0], 6, 2, lru())
|
|
cache.plan({1: 1})
|
|
counts = {0: 1, 1: 1, 2: 1, 3: 4}
|
|
self.assertEqual(cache.choose_cpu(counts, Miss()), set())
|
|
self.assertEqual(cache.choose_cpu(counts, Miss(mode="cpu")), {2, 3})
|
|
cpu = cache.choose_cpu(counts, Miss(mode="auto", cpu_max_tokens=2))
|
|
self.assertEqual(cpu, {2})
|
|
self.assertEqual(cache.plan(counts, 4, cpu), [({0: 0, 1: 1, 3: 2}, [(3, 2)])])
|
|
self.assertEqual((cache.cpu, cache.misses), (1, 2))
|
|
|
|
def test_random_plans_always_compute_every_expert_from_the_right_slot(self):
|
|
rng = random.Random(1)
|
|
for _ in range(300):
|
|
n = rng.randint(2, 12)
|
|
pinned = rng.sample(range(n), rng.randint(0, n // 2))
|
|
slots = rng.randint(1, 4)
|
|
policy = Policy(eviction=rng.choice(["lru", "lfu", "gdsf"]),
|
|
scan_tokens=rng.randint(1, 6), scan_keep=rng.randint(0, slots - 1),
|
|
promote=rng.randint(0, 3))
|
|
cache = SlotCache(pinned, n, slots, policy, [rng.choice([1.0, 4.0]) for _ in range(n)])
|
|
content = {slot: e for e, slot in cache.pinned.items()} # what each slot really holds
|
|
for _ in range(12):
|
|
if rng.random() < 0.3:
|
|
guesses = [(e, 0.5) for e in rng.sample(range(n), rng.randint(1, n))]
|
|
recent = set(cache.recent)
|
|
for e, slot in cache.prefetch(guesses, rng.randint(0, 3)):
|
|
self.assertNotIn(content.get(slot), recent - {None})
|
|
content[slot] = e
|
|
routed = rng.sample(range(n), rng.randint(1, n))
|
|
heat = {e: rng.randint(1, 9) for e in routed}
|
|
cpu = cache.choose_cpu(heat, Miss(mode=rng.choice(["copy", "auto"]), cpu_max_tokens=3))
|
|
computed = []
|
|
for slot_of, loads in cache.plan(heat, rng.randint(1, 8), cpu):
|
|
self.assertEqual(len(set(slot_of.values())), len(slot_of))
|
|
for e, slot in loads:
|
|
self.assertEqual(slot_of[e], slot)
|
|
content[slot] = e
|
|
for e, slot in slot_of.items():
|
|
self.assertEqual(content[slot], e)
|
|
computed += slot_of
|
|
self.assertEqual(sorted(computed + sorted(cpu)), sorted(routed))
|
|
held = list(cache.slots.values())
|
|
self.assertEqual(len(set(held)), len(held))
|
|
self.assertEqual(len(held) + len(cache.free), slots)
|
|
|
|
|
|
class Prefetch(TmpCase):
|
|
def test_prefetch_is_low_priority_and_never_evicts_the_last_forward(self):
|
|
cache = SlotCache([], 4, 2, Policy(eviction="lfu"))
|
|
cache.plan({0: 1})
|
|
self.assertEqual(cache.prefetch([(1, 0.5), (2, 0.4)], budget=2), [(1, 1)]) # 0 is recent
|
|
cache.plan({1: 1}) # the prefetch paid off
|
|
self.assertEqual(cache.prefetch([(3, 0.9)], budget=1), [(3, 0)]) # evicts 0, not 1
|
|
cache.plan({2: 1}) # 3 was never used: it goes first
|
|
self.assertEqual((cache.prefetches, cache.prefetch_hits, cache.prefetch_wasted), (2, 1, 1))
|
|
self.assertEqual(set(cache.slots), {1, 2})
|
|
path = os.path.join(self.dir, "stats.json")
|
|
write_stats(path, {7: cache})
|
|
with open(path) as f:
|
|
stats = json.load(f)["layers"]["7"]
|
|
self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]),
|
|
([1, 1, 1, 0], 1, 2))
|