diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index cbb77ab5..15f44407 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/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. diff --git a/inference_engine/distributed/cache_fill.py b/inference_engine/distributed/cache_fill.py new file mode 100644 index 00000000..6c1c81e7 --- /dev/null +++ b/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, + } diff --git a/inference_engine/distributed/capability.py b/inference_engine/distributed/capability.py index 7ee9d026..52c8a0c1 100644 --- a/inference_engine/distributed/capability.py +++ b/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, ) diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index 2386f7b2..a7ea7b91 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/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 diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py index e4b6b439..79ad9da6 100644 --- a/inference_engine/distributed/prefill_cache_runtime.py +++ b/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,7 +402,9 @@ 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, @@ -403,6 +412,25 @@ def _publish_boundary( 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) diff --git a/inference_engine/distributed/prefill_cache_service.py b/inference_engine/distributed/prefill_cache_service.py index 5168d180..a13d3c7f 100644 --- a/inference_engine/distributed/prefill_cache_service.py +++ b/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 diff --git a/inference_engine/network/api.py b/inference_engine/network/api.py index 0a56213a..2e84ef72 100644 --- a/inference_engine/network/api.py +++ b/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)], diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py index eb4bb8c9..f3d89237 100644 --- a/inference_engine/network/dashboard.py +++ b/inference_engine/network/dashboard.py @@ -32,7 +32,7 @@ def dashboard_html() -> str:

Kakeya Inference Network

P2P Prefill KV sharing across trusted inference nodes
-
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
0Remote prefill jobs
0Remote KV imports
0Tokens reused
+
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
0Remote prefill jobs
0Remote KV imports
0Tokens reused
0LRU evictions
0Publish failures

Online node distribution

Live KV discovery

Waiting for node telemetry
Capability gossip and prefix lookups appear here.

Cache capacity

0 / 0 GB
@@ -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 `
${x.region}
${x.alias}
`}).join(''); $('nodesBody').innerHTML=n.map(x=>`${x.alias}${x.role}${x.region}${x.cache?(x.cache.model_id+' / '+x.cache.format):'—'}${x.endpoint.network} ${x.endpoint.rtt_ms?x.endpoint.rtt_ms+'ms':''}${x.status}`).join(''); diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 39695d4a..0d7482f0 100644 --- a/inference_engine/network/state.py +++ b/inference_engine/network/state.py @@ -131,6 +131,9 @@ def nodes(self) -> list[dict[str, Any]]: "entry_count": cache.entry_count, "epoch": cache.cache_epoch, "tokens_served": cache.tokens_served, + "evictions": cache.evictions, + "bytes_evicted": cache.bytes_evicted, + "put_failures": cache.put_failures, "format": cache.compatibility.cache_format_version, "model_id": cache.compatibility.model_id, } @@ -221,6 +224,15 @@ def summary(self) -> dict[str, Any]: "local_lookup_hits": cache_stats.lookup_hits, "local_lookup_misses": cache_stats.lookup_misses, "local_tokens_served": cache_stats.tokens_served, + "cache_evictions": sum( + (node["cache"] or {}).get("evictions", 0) for node in nodes + ), + "cache_bytes_evicted": sum( + (node["cache"] or {}).get("bytes_evicted", 0) for node in nodes + ), + "cache_put_failures": sum( + (node["cache"] or {}).get("put_failures", 0) for node in nodes + ), "prefill": self.prefill_stats(), } diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py index 6d700b9d..eea5b029 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py @@ -24,19 +24,19 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1bkakeya/v1/distributed.proto\x12\tkakeya.v1\"}\n\x0fModelCapability\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\'\n\x04role\x18\x02 \x01(\x0e\x32\x19.kakeya.v1.CapabilityRole\x12\x14\n\x0cquantization\x18\x03 \x01(\t\x12\x19\n\x11tokens_per_second\x18\x04 \x01(\x01\"\x83\x03\n\x0eNodeCapability\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x14\n\x0cgrpc_address\x18\x02 \x01(\t\x12\x10\n\x08platform\x18\x03 \x01(\t\x12\x1c\n\x14unified_memory_bytes\x18\x04 \x01(\x04\x12\x13\n\x0bmlx_version\x18\x05 \x01(\t\x12*\n\x06models\x18\x06 \x03(\x0b\x32\x1a.kakeya.v1.ModelCapability\x12\x19\n\x11\x61nnounced_at_unix\x18\x07 \x01(\x01\x12\x13\n\x0bttl_seconds\x18\x08 \x01(\x01\x12\x14\n\x0cring_address\x18\t \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\n \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\x12*\n\tendpoints\x18\x0b \x03(\x0b\x32\x17.kakeya.v1.NodeEndpoint\x12;\n\x0fprefill_workers\x18\x0c \x03(\x0b\x32\".kakeya.v1.PrefillWorkerCapability\"[\n\x0cNodeEndpoint\x12\x0f\n\x07\x61\x64\x64ress\x18\x01 \x01(\t\x12\x0f\n\x07network\x18\x02 \x01(\t\x12\x10\n\x08priority\x18\x03 \x01(\r\x12\x17\n\x0fmeasured_rtt_ms\x18\x04 \x01(\x01\"\xad\x02\n\x12\x43\x61\x63heCompatibility\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0emodel_revision\x18\x02 \x01(\t\x12\x1a\n\x12tokenizer_revision\x18\x03 \x01(\t\x12\x1c\n\x14\x63\x61\x63he_format_version\x18\x04 \x01(\t\x12\x14\n\x0cquantization\x18\x05 \x01(\t\x12\x11\n\trope_hash\x18\x06 \x01(\t\x12\x1b\n\x13layer_geometry_hash\x18\x07 \x01(\t\x12\x10\n\x08kv_dtype\x18\x08 \x01(\t\x12\x19\n\x11\x62lock_size_tokens\x18\t \x01(\r\x12\x18\n\x10tenant_namespace\x18\n \x01(\t\x12\x11\n\tsink_size\x18\x0b \x01(\r\x12\x13\n\x0bwindow_size\x18\x0c \x01(\r\"\xcd\x02\n\x0f\x43\x61\x63heCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x15\n\rcache_address\x18\x02 \x01(\t\x12\x18\n\x10\x63\x61\x63he_bytes_used\x18\x03 \x01(\x04\x12\x18\n\x10\x63\x61\x63he_bytes_free\x18\x04 \x01(\x04\x12\x13\n\x0b\x65ntry_count\x18\x05 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x06 \x01(\x04\x12\x0c\n\x04load\x18\x07 \x01(\x01\x12\x15\n\rtokens_served\x18\x08 \x01(\x04\x12\x14\n\x0c\x62loom_filter\x18\t \x01(\x0c\x12\x38\n\x13\x64\x65\x66\x61ult_compression\x18\n \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\x12\x1a\n\x12replication_factor\x18\x0b \x01(\r\"\xae\x02\n\x17PrefillWorkerCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x16\n\x0eworker_address\x18\x02 \x01(\t\x12\x1b\n\x13max_concurrent_jobs\x18\x03 \x01(\r\x12\x15\n\rinflight_jobs\x18\x04 \x01(\r\x12\x13\n\x0bqueued_jobs\x18\x05 \x01(\r\x12\x0c\n\x04load\x18\x06 \x01(\x01\x12!\n\x19tokens_per_second_prefill\x18\x07 \x01(\x01\x12\x16\n\x0eram_bytes_free\x18\x08 \x01(\x04\x12\x1c\n\x14\x61\x63\x63\x65pts_compute_jobs\x18\t \x01(\x08\x12\x15\n\rqueued_tokens\x18\n \x01(\x04\"M\n\x1b\x45xchangeCapabilitiesRequest\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x1c\x45xchangeCapabilitiesResponse\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"\x1a\n\x18GetNodeCapabilityRequest\"D\n\x19GetNodeCapabilityResponse\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x16GetCacheSummaryRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\"V\n\x17GetCacheSummaryResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\x02 \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\"a\n\x13LookupPrefixRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x14\n\x0c\x62lock_hashes\x18\x02 \x03(\x0c\"\xcf\x01\n\x14LookupPrefixResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x17\n\x0fhit_block_count\x18\x02 \x01(\r\x12\x17\n\x0fhit_token_count\x18\x03 \x01(\x04\x12\x16\n\x0etransfer_bytes\x18\x04 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x05 \x01(\x04\x12\x10\n\x08lease_id\x18\x06 \x01(\t\x12\x1d\n\x15lease_expires_at_unix\x18\x07 \x01(\x01\x12\x16\n\x0epayload_sha256\x18\x08 \x01(\x0c\"&\n\x12\x46\x65tchBlocksRequest\x12\x10\n\x08lease_id\x18\x01 \x01(\t\"\xb7\x01\n\x13\x46\x65tchBlocksResponse\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\"\xed\x01\n\x13PublishBlockRequest\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\x12\x34\n\rcompatibility\x18\t \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\";\n\x14PublishBlockResponse\x12\x0e\n\x06stored\x18\x01 \x01(\x08\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x02 \x01(\x04\"\xf0\x01\n\x17SubmitPrefillJobRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\x12\x34\n\rcompatibility\x18\x03 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x11\n\ttoken_ids\x18\x04 \x03(\r\x12\x14\n\x0c\x62lock_hashes\x18\x05 \x03(\x0c\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x06 \x01(\r\x12:\n\x15preferred_compression\x18\x07 \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\"\x85\x01\n\x18SubmitPrefillJobResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x16\n\x0eworker_node_id\x18\x03 \x01(\t\x12\x14\n\x0cqueue_eta_ms\x18\x04 \x01(\x01\"?\n\x1aGetPrefillJobStatusRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"\x8c\x02\n\x1bGetPrefillJobStatusResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x17\n\x0ftokens_computed\x18\x03 \x01(\r\x12\x10\n\x08lease_id\x18\x04 \x01(\t\x12\x12\n\nblock_hash\x18\x05 \x01(\x0c\x12\x16\n\x0epayload_sha256\x18\x06 \x01(\x0c\x12\x16\n\x0etransfer_bytes\x18\x07 \x01(\x04\x12\x16\n\x0e\x66\x61ilure_reason\x18\x08 \x01(\t\x12\x12\n\ncompute_ms\x18\t \x01(\x01\x12\x15\n\rcache_address\x18\n \x01(\t\"<\n\x17\x43\x61ncelPrefillJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"-\n\x18\x43\x61ncelPrefillJobResponse\x12\x11\n\tcancelled\x18\x01 \x01(\x08\"k\n\x13ProposeBlockRequest\x12\x1b\n\x13\x63ommitted_token_ids\x18\x01 \x03(\r\x12\x12\n\nblock_size\x18\x02 \x01(\r\x12\x11\n\tnum_steps\x18\x03 \x01(\r\x12\x10\n\x08model_id\x18\x04 \x01(\t\"y\n\x14ProposeBlockResponse\x12\x11\n\ttoken_ids\x18\x01 \x03(\r\x12\x17\n\x0f\x64iffusion_steps\x18\x02 \x01(\r\x12\x16\n\x0e\x66orward_passes\x18\x03 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x04 \x01(\x04\"4\n\x06Tensor\x12\r\n\x05\x64type\x18\x01 \x01(\t\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"T\n\x07LayerKV\x12\r\n\x05layer\x18\x01 \x01(\x05\x12\x1c\n\x01k\x18\x02 \x01(\x0b\x32\x11.kakeya.v1.Tensor\x12\x1c\n\x01v\x18\x03 \x01(\x0b\x32\x11.kakeya.v1.Tensor\"\x84\x01\n\x0eRestoreRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x12\n\nprompt_ids\x18\x02 \x03(\r\x12\x0c\n\x04sink\x18\x03 \x01(\r\x12\x0e\n\x06window\x18\x04 \x01(\r\x12\x1a\n\x12s5_exact_full_attn\x18\x05 \x01(\x08\x12\x10\n\x08model_id\x18\x06 \x01(\t\"f\n\x0fRestoreResponse\x12$\n\x08restored\x18\x01 \x03(\x0b\x32\x12.kakeya.v1.LayerKV\x12\x19\n\x11\x65victed_positions\x18\x02 \x03(\x05\x12\x12\n\nprompt_len\x18\x03 \x01(\r\"[\n\x12SeedContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\"*\n\x13SeedContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\"h\n\x11\x44raftBlockRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x16\n\x0e\x62onus_token_id\x18\x02 \x01(\r\x12\x13\n\x0b\x63ontext_len\x18\x03 \x01(\r\x12\x12\n\nblock_size\x18\x04 \x01(\r\"d\n\x12\x44raftBlockResponse\x12\x17\n\x0f\x64raft_token_ids\x18\x01 \x03(\r\x12\x16\n\x0e\x66orward_passes\x18\x02 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x03 \x01(\x04\"]\n\x14\x45xtendContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\",\n\x15\x45xtendContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\">\n(DFlashProposerServiceCloseSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\"+\n)DFlashProposerServiceCloseSessionResponse*\xed\x01\n\x0e\x43\x61pabilityRole\x12\x1f\n\x1b\x43\x41PABILITY_ROLE_UNSPECIFIED\x10\x00\x12\x1c\n\x18\x43\x41PABILITY_ROLE_VERIFIER\x10\x01\x12\x1c\n\x18\x43\x41PABILITY_ROLE_PROPOSER\x10\x02\x12\x1c\n\x18\x43\x41PABILITY_ROLE_EMBEDDER\x10\x03\x12\x18\n\x14\x43\x41PABILITY_ROLE_TOOL\x10\x04\x12!\n\x1d\x43\x41PABILITY_ROLE_PREFILL_CACHE\x10\x05\x12#\n\x1f\x43\x41PABILITY_ROLE_PREFILL_COMPUTE\x10\x06*\x96\x01\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_NONE\x10\x01\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZLIB\x10\x02\x12\'\n#COMPRESSION_CODEC_KAKEYA_LATTICE_D4\x10\x03*\xd8\x01\n\x10PrefillJobStatus\x12\"\n\x1ePREFILL_JOB_STATUS_UNSPECIFIED\x10\x00\x12\x1d\n\x19PREFILL_JOB_STATUS_QUEUED\x10\x01\x12\x1e\n\x1aPREFILL_JOB_STATUS_RUNNING\x10\x02\x12 \n\x1cPREFILL_JOB_STATUS_COMPLETED\x10\x03\x12\x1d\n\x19PREFILL_JOB_STATUS_FAILED\x10\x04\x12 \n\x1cPREFILL_JOB_STATUS_CANCELLED\x10\x05\x32\xdc\x01\n\x11\x43\x61pabilityService\x12g\n\x14\x45xchangeCapabilities\x12&.kakeya.v1.ExchangeCapabilitiesRequest\x1a\'.kakeya.v1.ExchangeCapabilitiesResponse\x12^\n\x11GetNodeCapability\x12#.kakeya.v1.GetNodeCapabilityRequest\x1a$.kakeya.v1.GetNodeCapabilityResponse2b\n\x0fProposerService\x12O\n\x0cProposeBlock\x12\x1e.kakeya.v1.ProposeBlockRequest\x1a\x1f.kakeya.v1.ProposeBlockResponse2\xe3\x02\n\x13PrefillCacheService\x12X\n\x0fGetCacheSummary\x12!.kakeya.v1.GetCacheSummaryRequest\x1a\".kakeya.v1.GetCacheSummaryResponse\x12O\n\x0cLookupPrefix\x12\x1e.kakeya.v1.LookupPrefixRequest\x1a\x1f.kakeya.v1.LookupPrefixResponse\x12N\n\x0b\x46\x65tchBlocks\x12\x1d.kakeya.v1.FetchBlocksRequest\x1a\x1e.kakeya.v1.FetchBlocksResponse0\x01\x12Q\n\x0cPublishBlock\x12\x1e.kakeya.v1.PublishBlockRequest\x1a\x1f.kakeya.v1.PublishBlockResponse(\x01\x32\xb6\x02\n\x14PrefillWorkerService\x12[\n\x10SubmitPrefillJob\x12\".kakeya.v1.SubmitPrefillJobRequest\x1a#.kakeya.v1.SubmitPrefillJobResponse\x12\x64\n\x13GetPrefillJobStatus\x12%.kakeya.v1.GetPrefillJobStatusRequest\x1a&.kakeya.v1.GetPrefillJobStatusResponse\x12[\n\x10\x43\x61ncelPrefillJob\x12\".kakeya.v1.CancelPrefillJobRequest\x1a#.kakeya.v1.CancelPrefillJobResponse2\xc1\x03\n\x15\x44\x46lashProposerService\x12@\n\x07Restore\x12\x19.kakeya.v1.RestoreRequest\x1a\x1a.kakeya.v1.RestoreResponse\x12L\n\x0bSeedContext\x12\x1d.kakeya.v1.SeedContextRequest\x1a\x1e.kakeya.v1.SeedContextResponse\x12I\n\nDraftBlock\x12\x1c.kakeya.v1.DraftBlockRequest\x1a\x1d.kakeya.v1.DraftBlockResponse\x12R\n\rExtendContext\x12\x1f.kakeya.v1.ExtendContextRequest\x1a .kakeya.v1.ExtendContextResponse\x12y\n\x0c\x43loseSession\x12\x33.kakeya.v1.DFlashProposerServiceCloseSessionRequest\x1a\x34.kakeya.v1.DFlashProposerServiceCloseSessionResponseb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1bkakeya/v1/distributed.proto\x12\tkakeya.v1\"}\n\x0fModelCapability\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\'\n\x04role\x18\x02 \x01(\x0e\x32\x19.kakeya.v1.CapabilityRole\x12\x14\n\x0cquantization\x18\x03 \x01(\t\x12\x19\n\x11tokens_per_second\x18\x04 \x01(\x01\"\x83\x03\n\x0eNodeCapability\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x14\n\x0cgrpc_address\x18\x02 \x01(\t\x12\x10\n\x08platform\x18\x03 \x01(\t\x12\x1c\n\x14unified_memory_bytes\x18\x04 \x01(\x04\x12\x13\n\x0bmlx_version\x18\x05 \x01(\t\x12*\n\x06models\x18\x06 \x03(\x0b\x32\x1a.kakeya.v1.ModelCapability\x12\x19\n\x11\x61nnounced_at_unix\x18\x07 \x01(\x01\x12\x13\n\x0bttl_seconds\x18\x08 \x01(\x01\x12\x14\n\x0cring_address\x18\t \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\n \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\x12*\n\tendpoints\x18\x0b \x03(\x0b\x32\x17.kakeya.v1.NodeEndpoint\x12;\n\x0fprefill_workers\x18\x0c \x03(\x0b\x32\".kakeya.v1.PrefillWorkerCapability\"[\n\x0cNodeEndpoint\x12\x0f\n\x07\x61\x64\x64ress\x18\x01 \x01(\t\x12\x0f\n\x07network\x18\x02 \x01(\t\x12\x10\n\x08priority\x18\x03 \x01(\r\x12\x17\n\x0fmeasured_rtt_ms\x18\x04 \x01(\x01\"\xad\x02\n\x12\x43\x61\x63heCompatibility\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x16\n\x0emodel_revision\x18\x02 \x01(\t\x12\x1a\n\x12tokenizer_revision\x18\x03 \x01(\t\x12\x1c\n\x14\x63\x61\x63he_format_version\x18\x04 \x01(\t\x12\x14\n\x0cquantization\x18\x05 \x01(\t\x12\x11\n\trope_hash\x18\x06 \x01(\t\x12\x1b\n\x13layer_geometry_hash\x18\x07 \x01(\t\x12\x10\n\x08kv_dtype\x18\x08 \x01(\t\x12\x19\n\x11\x62lock_size_tokens\x18\t \x01(\r\x12\x18\n\x10tenant_namespace\x18\n \x01(\t\x12\x11\n\tsink_size\x18\x0b \x01(\r\x12\x13\n\x0bwindow_size\x18\x0c \x01(\r\"\x8d\x03\n\x0f\x43\x61\x63heCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x15\n\rcache_address\x18\x02 \x01(\t\x12\x18\n\x10\x63\x61\x63he_bytes_used\x18\x03 \x01(\x04\x12\x18\n\x10\x63\x61\x63he_bytes_free\x18\x04 \x01(\x04\x12\x13\n\x0b\x65ntry_count\x18\x05 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x06 \x01(\x04\x12\x0c\n\x04load\x18\x07 \x01(\x01\x12\x15\n\rtokens_served\x18\x08 \x01(\x04\x12\x14\n\x0c\x62loom_filter\x18\t \x01(\x0c\x12\x38\n\x13\x64\x65\x66\x61ult_compression\x18\n \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\x12\x1a\n\x12replication_factor\x18\x0b \x01(\r\x12\x11\n\tevictions\x18\x0c \x01(\x04\x12\x15\n\rbytes_evicted\x18\r \x01(\x04\x12\x14\n\x0cput_failures\x18\x0e \x01(\x04\"\xae\x02\n\x17PrefillWorkerCapability\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x16\n\x0eworker_address\x18\x02 \x01(\t\x12\x1b\n\x13max_concurrent_jobs\x18\x03 \x01(\r\x12\x15\n\rinflight_jobs\x18\x04 \x01(\r\x12\x13\n\x0bqueued_jobs\x18\x05 \x01(\r\x12\x0c\n\x04load\x18\x06 \x01(\x01\x12!\n\x19tokens_per_second_prefill\x18\x07 \x01(\x01\x12\x16\n\x0eram_bytes_free\x18\x08 \x01(\x04\x12\x1c\n\x14\x61\x63\x63\x65pts_compute_jobs\x18\t \x01(\x08\x12\x15\n\rqueued_tokens\x18\n \x01(\x04\"M\n\x1b\x45xchangeCapabilitiesRequest\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x1c\x45xchangeCapabilitiesResponse\x12.\n\x0bknown_nodes\x18\x01 \x03(\x0b\x32\x19.kakeya.v1.NodeCapability\"\x1a\n\x18GetNodeCapabilityRequest\"D\n\x19GetNodeCapabilityResponse\x12\'\n\x04node\x18\x01 \x01(\x0b\x32\x19.kakeya.v1.NodeCapability\"N\n\x16GetCacheSummaryRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\"V\n\x17GetCacheSummaryResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12*\n\x06\x63\x61\x63hes\x18\x02 \x03(\x0b\x32\x1a.kakeya.v1.CacheCapability\"a\n\x13LookupPrefixRequest\x12\x34\n\rcompatibility\x18\x01 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x14\n\x0c\x62lock_hashes\x18\x02 \x03(\x0c\"\xcf\x01\n\x14LookupPrefixResponse\x12\x0f\n\x07node_id\x18\x01 \x01(\t\x12\x17\n\x0fhit_block_count\x18\x02 \x01(\r\x12\x17\n\x0fhit_token_count\x18\x03 \x01(\x04\x12\x16\n\x0etransfer_bytes\x18\x04 \x01(\x04\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x05 \x01(\x04\x12\x10\n\x08lease_id\x18\x06 \x01(\t\x12\x1d\n\x15lease_expires_at_unix\x18\x07 \x01(\x01\x12\x16\n\x0epayload_sha256\x18\x08 \x01(\x0c\"&\n\x12\x46\x65tchBlocksRequest\x12\x10\n\x08lease_id\x18\x01 \x01(\t\"\xb7\x01\n\x13\x46\x65tchBlocksResponse\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\"\xed\x01\n\x13PublishBlockRequest\x12\x12\n\nblock_hash\x18\x01 \x01(\x0c\x12\x13\n\x0b\x62lock_index\x18\x02 \x01(\r\x12\x13\n\x0btoken_count\x18\x03 \x01(\r\x12\x13\n\x0b\x63hunk_index\x18\x04 \x01(\r\x12\x14\n\x0ctotal_chunks\x18\x05 \x01(\r\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x14\n\x0c\x62lock_sha256\x18\x07 \x01(\x0c\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x08 \x01(\x04\x12\x34\n\rcompatibility\x18\t \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\";\n\x14PublishBlockResponse\x12\x0e\n\x06stored\x18\x01 \x01(\x08\x12\x13\n\x0b\x63\x61\x63he_epoch\x18\x02 \x01(\x04\"\xf0\x01\n\x17SubmitPrefillJobRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\x12\x34\n\rcompatibility\x18\x03 \x01(\x0b\x32\x1d.kakeya.v1.CacheCompatibility\x12\x11\n\ttoken_ids\x18\x04 \x03(\r\x12\x14\n\x0c\x62lock_hashes\x18\x05 \x03(\x0c\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x06 \x01(\r\x12:\n\x15preferred_compression\x18\x07 \x01(\x0e\x32\x1b.kakeya.v1.CompressionCodec\"\x85\x01\n\x18SubmitPrefillJobResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x16\n\x0eworker_node_id\x18\x03 \x01(\t\x12\x14\n\x0cqueue_eta_ms\x18\x04 \x01(\x01\"?\n\x1aGetPrefillJobStatusRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"\x8c\x02\n\x1bGetPrefillJobStatusResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12+\n\x06status\x18\x02 \x01(\x0e\x32\x1b.kakeya.v1.PrefillJobStatus\x12\x17\n\x0ftokens_computed\x18\x03 \x01(\r\x12\x10\n\x08lease_id\x18\x04 \x01(\t\x12\x12\n\nblock_hash\x18\x05 \x01(\x0c\x12\x16\n\x0epayload_sha256\x18\x06 \x01(\x0c\x12\x16\n\x0etransfer_bytes\x18\x07 \x01(\x04\x12\x16\n\x0e\x66\x61ilure_reason\x18\x08 \x01(\t\x12\x12\n\ncompute_ms\x18\t \x01(\x01\x12\x15\n\rcache_address\x18\n \x01(\t\"<\n\x17\x43\x61ncelPrefillJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x11\n\ttenant_id\x18\x02 \x01(\t\"-\n\x18\x43\x61ncelPrefillJobResponse\x12\x11\n\tcancelled\x18\x01 \x01(\x08\"k\n\x13ProposeBlockRequest\x12\x1b\n\x13\x63ommitted_token_ids\x18\x01 \x03(\r\x12\x12\n\nblock_size\x18\x02 \x01(\r\x12\x11\n\tnum_steps\x18\x03 \x01(\r\x12\x10\n\x08model_id\x18\x04 \x01(\t\"y\n\x14ProposeBlockResponse\x12\x11\n\ttoken_ids\x18\x01 \x03(\r\x12\x17\n\x0f\x64iffusion_steps\x18\x02 \x01(\r\x12\x16\n\x0e\x66orward_passes\x18\x03 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x04 \x01(\x04\"4\n\x06Tensor\x12\r\n\x05\x64type\x18\x01 \x01(\t\x12\r\n\x05shape\x18\x02 \x03(\x03\x12\x0c\n\x04\x64\x61ta\x18\x03 \x01(\x0c\"T\n\x07LayerKV\x12\r\n\x05layer\x18\x01 \x01(\x05\x12\x1c\n\x01k\x18\x02 \x01(\x0b\x32\x11.kakeya.v1.Tensor\x12\x1c\n\x01v\x18\x03 \x01(\x0b\x32\x11.kakeya.v1.Tensor\"\x84\x01\n\x0eRestoreRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x12\n\nprompt_ids\x18\x02 \x03(\r\x12\x0c\n\x04sink\x18\x03 \x01(\r\x12\x0e\n\x06window\x18\x04 \x01(\r\x12\x1a\n\x12s5_exact_full_attn\x18\x05 \x01(\x08\x12\x10\n\x08model_id\x18\x06 \x01(\t\"f\n\x0fRestoreResponse\x12$\n\x08restored\x18\x01 \x03(\x0b\x32\x12.kakeya.v1.LayerKV\x12\x19\n\x11\x65victed_positions\x18\x02 \x03(\x05\x12\x12\n\nprompt_len\x18\x03 \x01(\r\"[\n\x12SeedContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\"*\n\x13SeedContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\"h\n\x11\x44raftBlockRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x16\n\x0e\x62onus_token_id\x18\x02 \x01(\r\x12\x13\n\x0b\x63ontext_len\x18\x03 \x01(\r\x12\x12\n\nblock_size\x18\x04 \x01(\r\"d\n\x12\x44raftBlockResponse\x12\x17\n\x0f\x64raft_token_ids\x18\x01 \x03(\r\x12\x16\n\x0e\x66orward_passes\x18\x02 \x01(\r\x12\x1d\n\x15peak_activation_bytes\x18\x03 \x01(\x04\"]\n\x14\x45xtendContextRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x1e\n\x03\x61ux\x18\x02 \x03(\x0b\x32\x11.kakeya.v1.Tensor\x12\x11\n\tpositions\x18\x03 \x03(\x05\",\n\x15\x45xtendContextResponse\x12\x13\n\x0b\x63ontext_len\x18\x01 \x01(\r\">\n(DFlashProposerServiceCloseSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\"+\n)DFlashProposerServiceCloseSessionResponse*\xed\x01\n\x0e\x43\x61pabilityRole\x12\x1f\n\x1b\x43\x41PABILITY_ROLE_UNSPECIFIED\x10\x00\x12\x1c\n\x18\x43\x41PABILITY_ROLE_VERIFIER\x10\x01\x12\x1c\n\x18\x43\x41PABILITY_ROLE_PROPOSER\x10\x02\x12\x1c\n\x18\x43\x41PABILITY_ROLE_EMBEDDER\x10\x03\x12\x18\n\x14\x43\x41PABILITY_ROLE_TOOL\x10\x04\x12!\n\x1d\x43\x41PABILITY_ROLE_PREFILL_CACHE\x10\x05\x12#\n\x1f\x43\x41PABILITY_ROLE_PREFILL_COMPUTE\x10\x06*\x96\x01\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_NONE\x10\x01\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZLIB\x10\x02\x12\'\n#COMPRESSION_CODEC_KAKEYA_LATTICE_D4\x10\x03*\xd8\x01\n\x10PrefillJobStatus\x12\"\n\x1ePREFILL_JOB_STATUS_UNSPECIFIED\x10\x00\x12\x1d\n\x19PREFILL_JOB_STATUS_QUEUED\x10\x01\x12\x1e\n\x1aPREFILL_JOB_STATUS_RUNNING\x10\x02\x12 \n\x1cPREFILL_JOB_STATUS_COMPLETED\x10\x03\x12\x1d\n\x19PREFILL_JOB_STATUS_FAILED\x10\x04\x12 \n\x1cPREFILL_JOB_STATUS_CANCELLED\x10\x05\x32\xdc\x01\n\x11\x43\x61pabilityService\x12g\n\x14\x45xchangeCapabilities\x12&.kakeya.v1.ExchangeCapabilitiesRequest\x1a\'.kakeya.v1.ExchangeCapabilitiesResponse\x12^\n\x11GetNodeCapability\x12#.kakeya.v1.GetNodeCapabilityRequest\x1a$.kakeya.v1.GetNodeCapabilityResponse2b\n\x0fProposerService\x12O\n\x0cProposeBlock\x12\x1e.kakeya.v1.ProposeBlockRequest\x1a\x1f.kakeya.v1.ProposeBlockResponse2\xe3\x02\n\x13PrefillCacheService\x12X\n\x0fGetCacheSummary\x12!.kakeya.v1.GetCacheSummaryRequest\x1a\".kakeya.v1.GetCacheSummaryResponse\x12O\n\x0cLookupPrefix\x12\x1e.kakeya.v1.LookupPrefixRequest\x1a\x1f.kakeya.v1.LookupPrefixResponse\x12N\n\x0b\x46\x65tchBlocks\x12\x1d.kakeya.v1.FetchBlocksRequest\x1a\x1e.kakeya.v1.FetchBlocksResponse0\x01\x12Q\n\x0cPublishBlock\x12\x1e.kakeya.v1.PublishBlockRequest\x1a\x1f.kakeya.v1.PublishBlockResponse(\x01\x32\xb6\x02\n\x14PrefillWorkerService\x12[\n\x10SubmitPrefillJob\x12\".kakeya.v1.SubmitPrefillJobRequest\x1a#.kakeya.v1.SubmitPrefillJobResponse\x12\x64\n\x13GetPrefillJobStatus\x12%.kakeya.v1.GetPrefillJobStatusRequest\x1a&.kakeya.v1.GetPrefillJobStatusResponse\x12[\n\x10\x43\x61ncelPrefillJob\x12\".kakeya.v1.CancelPrefillJobRequest\x1a#.kakeya.v1.CancelPrefillJobResponse2\xc1\x03\n\x15\x44\x46lashProposerService\x12@\n\x07Restore\x12\x19.kakeya.v1.RestoreRequest\x1a\x1a.kakeya.v1.RestoreResponse\x12L\n\x0bSeedContext\x12\x1d.kakeya.v1.SeedContextRequest\x1a\x1e.kakeya.v1.SeedContextResponse\x12I\n\nDraftBlock\x12\x1c.kakeya.v1.DraftBlockRequest\x1a\x1d.kakeya.v1.DraftBlockResponse\x12R\n\rExtendContext\x12\x1f.kakeya.v1.ExtendContextRequest\x1a .kakeya.v1.ExtendContextResponse\x12y\n\x0c\x43loseSession\x12\x33.kakeya.v1.DFlashProposerServiceCloseSessionRequest\x1a\x34.kakeya.v1.DFlashProposerServiceCloseSessionResponseb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'kakeya.v1.distributed_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None - _globals['_CAPABILITYROLE']._serialized_start=4889 - _globals['_CAPABILITYROLE']._serialized_end=5126 - _globals['_COMPRESSIONCODEC']._serialized_start=5129 - _globals['_COMPRESSIONCODEC']._serialized_end=5279 - _globals['_PREFILLJOBSTATUS']._serialized_start=5282 - _globals['_PREFILLJOBSTATUS']._serialized_end=5498 + _globals['_CAPABILITYROLE']._serialized_start=4953 + _globals['_CAPABILITYROLE']._serialized_end=5190 + _globals['_COMPRESSIONCODEC']._serialized_start=5193 + _globals['_COMPRESSIONCODEC']._serialized_end=5343 + _globals['_PREFILLJOBSTATUS']._serialized_start=5346 + _globals['_PREFILLJOBSTATUS']._serialized_end=5562 _globals['_MODELCAPABILITY']._serialized_start=42 _globals['_MODELCAPABILITY']._serialized_end=167 _globals['_NODECAPABILITY']._serialized_start=170 @@ -46,81 +46,81 @@ _globals['_CACHECOMPATIBILITY']._serialized_start=653 _globals['_CACHECOMPATIBILITY']._serialized_end=954 _globals['_CACHECAPABILITY']._serialized_start=957 - _globals['_CACHECAPABILITY']._serialized_end=1290 - _globals['_PREFILLWORKERCAPABILITY']._serialized_start=1293 - _globals['_PREFILLWORKERCAPABILITY']._serialized_end=1595 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=1597 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=1674 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=1676 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=1754 - _globals['_GETNODECAPABILITYREQUEST']._serialized_start=1756 - _globals['_GETNODECAPABILITYREQUEST']._serialized_end=1782 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=1784 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=1852 - _globals['_GETCACHESUMMARYREQUEST']._serialized_start=1854 - _globals['_GETCACHESUMMARYREQUEST']._serialized_end=1932 - _globals['_GETCACHESUMMARYRESPONSE']._serialized_start=1934 - _globals['_GETCACHESUMMARYRESPONSE']._serialized_end=2020 - _globals['_LOOKUPPREFIXREQUEST']._serialized_start=2022 - _globals['_LOOKUPPREFIXREQUEST']._serialized_end=2119 - _globals['_LOOKUPPREFIXRESPONSE']._serialized_start=2122 - _globals['_LOOKUPPREFIXRESPONSE']._serialized_end=2329 - _globals['_FETCHBLOCKSREQUEST']._serialized_start=2331 - _globals['_FETCHBLOCKSREQUEST']._serialized_end=2369 - _globals['_FETCHBLOCKSRESPONSE']._serialized_start=2372 - _globals['_FETCHBLOCKSRESPONSE']._serialized_end=2555 - _globals['_PUBLISHBLOCKREQUEST']._serialized_start=2558 - _globals['_PUBLISHBLOCKREQUEST']._serialized_end=2795 - _globals['_PUBLISHBLOCKRESPONSE']._serialized_start=2797 - _globals['_PUBLISHBLOCKRESPONSE']._serialized_end=2856 - _globals['_SUBMITPREFILLJOBREQUEST']._serialized_start=2859 - _globals['_SUBMITPREFILLJOBREQUEST']._serialized_end=3099 - _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_start=3102 - _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_end=3235 - _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_start=3237 - _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_end=3300 - _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_start=3303 - _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_end=3571 - _globals['_CANCELPREFILLJOBREQUEST']._serialized_start=3573 - _globals['_CANCELPREFILLJOBREQUEST']._serialized_end=3633 - _globals['_CANCELPREFILLJOBRESPONSE']._serialized_start=3635 - _globals['_CANCELPREFILLJOBRESPONSE']._serialized_end=3680 - _globals['_PROPOSEBLOCKREQUEST']._serialized_start=3682 - _globals['_PROPOSEBLOCKREQUEST']._serialized_end=3789 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=3791 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=3912 - _globals['_TENSOR']._serialized_start=3914 - _globals['_TENSOR']._serialized_end=3966 - _globals['_LAYERKV']._serialized_start=3968 - _globals['_LAYERKV']._serialized_end=4052 - _globals['_RESTOREREQUEST']._serialized_start=4055 - _globals['_RESTOREREQUEST']._serialized_end=4187 - _globals['_RESTORERESPONSE']._serialized_start=4189 - _globals['_RESTORERESPONSE']._serialized_end=4291 - _globals['_SEEDCONTEXTREQUEST']._serialized_start=4293 - _globals['_SEEDCONTEXTREQUEST']._serialized_end=4384 - _globals['_SEEDCONTEXTRESPONSE']._serialized_start=4386 - _globals['_SEEDCONTEXTRESPONSE']._serialized_end=4428 - _globals['_DRAFTBLOCKREQUEST']._serialized_start=4430 - _globals['_DRAFTBLOCKREQUEST']._serialized_end=4534 - _globals['_DRAFTBLOCKRESPONSE']._serialized_start=4536 - _globals['_DRAFTBLOCKRESPONSE']._serialized_end=4636 - _globals['_EXTENDCONTEXTREQUEST']._serialized_start=4638 - _globals['_EXTENDCONTEXTREQUEST']._serialized_end=4731 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=4733 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=4777 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=4779 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=4841 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=4843 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=4886 - _globals['_CAPABILITYSERVICE']._serialized_start=5501 - _globals['_CAPABILITYSERVICE']._serialized_end=5721 - _globals['_PROPOSERSERVICE']._serialized_start=5723 - _globals['_PROPOSERSERVICE']._serialized_end=5821 - _globals['_PREFILLCACHESERVICE']._serialized_start=5824 - _globals['_PREFILLCACHESERVICE']._serialized_end=6179 - _globals['_PREFILLWORKERSERVICE']._serialized_start=6182 - _globals['_PREFILLWORKERSERVICE']._serialized_end=6492 - _globals['_DFLASHPROPOSERSERVICE']._serialized_start=6495 - _globals['_DFLASHPROPOSERSERVICE']._serialized_end=6944 + _globals['_CACHECAPABILITY']._serialized_end=1354 + _globals['_PREFILLWORKERCAPABILITY']._serialized_start=1357 + _globals['_PREFILLWORKERCAPABILITY']._serialized_end=1659 + _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=1661 + _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=1738 + _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=1740 + _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=1818 + _globals['_GETNODECAPABILITYREQUEST']._serialized_start=1820 + _globals['_GETNODECAPABILITYREQUEST']._serialized_end=1846 + _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=1848 + _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=1916 + _globals['_GETCACHESUMMARYREQUEST']._serialized_start=1918 + _globals['_GETCACHESUMMARYREQUEST']._serialized_end=1996 + _globals['_GETCACHESUMMARYRESPONSE']._serialized_start=1998 + _globals['_GETCACHESUMMARYRESPONSE']._serialized_end=2084 + _globals['_LOOKUPPREFIXREQUEST']._serialized_start=2086 + _globals['_LOOKUPPREFIXREQUEST']._serialized_end=2183 + _globals['_LOOKUPPREFIXRESPONSE']._serialized_start=2186 + _globals['_LOOKUPPREFIXRESPONSE']._serialized_end=2393 + _globals['_FETCHBLOCKSREQUEST']._serialized_start=2395 + _globals['_FETCHBLOCKSREQUEST']._serialized_end=2433 + _globals['_FETCHBLOCKSRESPONSE']._serialized_start=2436 + _globals['_FETCHBLOCKSRESPONSE']._serialized_end=2619 + _globals['_PUBLISHBLOCKREQUEST']._serialized_start=2622 + _globals['_PUBLISHBLOCKREQUEST']._serialized_end=2859 + _globals['_PUBLISHBLOCKRESPONSE']._serialized_start=2861 + _globals['_PUBLISHBLOCKRESPONSE']._serialized_end=2920 + _globals['_SUBMITPREFILLJOBREQUEST']._serialized_start=2923 + _globals['_SUBMITPREFILLJOBREQUEST']._serialized_end=3163 + _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_start=3166 + _globals['_SUBMITPREFILLJOBRESPONSE']._serialized_end=3299 + _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_start=3301 + _globals['_GETPREFILLJOBSTATUSREQUEST']._serialized_end=3364 + _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_start=3367 + _globals['_GETPREFILLJOBSTATUSRESPONSE']._serialized_end=3635 + _globals['_CANCELPREFILLJOBREQUEST']._serialized_start=3637 + _globals['_CANCELPREFILLJOBREQUEST']._serialized_end=3697 + _globals['_CANCELPREFILLJOBRESPONSE']._serialized_start=3699 + _globals['_CANCELPREFILLJOBRESPONSE']._serialized_end=3744 + _globals['_PROPOSEBLOCKREQUEST']._serialized_start=3746 + _globals['_PROPOSEBLOCKREQUEST']._serialized_end=3853 + _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=3855 + _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=3976 + _globals['_TENSOR']._serialized_start=3978 + _globals['_TENSOR']._serialized_end=4030 + _globals['_LAYERKV']._serialized_start=4032 + _globals['_LAYERKV']._serialized_end=4116 + _globals['_RESTOREREQUEST']._serialized_start=4119 + _globals['_RESTOREREQUEST']._serialized_end=4251 + _globals['_RESTORERESPONSE']._serialized_start=4253 + _globals['_RESTORERESPONSE']._serialized_end=4355 + _globals['_SEEDCONTEXTREQUEST']._serialized_start=4357 + _globals['_SEEDCONTEXTREQUEST']._serialized_end=4448 + _globals['_SEEDCONTEXTRESPONSE']._serialized_start=4450 + _globals['_SEEDCONTEXTRESPONSE']._serialized_end=4492 + _globals['_DRAFTBLOCKREQUEST']._serialized_start=4494 + _globals['_DRAFTBLOCKREQUEST']._serialized_end=4598 + _globals['_DRAFTBLOCKRESPONSE']._serialized_start=4600 + _globals['_DRAFTBLOCKRESPONSE']._serialized_end=4700 + _globals['_EXTENDCONTEXTREQUEST']._serialized_start=4702 + _globals['_EXTENDCONTEXTREQUEST']._serialized_end=4795 + _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=4797 + _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=4841 + _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=4843 + _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=4905 + _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=4907 + _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=4950 + _globals['_CAPABILITYSERVICE']._serialized_start=5565 + _globals['_CAPABILITYSERVICE']._serialized_end=5785 + _globals['_PROPOSERSERVICE']._serialized_start=5787 + _globals['_PROPOSERSERVICE']._serialized_end=5885 + _globals['_PREFILLCACHESERVICE']._serialized_start=5888 + _globals['_PREFILLCACHESERVICE']._serialized_end=6243 + _globals['_PREFILLWORKERSERVICE']._serialized_start=6246 + _globals['_PREFILLWORKERSERVICE']._serialized_end=6556 + _globals['_DFLASHPROPOSERSERVICE']._serialized_start=6559 + _globals['_DFLASHPROPOSERSERVICE']._serialized_end=7008 # @@protoc_insertion_point(module_scope) diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi index 9f2f8cc1..5eb19ce2 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi @@ -131,7 +131,7 @@ class CacheCompatibility(_message.Message): def __init__(self, model_id: _Optional[str] = ..., model_revision: _Optional[str] = ..., tokenizer_revision: _Optional[str] = ..., cache_format_version: _Optional[str] = ..., quantization: _Optional[str] = ..., rope_hash: _Optional[str] = ..., layer_geometry_hash: _Optional[str] = ..., kv_dtype: _Optional[str] = ..., block_size_tokens: _Optional[int] = ..., tenant_namespace: _Optional[str] = ..., sink_size: _Optional[int] = ..., window_size: _Optional[int] = ...) -> None: ... class CacheCapability(_message.Message): - __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter", "default_compression", "replication_factor") + __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter", "default_compression", "replication_factor", "evictions", "bytes_evicted", "put_failures") COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int] CACHE_BYTES_USED_FIELD_NUMBER: _ClassVar[int] @@ -143,6 +143,9 @@ class CacheCapability(_message.Message): BLOOM_FILTER_FIELD_NUMBER: _ClassVar[int] DEFAULT_COMPRESSION_FIELD_NUMBER: _ClassVar[int] REPLICATION_FACTOR_FIELD_NUMBER: _ClassVar[int] + EVICTIONS_FIELD_NUMBER: _ClassVar[int] + BYTES_EVICTED_FIELD_NUMBER: _ClassVar[int] + PUT_FAILURES_FIELD_NUMBER: _ClassVar[int] compatibility: CacheCompatibility cache_address: str cache_bytes_used: int @@ -154,7 +157,10 @@ class CacheCapability(_message.Message): bloom_filter: bytes default_compression: CompressionCodec replication_factor: int - def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., cache_address: _Optional[str] = ..., cache_bytes_used: _Optional[int] = ..., cache_bytes_free: _Optional[int] = ..., entry_count: _Optional[int] = ..., cache_epoch: _Optional[int] = ..., load: _Optional[float] = ..., tokens_served: _Optional[int] = ..., bloom_filter: _Optional[bytes] = ..., default_compression: _Optional[_Union[CompressionCodec, str]] = ..., replication_factor: _Optional[int] = ...) -> None: ... + evictions: int + bytes_evicted: int + put_failures: int + def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., cache_address: _Optional[str] = ..., cache_bytes_used: _Optional[int] = ..., cache_bytes_free: _Optional[int] = ..., entry_count: _Optional[int] = ..., cache_epoch: _Optional[int] = ..., load: _Optional[float] = ..., tokens_served: _Optional[int] = ..., bloom_filter: _Optional[bytes] = ..., default_compression: _Optional[_Union[CompressionCodec, str]] = ..., replication_factor: _Optional[int] = ..., evictions: _Optional[int] = ..., bytes_evicted: _Optional[int] = ..., put_failures: _Optional[int] = ...) -> None: ... class PrefillWorkerCapability(_message.Message): __slots__ = ("compatibility", "worker_address", "max_concurrent_jobs", "inflight_jobs", "queued_jobs", "load", "tokens_per_second_prefill", "ram_bytes_free", "accepts_compute_jobs", "queued_tokens") diff --git a/inference_engine/session/coordinator.py b/inference_engine/session/coordinator.py index 5bdc8128..373739a3 100644 --- a/inference_engine/session/coordinator.py +++ b/inference_engine/session/coordinator.py @@ -48,7 +48,7 @@ from __future__ import annotations -from typing import Any, Iterable, List, Protocol +from typing import Any, Callable, Iterable, List, Protocol import torch @@ -145,6 +145,7 @@ def __init__( verifier: VerifierProtocol, resolver=None, prefill_cache: PrefillCacheHookProtocol | None = None, + on_first_append: Callable[[Session, list[int]], None] | None = None, ) -> None: self._store = store self._verifier = verifier @@ -153,6 +154,7 @@ def __init__( # for every session (v0.3 single-tenant behaviour, unchanged). self._resolver = resolver self._prefill_cache = prefill_cache + self._on_first_append = on_first_append def _verifier_for(self, session_id: str) -> "VerifierProtocol": return self._resolver(session_id) if self._resolver else self._verifier @@ -201,7 +203,8 @@ def append_tokens( # - cached_token_sequence is the post-trim parallel sequence # - next_global_position = sum of all tokens ever appended # - next_token_logits predicts position == next_global_position - if session.next_global_position == 0: + first_append = session.next_global_position == 0 + if first_append: if self._prefill_cache is not None: self._prefill_cache.prepare(verifier, token_list) else: @@ -241,5 +244,7 @@ def append_tokens( # so GetSessionInfo.kv_live_bytes reports physical bytes # rather than the slab's placeholder zero. PR-E1c. _sync_slab_bytes(session, verifier) + if first_append and self._on_first_append is not None: + self._on_first_append(session, token_list) return new_history_length diff --git a/proto/kakeya/v1/distributed.proto b/proto/kakeya/v1/distributed.proto index 702d1e72..229e6483 100644 --- a/proto/kakeya/v1/distributed.proto +++ b/proto/kakeya/v1/distributed.proto @@ -236,6 +236,9 @@ message CacheCapability { bytes bloom_filter = 9; CompressionCodec default_compression = 10; uint32 replication_factor = 11; + uint64 evictions = 12; + uint64 bytes_evicted = 13; + uint64 put_failures = 14; } message PrefillWorkerCapability { diff --git a/scripts/fill_prefill_cache_from_live_grpc.py b/scripts/fill_prefill_cache_from_live_grpc.py new file mode 100644 index 00000000..0dd872cb --- /dev/null +++ b/scripts/fill_prefill_cache_from_live_grpc.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""Fill distributed prefill caches from an in-memory live gRPC capture queue.""" +from __future__ import annotations + +import argparse +import json +import subprocess +import time +import urllib.request +import uuid +from pathlib import Path + + +def _request_json(url: str, *, api_key: str, body: dict | None = None): + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + url, + data=data, + headers={ + "Content-Type": "application/json", + "X-API-Key": api_key, + }, + ) + with urllib.request.urlopen(request, timeout=10) as response: + return json.load(response) + + +def _node_cache(nodes: list[dict], node_id: str) -> dict: + return next(node["cache"] for node in nodes if node["id"] == node_id) + + +def _ratio(cache: dict) -> float: + total = int(cache["bytes_used"]) + int(cache["bytes_free"]) + return int(cache["bytes_used"]) / total if total else 0.0 + + +def _memory_free_percent(ssh_target: str = "") -> int: + command = ["memory_pressure"] + if ssh_target: + command = ["ssh", "-o", "ConnectTimeout=5", ssh_target, "memory_pressure"] + output = subprocess.check_output(command, text=True, timeout=15) + marker = "System-wide memory free percentage:" + line = next(line for line in output.splitlines() if marker in line) + return int(line.split(marker, 1)[1].strip().rstrip("%")) + + +def _safe_report_item(item: dict, *, replay_tokens: int, wall_seconds: float) -> dict: + return { + "capture_id": item["capture_id"], + "captured_tokens": int(item["token_count"]), + "replay_tokens": int(replay_tokens), + "wall_seconds": wall_seconds, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--address", default="127.0.0.1:51051") + parser.add_argument("--dashboard", default="http://127.0.0.1:8090") + parser.add_argument("--api-key-file", default="~/.kakeya/network_api_key") + parser.add_argument("--tokenizer-id", required=True) + parser.add_argument("--head-node-id", default="head-runtime") + parser.add_argument("--cache-node-id", default="allens-mini") + parser.add_argument("--target-one", type=float, default=0.90) + parser.add_argument("--target-two", type=float, default=0.95) + parser.add_argument("--churn-gb", type=float, default=0.9) + parser.add_argument("--pause-seconds", type=float, default=300.0) + parser.add_argument("--poll-seconds", type=float, default=5.0) + parser.add_argument("--capture-batch", type=int, default=8) + parser.add_argument("--min-memory-free-percent", type=int, default=5) + parser.add_argument("--peer-ssh", default="allen@169.254.27.104") + parser.add_argument("--stop-file", default="/tmp/kakeya-cache-fill.stop") + parser.add_argument("--report", default="/tmp/kakeya-cache-fill-report.json") + args = parser.parse_args() + if not (0 < args.target_one <= args.target_two < 1): + raise SystemExit("targets must satisfy 0 < target-one <= target-two < 1") + + from kakeya import Client + from kakeya.errors import ResourceExhaustedError + from transformers import AutoTokenizer + from scripts.chat_grpc import _resolve_eos_token_ids + + api_key = Path(args.api_key_file).expanduser().read_text().strip() + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + eos_ids = _resolve_eos_token_ids(tokenizer) + stop_file = Path(args.stop_file) + report = { + "schema_version": 1, + "started_at": time.time(), + "items": [], + "stages": [], + } + captured: list[dict] = [] + replay_index = 0 + baseline_fallbacks = 0 + baseline_publish_failures = 0 + + def snapshot(): + nodes = _request_json(f"{args.dashboard}/v1/network/nodes", api_key=api_key) + prefill = _request_json(f"{args.dashboard}/v1/network/prefill", api_key=api_key) + head = _node_cache(nodes, args.head_node_id) + peer = _node_cache(nodes, args.cache_node_id) + return { + "head": head, + "peer": peer, + "head_ratio": _ratio(head), + "peer_ratio": _ratio(peer), + "prefill": prefill, + } + + def safety(current): + if stop_file.exists(): + raise RuntimeError(f"stop file present: {stop_file}") + if ( + int(current["prefill"].get("publish_failures", 0)) + > baseline_publish_failures + ): + raise RuntimeError(current["prefill"].get("last_publish_error", "publish failure")) + if int(current["prefill"].get("fallbacks", 0)) > baseline_fallbacks: + raise RuntimeError(current["prefill"].get("last_fallback_reason", "prefill fallback")) + local_free = _memory_free_percent() + peer_free = _memory_free_percent(args.peer_ssh) + if min(local_free, peer_free) < args.min_memory_free_percent: + raise RuntimeError( + f"memory pressure: head={local_free}% peer={peer_free}% free", + ) + + def get_captures(): + response = _request_json( + f"{args.dashboard}/v1/network/maintenance/capture/drain", + api_key=api_key, + body={"max_items": args.capture_batch}, + ) + captured.extend(response["items"]) + + def replay_one(): + nonlocal replay_index + if not captured: + get_captures() + if not captured: + time.sleep(args.poll_seconds) + return + item = captured[replay_index % len(captured)] + replay_index += 1 + nonce = tokenizer.encode( + f"cache-fill-{uuid.uuid4().hex} ", + add_special_tokens=False, + ) + token_ids = [*nonce, *item["token_ids"]] + started = time.perf_counter() + try: + with Client(args.address) as client: + with client.create_session( + eos_token_ids=eos_ids, + client_label=f"cache-fill-{replay_index}", + ) as session: + session.append(token_ids) + list(session.generate(max_tokens=1)) + except ResourceExhaustedError: + time.sleep(args.poll_seconds) + return + report["items"].append(_safe_report_item( + item, + replay_tokens=len(token_ids), + wall_seconds=time.perf_counter() - started, + )) + + def fill_to(name: str, target: float): + while True: + current = snapshot() + safety(current) + if ( + current["head_ratio"] >= target + and current["peer_ratio"] >= target + ): + report["stages"].append({ + "name": name, + "completed_at": time.time(), + "snapshot": current, + }) + Path(args.report).write_text(json.dumps(report, indent=2)) + return current + replay_one() + + initial = snapshot() + baseline_fallbacks = int(initial["prefill"].get("fallbacks", 0)) + baseline_publish_failures = int( + initial["prefill"].get("publish_failures", 0), + ) + report["baseline"] = initial + fill_to("target_one", args.target_one) + time.sleep(args.pause_seconds) + at_target_two = fill_to("target_two", args.target_two) + time.sleep(args.pause_seconds) + churn_target = int(at_target_two["peer"].get("bytes_evicted", 0)) + int( + args.churn_gb * (1 << 30), + ) + while True: + current = snapshot() + safety(current) + if int(current["peer"].get("bytes_evicted", 0)) >= churn_target: + report["stages"].append({ + "name": "lru_churn", + "completed_at": time.time(), + "snapshot": current, + }) + break + replay_one() + report["finished_at"] = time.time() + Path(args.report).write_text(json.dumps(report, indent=2)) + print(json.dumps({ + "ok": True, + "report": args.report, + "replays": len(report["items"]), + "final": report["stages"][-1], + }, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index 00192170..dccc740f 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -488,11 +488,30 @@ async def _serve(args: argparse.Namespace) -> int: slab_pool=pool, ) resolver = registry.get if registry is not None else None + cache_fill_capture = None + if args.cache_fill_capture_size: + from inference_engine.distributed.cache_fill import CacheFillCapture + cache_fill_capture = CacheFillCapture( + max_items=args.cache_fill_capture_size, + ) + _LOG.info( + "maintenance cache-fill capture enabled: max_items=%d", + args.cache_fill_capture_size, + ) append_coord = AppendTokensCoordinator( store, verifier, resolver=resolver, prefill_cache=prefill_hook, + on_first_append=( + ( + lambda session, tokens: cache_fill_capture.observe( + client_label=session.client_label, + token_ids=tokens, + ) + ) + if cache_fill_capture is not None else None + ), ) gen_coord = GenerationCoordinator( store, @@ -568,6 +587,7 @@ async def _serve(args: argparse.Namespace) -> int: create_network_app( network_state, api_key=args.network_api_key, + cache_fill_capture=cache_fill_capture, ), host=args.network_http_host, port=args.network_http_port, @@ -761,6 +781,13 @@ def main() -> int: ap.add_argument("--network-telemetry-url", default="", help="Optional POST endpoint receiving completed token counters.") ap.add_argument("--network-telemetry-api-key", default="") + ap.add_argument( + "--cache-fill-capture-size", + type=int, + default=0, + help="Maintenance-only in-memory first-append capture queue size; " + "0 disables capture.", + ) ap.add_argument("--skip-cache-check", action="store_true", help="Skip the HF-cache pre-flight assertion. By " "default the server fails fast if the verifier " diff --git a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts index df8d60ca..84ddf21b 100644 --- a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts +++ b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts @@ -327,6 +327,9 @@ export interface CacheCapability { bloomFilter: Uint8Array; defaultCompression: CompressionCodec; replicationFactor: number; + evictions: string; + bytesEvicted: string; + putFailures: string; } export interface PrefillWorkerCapability { @@ -1418,6 +1421,9 @@ function createBaseCacheCapability(): CacheCapability { bloomFilter: new Uint8Array(0), defaultCompression: 0, replicationFactor: 0, + evictions: "0", + bytesEvicted: "0", + putFailures: "0", }; } @@ -1456,6 +1462,15 @@ export const CacheCapability: MessageFns = { if (message.replicationFactor !== 0) { writer.uint32(88).uint32(message.replicationFactor); } + if (message.evictions !== "0") { + writer.uint32(96).uint64(message.evictions); + } + if (message.bytesEvicted !== "0") { + writer.uint32(104).uint64(message.bytesEvicted); + } + if (message.putFailures !== "0") { + writer.uint32(112).uint64(message.putFailures); + } return writer; }, @@ -1554,6 +1569,30 @@ export const CacheCapability: MessageFns = { message.replicationFactor = reader.uint32(); continue; } + case 12: { + if (tag !== 96) { + break; + } + + message.evictions = reader.uint64().toString(); + continue; + } + case 13: { + if (tag !== 104) { + break; + } + + message.bytesEvicted = reader.uint64().toString(); + continue; + } + case 14: { + if (tag !== 112) { + break; + } + + message.putFailures = reader.uint64().toString(); + continue; + } } if ((tag & 7) === 4 || tag === 0) { break; @@ -1612,6 +1651,17 @@ export const CacheCapability: MessageFns = { : isSet(object.replication_factor) ? globalThis.Number(object.replication_factor) : 0, + evictions: isSet(object.evictions) ? globalThis.String(object.evictions) : "0", + bytesEvicted: isSet(object.bytesEvicted) + ? globalThis.String(object.bytesEvicted) + : isSet(object.bytes_evicted) + ? globalThis.String(object.bytes_evicted) + : "0", + putFailures: isSet(object.putFailures) + ? globalThis.String(object.putFailures) + : isSet(object.put_failures) + ? globalThis.String(object.put_failures) + : "0", }; }, @@ -1650,6 +1700,15 @@ export const CacheCapability: MessageFns = { if (message.replicationFactor !== 0) { obj.replicationFactor = Math.round(message.replicationFactor); } + if (message.evictions !== "0") { + obj.evictions = message.evictions; + } + if (message.bytesEvicted !== "0") { + obj.bytesEvicted = message.bytesEvicted; + } + if (message.putFailures !== "0") { + obj.putFailures = message.putFailures; + } return obj; }, @@ -1671,6 +1730,9 @@ export const CacheCapability: MessageFns = { message.bloomFilter = object.bloomFilter ?? new Uint8Array(0); message.defaultCompression = object.defaultCompression ?? 0; message.replicationFactor = object.replicationFactor ?? 0; + message.evictions = object.evictions ?? "0"; + message.bytesEvicted = object.bytesEvicted ?? "0"; + message.putFailures = object.putFailures ?? "0"; return message; }, }; diff --git a/tests/inference_engine/bridge/test_cache_fill_script.py b/tests/inference_engine/bridge/test_cache_fill_script.py new file mode 100644 index 00000000..3c75899f --- /dev/null +++ b/tests/inference_engine/bridge/test_cache_fill_script.py @@ -0,0 +1,31 @@ +from scripts.fill_prefill_cache_from_live_grpc import ( + _node_cache, + _ratio, + _safe_report_item, +) + + +def test_capacity_helpers_are_node_specific(): + nodes = [ + {"id": "head", "cache": {"bytes_used": 9, "bytes_free": 1}}, + {"id": "peer", "cache": {"bytes_used": 8, "bytes_free": 2}}, + ] + assert _ratio(_node_cache(nodes, "head")) == 0.9 + assert _ratio(_node_cache(nodes, "peer")) == 0.8 + assert _ratio({"bytes_used": 0, "bytes_free": 0}) == 0.0 + + +def test_safe_report_never_contains_token_ids(): + item = { + "capture_id": "salted-id", + "token_count": 512, + "token_ids": [1, 2, 3], + } + report = _safe_report_item(item, replay_tokens=520, wall_seconds=1.25) + assert report == { + "capture_id": "salted-id", + "captured_tokens": 512, + "replay_tokens": 520, + "wall_seconds": 1.25, + } + assert "token_ids" not in report diff --git a/tests/inference_engine/distributed/test_cache_fill.py b/tests/inference_engine/distributed/test_cache_fill.py new file mode 100644 index 00000000..cf511049 --- /dev/null +++ b/tests/inference_engine/distributed/test_cache_fill.py @@ -0,0 +1,45 @@ +from inference_engine.distributed.cache_fill import CacheFillCapture + + +def test_capture_deduplicates_excludes_replay_and_drains(): + capture = CacheFillCapture(max_items=2) + assert capture.observe(client_label="live", token_ids=[1, 2, 3]) + assert not capture.observe(client_label="live", token_ids=[1, 2, 3]) + assert not capture.observe(client_label="cache-fill-1", token_ids=[4]) + assert capture.stats() == { + "queued": 1, + "captured": 1, + "duplicates": 1, + "dropped": 0, + "max_items": 2, + } + item = capture.drain(1)[0] + assert item.token_ids == (1, 2, 3) + assert item.token_count == 3 + assert capture.stats()["queued"] == 0 + + +def test_capture_is_bounded_and_allows_recapture_after_drain(): + capture = CacheFillCapture(max_items=1) + assert capture.observe(client_label="a", token_ids=[1]) + assert capture.observe(client_label="b", token_ids=[2]) + assert capture.stats()["dropped"] == 1 + assert capture.drain(1)[0].token_ids == (2,) + assert capture.observe(client_label="a", token_ids=[1]) + + +def test_capture_validates_limits_and_empty_input(): + try: + CacheFillCapture(max_items=0) + except ValueError: + pass + else: + raise AssertionError("expected max_items validation") + capture = CacheFillCapture() + assert not capture.observe(client_label="live", token_ids=[]) + try: + capture.drain(0) + except ValueError: + pass + else: + raise AssertionError("expected drain validation") diff --git a/tests/inference_engine/distributed/test_capability.py b/tests/inference_engine/distributed/test_capability.py index e730eed3..0ea9dcda 100644 --- a/tests/inference_engine/distributed/test_capability.py +++ b/tests/inference_engine/distributed/test_capability.py @@ -135,6 +135,9 @@ def test_cache_capability_and_endpoints_proto_round_trip(): bloom_filter=b"filter", default_compression=CompressionCodec.ZLIB, replication_factor=2, + evictions=5, + bytes_evicted=6, + put_failures=7, ), ), endpoints=( diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index 2aacbae8..01455a2e 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -61,6 +61,10 @@ def test_store_miss_expiry_collision_and_lru(): store.put(CacheBlock.create(hashes[0], 2, b"zz")) store.put(CacheBlock.create(hashes[1], 4, b"bbb")) assert hashes[0] not in store.block_hashes() + stats = store.stats() + assert stats.evictions == 1 + assert stats.bytes_evicted == 2 + assert stats.put_failures == 1 miss = store.lookup([hashes[0]], now=20.0) assert not miss.lease_id lease = store.lookup([hashes[1]], lease_seconds=1, now=20.0) @@ -83,6 +87,7 @@ def test_validation_and_stats(): stats = store.stats() assert stats.entry_count == 0 assert stats.max_bytes == 10 + assert stats.put_failures == 1 with pytest.raises(ValueError, match="one payload"): store.put_prefix([1, 2, 3], [b"only-one"]) with pytest.raises(ValueError, match="lease_seconds"): @@ -117,3 +122,4 @@ def test_put_rejects_when_active_lease_pins_capacity(): with pytest.raises(ValueError, match="pinned"): store.put(second) assert store.block_hashes() == (first.block_hash,) + assert store.stats().put_failures == 1 diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py index 5dd696d1..6755aea0 100644 --- a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py @@ -326,9 +326,17 @@ def test_publish_boundary_dispatches_selected_replica(monkeypatch): ) calls = [] + class Future: + def result(self): + return True + + def add_done_callback(self, callback): + callback(self) + class Publisher: def submit(self, fn, *args, **kwargs): calls.append((fn, args, kwargs)) + return Future() def shutdown(self, **kwargs): pass @@ -339,6 +347,17 @@ def shutdown(self, **kwargs): hashes = chained_block_hashes([1, 2], compatibility) hook._publish_boundary(verifier, [1, 2], hashes, 0, 2) assert calls and calls[0][1][0] == "peer:1" + assert hook.stats.publish_attempts == 1 + assert hook.stats.publish_successes == 1 + assert hook.stats.bytes_published > 0 + + class FailedFuture: + def result(self): + raise RuntimeError("publish failed") + + hook._publish_done(FailedFuture(), 10) + assert hook.stats.publish_failures == 1 + assert "publish failed" in hook.stats.last_publish_error hook.close() diff --git a/tests/inference_engine/network/test_network_api.py b/tests/inference_engine/network/test_network_api.py index 51a69402..0d87866a 100644 --- a/tests/inference_engine/network/test_network_api.py +++ b/tests/inference_engine/network/test_network_api.py @@ -7,6 +7,7 @@ CapabilityRegistry, NodeCapability, ) +from inference_engine.distributed.cache_fill import CacheFillCapture from inference_engine.distributed.prefill_cache import PrefixCacheStore from inference_engine.network.api import create_network_app from inference_engine.network.state import NetworkState @@ -24,8 +25,14 @@ def _client(tmp_path): "tokens_reused": 192, }, ) - client = TestClient(create_network_app(state, api_key="secret")) + capture = CacheFillCapture(max_items=4) + client = TestClient(create_network_app( + state, + api_key="secret", + cache_fill_capture=capture, + )) client.network_state = state + client.cache_fill_capture = capture return client @@ -35,6 +42,7 @@ def test_dashboard_health_and_read_apis(tmp_path): dashboard = client.get("/network").text assert "Kakeya Inference Network" in dashboard assert "Remote prefill jobs" in dashboard + assert "LRU evictions" in dashboard assert client.get("/healthz").json()["status"] == "ok" assert client.get("/v1/network/summary").json()["online_nodes"] == 1 assert len(client.get("/v1/network/nodes").json()) == 1 @@ -42,6 +50,7 @@ def test_dashboard_health_and_read_apis(tmp_path): assert "nodes" in client.get("/v1/network/topology").json() assert client.get("/v1/network/tokens").json()["completed"] == 0 assert client.get("/v1/network/prefill").json()["remote_jobs"] == 3 + assert client.get("/v1/network/maintenance/capture").status_code == 401 events = client.get("/v1/network/events?once=true") assert events.status_code == 200 assert "event: summary" in events.text @@ -71,6 +80,18 @@ def test_write_apis_require_key_and_update_state(tmp_path): ) assert telemetry.json()["status"] == "accepted" assert client.get("/v1/network/tokens").json()["kv_assisted"] == 7 + client.cache_fill_capture.observe(client_label="live", token_ids=[1, 2]) + status = client.get( + "/v1/network/maintenance/capture", + headers={"X-API-Key": "secret"}, + ) + assert status.json()["queued"] == 1 + drained = client.post( + "/v1/network/maintenance/capture/drain", + json={"max_items": 1}, + headers={"X-API-Key": "secret"}, + ) + assert drained.json()["items"][0]["token_ids"] == [1, 2] def test_write_error_mapping_and_event_stream(tmp_path, monkeypatch): @@ -106,3 +127,28 @@ def test_write_error_mapping_and_event_stream(tmp_path, monkeypatch): json={"node_id": "a", "completed": 1}, headers={"X-API-Key": "secret"}, ).status_code == 400 + + +def test_disabled_capture_and_missing_maintenance_key(tmp_path): + compatibility = CacheCompatibility(model_id="m") + state = NetworkState( + CapabilityRegistry(NodeCapability(node_id="head", grpc_address="head:1")), + PrefixCacheStore(compatibility, max_bytes=100, node_id="head"), + state_path=tmp_path / "disabled.json", + ) + without_capture = TestClient(create_network_app(state, api_key="secret")) + assert without_capture.get( + "/v1/network/maintenance/capture", + headers={"X-API-Key": "secret"}, + ).status_code == 404 + assert without_capture.post( + "/v1/network/maintenance/capture/drain", + json={"max_items": 1}, + headers={"X-API-Key": "secret"}, + ).status_code == 404 + without_key = TestClient(create_network_app( + state, + api_key="", + cache_fill_capture=CacheFillCapture(), + )) + assert without_key.get("/v1/network/maintenance/capture").status_code == 401 diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py index 8c969a9d..8e6c1649 100644 --- a/tests/inference_engine/network/test_network_state.py +++ b/tests/inference_engine/network/test_network_state.py @@ -26,6 +26,9 @@ def _state(tmp_path): compatibility, cache_address="head:2", cache_bytes_free=1000, + evictions=2, + bytes_evicted=300, + put_failures=1, ), ), endpoints=(NodeEndpoint("head:2", "thunderbolt", 100, 0.4),), @@ -64,6 +67,9 @@ def test_registration_groups_tokens_and_persistence(tmp_path): assert summary["completed_tokens"] == 100 assert summary["kv_hit_rate"] == 0.7 assert summary["prefill"]["remote_jobs"] == 2 + assert summary["cache_evictions"] == 2 + assert summary["cache_bytes_evicted"] == 300 + assert summary["cache_put_failures"] == 1 assert state.prefill_stats()["tokens_reused"] == 128 assert state.groups()[0]["id"] == group["id"] assert state.topology()["edges"][0]["target"] == "peer" diff --git a/tests/inference_engine/session/test_coordinator_validation.py b/tests/inference_engine/session/test_coordinator_validation.py index b147c5df..cf7b7c39 100644 --- a/tests/inference_engine/session/test_coordinator_validation.py +++ b/tests/inference_engine/session/test_coordinator_validation.py @@ -76,6 +76,46 @@ def test_constructor_stores_references_without_calling_them(): assert coord._verifier is sentinel_verifier +def test_first_append_callback_receives_session_and_tokens_once(): + class Verifier: + cached_token_sequence = [] + next_global_position = 0 + next_token_logits = None + + def prefill(self, tokens): + self.cached_token_sequence = list(tokens) + self.next_global_position = len(tokens) + + def forward_block(self, tokens): + self.cached_token_sequence.extend(tokens) + self.next_global_position += len(tokens) + return [type("Row", (), {"clone": lambda self: self})() for _ in tokens] + + def commit_or_truncate(self, *, forwarded, accepted): + assert forwarded == accepted + + def k_seq_length(self, _session): + return len(self.cached_token_sequence) + + def kv_live_bytes(self, _session): + return 0 + + verifier = Verifier() + store = SessionStore(capacity=1, cache_inspector=verifier) + session = store.create_session(client_label="live") + observed = [] + coord = AppendTokensCoordinator( + store, + verifier, + on_first_append=lambda sess, tokens: observed.append( + (sess.client_label, list(tokens)), + ), + ) + coord.append_tokens(session.session_id, [1, 2]) + coord.append_tokens(session.session_id, [3]) + assert observed == [("live", [1, 2])] + + # Note: tests for the ``_sync_slab_bytes`` helper (PR-E1c addition) # live in PR-E1c's own commit. PR-N1 is branched off main; once # PR-E1c merges, a follow-up will add the helper's None-branch test