feat(prefill): add maintenance cache saturation harness

docs/ops/distributed-prefill-kv-network.md

@@ -233,6 +233,33 @@ separate Worker A/B/C path; that mode additionally requires `remote_jobs`.
 Decode throughput is reported separately because all autoregressive decode
 remains on the primary.
 
+## Maintenance cache saturation
+
+Enable the bounded, memory-only first-append capture queue on the primary:
+
+```text
+--cache-fill-capture-size 256
+```
+
+The maintenance endpoints require the network API key even when other read
+endpoints are public. Captured token IDs remain in process memory and are
+removed when drained; reports contain only salted capture IDs and token counts.
+
+During a maintenance window, start real gRPC chat sessions and run:
+
+```bash
+PYTHONPATH=.:sdks/python python scripts/fill_prefill_cache_from_live_grpc.py \
+  --tokenizer-id ~/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit \
+  --target-one 0.90 \
+  --target-two 0.95 \
+  --churn-gb 0.9
+```
+
+The harness controls head and allens independently, stops on memory pressure,
+publish failures, fallbacks, or `/tmp/kakeya-cache-fill.stop`, and never expects
+resident fleet usage to exceed the configured 1+8 GiB ceiling. Churn is accepted
+when `bytes_evicted` increases while resident bytes remain bounded.
+
 ## Rollback
 
 The cache is an optimization; inference correctness does not depend on it.

inference_engine/distributed/cache_fill.py

@@ -0,0 +1,88 @@
+"""Bounded, in-memory capture queue for maintenance cache-fill replays."""
+from __future__ import annotations
+
+import hashlib
+import secrets
+import threading
+from collections import deque
+from dataclasses import dataclass
+from typing import Iterable
+
+
+@dataclass(frozen=True)
+class CapturedPrefix:
+    capture_id: str
+    token_ids: tuple[int, ...]
+    token_count: int
+
+
+class CacheFillCapture:
+    """Capture first appends without persisting prompt or token content."""
+
+    def __init__(
+        self,
+        *,
+        max_items: int = 256,
+        excluded_label_prefix: str = "cache-fill-",
+    ) -> None:
+        if max_items <= 0:
+            raise ValueError("max_items must be > 0")
+        self.max_items = int(max_items)
+        self.excluded_label_prefix = excluded_label_prefix
+        self._salt = secrets.token_bytes(32)
+        self._items: deque[CapturedPrefix] = deque()
+        self._seen: set[bytes] = set()
+        self._lock = threading.Lock()
+        self.captured = 0
+        self.duplicates = 0
+        self.dropped = 0
+
+    def observe(
+        self,
+        *,
+        client_label: str,
+        token_ids: Iterable[int],
+    ) -> bool:
+        if client_label.startswith(self.excluded_label_prefix):
+            return False
+        tokens = tuple(int(token) for token in token_ids)
+        if not tokens:
+            return False
+        digest = hashlib.sha256(
+            self._salt
+            + b"".join(token.to_bytes(4, "little", signed=False) for token in tokens)
+        ).digest()
+        with self._lock:
+            if digest in self._seen:
+                self.duplicates += 1
+                return False
+            if len(self._items) >= self.max_items:
+                evicted = self._items.popleft()
+                self._seen.discard(bytes.fromhex(evicted.capture_id))
+                self.dropped += 1
+            item = CapturedPrefix(digest.hex(), tokens, len(tokens))
+            self._items.append(item)
+            self._seen.add(digest)
+            self.captured += 1
+            return True
+
+    def drain(self, max_items: int) -> list[CapturedPrefix]:
+        if max_items <= 0:
+            raise ValueError("max_items must be > 0")
+        output = []
+        with self._lock:
+            while self._items and len(output) < max_items:
+                item = self._items.popleft()
+                self._seen.discard(bytes.fromhex(item.capture_id))
+                output.append(item)
+        return output
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return {
+                "queued": len(self._items),
+                "captured": self.captured,
+                "duplicates": self.duplicates,
+                "dropped": self.dropped,
+                "max_items": self.max_items,
+            }

inference_engine/distributed/capability.py

@@ -181,6 +181,9 @@ class CacheCapability:
     bloom_filter: bytes = b""
     default_compression: CompressionCodec = CompressionCodec.NONE
     replication_factor: int = 1
+    evictions: int = 0
+    bytes_evicted: int = 0
+    put_failures: int = 0
 
     def to_proto(self) -> distributed_pb2.CacheCapability:
         return distributed_pb2.CacheCapability(
@@ -195,6 +198,9 @@ def to_proto(self) -> distributed_pb2.CacheCapability:
             bloom_filter=self.bloom_filter,
             default_compression=int(self.default_compression),
             replication_factor=self.replication_factor,
+            evictions=self.evictions,
+            bytes_evicted=self.bytes_evicted,
+            put_failures=self.put_failures,
         )
 
     @classmethod
@@ -211,6 +217,9 @@ def from_proto(cls, msg: distributed_pb2.CacheCapability) -> "CacheCapability":
             bloom_filter=msg.bloom_filter,
             default_compression=CompressionCodec(msg.default_compression),
             replication_factor=msg.replication_factor,
+            evictions=msg.evictions,
+            bytes_evicted=msg.bytes_evicted,
+            put_failures=msg.put_failures,
         )
 
 

inference_engine/distributed/prefill_cache.py

@@ -126,6 +126,9 @@ class CacheStats:
     lookup_misses: int
     tokens_served: int
     bytes_served: int
+    evictions: int
+    bytes_evicted: int
+    put_failures: int
 
 
 class PrefixCacheStore:
@@ -153,17 +156,23 @@ def __init__(
         self._lookup_misses = 0
         self._tokens_served = 0
         self._bytes_served = 0
+        self._evictions = 0
+        self._bytes_evicted = 0
+        self._put_failures = 0
         self._lock = threading.RLock()
 
     def put(self, block: CacheBlock) -> bool:
         """Publish one immutable block. Returns False for an identical hit."""
         if block.nbytes > self.max_bytes:
+            with self._lock:
+                self._put_failures += 1
             raise ValueError("block payload exceeds cache capacity")
         with self._lock:
             self._expire_leases(time.time())
             existing = self._blocks.get(block.block_hash)
             if existing is not None:
                 if existing.payload_sha256 != block.payload_sha256:
+                    self._put_failures += 1
                     raise ValueError("content-address collision with different payload")
                 self._blocks.move_to_end(block.block_hash)
                 return False
@@ -172,6 +181,7 @@ def put(self, block: CacheBlock) -> bool:
             self._epoch += 1
             self._evict_to_budget()
             if block.block_hash not in self._blocks:
+                self._put_failures += 1
                 raise ValueError(
                     "cache capacity is pinned by active leases",
                 )
@@ -259,6 +269,9 @@ def stats(self) -> CacheStats:
                 lookup_misses=self._lookup_misses,
                 tokens_served=self._tokens_served,
                 bytes_served=self._bytes_served,
+                evictions=self._evictions,
+                bytes_evicted=self._bytes_evicted,
+                put_failures=self._put_failures,
             )
 
     def block_hashes(self) -> tuple[bytes, ...]:
@@ -298,6 +311,8 @@ def _evict_to_budget(self) -> None:
                 break
             block = self._blocks.pop(victim)
             self._bytes_used -= block.nbytes
+            self._evictions += 1
+            self._bytes_evicted += block.nbytes
             self._epoch += 1
 
 

inference_engine/distributed/prefill_cache_runtime.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import hashlib
+import threading
 import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -72,6 +73,11 @@ class PrefillReuseStats:
     remote_job_failures: int = 0
     fallbacks: int = 0
     last_fallback_reason: str = ""
+    publish_attempts: int = 0
+    publish_successes: int = 0
+    publish_failures: int = 0
+    bytes_published: int = 0
+    last_publish_error: str = ""
 
 
 @dataclass(frozen=True)
@@ -139,6 +145,7 @@ def __init__(
         self.auth = auth
         self._hash_key = auth.tenant_hash_key() if auth is not None else b""
         self.stats = PrefillReuseStats()
+        self._stats_lock = threading.Lock()
         self._on_reuse = on_reuse
         self._publisher = ThreadPoolExecutor(
             max_workers=4,
@@ -395,14 +402,35 @@ def _publish_boundary(
                 publish_block_sync,
             )
             for peer in peers:
-                self._publisher.submit(
+                with self._stats_lock:
+                    self.stats.publish_attempts += 1
+                future = self._publisher.submit(
                     publish_block_sync,
                     peer,
                     self.compatibility,
                     block,
                     timeout_s=self.fetch_timeout_s,
                     auth=self.auth,
                 )
+                future.add_done_callback(
+                    lambda completed, nbytes=block.nbytes: self._publish_done(
+                        completed,
+                        nbytes,
+                    ),
+                )
+
+    def _publish_done(self, future, nbytes: int) -> None:
+        try:
+            stored = bool(future.result())
+        except Exception as exc:
+            with self._stats_lock:
+                self.stats.publish_failures += 1
+                self.stats.last_publish_error = f"{type(exc).__name__}: {exc}"
+            return
+        with self._stats_lock:
+            self.stats.publish_successes += 1
+            if stored:
+                self.stats.bytes_published += int(nbytes)
 
     def close(self) -> None:
         self._publisher.shutdown(wait=False, cancel_futures=True)

inference_engine/distributed/prefill_cache_service.py

@@ -52,6 +52,9 @@ def cache_capability(
         cache_epoch=stats.cache_epoch,
         load=load,
         tokens_served=stats.tokens_served,
+        evictions=stats.evictions,
+        bytes_evicted=stats.bytes_evicted,
+        put_failures=stats.put_failures,
         default_compression=(
             CompressionCodec.NONE
             if default_compression is None

inference_engine/network/api.py

@@ -5,7 +5,7 @@
 import asyncio
 import json
 import os
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from fastapi import Depends, FastAPI, Header, HTTPException
 from fastapi.responses import HTMLResponse, StreamingResponse
@@ -14,6 +14,9 @@
 from inference_engine.network.dashboard import dashboard_html
 from inference_engine.network.state import NetworkState
 
+if TYPE_CHECKING:
+    from inference_engine.distributed.cache_fill import CacheFillCapture
+
 
 class RegisterNodeRequest(BaseModel):
     alias: str = Field(min_length=1, max_length=100)
@@ -33,10 +36,15 @@ class TokenTelemetryRequest(BaseModel):
     kv_assisted: int = Field(default=0, ge=0)
 
 
+class DrainCaptureRequest(BaseModel):
+    max_items: int = Field(default=8, ge=1, le=64)
+
+
 def create_network_app(
     state: NetworkState,
     *,
     api_key: Optional[str] = None,
+    cache_fill_capture: Optional["CacheFillCapture"] = None,
 ) -> FastAPI:
     key = (api_key if api_key is not None else os.environ.get(
         "KAKEYA_NETWORK_API_KEY", "",
@@ -49,6 +57,12 @@ def require_key(
         if key and x_api_key != key:
             raise HTTPException(status_code=401, detail="invalid X-API-Key")
 
+    def require_maintenance_key(
+        x_api_key: Optional[str] = Header(default=None),
+    ) -> None:
+        if not key or x_api_key != key:
+            raise HTTPException(status_code=401, detail="maintenance API key required")
+
     @app.get("/", response_class=HTMLResponse)
     @app.get("/network", response_class=HTMLResponse)
     def dashboard() -> str:
@@ -106,6 +120,33 @@ def tokens():
     def prefill():
         return state.prefill_stats()
 
+    @app.get(
+        "/v1/network/maintenance/capture",
+        dependencies=[Depends(require_maintenance_key)],
+    )
+    def capture_status():
+        if cache_fill_capture is None:
+            raise HTTPException(status_code=404, detail="cache-fill capture disabled")
+        return cache_fill_capture.stats()
+
+    @app.post(
+        "/v1/network/maintenance/capture/drain",
+        dependencies=[Depends(require_maintenance_key)],
+    )
+    def drain_capture(request: DrainCaptureRequest):
+        if cache_fill_capture is None:
+            raise HTTPException(status_code=404, detail="cache-fill capture disabled")
+        return {
+            "items": [
+                {
+                    "capture_id": item.capture_id,
+                    "token_count": item.token_count,
+                    "token_ids": list(item.token_ids),
+                }
+                for item in cache_fill_capture.drain(request.max_items)
+            ],
+        }
+
     @app.post(
         "/v1/network/telemetry/tokens",
         dependencies=[Depends(require_key)],

inference_engine/network/dashboard.py

@@ -32,7 +32,7 @@ def dashboard_html() -> str:
 <div class="top"><div><h1>Kakeya Inference Network</h1><div class="sub">P2P Prefill KV sharing across trusted inference nodes</div></div>
 <div class="tabs"><button class="active" data-tab="overview">Overview</button><button data-tab="nodes">Nodes</button><button data-tab="groups">Groups</button><button class="primary" id="registerBtn">Register node</button></div></div>
 <section id="register" class="card register hidden"><h3>Register inference node</h3><div class="grid3"><label>Alias<input id="alias" placeholder="prefill-worker-tb"></label><label>Address<input id="address" placeholder="169.254.27.104:53051"></label><label>Region<input id="region" placeholder="Hong Kong"></label></div><label>Admin API key<input id="adminKey" type="password" placeholder="Required for network changes"></label><button class="primary" id="createRegistration">Create pairing token</button><code id="pairing" class="hidden"></code></section>
-<section class="stats"><div class="card stat"><b id="online">0</b><span>Online nodes</span></div><div class="card stat"><b id="groupCount">0</b><span>Inference groups</span></div><div class="card stat"><b id="tokens">0</b><span>Completed tokens</span></div><div class="card stat"><b id="hitRate">0%</b><span>KV-assisted tokens</span></div><div class="card stat"><b id="cache">0 GB</b><span>Shared cache online</span></div><div class="card stat"><b id="remoteJobs">0</b><span>Remote prefill jobs</span></div><div class="card stat"><b id="remoteHits">0</b><span>Remote KV imports</span></div><div class="card stat"><b id="reusedTokens">0</b><span>Tokens reused</span></div></section>
+<section class="stats"><div class="card stat"><b id="online">0</b><span>Online nodes</span></div><div class="card stat"><b id="groupCount">0</b><span>Inference groups</span></div><div class="card stat"><b id="tokens">0</b><span>Completed tokens</span></div><div class="card stat"><b id="hitRate">0%</b><span>KV-assisted tokens</span></div><div class="card stat"><b id="cache">0 GB</b><span>Shared cache online</span></div><div class="card stat"><b id="remoteJobs">0</b><span>Remote prefill jobs</span></div><div class="card stat"><b id="remoteHits">0</b><span>Remote KV imports</span></div><div class="card stat"><b id="reusedTokens">0</b><span>Tokens reused</span></div><div class="card stat"><b id="evictions">0</b><span>LRU evictions</span></div><div class="card stat"><b id="publishFailures">0</b><span>Publish failures</span></div></section>
 <section id="overview" class="tab">
 <div class="grid2"><div><h2>Online node distribution</h2><div class="card map" id="map"></div></div><div><h2>Live KV discovery</h2><div class="card" id="events"><div class="event"><time>live</time><div><b>Waiting for node telemetry</b><div class="muted">Capability gossip and prefix lookups appear here.</div></div></div></div><h2>Cache capacity</h2><div class="card"><span id="capacityLabel">0 / 0 GB</span><div class="bar"><i id="capacityBar" style="width:0%"></i></div></div></div></div>
 </section>
@@ -49,7 +49,7 @@ def dashboard_html() -> str:
 $('createGroup').onclick=async()=>{await fetch('/v1/network/groups',{method:'POST',headers:writeHeaders(),body:JSON.stringify({name:$('groupName').value,node_ids:$('groupNodes').value.split(',').map(x=>x.trim()).filter(Boolean)})});load()};
 function nodePosition(i,total){let a=(i/Math.max(total,1))*Math.PI*2;return {x:50+38*Math.cos(a),y:53+35*Math.sin(a)}}
 async function load(){let [s,n,g]=await Promise.all([fetch('/v1/network/summary').then(r=>r.json()),fetch('/v1/network/nodes').then(r=>r.json()),fetch('/v1/network/groups').then(r=>r.json())]);
-$('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB';let p=s.prefill||{};$('remoteJobs').textContent=fmt(p.remote_jobs);$('remoteHits').textContent=fmt(p.remote_hits);$('reusedTokens').textContent=fmt(p.tokens_reused);
+$('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB';let p=s.prefill||{};$('remoteJobs').textContent=fmt(p.remote_jobs);$('remoteHits').textContent=fmt(p.remote_hits);$('reusedTokens').textContent=fmt(p.tokens_reused);$('evictions').textContent=fmt(s.cache_evictions);$('publishFailures').textContent=fmt(p.publish_failures);
 let total=s.cache_bytes_used+s.cache_bytes_free,pct=total?s.cache_bytes_used/total*100:0;$('capacityLabel').textContent=`${gb(s.cache_bytes_used)} / ${gb(total)} GB`;$('capacityBar').style.width=pct+'%';
 $('map').innerHTML=n.map((x,i)=>{let p=nodePosition(i,n.length);return `<div class="dot ${x.role.includes('head')?'head':''}" style="left:${p.x}%;top:${p.y}%"><i></i><small>${x.region}<br>${x.alias}</small></div>`}).join('');
 $('nodesBody').innerHTML=n.map(x=>`<tr><td>${x.alias}</td><td>${x.role}</td><td>${x.region}</td><td>${x.cache?(x.cache.model_id+' / '+x.cache.format):'—'}</td><td>${x.endpoint.network} ${x.endpoint.rtt_ms?x.endpoint.rtt_ms+'ms':''}</td><td class="${x.status}">${x.status}</td></tr>`).join('');