diff --git a/src/kernbench/policy/routing/router.py b/src/kernbench/policy/routing/router.py index 0869079..d4ec39b 100644 --- a/src/kernbench/policy/routing/router.py +++ b/src/kernbench/policy/routing/router.py @@ -111,6 +111,13 @@ class PathRouter: self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list) self._adj_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list) self._adj_local: dict[str, list[tuple[str, float]]] = defaultdict(list) + # Memoize path lookups: adj dicts are built once here and never + # mutated (topology is static per ADR-0006 / SPEC ยง0.1), so + # (id(adj), start, goal) is a stable cache key for the router's + # lifetime. Callers use the returned path list read-only. + self._path_cache: dict[ + tuple[int, str, str], tuple[list[str], float] + ] = {} for e in graph.edges: w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm self._adj_all[e.src].append((e.dst, w)) @@ -185,8 +192,14 @@ class PathRouter: start: str, goal: str, ) -> tuple[list[str], float]: + cache_key = (id(adj), start, goal) + cached = self._path_cache.get(cache_key) + if cached is not None: + return cached if start == goal: - return [start], 0.0 + result = ([start], 0.0) + self._path_cache[cache_key] = result + return result best: dict[str, float] = {start: 0.0} prev: dict[str, str] = {} heap: list[tuple[float, str]] = [(0.0, start)] @@ -200,7 +213,9 @@ class PathRouter: cur = prev[cur] path.append(start) path.reverse() - return path, d + result = (path, d) + self._path_cache[cache_key] = result + return result if d > best.get(node, float("inf")): continue for neighbor, edge_dist in adj[node]: