From 0be2d4f38632357b8d085d90746537e13077e242 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 11 Jul 2026 15:09:11 +0000 Subject: [PATCH 01/11] docs(architecture): adopt primary-decode distributed-prefill workers Co-authored-by: FluffyAIcode --- README.md | 53 +++---- ...17-prefill-compute-worker-orchestration.md | 131 ++++++++++++++++++ docs/adr/README.md | 9 +- docs/ops/distributed-prefill-kv-network.md | 72 ++++++++-- 4 files changed, 227 insertions(+), 38 deletions(-) create mode 100644 docs/adr/0017-prefill-compute-worker-orchestration.md diff --git a/README.md b/README.md index 3b46af79..c105ef8e 100644 --- a/README.md +++ b/README.md @@ -9,17 +9,19 @@ Kakeya is a **memory-bounded local agent runtime**: a long-running inference server that holds session state on the server side, exposes a gRPC `RuntimeService`, and bounds **KV memory + per-turn latency** on long -conversations — its KV footprint **does not grow with the conversation** (see -[Kakeya Attention](#how-this-differs--kakeya-attention-vs-pagedattention--radixattention)). - -**v0.4** pairs a frozen **AR verifier** (Gemma-4 26B-A4B) with a **dLLM proposer** -(DFlash) and a trained projection **f_θ**: a sliding-window-bounded KV cache plus -**K/V restoration** reconstructs evicted context on demand, so memory is bounded -**without trading away recall, throughput, or context length**. It ships per -platform — **`v0.4-mac`** (Apple-Silicon MLX) and **`v0.4-cuda`** (NVIDIA) — atop -the session-bound gRPC runtime whose foundation landed in June 2026 -([ADR 0008](docs/adr/0008-session-bound-runtime-and-grpc-protocol.md): 9 ms -latency drift over a 4-hour, 480-turn Mac M4 run; bounded memory). +conversations — its active decode KV footprint **does not grow with the +conversation**. + +**Current architecture (ADR 0017):** one primary Mac owns RuntimeService, +sessions and every decode step. Compatible peer Macs are either +`PREFILL_COMPUTE` workers (same model, prefill only) or `PREFILL_CACHE` nodes +(RAM-only). The primary imports one immutable longest-prefix KV snapshot, +computes a missing suffix locally, and decodes entirely locally. + +> **Legacy notice:** DFlash, f_θ and fused speculative-decode experiments are +> retained for reproducibility, but they are no longer the product architecture +> or a completeness criterion. Historical v0.4 sections below are labeled as +> such. ``` ┌────────────────────┐ gRPC bidi-stream ┌────────────────────────┐ @@ -33,27 +35,26 @@ latency drift over a 4-hour, 480-turn Mac M4 run; bounded memory). │ └────────┬─────────┘ │ │ ▼ │ │ ┌──────────────────┐ │ - │ │ Restored verifier│ │ + │ │ Primary verifier│ │ │ │ Gemma-4 26B (AR) │ │ - │ │ + DFlash proposer│ │ - │ │ + f_θ / S5 │ │ │ │ bounded sink+win │ │ + │ │ peer-prefill KV │ │ │ │ per-session bind │ │ │ └──────────────────┘ │ └────────────────────────┘ ``` -> The verifier slot is pluggable: a small **Qwen3 (CPU/MLX)** sink+window -> verifier for lightweight serving, or the **restored Gemma-4 26B** path -> (proposer + f_θ/S5) for the memory-bounded, recall-preserving engine below. +> The verifier slot remains pluggable. All nodes in one prefill group must use +> the exact same model/tokenizer/quantization/RoPE/cache geometry. ## Distributed Prefill KV Cache Network -Kakeya can use trusted peer Mac minis as an **immutable prefill-cache tier**. -Every node advertises model/cache compatibility through the existing P2P -`CapabilityService`; a cold inference node queries local and remote caches in -parallel, imports the longest valid token-prefix snapshot once, computes only -the missing suffix, and keeps autoregressive decode entirely local. +Kakeya uses trusted peer Mac minis as an **immutable prefill-cache and compute +tier**. Every node advertises exact compatibility through P2P +`CapabilityService`. A cold primary queries compatible snapshots; on a miss it +can submit the prompt to a load-aware `PREFILL_COMPUTE` worker. The worker +prefills with the same model and retains the snapshot in RAM. The primary +imports it once and keeps autoregressive decode entirely local. This is not remote attention and not coherent shared RAM: @@ -73,6 +74,9 @@ Key properties: - exact model/tokenizer/quantization/RoPE/cache-format compatibility; - longest **contiguous** prefix reuse — arbitrary holes are never reused; - memory-bounded LRU storage with leases and cache epochs; +- load/cost-aware prefill workers, idempotent jobs and local fallback; +- gossip-driven worker/cache discovery and deterministic bounded replication; +- optional zlib framing, fleet-PSK authentication and tenant-HMAC prefix hashes; - point-to-point chunked gRPC publish/fetch with SHA-256 validation; - failure-safe fallback to local prefill; - Thunderbolt/LAN/Tailscale endpoint priority; @@ -91,6 +95,7 @@ observed speedup ≈97× Architecture and operations: - [ADR 0016 — Distributed Prefill KV Cache Network](docs/adr/0016-distributed-prefill-kv-cache-network.md) +- [ADR 0017 — Primary-decode / distributed-prefill orchestration](docs/adr/0017-prefill-compute-worker-orchestration.md) - [Two-Mac live report](docs/reports/distributed-prefill-kv-mac-thunderbolt.md) - [Operator runbook](docs/ops/distributed-prefill-kv-network.md) @@ -136,7 +141,7 @@ For a full 10-minute walkthrough — Mac vs Linux setup, troubleshooting, HuggingFace cache pre-warm, mainland-China mirror routing, gRPC SDK patterns — see [`docs/quickstart.md`](docs/quickstart.md). -## What's in the v0.4 architecture +## Historical: v0.4 DFlash/f_θ architecture (no longer the product path) | Component | What it does | Where | | --- | --- | --- | @@ -451,7 +456,7 @@ H200, gemma-4-26B-A4B, recall **1.0**: > models — the large ~6× memory differentiator, where a bounded window *without* > restoration would destroy recall — is the **v0.6** roadmap item. -## v0.4 for Mac — MLX speculative-decode port (the journey to parity) +## Historical: v0.4 for Mac — MLX speculative-decode research After the **CUDA** path (f_θ + S5 K/V-restoration verifier, **fused DFlash spec-decode at 1.79–2.06× AR, recall 1.0 on Gemma-4-26B-A4B / H200**), the engine diff --git a/docs/adr/0017-prefill-compute-worker-orchestration.md b/docs/adr/0017-prefill-compute-worker-orchestration.md new file mode 100644 index 00000000..662c24e4 --- /dev/null +++ b/docs/adr/0017-prefill-compute-worker-orchestration.md @@ -0,0 +1,131 @@ +# ADR 0017 — Primary-decode / distributed-prefill worker orchestration + +- **Status:** Accepted / implementation +- **Date:** 2026-07-11 +- **Supersedes for product architecture:** DFlash/f_θ/fused speculative-decode + serving paths +- **Extends:** ADR 0009 (capability gossip), ADR 0016 (distributed prefill KV) + +## Context + +The product topology is one primary Mac mini that owns RuntimeService, session +state and every autoregressive decode step. Peer Mac minis run the exact same +model but are restricted to prefill work. They retain immutable prefill K/V +snapshots in unified memory and answer lookup/fetch requests from the primary. +Cache-only peers remain useful when maximizing snapshot capacity matters more +than compute. + +ADR 0016 implemented the cache half of this design, but its deployed peer is +cache-only: the primary computes every miss and replicates snapshots. It does +not let a peer compute a cold prefix. Peer addresses are also static CLI flags, +and a fetch/import failure can escape instead of falling back to local prefill. + +## Decision + +Kakeya adopts three explicit fleet roles: + +1. **PRIMARY_DECODE** (`VERIFIER` on the wire): serves users and performs all + decode. No remote RPC is allowed in the token loop. +2. **PREFILL_COMPUTE:** loads the same model, accepts bounded/idempotent prefill + jobs, writes snapshots to its local `PrefixCacheStore`, and never serves + user decode. +3. **PREFILL_CACHE:** loads no model; spends RAM on immutable snapshots only. + +### Request flow + +For a cold append, the primary: + +1. computes tenant-HMAC chained prefix hashes; +2. queries its local cache and compatible live cache/worker cards from gossip; +3. imports the best cache hit when transfer cost beats local recomputation; +4. on a miss, submits an idempotent job to the least-loaded compatible + `PREFILL_COMPUTE` worker when remote compute+transfer is cheaper; +5. fetches and imports the completed snapshot once; +6. computes any missing suffix locally; +7. decodes entirely locally; +8. publishes snapshots to a deterministic subset of cache peers. + +Every remote error (lookup, job, lease, fetch, checksum, decompress, import) +resets the verifier and falls back to full local prefill. Cache availability +must never determine request correctness. + +### Discovery and placement + +Capability gossip is the only membership source. Static `--cache-peer` and +`--prefill-worker` flags remain emergency/operator overrides. Cards advertise +exact cache compatibility, cache address, compute address, queue depth, +inflight jobs, measured prefill throughput, free RAM and endpoint RTT. + +The scheduler is deterministic and cost-aware: + +```text +local_ms = missing_tokens / local_prefill_tps * 1000 +import_ms = endpoint_rtt_ms + transfer_bytes / link_bytes_per_ms +remote_ms = queue_eta_ms + prompt_tokens / worker_tps * 1000 + import_ms +``` + +The longest safe hit is preferred only when `import_ms < local_ms`. A worker is +used only when `remote_ms < local_ms`. Unknown metrics use conservative +operator-configured defaults. + +### Memory/storage policy + +- Decode KV remains local to the primary. +- Peer memory is a pre-decode snapshot tier, not coherent remote attention RAM. +- Snapshot payloads support zlib framing and retain SHA-256 of the uncompressed + bytes. +- Replication uses rendezvous hashing and a bounded replication factor instead + of publishing every snapshot to every peer. +- Block size controls checkpoint sparsity. Cache-only nodes should use larger + blocks for long prompts; delta snapshots remain a future format revision. +- Import checks advertised transfer size against a configurable byte budget + before allocating/reassembling. + +### Security and tenant isolation + +Trusted-LAN mode remains available, but production mode uses a fleet PSK: + +- request metadata is HMAC-SHA256 signed with timestamp and node/tenant id; +- replay window is bounded; +- prefix hashes are HMACed per tenant, preventing cross-tenant prefix probing; +- tenant namespace is part of `CacheCompatibility`; +- cache/worker services reject unauthenticated requests before allocation. + +mTLS can replace PSK transport later without changing the application protocol. + +## Correctness gates + +Blocking tests cover: + +- scheduler decisions, worker job lifecycle and idempotency; +- auth/replay/tenant isolation; +- compression/checksum and replication placement; +- every remote failure falling back to local prefill; +- real MLX local-prefill vs remote-prefill-import continuation logits and + argmax equivalence; +- zero prefill/cache RPCs during `Generate`. + +## Consequences + +Positive: + +- the primary dedicates compute to decode; +- peer compute and RAM scale independently; +- repeated system/RAG prefixes survive primary restarts; +- peers join/leave through existing gossip/TTL; +- local privacy is preserved inside the authenticated fleet. + +Costs: + +- every compute worker duplicates model weights, reducing RAM available for KV; +- first-use remote prefill only wins when worker compute plus transfer beats the + primary; +- snapshots still duplicate bounded state at checkpoint boundaries; +- MLX worker execution is serialized per loaded verifier in the MVP. + +## Legacy paths + +DFlash, f_θ and fused speculative-decode modules remain temporarily under +`research/legacy` compatibility surfaces for reproducibility. They are not +product completeness criteria and must not be wired into RuntimeService. + diff --git a/docs/adr/README.md b/docs/adr/README.md index bd7a7410..cf487188 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -1,7 +1,7 @@ # Architecture Decision Records This directory contains Architecture Decision Records (ADRs) for the -DLM-proposer + AR-verifier project. Each ADR captures a single architectural +Kakeya distributed inference runtime. Each ADR captures a single architectural decision, the context that led to it, the alternatives considered, and the consequences of choosing one path over another. @@ -40,12 +40,13 @@ reader what was *not* chosen. | 0006 | [Project positioning as local agent infrastructure](0006-local-agent-infrastructure-positioning.md) | Accepted | | 0007 | [Cross-request KV cache reuse for long sessions](0007-cross-request-kv-reuse.md) | Superseded by 0008 | | 0008 | [Session-bound runtime + gRPC protocol](0008-session-bound-runtime-and-grpc-protocol.md) | Accepted | -| 0009 | [Multi-host milestone: AR-verifier / dLM-proposer on mlx.distributed + agent capability exchange](0009-mlx-distributed-spec-decode-and-capability-exchange.md) | Accepted | +| 0009 | [Multi-host milestone: AR-verifier / dLM-proposer on mlx.distributed + agent capability exchange](0009-mlx-distributed-spec-decode-and-capability-exchange.md) | Capability gossip retained; proposer path legacy | | 0012 | [Proposer/verifier value proposition: bounded-memory + recall, platform-forked throughput](0012-proposer-verifier-value-proposition.md) | Accepted | | 0013 | [Distributed inference topology: what AR sequentiality allows](0013-distributed-inference-topology.md) | Accepted | | 0014 | [Agent-connection capacity & cross-host proposer/verifier topology: test plan & results](0014-agent-connection-capacity-and-cross-host-topology-tests.md) | Accepted | -| 0015 | [Kakeya Inference Engine: a product-grade vLLM replacement, Kakeya Attention native](0015-kakeya-attention-and-engine-substrate.md) | Accepted | -| 0016 | [Distributed Prefill KV Cache Network](0016-distributed-prefill-kv-cache-network.md) | Proposed / MVP | +| 0015 | [Kakeya Inference Engine: a product-grade vLLM replacement, Kakeya Attention native](0015-kakeya-attention-and-engine-substrate.md) | Superseded by 0017 for product architecture | +| 0016 | [Distributed Prefill KV Cache Network](0016-distributed-prefill-kv-cache-network.md) | Accepted foundation | +| 0017 | [Primary-decode / distributed-prefill worker orchestration](0017-prefill-compute-worker-orchestration.md) | Accepted / implementation | Note: ADR numbering is monotonically increasing; in-flight or planned numbers (0005) appear in the index so readers can diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 50c69535..93f26b7d 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -1,8 +1,8 @@ # Distributed Prefill KV Cache Network — Operator Runbook -This runbook deploys one inference head and one cache peer over a private -Thunderbolt Bridge. The same services work over LAN/Tailscale with lower -endpoint priority. +This runbook deploys one primary decode node, prefill-compute workers and +optional cache-only peers over a private Thunderbolt Bridge. The same services +work over LAN/Tailscale with lower endpoint priority. ## Production surfaces @@ -17,13 +17,15 @@ require `X-API-Key`. ## Components -| Component | Head | Cache peer | -|---|---|---| -| Kakeya RuntimeService | `127.0.0.1:51051` | optional | -| PrefillCacheService | runtime + `:52051` control node | `169.254.27.104:52051` | -| CapabilityService gossip | enabled | enabled / pull-only if macOS blocks outbound Python sockets | -| Dashboard/API | `127.0.0.1:8090` | no | -| Cloudflare public edge | `kakeya.ai/*` Worker | no | +| Component | Primary | Prefill worker | Cache-only peer | +|---|---|---|---| +| Kakeya RuntimeService / decode | `127.0.0.1:51051` | no | no | +| Same MLX model loaded | yes | yes | no | +| PrefillWorkerService | no | `:53051` | no | +| PrefillCacheService | runtime/local LRU | co-located | `:52051` | +| CapabilityService gossip | enabled | enabled | enabled / pull-only | +| Dashboard/API | `127.0.0.1:8090` | no | no | +| Cloudflare public edge | `kakeya.ai/*` Worker | no | no | ## Compatibility lock @@ -101,6 +103,52 @@ If `nc` works but Python/gRPC outbound calls return `Errno 65`, grant Local Network access to that Python executable in macOS Privacy & Security. Head→peer lookup/publish/fetch remains usable while reverse gossip is disabled. +## Prefill-compute worker + +The worker loads the exact same MLX model as the primary, accepts queued +prefill-only jobs, writes immutable snapshots into its co-located RAM cache and +never serves user decode. + +Create a fleet PSK once and copy it to every trusted node: + +```bash +openssl rand -hex 32 > ~/.kakeya/fleet.psk +chmod 600 ~/.kakeya/fleet.psk +``` + +Install the worker: + +```bash +export KAKEYA_WORKER_REPO="$HOME/Kakeya-LLM-Inference-engine" +export KAKEYA_WORKER_PYTHON="$HOME/kakeya-venv/bin/python" +export KAKEYA_WORKER_MODEL="$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit" +export KAKEYA_CACHE_MODEL_ID="gemma-4-26B-A4B-it-mlx-4bit" +export KAKEYA_MODEL_REVISION="local-4bit-v1" +export KAKEYA_TOKENIZER_REVISION="gemma4-v1" +export KAKEYA_WORKER_NODE_ID="prefill-mini-1" +export KAKEYA_WORKER_ADVERTISE="169.254.27.104:53051" +export KAKEYA_LAYER_GEOMETRY_HASH="" +export KAKEYA_FLEET_PSK_FILE="$HOME/.kakeya/fleet.psk" +export KAKEYA_TENANT_ID="private-fleet" +bash deploy/install_prefill_worker_launchd.sh +``` + +The primary must use the same compatibility and auth values: + +```text +--enable-prefill-cache +--enable-capability-exchange +--peer 169.254.27.104:53051 +--cache-tenant-id private-fleet +--fleet-psk-file ~/.kakeya/fleet.psk +--cache-compression zlib +--cache-replication-factor 1 +``` + +`--cache-peer` remains an emergency static override. Normal worker/cache +selection is derived from compatible live capability cards and their TTL/load +metrics. + ## Node registration and groups Create a registration: @@ -132,6 +180,10 @@ Expected invariants: - imported snapshot checksum and compatibility fingerprint match; - remote failure falls back to local prefill; - no remote RPC occurs in autoregressive decode; +- a cache miss is assigned to a compatible `PREFILL_COMPUTE` worker when its + queue+compute+transfer estimate beats blocking the primary; +- worker failure/timeout/import rejection resets the verifier and performs full + local prefill; - completed and KV-assisted token counters increase after live calls. Minimal acceptance: From 7d436591c662e8c5ad5b2dd905a4c6aa321e0217 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 11 Jul 2026 15:09:19 +0000 Subject: [PATCH 02/11] feat(distributed): orchestrate authenticated prefill-compute workers Co-authored-by: FluffyAIcode --- deploy/install_prefill_worker_launchd.sh | 69 + .../backends/mlx/prefill_snapshot.py | 4 + .../backends/mlx/prefill_worker.py | 95 + inference_engine/distributed/__init__.py | 23 +- inference_engine/distributed/capability.py | 77 +- inference_engine/distributed/prefill_auth.py | 143 ++ inference_engine/distributed/prefill_cache.py | 31 +- .../distributed/prefill_cache_runtime.py | 408 +++- .../distributed/prefill_cache_service.py | 131 +- .../distributed/prefill_compression.py | 87 + .../distributed/prefill_scheduler.py | 190 ++ .../distributed/prefill_worker.py | 475 +++++ inference_engine/server/grpc_app.py | 2 + .../proto_gen/kakeya/v1/distributed_pb2.py | 160 +- .../proto_gen/kakeya/v1/distributed_pb2.pyi | 150 +- .../kakeya/v1/distributed_pb2_grpc.py | 167 ++ proto/kakeya/v1/distributed.proto | 102 + scripts/start_grpc_runtime_server.py | 83 +- scripts/start_prefill_cache_node.py | 34 +- scripts/start_prefill_worker_node.py | 252 +++ .../src/proto_gen/kakeya/v1/distributed.ts | 1746 +++++++++++++++-- 21 files changed, 4141 insertions(+), 288 deletions(-) create mode 100755 deploy/install_prefill_worker_launchd.sh create mode 100644 inference_engine/backends/mlx/prefill_worker.py create mode 100644 inference_engine/distributed/prefill_auth.py create mode 100644 inference_engine/distributed/prefill_compression.py create mode 100644 inference_engine/distributed/prefill_scheduler.py create mode 100644 inference_engine/distributed/prefill_worker.py create mode 100644 scripts/start_prefill_worker_node.py diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh new file mode 100755 index 00000000..f55c0bf0 --- /dev/null +++ b/deploy/install_prefill_worker_launchd.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Install the ADR 0017 prefill-compute worker as a per-user LaunchAgent. +set -euo pipefail + +: "${KAKEYA_WORKER_REPO:?set KAKEYA_WORKER_REPO}" +: "${KAKEYA_WORKER_PYTHON:?set KAKEYA_WORKER_PYTHON}" +: "${KAKEYA_WORKER_MODEL:?set KAKEYA_WORKER_MODEL}" +: "${KAKEYA_WORKER_NODE_ID:?set KAKEYA_WORKER_NODE_ID}" +: "${KAKEYA_WORKER_ADVERTISE:?set KAKEYA_WORKER_ADVERTISE (host:port)}" +: "${KAKEYA_LAYER_GEOMETRY_HASH:?set KAKEYA_LAYER_GEOMETRY_HASH}" + +BIND="${KAKEYA_WORKER_BIND:-0.0.0.0:53051}" +TENANT="${KAKEYA_TENANT_ID:-default}" +CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}" +PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}" +CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}" +MODEL_REVISION="${KAKEYA_MODEL_REVISION:-}" +TOKENIZER_REVISION="${KAKEYA_TOKENIZER_REVISION:-}" +QUANTIZATION="${KAKEYA_CACHE_QUANTIZATION:-4bit-mlx}" +ROPE_HASH="${KAKEYA_ROPE_HASH:-}" +LABEL="ai.kakeya.prefill-worker" +PLIST="$HOME/Library/LaunchAgents/$LABEL.plist" +LOG_DIR="$HOME/.kakeya" +mkdir -p "$(dirname "$PLIST")" "$LOG_DIR" + +psk_xml="" +if [[ -n "$PSK_FILE" ]]; then + psk_xml="--fleet-psk-file$PSK_FILE" +fi + +cat > "$PLIST" < + + + Label$LABEL + ProgramArguments + $KAKEYA_WORKER_PYTHON + $KAKEYA_WORKER_REPO/scripts/start_prefill_worker_node.py + --node-id$KAKEYA_WORKER_NODE_ID + --bind$BIND + --advertise$KAKEYA_WORKER_ADVERTISE + --model-id$KAKEYA_WORKER_MODEL + --cache-model-id$CACHE_MODEL_ID + --model-revision$MODEL_REVISION + --tokenizer-revision$TOKENIZER_REVISION + --quantization$QUANTIZATION + --rope-hash$ROPE_HASH + --layer-geometry-hash$KAKEYA_LAYER_GEOMETRY_HASH + --tenant-id$TENANT + --cache-gb$CACHE_GB + $psk_xml + + WorkingDirectory$KAKEYA_WORKER_REPO + EnvironmentVariables + PYTHONPATH$KAKEYA_WORKER_REPO:$KAKEYA_WORKER_REPO/sdks/python + + RunAtLoad + KeepAlive + ProcessTypeInteractive + StandardOutPath$LOG_DIR/prefill-worker.log + StandardErrorPath$LOG_DIR/prefill-worker.log + +EOF + +launchctl bootout "gui/$(id -u)/$LABEL" 2>/dev/null || true +launchctl bootstrap "gui/$(id -u)" "$PLIST" +echo "installed $LABEL -> $PLIST" + diff --git a/inference_engine/backends/mlx/prefill_snapshot.py b/inference_engine/backends/mlx/prefill_snapshot.py index 969b9afd..1fa5abdf 100644 --- a/inference_engine/backends/mlx/prefill_snapshot.py +++ b/inference_engine/backends/mlx/prefill_snapshot.py @@ -35,6 +35,7 @@ class ImportedPrefillSnapshot: token_count: int cached_token_ids: tuple[int, ...] next_token_logits: Any | None + block_hash: bytes = b"" def export_mlx_prefill_snapshot( @@ -44,6 +45,7 @@ def export_mlx_prefill_snapshot( cached_token_ids: Sequence[int], compatibility: CacheCompatibility, next_token_logits: Any | None = None, + block_hash: bytes = b"", ) -> bytes: """Serialize current MLX cache state at one prefix boundary.""" if token_count <= 0: @@ -68,6 +70,7 @@ def export_mlx_prefill_snapshot( "token_count": int(token_count), "cached_token_ids": [int(token) for token in cached_token_ids], "layer_count": len(cache), + "block_hash": bytes(block_hash).hex(), }, ) @@ -108,6 +111,7 @@ def import_mlx_prefill_snapshot( token_count=token_count, cached_token_ids=tuple(int(t) for t in metadata["cached_token_ids"]), next_token_logits=next_logits, + block_hash=bytes.fromhex(metadata.get("block_hash", "")), ) diff --git a/inference_engine/backends/mlx/prefill_worker.py b/inference_engine/backends/mlx/prefill_worker.py new file mode 100644 index 00000000..6b8f1942 --- /dev/null +++ b/inference_engine/backends/mlx/prefill_worker.py @@ -0,0 +1,95 @@ +"""MLX prefill-only compute engine for ADR 0017 worker nodes.""" +from __future__ import annotations + +import threading +from typing import Sequence + +from inference_engine.backends.mlx.prefill_snapshot import ( + export_mlx_prefill_snapshot, +) +from inference_engine.distributed.capability import ( + CacheCompatibility, + CompressionCodec, +) +from inference_engine.distributed.prefill_cache import CacheBlock +from inference_engine.distributed.prefill_compression import compress_payload + + +class MLXPrefillComputeEngine: + """Serially runs prefill with a loaded MLX verifier and exports one snapshot.""" + + def __init__(self, verifier, compatibility: CacheCompatibility) -> None: + self.verifier = verifier + self.compatibility = compatibility + self._lock = threading.Lock() + + def compute_prefill( + self, + token_ids: Sequence[int], + block_hashes: Sequence[bytes], + *, + compression: CompressionCodec, + cancelled: threading.Event, + ) -> Sequence[CacheBlock]: + tokens = [int(token) for token in token_ids] + if not tokens or not block_hashes: + raise ValueError("token_ids and block_hashes must be non-empty") + size = self.compatibility.block_size_tokens + expected_blocks = (len(tokens) + size - 1) // size + if len(block_hashes) != expected_blocks: + raise ValueError( + f"expected {expected_blocks} block hashes, got {len(block_hashes)}", + ) + with self._lock: + if cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + first_end = min(size, len(tokens)) + self.verifier.prefill(tokens[:first_end]) + blocks: list[CacheBlock] = [ + self._snapshot( + token_count=first_end, + block_hash=block_hashes[0], + compression=compression, + ), + ] + for block_index, start in enumerate( + range(first_end, len(tokens), size), + start=1, + ): + if cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + block = tokens[start:start + size] + logits = self.verifier.forward_block(block) + self.verifier.commit_or_truncate( + forwarded=len(block), + accepted=len(block), + ) + self.verifier.next_token_logits = logits[-1].clone() + blocks.append(self._snapshot( + token_count=min(start + size, len(tokens)), + block_hash=block_hashes[block_index], + compression=compression, + )) + return blocks + + def _snapshot( + self, + *, + token_count: int, + block_hash: bytes, + compression: CompressionCodec, + ) -> CacheBlock: + raw = export_mlx_prefill_snapshot( + self.verifier.cache, + token_count=token_count, + cached_token_ids=self.verifier.cached_token_sequence, + compatibility=self.compatibility, + next_token_logits=self.verifier.next_token_logits, + block_hash=block_hash, + ) + return CacheBlock.create( + bytes(block_hash), + token_count, + compress_payload(raw, compression), + ) + diff --git a/inference_engine/distributed/__init__.py b/inference_engine/distributed/__init__.py index 110bcbe5..d0832c0b 100644 --- a/inference_engine/distributed/__init__.py +++ b/inference_engine/distributed/__init__.py @@ -1,20 +1,19 @@ -"""Multi-host plane for Kakeya (ADR 0009, v0.5-M1). +"""Multi-host plane for Kakeya (ADR 0009 / 0016 / 0017). Subpackage layout: - :mod:`capability` — capability cards + the converging gossip registry (``NodeCapability`` / ``ModelCapability`` / ``CapabilityRegistry``). -- :mod:`placement` — deterministic spec-decode placement over a - fleet snapshot. +- :mod:`prefill_worker` — queued prefill-only compute workers. +- :mod:`prefill_cache` / :mod:`prefill_cache_service` — immutable RAM cache. +- :mod:`prefill_cache_runtime` — primary-side orchestration and fallback. +- :mod:`prefill_scheduler` — load/cost-aware worker and replica placement. +- :mod:`prefill_auth` — fleet-PSK authentication and tenant hash isolation. - :mod:`exchange` — ``CapabilityService`` gRPC servicer + the ``exchange_once`` gossip client. -- :mod:`ngram` — model-free prompt-lookup proposer (the always- - available proposer capability every node can advertise). -- :mod:`proposer_service` — ``ProposerService`` gRPC servicer + - ``RemoteProposer`` client (drop-in ``DLMProposer`` substitute). -- :mod:`spec_decode` — pure greedy accept rule + - ``DistributedSpeculativeDecoder``. +- :mod:`ngram`, :mod:`proposer_service`, :mod:`spec_decode` — legacy research + proposer paths retained for reproducibility; not the product architecture. - :mod:`mlx_ring` — optional ``mlx.distributed`` ring probe (bulk-tensor data plane advertisement, ADR 0009 §4 item 4). @@ -26,8 +25,11 @@ from inference_engine.distributed.capability import ( CapabilityRegistry, CapabilityRole, + CacheCompatibility, + CompressionCodec, ModelCapability, NodeCapability, + PrefillWorkerCapability, ) from inference_engine.distributed.placement import ( PlacementError, @@ -38,8 +40,11 @@ __all__ = [ "CapabilityRegistry", "CapabilityRole", + "CacheCompatibility", + "CompressionCodec", "ModelCapability", "NodeCapability", + "PrefillWorkerCapability", "PlacementError", "SpecDecodePlacement", "plan_spec_decode_placement", diff --git a/inference_engine/distributed/capability.py b/inference_engine/distributed/capability.py index 627396ea..b4a4ef83 100644 --- a/inference_engine/distributed/capability.py +++ b/inference_engine/distributed/capability.py @@ -49,6 +49,13 @@ class CapabilityRole(enum.IntEnum): EMBEDDER = 3 TOOL = 4 PREFILL_CACHE = 5 + PREFILL_COMPUTE = 6 + + +class CompressionCodec(enum.IntEnum): + UNSPECIFIED = 0 + NONE = 1 + ZLIB = 2 @dataclass(frozen=True) @@ -112,12 +119,15 @@ class CacheCompatibility: model_id: str model_revision: str = "" tokenizer_revision: str = "" - cache_format_version: str = "kv-v1" + cache_format_version: str = "kakeya-prefill-v2-zlib" quantization: str = "" rope_hash: str = "" layer_geometry_hash: str = "" kv_dtype: str = "" block_size_tokens: int = 64 + tenant_namespace: str = "" + sink_size: int = 4 + window_size: int = 64 def to_proto(self) -> distributed_pb2.CacheCompatibility: return distributed_pb2.CacheCompatibility( @@ -130,6 +140,9 @@ def to_proto(self) -> distributed_pb2.CacheCompatibility: layer_geometry_hash=self.layer_geometry_hash, kv_dtype=self.kv_dtype, block_size_tokens=self.block_size_tokens, + tenant_namespace=self.tenant_namespace, + sink_size=self.sink_size, + window_size=self.window_size, ) @classmethod @@ -146,6 +159,9 @@ def from_proto( layer_geometry_hash=msg.layer_geometry_hash, kv_dtype=msg.kv_dtype, block_size_tokens=msg.block_size_tokens, + tenant_namespace=msg.tenant_namespace, + sink_size=msg.sink_size, + window_size=msg.window_size, ) @@ -162,6 +178,8 @@ class CacheCapability: load: float = 0.0 tokens_served: int = 0 bloom_filter: bytes = b"" + default_compression: CompressionCodec = CompressionCodec.NONE + replication_factor: int = 1 def to_proto(self) -> distributed_pb2.CacheCapability: return distributed_pb2.CacheCapability( @@ -174,6 +192,8 @@ def to_proto(self) -> distributed_pb2.CacheCapability: load=self.load, tokens_served=self.tokens_served, bloom_filter=self.bloom_filter, + default_compression=int(self.default_compression), + replication_factor=self.replication_factor, ) @classmethod @@ -188,6 +208,55 @@ def from_proto(cls, msg: distributed_pb2.CacheCapability) -> "CacheCapability": load=msg.load, tokens_served=msg.tokens_served, bloom_filter=msg.bloom_filter, + default_compression=CompressionCodec(msg.default_compression), + replication_factor=msg.replication_factor, + ) + + +@dataclass(frozen=True) +class PrefillWorkerCapability: + """One prefill-only compute offering on a node.""" + + compatibility: CacheCompatibility + worker_address: str = "" + max_concurrent_jobs: int = 1 + inflight_jobs: int = 0 + queued_jobs: int = 0 + load: float = 0.0 + tokens_per_second_prefill: float = 0.0 + ram_bytes_free: int = 0 + accepts_compute_jobs: bool = True + queued_tokens: int = 0 + + def to_proto(self) -> distributed_pb2.PrefillWorkerCapability: + return distributed_pb2.PrefillWorkerCapability( + compatibility=self.compatibility.to_proto(), + worker_address=self.worker_address, + max_concurrent_jobs=self.max_concurrent_jobs, + inflight_jobs=self.inflight_jobs, + queued_jobs=self.queued_jobs, + load=self.load, + tokens_per_second_prefill=self.tokens_per_second_prefill, + ram_bytes_free=self.ram_bytes_free, + accepts_compute_jobs=self.accepts_compute_jobs, + queued_tokens=self.queued_tokens, + ) + + @classmethod + def from_proto( + cls, msg: distributed_pb2.PrefillWorkerCapability, + ) -> "PrefillWorkerCapability": + return cls( + compatibility=CacheCompatibility.from_proto(msg.compatibility), + worker_address=msg.worker_address, + max_concurrent_jobs=msg.max_concurrent_jobs, + inflight_jobs=msg.inflight_jobs, + queued_jobs=msg.queued_jobs, + load=msg.load, + tokens_per_second_prefill=msg.tokens_per_second_prefill, + ram_bytes_free=msg.ram_bytes_free, + accepts_compute_jobs=msg.accepts_compute_jobs, + queued_tokens=msg.queued_tokens, ) @@ -206,6 +275,7 @@ class NodeCapability: ring_address: str = "" caches: Tuple[CacheCapability, ...] = () endpoints: Tuple[NodeEndpoint, ...] = () + prefill_workers: Tuple[PrefillWorkerCapability, ...] = () def __post_init__(self) -> None: if not self.node_id: @@ -239,6 +309,7 @@ def to_proto(self) -> distributed_pb2.NodeCapability: ring_address=self.ring_address, caches=[c.to_proto() for c in self.caches], endpoints=[e.to_proto() for e in self.endpoints], + prefill_workers=[w.to_proto() for w in self.prefill_workers], ) @classmethod @@ -255,6 +326,10 @@ def from_proto(cls, msg: distributed_pb2.NodeCapability) -> "NodeCapability": ring_address=msg.ring_address, caches=tuple(CacheCapability.from_proto(c) for c in msg.caches), endpoints=tuple(NodeEndpoint.from_proto(e) for e in msg.endpoints), + prefill_workers=tuple( + PrefillWorkerCapability.from_proto(w) + for w in msg.prefill_workers + ), ) diff --git a/inference_engine/distributed/prefill_auth.py b/inference_engine/distributed/prefill_auth.py new file mode 100644 index 00000000..db30e720 --- /dev/null +++ b/inference_engine/distributed/prefill_auth.py @@ -0,0 +1,143 @@ +"""Fleet-PSK authentication and tenant isolation for prefill RPCs. + +The transport can still be a private-network insecure gRPC channel, but every +request is authenticated before allocation/compute. The signature covers the +deterministic protobuf bytes plus caller/tenant/timestamp metadata. Prefix +hashes use a tenant-derived HMAC key so one tenant cannot probe another +tenant's prompt prefixes. +""" +from __future__ import annotations + +import hashlib +import hmac +import time +from dataclasses import dataclass +from typing import Iterable, Sequence, Tuple + +AUTH_TENANT = "x-kakeya-tenant-id" +AUTH_NODE = "x-kakeya-node-id" +AUTH_TS = "x-kakeya-auth-ts" +AUTH_MAC = "x-kakeya-auth-mac" + + +class PrefillAuthError(ValueError): + """Authentication or replay-window failure.""" + + +@dataclass(frozen=True) +class FleetAuthConfig: + psk: bytes + tenant_id: str + node_id: str + max_clock_skew_s: float = 60.0 + + def __post_init__(self) -> None: + if len(self.psk) < 16: + raise ValueError("fleet PSK must be at least 16 bytes") + if not self.tenant_id: + raise ValueError("tenant_id must be non-empty") + if not self.node_id: + raise ValueError("node_id must be non-empty") + if self.max_clock_skew_s <= 0: + raise ValueError("max_clock_skew_s must be > 0") + + @classmethod + def from_file( + cls, path: str, *, tenant_id: str, node_id: str, + max_clock_skew_s: float = 60.0, + ) -> "FleetAuthConfig": + with open(path, "rb") as fh: + secret = fh.read().strip() + return cls(secret, tenant_id, node_id, max_clock_skew_s) + + def tenant_hash_key(self) -> bytes: + return hmac.new( + self.psk, + b"kakeya-prefill-tenant\0" + self.tenant_id.encode(), + hashlib.sha256, + ).digest() + + +def _request_bytes(request) -> bytes: + serialize = getattr(request, "SerializeToString", None) + if serialize is None: + raise TypeError("authenticated request must be a protobuf message") + return serialize(deterministic=True) + + +def _mac( + request, *, psk: bytes, tenant_id: str, node_id: str, timestamp: str, +) -> str: + body_hash = hashlib.sha256(_request_bytes(request)).digest() + payload = b"\0".join(( + tenant_id.encode(), + node_id.encode(), + timestamp.encode(), + body_hash, + )) + return hmac.new(psk, payload, hashlib.sha256).hexdigest() + + +def signed_metadata( + request, config: FleetAuthConfig, *, now: float | None = None, +) -> Tuple[Tuple[str, str], ...]: + timestamp = str(int(time.time() if now is None else now)) + return ( + (AUTH_TENANT, config.tenant_id), + (AUTH_NODE, config.node_id), + (AUTH_TS, timestamp), + (AUTH_MAC, _mac( + request, + psk=config.psk, + tenant_id=config.tenant_id, + node_id=config.node_id, + timestamp=timestamp, + )), + ) + + +def verify_metadata( + metadata: Iterable[Tuple[str, str]], + request, + config: FleetAuthConfig, + *, + now: float | None = None, +) -> Tuple[str, str]: + values = {key.lower(): value for key, value in metadata} + tenant = values.get(AUTH_TENANT, "") + node = values.get(AUTH_NODE, "") + timestamp = values.get(AUTH_TS, "") + supplied = values.get(AUTH_MAC, "") + if tenant != config.tenant_id: + raise PrefillAuthError("tenant mismatch") + if not node or not timestamp or not supplied: + raise PrefillAuthError("missing prefill authentication metadata") + try: + ts = int(timestamp) + except ValueError as exc: + raise PrefillAuthError("invalid authentication timestamp") from exc + current = time.time() if now is None else now + if abs(current - ts) > config.max_clock_skew_s: + raise PrefillAuthError("authentication timestamp outside replay window") + expected = _mac( + request, + psk=config.psk, + tenant_id=tenant, + node_id=node, + timestamp=timestamp, + ) + if not hmac.compare_digest(supplied, expected): + raise PrefillAuthError("invalid prefill authentication MAC") + return tenant, node + + +def metadata_pairs(metadata: Sequence) -> Tuple[Tuple[str, str], ...]: + """Normalize grpc metadata objects or plain pairs for verification.""" + result = [] + for item in metadata: + if hasattr(item, "key") and hasattr(item, "value"): + result.append((item.key, item.value)) + else: + result.append((item[0], item[1])) + return tuple(result) + diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index d931d924..2386f7b2 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/inference_engine/distributed/prefill_cache.py @@ -13,6 +13,7 @@ from __future__ import annotations import hashlib +import hmac import json import secrets import threading @@ -37,7 +38,10 @@ def compatibility_fingerprint(compatibility: CacheCompatibility) -> bytes: "model_revision": compatibility.model_revision, "quantization": compatibility.quantization, "rope_hash": compatibility.rope_hash, + "sink_size": compatibility.sink_size, + "tenant_namespace": compatibility.tenant_namespace, "tokenizer_revision": compatibility.tokenizer_revision, + "window_size": compatibility.window_size, } return hashlib.sha256( json.dumps(payload, sort_keys=True, separators=(",", ":")).encode(), @@ -47,6 +51,8 @@ def compatibility_fingerprint(compatibility: CacheCompatibility) -> bytes: def chained_block_hashes( token_ids: Sequence[int], compatibility: CacheCompatibility, + *, + hmac_key: bytes = b"", ) -> list[bytes]: """Hash fixed-size token blocks, chaining each hash to its predecessor. @@ -62,7 +68,12 @@ def chained_block_hashes( for start in range(0, len(token_ids), size): block = token_ids[start:start + size] encoded = b"".join(int(t).to_bytes(4, "little", signed=False) for t in block) - previous = hashlib.sha256(namespace + previous + encoded).digest() + material = namespace + previous + encoded + previous = ( + hmac.new(hmac_key, material, hashlib.sha256).digest() + if hmac_key + else hashlib.sha256(material).digest() + ) hashes.append(previous) return hashes @@ -149,6 +160,7 @@ def put(self, block: CacheBlock) -> bool: if block.nbytes > self.max_bytes: 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: @@ -159,6 +171,10 @@ def put(self, block: CacheBlock) -> bool: self._bytes_used += block.nbytes self._epoch += 1 self._evict_to_budget() + if block.block_hash not in self._blocks: + raise ValueError( + "cache capacity is pinned by active leases", + ) return True def put_prefix( @@ -249,6 +265,19 @@ def block_hashes(self) -> tuple[bytes, ...]: with self._lock: return tuple(self._blocks) + def invalidate(self, block_hash: bytes) -> bool: + """Drop a corrupt/rejected snapshot so local recompute can replace it.""" + with self._lock: + block = self._blocks.pop(bytes(block_hash), None) + if block is None: + return False + self._bytes_used -= block.nbytes + self._epoch += 1 + for lease_id, lease in list(self._leases.items()): + if bytes(block_hash) in lease.block_hashes: + del self._leases[lease_id] + return True + def _expire_leases(self, now: float) -> None: for lease_id, lease in list(self._leases.items()): if now > lease.expires_at_unix: diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py index 0b2c5484..e4b6b439 100644 --- a/inference_engine/distributed/prefill_cache_runtime.py +++ b/inference_engine/distributed/prefill_cache_runtime.py @@ -9,9 +9,11 @@ from __future__ import annotations import hashlib +import time +import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass -from typing import Any, Sequence +from typing import Any, Callable, Sequence import grpc @@ -19,12 +21,39 @@ export_mlx_prefill_snapshot, import_mlx_prefill_snapshot, ) -from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.capability import ( + CacheCompatibility, + CompressionCodec, + NodeCapability, +) +from inference_engine.distributed.prefill_auth import ( + FleetAuthConfig, + signed_metadata, +) from inference_engine.distributed.prefill_cache import ( CacheBlock, PrefixCacheStore, chained_block_hashes, ) +from inference_engine.distributed.prefill_compression import ( + compress_payload, + decompress_payload, +) +from inference_engine.distributed.prefill_scheduler import ( + PrefillCostConfig, + choose_prefill_worker, + remote_import_wins, + select_cache_replicas, +) +from inference_engine.distributed.prefill_worker import ( + PrefillJobState, + cancel_prefill_job_sync, + get_prefill_job_sync, + submit_prefill_job_sync, +) +from inference_engine.distributed.prefill_cache_service import ( + compatible_cache_peers, +) from inference_engine.server.proto_gen.kakeya.v1 import ( distributed_pb2, distributed_pb2_grpc, @@ -39,6 +68,10 @@ class PrefillReuseStats: tokens_reused: int = 0 tokens_computed: int = 0 bytes_received: int = 0 + remote_jobs: int = 0 + remote_job_failures: int = 0 + fallbacks: int = 0 + last_fallback_reason: str = "" @dataclass(frozen=True) @@ -49,6 +82,9 @@ class _Hit: hit_tokens: int transfer_bytes: int payload: bytes | None = None + rtt_ms: float = 0.0 + block_hash: bytes = b"" + payload_sha256: bytes = b"" class DistributedPrefillCacheHook: @@ -59,19 +95,53 @@ def __init__( local_store: PrefixCacheStore, *, peers: Sequence[str] = (), + registry_provider: Callable[[], Sequence[NodeCapability]] | None = None, lookup_timeout_s: float = 2.0, fetch_timeout_s: float = 30.0, + worker_timeout_s: float = 120.0, + worker_poll_interval_s: float = 0.05, + remote_compute_min_tokens: int = 128, + max_import_bytes: int = 1 << 30, + estimated_snapshot_bytes_per_token: int = 400_000, + compression: CompressionCodec = CompressionCodec.ZLIB, + replication_factor: int = 1, + cost_config: PrefillCostConfig | None = None, + auth: FleetAuthConfig | None = None, on_reuse=None, ) -> None: + if min( + lookup_timeout_s, + fetch_timeout_s, + worker_timeout_s, + worker_poll_interval_s, + max_import_bytes, + estimated_snapshot_bytes_per_token, + ) <= 0: + raise ValueError("prefill runtime limits must be > 0") + if remote_compute_min_tokens < 0 or replication_factor < 0: + raise ValueError("prefill thresholds must be >= 0") self.local_store = local_store self.compatibility = local_store.compatibility self.peers = tuple(dict.fromkeys(peer for peer in peers if peer)) + self.registry_provider = registry_provider self.lookup_timeout_s = float(lookup_timeout_s) self.fetch_timeout_s = float(fetch_timeout_s) + self.worker_timeout_s = float(worker_timeout_s) + self.worker_poll_interval_s = float(worker_poll_interval_s) + self.remote_compute_min_tokens = int(remote_compute_min_tokens) + self.max_import_bytes = int(max_import_bytes) + self.estimated_snapshot_bytes_per_token = int( + estimated_snapshot_bytes_per_token, + ) + self.compression = CompressionCodec(compression) + self.replication_factor = int(replication_factor) + self.cost_config = cost_config or PrefillCostConfig() + self.auth = auth + self._hash_key = auth.tenant_hash_key() if auth is not None else b"" self.stats = PrefillReuseStats() self._on_reuse = on_reuse self._publisher = ThreadPoolExecutor( - max_workers=max(1, min(4, len(self.peers))), + max_workers=4, thread_name_prefix="prefill-kv-publish", ) @@ -83,18 +153,82 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int: tokens = [int(token) for token in token_ids] if not tokens: return 0 - hashes = chained_block_hashes(tokens, self.compatibility) + hashes = chained_block_hashes( + tokens, + self.compatibility, + hmac_key=self._hash_key, + ) hit = self._best_hit(hashes) reused = 0 if hit is not None: - payload = hit.payload if hit.payload is not None else self._fetch_remote(hit) + reused = self._try_import(verifier, tokens, hit) + elif len(tokens) >= self.remote_compute_min_tokens: + remote_hit = self._compute_remote(tokens, hashes) + if remote_hit is not None: + reused = self._try_import(verifier, tokens, remote_hit) + if reused == 0: + self.stats.misses += 1 + + self._compute_and_publish(verifier, tokens, hashes, reused) + return reused + + def _try_import( + self, + verifier: Any, + tokens: list[int], + hit: _Hit, + ) -> int: + try: + expected_hit_tokens = min( + hit.hit_blocks * self.compatibility.block_size_tokens, + len(tokens), + ) + if ( + hit.hit_blocks <= 0 + or hit.hit_tokens != expected_hit_tokens + ): + raise ValueError("prefill hit does not match a token boundary") + if ( + hit.payload is None + and hit.transfer_bytes > self.max_import_bytes + ): + raise ValueError("prefill snapshot exceeds wire import budget") + payload = ( + hit.payload + if hit.payload is not None + else self._fetch_remote(hit) + ) + if len(payload) > self.max_import_bytes: + raise ValueError("prefill snapshot exceeds wire import budget") + raw = decompress_payload( + payload, + max_uncompressed_bytes=self.max_import_bytes, + ) verifier.reset() imported = import_mlx_prefill_snapshot( - payload, + raw, verifier.cache, compatibility=self.compatibility, ) + if ( + not (0 < imported.token_count <= len(tokens)) + or imported.token_count != hit.hit_tokens + ): + raise ValueError("prefill snapshot token_count is invalid") + if hit.block_hash and imported.block_hash != hit.block_hash: + raise ValueError("prefill snapshot block hash mismatch") + if imported.next_token_logits is None: + raise ValueError("prefill snapshot is missing continuation logits") reused = min(imported.token_count, len(tokens)) + expected_prefix = tokens[:reused] + sink_window = getattr(verifier, "_sink_window_slice", None) + expected_cached = ( + list(sink_window(expected_prefix)) + if callable(sink_window) + else expected_prefix + ) + if list(imported.cached_token_ids) != expected_cached: + raise ValueError("prefill snapshot cached token sequence mismatch") verifier.cached_token_sequence = list(imported.cached_token_ids) verifier.next_global_position = reused if imported.next_token_logits is not None: @@ -106,11 +240,106 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int: self.stats.local_hits += 1 else: self.stats.remote_hits += 1 - else: - self.stats.misses += 1 + return reused + except Exception as exc: + # Cache is an optimization. A corrupt/expired/unreachable hit must + # never determine request correctness. + self.stats.fallbacks += 1 + self.stats.last_fallback_reason = f"{type(exc).__name__}: {exc}" + if hit.source == "local" and hit.block_hash: + self.local_store.invalidate(hit.block_hash) + verifier.reset() + return 0 - self._compute_and_publish(verifier, tokens, hashes, reused) - return reused + def _compute_remote( + self, + tokens: list[int], + hashes: list[bytes], + ) -> _Hit | None: + cards = tuple( + card for card in self._cards() + if card.node_id != self.local_store.node_id + ) + target = choose_prefill_worker( + cards, + self.compatibility, + prompt_tokens=len(tokens), + estimated_snapshot_bytes=( + len(tokens) * self.estimated_snapshot_bytes_per_token + ), + config=self.cost_config, + ) + if target is None: + return None + request = distributed_pb2.SubmitPrefillJobRequest( + request_id=uuid.uuid4().hex, + tenant_id=self.compatibility.tenant_namespace or "default", + compatibility=self.compatibility.to_proto(), + token_ids=tokens, + block_hashes=hashes, + deadline_ms=int(self.worker_timeout_s * 1000), + preferred_compression=int(self.compression), + ) + try: + response = None + response = submit_prefill_job_sync( + target.address, + request, + timeout_s=self.lookup_timeout_s, + auth=self.auth, + ) + self.stats.remote_jobs += 1 + deadline = time.monotonic() + self.worker_timeout_s + while time.monotonic() < deadline: + status_request = distributed_pb2.GetPrefillJobStatusRequest( + job_id=response.job_id, + tenant_id=request.tenant_id, + ) + status = get_prefill_job_sync( + target.address, + status_request, + timeout_s=self.lookup_timeout_s, + auth=self.auth, + ) + if status.status == int(PrefillJobState.COMPLETED): + return _Hit( + source=status.cache_address or target.address, + lease_id=status.lease_id, + hit_blocks=len(hashes), + hit_tokens=status.tokens_computed, + transfer_bytes=status.transfer_bytes, + rtt_ms=target.rtt_ms, + block_hash=hashes[-1], + payload_sha256=bytes(status.payload_sha256), + ) + if status.status in ( + int(PrefillJobState.FAILED), + int(PrefillJobState.CANCELLED), + ): + raise RuntimeError( + status.failure_reason or "remote prefill job failed", + ) + time.sleep(self.worker_poll_interval_s) + raise TimeoutError("remote prefill job timed out") + except Exception as exc: + if response is not None: + try: + cancel_request = distributed_pb2.CancelPrefillJobRequest( + job_id=response.job_id, + tenant_id=request.tenant_id, + ) + cancel_prefill_job_sync( + target.address, + cancel_request, + timeout_s=self.lookup_timeout_s, + auth=self.auth, + ) + except Exception: + pass + self.stats.remote_job_failures += 1 + self.stats.fallbacks += 1 + self.stats.last_fallback_reason = f"{type(exc).__name__}: {exc}" + return None def _compute_and_publish( self, @@ -120,21 +349,17 @@ def _compute_and_publish( reused: int, ) -> None: size = self.compatibility.block_size_tokens - start_block = reused // size if reused == 0: first_end = min(size, len(tokens)) verifier.prefill(tokens[:first_end]) self.stats.tokens_computed += first_end self._publish_boundary(verifier, tokens, hashes, 0, first_end) - start_block = 1 - for block_index in range(start_block, len(hashes)): - start = block_index * size - if start < reused: - continue - end = min(start + size, len(tokens)) - block_tokens = tokens[start:end] - if not block_tokens: - continue + reused = first_end + cursor = reused + while cursor < len(tokens): + block_index = cursor // size + end = min((block_index + 1) * size, len(tokens)) + block_tokens = tokens[cursor:end] logits = verifier.forward_block(block_tokens) verifier.commit_or_truncate( forwarded=len(block_tokens), @@ -143,6 +368,7 @@ def _compute_and_publish( verifier.next_token_logits = logits[-1].clone() self.stats.tokens_computed += len(block_tokens) self._publish_boundary(verifier, tokens, hashes, block_index, end) + cursor = end def _publish_boundary( self, @@ -152,31 +378,69 @@ def _publish_boundary( block_index: int, prefix_end: int, ) -> None: - payload = export_mlx_prefill_snapshot( + raw_payload = export_mlx_prefill_snapshot( verifier.cache, token_count=prefix_end, cached_token_ids=verifier.cached_token_sequence, compatibility=self.compatibility, next_token_logits=verifier.next_token_logits, + block_hash=hashes[block_index], ) + payload = compress_payload(raw_payload, self.compression) block = CacheBlock.create(hashes[block_index], prefix_end, payload) self.local_store.put(block) - if self.peers: + peers = self._publish_peers(block.block_hash) + if peers: from inference_engine.distributed.prefill_cache_service import ( publish_block_sync, ) - for peer in self.peers: + for peer in peers: self._publisher.submit( publish_block_sync, peer, self.compatibility, block, timeout_s=self.fetch_timeout_s, + auth=self.auth, ) def close(self) -> None: self._publisher.shutdown(wait=False, cancel_futures=True) + def _cards(self) -> tuple[NodeCapability, ...]: + if self.registry_provider is None: + return () + try: + return tuple(self.registry_provider()) + except Exception: + return () + + def _cache_peers(self) -> tuple[str, ...]: + dynamic = compatible_cache_peers( + tuple( + card for card in self._cards() + if card.node_id != self.local_store.node_id + ), + self.compatibility, + ) + return tuple(dict.fromkeys((*self.peers, *dynamic))) + + def _publish_peers(self, block_hash: bytes) -> tuple[str, ...]: + cards = tuple( + card for card in self._cards() + if card.node_id != self.local_store.node_id + ) + if cards: + selected = select_cache_replicas( + cards, + self.compatibility, + block_hash=block_hash, + replication_factor=self.replication_factor, + ) + if selected: + return tuple(selected) + return self._cache_peers()[:self.replication_factor] + def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None: candidates: list[_Hit] = [] local = self.local_store.lookup(hashes) @@ -189,12 +453,15 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None: hit_tokens=local.hit_token_count, transfer_bytes=local.transfer_bytes, payload=blocks[-1].payload, + block_hash=blocks[-1].block_hash, + payload_sha256=blocks[-1].payload_sha256, )) - if self.peers: - with ThreadPoolExecutor(max_workers=min(8, len(self.peers))) as pool: + peers = self._cache_peers() + if peers: + with ThreadPoolExecutor(max_workers=min(8, len(peers))) as pool: futures = { pool.submit(self._lookup_peer, peer, hashes): peer - for peer in self.peers + for peer in peers } for future in as_completed(futures): hit = future.result() @@ -202,7 +469,7 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None: candidates.append(hit) if not candidates: return None - return max( + best = max( candidates, key=lambda hit: ( hit.hit_tokens, @@ -210,48 +477,119 @@ def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None: -hit.transfer_bytes, ), ) + if ( + best.source != "local" + and not remote_import_wins( + hit_tokens=best.hit_tokens, + transfer_bytes=best.transfer_bytes, + rtt_ms=best.rtt_ms, + config=self.cost_config, + ) + ): + local = next( + (candidate for candidate in candidates if candidate.source == "local"), + None, + ) + return local + return best def _lookup_peer(self, peer: str, hashes: Sequence[bytes]) -> _Hit | None: try: with grpc.insecure_channel(peer) as channel: stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + request = distributed_pb2.LookupPrefixRequest( + compatibility=self.compatibility.to_proto(), + block_hashes=hashes, + ) response = stub.LookupPrefix( - distributed_pb2.LookupPrefixRequest( - compatibility=self.compatibility.to_proto(), - block_hashes=hashes, - ), + request, timeout=self.lookup_timeout_s, + metadata=( + signed_metadata(request, self.auth) + if self.auth is not None else None + ), ) except grpc.RpcError: return None if not response.lease_id or response.hit_block_count == 0: return None + if response.hit_block_count > len(hashes): + return None return _Hit( source=peer, lease_id=response.lease_id, hit_blocks=response.hit_block_count, hit_tokens=response.hit_token_count, transfer_bytes=response.transfer_bytes, + rtt_ms=self._peer_rtt(peer), + block_hash=bytes(hashes[response.hit_block_count - 1]), + payload_sha256=bytes(response.payload_sha256), ) + def _peer_rtt(self, address: str) -> float: + for card in self._cards(): + for endpoint in card.endpoints: + if endpoint.address == address: + return endpoint.measured_rtt_ms + return 0.0 + def _fetch_remote(self, hit: _Hit) -> bytes: + if hit.transfer_bytes > self.max_import_bytes: + raise RuntimeError("remote prefill payload exceeds import budget") parts: dict[int, bytes] = {} expected_chunks = 0 expected_sha = b"" + expected_block_hash = b"" + received = 0 try: with grpc.insecure_channel(hit.source) as channel: stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + request = distributed_pb2.FetchBlocksRequest( + lease_id=hit.lease_id, + ) for chunk in stub.FetchBlocks( - distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id), + request, timeout=self.fetch_timeout_s, + metadata=( + signed_metadata(request, self.auth) + if self.auth is not None else None + ), ): - parts[chunk.chunk_index] = bytes(chunk.data) - expected_chunks = chunk.total_chunks - expected_sha = bytes(chunk.block_sha256) + if expected_chunks == 0: + expected_chunks = chunk.total_chunks + if expected_chunks > 65_536: + raise RuntimeError( + "remote prefill cache chunk count exceeds limit", + ) + expected_sha = bytes(chunk.block_sha256) + expected_block_hash = bytes(chunk.block_hash) + elif ( + chunk.total_chunks != expected_chunks + or bytes(chunk.block_sha256) != expected_sha + or bytes(chunk.block_hash) != expected_block_hash + ): + raise RuntimeError( + "remote prefill cache chunk metadata changed", + ) + if ( + chunk.chunk_index < 0 + or chunk.chunk_index >= expected_chunks + or chunk.chunk_index in parts + ): + raise RuntimeError("invalid or duplicate prefill chunk") + data = bytes(chunk.data) + received += len(data) + if received > self.max_import_bytes: + raise RuntimeError( + "remote prefill payload exceeds import budget", + ) + parts[chunk.chunk_index] = data except grpc.RpcError as exc: raise RuntimeError(f"remote prefill cache fetch failed: {exc}") from exc if expected_chunks <= 0 or len(parts) != expected_chunks: raise RuntimeError("remote prefill cache stream was incomplete") + if hit.payload_sha256 and expected_sha != hit.payload_sha256: + raise RuntimeError("remote prefill cache lease checksum changed") payload = b"".join(parts[index] for index in range(expected_chunks)) if hashlib.sha256(payload).digest() != expected_sha: raise RuntimeError("remote prefill cache checksum mismatch") diff --git a/inference_engine/distributed/prefill_cache_service.py b/inference_engine/distributed/prefill_cache_service.py index ec70056c..5168d180 100644 --- a/inference_engine/distributed/prefill_cache_service.py +++ b/inference_engine/distributed/prefill_cache_service.py @@ -14,6 +14,12 @@ CacheCompatibility, NodeCapability, ) +from inference_engine.distributed.prefill_auth import ( + FleetAuthConfig, + PrefillAuthError, + signed_metadata, + verify_metadata, +) from inference_engine.distributed.prefill_cache import ( CacheBlock, PrefixCacheStore, @@ -32,7 +38,10 @@ def cache_capability( *, cache_address: str, load: float = 0.0, + default_compression=None, + replication_factor: int = 1, ) -> CacheCapability: + from inference_engine.distributed.capability import CompressionCodec stats = store.stats() return CacheCapability( compatibility=store.compatibility, @@ -43,6 +52,12 @@ def cache_capability( cache_epoch=stats.cache_epoch, load=load, tokens_served=stats.tokens_served, + default_compression=( + CompressionCodec.NONE + if default_compression is None + else CompressionCodec(default_compression) + ), + replication_factor=replication_factor, ) @@ -55,19 +70,48 @@ def __init__( *, cache_address: str, chunk_bytes: int = DEFAULT_CHUNK_BYTES, + auth: FleetAuthConfig | None = None, + max_payload_bytes: int | None = None, ) -> None: if chunk_bytes <= 0: raise ValueError("chunk_bytes must be > 0") self.store = store self.cache_address = cache_address self.chunk_bytes = int(chunk_bytes) + self.auth = auth + self.max_payload_bytes = ( + int(max_payload_bytes) + if max_payload_bytes is not None else store.max_bytes + ) + if self.max_payload_bytes <= 0: + raise ValueError("max_payload_bytes must be > 0") + + async def _authenticate(self, request, context) -> None: + if self.auth is None: + return + try: + verify_metadata(context.invocation_metadata(), request, self.auth) + except PrefillAuthError as exc: + await context.abort(grpc.StatusCode.UNAUTHENTICATED, str(exc)) + + async def _authorize_compatibility(self, compatibility, context) -> None: + if ( + self.auth is not None + and compatibility.tenant_namespace != self.auth.tenant_id + ): + await context.abort( + grpc.StatusCode.PERMISSION_DENIED, + "cache tenant namespace mismatch", + ) async def GetCacheSummary( # noqa: N802 self, request: distributed_pb2.GetCacheSummaryRequest, context: grpc.aio.ServicerContext, ) -> distributed_pb2.GetCacheSummaryResponse: + await self._authenticate(request, context) requested = CacheCompatibility.from_proto(request.compatibility) + await self._authorize_compatibility(requested, context) caches = [] if requested == self.store.compatibility: caches.append( @@ -86,7 +130,9 @@ async def LookupPrefix( # noqa: N802 request: distributed_pb2.LookupPrefixRequest, context: grpc.aio.ServicerContext, ) -> distributed_pb2.LookupPrefixResponse: + await self._authenticate(request, context) requested = CacheCompatibility.from_proto(request.compatibility) + await self._authorize_compatibility(requested, context) if requested != self.store.compatibility: return distributed_pb2.LookupPrefixResponse( node_id=self.store.node_id, @@ -100,6 +146,7 @@ async def FetchBlocks( # noqa: N802 request: distributed_pb2.FetchBlocksRequest, context: grpc.aio.ServicerContext, ): + await self._authenticate(request, context) try: blocks = self.store.fetch(request.lease_id) except KeyError as exc: @@ -130,9 +177,23 @@ async def PublishBlock( # noqa: N802 ) -> distributed_pb2.PublishBlockResponse: parts: dict[int, bytes] = {} first = None + received = 0 async for chunk in request_iterator: if first is None: first = chunk + await self._authenticate(first, context) + requested = CacheCompatibility.from_proto(first.compatibility) + await self._authorize_compatibility(requested, context) + if ( + first.total_chunks <= 0 + or first.total_chunks + > (self.max_payload_bytes + self.chunk_bytes - 1) + // self.chunk_bytes + ): + await context.abort( + grpc.StatusCode.RESOURCE_EXHAUSTED, + "publish stream exceeds payload budget", + ) elif ( chunk.block_hash != first.block_hash or chunk.total_chunks != first.total_chunks @@ -142,7 +203,23 @@ async def PublishBlock( # noqa: N802 grpc.StatusCode.INVALID_ARGUMENT, "inconsistent publish chunk metadata", ) - parts[chunk.chunk_index] = bytes(chunk.data) + if ( + chunk.chunk_index >= chunk.total_chunks + or chunk.chunk_index in parts + or len(chunk.data) > self.chunk_bytes + ): + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "invalid or duplicate publish chunk", + ) + data = bytes(chunk.data) + received += len(data) + if received > self.max_payload_bytes: + await context.abort( + grpc.StatusCode.RESOURCE_EXHAUSTED, + "publish payload exceeds cache allocation budget", + ) + parts[chunk.chunk_index] = data if first is None: await context.abort( grpc.StatusCode.INVALID_ARGUMENT, @@ -185,11 +262,15 @@ def add_prefill_cache_service( *, cache_address: str, chunk_bytes: int = DEFAULT_CHUNK_BYTES, + auth: FleetAuthConfig | None = None, + max_payload_bytes: int | None = None, ) -> PrefillCacheServiceServicer: servicer = PrefillCacheServiceServicer( store, cache_address=cache_address, chunk_bytes=chunk_bytes, + auth=auth, + max_payload_bytes=max_payload_bytes, ) distributed_pb2_grpc.add_PrefillCacheServiceServicer_to_server( servicer, server, @@ -220,16 +301,19 @@ async def lookup_peer( block_hashes: Sequence[bytes], *, timeout_s: float = 3.0, + auth: FleetAuthConfig | None = None, ) -> RemotePrefixHit: try: async with grpc.aio.insecure_channel(address) as channel: stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + request = distributed_pb2.LookupPrefixRequest( + compatibility=compatibility.to_proto(), + block_hashes=block_hashes, + ) response = await stub.LookupPrefix( - distributed_pb2.LookupPrefixRequest( - compatibility=compatibility.to_proto(), - block_hashes=block_hashes, - ), + request, timeout=timeout_s, + metadata=signed_metadata(request, auth) if auth else None, ) except grpc.aio.AioRpcError: return RemotePrefixHit(address, "", "", 0, 0, 0, 0, 0.0, b"") @@ -252,6 +336,7 @@ async def lookup_best_peer( block_hashes: Sequence[bytes], *, timeout_s: float = 3.0, + auth: FleetAuthConfig | None = None, ) -> RemotePrefixHit | None: """Fan out concurrently and choose longest hit, then smallest transfer.""" if not peers: @@ -262,6 +347,7 @@ async def lookup_best_peer( compatibility, block_hashes, timeout_s=timeout_s, + auth=auth, ) for peer in peers )) @@ -282,12 +368,15 @@ async def fetch_remote_blocks( hit: RemotePrefixHit, *, timeout_s: float = 30.0, + auth: FleetAuthConfig | None = None, ) -> list[distributed_pb2.FetchBlocksResponse]: async with grpc.aio.insecure_channel(hit.address) as channel: stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + request = distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id) stream = stub.FetchBlocks( - distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id), + request, timeout=timeout_s, + metadata=signed_metadata(request, auth) if auth else None, ) return [chunk async for chunk in stream] @@ -299,6 +388,7 @@ def publish_block_sync( *, timeout_s: float = 30.0, chunk_bytes: int = DEFAULT_CHUNK_BYTES, + auth: FleetAuthConfig | None = None, ) -> bool: """Publish one immutable snapshot to a peer (used by background workers).""" try: @@ -306,24 +396,31 @@ def publish_block_sync( 1, (block.nbytes + chunk_bytes - 1) // chunk_bytes, ) + chunk_messages = [] + for chunk_index in range(total_chunks): + start = chunk_index * chunk_bytes + chunk_messages.append(distributed_pb2.PublishBlockRequest( + block_hash=block.block_hash, + token_count=block.token_count, + chunk_index=chunk_index, + total_chunks=total_chunks, + data=block.payload[start:start + chunk_bytes], + block_sha256=block.payload_sha256, + compatibility=compatibility.to_proto(), + )) + def chunks(): - for chunk_index in range(total_chunks): - start = chunk_index * chunk_bytes - yield distributed_pb2.PublishBlockRequest( - block_hash=block.block_hash, - token_count=block.token_count, - chunk_index=chunk_index, - total_chunks=total_chunks, - data=block.payload[start:start + chunk_bytes], - block_sha256=block.payload_sha256, - compatibility=compatibility.to_proto(), - ) + yield from chunk_messages with grpc.insecure_channel(address) as channel: stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) response = stub.PublishBlock( chunks(), timeout=timeout_s, + metadata=( + signed_metadata(chunk_messages[0], auth) + if auth is not None else None + ), ) except grpc.RpcError: return False diff --git a/inference_engine/distributed/prefill_compression.py b/inference_engine/distributed/prefill_compression.py new file mode 100644 index 00000000..23f7ef2c --- /dev/null +++ b/inference_engine/distributed/prefill_compression.py @@ -0,0 +1,87 @@ +"""Snapshot payload framing and compression for distributed prefill K/V.""" +from __future__ import annotations + +import hashlib +import struct +import zlib + +from inference_engine.distributed.capability import CompressionCodec + +_MAGIC = b"KPC1" +_HEADER = struct.Struct("<4sBQ32s") + + +def compress_payload( + payload: bytes, + codec: CompressionCodec, + *, + level: int = 3, +) -> bytes: + raw = bytes(payload) + if codec in (CompressionCodec.UNSPECIFIED, CompressionCodec.NONE): + return raw + if codec != CompressionCodec.ZLIB: + raise ValueError(f"unsupported compression codec {codec!r}") + if not (0 <= level <= 9): + raise ValueError("zlib level must be in [0, 9]") + compressed = zlib.compress(raw, level) + return _HEADER.pack( + _MAGIC, + int(codec), + len(raw), + hashlib.sha256(raw).digest(), + ) + compressed + + +def decompress_payload( + payload: bytes, + *, + max_uncompressed_bytes: int, +) -> bytes: + data = bytes(payload) + if not data.startswith(_MAGIC): + if len(data) > max_uncompressed_bytes: + raise ValueError("uncompressed prefill payload exceeds import budget") + return data + if len(data) < _HEADER.size: + raise ValueError("truncated compressed prefill payload header") + _magic, raw_codec, expected_size, expected_sha = _HEADER.unpack( + data[:_HEADER.size], + ) + if expected_size > max_uncompressed_bytes: + raise ValueError("prefill payload exceeds uncompressed import budget") + try: + codec = CompressionCodec(raw_codec) + except ValueError as exc: + raise ValueError(f"unsupported compression codec {raw_codec}") from exc + if codec != CompressionCodec.ZLIB: + raise ValueError(f"unsupported framed compression codec {codec.name}") + decompressor = zlib.decompressobj() + raw = decompressor.decompress( + data[_HEADER.size:], + max_uncompressed_bytes + 1, + ) + if decompressor.unconsumed_tail or len(raw) > max_uncompressed_bytes: + raise ValueError("decompressed payload exceeds import budget") + remaining = max_uncompressed_bytes - len(raw) + raw += decompressor.flush(remaining + 1) + if len(raw) > max_uncompressed_bytes: # pragma: no cover - zlib flush guard + raise ValueError("decompressed payload exceeds import budget") + if decompressor.unused_data: + raise ValueError("compressed prefill payload has trailing data") + if len(raw) != expected_size: + raise ValueError( + f"decompressed payload size {len(raw)} != expected {expected_size}", + ) + if hashlib.sha256(raw).digest() != expected_sha: + raise ValueError("decompressed prefill payload checksum mismatch") + return raw + + +def payload_sizes(payload: bytes) -> tuple[int, int]: + """Return ``(wire_bytes, uncompressed_bytes)`` without decompressing.""" + data = bytes(payload) + if data.startswith(_MAGIC) and len(data) >= _HEADER.size: + return len(data), int(_HEADER.unpack(data[:_HEADER.size])[2]) + return len(data), len(data) + diff --git a/inference_engine/distributed/prefill_scheduler.py b/inference_engine/distributed/prefill_scheduler.py new file mode 100644 index 00000000..c8cd8475 --- /dev/null +++ b/inference_engine/distributed/prefill_scheduler.py @@ -0,0 +1,190 @@ +"""Pure placement/cost functions for distributed prefill compute and storage.""" +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from enum import Enum +from typing import Sequence + +from inference_engine.distributed.capability import ( + CacheCompatibility, + NodeCapability, + PrefillWorkerCapability, +) + + +class PrefillAction(str, Enum): + IMPORT = "import" + REMOTE_COMPUTE = "remote_compute" + LOCAL_COMPUTE = "local_compute" + + +@dataclass(frozen=True) +class PrefillCostConfig: + local_prefill_tps: float = 20.0 + default_worker_tps: float = 20.0 + link_mbps: float = 1000.0 + default_rtt_ms: float = 2.0 + minimum_savings_ratio: float = 0.10 + primary_compute_penalty_ms: float = 0.0 + + def __post_init__(self) -> None: + if min( + self.local_prefill_tps, + self.default_worker_tps, + self.link_mbps, + self.default_rtt_ms, + ) <= 0: + raise ValueError("prefill cost metrics must be > 0") + if not (0 <= self.minimum_savings_ratio < 1): + raise ValueError("minimum_savings_ratio must be in [0, 1)") + if self.primary_compute_penalty_ms < 0: + raise ValueError("primary_compute_penalty_ms must be >= 0") + + +@dataclass(frozen=True) +class WorkerTarget: + node_id: str + address: str + capability: PrefillWorkerCapability + rtt_ms: float + + @property + def queue_eta_ms(self) -> float: + tps = self.capability.tokens_per_second_prefill + if tps <= 0: + return 0.0 + return self.capability.queued_tokens * 1000.0 / tps + + +def estimate_local_prefill_ms(tokens: int, config: PrefillCostConfig) -> float: + return ( + max(0, tokens) / config.local_prefill_tps * 1000.0 + + config.primary_compute_penalty_ms + ) + + +def estimate_import_ms( + transfer_bytes: int, + *, + rtt_ms: float, + config: PrefillCostConfig, +) -> float: + bytes_per_ms = config.link_mbps * 1_000_000.0 / 8.0 / 1000.0 + return max(rtt_ms, 0.0) + max(transfer_bytes, 0) / bytes_per_ms + + +def remote_import_wins( + *, + hit_tokens: int, + transfer_bytes: int, + rtt_ms: float, + config: PrefillCostConfig, +) -> bool: + local = estimate_local_prefill_ms(hit_tokens, config) + remote = estimate_import_ms( + transfer_bytes, + rtt_ms=rtt_ms or config.default_rtt_ms, + config=config, + ) + return remote <= local * (1.0 - config.minimum_savings_ratio) + + +def compatible_prefill_workers( + cards: Sequence[NodeCapability], + compatibility: CacheCompatibility, +) -> list[WorkerTarget]: + targets: list[WorkerTarget] = [] + for card in cards: + rtt = min( + ( + endpoint.measured_rtt_ms + for endpoint in card.endpoints + if endpoint.measured_rtt_ms > 0 + ), + default=0.0, + ) + for worker in card.prefill_workers: + if ( + worker.accepts_compute_jobs + and worker.compatibility == compatibility + and (worker.worker_address or card.grpc_address) + ): + targets.append(WorkerTarget( + card.node_id, + worker.worker_address or card.grpc_address, + worker, + rtt, + )) + return targets + + +def choose_prefill_worker( + cards: Sequence[NodeCapability], + compatibility: CacheCompatibility, + *, + prompt_tokens: int, + estimated_snapshot_bytes: int, + config: PrefillCostConfig, +) -> WorkerTarget | None: + candidates = compatible_prefill_workers(cards, compatibility) + if not candidates: + return None + local_ms = estimate_local_prefill_ms(prompt_tokens, config) + + def cost(target: WorkerTarget) -> float: + tps = ( + target.capability.tokens_per_second_prefill + or config.default_worker_tps + ) + compute_ms = prompt_tokens / tps * 1000.0 + import_ms = estimate_import_ms( + estimated_snapshot_bytes, + rtt_ms=target.rtt_ms or config.default_rtt_ms, + config=config, + ) + load_penalty = max(0.0, target.capability.load) * compute_ms + return target.queue_eta_ms + compute_ms + import_ms + load_penalty + + best = min(candidates, key=lambda target: ( + cost(target), + -target.capability.ram_bytes_free, + target.node_id, + )) + if cost(best) > local_ms * (1.0 - config.minimum_savings_ratio): + return None + return best + + +def select_cache_replicas( + cards: Sequence[NodeCapability], + compatibility: CacheCompatibility, + *, + block_hash: bytes, + replication_factor: int, +) -> list[str]: + """Deterministic rendezvous placement avoids publishing to every peer.""" + if replication_factor <= 0: + return [] + candidates: list[tuple[int, int, str]] = [] + for card in cards: + for cache in card.caches: + if cache.compatibility != compatibility or not cache.cache_address: + continue + score = int.from_bytes(hashlib.sha256( + bytes(block_hash) + card.node_id.encode(), + ).digest(), "big") + candidates.append(( + score, + cache.cache_bytes_free, + cache.cache_address, + )) + candidates.sort(reverse=True) + result: list[str] = [] + for _score, _free, address in candidates: + if address not in result: + result.append(address) + if len(result) >= replication_factor: + break + return result + diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py new file mode 100644 index 00000000..e4b41e36 --- /dev/null +++ b/inference_engine/distributed/prefill_worker.py @@ -0,0 +1,475 @@ +"""Queued prefill-only worker jobs and gRPC service (ADR 0017).""" +from __future__ import annotations + +import asyncio +import hashlib +import threading +import time +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import IntEnum +from typing import Protocol, Sequence + +import grpc + +from inference_engine.distributed.capability import ( + CacheCompatibility, + CompressionCodec, +) +from inference_engine.distributed.prefill_auth import ( + FleetAuthConfig, + PrefillAuthError, + signed_metadata, + verify_metadata, +) +from inference_engine.distributed.prefill_cache import ( + CacheBlock, + PrefixCacheStore, + compatibility_fingerprint, +) +from inference_engine.server.proto_gen.kakeya.v1 import ( + distributed_pb2, + distributed_pb2_grpc, +) + + +class PrefillJobState(IntEnum): + UNSPECIFIED = 0 + QUEUED = 1 + RUNNING = 2 + COMPLETED = 3 + FAILED = 4 + CANCELLED = 5 + + +class PrefillComputeEngine(Protocol): + """Model-specific worker that produces the final restorable snapshot.""" + + def compute_prefill( + self, + token_ids: Sequence[int], + block_hashes: Sequence[bytes], + *, + compression: CompressionCodec, + cancelled: threading.Event, + ) -> Sequence[CacheBlock]: ... + + +@dataclass +class PrefillJob: + job_id: str + request_id: str + tenant_id: str + token_ids: tuple[int, ...] + block_hashes: tuple[bytes, ...] + compression: CompressionCodec + state: PrefillJobState = PrefillJobState.QUEUED + tokens_computed: int = 0 + lease_id: str = "" + block_hash: bytes = b"" + payload_sha256: bytes = b"" + transfer_bytes: int = 0 + failure_reason: str = "" + compute_ms: float = 0.0 + created_at: float = field(default_factory=time.time) + finished_at: float = 0.0 + cancelled: threading.Event = field(default_factory=threading.Event) + future: Future | None = field(default=None, repr=False) + request_digest: bytes = b"" + deadline_at: float = 0.0 + + +class PrefillJobStore: + """Bounded, idempotent job queue around one or more prefill engines.""" + + def __init__( + self, + engine: PrefillComputeEngine, + cache_store: PrefixCacheStore, + *, + max_concurrent_jobs: int = 1, + max_jobs: int = 128, + completed_ttl_s: float = 600.0, + max_prompt_tokens: int = 131_072, + ) -> None: + if min( + max_concurrent_jobs, + max_jobs, + completed_ttl_s, + max_prompt_tokens, + ) <= 0: + raise ValueError("worker limits must be > 0") + self.engine = engine + self.cache_store = cache_store + self.max_concurrent_jobs = int(max_concurrent_jobs) + self.max_jobs = int(max_jobs) + self.completed_ttl_s = float(completed_ttl_s) + self.max_prompt_tokens = int(max_prompt_tokens) + self._jobs: dict[str, PrefillJob] = {} + self._requests: dict[tuple[str, str], str] = {} + self._lock = threading.RLock() + self._executor = ThreadPoolExecutor( + max_workers=self.max_concurrent_jobs, + thread_name_prefix="kakeya-prefill-worker", + ) + + def submit( + self, + *, + request_id: str, + tenant_id: str, + token_ids: Sequence[int], + block_hashes: Sequence[bytes], + compatibility: CacheCompatibility, + compression: CompressionCodec, + deadline_ms: int = 0, + ) -> PrefillJob: + if not request_id or not tenant_id: + raise ValueError("request_id and tenant_id must be non-empty") + if deadline_ms < 0: + raise ValueError("deadline_ms must be >= 0") + if compatibility != self.cache_store.compatibility: + raise ValueError("prefill worker compatibility mismatch") + if not token_ids or not block_hashes: + raise ValueError("token_ids and block_hashes must be non-empty") + if len(token_ids) > self.max_prompt_tokens: + raise ValueError("prefill prompt exceeds worker token limit") + expected_blocks = ( + len(token_ids) + compatibility.block_size_tokens - 1 + ) // compatibility.block_size_tokens + if len(block_hashes) != expected_blocks: + raise ValueError( + f"expected {expected_blocks} block hashes, got {len(block_hashes)}", + ) + if any(len(bytes(block_hash)) != 32 for block_hash in block_hashes): + raise ValueError("every prefill block hash must be SHA-256 (32 bytes)") + digest = hashlib.sha256( + compatibility_fingerprint(compatibility) + + b"".join(int(token).to_bytes(4, "little") for token in token_ids) + + b"".join(bytes(h) for h in block_hashes) + + int(compression).to_bytes(2, "little") + + int(deadline_ms).to_bytes(8, "little", signed=False) + ).digest() + with self._lock: + self._gc_locked() + request_key = (tenant_id, request_id) + existing_id = self._requests.get(request_key) + if existing_id is not None: + existing = self._jobs[existing_id] + if existing.request_digest != digest: + raise ValueError( + "idempotency key reused with different prefill request", + ) + return existing + active = sum( + job.state in (PrefillJobState.QUEUED, PrefillJobState.RUNNING) + for job in self._jobs.values() + ) + if active >= self.max_jobs: + raise ValueError("prefill worker job queue is full") + job = PrefillJob( + job_id=uuid.uuid4().hex, + request_id=request_id, + tenant_id=tenant_id, + token_ids=tuple(int(token) for token in token_ids), + block_hashes=tuple(bytes(h) for h in block_hashes), + compression=compression, + request_digest=digest, + deadline_at=( + time.time() + deadline_ms / 1000.0 + if deadline_ms > 0 else 0.0 + ), + ) + self._jobs[job.job_id] = job + self._requests[request_key] = job.job_id + job.future = self._executor.submit(self._run, job.job_id) + return job + + def get(self, job_id: str, tenant_id: str) -> PrefillJob: + with self._lock: + self._gc_locked() + job = self._jobs.get(job_id) + if job is None or job.tenant_id != tenant_id: + raise KeyError(job_id) + return job + + def cancel(self, job_id: str, tenant_id: str) -> bool: + with self._lock: + job = self.get(job_id, tenant_id) + if job.state in ( + PrefillJobState.COMPLETED, + PrefillJobState.FAILED, + PrefillJobState.CANCELLED, + ): + return False + job.cancelled.set() + if job.future is not None and job.future.cancel(): + job.state = PrefillJobState.CANCELLED + job.finished_at = time.time() + return True + + def stats(self) -> tuple[int, int, float, int]: + with self._lock: + running = sum(j.state == PrefillJobState.RUNNING for j in self._jobs.values()) + queued = sum(j.state == PrefillJobState.QUEUED for j in self._jobs.values()) + queued_tokens = sum( + len(j.token_ids) + for j in self._jobs.values() + if j.state in (PrefillJobState.QUEUED, PrefillJobState.RUNNING) + ) + load = min(1.0, (running + queued) / self.max_concurrent_jobs) + return running, queued, load, queued_tokens + + def close(self) -> None: + self._executor.shutdown(wait=False, cancel_futures=True) + + def _run(self, job_id: str) -> None: + with self._lock: + job = self._jobs[job_id] + if job.cancelled.is_set(): + job.state = PrefillJobState.CANCELLED + job.finished_at = time.time() + return + job.state = PrefillJobState.RUNNING + started = time.perf_counter() + timer = None + if job.deadline_at: + remaining = job.deadline_at - time.time() + if remaining <= 0: + job.cancelled.set() + else: + timer = threading.Timer(remaining, job.cancelled.set) + timer.daemon = True + timer.start() + try: + blocks = tuple(self.engine.compute_prefill( + job.token_ids, + job.block_hashes, + compression=job.compression, + cancelled=job.cancelled, + )) + if job.cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + if len(blocks) != len(job.block_hashes): + raise RuntimeError( + "prefill engine must return one snapshot per block hash", + ) + for block in blocks: + if job.cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + self.cache_store.put(block) + lease = self.cache_store.lookup(job.block_hashes) + if not lease.lease_id: + raise RuntimeError("computed snapshot was not discoverable") + if job.cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + with self._lock: + if job.cancelled.is_set(): + raise InterruptedError("prefill job cancelled") + job.state = PrefillJobState.COMPLETED + job.tokens_computed = len(job.token_ids) + job.lease_id = lease.lease_id + final = blocks[-1] + job.block_hash = final.block_hash + job.payload_sha256 = final.payload_sha256 + job.transfer_bytes = final.nbytes + except InterruptedError as exc: + with self._lock: + job.state = PrefillJobState.CANCELLED + job.failure_reason = str(exc) + except Exception as exc: + with self._lock: + job.state = PrefillJobState.FAILED + job.failure_reason = f"{type(exc).__name__}: {exc}" + finally: + if timer is not None: + timer.cancel() + with self._lock: + job.compute_ms = (time.perf_counter() - started) * 1000.0 + job.finished_at = time.time() + + def _gc_locked(self) -> None: + cutoff = time.time() - self.completed_ttl_s + for job_id, job in list(self._jobs.items()): + if job.finished_at and job.finished_at < cutoff: + self._jobs.pop(job_id, None) + self._requests.pop((job.tenant_id, job.request_id), None) + overflow = len(self._jobs) - self.max_jobs * 2 + if overflow > 0: + finished = sorted( + ( + job for job in self._jobs.values() + if job.finished_at + ), + key=lambda job: job.finished_at, + ) + for job in finished[:overflow]: + self._jobs.pop(job.job_id, None) + self._requests.pop((job.tenant_id, job.request_id), None) + + +class PrefillWorkerServiceServicer( + distributed_pb2_grpc.PrefillWorkerServiceServicer, +): + def __init__( + self, + jobs: PrefillJobStore, + *, + node_id: str, + cache_address: str, + auth: FleetAuthConfig | None = None, + tokens_per_second_prefill: float = 0.0, + ) -> None: + self.jobs = jobs + self.node_id = node_id + self.cache_address = cache_address + self.auth = auth + self.tokens_per_second_prefill = float(tokens_per_second_prefill) + + async def _authenticate(self, request, context) -> None: + if self.auth is None: + return + try: + verify_metadata( + context.invocation_metadata(), + request, + self.auth, + ) + except PrefillAuthError as exc: + await context.abort(grpc.StatusCode.UNAUTHENTICATED, str(exc)) + + async def SubmitPrefillJob(self, request, context): # noqa: N802 + await self._authenticate(request, context) + if self.auth is not None and request.tenant_id != self.auth.tenant_id: + await context.abort( + grpc.StatusCode.PERMISSION_DENIED, + "prefill job tenant mismatch", + ) + try: + job = await asyncio.to_thread( + self.jobs.submit, + request_id=request.request_id, + tenant_id=request.tenant_id, + token_ids=list(request.token_ids), + block_hashes=list(request.block_hashes), + compatibility=CacheCompatibility.from_proto(request.compatibility), + compression=CompressionCodec(request.preferred_compression), + deadline_ms=request.deadline_ms, + ) + except ValueError as exc: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(exc)) + return distributed_pb2.SubmitPrefillJobResponse( + job_id=job.job_id, + status=int(job.state), + worker_node_id=self.node_id, + queue_eta_ms=( + self.jobs.stats()[3] / self.tokens_per_second_prefill * 1000.0 + if self.tokens_per_second_prefill > 0 else 0.0 + ), + ) + + async def GetPrefillJobStatus(self, request, context): # noqa: N802 + await self._authenticate(request, context) + if self.auth is not None and request.tenant_id != self.auth.tenant_id: + await context.abort(grpc.StatusCode.PERMISSION_DENIED, "tenant mismatch") + try: + job = self.jobs.get(request.job_id, request.tenant_id) + except KeyError as exc: + await context.abort(grpc.StatusCode.NOT_FOUND, str(exc)) + return _job_status_proto(job, self.cache_address) + + async def CancelPrefillJob(self, request, context): # noqa: N802 + await self._authenticate(request, context) + if self.auth is not None and request.tenant_id != self.auth.tenant_id: + await context.abort(grpc.StatusCode.PERMISSION_DENIED, "tenant mismatch") + try: + cancelled = self.jobs.cancel(request.job_id, request.tenant_id) + except KeyError as exc: + await context.abort(grpc.StatusCode.NOT_FOUND, str(exc)) + return distributed_pb2.CancelPrefillJobResponse(cancelled=cancelled) + + +def add_prefill_worker_service( + server: grpc.aio.Server, + jobs: PrefillJobStore, + *, + node_id: str, + cache_address: str, + auth: FleetAuthConfig | None = None, + tokens_per_second_prefill: float = 0.0, +) -> PrefillWorkerServiceServicer: + servicer = PrefillWorkerServiceServicer( + jobs, + node_id=node_id, + cache_address=cache_address, + auth=auth, + tokens_per_second_prefill=tokens_per_second_prefill, + ) + distributed_pb2_grpc.add_PrefillWorkerServiceServicer_to_server( + servicer, + server, + ) + return servicer + + +def submit_prefill_job_sync( + address: str, + request: distributed_pb2.SubmitPrefillJobRequest, + *, + timeout_s: float, + auth: FleetAuthConfig | None = None, +) -> distributed_pb2.SubmitPrefillJobResponse: + metadata = signed_metadata(request, auth) if auth is not None else None + with grpc.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel) + return stub.SubmitPrefillJob(request, timeout=timeout_s, metadata=metadata) + + +def get_prefill_job_sync( + address: str, + request: distributed_pb2.GetPrefillJobStatusRequest, + *, + timeout_s: float, + auth: FleetAuthConfig | None = None, +) -> distributed_pb2.GetPrefillJobStatusResponse: + metadata = signed_metadata(request, auth) if auth is not None else None + with grpc.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel) + return stub.GetPrefillJobStatus(request, timeout=timeout_s, metadata=metadata) + + +def cancel_prefill_job_sync( + address: str, + request: distributed_pb2.CancelPrefillJobRequest, + *, + timeout_s: float, + auth: FleetAuthConfig | None = None, +) -> distributed_pb2.CancelPrefillJobResponse: + metadata = signed_metadata(request, auth) if auth is not None else None + with grpc.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillWorkerServiceStub(channel) + return stub.CancelPrefillJob( + request, + timeout=timeout_s, + metadata=metadata, + ) + + +def _job_status_proto( + job: PrefillJob, + cache_address: str, +) -> distributed_pb2.GetPrefillJobStatusResponse: + return distributed_pb2.GetPrefillJobStatusResponse( + job_id=job.job_id, + status=int(job.state), + tokens_computed=job.tokens_computed, + lease_id=job.lease_id, + block_hash=job.block_hash, + payload_sha256=job.payload_sha256, + transfer_bytes=job.transfer_bytes, + failure_reason=job.failure_reason, + compute_ms=job.compute_ms, + cache_address=cache_address, + ) + diff --git a/inference_engine/server/grpc_app.py b/inference_engine/server/grpc_app.py index 02a093ad..3a7f45af 100644 --- a/inference_engine/server/grpc_app.py +++ b/inference_engine/server/grpc_app.py @@ -382,6 +382,7 @@ def create_grpc_server( default_proposer_model_id: str = "", prefill_cache_store: Optional[object] = None, prefill_cache_address: str = "", + prefill_auth: Optional[object] = None, ) -> grpc.aio.Server: """Build, but do not start, a configured gRPC asyncio server. @@ -455,6 +456,7 @@ def create_grpc_server( server, prefill_cache_store, cache_address=prefill_cache_address or config.bind_address, + auth=prefill_auth, ) _logger.info( "gRPC PrefillCacheService enabled at %s", 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 699e21ec..7fef240f 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py @@ -24,83 +24,103 @@ -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\"\xc6\x02\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\"[\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\"\xeb\x01\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\"\xf7\x01\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\"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\"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*\xc8\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\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\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\"\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*m\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*\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=3547 - _globals['_CAPABILITYROLE']._serialized_end=3747 + _globals['_CAPABILITYROLE']._serialized_start=4889 + _globals['_CAPABILITYROLE']._serialized_end=5126 + _globals['_COMPRESSIONCODEC']._serialized_start=5128 + _globals['_COMPRESSIONCODEC']._serialized_end=5237 + _globals['_PREFILLJOBSTATUS']._serialized_start=5240 + _globals['_PREFILLJOBSTATUS']._serialized_end=5456 _globals['_MODELCAPABILITY']._serialized_start=42 _globals['_MODELCAPABILITY']._serialized_end=167 _globals['_NODECAPABILITY']._serialized_start=170 - _globals['_NODECAPABILITY']._serialized_end=496 - _globals['_NODEENDPOINT']._serialized_start=498 - _globals['_NODEENDPOINT']._serialized_end=589 - _globals['_CACHECOMPATIBILITY']._serialized_start=592 - _globals['_CACHECOMPATIBILITY']._serialized_end=827 - _globals['_CACHECAPABILITY']._serialized_start=830 - _globals['_CACHECAPABILITY']._serialized_end=1077 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=1079 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=1156 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=1158 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=1236 - _globals['_GETNODECAPABILITYREQUEST']._serialized_start=1238 - _globals['_GETNODECAPABILITYREQUEST']._serialized_end=1264 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=1266 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=1334 - _globals['_GETCACHESUMMARYREQUEST']._serialized_start=1336 - _globals['_GETCACHESUMMARYREQUEST']._serialized_end=1414 - _globals['_GETCACHESUMMARYRESPONSE']._serialized_start=1416 - _globals['_GETCACHESUMMARYRESPONSE']._serialized_end=1502 - _globals['_LOOKUPPREFIXREQUEST']._serialized_start=1504 - _globals['_LOOKUPPREFIXREQUEST']._serialized_end=1601 - _globals['_LOOKUPPREFIXRESPONSE']._serialized_start=1604 - _globals['_LOOKUPPREFIXRESPONSE']._serialized_end=1811 - _globals['_FETCHBLOCKSREQUEST']._serialized_start=1813 - _globals['_FETCHBLOCKSREQUEST']._serialized_end=1851 - _globals['_FETCHBLOCKSRESPONSE']._serialized_start=1854 - _globals['_FETCHBLOCKSRESPONSE']._serialized_end=2037 - _globals['_PUBLISHBLOCKREQUEST']._serialized_start=2040 - _globals['_PUBLISHBLOCKREQUEST']._serialized_end=2277 - _globals['_PUBLISHBLOCKRESPONSE']._serialized_start=2279 - _globals['_PUBLISHBLOCKRESPONSE']._serialized_end=2338 - _globals['_PROPOSEBLOCKREQUEST']._serialized_start=2340 - _globals['_PROPOSEBLOCKREQUEST']._serialized_end=2447 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=2449 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=2570 - _globals['_TENSOR']._serialized_start=2572 - _globals['_TENSOR']._serialized_end=2624 - _globals['_LAYERKV']._serialized_start=2626 - _globals['_LAYERKV']._serialized_end=2710 - _globals['_RESTOREREQUEST']._serialized_start=2713 - _globals['_RESTOREREQUEST']._serialized_end=2845 - _globals['_RESTORERESPONSE']._serialized_start=2847 - _globals['_RESTORERESPONSE']._serialized_end=2949 - _globals['_SEEDCONTEXTREQUEST']._serialized_start=2951 - _globals['_SEEDCONTEXTREQUEST']._serialized_end=3042 - _globals['_SEEDCONTEXTRESPONSE']._serialized_start=3044 - _globals['_SEEDCONTEXTRESPONSE']._serialized_end=3086 - _globals['_DRAFTBLOCKREQUEST']._serialized_start=3088 - _globals['_DRAFTBLOCKREQUEST']._serialized_end=3192 - _globals['_DRAFTBLOCKRESPONSE']._serialized_start=3194 - _globals['_DRAFTBLOCKRESPONSE']._serialized_end=3294 - _globals['_EXTENDCONTEXTREQUEST']._serialized_start=3296 - _globals['_EXTENDCONTEXTREQUEST']._serialized_end=3389 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=3391 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=3435 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=3437 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=3499 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=3501 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=3544 - _globals['_CAPABILITYSERVICE']._serialized_start=3750 - _globals['_CAPABILITYSERVICE']._serialized_end=3970 - _globals['_PROPOSERSERVICE']._serialized_start=3972 - _globals['_PROPOSERSERVICE']._serialized_end=4070 - _globals['_PREFILLCACHESERVICE']._serialized_start=4073 - _globals['_PREFILLCACHESERVICE']._serialized_end=4428 - _globals['_DFLASHPROPOSERSERVICE']._serialized_start=4431 - _globals['_DFLASHPROPOSERSERVICE']._serialized_end=4880 + _globals['_NODECAPABILITY']._serialized_end=557 + _globals['_NODEENDPOINT']._serialized_start=559 + _globals['_NODEENDPOINT']._serialized_end=650 + _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=5459 + _globals['_CAPABILITYSERVICE']._serialized_end=5679 + _globals['_PROPOSERSERVICE']._serialized_start=5681 + _globals['_PROPOSERSERVICE']._serialized_end=5779 + _globals['_PREFILLCACHESERVICE']._serialized_start=5782 + _globals['_PREFILLCACHESERVICE']._serialized_end=6137 + _globals['_PREFILLWORKERSERVICE']._serialized_start=6140 + _globals['_PREFILLWORKERSERVICE']._serialized_end=6450 + _globals['_DFLASHPROPOSERSERVICE']._serialized_start=6453 + _globals['_DFLASHPROPOSERSERVICE']._serialized_end=6902 # @@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 8141a68c..02da80a0 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi @@ -15,12 +15,38 @@ class CapabilityRole(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): CAPABILITY_ROLE_EMBEDDER: _ClassVar[CapabilityRole] CAPABILITY_ROLE_TOOL: _ClassVar[CapabilityRole] CAPABILITY_ROLE_PREFILL_CACHE: _ClassVar[CapabilityRole] + CAPABILITY_ROLE_PREFILL_COMPUTE: _ClassVar[CapabilityRole] + +class CompressionCodec(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + COMPRESSION_CODEC_UNSPECIFIED: _ClassVar[CompressionCodec] + COMPRESSION_CODEC_NONE: _ClassVar[CompressionCodec] + COMPRESSION_CODEC_ZLIB: _ClassVar[CompressionCodec] + +class PrefillJobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + PREFILL_JOB_STATUS_UNSPECIFIED: _ClassVar[PrefillJobStatus] + PREFILL_JOB_STATUS_QUEUED: _ClassVar[PrefillJobStatus] + PREFILL_JOB_STATUS_RUNNING: _ClassVar[PrefillJobStatus] + PREFILL_JOB_STATUS_COMPLETED: _ClassVar[PrefillJobStatus] + PREFILL_JOB_STATUS_FAILED: _ClassVar[PrefillJobStatus] + PREFILL_JOB_STATUS_CANCELLED: _ClassVar[PrefillJobStatus] CAPABILITY_ROLE_UNSPECIFIED: CapabilityRole CAPABILITY_ROLE_VERIFIER: CapabilityRole CAPABILITY_ROLE_PROPOSER: CapabilityRole CAPABILITY_ROLE_EMBEDDER: CapabilityRole CAPABILITY_ROLE_TOOL: CapabilityRole CAPABILITY_ROLE_PREFILL_CACHE: CapabilityRole +CAPABILITY_ROLE_PREFILL_COMPUTE: CapabilityRole +COMPRESSION_CODEC_UNSPECIFIED: CompressionCodec +COMPRESSION_CODEC_NONE: CompressionCodec +COMPRESSION_CODEC_ZLIB: CompressionCodec +PREFILL_JOB_STATUS_UNSPECIFIED: PrefillJobStatus +PREFILL_JOB_STATUS_QUEUED: PrefillJobStatus +PREFILL_JOB_STATUS_RUNNING: PrefillJobStatus +PREFILL_JOB_STATUS_COMPLETED: PrefillJobStatus +PREFILL_JOB_STATUS_FAILED: PrefillJobStatus +PREFILL_JOB_STATUS_CANCELLED: PrefillJobStatus class ModelCapability(_message.Message): __slots__ = ("model_id", "role", "quantization", "tokens_per_second") @@ -35,7 +61,7 @@ class ModelCapability(_message.Message): def __init__(self, model_id: _Optional[str] = ..., role: _Optional[_Union[CapabilityRole, str]] = ..., quantization: _Optional[str] = ..., tokens_per_second: _Optional[float] = ...) -> None: ... class NodeCapability(_message.Message): - __slots__ = ("node_id", "grpc_address", "platform", "unified_memory_bytes", "mlx_version", "models", "announced_at_unix", "ttl_seconds", "ring_address", "caches", "endpoints") + __slots__ = ("node_id", "grpc_address", "platform", "unified_memory_bytes", "mlx_version", "models", "announced_at_unix", "ttl_seconds", "ring_address", "caches", "endpoints", "prefill_workers") NODE_ID_FIELD_NUMBER: _ClassVar[int] GRPC_ADDRESS_FIELD_NUMBER: _ClassVar[int] PLATFORM_FIELD_NUMBER: _ClassVar[int] @@ -47,6 +73,7 @@ class NodeCapability(_message.Message): RING_ADDRESS_FIELD_NUMBER: _ClassVar[int] CACHES_FIELD_NUMBER: _ClassVar[int] ENDPOINTS_FIELD_NUMBER: _ClassVar[int] + PREFILL_WORKERS_FIELD_NUMBER: _ClassVar[int] node_id: str grpc_address: str platform: str @@ -58,7 +85,8 @@ class NodeCapability(_message.Message): ring_address: str caches: _containers.RepeatedCompositeFieldContainer[CacheCapability] endpoints: _containers.RepeatedCompositeFieldContainer[NodeEndpoint] - def __init__(self, node_id: _Optional[str] = ..., grpc_address: _Optional[str] = ..., platform: _Optional[str] = ..., unified_memory_bytes: _Optional[int] = ..., mlx_version: _Optional[str] = ..., models: _Optional[_Iterable[_Union[ModelCapability, _Mapping]]] = ..., announced_at_unix: _Optional[float] = ..., ttl_seconds: _Optional[float] = ..., ring_address: _Optional[str] = ..., caches: _Optional[_Iterable[_Union[CacheCapability, _Mapping]]] = ..., endpoints: _Optional[_Iterable[_Union[NodeEndpoint, _Mapping]]] = ...) -> None: ... + prefill_workers: _containers.RepeatedCompositeFieldContainer[PrefillWorkerCapability] + def __init__(self, node_id: _Optional[str] = ..., grpc_address: _Optional[str] = ..., platform: _Optional[str] = ..., unified_memory_bytes: _Optional[int] = ..., mlx_version: _Optional[str] = ..., models: _Optional[_Iterable[_Union[ModelCapability, _Mapping]]] = ..., announced_at_unix: _Optional[float] = ..., ttl_seconds: _Optional[float] = ..., ring_address: _Optional[str] = ..., caches: _Optional[_Iterable[_Union[CacheCapability, _Mapping]]] = ..., endpoints: _Optional[_Iterable[_Union[NodeEndpoint, _Mapping]]] = ..., prefill_workers: _Optional[_Iterable[_Union[PrefillWorkerCapability, _Mapping]]] = ...) -> None: ... class NodeEndpoint(_message.Message): __slots__ = ("address", "network", "priority", "measured_rtt_ms") @@ -73,7 +101,7 @@ class NodeEndpoint(_message.Message): def __init__(self, address: _Optional[str] = ..., network: _Optional[str] = ..., priority: _Optional[int] = ..., measured_rtt_ms: _Optional[float] = ...) -> None: ... class CacheCompatibility(_message.Message): - __slots__ = ("model_id", "model_revision", "tokenizer_revision", "cache_format_version", "quantization", "rope_hash", "layer_geometry_hash", "kv_dtype", "block_size_tokens") + __slots__ = ("model_id", "model_revision", "tokenizer_revision", "cache_format_version", "quantization", "rope_hash", "layer_geometry_hash", "kv_dtype", "block_size_tokens", "tenant_namespace", "sink_size", "window_size") MODEL_ID_FIELD_NUMBER: _ClassVar[int] MODEL_REVISION_FIELD_NUMBER: _ClassVar[int] TOKENIZER_REVISION_FIELD_NUMBER: _ClassVar[int] @@ -83,6 +111,9 @@ class CacheCompatibility(_message.Message): LAYER_GEOMETRY_HASH_FIELD_NUMBER: _ClassVar[int] KV_DTYPE_FIELD_NUMBER: _ClassVar[int] BLOCK_SIZE_TOKENS_FIELD_NUMBER: _ClassVar[int] + TENANT_NAMESPACE_FIELD_NUMBER: _ClassVar[int] + SINK_SIZE_FIELD_NUMBER: _ClassVar[int] + WINDOW_SIZE_FIELD_NUMBER: _ClassVar[int] model_id: str model_revision: str tokenizer_revision: str @@ -92,10 +123,13 @@ class CacheCompatibility(_message.Message): layer_geometry_hash: str kv_dtype: str block_size_tokens: int - 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] = ...) -> None: ... + tenant_namespace: str + sink_size: int + window_size: int + 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") + __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter", "default_compression", "replication_factor") COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int] CACHE_BYTES_USED_FIELD_NUMBER: _ClassVar[int] @@ -105,6 +139,8 @@ class CacheCapability(_message.Message): LOAD_FIELD_NUMBER: _ClassVar[int] TOKENS_SERVED_FIELD_NUMBER: _ClassVar[int] BLOOM_FILTER_FIELD_NUMBER: _ClassVar[int] + DEFAULT_COMPRESSION_FIELD_NUMBER: _ClassVar[int] + REPLICATION_FACTOR_FIELD_NUMBER: _ClassVar[int] compatibility: CacheCompatibility cache_address: str cache_bytes_used: int @@ -114,7 +150,33 @@ class CacheCapability(_message.Message): load: float tokens_served: int bloom_filter: bytes - 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] = ...) -> None: ... + 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: ... + +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") + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + WORKER_ADDRESS_FIELD_NUMBER: _ClassVar[int] + MAX_CONCURRENT_JOBS_FIELD_NUMBER: _ClassVar[int] + INFLIGHT_JOBS_FIELD_NUMBER: _ClassVar[int] + QUEUED_JOBS_FIELD_NUMBER: _ClassVar[int] + LOAD_FIELD_NUMBER: _ClassVar[int] + TOKENS_PER_SECOND_PREFILL_FIELD_NUMBER: _ClassVar[int] + RAM_BYTES_FREE_FIELD_NUMBER: _ClassVar[int] + ACCEPTS_COMPUTE_JOBS_FIELD_NUMBER: _ClassVar[int] + QUEUED_TOKENS_FIELD_NUMBER: _ClassVar[int] + compatibility: CacheCompatibility + worker_address: str + max_concurrent_jobs: int + inflight_jobs: int + queued_jobs: int + load: float + tokens_per_second_prefill: float + ram_bytes_free: int + accepts_compute_jobs: bool + queued_tokens: int + def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., worker_address: _Optional[str] = ..., max_concurrent_jobs: _Optional[int] = ..., inflight_jobs: _Optional[int] = ..., queued_jobs: _Optional[int] = ..., load: _Optional[float] = ..., tokens_per_second_prefill: _Optional[float] = ..., ram_bytes_free: _Optional[int] = ..., accepts_compute_jobs: _Optional[bool] = ..., queued_tokens: _Optional[int] = ...) -> None: ... class ExchangeCapabilitiesRequest(_message.Message): __slots__ = ("known_nodes",) @@ -236,6 +298,82 @@ class PublishBlockResponse(_message.Message): cache_epoch: int def __init__(self, stored: _Optional[bool] = ..., cache_epoch: _Optional[int] = ...) -> None: ... +class SubmitPrefillJobRequest(_message.Message): + __slots__ = ("request_id", "tenant_id", "compatibility", "token_ids", "block_hashes", "deadline_ms", "preferred_compression") + REQUEST_ID_FIELD_NUMBER: _ClassVar[int] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + TOKEN_IDS_FIELD_NUMBER: _ClassVar[int] + BLOCK_HASHES_FIELD_NUMBER: _ClassVar[int] + DEADLINE_MS_FIELD_NUMBER: _ClassVar[int] + PREFERRED_COMPRESSION_FIELD_NUMBER: _ClassVar[int] + request_id: str + tenant_id: str + compatibility: CacheCompatibility + token_ids: _containers.RepeatedScalarFieldContainer[int] + block_hashes: _containers.RepeatedScalarFieldContainer[bytes] + deadline_ms: int + preferred_compression: CompressionCodec + def __init__(self, request_id: _Optional[str] = ..., tenant_id: _Optional[str] = ..., compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., token_ids: _Optional[_Iterable[int]] = ..., block_hashes: _Optional[_Iterable[bytes]] = ..., deadline_ms: _Optional[int] = ..., preferred_compression: _Optional[_Union[CompressionCodec, str]] = ...) -> None: ... + +class SubmitPrefillJobResponse(_message.Message): + __slots__ = ("job_id", "status", "worker_node_id", "queue_eta_ms") + JOB_ID_FIELD_NUMBER: _ClassVar[int] + STATUS_FIELD_NUMBER: _ClassVar[int] + WORKER_NODE_ID_FIELD_NUMBER: _ClassVar[int] + QUEUE_ETA_MS_FIELD_NUMBER: _ClassVar[int] + job_id: str + status: PrefillJobStatus + worker_node_id: str + queue_eta_ms: float + def __init__(self, job_id: _Optional[str] = ..., status: _Optional[_Union[PrefillJobStatus, str]] = ..., worker_node_id: _Optional[str] = ..., queue_eta_ms: _Optional[float] = ...) -> None: ... + +class GetPrefillJobStatusRequest(_message.Message): + __slots__ = ("job_id", "tenant_id") + JOB_ID_FIELD_NUMBER: _ClassVar[int] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + job_id: str + tenant_id: str + def __init__(self, job_id: _Optional[str] = ..., tenant_id: _Optional[str] = ...) -> None: ... + +class GetPrefillJobStatusResponse(_message.Message): + __slots__ = ("job_id", "status", "tokens_computed", "lease_id", "block_hash", "payload_sha256", "transfer_bytes", "failure_reason", "compute_ms", "cache_address") + JOB_ID_FIELD_NUMBER: _ClassVar[int] + STATUS_FIELD_NUMBER: _ClassVar[int] + TOKENS_COMPUTED_FIELD_NUMBER: _ClassVar[int] + LEASE_ID_FIELD_NUMBER: _ClassVar[int] + BLOCK_HASH_FIELD_NUMBER: _ClassVar[int] + PAYLOAD_SHA256_FIELD_NUMBER: _ClassVar[int] + TRANSFER_BYTES_FIELD_NUMBER: _ClassVar[int] + FAILURE_REASON_FIELD_NUMBER: _ClassVar[int] + COMPUTE_MS_FIELD_NUMBER: _ClassVar[int] + CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int] + job_id: str + status: PrefillJobStatus + tokens_computed: int + lease_id: str + block_hash: bytes + payload_sha256: bytes + transfer_bytes: int + failure_reason: str + compute_ms: float + cache_address: str + def __init__(self, job_id: _Optional[str] = ..., status: _Optional[_Union[PrefillJobStatus, str]] = ..., tokens_computed: _Optional[int] = ..., lease_id: _Optional[str] = ..., block_hash: _Optional[bytes] = ..., payload_sha256: _Optional[bytes] = ..., transfer_bytes: _Optional[int] = ..., failure_reason: _Optional[str] = ..., compute_ms: _Optional[float] = ..., cache_address: _Optional[str] = ...) -> None: ... + +class CancelPrefillJobRequest(_message.Message): + __slots__ = ("job_id", "tenant_id") + JOB_ID_FIELD_NUMBER: _ClassVar[int] + TENANT_ID_FIELD_NUMBER: _ClassVar[int] + job_id: str + tenant_id: str + def __init__(self, job_id: _Optional[str] = ..., tenant_id: _Optional[str] = ...) -> None: ... + +class CancelPrefillJobResponse(_message.Message): + __slots__ = ("cancelled",) + CANCELLED_FIELD_NUMBER: _ClassVar[int] + cancelled: bool + def __init__(self, cancelled: _Optional[bool] = ...) -> None: ... + class ProposeBlockRequest(_message.Message): __slots__ = ("committed_token_ids", "block_size", "num_steps", "model_id") COMMITTED_TOKEN_IDS_FIELD_NUMBER: _ClassVar[int] diff --git a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py index 8885f712..dcfd367b 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2_grpc.py @@ -468,6 +468,173 @@ def PublishBlock(request_iterator, _registered_method=True) +class PrefillWorkerServiceStub: + """PrefillWorkerService executes model prefill only. Workers load the same model + as the primary, materialize immutable snapshots into their co-located + PrefillCacheService, and never participate in autoregressive decode. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.SubmitPrefillJob = channel.unary_unary( + '/kakeya.v1.PrefillWorkerService/SubmitPrefillJob', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.FromString, + _registered_method=True) + self.GetPrefillJobStatus = channel.unary_unary( + '/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.FromString, + _registered_method=True) + self.CancelPrefillJob = channel.unary_unary( + '/kakeya.v1.PrefillWorkerService/CancelPrefillJob', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.FromString, + _registered_method=True) + + +class PrefillWorkerServiceServicer: + """PrefillWorkerService executes model prefill only. Workers load the same model + as the primary, materialize immutable snapshots into their co-located + PrefillCacheService, and never participate in autoregressive decode. + """ + + def SubmitPrefillJob(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetPrefillJobStatus(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CancelPrefillJob(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_PrefillWorkerServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'SubmitPrefillJob': grpc.unary_unary_rpc_method_handler( + servicer.SubmitPrefillJob, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.SerializeToString, + ), + 'GetPrefillJobStatus': grpc.unary_unary_rpc_method_handler( + servicer.GetPrefillJobStatus, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.SerializeToString, + ), + 'CancelPrefillJob': grpc.unary_unary_rpc_method_handler( + servicer.CancelPrefillJob, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'kakeya.v1.PrefillWorkerService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('kakeya.v1.PrefillWorkerService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class PrefillWorkerService: + """PrefillWorkerService executes model prefill only. Workers load the same model + as the primary, materialize immutable snapshots into their co-located + PrefillCacheService, and never participate in autoregressive decode. + """ + + @staticmethod + def SubmitPrefillJob(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/kakeya.v1.PrefillWorkerService/SubmitPrefillJob', + kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.SubmitPrefillJobResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def GetPrefillJobStatus(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus', + kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.GetPrefillJobStatusResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def CancelPrefillJob(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/kakeya.v1.PrefillWorkerService/CancelPrefillJob', + kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.CancelPrefillJobResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + class DFlashProposerServiceStub: """DFlashProposerService: stateful remote DFlash drafter + f_θ restoration. Per turn: Restore (prompt -> f_θ-projected verifier K/V) then SeedContext diff --git a/proto/kakeya/v1/distributed.proto b/proto/kakeya/v1/distributed.proto index 8e51fa6b..cf9b679d 100644 --- a/proto/kakeya/v1/distributed.proto +++ b/proto/kakeya/v1/distributed.proto @@ -72,6 +72,15 @@ service PrefillCacheService { rpc PublishBlock(stream PublishBlockRequest) returns (PublishBlockResponse); } +// PrefillWorkerService executes model prefill only. Workers load the same model +// as the primary, materialize immutable snapshots into their co-located +// PrefillCacheService, and never participate in autoregressive decode. +service PrefillWorkerService { + rpc SubmitPrefillJob(SubmitPrefillJobRequest) returns (SubmitPrefillJobResponse); + rpc GetPrefillJobStatus(GetPrefillJobStatusRequest) returns (GetPrefillJobStatusResponse); + rpc CancelPrefillJob(CancelPrefillJobRequest) returns (CancelPrefillJobResponse); +} + // ----------------------------------------------------------------------------- // Capability messages // ----------------------------------------------------------------------------- @@ -96,6 +105,24 @@ enum CapabilityRole { // The node can answer PrefillCacheService lookups and stream compatible // immutable prefill K/V blocks. CAPABILITY_ROLE_PREFILL_CACHE = 5; + + // The node has the compatible model loaded and accepts prefill-only jobs. + CAPABILITY_ROLE_PREFILL_COMPUTE = 6; +} + +enum CompressionCodec { + COMPRESSION_CODEC_UNSPECIFIED = 0; + COMPRESSION_CODEC_NONE = 1; + COMPRESSION_CODEC_ZLIB = 2; +} + +enum PrefillJobStatus { + PREFILL_JOB_STATUS_UNSPECIFIED = 0; + PREFILL_JOB_STATUS_QUEUED = 1; + PREFILL_JOB_STATUS_RUNNING = 2; + PREFILL_JOB_STATUS_COMPLETED = 3; + PREFILL_JOB_STATUS_FAILED = 4; + PREFILL_JOB_STATUS_CANCELLED = 5; } // ModelCapability is one (model, role) a node offers. @@ -164,6 +191,9 @@ message NodeCapability { // Reachable interfaces ordered by operator preference (Thunderbolt before // LAN/Tailscale). grpc_address remains the compatibility default. repeated NodeEndpoint endpoints = 11; + + // Prefill-only compute offerings. Cache-only nodes leave this empty. + repeated PrefillWorkerCapability prefill_workers = 12; } message NodeEndpoint { @@ -183,6 +213,12 @@ message CacheCompatibility { string layer_geometry_hash = 7; string kv_dtype = 8; uint32 block_size_tokens = 9; + // Tenant namespace is part of the compatibility fingerprint. It prevents a + // snapshot from being reused across isolation domains even with identical + // model/tokenizer geometry. + string tenant_namespace = 10; + uint32 sink_size = 11; + uint32 window_size = 12; } message CacheCapability { @@ -196,6 +232,21 @@ message CacheCapability { uint64 tokens_served = 8; // Compact probabilistic summary. Empty means "query me directly". bytes bloom_filter = 9; + CompressionCodec default_compression = 10; + uint32 replication_factor = 11; +} + +message PrefillWorkerCapability { + CacheCompatibility compatibility = 1; + string worker_address = 2; + uint32 max_concurrent_jobs = 3; + uint32 inflight_jobs = 4; + uint32 queued_jobs = 5; + double load = 6; + double tokens_per_second_prefill = 7; + uint64 ram_bytes_free = 8; + bool accepts_compute_jobs = 9; + uint64 queued_tokens = 10; } message ExchangeCapabilitiesRequest { @@ -277,6 +328,57 @@ message PublishBlockResponse { uint64 cache_epoch = 2; } +// ----------------------------------------------------------------------------- +// Prefill-only compute jobs +// ----------------------------------------------------------------------------- + +message SubmitPrefillJobRequest { + // Caller-generated idempotency key. Repeating it returns the existing job. + string request_id = 1; + string tenant_id = 2; + CacheCompatibility compatibility = 3; + repeated uint32 token_ids = 4; + // Chained hashes computed by the primary (tenant-HMACed in authenticated + // mode). One hash per token block. + repeated bytes block_hashes = 5; + uint32 deadline_ms = 6; + CompressionCodec preferred_compression = 7; +} + +message SubmitPrefillJobResponse { + string job_id = 1; + PrefillJobStatus status = 2; + string worker_node_id = 3; + double queue_eta_ms = 4; +} + +message GetPrefillJobStatusRequest { + string job_id = 1; + string tenant_id = 2; +} + +message GetPrefillJobStatusResponse { + string job_id = 1; + PrefillJobStatus status = 2; + uint32 tokens_computed = 3; + string lease_id = 4; + bytes block_hash = 5; + bytes payload_sha256 = 6; + uint64 transfer_bytes = 7; + string failure_reason = 8; + double compute_ms = 9; + string cache_address = 10; +} + +message CancelPrefillJobRequest { + string job_id = 1; + string tenant_id = 2; +} + +message CancelPrefillJobResponse { + bool cancelled = 1; +} + // ----------------------------------------------------------------------------- // Remote proposal messages // ----------------------------------------------------------------------------- diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index d5529ade..acda21bd 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -213,6 +213,10 @@ def _build_capability_registry( cache_capability( cache_store, cache_address=args.cache_advertise or args.advertise or args.bind, + default_compression=( + 2 if args.cache_compression == "zlib" else 1 + ), + replication_factor=args.cache_replication_factor, ), ) @@ -253,6 +257,8 @@ async def _exchange_loop( *, cache_store=None, cache_address: str = "", + cache_compression: int = 1, + cache_replication_factor: int = 1, ) -> None: """Periodic gossip with seed peers until the task is cancelled.""" from inference_engine.distributed.exchange import exchange_once @@ -269,6 +275,8 @@ async def _exchange_loop( cache_capability( cache_store, cache_address=cache_address, + default_compression=cache_compression, + replication_factor=cache_replication_factor, ), ), ) @@ -369,15 +377,24 @@ async def _serve(args: argparse.Namespace) -> int: prefill_store = None prefill_hook = None + prefill_auth = None + capability_registry_holder = [None] if args.enable_prefill_cache: if args.backend != "mlx": raise SystemExit("--enable-prefill-cache currently requires --backend mlx") import hashlib - from inference_engine.distributed.capability import CacheCompatibility + from inference_engine.distributed.capability import ( + CacheCompatibility, + CompressionCodec, + ) + from inference_engine.distributed.prefill_auth import FleetAuthConfig from inference_engine.distributed.prefill_cache import PrefixCacheStore from inference_engine.distributed.prefill_cache_runtime import ( DistributedPrefillCacheHook, ) + from inference_engine.distributed.prefill_scheduler import ( + PrefillCostConfig, + ) geometry = f"{num_layers}:{num_kv_heads}:{head_dim}" compatibility = CacheCompatibility( @@ -390,6 +407,9 @@ async def _serve(args: argparse.Namespace) -> int: layer_geometry_hash=hashlib.sha256(geometry.encode()).hexdigest(), kv_dtype=args.cache_kv_dtype, block_size_tokens=args.cache_block_tokens, + tenant_namespace=args.cache_tenant_id, + sink_size=args.sink, + window_size=args.window, ) prefill_store = PrefixCacheStore( compatibility, @@ -397,11 +417,40 @@ async def _serve(args: argparse.Namespace) -> int: node_id=args.node_id or (__import__("platform").node() or "localhost"), ) telemetry_callback = _build_token_telemetry_callback(args) + if args.fleet_psk_file: + prefill_auth = FleetAuthConfig.from_file( + args.fleet_psk_file, + tenant_id=args.cache_tenant_id, + node_id=args.node_id or "primary", + ) prefill_hook = DistributedPrefillCacheHook( prefill_store, peers=args.cache_peer, + registry_provider=lambda: ( + capability_registry_holder[0].snapshot() + if capability_registry_holder[0] is not None else () + ), lookup_timeout_s=args.cache_lookup_timeout_s, fetch_timeout_s=args.cache_fetch_timeout_s, + worker_timeout_s=args.prefill_worker_timeout_s, + remote_compute_min_tokens=args.remote_prefill_min_tokens, + max_import_bytes=int(args.cache_max_import_gb * (1 << 30)), + estimated_snapshot_bytes_per_token=args.cache_estimated_bytes_per_token, + compression=( + CompressionCodec.ZLIB + if args.cache_compression == "zlib" + else CompressionCodec.NONE + ), + replication_factor=args.cache_replication_factor, + cost_config=PrefillCostConfig( + local_prefill_tps=args.local_prefill_tps, + default_worker_tps=args.worker_prefill_tps, + link_mbps=args.cache_link_mbps, + default_rtt_ms=args.cache_default_rtt_ms, + minimum_savings_ratio=args.prefill_min_savings_ratio, + primary_compute_penalty_ms=args.primary_prefill_penalty_ms, + ), + auth=prefill_auth, on_reuse=( (lambda count: telemetry_callback(count, count)) if telemetry_callback is not None else None @@ -467,6 +516,7 @@ async def _serve(args: argparse.Namespace) -> int: backend=args.backend, cache_store=prefill_store, ) + capability_registry_holder[0] = registry if args.serve_ngram_proposer: from inference_engine.distributed.capability import NGRAM_MODEL_ID from inference_engine.distributed.ngram import NGramProposer @@ -486,6 +536,7 @@ async def _serve(args: argparse.Namespace) -> int: proposers=proposers, prefill_cache_store=prefill_store, prefill_cache_address=args.cache_advertise or args.advertise or args.bind, + prefill_auth=prefill_auth, ) await server.start() @@ -532,6 +583,10 @@ async def _serve(args: argparse.Namespace) -> int: args.exchange_interval_s, cache_store=prefill_store, cache_address=args.cache_advertise or args.advertise or args.bind, + cache_compression=( + 2 if args.cache_compression == "zlib" else 1 + ), + cache_replication_factor=args.cache_replication_factor, ), ) _LOG.info( @@ -563,6 +618,8 @@ def _on_signal(sig: int) -> None: http_server.should_exit = True if http_task is not None: await http_task + if prefill_hook is not None: + prefill_hook.close() await server.stop(grace=args.shutdown_grace_s) _LOG.info("kakeya gRPC RuntimeService stopped cleanly") return 0 @@ -648,7 +705,7 @@ def main() -> int: "Defaults to --advertise/--bind.") ap.add_argument("--cache-block-tokens", type=int, default=64, help="Token boundary interval for restorable snapshots.") - ap.add_argument("--cache-format-version", default="kakeya-prefill-v1") + ap.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib") ap.add_argument("--cache-model-id", default="", help="Logical cache model id; defaults to --verifier-id. " "Use this when verifier-id is a host-specific path.") @@ -663,6 +720,28 @@ def main() -> int: help="RoPE/position configuration fingerprint.") ap.add_argument("--cache-lookup-timeout-s", type=float, default=2.0) ap.add_argument("--cache-fetch-timeout-s", type=float, default=30.0) + ap.add_argument("--cache-tenant-id", default="default", + help="Tenant namespace included in cache compatibility.") + ap.add_argument("--fleet-psk-file", default="", + help="Optional fleet PSK file for authenticated prefill RPCs " + "and tenant-HMAC prefix hashes.") + ap.add_argument("--cache-compression", choices=["none", "zlib"], + default="zlib") + ap.add_argument("--cache-replication-factor", type=int, default=1) + ap.add_argument("--cache-max-import-gb", type=float, default=1.0, + help="Reject remote snapshots whose wire or expanded size " + "would exceed this import budget.") + ap.add_argument("--cache-estimated-bytes-per-token", type=int, default=400000) + ap.add_argument("--cache-link-mbps", type=float, default=1000.0) + ap.add_argument("--cache-default-rtt-ms", type=float, default=2.0) + ap.add_argument("--local-prefill-tps", type=float, default=20.0) + ap.add_argument("--worker-prefill-tps", type=float, default=20.0) + ap.add_argument("--prefill-min-savings-ratio", type=float, default=0.10) + ap.add_argument("--primary-prefill-penalty-ms", type=float, default=5000.0, + help="Opportunity cost assigned to blocking the primary " + "with prefill; drives work to compute peers.") + ap.add_argument("--remote-prefill-min-tokens", type=int, default=128) + ap.add_argument("--prefill-worker-timeout-s", type=float, default=120.0) ap.add_argument("--network-label", default="lan", help="Advertised interface: thunderbolt|lan|tailscale|public.") ap.add_argument("--network-priority", type=int, default=50) diff --git a/scripts/start_prefill_cache_node.py b/scripts/start_prefill_cache_node.py index 4a55f3ae..6778e5d3 100644 --- a/scripts/start_prefill_cache_node.py +++ b/scripts/start_prefill_cache_node.py @@ -24,6 +24,7 @@ from inference_engine.distributed.capability import ( CacheCompatibility, + CompressionCodec, CapabilityRegistry, CapabilityRole, ModelCapability, @@ -35,6 +36,7 @@ exchange_once, ) from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.distributed.prefill_auth import FleetAuthConfig from inference_engine.distributed.prefill_cache_service import ( add_prefill_cache_service, cache_capability, @@ -62,12 +64,23 @@ async def serve(args) -> None: layer_geometry_hash=args.layer_geometry_hash, kv_dtype=args.kv_dtype, block_size_tokens=args.block_size_tokens, + tenant_namespace=args.tenant_id, + sink_size=args.sink, + window_size=args.window, ) store = PrefixCacheStore( compatibility, max_bytes=int(args.cache_gb * (1 << 30)), node_id=args.node_id, ) + auth = ( + FleetAuthConfig.from_file( + args.fleet_psk_file, + tenant_id=args.tenant_id, + node_id=args.node_id, + ) + if args.fleet_psk_file else None + ) card = NodeCapability( node_id=args.node_id, grpc_address=args.advertise, @@ -82,7 +95,16 @@ async def serve(args) -> None: ), announced_at_unix=time.time(), ttl_seconds=args.ttl_seconds, - caches=(cache_capability(store, cache_address=args.advertise),), + caches=(cache_capability( + store, + cache_address=args.advertise, + default_compression=( + CompressionCodec.ZLIB + if args.cache_compression == "zlib" + else CompressionCodec.NONE + ), + replication_factor=args.replication_factor, + ),), endpoints=( NodeEndpoint( args.advertise, @@ -99,6 +121,7 @@ async def serve(args) -> None: grpc_server, store, cache_address=args.advertise, + auth=auth, ) grpc_server.add_insecure_port(args.bind) await grpc_server.start() @@ -169,12 +192,19 @@ def main() -> None: ap.add_argument("--model-id", required=True) ap.add_argument("--model-revision", default="") ap.add_argument("--tokenizer-revision", default="") - ap.add_argument("--cache-format-version", default="kakeya-prefill-v1") + ap.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib") ap.add_argument("--quantization", default="") ap.add_argument("--rope-hash", default="") ap.add_argument("--layer-geometry-hash", default="") ap.add_argument("--kv-dtype", default="bfloat16") ap.add_argument("--block-size-tokens", type=int, default=64) + ap.add_argument("--tenant-id", default="default") + ap.add_argument("--sink", type=int, default=4) + ap.add_argument("--window", type=int, default=64) + ap.add_argument("--fleet-psk-file", default="") + ap.add_argument("--cache-compression", choices=["none", "zlib"], + default="zlib") + ap.add_argument("--replication-factor", type=int, default=1) ap.add_argument("--cache-gb", type=float, default=4) ap.add_argument("--platform", default="") ap.add_argument("--memory-bytes", type=int, default=0) diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py new file mode 100644 index 00000000..dda9138c --- /dev/null +++ b/scripts/start_prefill_worker_node.py @@ -0,0 +1,252 @@ +"""Start a prefill-only MLX compute worker with a co-located RAM cache. + +The worker loads the same model as the primary, accepts PrefillWorkerService +jobs, stores immutable snapshots, and never serves user decode. +""" +from __future__ import annotations + +import argparse +import asyncio +import logging +import platform +import signal +import sys +import time +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import grpc +import torch + +from inference_engine.backends.mlx.prefill_worker import MLXPrefillComputeEngine +from inference_engine.backends.mlx.verifier import MLXSinkWindowVerifier +from inference_engine.distributed.capability import ( + CacheCompatibility, + CapabilityRegistry, + CapabilityRole, + CompressionCodec, + ModelCapability, + NodeCapability, + NodeEndpoint, + PrefillWorkerCapability, +) +from inference_engine.distributed.exchange import ( + add_capability_service, + exchange_once, +) +from inference_engine.distributed.prefill_auth import FleetAuthConfig +from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.distributed.prefill_cache_service import ( + add_prefill_cache_service, + cache_capability, +) +from inference_engine.distributed.prefill_worker import ( + PrefillJobStore, + add_prefill_worker_service, +) +from kv_cache_proposer.verifier import VerifierConfig + +_LOG = logging.getLogger("kakeya.prefill-worker") + + +def physical_memory_bytes() -> int: + try: + import os + return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) + except (ValueError, OSError, AttributeError): + return 0 + + +async def serve(args) -> None: + compatibility = CacheCompatibility( + model_id=args.cache_model_id or args.model_id, + model_revision=args.model_revision, + tokenizer_revision=args.tokenizer_revision, + cache_format_version=args.cache_format_version, + quantization=args.quantization, + rope_hash=args.rope_hash, + layer_geometry_hash=args.layer_geometry_hash, + kv_dtype=args.kv_dtype, + block_size_tokens=args.block_size_tokens, + tenant_namespace=args.tenant_id, + sink_size=args.sink, + window_size=args.window, + ) + auth = ( + FleetAuthConfig.from_file( + args.fleet_psk_file, + tenant_id=args.tenant_id, + node_id=args.node_id, + ) + if args.fleet_psk_file else None + ) + verifier = MLXSinkWindowVerifier(VerifierConfig( + model_id=args.model_id, + sink_size=args.sink, + window_size=args.window, + dtype=torch.bfloat16, + device="cpu", + )) + store = PrefixCacheStore( + compatibility, + max_bytes=int(args.cache_gb * (1 << 30)), + node_id=args.node_id, + ) + engine = MLXPrefillComputeEngine(verifier, compatibility) + jobs = PrefillJobStore( + engine, + store, + max_concurrent_jobs=args.max_concurrent_jobs, + max_jobs=args.max_jobs, + completed_ttl_s=args.job_ttl_s, + max_prompt_tokens=args.max_prompt_tokens, + ) + + def card() -> NodeCapability: + inflight, queued, load, queued_tokens = jobs.stats() + worker = PrefillWorkerCapability( + compatibility=compatibility, + worker_address=args.advertise, + max_concurrent_jobs=args.max_concurrent_jobs, + inflight_jobs=inflight, + queued_jobs=queued, + load=load, + tokens_per_second_prefill=args.prefill_tps, + ram_bytes_free=max( + 0, + physical_memory_bytes() - store.stats().bytes_used, + ), + queued_tokens=queued_tokens, + ) + return NodeCapability( + node_id=args.node_id, + grpc_address=args.advertise, + platform=f"{platform.system()}-{platform.machine()}", + unified_memory_bytes=physical_memory_bytes(), + models=( + ModelCapability( + args.cache_model_id or args.model_id, + CapabilityRole.PREFILL_COMPUTE, + args.quantization, + args.prefill_tps, + ), + ), + announced_at_unix=time.time(), + ttl_seconds=args.ttl_seconds, + caches=(cache_capability( + store, + cache_address=args.advertise, + load=load, + default_compression=( + CompressionCodec.ZLIB + if args.cache_compression == "zlib" + else CompressionCodec.NONE + ), + replication_factor=args.replication_factor, + ),), + endpoints=( + NodeEndpoint( + args.advertise, + args.network, + args.priority, + args.rtt_ms, + ), + ), + prefill_workers=(worker,), + ) + + registry = CapabilityRegistry(card()) + server = grpc.aio.server( + maximum_concurrent_rpcs=args.max_concurrent_rpcs, + ) + add_capability_service(server, registry) + add_prefill_cache_service( + server, + store, + cache_address=args.advertise, + auth=auth, + ) + add_prefill_worker_service( + server, + jobs, + node_id=args.node_id, + cache_address=args.advertise, + auth=auth, + tokens_per_second_prefill=args.prefill_tps, + ) + server.add_insecure_port(args.bind) + await server.start() + _LOG.info("prefill worker ready on %s for %s", args.bind, args.model_id) + + stop = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, stop.set) + + async def gossip() -> None: + while not stop.is_set(): + registry.self_card = replace(card(), announced_at_unix=time.time()) + if args.peer: + await exchange_once(registry, args.peer, timeout_s=args.gossip_timeout_s) + try: + await asyncio.wait_for(stop.wait(), timeout=args.gossip_interval_s) + except asyncio.TimeoutError: + pass + + gossip_task = asyncio.create_task(gossip()) + await stop.wait() + gossip_task.cancel() + jobs.close() + await server.stop(grace=2.0) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--node-id", required=True) + parser.add_argument("--bind", default="127.0.0.1:53051") + parser.add_argument("--advertise", default="127.0.0.1:53051") + parser.add_argument("--peer", action="append", default=[]) + parser.add_argument("--model-id", required=True) + parser.add_argument("--model-revision", default="") + parser.add_argument("--tokenizer-revision", default="") + parser.add_argument("--cache-format-version", default="kakeya-prefill-v2-zlib") + parser.add_argument("--cache-model-id", default="", + help="Logical model id used for compatibility; defaults " + "to --model-id (which may be a host-specific path).") + parser.add_argument("--quantization", default="4bit-mlx") + parser.add_argument("--rope-hash", default="") + parser.add_argument("--layer-geometry-hash", required=True) + parser.add_argument("--kv-dtype", default="bfloat16") + parser.add_argument("--block-size-tokens", type=int, default=64) + parser.add_argument("--tenant-id", default="default") + parser.add_argument("--fleet-psk-file") + parser.add_argument("--sink", type=int, default=4) + parser.add_argument("--window", type=int, default=64) + parser.add_argument("--cache-gb", type=float, default=4.0) + parser.add_argument("--cache-compression", choices=["none", "zlib"], + default="zlib") + parser.add_argument("--replication-factor", type=int, default=1) + parser.add_argument("--max-concurrent-jobs", type=int, default=1) + parser.add_argument("--max-jobs", type=int, default=128) + parser.add_argument("--max-prompt-tokens", type=int, default=131072) + parser.add_argument("--job-ttl-s", type=float, default=600.0) + parser.add_argument("--prefill-tps", type=float, default=20.0) + parser.add_argument("--max-concurrent-rpcs", type=int, default=32) + parser.add_argument("--ttl-seconds", type=float, default=120.0) + parser.add_argument("--gossip-interval-s", type=float, default=10.0) + parser.add_argument("--gossip-timeout-s", type=float, default=3.0) + parser.add_argument("--network", default="lan") + parser.add_argument("--priority", type=int, default=50) + parser.add_argument("--rtt-ms", type=float, default=1.0) + parser.add_argument("--log-level", default="INFO") + args = parser.parse_args() + logging.basicConfig(level=getattr(logging, args.log_level.upper())) + asyncio.run(serve(args)) + + +if __name__ == "__main__": + main() + diff --git a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts index e35e7130..868f3d3b 100644 --- a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts +++ b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts @@ -47,6 +47,8 @@ export enum CapabilityRole { * immutable prefill K/V blocks. */ PREFILL_CACHE = 5, + /** PREFILL_COMPUTE - The node has the compatible model loaded and accepts prefill-only jobs. */ + PREFILL_COMPUTE = 6, UNRECOGNIZED = -1, } @@ -70,6 +72,9 @@ export function capabilityRoleFromJSON(object: any): CapabilityRole { case 5: case "CAPABILITY_ROLE_PREFILL_CACHE": return CapabilityRole.PREFILL_CACHE; + case 6: + case "CAPABILITY_ROLE_PREFILL_COMPUTE": + return CapabilityRole.PREFILL_COMPUTE; case -1: case "UNRECOGNIZED": default: @@ -91,12 +96,110 @@ export function capabilityRoleToJSON(object: CapabilityRole): string { return "CAPABILITY_ROLE_TOOL"; case CapabilityRole.PREFILL_CACHE: return "CAPABILITY_ROLE_PREFILL_CACHE"; + case CapabilityRole.PREFILL_COMPUTE: + return "CAPABILITY_ROLE_PREFILL_COMPUTE"; case CapabilityRole.UNRECOGNIZED: default: return "UNRECOGNIZED"; } } +export enum CompressionCodec { + UNSPECIFIED = 0, + NONE = 1, + ZLIB = 2, + UNRECOGNIZED = -1, +} + +export function compressionCodecFromJSON(object: any): CompressionCodec { + switch (object) { + case 0: + case "COMPRESSION_CODEC_UNSPECIFIED": + return CompressionCodec.UNSPECIFIED; + case 1: + case "COMPRESSION_CODEC_NONE": + return CompressionCodec.NONE; + case 2: + case "COMPRESSION_CODEC_ZLIB": + return CompressionCodec.ZLIB; + case -1: + case "UNRECOGNIZED": + default: + return CompressionCodec.UNRECOGNIZED; + } +} + +export function compressionCodecToJSON(object: CompressionCodec): string { + switch (object) { + case CompressionCodec.UNSPECIFIED: + return "COMPRESSION_CODEC_UNSPECIFIED"; + case CompressionCodec.NONE: + return "COMPRESSION_CODEC_NONE"; + case CompressionCodec.ZLIB: + return "COMPRESSION_CODEC_ZLIB"; + case CompressionCodec.UNRECOGNIZED: + default: + return "UNRECOGNIZED"; + } +} + +export enum PrefillJobStatus { + UNSPECIFIED = 0, + QUEUED = 1, + RUNNING = 2, + COMPLETED = 3, + FAILED = 4, + CANCELLED = 5, + UNRECOGNIZED = -1, +} + +export function prefillJobStatusFromJSON(object: any): PrefillJobStatus { + switch (object) { + case 0: + case "PREFILL_JOB_STATUS_UNSPECIFIED": + return PrefillJobStatus.UNSPECIFIED; + case 1: + case "PREFILL_JOB_STATUS_QUEUED": + return PrefillJobStatus.QUEUED; + case 2: + case "PREFILL_JOB_STATUS_RUNNING": + return PrefillJobStatus.RUNNING; + case 3: + case "PREFILL_JOB_STATUS_COMPLETED": + return PrefillJobStatus.COMPLETED; + case 4: + case "PREFILL_JOB_STATUS_FAILED": + return PrefillJobStatus.FAILED; + case 5: + case "PREFILL_JOB_STATUS_CANCELLED": + return PrefillJobStatus.CANCELLED; + case -1: + case "UNRECOGNIZED": + default: + return PrefillJobStatus.UNRECOGNIZED; + } +} + +export function prefillJobStatusToJSON(object: PrefillJobStatus): string { + switch (object) { + case PrefillJobStatus.UNSPECIFIED: + return "PREFILL_JOB_STATUS_UNSPECIFIED"; + case PrefillJobStatus.QUEUED: + return "PREFILL_JOB_STATUS_QUEUED"; + case PrefillJobStatus.RUNNING: + return "PREFILL_JOB_STATUS_RUNNING"; + case PrefillJobStatus.COMPLETED: + return "PREFILL_JOB_STATUS_COMPLETED"; + case PrefillJobStatus.FAILED: + return "PREFILL_JOB_STATUS_FAILED"; + case PrefillJobStatus.CANCELLED: + return "PREFILL_JOB_STATUS_CANCELLED"; + case PrefillJobStatus.UNRECOGNIZED: + default: + return "UNRECOGNIZED"; + } +} + /** ModelCapability is one (model, role) a node offers. */ export interface ModelCapability { /** @@ -172,6 +275,8 @@ export interface NodeCapability { * LAN/Tailscale). grpc_address remains the compatibility default. */ endpoints: NodeEndpoint[]; + /** Prefill-only compute offerings. Cache-only nodes leave this empty. */ + prefillWorkers: PrefillWorkerCapability[]; } export interface NodeEndpoint { @@ -192,6 +297,14 @@ export interface CacheCompatibility { layerGeometryHash: string; kvDtype: string; blockSizeTokens: number; + /** + * Tenant namespace is part of the compatibility fingerprint. It prevents a + * snapshot from being reused across isolation domains even with identical + * model/tokenizer geometry. + */ + tenantNamespace: string; + sinkSize: number; + windowSize: number; } export interface CacheCapability { @@ -205,6 +318,21 @@ export interface CacheCapability { tokensServed: string; /** Compact probabilistic summary. Empty means "query me directly". */ bloomFilter: Uint8Array; + defaultCompression: CompressionCodec; + replicationFactor: number; +} + +export interface PrefillWorkerCapability { + compatibility?: CacheCompatibility | undefined; + workerAddress: string; + maxConcurrentJobs: number; + inflightJobs: number; + queuedJobs: number; + load: number; + tokensPerSecondPrefill: number; + ramBytesFree: string; + acceptsComputeJobs: boolean; + queuedTokens: string; } export interface ExchangeCapabilitiesRequest { @@ -285,6 +413,55 @@ export interface PublishBlockResponse { cacheEpoch: string; } +export interface SubmitPrefillJobRequest { + /** Caller-generated idempotency key. Repeating it returns the existing job. */ + requestId: string; + tenantId: string; + compatibility?: CacheCompatibility | undefined; + tokenIds: number[]; + /** + * Chained hashes computed by the primary (tenant-HMACed in authenticated + * mode). One hash per token block. + */ + blockHashes: Uint8Array[]; + deadlineMs: number; + preferredCompression: CompressionCodec; +} + +export interface SubmitPrefillJobResponse { + jobId: string; + status: PrefillJobStatus; + workerNodeId: string; + queueEtaMs: number; +} + +export interface GetPrefillJobStatusRequest { + jobId: string; + tenantId: string; +} + +export interface GetPrefillJobStatusResponse { + jobId: string; + status: PrefillJobStatus; + tokensComputed: number; + leaseId: string; + blockHash: Uint8Array; + payloadSha256: Uint8Array; + transferBytes: string; + failureReason: string; + computeMs: number; + cacheAddress: string; +} + +export interface CancelPrefillJobRequest { + jobId: string; + tenantId: string; +} + +export interface CancelPrefillJobResponse { + cancelled: boolean; +} + export interface ProposeBlockRequest { /** * The committed prefix (prompt + accepted tokens), raw token ids in @@ -542,6 +719,7 @@ function createBaseNodeCapability(): NodeCapability { ringAddress: "", caches: [], endpoints: [], + prefillWorkers: [], }; } @@ -580,6 +758,9 @@ export const NodeCapability: MessageFns = { for (const v of message.endpoints) { NodeEndpoint.encode(v!, writer.uint32(90).fork()).join(); } + for (const v of message.prefillWorkers) { + PrefillWorkerCapability.encode(v!, writer.uint32(98).fork()).join(); + } return writer; }, @@ -678,6 +859,14 @@ export const NodeCapability: MessageFns = { message.endpoints.push(NodeEndpoint.decode(reader, reader.uint32())); continue; } + case 12: { + if (tag !== 98) { + break; + } + + message.prefillWorkers.push(PrefillWorkerCapability.decode(reader, reader.uint32())); + continue; + } } if ((tag & 7) === 4 || tag === 0) { break; @@ -734,6 +923,11 @@ export const NodeCapability: MessageFns = { endpoints: globalThis.Array.isArray(object?.endpoints) ? object.endpoints.map((e: any) => NodeEndpoint.fromJSON(e)) : [], + prefillWorkers: globalThis.Array.isArray(object?.prefillWorkers) + ? object.prefillWorkers.map((e: any) => PrefillWorkerCapability.fromJSON(e)) + : globalThis.Array.isArray(object?.prefill_workers) + ? object.prefill_workers.map((e: any) => PrefillWorkerCapability.fromJSON(e)) + : [], }; }, @@ -772,6 +966,9 @@ export const NodeCapability: MessageFns = { if (message.endpoints?.length) { obj.endpoints = message.endpoints.map((e) => NodeEndpoint.toJSON(e)); } + if (message.prefillWorkers?.length) { + obj.prefillWorkers = message.prefillWorkers.map((e) => PrefillWorkerCapability.toJSON(e)); + } return obj; }, @@ -791,6 +988,7 @@ export const NodeCapability: MessageFns = { message.ringAddress = object.ringAddress ?? ""; message.caches = object.caches?.map((e) => CacheCapability.fromPartial(e)) || []; message.endpoints = object.endpoints?.map((e) => NodeEndpoint.fromPartial(e)) || []; + message.prefillWorkers = object.prefillWorkers?.map((e) => PrefillWorkerCapability.fromPartial(e)) || []; return message; }, }; @@ -918,6 +1116,9 @@ function createBaseCacheCompatibility(): CacheCompatibility { layerGeometryHash: "", kvDtype: "", blockSizeTokens: 0, + tenantNamespace: "", + sinkSize: 0, + windowSize: 0, }; } @@ -950,6 +1151,15 @@ export const CacheCompatibility: MessageFns = { if (message.blockSizeTokens !== 0) { writer.uint32(72).uint32(message.blockSizeTokens); } + if (message.tenantNamespace !== "") { + writer.uint32(82).string(message.tenantNamespace); + } + if (message.sinkSize !== 0) { + writer.uint32(88).uint32(message.sinkSize); + } + if (message.windowSize !== 0) { + writer.uint32(96).uint32(message.windowSize); + } return writer; }, @@ -1032,6 +1242,30 @@ export const CacheCompatibility: MessageFns = { message.blockSizeTokens = reader.uint32(); continue; } + case 10: { + if (tag !== 82) { + break; + } + + message.tenantNamespace = reader.string(); + continue; + } + case 11: { + if (tag !== 88) { + break; + } + + message.sinkSize = reader.uint32(); + continue; + } + case 12: { + if (tag !== 96) { + break; + } + + message.windowSize = reader.uint32(); + continue; + } } if ((tag & 7) === 4 || tag === 0) { break; @@ -1084,6 +1318,21 @@ export const CacheCompatibility: MessageFns = { : isSet(object.block_size_tokens) ? globalThis.Number(object.block_size_tokens) : 0, + tenantNamespace: isSet(object.tenantNamespace) + ? globalThis.String(object.tenantNamespace) + : isSet(object.tenant_namespace) + ? globalThis.String(object.tenant_namespace) + : "", + sinkSize: isSet(object.sinkSize) + ? globalThis.Number(object.sinkSize) + : isSet(object.sink_size) + ? globalThis.Number(object.sink_size) + : 0, + windowSize: isSet(object.windowSize) + ? globalThis.Number(object.windowSize) + : isSet(object.window_size) + ? globalThis.Number(object.window_size) + : 0, }; }, @@ -1116,6 +1365,15 @@ export const CacheCompatibility: MessageFns = { if (message.blockSizeTokens !== 0) { obj.blockSizeTokens = Math.round(message.blockSizeTokens); } + if (message.tenantNamespace !== "") { + obj.tenantNamespace = message.tenantNamespace; + } + if (message.sinkSize !== 0) { + obj.sinkSize = Math.round(message.sinkSize); + } + if (message.windowSize !== 0) { + obj.windowSize = Math.round(message.windowSize); + } return obj; }, @@ -1133,6 +1391,9 @@ export const CacheCompatibility: MessageFns = { message.layerGeometryHash = object.layerGeometryHash ?? ""; message.kvDtype = object.kvDtype ?? ""; message.blockSizeTokens = object.blockSizeTokens ?? 0; + message.tenantNamespace = object.tenantNamespace ?? ""; + message.sinkSize = object.sinkSize ?? 0; + message.windowSize = object.windowSize ?? 0; return message; }, }; @@ -1148,6 +1409,8 @@ function createBaseCacheCapability(): CacheCapability { load: 0, tokensServed: "0", bloomFilter: new Uint8Array(0), + defaultCompression: 0, + replicationFactor: 0, }; } @@ -1180,6 +1443,12 @@ export const CacheCapability: MessageFns = { if (message.bloomFilter.length !== 0) { writer.uint32(74).bytes(message.bloomFilter); } + if (message.defaultCompression !== 0) { + writer.uint32(80).int32(message.defaultCompression); + } + if (message.replicationFactor !== 0) { + writer.uint32(88).uint32(message.replicationFactor); + } return writer; }, @@ -1262,6 +1531,22 @@ export const CacheCapability: MessageFns = { message.bloomFilter = reader.bytes(); continue; } + case 10: { + if (tag !== 80) { + break; + } + + message.defaultCompression = reader.int32() as any; + continue; + } + case 11: { + if (tag !== 88) { + break; + } + + message.replicationFactor = reader.uint32(); + continue; + } } if ((tag & 7) === 4 || tag === 0) { break; @@ -1310,6 +1595,16 @@ export const CacheCapability: MessageFns = { : isSet(object.bloom_filter) ? bytesFromBase64(object.bloom_filter) : new Uint8Array(0), + defaultCompression: isSet(object.defaultCompression) + ? compressionCodecFromJSON(object.defaultCompression) + : isSet(object.default_compression) + ? compressionCodecFromJSON(object.default_compression) + : 0, + replicationFactor: isSet(object.replicationFactor) + ? globalThis.Number(object.replicationFactor) + : isSet(object.replication_factor) + ? globalThis.Number(object.replication_factor) + : 0, }; }, @@ -1342,6 +1637,12 @@ export const CacheCapability: MessageFns = { if (message.bloomFilter.length !== 0) { obj.bloomFilter = base64FromBytes(message.bloomFilter); } + if (message.defaultCompression !== 0) { + obj.defaultCompression = compressionCodecToJSON(message.defaultCompression); + } + if (message.replicationFactor !== 0) { + obj.replicationFactor = Math.round(message.replicationFactor); + } return obj; }, @@ -1361,26 +1662,66 @@ export const CacheCapability: MessageFns = { message.load = object.load ?? 0; message.tokensServed = object.tokensServed ?? "0"; message.bloomFilter = object.bloomFilter ?? new Uint8Array(0); + message.defaultCompression = object.defaultCompression ?? 0; + message.replicationFactor = object.replicationFactor ?? 0; return message; }, }; -function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest { - return { knownNodes: [] }; +function createBasePrefillWorkerCapability(): PrefillWorkerCapability { + return { + compatibility: undefined, + workerAddress: "", + maxConcurrentJobs: 0, + inflightJobs: 0, + queuedJobs: 0, + load: 0, + tokensPerSecondPrefill: 0, + ramBytesFree: "0", + acceptsComputeJobs: false, + queuedTokens: "0", + }; } -export const ExchangeCapabilitiesRequest: MessageFns = { - encode(message: ExchangeCapabilitiesRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { - for (const v of message.knownNodes) { - NodeCapability.encode(v!, writer.uint32(10).fork()).join(); +export const PrefillWorkerCapability: MessageFns = { + encode(message: PrefillWorkerCapability, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(10).fork()).join(); + } + if (message.workerAddress !== "") { + writer.uint32(18).string(message.workerAddress); + } + if (message.maxConcurrentJobs !== 0) { + writer.uint32(24).uint32(message.maxConcurrentJobs); + } + if (message.inflightJobs !== 0) { + writer.uint32(32).uint32(message.inflightJobs); + } + if (message.queuedJobs !== 0) { + writer.uint32(40).uint32(message.queuedJobs); + } + if (message.load !== 0) { + writer.uint32(49).double(message.load); + } + if (message.tokensPerSecondPrefill !== 0) { + writer.uint32(57).double(message.tokensPerSecondPrefill); + } + if (message.ramBytesFree !== "0") { + writer.uint32(64).uint64(message.ramBytesFree); + } + if (message.acceptsComputeJobs !== false) { + writer.uint32(72).bool(message.acceptsComputeJobs); + } + if (message.queuedTokens !== "0") { + writer.uint32(80).uint64(message.queuedTokens); } return writer; }, - decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesRequest { + decode(input: BinaryReader | Uint8Array, length?: number): PrefillWorkerCapability { const reader = input instanceof BinaryReader ? input : new BinaryReader(input); const end = length === undefined ? reader.len : reader.pos + length; - const message = createBaseExchangeCapabilitiesRequest(); + const message = createBasePrefillWorkerCapability(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { @@ -1389,7 +1730,79 @@ export const ExchangeCapabilitiesRequest: MessageFns NodeCapability.fromJSON(e)) - : globalThis.Array.isArray(object?.known_nodes) - ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e)) - : [], + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + workerAddress: isSet(object.workerAddress) + ? globalThis.String(object.workerAddress) + : isSet(object.worker_address) + ? globalThis.String(object.worker_address) + : "", + maxConcurrentJobs: isSet(object.maxConcurrentJobs) + ? globalThis.Number(object.maxConcurrentJobs) + : isSet(object.max_concurrent_jobs) + ? globalThis.Number(object.max_concurrent_jobs) + : 0, + inflightJobs: isSet(object.inflightJobs) + ? globalThis.Number(object.inflightJobs) + : isSet(object.inflight_jobs) + ? globalThis.Number(object.inflight_jobs) + : 0, + queuedJobs: isSet(object.queuedJobs) + ? globalThis.Number(object.queuedJobs) + : isSet(object.queued_jobs) + ? globalThis.Number(object.queued_jobs) + : 0, + load: isSet(object.load) ? globalThis.Number(object.load) : 0, + tokensPerSecondPrefill: isSet(object.tokensPerSecondPrefill) + ? globalThis.Number(object.tokensPerSecondPrefill) + : isSet(object.tokens_per_second_prefill) + ? globalThis.Number(object.tokens_per_second_prefill) + : 0, + ramBytesFree: isSet(object.ramBytesFree) + ? globalThis.String(object.ramBytesFree) + : isSet(object.ram_bytes_free) + ? globalThis.String(object.ram_bytes_free) + : "0", + acceptsComputeJobs: isSet(object.acceptsComputeJobs) + ? globalThis.Boolean(object.acceptsComputeJobs) + : isSet(object.accepts_compute_jobs) + ? globalThis.Boolean(object.accepts_compute_jobs) + : false, + queuedTokens: isSet(object.queuedTokens) + ? globalThis.String(object.queuedTokens) + : isSet(object.queued_tokens) + ? globalThis.String(object.queued_tokens) + : "0", }; }, - toJSON(message: ExchangeCapabilitiesRequest): unknown { + toJSON(message: PrefillWorkerCapability): unknown { const obj: any = {}; - if (message.knownNodes?.length) { - obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e)); + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + if (message.workerAddress !== "") { + obj.workerAddress = message.workerAddress; + } + if (message.maxConcurrentJobs !== 0) { + obj.maxConcurrentJobs = Math.round(message.maxConcurrentJobs); + } + if (message.inflightJobs !== 0) { + obj.inflightJobs = Math.round(message.inflightJobs); + } + if (message.queuedJobs !== 0) { + obj.queuedJobs = Math.round(message.queuedJobs); + } + if (message.load !== 0) { + obj.load = message.load; + } + if (message.tokensPerSecondPrefill !== 0) { + obj.tokensPerSecondPrefill = message.tokensPerSecondPrefill; + } + if (message.ramBytesFree !== "0") { + obj.ramBytesFree = message.ramBytesFree; + } + if (message.acceptsComputeJobs !== false) { + obj.acceptsComputeJobs = message.acceptsComputeJobs; + } + if (message.queuedTokens !== "0") { + obj.queuedTokens = message.queuedTokens; } return obj; }, - create, I>>(base?: I): ExchangeCapabilitiesRequest { - return ExchangeCapabilitiesRequest.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): PrefillWorkerCapability { + return PrefillWorkerCapability.fromPartial(base ?? ({} as any)); }, - fromPartial, I>>(object: I): ExchangeCapabilitiesRequest { - const message = createBaseExchangeCapabilitiesRequest(); - message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || []; + fromPartial, I>>(object: I): PrefillWorkerCapability { + const message = createBasePrefillWorkerCapability(); + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + message.workerAddress = object.workerAddress ?? ""; + message.maxConcurrentJobs = object.maxConcurrentJobs ?? 0; + message.inflightJobs = object.inflightJobs ?? 0; + message.queuedJobs = object.queuedJobs ?? 0; + message.load = object.load ?? 0; + message.tokensPerSecondPrefill = object.tokensPerSecondPrefill ?? 0; + message.ramBytesFree = object.ramBytesFree ?? "0"; + message.acceptsComputeJobs = object.acceptsComputeJobs ?? false; + message.queuedTokens = object.queuedTokens ?? "0"; return message; }, }; -function createBaseExchangeCapabilitiesResponse(): ExchangeCapabilitiesResponse { +function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest { return { knownNodes: [] }; } -export const ExchangeCapabilitiesResponse: MessageFns = { +export const ExchangeCapabilitiesRequest: MessageFns = { + encode(message: ExchangeCapabilitiesRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + for (const v of message.knownNodes) { + NodeCapability.encode(v!, writer.uint32(10).fork()).join(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseExchangeCapabilitiesRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.knownNodes.push(NodeCapability.decode(reader, reader.uint32())); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): ExchangeCapabilitiesRequest { + return { + knownNodes: globalThis.Array.isArray(object?.knownNodes) + ? object.knownNodes.map((e: any) => NodeCapability.fromJSON(e)) + : globalThis.Array.isArray(object?.known_nodes) + ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e)) + : [], + }; + }, + + toJSON(message: ExchangeCapabilitiesRequest): unknown { + const obj: any = {}; + if (message.knownNodes?.length) { + obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e)); + } + return obj; + }, + + create, I>>(base?: I): ExchangeCapabilitiesRequest { + return ExchangeCapabilitiesRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): ExchangeCapabilitiesRequest { + const message = createBaseExchangeCapabilitiesRequest(); + message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || []; + return message; + }, +}; + +function createBaseExchangeCapabilitiesResponse(): ExchangeCapabilitiesResponse { + return { knownNodes: [] }; +} + +export const ExchangeCapabilitiesResponse: MessageFns = { encode(message: ExchangeCapabilitiesResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { for (const v of message.knownNodes) { NodeCapability.encode(v!, writer.uint32(10).fork()).join(); @@ -2339,25 +2891,743 @@ export const PublishBlockRequest: MessageFns = { if (message.totalChunks !== 0) { writer.uint32(40).uint32(message.totalChunks); } - if (message.data.length !== 0) { - writer.uint32(50).bytes(message.data); + if (message.data.length !== 0) { + writer.uint32(50).bytes(message.data); + } + if (message.blockSha256.length !== 0) { + writer.uint32(58).bytes(message.blockSha256); + } + if (message.cacheEpoch !== "0") { + writer.uint32(64).uint64(message.cacheEpoch); + } + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(74).fork()).join(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBasePublishBlockRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.blockHash = reader.bytes(); + continue; + } + case 2: { + if (tag !== 16) { + break; + } + + message.blockIndex = reader.uint32(); + continue; + } + case 3: { + if (tag !== 24) { + break; + } + + message.tokenCount = reader.uint32(); + continue; + } + case 4: { + if (tag !== 32) { + break; + } + + message.chunkIndex = reader.uint32(); + continue; + } + case 5: { + if (tag !== 40) { + break; + } + + message.totalChunks = reader.uint32(); + continue; + } + case 6: { + if (tag !== 50) { + break; + } + + message.data = reader.bytes(); + continue; + } + case 7: { + if (tag !== 58) { + break; + } + + message.blockSha256 = reader.bytes(); + continue; + } + case 8: { + if (tag !== 64) { + break; + } + + message.cacheEpoch = reader.uint64().toString(); + continue; + } + case 9: { + if (tag !== 74) { + break; + } + + message.compatibility = CacheCompatibility.decode(reader, reader.uint32()); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): PublishBlockRequest { + return { + blockHash: isSet(object.blockHash) + ? bytesFromBase64(object.blockHash) + : isSet(object.block_hash) + ? bytesFromBase64(object.block_hash) + : new Uint8Array(0), + blockIndex: isSet(object.blockIndex) + ? globalThis.Number(object.blockIndex) + : isSet(object.block_index) + ? globalThis.Number(object.block_index) + : 0, + tokenCount: isSet(object.tokenCount) + ? globalThis.Number(object.tokenCount) + : isSet(object.token_count) + ? globalThis.Number(object.token_count) + : 0, + chunkIndex: isSet(object.chunkIndex) + ? globalThis.Number(object.chunkIndex) + : isSet(object.chunk_index) + ? globalThis.Number(object.chunk_index) + : 0, + totalChunks: isSet(object.totalChunks) + ? globalThis.Number(object.totalChunks) + : isSet(object.total_chunks) + ? globalThis.Number(object.total_chunks) + : 0, + data: isSet(object.data) ? bytesFromBase64(object.data) : new Uint8Array(0), + blockSha256: isSet(object.blockSha256) + ? bytesFromBase64(object.blockSha256) + : isSet(object.block_sha256) + ? bytesFromBase64(object.block_sha256) + : new Uint8Array(0), + cacheEpoch: isSet(object.cacheEpoch) + ? globalThis.String(object.cacheEpoch) + : isSet(object.cache_epoch) + ? globalThis.String(object.cache_epoch) + : "0", + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + }; + }, + + toJSON(message: PublishBlockRequest): unknown { + const obj: any = {}; + if (message.blockHash.length !== 0) { + obj.blockHash = base64FromBytes(message.blockHash); + } + if (message.blockIndex !== 0) { + obj.blockIndex = Math.round(message.blockIndex); + } + if (message.tokenCount !== 0) { + obj.tokenCount = Math.round(message.tokenCount); + } + if (message.chunkIndex !== 0) { + obj.chunkIndex = Math.round(message.chunkIndex); + } + if (message.totalChunks !== 0) { + obj.totalChunks = Math.round(message.totalChunks); + } + if (message.data.length !== 0) { + obj.data = base64FromBytes(message.data); + } + if (message.blockSha256.length !== 0) { + obj.blockSha256 = base64FromBytes(message.blockSha256); + } + if (message.cacheEpoch !== "0") { + obj.cacheEpoch = message.cacheEpoch; + } + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + return obj; + }, + + create, I>>(base?: I): PublishBlockRequest { + return PublishBlockRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): PublishBlockRequest { + const message = createBasePublishBlockRequest(); + message.blockHash = object.blockHash ?? new Uint8Array(0); + message.blockIndex = object.blockIndex ?? 0; + message.tokenCount = object.tokenCount ?? 0; + message.chunkIndex = object.chunkIndex ?? 0; + message.totalChunks = object.totalChunks ?? 0; + message.data = object.data ?? new Uint8Array(0); + message.blockSha256 = object.blockSha256 ?? new Uint8Array(0); + message.cacheEpoch = object.cacheEpoch ?? "0"; + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + return message; + }, +}; + +function createBasePublishBlockResponse(): PublishBlockResponse { + return { stored: false, cacheEpoch: "0" }; +} + +export const PublishBlockResponse: MessageFns = { + encode(message: PublishBlockResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.stored !== false) { + writer.uint32(8).bool(message.stored); + } + if (message.cacheEpoch !== "0") { + writer.uint32(16).uint64(message.cacheEpoch); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBasePublishBlockResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 8) { + break; + } + + message.stored = reader.bool(); + continue; + } + case 2: { + if (tag !== 16) { + break; + } + + message.cacheEpoch = reader.uint64().toString(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): PublishBlockResponse { + return { + stored: isSet(object.stored) ? globalThis.Boolean(object.stored) : false, + cacheEpoch: isSet(object.cacheEpoch) + ? globalThis.String(object.cacheEpoch) + : isSet(object.cache_epoch) + ? globalThis.String(object.cache_epoch) + : "0", + }; + }, + + toJSON(message: PublishBlockResponse): unknown { + const obj: any = {}; + if (message.stored !== false) { + obj.stored = message.stored; + } + if (message.cacheEpoch !== "0") { + obj.cacheEpoch = message.cacheEpoch; + } + return obj; + }, + + create, I>>(base?: I): PublishBlockResponse { + return PublishBlockResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): PublishBlockResponse { + const message = createBasePublishBlockResponse(); + message.stored = object.stored ?? false; + message.cacheEpoch = object.cacheEpoch ?? "0"; + return message; + }, +}; + +function createBaseSubmitPrefillJobRequest(): SubmitPrefillJobRequest { + return { + requestId: "", + tenantId: "", + compatibility: undefined, + tokenIds: [], + blockHashes: [], + deadlineMs: 0, + preferredCompression: 0, + }; +} + +export const SubmitPrefillJobRequest: MessageFns = { + encode(message: SubmitPrefillJobRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.requestId !== "") { + writer.uint32(10).string(message.requestId); + } + if (message.tenantId !== "") { + writer.uint32(18).string(message.tenantId); + } + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(26).fork()).join(); + } + writer.uint32(34).fork(); + for (const v of message.tokenIds) { + writer.uint32(v); + } + writer.join(); + for (const v of message.blockHashes) { + writer.uint32(42).bytes(v!); + } + if (message.deadlineMs !== 0) { + writer.uint32(48).uint32(message.deadlineMs); + } + if (message.preferredCompression !== 0) { + writer.uint32(56).int32(message.preferredCompression); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): SubmitPrefillJobRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseSubmitPrefillJobRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.requestId = reader.string(); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.tenantId = reader.string(); + continue; + } + case 3: { + if (tag !== 26) { + break; + } + + message.compatibility = CacheCompatibility.decode(reader, reader.uint32()); + continue; + } + case 4: { + if (tag === 32) { + message.tokenIds.push(reader.uint32()); + + continue; + } + + if (tag === 34) { + const end2 = reader.uint32() + reader.pos; + while (reader.pos < end2) { + message.tokenIds.push(reader.uint32()); + } + + continue; + } + + break; + } + case 5: { + if (tag !== 42) { + break; + } + + message.blockHashes.push(reader.bytes()); + continue; + } + case 6: { + if (tag !== 48) { + break; + } + + message.deadlineMs = reader.uint32(); + continue; + } + case 7: { + if (tag !== 56) { + break; + } + + message.preferredCompression = reader.int32() as any; + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): SubmitPrefillJobRequest { + return { + requestId: isSet(object.requestId) + ? globalThis.String(object.requestId) + : isSet(object.request_id) + ? globalThis.String(object.request_id) + : "", + tenantId: isSet(object.tenantId) + ? globalThis.String(object.tenantId) + : isSet(object.tenant_id) + ? globalThis.String(object.tenant_id) + : "", + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + tokenIds: globalThis.Array.isArray(object?.tokenIds) + ? object.tokenIds.map((e: any) => globalThis.Number(e)) + : globalThis.Array.isArray(object?.token_ids) + ? object.token_ids.map((e: any) => globalThis.Number(e)) + : [], + blockHashes: globalThis.Array.isArray(object?.blockHashes) + ? object.blockHashes.map((e: any) => bytesFromBase64(e)) + : globalThis.Array.isArray(object?.block_hashes) + ? object.block_hashes.map((e: any) => bytesFromBase64(e)) + : [], + deadlineMs: isSet(object.deadlineMs) + ? globalThis.Number(object.deadlineMs) + : isSet(object.deadline_ms) + ? globalThis.Number(object.deadline_ms) + : 0, + preferredCompression: isSet(object.preferredCompression) + ? compressionCodecFromJSON(object.preferredCompression) + : isSet(object.preferred_compression) + ? compressionCodecFromJSON(object.preferred_compression) + : 0, + }; + }, + + toJSON(message: SubmitPrefillJobRequest): unknown { + const obj: any = {}; + if (message.requestId !== "") { + obj.requestId = message.requestId; + } + if (message.tenantId !== "") { + obj.tenantId = message.tenantId; + } + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + if (message.tokenIds?.length) { + obj.tokenIds = message.tokenIds.map((e) => Math.round(e)); + } + if (message.blockHashes?.length) { + obj.blockHashes = message.blockHashes.map((e) => base64FromBytes(e)); + } + if (message.deadlineMs !== 0) { + obj.deadlineMs = Math.round(message.deadlineMs); + } + if (message.preferredCompression !== 0) { + obj.preferredCompression = compressionCodecToJSON(message.preferredCompression); + } + return obj; + }, + + create, I>>(base?: I): SubmitPrefillJobRequest { + return SubmitPrefillJobRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): SubmitPrefillJobRequest { + const message = createBaseSubmitPrefillJobRequest(); + message.requestId = object.requestId ?? ""; + message.tenantId = object.tenantId ?? ""; + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + message.tokenIds = object.tokenIds?.map((e) => e) || []; + message.blockHashes = object.blockHashes?.map((e) => e) || []; + message.deadlineMs = object.deadlineMs ?? 0; + message.preferredCompression = object.preferredCompression ?? 0; + return message; + }, +}; + +function createBaseSubmitPrefillJobResponse(): SubmitPrefillJobResponse { + return { jobId: "", status: 0, workerNodeId: "", queueEtaMs: 0 }; +} + +export const SubmitPrefillJobResponse: MessageFns = { + encode(message: SubmitPrefillJobResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.jobId !== "") { + writer.uint32(10).string(message.jobId); + } + if (message.status !== 0) { + writer.uint32(16).int32(message.status); + } + if (message.workerNodeId !== "") { + writer.uint32(26).string(message.workerNodeId); + } + if (message.queueEtaMs !== 0) { + writer.uint32(33).double(message.queueEtaMs); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): SubmitPrefillJobResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseSubmitPrefillJobResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.jobId = reader.string(); + continue; + } + case 2: { + if (tag !== 16) { + break; + } + + message.status = reader.int32() as any; + continue; + } + case 3: { + if (tag !== 26) { + break; + } + + message.workerNodeId = reader.string(); + continue; + } + case 4: { + if (tag !== 33) { + break; + } + + message.queueEtaMs = reader.double(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): SubmitPrefillJobResponse { + return { + jobId: isSet(object.jobId) + ? globalThis.String(object.jobId) + : isSet(object.job_id) + ? globalThis.String(object.job_id) + : "", + status: isSet(object.status) ? prefillJobStatusFromJSON(object.status) : 0, + workerNodeId: isSet(object.workerNodeId) + ? globalThis.String(object.workerNodeId) + : isSet(object.worker_node_id) + ? globalThis.String(object.worker_node_id) + : "", + queueEtaMs: isSet(object.queueEtaMs) + ? globalThis.Number(object.queueEtaMs) + : isSet(object.queue_eta_ms) + ? globalThis.Number(object.queue_eta_ms) + : 0, + }; + }, + + toJSON(message: SubmitPrefillJobResponse): unknown { + const obj: any = {}; + if (message.jobId !== "") { + obj.jobId = message.jobId; + } + if (message.status !== 0) { + obj.status = prefillJobStatusToJSON(message.status); + } + if (message.workerNodeId !== "") { + obj.workerNodeId = message.workerNodeId; + } + if (message.queueEtaMs !== 0) { + obj.queueEtaMs = message.queueEtaMs; + } + return obj; + }, + + create, I>>(base?: I): SubmitPrefillJobResponse { + return SubmitPrefillJobResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): SubmitPrefillJobResponse { + const message = createBaseSubmitPrefillJobResponse(); + message.jobId = object.jobId ?? ""; + message.status = object.status ?? 0; + message.workerNodeId = object.workerNodeId ?? ""; + message.queueEtaMs = object.queueEtaMs ?? 0; + return message; + }, +}; + +function createBaseGetPrefillJobStatusRequest(): GetPrefillJobStatusRequest { + return { jobId: "", tenantId: "" }; +} + +export const GetPrefillJobStatusRequest: MessageFns = { + encode(message: GetPrefillJobStatusRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.jobId !== "") { + writer.uint32(10).string(message.jobId); + } + if (message.tenantId !== "") { + writer.uint32(18).string(message.tenantId); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): GetPrefillJobStatusRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseGetPrefillJobStatusRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.jobId = reader.string(); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.tenantId = reader.string(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): GetPrefillJobStatusRequest { + return { + jobId: isSet(object.jobId) + ? globalThis.String(object.jobId) + : isSet(object.job_id) + ? globalThis.String(object.job_id) + : "", + tenantId: isSet(object.tenantId) + ? globalThis.String(object.tenantId) + : isSet(object.tenant_id) + ? globalThis.String(object.tenant_id) + : "", + }; + }, + + toJSON(message: GetPrefillJobStatusRequest): unknown { + const obj: any = {}; + if (message.jobId !== "") { + obj.jobId = message.jobId; + } + if (message.tenantId !== "") { + obj.tenantId = message.tenantId; + } + return obj; + }, + + create, I>>(base?: I): GetPrefillJobStatusRequest { + return GetPrefillJobStatusRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): GetPrefillJobStatusRequest { + const message = createBaseGetPrefillJobStatusRequest(); + message.jobId = object.jobId ?? ""; + message.tenantId = object.tenantId ?? ""; + return message; + }, +}; + +function createBaseGetPrefillJobStatusResponse(): GetPrefillJobStatusResponse { + return { + jobId: "", + status: 0, + tokensComputed: 0, + leaseId: "", + blockHash: new Uint8Array(0), + payloadSha256: new Uint8Array(0), + transferBytes: "0", + failureReason: "", + computeMs: 0, + cacheAddress: "", + }; +} + +export const GetPrefillJobStatusResponse: MessageFns = { + encode(message: GetPrefillJobStatusResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.jobId !== "") { + writer.uint32(10).string(message.jobId); + } + if (message.status !== 0) { + writer.uint32(16).int32(message.status); + } + if (message.tokensComputed !== 0) { + writer.uint32(24).uint32(message.tokensComputed); + } + if (message.leaseId !== "") { + writer.uint32(34).string(message.leaseId); + } + if (message.blockHash.length !== 0) { + writer.uint32(42).bytes(message.blockHash); + } + if (message.payloadSha256.length !== 0) { + writer.uint32(50).bytes(message.payloadSha256); } - if (message.blockSha256.length !== 0) { - writer.uint32(58).bytes(message.blockSha256); + if (message.transferBytes !== "0") { + writer.uint32(56).uint64(message.transferBytes); } - if (message.cacheEpoch !== "0") { - writer.uint32(64).uint64(message.cacheEpoch); + if (message.failureReason !== "") { + writer.uint32(66).string(message.failureReason); } - if (message.compatibility !== undefined) { - CacheCompatibility.encode(message.compatibility, writer.uint32(74).fork()).join(); + if (message.computeMs !== 0) { + writer.uint32(73).double(message.computeMs); + } + if (message.cacheAddress !== "") { + writer.uint32(82).string(message.cacheAddress); } return writer; }, - decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockRequest { + decode(input: BinaryReader | Uint8Array, length?: number): GetPrefillJobStatusResponse { const reader = input instanceof BinaryReader ? input : new BinaryReader(input); const end = length === undefined ? reader.len : reader.pos + length; - const message = createBasePublishBlockRequest(); + const message = createBaseGetPrefillJobStatusResponse(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { @@ -2366,7 +3636,7 @@ export const PublishBlockRequest: MessageFns = { break; } - message.blockHash = reader.bytes(); + message.jobId = reader.string(); continue; } case 2: { @@ -2374,7 +3644,7 @@ export const PublishBlockRequest: MessageFns = { break; } - message.blockIndex = reader.uint32(); + message.status = reader.int32() as any; continue; } case 3: { @@ -2382,23 +3652,23 @@ export const PublishBlockRequest: MessageFns = { break; } - message.tokenCount = reader.uint32(); + message.tokensComputed = reader.uint32(); continue; } case 4: { - if (tag !== 32) { + if (tag !== 34) { break; } - message.chunkIndex = reader.uint32(); + message.leaseId = reader.string(); continue; } case 5: { - if (tag !== 40) { + if (tag !== 42) { break; } - message.totalChunks = reader.uint32(); + message.blockHash = reader.bytes(); continue; } case 6: { @@ -2406,31 +3676,39 @@ export const PublishBlockRequest: MessageFns = { break; } - message.data = reader.bytes(); + message.payloadSha256 = reader.bytes(); continue; } case 7: { - if (tag !== 58) { + if (tag !== 56) { break; } - message.blockSha256 = reader.bytes(); + message.transferBytes = reader.uint64().toString(); continue; } case 8: { - if (tag !== 64) { + if (tag !== 66) { break; } - message.cacheEpoch = reader.uint64().toString(); + message.failureReason = reader.string(); continue; } case 9: { - if (tag !== 74) { + if (tag !== 73) { break; } - message.compatibility = CacheCompatibility.decode(reader, reader.uint32()); + message.computeMs = reader.double(); + continue; + } + case 10: { + if (tag !== 82) { + break; + } + + message.cacheAddress = reader.string(); continue; } } @@ -2442,136 +3720,147 @@ export const PublishBlockRequest: MessageFns = { return message; }, - fromJSON(object: any): PublishBlockRequest { + fromJSON(object: any): GetPrefillJobStatusResponse { return { + jobId: isSet(object.jobId) + ? globalThis.String(object.jobId) + : isSet(object.job_id) + ? globalThis.String(object.job_id) + : "", + status: isSet(object.status) ? prefillJobStatusFromJSON(object.status) : 0, + tokensComputed: isSet(object.tokensComputed) + ? globalThis.Number(object.tokensComputed) + : isSet(object.tokens_computed) + ? globalThis.Number(object.tokens_computed) + : 0, + leaseId: isSet(object.leaseId) + ? globalThis.String(object.leaseId) + : isSet(object.lease_id) + ? globalThis.String(object.lease_id) + : "", blockHash: isSet(object.blockHash) ? bytesFromBase64(object.blockHash) : isSet(object.block_hash) ? bytesFromBase64(object.block_hash) : new Uint8Array(0), - blockIndex: isSet(object.blockIndex) - ? globalThis.Number(object.blockIndex) - : isSet(object.block_index) - ? globalThis.Number(object.block_index) - : 0, - tokenCount: isSet(object.tokenCount) - ? globalThis.Number(object.tokenCount) - : isSet(object.token_count) - ? globalThis.Number(object.token_count) - : 0, - chunkIndex: isSet(object.chunkIndex) - ? globalThis.Number(object.chunkIndex) - : isSet(object.chunk_index) - ? globalThis.Number(object.chunk_index) - : 0, - totalChunks: isSet(object.totalChunks) - ? globalThis.Number(object.totalChunks) - : isSet(object.total_chunks) - ? globalThis.Number(object.total_chunks) - : 0, - data: isSet(object.data) ? bytesFromBase64(object.data) : new Uint8Array(0), - blockSha256: isSet(object.blockSha256) - ? bytesFromBase64(object.blockSha256) - : isSet(object.block_sha256) - ? bytesFromBase64(object.block_sha256) + payloadSha256: isSet(object.payloadSha256) + ? bytesFromBase64(object.payloadSha256) + : isSet(object.payload_sha256) + ? bytesFromBase64(object.payload_sha256) : new Uint8Array(0), - cacheEpoch: isSet(object.cacheEpoch) - ? globalThis.String(object.cacheEpoch) - : isSet(object.cache_epoch) - ? globalThis.String(object.cache_epoch) + transferBytes: isSet(object.transferBytes) + ? globalThis.String(object.transferBytes) + : isSet(object.transfer_bytes) + ? globalThis.String(object.transfer_bytes) : "0", - compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + failureReason: isSet(object.failureReason) + ? globalThis.String(object.failureReason) + : isSet(object.failure_reason) + ? globalThis.String(object.failure_reason) + : "", + computeMs: isSet(object.computeMs) + ? globalThis.Number(object.computeMs) + : isSet(object.compute_ms) + ? globalThis.Number(object.compute_ms) + : 0, + cacheAddress: isSet(object.cacheAddress) + ? globalThis.String(object.cacheAddress) + : isSet(object.cache_address) + ? globalThis.String(object.cache_address) + : "", }; }, - toJSON(message: PublishBlockRequest): unknown { + toJSON(message: GetPrefillJobStatusResponse): unknown { const obj: any = {}; - if (message.blockHash.length !== 0) { - obj.blockHash = base64FromBytes(message.blockHash); + if (message.jobId !== "") { + obj.jobId = message.jobId; } - if (message.blockIndex !== 0) { - obj.blockIndex = Math.round(message.blockIndex); + if (message.status !== 0) { + obj.status = prefillJobStatusToJSON(message.status); } - if (message.tokenCount !== 0) { - obj.tokenCount = Math.round(message.tokenCount); + if (message.tokensComputed !== 0) { + obj.tokensComputed = Math.round(message.tokensComputed); } - if (message.chunkIndex !== 0) { - obj.chunkIndex = Math.round(message.chunkIndex); + if (message.leaseId !== "") { + obj.leaseId = message.leaseId; } - if (message.totalChunks !== 0) { - obj.totalChunks = Math.round(message.totalChunks); + if (message.blockHash.length !== 0) { + obj.blockHash = base64FromBytes(message.blockHash); } - if (message.data.length !== 0) { - obj.data = base64FromBytes(message.data); + if (message.payloadSha256.length !== 0) { + obj.payloadSha256 = base64FromBytes(message.payloadSha256); } - if (message.blockSha256.length !== 0) { - obj.blockSha256 = base64FromBytes(message.blockSha256); + if (message.transferBytes !== "0") { + obj.transferBytes = message.transferBytes; } - if (message.cacheEpoch !== "0") { - obj.cacheEpoch = message.cacheEpoch; + if (message.failureReason !== "") { + obj.failureReason = message.failureReason; } - if (message.compatibility !== undefined) { - obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + if (message.computeMs !== 0) { + obj.computeMs = message.computeMs; + } + if (message.cacheAddress !== "") { + obj.cacheAddress = message.cacheAddress; } return obj; }, - create, I>>(base?: I): PublishBlockRequest { - return PublishBlockRequest.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): GetPrefillJobStatusResponse { + return GetPrefillJobStatusResponse.fromPartial(base ?? ({} as any)); }, - fromPartial, I>>(object: I): PublishBlockRequest { - const message = createBasePublishBlockRequest(); + fromPartial, I>>(object: I): GetPrefillJobStatusResponse { + const message = createBaseGetPrefillJobStatusResponse(); + message.jobId = object.jobId ?? ""; + message.status = object.status ?? 0; + message.tokensComputed = object.tokensComputed ?? 0; + message.leaseId = object.leaseId ?? ""; message.blockHash = object.blockHash ?? new Uint8Array(0); - message.blockIndex = object.blockIndex ?? 0; - message.tokenCount = object.tokenCount ?? 0; - message.chunkIndex = object.chunkIndex ?? 0; - message.totalChunks = object.totalChunks ?? 0; - message.data = object.data ?? new Uint8Array(0); - message.blockSha256 = object.blockSha256 ?? new Uint8Array(0); - message.cacheEpoch = object.cacheEpoch ?? "0"; - message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) - ? CacheCompatibility.fromPartial(object.compatibility) - : undefined; + message.payloadSha256 = object.payloadSha256 ?? new Uint8Array(0); + message.transferBytes = object.transferBytes ?? "0"; + message.failureReason = object.failureReason ?? ""; + message.computeMs = object.computeMs ?? 0; + message.cacheAddress = object.cacheAddress ?? ""; return message; }, }; -function createBasePublishBlockResponse(): PublishBlockResponse { - return { stored: false, cacheEpoch: "0" }; +function createBaseCancelPrefillJobRequest(): CancelPrefillJobRequest { + return { jobId: "", tenantId: "" }; } -export const PublishBlockResponse: MessageFns = { - encode(message: PublishBlockResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { - if (message.stored !== false) { - writer.uint32(8).bool(message.stored); +export const CancelPrefillJobRequest: MessageFns = { + encode(message: CancelPrefillJobRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.jobId !== "") { + writer.uint32(10).string(message.jobId); } - if (message.cacheEpoch !== "0") { - writer.uint32(16).uint64(message.cacheEpoch); + if (message.tenantId !== "") { + writer.uint32(18).string(message.tenantId); } return writer; }, - decode(input: BinaryReader | Uint8Array, length?: number): PublishBlockResponse { + decode(input: BinaryReader | Uint8Array, length?: number): CancelPrefillJobRequest { const reader = input instanceof BinaryReader ? input : new BinaryReader(input); const end = length === undefined ? reader.len : reader.pos + length; - const message = createBasePublishBlockResponse(); + const message = createBaseCancelPrefillJobRequest(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { case 1: { - if (tag !== 8) { + if (tag !== 10) { break; } - message.stored = reader.bool(); + message.jobId = reader.string(); continue; } case 2: { - if (tag !== 16) { + if (tag !== 18) { break; } - message.cacheEpoch = reader.uint64().toString(); + message.tenantId = reader.string(); continue; } } @@ -2583,35 +3872,97 @@ export const PublishBlockResponse: MessageFns = { return message; }, - fromJSON(object: any): PublishBlockResponse { + fromJSON(object: any): CancelPrefillJobRequest { return { - stored: isSet(object.stored) ? globalThis.Boolean(object.stored) : false, - cacheEpoch: isSet(object.cacheEpoch) - ? globalThis.String(object.cacheEpoch) - : isSet(object.cache_epoch) - ? globalThis.String(object.cache_epoch) - : "0", + jobId: isSet(object.jobId) + ? globalThis.String(object.jobId) + : isSet(object.job_id) + ? globalThis.String(object.job_id) + : "", + tenantId: isSet(object.tenantId) + ? globalThis.String(object.tenantId) + : isSet(object.tenant_id) + ? globalThis.String(object.tenant_id) + : "", }; }, - toJSON(message: PublishBlockResponse): unknown { + toJSON(message: CancelPrefillJobRequest): unknown { const obj: any = {}; - if (message.stored !== false) { - obj.stored = message.stored; + if (message.jobId !== "") { + obj.jobId = message.jobId; } - if (message.cacheEpoch !== "0") { - obj.cacheEpoch = message.cacheEpoch; + if (message.tenantId !== "") { + obj.tenantId = message.tenantId; } return obj; }, - create, I>>(base?: I): PublishBlockResponse { - return PublishBlockResponse.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): CancelPrefillJobRequest { + return CancelPrefillJobRequest.fromPartial(base ?? ({} as any)); }, - fromPartial, I>>(object: I): PublishBlockResponse { - const message = createBasePublishBlockResponse(); - message.stored = object.stored ?? false; - message.cacheEpoch = object.cacheEpoch ?? "0"; + fromPartial, I>>(object: I): CancelPrefillJobRequest { + const message = createBaseCancelPrefillJobRequest(); + message.jobId = object.jobId ?? ""; + message.tenantId = object.tenantId ?? ""; + return message; + }, +}; + +function createBaseCancelPrefillJobResponse(): CancelPrefillJobResponse { + return { cancelled: false }; +} + +export const CancelPrefillJobResponse: MessageFns = { + encode(message: CancelPrefillJobResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.cancelled !== false) { + writer.uint32(8).bool(message.cancelled); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): CancelPrefillJobResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseCancelPrefillJobResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 8) { + break; + } + + message.cancelled = reader.bool(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): CancelPrefillJobResponse { + return { cancelled: isSet(object.cancelled) ? globalThis.Boolean(object.cancelled) : false }; + }, + + toJSON(message: CancelPrefillJobResponse): unknown { + const obj: any = {}; + if (message.cancelled !== false) { + obj.cancelled = message.cancelled; + } + return obj; + }, + + create, I>>(base?: I): CancelPrefillJobResponse { + return CancelPrefillJobResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): CancelPrefillJobResponse { + const message = createBaseCancelPrefillJobResponse(); + message.cancelled = object.cancelled ?? false; return message; }, }; @@ -4373,6 +5724,111 @@ export const PrefillCacheServiceClient = makeGenericClientConstructor( serviceName: string; }; +/** + * PrefillWorkerService executes model prefill only. Workers load the same model + * as the primary, materialize immutable snapshots into their co-located + * PrefillCacheService, and never participate in autoregressive decode. + */ +export type PrefillWorkerServiceService = typeof PrefillWorkerServiceService; +export const PrefillWorkerServiceService = { + submitPrefillJob: { + path: "/kakeya.v1.PrefillWorkerService/SubmitPrefillJob" as const, + requestStream: false as const, + responseStream: false as const, + requestSerialize: (value: SubmitPrefillJobRequest): Buffer => + Buffer.from(SubmitPrefillJobRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): SubmitPrefillJobRequest => SubmitPrefillJobRequest.decode(value), + responseSerialize: (value: SubmitPrefillJobResponse): Buffer => + Buffer.from(SubmitPrefillJobResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): SubmitPrefillJobResponse => SubmitPrefillJobResponse.decode(value), + }, + getPrefillJobStatus: { + path: "/kakeya.v1.PrefillWorkerService/GetPrefillJobStatus" as const, + requestStream: false as const, + responseStream: false as const, + requestSerialize: (value: GetPrefillJobStatusRequest): Buffer => + Buffer.from(GetPrefillJobStatusRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): GetPrefillJobStatusRequest => GetPrefillJobStatusRequest.decode(value), + responseSerialize: (value: GetPrefillJobStatusResponse): Buffer => + Buffer.from(GetPrefillJobStatusResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): GetPrefillJobStatusResponse => GetPrefillJobStatusResponse.decode(value), + }, + cancelPrefillJob: { + path: "/kakeya.v1.PrefillWorkerService/CancelPrefillJob" as const, + requestStream: false as const, + responseStream: false as const, + requestSerialize: (value: CancelPrefillJobRequest): Buffer => + Buffer.from(CancelPrefillJobRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): CancelPrefillJobRequest => CancelPrefillJobRequest.decode(value), + responseSerialize: (value: CancelPrefillJobResponse): Buffer => + Buffer.from(CancelPrefillJobResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): CancelPrefillJobResponse => CancelPrefillJobResponse.decode(value), + }, +} as const; + +export interface PrefillWorkerServiceServer extends UntypedServiceImplementation { + submitPrefillJob: handleUnaryCall; + getPrefillJobStatus: handleUnaryCall; + cancelPrefillJob: handleUnaryCall; +} + +export interface PrefillWorkerServiceClient extends Client { + submitPrefillJob( + request: SubmitPrefillJobRequest, + callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void, + ): ClientUnaryCall; + submitPrefillJob( + request: SubmitPrefillJobRequest, + metadata: Metadata, + callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void, + ): ClientUnaryCall; + submitPrefillJob( + request: SubmitPrefillJobRequest, + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: SubmitPrefillJobResponse) => void, + ): ClientUnaryCall; + getPrefillJobStatus( + request: GetPrefillJobStatusRequest, + callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void, + ): ClientUnaryCall; + getPrefillJobStatus( + request: GetPrefillJobStatusRequest, + metadata: Metadata, + callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void, + ): ClientUnaryCall; + getPrefillJobStatus( + request: GetPrefillJobStatusRequest, + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: GetPrefillJobStatusResponse) => void, + ): ClientUnaryCall; + cancelPrefillJob( + request: CancelPrefillJobRequest, + callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void, + ): ClientUnaryCall; + cancelPrefillJob( + request: CancelPrefillJobRequest, + metadata: Metadata, + callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void, + ): ClientUnaryCall; + cancelPrefillJob( + request: CancelPrefillJobRequest, + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: CancelPrefillJobResponse) => void, + ): ClientUnaryCall; +} + +export const PrefillWorkerServiceClient = makeGenericClientConstructor( + PrefillWorkerServiceService, + "kakeya.v1.PrefillWorkerService", +) as unknown as { + new (address: string, credentials: ChannelCredentials, options?: Partial): PrefillWorkerServiceClient; + service: typeof PrefillWorkerServiceService; + serviceName: string; +}; + /** * DFlashProposerService: stateful remote DFlash drafter + f_θ restoration. * Per turn: Restore (prompt -> f_θ-projected verifier K/V) then SeedContext From 6c49a4cc2c49b81363ccd761068400cc69c6cd7f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 11 Jul 2026 15:09:26 +0000 Subject: [PATCH 03/11] test(distributed): gate worker orchestration and MLX cache equivalence Co-authored-by: FluffyAIcode --- .github/workflows/ci.yaml | 4 +- .github/workflows/integration.yaml | 35 ++ .../distributed/test_capability.py | 18 + .../distributed/test_prefill_auth.py | 108 ++++ .../distributed/test_prefill_cache.py | 13 + .../test_prefill_cache_runtime_fallback.py | 593 ++++++++++++++++++ .../distributed/test_prefill_cache_service.py | 120 ++++ .../distributed/test_prefill_compression.py | 95 +++ .../test_prefill_orchestrator_e2e.py | 165 +++++ .../distributed/test_prefill_scheduler.py | 160 +++++ .../distributed/test_prefill_worker.py | 496 +++++++++++++++ .../test_prefill_snapshot_mlx_equivalence.py | 77 +++ 12 files changed, 1881 insertions(+), 3 deletions(-) create mode 100644 tests/inference_engine/distributed/test_prefill_auth.py create mode 100644 tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py create mode 100644 tests/inference_engine/distributed/test_prefill_compression.py create mode 100644 tests/inference_engine/distributed/test_prefill_orchestrator_e2e.py create mode 100644 tests/inference_engine/distributed/test_prefill_scheduler.py create mode 100644 tests/inference_engine/distributed/test_prefill_worker.py create mode 100644 tests/integration/test_prefill_snapshot_mlx_equivalence.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 3c6c9dda..ce5a63d6 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -105,11 +105,9 @@ jobs: -v coverage report \ --include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*' \ - --omit='inference_engine/distributed/prefill_cache_runtime.py' \ --fail-under=100 coverage xml -o coverage.xml \ - --include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*' \ - --omit='inference_engine/distributed/prefill_cache_runtime.py' + --include='inference_engine/server/auth.py,inference_engine/server/config.py,inference_engine/server/errors.py,inference_engine/server/grpc_app.py,inference_engine/server/metrics.py,inference_engine/server/schemas.py,inference_engine/server/proto_gen/**/*.py,inference_engine/memory/*,inference_engine/bridge/*,inference_engine/distributed/*,inference_engine/network/*,inference_engine/scheduler/config.py,inference_engine/scheduler/session.py,inference_engine/pipeline/*,inference_engine/session/store.py,inference_engine/setup/*,sdks/python/kakeya/__init__.py,sdks/python/kakeya/errors.py,training/repr_align/*' - name: Upload coverage artifact if: always() diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index e4e2cebd..5902b554 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -80,6 +80,41 @@ jobs: fi echo "Found $MODEL_DIR" + - name: Gate real MLX distributed-prefill continuation equivalence + run: | + set -euo pipefail + if [ -z "${KAKEYA_MAC_VERIFIER_PATH:-}" ] || [ ! -d "$KAKEYA_MAC_VERIFIER_PATH" ]; then + echo "::error::KAKEYA_MAC_VERIFIER_PATH must point to the pre-warmed MLX verifier." + exit 1 + fi + PYBIN="$( + python3 - <<'PY' + import os, shutil, subprocess + candidates = [ + os.environ.get("KAKEYA_MAC_PYTHON"), + os.path.expanduser("~/kakeya-venv/bin/python"), + os.path.expanduser("~/.venv/bin/python"), + shutil.which("python3.13"), + shutil.which("python3"), + ] + for candidate in candidates: + if not candidate: + continue + if subprocess.run( + [candidate, "-c", "import mlx_lm, torch, pytest"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode == 0: + print(candidate) + break + PY + )" + test -n "$PYBIN" + PYTHONPATH=.:sdks/python "$PYBIN" -m pytest \ + -m integration \ + tests/integration/test_prefill_snapshot_mlx_equivalence.py \ + -q + - name: Install Python dependencies run: | # The runner is expected to have a long-lived venv. diff --git a/tests/inference_engine/distributed/test_capability.py b/tests/inference_engine/distributed/test_capability.py index c8978ef8..e730eed3 100644 --- a/tests/inference_engine/distributed/test_capability.py +++ b/tests/inference_engine/distributed/test_capability.py @@ -18,9 +18,11 @@ CapabilityRole, CacheCapability, CacheCompatibility, + CompressionCodec, ModelCapability, NodeCapability, NodeEndpoint, + PrefillWorkerCapability, ) T0 = 1_000_000.0 @@ -115,6 +117,7 @@ def test_cache_capability_and_endpoints_proto_round_trip(): layer_geometry_hash="geometry", kv_dtype="bfloat16", block_size_tokens=64, + tenant_namespace="tenant", ) card = NodeCapability( node_id="cache-peer", @@ -130,6 +133,8 @@ def test_cache_capability_and_endpoints_proto_round_trip(): load=0.5, tokens_served=100, bloom_filter=b"filter", + default_compression=CompressionCodec.ZLIB, + replication_factor=2, ), ), endpoints=( @@ -140,6 +145,19 @@ def test_cache_capability_and_endpoints_proto_round_trip(): 0.45, ), ), + prefill_workers=( + PrefillWorkerCapability( + compatibility, + worker_address="169.254.27.104:53051", + max_concurrent_jobs=1, + inflight_jobs=1, + queued_jobs=2, + load=0.5, + tokens_per_second_prefill=33.0, + ram_bytes_free=1234, + queued_tokens=456, + ), + ), ) assert NodeCapability.from_proto(card.to_proto()) == card assert CapabilityRole.PREFILL_CACHE.value == 5 diff --git a/tests/inference_engine/distributed/test_prefill_auth.py b/tests/inference_engine/distributed/test_prefill_auth.py new file mode 100644 index 00000000..bea8fee6 --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_auth.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import pytest + +from inference_engine.distributed.prefill_auth import ( + FleetAuthConfig, + PrefillAuthError, + metadata_pairs, + signed_metadata, + verify_metadata, +) +from inference_engine.server.proto_gen.kakeya.v1 import distributed_pb2 + + +def _request(): + return distributed_pb2.FetchBlocksRequest(lease_id="lease") + + +def _config(): + return FleetAuthConfig(b"x" * 32, "tenant-a", "node-a", 30) + + +def test_sign_and_verify_round_trip(): + request = _request() + metadata = signed_metadata(request, _config(), now=100) + assert verify_metadata(metadata, request, _config(), now=110) == ( + "tenant-a", "node-a", + ) + + +def test_tenant_hash_keys_are_isolated(): + a = _config().tenant_hash_key() + b = FleetAuthConfig(b"x" * 32, "tenant-b", "node-a").tenant_hash_key() + assert a != b + + +@pytest.mark.parametrize("mutator,match", [ + (lambda md: tuple((k, "tenant-b" if k.endswith("tenant-id") else v) + for k, v in md), "tenant mismatch"), + (lambda md: tuple((k, "bad" if k.endswith("auth-mac") else v) + for k, v in md), "invalid prefill authentication MAC"), + (lambda md: tuple((k, v) for k, v in md if not k.endswith("auth-mac")), + "missing prefill authentication metadata"), +]) +def test_auth_rejects_bad_metadata(mutator, match): + request = _request() + with pytest.raises(PrefillAuthError, match=match): + verify_metadata( + mutator(signed_metadata(request, _config(), now=100)), + request, + _config(), + now=100, + ) + + +def test_auth_rejects_replay_and_bad_timestamp(): + request = _request() + metadata = signed_metadata(request, _config(), now=100) + with pytest.raises(PrefillAuthError, match="replay window"): + verify_metadata(metadata, request, _config(), now=200) + bad = tuple(("x-kakeya-auth-ts", "nan") if k == "x-kakeya-auth-ts" + else (k, v) for k, v in metadata) + with pytest.raises(PrefillAuthError, match="invalid authentication timestamp"): + verify_metadata(bad, request, _config(), now=100) + + +def test_auth_rejects_tampered_request(): + request = _request() + metadata = signed_metadata(request, _config(), now=100) + with pytest.raises(PrefillAuthError, match="invalid prefill authentication MAC"): + verify_metadata( + metadata, + distributed_pb2.FetchBlocksRequest(lease_id="other"), + _config(), + now=100, + ) + + +def test_config_validation_and_file(tmp_path): + with pytest.raises(ValueError): + FleetAuthConfig(b"short", "t", "n") + with pytest.raises(ValueError): + FleetAuthConfig(b"x" * 32, "", "n") + with pytest.raises(ValueError): + FleetAuthConfig(b"x" * 32, "t", "") + with pytest.raises(ValueError): + FleetAuthConfig(b"x" * 32, "t", "n", 0) + path = tmp_path / "psk" + path.write_bytes(b"y" * 32 + b"\n") + assert FleetAuthConfig.from_file( + str(path), tenant_id="t", node_id="n", + ).psk == b"y" * 32 + + +def test_metadata_pairs_accepts_tuples(): + assert metadata_pairs((("a", "b"),)) == (("a", "b"),) + + class Item: + key = "c" + value = "d" + + assert metadata_pairs((Item(),)) == (("c", "d"),) + + +def test_sign_rejects_non_protobuf(): + with pytest.raises(TypeError): + signed_metadata(object(), _config()) + diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index 5295080b..2aacbae8 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -104,3 +104,16 @@ def test_pinned_eviction_and_missing_leased_block_guards(): store._blocks.pop(block.block_hash) with pytest.raises(KeyError, match="evicted"): store.fetch(lease.lease_id, now=1) + + +def test_put_rejects_when_active_lease_pins_capacity(): + import time + + store = PrefixCacheStore(_compat(), max_bytes=5, node_id="x") + first = CacheBlock.create(bytes(32), 1, b"12345") + store.put(first) + store.lookup([first.block_hash], now=time.time()) + second = CacheBlock.create(bytes.fromhex("01" * 32), 1, b"abc") + with pytest.raises(ValueError, match="pinned"): + store.put(second) + assert store.block_hashes() == (first.block_hash,) diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py new file mode 100644 index 00000000..5dd696d1 --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py @@ -0,0 +1,593 @@ +from __future__ import annotations + +from inference_engine.distributed.capability import ( + CacheCapability, + CacheCompatibility, + CompressionCodec, + NodeCapability, + NodeEndpoint, + PrefillWorkerCapability, +) +from inference_engine.distributed.prefill_cache import ( + CacheBlock, + PrefixCacheStore, + chained_block_hashes, +) +from inference_engine.distributed.prefill_cache_runtime import ( + DistributedPrefillCacheHook, + _Hit, +) +from inference_engine.distributed.prefill_scheduler import PrefillCostConfig +from inference_engine.server.proto_gen.kakeya.v1 import distributed_pb2 + + +class _Verifier: + def __init__(self): + self.cache = [] + self.cached_token_sequence = [] + self.next_global_position = 0 + self.next_token_logits = None + self.prefill_calls = 0 + + def reset(self): + self.cache = [] + self.cached_token_sequence = [] + self.next_global_position = 0 + + def prefill(self, tokens): + self.reset() + self.prefill_calls += 1 + self.cached_token_sequence = list(tokens) + self.next_global_position = len(tokens) + self.next_token_logits = _Row(tokens[-1]) + + def forward_block(self, tokens): + self.cached_token_sequence.extend(tokens) + self.next_global_position += len(tokens) + return [_Row(token) for token in tokens] + + def commit_or_truncate(self, *, forwarded, accepted): + assert forwarded == accepted + + +class _Row: + def __init__(self, value): + self.value = value + + def clone(self): + return _Row(self.value) + + +def test_import_failure_always_falls_back_to_full_local_prefill(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hashes = chained_block_hashes([1, 2], compatibility) + store.put(CacheBlock.create(hashes[0], 2, b"corrupt")) + hook = DistributedPrefillCacheHook( + store, + compression=CompressionCodec.NONE, + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "import_mlx_prefill_snapshot", + lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("bad snapshot")), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "export_mlx_prefill_snapshot", + lambda *args, **kwargs: b"fresh", + ) + verifier = _Verifier() + assert hook.prepare(verifier, [1, 2]) == 0 + assert verifier.prefill_calls == 1 + assert verifier.cached_token_sequence == [1, 2] + assert hook.stats.fallbacks == 1 + assert "bad snapshot" in hook.stats.last_fallback_reason + assert not store.invalidate(b"missing") + hook.close() + + +def test_tenant_hmac_changes_hash_namespace(monkeypatch, tmp_path): + from inference_engine.distributed.prefill_auth import FleetAuthConfig + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "export_mlx_prefill_snapshot", + lambda *args, **kwargs: b"snapshot", + ) + a = CacheCompatibility( + model_id="m", block_size_tokens=2, tenant_namespace="a", + ) + b = CacheCompatibility( + model_id="m", block_size_tokens=2, tenant_namespace="b", + ) + store_a = PrefixCacheStore(a, max_bytes=1024, node_id="a") + store_b = PrefixCacheStore(b, max_bytes=1024, node_id="b") + auth_a = FleetAuthConfig(b"x" * 32, "a", "node") + auth_b = FleetAuthConfig(b"x" * 32, "b", "node") + hook_a = DistributedPrefillCacheHook( + store_a, auth=auth_a, compression=CompressionCodec.NONE, + ) + hook_b = DistributedPrefillCacheHook( + store_b, auth=auth_b, compression=CompressionCodec.NONE, + ) + hook_a.prepare(_Verifier(), [1, 2]) + hook_b.prepare(_Verifier(), [1, 2]) + assert store_a.block_hashes() != store_b.block_hashes() + hook_a.close() + hook_b.close() + + +def test_runtime_validation_empty_and_provider_failure(): + store = PrefixCacheStore( + CacheCompatibility(model_id="m"), + max_bytes=1024, + node_id="head", + ) + with __import__("pytest").raises(ValueError): + DistributedPrefillCacheHook(store, lookup_timeout_s=0) + with __import__("pytest").raises(ValueError): + DistributedPrefillCacheHook(store, replication_factor=-1) + hook = DistributedPrefillCacheHook( + store, + registry_provider=lambda: (_ for _ in ()).throw(RuntimeError("boom")), + ) + assert hook.prepare(_Verifier(), []) == 0 + assert hook._cards() == () + hook.close() + + +def test_successful_local_import_suffix_and_on_reuse(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hashes = chained_block_hashes([1, 2, 3, 4], compatibility) + store.put(CacheBlock.create(hashes[0], 2, b"snapshot")) + imported = type("Imported", (), { + "token_count": 2, + "cached_token_ids": (1, 2), + "next_token_logits": _Row(2), + "block_hash": hashes[0], + })() + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "import_mlx_prefill_snapshot", + lambda *args, **kwargs: imported, + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "export_mlx_prefill_snapshot", + lambda *args, **kwargs: b"fresh", + ) + reused = [] + hook = DistributedPrefillCacheHook( + store, + compression=CompressionCodec.NONE, + on_reuse=reused.append, + ) + verifier = _Verifier() + assert hook.prepare(verifier, [1, 2, 3, 4]) == 2 + assert verifier.cached_token_sequence == [1, 2, 3, 4] + assert hook.stats.local_hits == 1 + assert reused == [2] + hook.close() + + +def test_import_budget_and_remote_worker_failures(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hook = DistributedPrefillCacheHook( + store, + max_import_bytes=2, + remote_compute_min_tokens=1, + worker_timeout_s=0.001, + worker_poll_interval_s=0.001, + ) + verifier = _Verifier() + assert hook._try_import( + verifier, [1, 2], _Hit("local", "l", 1, 1, 1, b"x"), + ) == 0 + assert hook._try_import( + verifier, [1, 2], _Hit("local", "l", 1, 2, 3, b"big"), + ) == 0 + assert hook._try_import( + verifier, [1, 2], _Hit("peer", "l", 1, 2, 3), + ) == 0 + + invalid = [ + type("Imported", (), { + "token_count": 0, + "cached_token_ids": (), + "next_token_logits": _Row(1), + "block_hash": b"h" * 32, + })(), + type("Imported", (), { + "token_count": 2, + "cached_token_ids": (1, 2), + "next_token_logits": _Row(1), + "block_hash": b"x" * 32, + })(), + type("Imported", (), { + "token_count": 2, + "cached_token_ids": (1, 2), + "next_token_logits": None, + "block_hash": b"h" * 32, + })(), + type("Imported", (), { + "token_count": 2, + "cached_token_ids": (9, 9), + "next_token_logits": _Row(1), + "block_hash": b"h" * 32, + })(), + ] + monkeypatch.setattr( + hook, + "_fetch_remote", + lambda hit: b"ok", + ) + for imported in invalid: + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "import_mlx_prefill_snapshot", + lambda *args, _imported=imported, **kwargs: _imported, + ) + assert hook._try_import( + verifier, + [1, 2], + _Hit("peer", "l", 1, 2, 2, block_hash=b"h" * 32), + ) == 0 + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "choose_prefill_worker", + lambda *args, **kwargs: None, + ) + assert hook._compute_remote([1, 2], [b"a" * 32]) is None + + target = type("Target", (), { + "address": "worker:1", + "rtt_ms": 1.0, + })() + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "choose_prefill_worker", + lambda *args, **kwargs: target, + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "submit_prefill_job_sync", + lambda *args, **kwargs: type("Response", (), {"job_id": "j"})(), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "get_prefill_job_sync", + lambda *args, **kwargs: type("Status", (), { + "status": 4, + "failure_reason": "failed", + })(), + ) + assert hook._compute_remote([1, 2], [b"a" * 32]) is None + assert hook.stats.remote_job_failures == 1 + + # A perpetually queued job reaches the bounded deadline and falls back. + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "get_prefill_job_sync", + lambda *args, **kwargs: type("Status", (), { + "status": 1, + "failure_reason": "", + })(), + ) + assert hook._compute_remote([1, 2], [b"a" * 32]) is None + assert hook.stats.remote_job_failures == 2 + hook.close() + + +def test_dynamic_replica_selection_and_cost_reject(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + card = NodeCapability( + "peer", + "peer:1", + caches=(CacheCapability(compatibility, "peer:2"),), + endpoints=(NodeEndpoint("peer:2", "lan", 1, 5),), + ) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hook = DistributedPrefillCacheHook( + store, + registry_provider=lambda: (card,), + replication_factor=1, + cost_config=PrefillCostConfig( + local_prefill_tps=10000, + default_worker_tps=1, + link_mbps=1, + default_rtt_ms=10, + minimum_savings_ratio=0.5, + ), + ) + assert hook._publish_peers(b"h" * 32) == ("peer:2",) + monkeypatch.setattr( + hook, + "_lookup_peer", + lambda peer, hashes: _Hit(peer, "l", 1, 1, 10_000_000), + ) + assert hook._best_hit([b"h" * 32]) is None + assert hook._peer_rtt("peer:2") == 5 + assert hook._peer_rtt("unknown") == 0 + hook.close() + + +def test_publish_boundary_dispatches_selected_replica(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hook = DistributedPrefillCacheHook(store, peers=("peer:1",)) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "export_mlx_prefill_snapshot", + lambda *args, **kwargs: b"snapshot", + ) + calls = [] + + class Publisher: + def submit(self, fn, *args, **kwargs): + calls.append((fn, args, kwargs)) + + def shutdown(self, **kwargs): + pass + + hook._publisher = Publisher() + verifier = _Verifier() + verifier.prefill([1, 2]) + 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" + hook.close() + + +def test_unaligned_reused_prefix_computes_boundary_remainder(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hook = DistributedPrefillCacheHook(store) + monkeypatch.setattr(hook, "_publish_boundary", lambda *args, **kwargs: None) + verifier = _Verifier() + verifier.prefill([1, 2, 3]) + hashes = chained_block_hashes([1, 2, 3, 4, 5, 6], compatibility) + hook._compute_and_publish(verifier, [1, 2, 3, 4, 5, 6], hashes, 3) + # Snapshot ends at token 3; token 4 completes that block before [5,6]. + assert verifier.cached_token_sequence[-3:] == [4, 5, 6] + hook.close() + + +class _Context: + def __init__(self, stub): + self.stub = stub + + def __enter__(self): + return type("Channel", (), { + "unary_unary": lambda *args, **kwargs: None, + })() + + def __exit__(self, *args): + return False + + +def test_fetch_remote_validates_stream(monkeypatch): + import hashlib + + compatibility = CacheCompatibility(model_id="m") + hook = DistributedPrefillCacheHook( + PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"), + ) + payload = b"payload" + + class Stub: + def __init__(self, channel): + pass + + def FetchBlocks(self, request, **kwargs): + return [distributed_pb2.FetchBlocksResponse( + chunk_index=0, + total_chunks=1, + data=payload, + block_sha256=hashlib.sha256(payload).digest(), + )] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime.grpc.insecure_channel", + lambda address: _Context(Stub), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + Stub, + ) + assert hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7)) == payload + + class EmptyStub(Stub): + def FetchBlocks(self, request, **kwargs): + return [] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + EmptyStub, + ) + with __import__("pytest").raises(RuntimeError, match="incomplete"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7)) + + class ErrorStub(Stub): + def FetchBlocks(self, request, **kwargs): + import grpc + raise grpc.RpcError("down") + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + ErrorStub, + ) + with __import__("pytest").raises(RuntimeError, match="fetch failed"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7)) + + class CorruptStub(Stub): + def FetchBlocks(self, request, **kwargs): + return [distributed_pb2.FetchBlocksResponse( + chunk_index=0, + total_chunks=1, + data=payload, + block_sha256=b"x" * 32, + )] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + CorruptStub, + ) + with __import__("pytest").raises(RuntimeError, match="checksum"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 7)) + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + Stub, + ) + with __import__("pytest").raises(RuntimeError, match="lease checksum"): + hook._fetch_remote(_Hit( + "peer", + "lease", + 1, + 1, + 7, + payload_sha256=b"x" * 32, + )) + + with __import__("pytest").raises(RuntimeError, match="import budget"): + hook._fetch_remote( + _Hit( + "peer", "lease", 1, 1, hook.max_import_bytes + 1, + ), + ) + + class ChangedMetadataStub(Stub): + def FetchBlocks(self, request, **kwargs): + digest = hashlib.sha256(b"ab").digest() + return [ + distributed_pb2.FetchBlocksResponse( + chunk_index=0, total_chunks=2, data=b"a", + block_hash=b"h", block_sha256=digest, + ), + distributed_pb2.FetchBlocksResponse( + chunk_index=1, total_chunks=3, data=b"b", + block_hash=b"h", block_sha256=digest, + ), + ] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + ChangedMetadataStub, + ) + with __import__("pytest").raises(RuntimeError, match="metadata changed"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 2)) + + class TooManyChunksStub(Stub): + def FetchBlocks(self, request, **kwargs): + return [distributed_pb2.FetchBlocksResponse( + chunk_index=0, + total_chunks=65_537, + data=b"", + block_hash=b"h", + block_sha256=hashlib.sha256(b"").digest(), + )] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + TooManyChunksStub, + ) + with __import__("pytest").raises(RuntimeError, match="chunk count"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 0)) + + class DuplicateStub(Stub): + def FetchBlocks(self, request, **kwargs): + digest = hashlib.sha256(b"aa").digest() + return [ + distributed_pb2.FetchBlocksResponse( + chunk_index=0, total_chunks=2, data=b"a", + block_hash=b"h", block_sha256=digest, + ), + distributed_pb2.FetchBlocksResponse( + chunk_index=0, total_chunks=2, data=b"a", + block_hash=b"h", block_sha256=digest, + ), + ] + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + DuplicateStub, + ) + with __import__("pytest").raises(RuntimeError, match="duplicate"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 2)) + + hook.max_import_bytes = 1 + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + Stub, + ) + with __import__("pytest").raises(RuntimeError, match="import budget"): + hook._fetch_remote(_Hit("peer", "lease", 1, 1, 1)) + hook.close() + + +def test_lookup_peer_success_miss_and_rpc_error(monkeypatch): + import grpc + + compatibility = CacheCompatibility(model_id="m") + hook = DistributedPrefillCacheHook( + PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"), + ) + + class LookupStub: + response = distributed_pb2.LookupPrefixResponse( + lease_id="lease", + hit_block_count=1, + hit_token_count=5, + transfer_bytes=7, + ) + + def __init__(self, channel): + pass + + def LookupPrefix(self, request, **kwargs): + return self.response + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime.grpc.insecure_channel", + lambda address: _Context(LookupStub), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + LookupStub, + ) + hit = hook._lookup_peer("peer", [b"h" * 32]) + assert hit is not None and hit.hit_tokens == 5 + LookupStub.response = distributed_pb2.LookupPrefixResponse() + assert hook._lookup_peer("peer", [b"h" * 32]) is None + LookupStub.response = distributed_pb2.LookupPrefixResponse( + lease_id="lease", + hit_block_count=2, + ) + assert hook._lookup_peer("peer", [b"h" * 32]) is None + + class ErrorStub(LookupStub): + def LookupPrefix(self, request, **kwargs): + raise grpc.RpcError("down") + + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "distributed_pb2_grpc.PrefillCacheServiceStub", + ErrorStub, + ) + assert hook._lookup_peer("peer", [b"h" * 32]) is None + hook.close() + diff --git a/tests/inference_engine/distributed/test_prefill_cache_service.py b/tests/inference_engine/distributed/test_prefill_cache_service.py index e4f4cabd..ac093ab5 100644 --- a/tests/inference_engine/distributed/test_prefill_cache_service.py +++ b/tests/inference_engine/distributed/test_prefill_cache_service.py @@ -13,6 +13,10 @@ ) from inference_engine.distributed.prefill_cache import PrefixCacheStore from inference_engine.distributed.prefill_cache import CacheBlock +from inference_engine.distributed.prefill_auth import ( + FleetAuthConfig, + signed_metadata, +) from inference_engine.distributed.prefill_cache_service import ( PrefillCacheServiceServicer, add_prefill_cache_service, @@ -67,6 +71,54 @@ async def test_lookup_and_fetch_over_real_grpc(): await server.stop(0) +@pytest.mark.asyncio +async def test_authenticated_cache_service_rejects_unsigned_requests(): + compatibility = CacheCompatibility( + model_id="m", + block_size_tokens=2, + tenant_namespace="tenant", + ) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="peer") + hashes = store.put_prefix([1, 2], [b"payload"]) + auth = FleetAuthConfig(b"k" * 32, "tenant", "client") + server = grpc.aio.server() + port = server.add_insecure_port("127.0.0.1:0") + address = f"127.0.0.1:{port}" + add_prefill_cache_service( + server, store, cache_address=address, auth=auth, + ) + await server.start() + try: + request = distributed_pb2.LookupPrefixRequest( + compatibility=_compat().to_proto(), + # overwritten below to use the authenticated tenant namespace + block_hashes=hashes, + ) + request.compatibility.CopyFrom(compatibility.to_proto()) + async with grpc.aio.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + with pytest.raises(grpc.aio.AioRpcError) as error: + await stub.LookupPrefix(request) + assert error.value.code() == grpc.StatusCode.UNAUTHENTICATED + response = await stub.LookupPrefix( + request, + metadata=signed_metadata(request, auth), + ) + assert response.hit_block_count == 1 + wrong = distributed_pb2.LookupPrefixRequest( + compatibility=_compat().to_proto(), + block_hashes=hashes, + ) + with pytest.raises(grpc.aio.AioRpcError) as error: + await stub.LookupPrefix( + wrong, + metadata=signed_metadata(wrong, auth), + ) + assert error.value.code() == grpc.StatusCode.PERMISSION_DENIED + finally: + await server.stop(0) + + @pytest.mark.asyncio async def test_incompatible_and_dead_peers_are_misses(): store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer") @@ -224,5 +276,73 @@ def test_service_validation_and_dead_publish(): cache_address="peer:1", chunk_bytes=0, ) + with pytest.raises(ValueError, match="max_payload_bytes"): + PrefillCacheServiceServicer( + store, + cache_address="peer:1", + max_payload_bytes=0, + ) block = CacheBlock.create(bytes(32), 1, b"x") assert not publish_block_sync("127.0.0.1:1", _compat(), block, timeout_s=0.1) + + +@pytest.mark.asyncio +async def test_publish_rejects_duplicate_over_budget_and_incomplete_streams(): + store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer") + server = grpc.aio.server() + port = server.add_insecure_port("127.0.0.1:0") + address = f"127.0.0.1:{port}" + add_prefill_cache_service( + server, + store, + cache_address=address, + chunk_bytes=3, + max_payload_bytes=4, + ) + await server.start() + + async def call(chunks): + async with grpc.aio.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + return await stub.PublishBlock(iter(chunks)) + + import hashlib + base = dict( + block_hash=bytes(32), + token_count=2, + total_chunks=2, + block_sha256=hashlib.sha256(b"abcdef").digest(), + compatibility=_compat().to_proto(), + ) + try: + with pytest.raises(grpc.aio.AioRpcError) as exc: + await call([ + distributed_pb2.PublishBlockRequest( + **base, chunk_index=0, data=b"abc", + ), + distributed_pb2.PublishBlockRequest( + **base, chunk_index=0, data=b"abc", + ), + ]) + assert exc.value.code() == grpc.StatusCode.INVALID_ARGUMENT + + with pytest.raises(grpc.aio.AioRpcError) as exc: + await call([ + distributed_pb2.PublishBlockRequest( + **base, chunk_index=0, data=b"abc", + ), + distributed_pb2.PublishBlockRequest( + **base, chunk_index=1, data=b"def", + ), + ]) + assert exc.value.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + + with pytest.raises(grpc.aio.AioRpcError) as exc: + await call([ + distributed_pb2.PublishBlockRequest( + **base, chunk_index=0, data=b"abc", + ), + ]) + assert exc.value.code() == grpc.StatusCode.DATA_LOSS + finally: + await server.stop(0) diff --git a/tests/inference_engine/distributed/test_prefill_compression.py b/tests/inference_engine/distributed/test_prefill_compression.py new file mode 100644 index 00000000..027b8e92 --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_compression.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import pytest + +from inference_engine.distributed.capability import CompressionCodec +from inference_engine.distributed.prefill_compression import ( + compress_payload, + decompress_payload, + payload_sizes, +) + + +def test_zlib_round_trip_and_sizes(): + raw = (b"prefill-kv-" * 1000) + compressed = compress_payload(raw, CompressionCodec.ZLIB) + assert len(compressed) < len(raw) + assert payload_sizes(compressed) == (len(compressed), len(raw)) + assert decompress_payload( + compressed, + max_uncompressed_bytes=len(raw), + ) == raw + + +def test_none_is_backward_compatible_raw_payload(): + raw = b"legacy-kpkv" + assert compress_payload(raw, CompressionCodec.NONE) == raw + assert decompress_payload(raw, max_uncompressed_bytes=100) == raw + assert payload_sizes(raw) == (len(raw), len(raw)) + + +def test_compression_and_import_limits_validate(): + with pytest.raises(ValueError, match="zlib level"): + compress_payload(b"x", CompressionCodec.ZLIB, level=10) + with pytest.raises(ValueError, match="unsupported compression"): + compress_payload(b"x", 99) # type: ignore[arg-type] + framed = compress_payload(b"x" * 100, CompressionCodec.ZLIB) + with pytest.raises(ValueError, match="import budget"): + decompress_payload(framed, max_uncompressed_bytes=10) + with pytest.raises(ValueError, match="import budget"): + decompress_payload(b"x" * 11, max_uncompressed_bytes=10) + + +def test_truncated_and_corrupt_payloads_fail(): + import hashlib + import struct + + framed = compress_payload(b"x" * 100, CompressionCodec.ZLIB) + with pytest.raises(ValueError): + decompress_payload(framed[:10], max_uncompressed_bytes=1000) + damaged = framed[:-1] + bytes([framed[-1] ^ 1]) + with pytest.raises((ValueError, __import__("zlib").error)): + decompress_payload(damaged, max_uncompressed_bytes=1000) + magic = b"KPC1" + with pytest.raises(ValueError, match="unsupported compression codec"): + decompress_payload( + struct.pack("<4sBQ32s", magic, 99, 1, hashlib.sha256(b"x").digest()) + + b"x", + max_uncompressed_bytes=100, + ) + with pytest.raises(ValueError, match="unsupported framed"): + decompress_payload( + struct.pack( + "<4sBQ32s", + magic, + int(CompressionCodec.NONE), + 1, + hashlib.sha256(b"x").digest(), + ) + b"x", + max_uncompressed_bytes=100, + ) + with pytest.raises(ValueError, match="decompressed payload size"): + decompress_payload( + framed[:5] + struct.pack("= 5 + + def set(self): + self.calls = 5 + + final_store = PrefixCacheStore(COMPAT, max_bytes=4096, node_id="final") + final_jobs = PrefillJobStore(_Engine(), final_store) + final_job = PrefillJob( + "final", "final", "tenant", (1, 2), (b"a" * 32,), + CompressionCodec.NONE, + ) + final_job.cancelled = SequenceEvent() + final_jobs._jobs[final_job.job_id] = final_job + final_jobs._run(final_job.job_id) + assert final_job.state == PrefillJobState.CANCELLED + final_jobs.close() + + +@pytest.mark.asyncio +async def test_worker_service_tenant_and_not_found_errors(worker): + address, _, _, _ = worker + wrong_tenant = _submit("wrong-tenant") + wrong_tenant.tenant_id = "other" + with pytest.raises(grpc.aio.AioRpcError) as exc: + await _rpc(address, "SubmitPrefillJob", wrong_tenant) + assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED + wrong_get = distributed_pb2.GetPrefillJobStatusRequest( + job_id="anything", tenant_id="other", + ) + with pytest.raises(grpc.aio.AioRpcError) as exc: + await _rpc(address, "GetPrefillJobStatus", wrong_get) + assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED + wrong_cancel = distributed_pb2.CancelPrefillJobRequest( + job_id="anything", tenant_id="other", + ) + with pytest.raises(grpc.aio.AioRpcError) as exc: + await _rpc(address, "CancelPrefillJob", wrong_cancel) + assert exc.value.code() == grpc.StatusCode.PERMISSION_DENIED + missing = distributed_pb2.GetPrefillJobStatusRequest( + job_id="missing", + tenant_id="tenant", + ) + with pytest.raises(grpc.aio.AioRpcError) as exc: + await _rpc(address, "GetPrefillJobStatus", missing) + assert exc.value.code() == grpc.StatusCode.NOT_FOUND + cancel = distributed_pb2.CancelPrefillJobRequest( + job_id="missing", + tenant_id="tenant", + ) + with pytest.raises(grpc.aio.AioRpcError) as exc: + await _rpc(address, "CancelPrefillJob", cancel) + assert exc.value.code() == grpc.StatusCode.NOT_FOUND + + +@pytest.mark.asyncio +async def test_sync_clients_work_without_auth(): + cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="worker") + jobs = PrefillJobStore(_Engine(), cache) + server = grpc.aio.server() + add_prefill_worker_service( + server, + jobs, + node_id="worker", + cache_address="cache:1", + ) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + address = f"127.0.0.1:{port}" + try: + submitted = await asyncio.to_thread( + submit_prefill_job_sync, + address, + _submit("no-auth"), + timeout_s=2, + ) + request = distributed_pb2.GetPrefillJobStatusRequest( + job_id=submitted.job_id, + tenant_id="tenant", + ) + status = await asyncio.to_thread( + get_prefill_job_sync, + address, + request, + timeout_s=2, + ) + assert status.job_id == submitted.job_id + finally: + jobs.close() + await server.stop(0) + diff --git a/tests/integration/test_prefill_snapshot_mlx_equivalence.py b/tests/integration/test_prefill_snapshot_mlx_equivalence.py new file mode 100644 index 00000000..936af22f --- /dev/null +++ b/tests/integration/test_prefill_snapshot_mlx_equivalence.py @@ -0,0 +1,77 @@ +"""Real-MLX gate: imported prefill state must preserve continuation logits.""" +from __future__ import annotations + +import os + +import pytest + + +def test_real_mlx_prefill_snapshot_preserves_continuation_logits(): + pytest.importorskip("mlx.core") + torch = pytest.importorskip("torch") + + from inference_engine.backends.mlx.verifier import MLXSinkWindowVerifier + from inference_engine.distributed.capability import ( + CacheCompatibility, + CompressionCodec, + ) + from inference_engine.distributed.prefill_cache import PrefixCacheStore + from inference_engine.distributed.prefill_cache_runtime import ( + DistributedPrefillCacheHook, + ) + from kv_cache_proposer.verifier import VerifierConfig + + model_path = os.environ["KAKEYA_MAC_VERIFIER_PATH"] + verifier = MLXSinkWindowVerifier(VerifierConfig( + model_id=model_path, + dtype=torch.bfloat16, + device="cpu", + sink_size=4, + window_size=64, + )) + prompt = ( + "Kakeya distributed prefill equivalence test. " + "The imported cache must produce exactly the same continuation logits. " + ) * 8 + token_ids = verifier.tokenizer.encode(prompt) + compatibility = CacheCompatibility( + model_id=os.environ.get("KAKEYA_CACHE_MODEL_ID", model_path), + model_revision=os.environ.get("KAKEYA_MODEL_REVISION", ""), + tokenizer_revision=os.environ.get("KAKEYA_TOKENIZER_REVISION", ""), + cache_format_version="kakeya-prefill-v2-zlib", + quantization="4bit-mlx", + layer_geometry_hash=os.environ["KAKEYA_LAYER_GEOMETRY_HASH"], + kv_dtype="bfloat16", + block_size_tokens=64, + tenant_namespace="integration", + ) + store = PrefixCacheStore( + compatibility, + max_bytes=1 << 30, + node_id="integration-head", + ) + hook = DistributedPrefillCacheHook( + store, + compression=CompressionCodec.ZLIB, + ) + try: + assert hook.prepare(verifier, token_ids) == 0 + baseline_logits = verifier.next_token_logits.clone() + baseline_argmax = int(torch.argmax(baseline_logits).item()) + baseline_tokens = list(verifier.cached_token_sequence) + + reused = hook.prepare(verifier, token_ids) + assert reused == len(token_ids) + assert int(torch.argmax(verifier.next_token_logits).item()) == baseline_argmax + assert torch.equal(verifier.next_token_logits, baseline_logits) + assert verifier.cached_token_sequence == baseline_tokens + + # One real decode step must remain bit-identical after re-import. + next_token = baseline_argmax + imported_row = verifier.forward_block([next_token])[-1].clone() + hook.prepare(verifier, token_ids) + local_row = verifier.forward_block([next_token])[-1].clone() + assert torch.equal(imported_row, local_row) + finally: + hook.close() + From 8486d36e98c7d93ba7deec035f75e71bb6020ec0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 11 Jul 2026 15:11:23 +0000 Subject: [PATCH 04/11] fix(ci): trigger Mac gate directly from runtime path changes Co-authored-by: FluffyAIcode --- .github/workflows/integration.yaml | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index 5902b554..d1585da6 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -4,12 +4,10 @@ name: Integration (Mac M4) # tests/integration/ against real Qwen3-0.6B on Apple Silicon. # # Trigger model: -# - Pull-request events. Only fires when the PR carries the -# ``needs-mac-m4`` label (auto-applied by .github/workflows/ -# auto-label-mac.yaml when a PR touches inference_engine/, -# sdks/, proto/, or tests/integration/). PRs that don't touch -# verifier-dependent code skip this gate entirely so the runner -# pool isn't burned on doc-only or CI-only PRs. +# - Pull-request events touching runtime/model/proto/integration paths. +# Path filtering is used directly instead of depending on an auto-label: +# workflows triggered with GITHUB_TOKEN do not recursively trigger the +# ``labeled`` event, which previously made the first Mac gate skip. # - Manual workflow_dispatch for re-runs from the Actions UI. # # Runner requirements (self-hosted): @@ -27,6 +25,14 @@ on: # Only run on PR events for branches targeting main. types: [opened, synchronize, reopened, labeled] branches: [main] + paths: + - "inference_engine/**" + - "kv_cache_proposer/**" + - "proto/**" + - "sdks/**" + - "tests/integration/**" + - "tests/backends/mlx/**" + - ".github/workflows/integration.yaml" workflow_dispatch: {} # Cancel superseded runs on the same PR — saves runner time when @@ -39,13 +45,6 @@ concurrency: jobs: integration: name: pytest -m integration on Mac M4 - # Only fire on labeled PRs (this saves the runner pool from - # doc-only / CI-only PRs that don't touch verifier-dependent - # code). The auto-label workflow adds 'needs-mac-m4' on file - # paths that warrant the GA gate. - if: | - github.event_name == 'workflow_dispatch' || - contains(github.event.pull_request.labels.*.name, 'needs-mac-m4') runs-on: [self-hosted, macOS, ARM64, kakeya-mac-m4] timeout-minutes: 90 steps: From f70f1ff14e4e8c36cb055675e72bee0d9acbae61 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 11 Jul 2026 15:21:01 +0000 Subject: [PATCH 05/11] feat(ops): expose worker telemetry and Mac MLX equivalence preset Co-authored-by: FluffyAIcode --- inference_engine/bridge/manifest.py | 16 ++++++++++++++++ inference_engine/network/state.py | 15 +++++++++++++++ tests/inference_engine/bridge/test_manifest.py | 1 + .../network/test_network_state.py | 7 +++++++ 4 files changed, 39 insertions(+) diff --git a/inference_engine/bridge/manifest.py b/inference_engine/bridge/manifest.py index 483d63db..73535794 100644 --- a/inference_engine/bridge/manifest.py +++ b/inference_engine/bridge/manifest.py @@ -104,6 +104,22 @@ def _harness_preset( PRESETS: Dict[str, Preset] = { p.name: p for p in ( + Preset( + name="mlx-prefill-snapshot-equivalence", + description="ADR 0017 real-model gate: local MLX prefill vs " + "snapshot export/import must preserve continuation " + "logits, argmax, retained KV tokens and one decode step.", + command_templates=( + ( + "python3", "-m", "pytest", + "-m", "integration", + "tests/integration/test_prefill_snapshot_mlx_equivalence.py", + "-q", + ), + ), + timeout_minutes=45, + validate_reports=False, + ), Preset( name="mlx-distributed-dflash-e2e-inproc", description="Real-model distributed DFlash+f_θ E2E (in-process): loads " diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 82a1be2f..0e194804 100644 --- a/inference_engine/network/state.py +++ b/inference_engine/network/state.py @@ -95,6 +95,7 @@ def nodes(self) -> list[dict[str, Any]]: for card in self.registry.snapshot(): registration = registrations.get(card.node_id, {}) cache = card.caches[0] if card.caches else None + worker = card.prefill_workers[0] if card.prefill_workers else None endpoint = sorted( card.endpoints, key=lambda item: item.priority, @@ -132,6 +133,19 @@ def nodes(self) -> list[dict[str, Any]]: } if cache else None ), + "prefill_worker": ( + { + "address": worker.worker_address, + "max_concurrent_jobs": worker.max_concurrent_jobs, + "inflight_jobs": worker.inflight_jobs, + "queued_jobs": worker.queued_jobs, + "queued_tokens": worker.queued_tokens, + "load": worker.load, + "tokens_per_second": worker.tokens_per_second_prefill, + "ram_bytes_free": worker.ram_bytes_free, + } + if worker else None + ), "endpoint": ( { "address": endpoint[0].address, @@ -160,6 +174,7 @@ def nodes(self) -> list[dict[str, Any]]: "memory_bytes": 0, "models": [], "cache": None, + "prefill_worker": None, "endpoint": { "address": registration["address"], "network": "pending", diff --git a/tests/inference_engine/bridge/test_manifest.py b/tests/inference_engine/bridge/test_manifest.py index ce623c9e..02b51f59 100644 --- a/tests/inference_engine/bridge/test_manifest.py +++ b/tests/inference_engine/bridge/test_manifest.py @@ -96,6 +96,7 @@ def test_allowlist_contains_exactly_the_documented_presets(): "mlx-kakeya-launcher-full", "mlx-kakeya-launcher-smoke", "mlx-multitenant-pressure", + "mlx-prefill-snapshot-equivalence", "mlx-upgrade", "mlx-upstream-batch-probe", "pytest-path", diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py index c92cc96d..418671eb 100644 --- a/tests/inference_engine/network/test_network_state.py +++ b/tests/inference_engine/network/test_network_state.py @@ -6,6 +6,7 @@ CapabilityRegistry, NodeCapability, NodeEndpoint, + PrefillWorkerCapability, ) from inference_engine.distributed.prefill_cache import PrefixCacheStore from inference_engine.network.state import NetworkState @@ -27,6 +28,12 @@ def _state(tmp_path): ), ), endpoints=(NodeEndpoint("head:2", "thunderbolt", 100, 0.4),), + prefill_workers=(PrefillWorkerCapability( + compatibility, + worker_address="head:3", + queued_tokens=128, + tokens_per_second_prefill=32, + ),), ) return NetworkState( CapabilityRegistry(self_card=card), From cfb2b6cf31fc52a0e180af9568ee9b686fd0e57c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 03:51:36 +0000 Subject: [PATCH 06/11] fix(mac-runner): recover headlessly and diagnose offline queues Co-authored-by: FluffyAIcode --- docs/ops/mac-m4-runner-setup.md | 41 +++++++--- .../install_autorecover_launchagent.sh | 59 +++++++++++--- .../mac_bridge/recover_runner_after_reboot.sh | 80 +++++++++++++++---- scripts/mac_bridge/runner_healthcheck.sh | 47 +++++++++++ 4 files changed, 187 insertions(+), 40 deletions(-) create mode 100755 scripts/mac_bridge/runner_healthcheck.sh diff --git a/docs/ops/mac-m4-runner-setup.md b/docs/ops/mac-m4-runner-setup.md index 2cb233c2..ac759646 100644 --- a/docs/ops/mac-m4-runner-setup.md +++ b/docs/ops/mac-m4-runner-setup.md @@ -1,11 +1,8 @@ # Mac M4 self-hosted runner setup -This runner backs the **Integration (Mac M4)** GitHub Actions workflow -(`.github/workflows/integration.yaml`). It runs `pytest -m integration` -against real Qwen3-0.6B on every PR labelled `needs-mac-m4` -(auto-applied by `.github/workflows/auto-label-mac.yaml` when a PR -touches `inference_engine/`, `sdks/`, `proto/`, `tests/integration/`, -or `kv_cache_proposer/`). +This runner backs the **Integration (Mac M4)** and **Mac bridge** workflows. +Integration is triggered directly by runtime/model/proto/integration path +changes; it does not depend on a label being applied by another workflow. ## Hardware requirements @@ -66,13 +63,11 @@ pyenv global 3.12.7 Confirm `python3 --version` returns 3.12.x and `python3 -c 'import platform; print(platform.machine())'` returns `arm64`. -### 4. (Optional) long-lived venv +### 4. Pin the MLX workload venv -The workflow currently does `pip install -e .` per run, which is -~30 s on a warm pip cache. If you want to skip even that, create a -venv at `~/kakeya-runner-venv` and add a step to the workflow that -activates it before `pytest`. v0.3 keeps the per-run install for -simplicity. +Set `KAKEYA_MAC_PYTHON` in the runner service environment to the long-lived +venv that imports `mlx_lm`, `torch` and `pytest`. Real MLX gates resolve this +interpreter before the legacy Qwen integration environment is installed. ## Runtime expectations @@ -123,12 +118,14 @@ cd ~/actions-runner Workflow failures are visible at `Actions → Integration (Mac M4)`. The "Surface failure summary" step inlines the test names + first-line error messages so triage doesn't require downloading the JUnit XML. -If the runner itself is offline (queue depth grows, no jobs pick up), check on the Mac: +If the runner itself is offline (multiple differently-labelled jobs remain +`queued` and no first step starts), check on the Mac: ```bash cd ~/actions-runner sudo ./svc.sh status tail -200 ~/Library/Logs/actions-runner/Runner_*.log +bash /path/to/repo/scripts/mac_bridge/runner_healthcheck.sh ``` Common causes: @@ -136,6 +133,24 @@ Common causes: - HF cache was purged; the verify step fails. Re-warm. - Disk full from accumulated pip downloads; clear cache. +### Headless reboot recovery + +A user LaunchAgent runs only after GUI login. For an unattended Mac mini, +install the system watchdog once so the runner recovers before login: + +```bash +cd /path/to/Kakeya-LLM-Inference-engine +sudo "$HOME/actions-runner/svc.sh" install 2>/dev/null || true +sudo "$HOME/actions-runner/svc.sh" start 2>/dev/null || true +bash scripts/mac_bridge/install_autorecover_launchagent.sh --system +bash scripts/mac_bridge/recover_runner_after_reboot.sh +``` + +For jobs that are already queued, the last command is the immediate recovery +action. GitHub assigns queued jobs automatically once `Runner.Listener` +reconnects. The system LaunchDaemon retries every 60 seconds and does not +require a logged-in desktop session. + ## Mac bridge (cloud-agent access) The same runner also serves the **Mac bridge** diff --git a/scripts/mac_bridge/install_autorecover_launchagent.sh b/scripts/mac_bridge/install_autorecover_launchagent.sh index a102c53b..14ac198e 100755 --- a/scripts/mac_bridge/install_autorecover_launchagent.sh +++ b/scripts/mac_bridge/install_autorecover_launchagent.sh @@ -1,20 +1,29 @@ #!/usr/bin/env bash -# Install a user LaunchAgent that re-checks the mac-bridge runner -# after reboot and periodically self-heals it. +# Install a runner watchdog. `--system` installs a headless LaunchDaemon that +# runs before user login (recommended for remote Mac minis). Without it, a +# per-user LaunchAgent is installed as a fallback. set -euo pipefail REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)" RECOVER_SCRIPT="${REPO_ROOT}/scripts/mac_bridge/recover_runner_after_reboot.sh" -PLIST_DIR="${HOME}/Library/LaunchAgents" -PLIST_PATH="${PLIST_DIR}/com.kakeya.mac-bridge-runner-autorecover.plist" LABEL="com.kakeya.mac-bridge-runner-autorecover" UID_NUM="$(id -u)" +USER_NAME="$(id -un)" +MODE="user" +if [ "${1:-}" = "--system" ]; then + MODE="system" +elif [ -n "${1:-}" ]; then + echo "usage: $0 [--system]" >&2 + exit 2 +fi -mkdir -p "$PLIST_DIR" [ -x "$RECOVER_SCRIPT" ] || chmod +x "$RECOVER_SCRIPT" +mkdir -p "${HOME}/actions-runner/_diag" +TMP_PLIST="$(mktemp)" +trap 'rm -f "$TMP_PLIST"' EXIT -cat >"$PLIST_PATH" <"$TMP_PLIST" < @@ -31,8 +40,18 @@ cat >"$PLIST_PATH" <WorkingDirectory ${REPO_ROOT} + EnvironmentVariables + + HOME${HOME} + RUNNER_DIR${HOME}/actions-runner + + RunAtLoad + KeepAlive + + SuccessfulExit + StartInterval 60 @@ -44,10 +63,28 @@ cat >"$PLIST_PATH" < EOF -launchctl bootout "gui/${UID_NUM}" "${PLIST_PATH}" >/dev/null 2>&1 || true -launchctl bootstrap "gui/${UID_NUM}" "${PLIST_PATH}" -launchctl enable "gui/${UID_NUM}/${LABEL}" || true -launchctl kickstart -k "gui/${UID_NUM}/${LABEL}" || true +if [ "$MODE" = "system" ]; then + PLIST_PATH="/Library/LaunchDaemons/${LABEL}.plist" + # UserName lets the runner retain access to its registration, models and + # workspace while the daemon itself starts before GUI login. + /usr/libexec/PlistBuddy -c "Add :UserName string ${USER_NAME}" "$TMP_PLIST" + sudo install -o root -g wheel -m 644 "$TMP_PLIST" "$PLIST_PATH" + sudo launchctl bootout system "$PLIST_PATH" >/dev/null 2>&1 || true + sudo launchctl bootstrap system "$PLIST_PATH" + sudo launchctl enable "system/${LABEL}" || true + sudo launchctl kickstart -k "system/${LABEL}" || true + DOMAIN="system" +else + PLIST_DIR="${HOME}/Library/LaunchAgents" + PLIST_PATH="${PLIST_DIR}/${LABEL}.plist" + mkdir -p "$PLIST_DIR" + install -m 644 "$TMP_PLIST" "$PLIST_PATH" + launchctl bootout "gui/${UID_NUM}" "$PLIST_PATH" >/dev/null 2>&1 || true + launchctl bootstrap "gui/${UID_NUM}" "$PLIST_PATH" + launchctl enable "gui/${UID_NUM}/${LABEL}" || true + launchctl kickstart -k "gui/${UID_NUM}/${LABEL}" || true + DOMAIN="gui/${UID_NUM}" +fi echo "[mac-bridge-autorecover] installed: ${PLIST_PATH}" -echo "[mac-bridge-autorecover] label: ${LABEL}" +echo "[mac-bridge-autorecover] label: ${DOMAIN}/${LABEL}" diff --git a/scripts/mac_bridge/recover_runner_after_reboot.sh b/scripts/mac_bridge/recover_runner_after_reboot.sh index c80e9204..6497d00e 100755 --- a/scripts/mac_bridge/recover_runner_after_reboot.sh +++ b/scripts/mac_bridge/recover_runner_after_reboot.sh @@ -12,37 +12,85 @@ set -euo pipefail RUNNER_DIR="${RUNNER_DIR:-$HOME/actions-runner}" LOG_DIR="${RUNNER_LOG_DIR:-$HOME/actions-runner/_diag}" mkdir -p "$LOG_DIR" +LOCK_DIR="${TMPDIR:-/tmp}/kakeya-runner-recover.lock" log() { echo "[mac-bridge-recover] $*" >&2; } -if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then +listener_running() { + pgrep -f "${RUNNER_DIR}/bin/Runner.Listener" >/dev/null 2>&1 \ + || pgrep -f "Runner.Listener run" >/dev/null 2>&1 +} + +if ! mkdir "$LOCK_DIR" 2>/dev/null; then + lock_pid="$(cat "$LOCK_DIR/pid" 2>/dev/null || true)" + if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then + log "another recovery attempt is active (pid=$lock_pid)." + exit 0 + fi + rm -rf "$LOCK_DIR" + mkdir "$LOCK_DIR" +fi +echo "$$" >"$LOCK_DIR/pid" +trap 'rm -rf "$LOCK_DIR" 2>/dev/null || true' EXIT + +if listener_running; then log "Runner.Listener already running." exit 0 fi if [ ! -x "${RUNNER_DIR}/run.sh" ]; then - log "runner not found at ${RUNNER_DIR}; skipping." - exit 0 + log "runner not found at ${RUNNER_DIR}." + exit 1 +fi +if [ ! -s "${RUNNER_DIR}/.runner" ]; then + log "runner exists but is not registered (.runner missing)." + exit 1 fi -if [ -x "${RUNNER_DIR}/svc.sh" ] && "${RUNNER_DIR}/svc.sh" status 2>/dev/null | grep -q "installed"; then - log "service installed; attempting svc.sh start" - if "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1; then - sleep 2 - if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then - log "runner started via service." - exit 0 +if ! curl -fsSI --max-time 10 https://github.com/ >/dev/null 2>&1; then + log "github.com is unreachable; leaving runner stopped for the next retry." + exit 1 +fi + +if [ -x "${RUNNER_DIR}/svc.sh" ]; then + status="$("${RUNNER_DIR}/svc.sh" status 2>&1 || true)" + # Do not match the phrase "not installed". + if echo "$status" | grep -qi "installed" \ + && ! echo "$status" | grep -qi "not installed"; then + log "official service is installed; attempting svc.sh start" + if [ "$(id -u)" -eq 0 ]; then + "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1 || true + elif sudo -n true >/dev/null 2>&1; then + sudo -n "${RUNNER_DIR}/svc.sh" start >/dev/null 2>&1 || true fi + for _ in 1 2 3 4 5; do + sleep 2 + if listener_running; then + log "runner started via official service." + exit 0 + fi + done + log "official service did not bring up listener; using direct fallback." fi - log "service start did not bring up listener, falling back to run.sh" fi ts="$(date +%Y%m%d_%H%M%S)" -nohup "${RUNNER_DIR}/run.sh" >"${LOG_DIR}/runner-nohup-${ts}.log" 2>&1 & -sleep 2 -if pgrep -f "Runner.Listener.*${RUNNER_DIR}" >/dev/null 2>&1; then - log "runner started via nohup run.sh fallback." - exit 0 +( + cd "$RUNNER_DIR" + nohup ./run.sh >"${LOG_DIR}/runner-nohup-${ts}.log" 2>&1 & +) +for _ in 1 2 3 4 5 6 7 8; do + sleep 2 + if listener_running; then + log "runner started via direct run.sh fallback." + exit 0 + fi +done + +latest_log="$(ls -t "$LOG_DIR"/runner-nohup-*.log 2>/dev/null | head -1 || true)" +if [ -n "$latest_log" ]; then + log "last runner output:" + tail -20 "$latest_log" >&2 || true fi log "failed to start runner listener." diff --git a/scripts/mac_bridge/runner_healthcheck.sh b/scripts/mac_bridge/runner_healthcheck.sh new file mode 100755 index 00000000..e49967ce --- /dev/null +++ b/scripts/mac_bridge/runner_healthcheck.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Diagnose why GitHub Actions jobs are queued for kakeya-mac-m4. +set -uo pipefail + +RUNNER_DIR="${RUNNER_DIR:-$HOME/actions-runner}" +LABEL="com.kakeya.mac-bridge-runner-autorecover" +failed=0 + +check() { + local name="$1"; shift + if "$@" >/dev/null 2>&1; then + printf 'PASS %s\n' "$name" + else + printf 'FAIL %s\n' "$name" + failed=1 + fi +} + +echo "runner_dir=$RUNNER_DIR" +check "runner registration (.runner)" test -s "$RUNNER_DIR/.runner" +check "runner executable" test -x "$RUNNER_DIR/run.sh" +check "github.com connectivity" curl -fsSI --max-time 10 https://github.com/ +check "Runner.Listener process" sh -c \ + "pgrep -f '$RUNNER_DIR/bin/Runner.Listener' >/dev/null || pgrep -f 'Runner.Listener run' >/dev/null" +check "disk has >=10 GiB free" sh -c \ + "[ \$(df -k '$RUNNER_DIR' | awk 'NR==2 {print \$4}') -ge 10485760 ]" + +echo +echo "official service:" +"$RUNNER_DIR/svc.sh" status 2>&1 || true +echo +echo "user watchdog:" +launchctl print "gui/$(id -u)/$LABEL" 2>&1 | head -30 || true +echo +echo "system watchdog:" +sudo -n launchctl print "system/$LABEL" 2>&1 | head -30 || true +echo +echo "recent runner diagnostics:" +ls -t "$RUNNER_DIR"/_diag/Runner_*.log \ + "$RUNNER_DIR"/_diag/runner-nohup-*.log 2>/dev/null \ + | head -2 | while read -r log; do + echo "=== $log ===" + tail -20 "$log" + done + +exit "$failed" + From bd03913e6a0005c6a0f195edc0186218767c4fc1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 04:17:47 +0000 Subject: [PATCH 07/11] fix(mac-bridge): bootstrap git-lfs after pointer checkout Co-authored-by: FluffyAIcode --- .github/workflows/mac-bridge.yaml | 33 ++++++++++++++++++++++++------- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/.github/workflows/mac-bridge.yaml b/.github/workflows/mac-bridge.yaml index 1eb13c64..a3c66d56 100644 --- a/.github/workflows/mac-bridge.yaml +++ b/.github/workflows/mac-bridge.yaml @@ -55,13 +55,32 @@ jobs: with: # Push results back to the request branch. persist-credentials: true - # k3-* presets load LFS-tracked checkpoints from the repo - # (e.g. results/research/f_theta_v5_s5_sliding/ - # f_theta_weights.pt). Without lfs:true the workspace holds - # pointer files and torch.load fails with the cryptic - # "Unsupported operand 118" (ASCII 'v' = the first byte of - # an LFS pointer). - lfs: true + + - name: Ensure git-lfs is available + # Do this AFTER a normal checkout. checkout@v4 with lfs:true invokes + # git-lfs before any workflow step can repair PATH, which makes a + # recovered/minimal runner fail at step 2 with "Unable to locate + # executable file: git-lfs". + run: | + set -euo pipefail + if ! command -v git-lfs >/dev/null 2>&1; then + for candidate in /opt/homebrew/bin/git-lfs /usr/local/bin/git-lfs; do + if [ -x "$candidate" ]; then + echo "$(dirname "$candidate")" >> "$GITHUB_PATH" + export PATH="$(dirname "$candidate"):$PATH" + break + fi + done + fi + if ! command -v git-lfs >/dev/null 2>&1; then + if command -v brew >/dev/null 2>&1; then + brew install git-lfs + else + echo "::error::git-lfs is missing and Homebrew is unavailable." + exit 1 + fi + fi + git lfs version - name: Show request run: | From 67a3eba3815912731d1a164f86b380ffbbf31ae1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 04:20:11 +0000 Subject: [PATCH 08/11] fix(mac-ci): bootstrap git-lfs before reused-worktree checkout Co-authored-by: FluffyAIcode --- .github/workflows/integration.yaml | 24 ++++++++++++++ .github/workflows/mac-bridge.yaml | 51 +++++++++++++++--------------- 2 files changed, 49 insertions(+), 26 deletions(-) diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index d1585da6..02bcfe9c 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -48,6 +48,30 @@ jobs: runs-on: [self-hosted, macOS, ARM64, kakeya-mac-m4] timeout-minutes: 90 steps: + - name: Bootstrap git-lfs before checkout + # Reused worktrees can retain a git-lfs post-checkout hook. The hook + # executes inside checkout@v4, so repair PATH before checkout. + run: | + set -euo pipefail + if [ -x /opt/homebrew/bin/git-lfs ]; then + echo "/opt/homebrew/bin" >> "$GITHUB_PATH" + exit 0 + fi + if [ -x /usr/local/bin/git-lfs ]; then + echo "/usr/local/bin" >> "$GITHUB_PATH" + exit 0 + fi + brew_bin="" + for candidate in /opt/homebrew/bin/brew /usr/local/bin/brew; do + if [ -x "$candidate" ]; then brew_bin="$candidate"; break; fi + done + if [ -z "$brew_bin" ]; then + echo "::error::git-lfs is missing and Homebrew is unavailable." + exit 1 + fi + "$brew_bin" install git-lfs + echo "$(dirname "$brew_bin")" >> "$GITHUB_PATH" + - uses: actions/checkout@v4 with: # Full history so the runner can compare against base for diff --git a/.github/workflows/mac-bridge.yaml b/.github/workflows/mac-bridge.yaml index a3c66d56..07e63fb3 100644 --- a/.github/workflows/mac-bridge.yaml +++ b/.github/workflows/mac-bridge.yaml @@ -51,36 +51,35 @@ jobs: runs-on: [self-hosted, macOS, ARM64, kakeya-mac-m4] timeout-minutes: 150 steps: - - uses: actions/checkout@v4 - with: - # Push results back to the request branch. - persist-credentials: true - - - name: Ensure git-lfs is available - # Do this AFTER a normal checkout. checkout@v4 with lfs:true invokes - # git-lfs before any workflow step can repair PATH, which makes a - # recovered/minimal runner fail at step 2 with "Unable to locate - # executable file: git-lfs". + - name: Bootstrap git-lfs before checkout + # A previous `git lfs install --local` leaves a post-checkout hook in + # the reused worktree. That hook runs inside actions/checkout itself, + # even with lfs:false, so git-lfs must be available BEFORE checkout. run: | set -euo pipefail - if ! command -v git-lfs >/dev/null 2>&1; then - for candidate in /opt/homebrew/bin/git-lfs /usr/local/bin/git-lfs; do - if [ -x "$candidate" ]; then - echo "$(dirname "$candidate")" >> "$GITHUB_PATH" - export PATH="$(dirname "$candidate"):$PATH" - break - fi - done + if [ -x /opt/homebrew/bin/git-lfs ]; then + echo "/opt/homebrew/bin" >> "$GITHUB_PATH" + exit 0 fi - if ! command -v git-lfs >/dev/null 2>&1; then - if command -v brew >/dev/null 2>&1; then - brew install git-lfs - else - echo "::error::git-lfs is missing and Homebrew is unavailable." - exit 1 - fi + if [ -x /usr/local/bin/git-lfs ]; then + echo "/usr/local/bin" >> "$GITHUB_PATH" + exit 0 fi - git lfs version + brew_bin="" + for candidate in /opt/homebrew/bin/brew /usr/local/bin/brew; do + if [ -x "$candidate" ]; then brew_bin="$candidate"; break; fi + done + if [ -z "$brew_bin" ]; then + echo "::error::git-lfs is missing and Homebrew is unavailable." + exit 1 + fi + "$brew_bin" install git-lfs + echo "$(dirname "$brew_bin")" >> "$GITHUB_PATH" + + - uses: actions/checkout@v4 + with: + # Push results back to the request branch. + persist-credentials: true - name: Show request run: | From a781d939e2004b287bc51bb42c01bbebce5d5af3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 04:24:02 +0000 Subject: [PATCH 09/11] fix(mac-ci): resolve the working MLX venv and verifier path Co-authored-by: FluffyAIcode --- .github/workflows/integration.yaml | 17 ++++++++++++++++- inference_engine/bridge/runner_python.py | 8 ++++++++ .../bridge/test_runner_python.py | 4 ++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index 02bcfe9c..def81123 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -104,9 +104,16 @@ jobs: echo "Found $MODEL_DIR" - name: Gate real MLX distributed-prefill continuation equivalence + env: + KAKEYA_MAC_VERIFIER_PATH_VAR: ${{ vars.KAKEYA_MAC_VERIFIER_PATH || '' }} run: | set -euo pipefail - if [ -z "${KAKEYA_MAC_VERIFIER_PATH:-}" ] || [ ! -d "$KAKEYA_MAC_VERIFIER_PATH" ]; then + default_verifier="$HOME/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit" + if [ ! -d "$default_verifier" ]; then + default_verifier="models/gemma-4-26B-A4B-it-mlx-4bit" + fi + export KAKEYA_MAC_VERIFIER_PATH="${KAKEYA_MAC_VERIFIER_PATH_VAR:-$default_verifier}" + if [ ! -d "$KAKEYA_MAC_VERIFIER_PATH" ]; then echo "::error::KAKEYA_MAC_VERIFIER_PATH must point to the pre-warmed MLX verifier." exit 1 fi @@ -117,6 +124,14 @@ jobs: os.environ.get("KAKEYA_MAC_PYTHON"), os.path.expanduser("~/kakeya-venv/bin/python"), os.path.expanduser("~/.venv/bin/python"), + os.path.expanduser( + "~/Documents/Kakeya-LLM-Inference-engine-pr109/" + ".venv-mac/bin/python3.13" + ), + os.path.expanduser( + "~/Documents/Kakeya-LLM-Inference-engine-pr109/" + ".venv-mac/bin/python" + ), shutil.which("python3.13"), shutil.which("python3"), ] diff --git a/inference_engine/bridge/runner_python.py b/inference_engine/bridge/runner_python.py index 1b0c27b4..84eb848e 100644 --- a/inference_engine/bridge/runner_python.py +++ b/inference_engine/bridge/runner_python.py @@ -55,6 +55,14 @@ def workload_python_candidates( environ.get("KAKEYA_MAC_PYTHON"), expanduser("~/kakeya-venv/bin/python"), expanduser("~/.venv/bin/python"), + expanduser( + "~/Documents/Kakeya-LLM-Inference-engine-pr109/" + ".venv-mac/bin/python3.13", + ), + expanduser( + "~/Documents/Kakeya-LLM-Inference-engine-pr109/" + ".venv-mac/bin/python", + ), which("python3.13"), which("python3"), ] diff --git a/tests/inference_engine/bridge/test_runner_python.py b/tests/inference_engine/bridge/test_runner_python.py index 8a0a6bb5..acf3c5a4 100644 --- a/tests/inference_engine/bridge/test_runner_python.py +++ b/tests/inference_engine/bridge/test_runner_python.py @@ -30,6 +30,8 @@ def test_candidates_prioritise_pin_then_venvs_then_path(): "/pin/bin/python", "/home/me/kakeya-venv/bin/python", "/home/me/.venv/bin/python", + "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python3.13", + "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python", "/usr/bin/python3.13", "/usr/bin/python3", ] @@ -44,6 +46,8 @@ def test_candidates_drop_empty_and_dedupe(): assert cands == [ "/home/me/kakeya-venv/bin/python", "/home/me/.venv/bin/python", + "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python3.13", + "/home/me/Documents/Kakeya-LLM-Inference-engine-pr109/.venv-mac/bin/python", ] assert None not in cands From f20ae90d33be6f4ad71ef03b2b107588aa81f665 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 04:27:33 +0000 Subject: [PATCH 10/11] fix(mac-test): skip absent interpreters and derive local geometry namespace Co-authored-by: FluffyAIcode --- .github/workflows/integration.yaml | 2 +- .../integration/test_prefill_snapshot_mlx_equivalence.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index def81123..787e3e69 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -136,7 +136,7 @@ jobs: shutil.which("python3"), ] for candidate in candidates: - if not candidate: + if not candidate or not os.path.isfile(candidate): continue if subprocess.run( [candidate, "-c", "import mlx_lm, torch, pytest"], diff --git a/tests/integration/test_prefill_snapshot_mlx_equivalence.py b/tests/integration/test_prefill_snapshot_mlx_equivalence.py index 936af22f..d499731f 100644 --- a/tests/integration/test_prefill_snapshot_mlx_equivalence.py +++ b/tests/integration/test_prefill_snapshot_mlx_equivalence.py @@ -40,7 +40,13 @@ def test_real_mlx_prefill_snapshot_preserves_continuation_logits(): tokenizer_revision=os.environ.get("KAKEYA_TOKENIZER_REVISION", ""), cache_format_version="kakeya-prefill-v2-zlib", quantization="4bit-mlx", - layer_geometry_hash=os.environ["KAKEYA_LAYER_GEOMETRY_HASH"], + # This is a single-verifier round-trip gate; production head/workers + # derive and compare the real geometry hash separately. Keep the + # namespace stable without requiring another machine-local env var. + layer_geometry_hash=os.environ.get( + "KAKEYA_LAYER_GEOMETRY_HASH", + "integration-single-verifier", + ), kv_dtype="bfloat16", block_size_tokens=64, tenant_namespace="integration", From 159557b7f9bc506a772b9ab74616e7c5d513a032 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 12 Jul 2026 04:31:37 +0000 Subject: [PATCH 11/11] fix(mac-ci): isolate legacy integration deps from managed Python Co-authored-by: FluffyAIcode --- .github/workflows/integration.yaml | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index 787e3e69..193407d4 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -155,21 +155,29 @@ jobs: - name: Install Python dependencies run: | - # The runner is expected to have a long-lived venv. - # If a per-run venv is preferred, swap to ``python3 -m venv .venv``. - python3 -m pip install --upgrade pip + # Keep the legacy Qwen/Transformers-4 suite isolated from both + # Homebrew's PEP-668-managed system Python and the Transformers-5 MLX + # production venv used by the preceding real-model gate. + LEGACY_VENV="$HOME/.kakeya/integration-legacy-venv" + if [ ! -x "$LEGACY_VENV/bin/python" ]; then + mkdir -p "$(dirname "$LEGACY_VENV")" + python3 -m venv "$LEGACY_VENV" + fi + LEGACY_PY="$LEGACY_VENV/bin/python" + "$LEGACY_PY" -m pip install --upgrade pip # The repo runs via PYTHONPATH (see ci.yaml) — it is NOT a pip package # (no setup.py/pyproject.toml), so install runtime deps from # requirements.txt rather than an editable `-e .` (which errors with # "does not appear to be a Python project"). - python3 -m pip install -r requirements.txt + "$LEGACY_PY" -m pip install -r requirements.txt # The integration suite exercises the legacy dllm-hub Qwen proposer, # whose remote modeling file depends on the Transformers 4.x # decoder_layer.attention_type API. Keep this runner in the dedicated # legacy range; K3/Gemma production paths use requirements.txt's # unbounded Transformers 5.x-compatible environment. - python3 -m pip install 'transformers>=4.45,<5.0' - python3 -m pip install pytest pytest-asyncio pytest-timeout coverage + "$LEGACY_PY" -m pip install 'transformers>=4.45,<5.0' + "$LEGACY_PY" -m pip install pytest pytest-asyncio pytest-timeout coverage + echo "KAKEYA_INTEGRATION_PY=$LEGACY_PY" >> "$GITHUB_ENV" - name: Run integration suite env: @@ -180,7 +188,7 @@ jobs: run: | mkdir -p results/platform-tests stamp=$(date +%s) - python3 -m pytest \ + "$KAKEYA_INTEGRATION_PY" -m pytest \ -m integration \ tests/integration/ \ --junitxml="results/platform-tests/integration-mac-m4-${stamp}.junit.xml" \