diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 9eb50cc1..3c6c9dda 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -97,16 +97,19 @@ jobs: tests/inference_engine/setup/ \ tests/inference_engine/bridge/ \ tests/inference_engine/distributed/ \ + tests/inference_engine/network/ \ tests/sdk/python/ \ tests/training/repr_align/ \ tests/backends/mlx/test_env.py \ --junitxml=junit.xml \ -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/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/*' \ + --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/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/*' + --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' - name: Upload coverage artifact if: always() diff --git a/.github/workflows/integration.yaml b/.github/workflows/integration.yaml index 6bb7c60e..e4e2cebd 100644 --- a/.github/workflows/integration.yaml +++ b/.github/workflows/integration.yaml @@ -90,6 +90,12 @@ jobs: # 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 + # 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 - name: Run integration suite diff --git a/README.md b/README.md index 47255ce9..3b46af79 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,53 @@ latency drift over a 4-hour, 480-turn Mac M4 run; bounded memory). > verifier for lightweight serving, or the **restored Gemma-4 26B** path > (proposer + f_θ/S5) for the memory-bounded, recall-preserving engine below. +## 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. + +This is not remote attention and not coherent shared RAM: + +```text +tokenize + chained prefix hashes + │ + ├── local lookup ──────────────┐ + └── P2P lookup over gossip ────┤ choose longest compatible prefix + ▼ + stream one immutable KV snapshot + ▼ + local suffix prefill → local decode +``` + +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; +- point-to-point chunked gRPC publish/fetch with SHA-256 validation; +- failure-safe fallback to local prefill; +- Thunderbolt/LAN/Tailscale endpoint priority; +- node registration, inference groups, token accounting and topology UI. + +Live product dashboard: **[https://kakeya.ai](https://kakeya.ai)**. + +Two-Mac measured evidence (Gemma 26B MLX 4-bit, 93-token prompt): + +```text +cold local prefill 5.926 s +remote Thunderbolt hit 0.061 s +observed speedup ≈97× +``` + +Architecture and operations: + +- [ADR 0016 — Distributed Prefill KV Cache Network](docs/adr/0016-distributed-prefill-kv-cache-network.md) +- [Two-Mac live report](docs/reports/distributed-prefill-kv-mac-thunderbolt.md) +- [Operator runbook](docs/ops/distributed-prefill-kv-network.md) + ## Quickstart (5 minutes on Mac M4 / Linux x86) > **Status — v0.4** (`v0.4-mac` / `v0.4-cuda` tags). Ships from source; PyPI + @@ -102,6 +149,7 @@ patterns — see [`docs/quickstart.md`](docs/quickstart.md). | `AppendTokens` / `Generation` coordinators | Drive prefill / incremental forward / greedy decode; route per-session (multi-tenant) or single. | `inference_engine.session.{coordinator,generator}` | | Python / TypeScript SDKs | `kakeya.Client` / `Session` (sync gRPC); `@kakeya/runtime` (Node 20+). | [`sdks/`](sdks/) | | HTTP shim (deprecated) | OpenAI-compatible `/v1/chat/completions`; `Deprecation` + `Sunset` headers. | `inference_engine.server.app` | +| **Distributed Prefill KV Cache** | P2P capability gossip, exact compatibility locks, chained longest-prefix lookup, chunked snapshot publish/fetch, local suffix prefill and local decode. | `inference_engine.distributed.prefill_cache*`, `inference_engine.network` | ## Runtime evidence (foundational, carried into v0.4) @@ -703,7 +751,7 @@ scripts/ | v0.5 GA multi-host hardening | queued | mTLS node identity, Bonjour seed discovery, K3 DFlash hidden-state flow over the mlx.distributed ring | | Async continuous batching | designing | Dynamic mid-flight arrival + ragged-length cohorts under the async gRPC `Generate` handlers (current batcher is fixed-cohort) | | Deployment polish | queued | PyPI + npm publishing, GHCR Docker image, `kakeya prewarm` CLI, `kakeya chat` REPL | -| Cross-request KV reuse | designing | Sessions survive across requests on gRPC; turns intra-session drift into 0 ms inter-request drift | +| **Distributed Prefill KV reuse** | ✅ live MVP | Cross-node immutable snapshots, chained longest-prefix matching, Thunderbolt gRPC transfer and public fleet dashboard ([ADR 0016](docs/adr/0016-distributed-prefill-kv-cache-network.md)) | ## Continuous integration diff --git a/deploy/cloudflare-worker/.gitignore b/deploy/cloudflare-worker/.gitignore new file mode 100644 index 00000000..0dcc8a41 --- /dev/null +++ b/deploy/cloudflare-worker/.gitignore @@ -0,0 +1,2 @@ +node_modules/ +.wrangler/ diff --git a/deploy/cloudflare-worker/README.md b/deploy/cloudflare-worker/README.md new file mode 100644 index 00000000..f6babf36 --- /dev/null +++ b/deploy/cloudflare-worker/README.md @@ -0,0 +1,62 @@ +# kakeya.ai Cloudflare Worker + +`kakeya-inference-network` owns `kakeya.ai/*` and forwards the public product +surface to the direct origin at `agent.kakeya.ai`. + +Routing: + +- `/` and `/network` → `/network` +- `/v1/network/*` → same API path +- `/healthz` → `/v1/network/summary` +- unknown browser paths → dashboard + +## Validate + +```bash +npm install +npm audit --omit=dev +npm run check +``` + +## Deploy + +```bash +npx wrangler login +npm run deploy +``` + +Expected route: + +```text +kakeya.ai/* (zone kakeya.ai) +``` + +## Verify + +```bash +curl -fsS https://kakeya.ai/ | grep "Kakeya Inference Network" +curl -fsS https://kakeya.ai/healthz +curl -fsS https://kakeya.ai/v1/network/nodes +``` + +Responses carry: + +```text +X-Kakeya-Surface: inference-network +``` + +## Rollback + +List versions/deployments: + +```bash +npx wrangler versions list +npx wrangler deployments list +``` + +Roll back with Wrangler's version rollback/deployment command, or remove the +`kakeya.ai/*` route. The direct origin remains available at: + +```text +https://agent.kakeya.ai/network +``` diff --git a/deploy/cloudflare-worker/package-lock.json b/deploy/cloudflare-worker/package-lock.json new file mode 100644 index 00000000..710fee4f --- /dev/null +++ b/deploy/cloudflare-worker/package-lock.json @@ -0,0 +1,1503 @@ +{ + "name": "kakeya-inference-network-worker", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "kakeya-inference-network-worker", + "devDependencies": { + "wrangler": "^4.0.0" + } + }, + "node_modules/@cloudflare/kv-asset-handler": { + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/@cloudflare/kv-asset-handler/-/kv-asset-handler-0.5.0.tgz", + "integrity": "sha512-jxQYkj8dSIzc0cD6cMMNdOc1UVjqSqu8BZdor5s8cGjW2I8BjODt/kWPVdY+u9zj3ms75Q5qaZgnxUad83+eAg==", + "dev": true, + "license": "MIT OR Apache-2.0", + "engines": { + "node": ">=22.0.0" + } + }, + "node_modules/@cloudflare/unenv-preset": { + "version": "2.16.1", + "resolved": "https://registry.npmjs.org/@cloudflare/unenv-preset/-/unenv-preset-2.16.1.tgz", + "integrity": "sha512-ECxObrMfyTl5bhQf/lZCXwo5G6xX9IAUo+nDMKK4SZ8m4Jvvxp52vilxyySSWh2YTZz8+HQ07qGH/2rEom1vDw==", + "dev": true, + "license": "MIT OR Apache-2.0", + "peerDependencies": { + "unenv": "2.0.0-rc.24", + "workerd": ">1.20260305.0 <2.0.0-0" + }, + "peerDependenciesMeta": { + "workerd": { + "optional": true + } + } + }, + "node_modules/@cloudflare/workerd-darwin-64": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/@cloudflare/workerd-darwin-64/-/workerd-darwin-64-1.20260708.1.tgz", + "integrity": "sha512-HXFCvhS1wpg3uXO0CLUwmwC41i2loM5FSK69EUchOBpmYBAXxT1oHLm6EOA5lqhTk5Mu9kjRiQYxa1GwKPwfJg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=16" + } + }, + "node_modules/@cloudflare/workerd-darwin-arm64": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/@cloudflare/workerd-darwin-arm64/-/workerd-darwin-arm64-1.20260708.1.tgz", + "integrity": "sha512-JVlJaKDoRTVKSroHIlf8g3UCPjKj4iDbMZE2CNYht5qQ+2rL0FAUiVlV82G3BqKnnw9kHYnnsMzC08b9zVtdzA==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=16" + } + }, + "node_modules/@cloudflare/workerd-linux-64": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/@cloudflare/workerd-linux-64/-/workerd-linux-64-1.20260708.1.tgz", + "integrity": "sha512-3daE60YdD7YX0Jtuzc9DE/r/qMkmx8ZvHTkF8Mzmp3F5tbzlV0DAzmu5PFUPF2WuvtKbAhZKbvC2cHmWpQYxnA==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=16" + } + }, + "node_modules/@cloudflare/workerd-linux-arm64": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/@cloudflare/workerd-linux-arm64/-/workerd-linux-arm64-1.20260708.1.tgz", + "integrity": "sha512-VLdNYOx5Hj+9C6isy0ACWZsbMtSxex2DIJWEe7cZxUdlphZ58ZT8zxNXK8yunFiowd34hn3VwGMopdvdj8lvmA==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=16" + } + }, + "node_modules/@cloudflare/workerd-windows-64": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/@cloudflare/workerd-windows-64/-/workerd-windows-64-1.20260708.1.tgz", + "integrity": "sha512-bC/aSAwLy16Vjo24i9XU3aWH+eRgz7NeR5xPKavGbembO18ZywYTQbXh14eXtY6fAqN3RzRG8psijTdhX4xydA==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=16" + } + }, + "node_modules/@cspotcode/source-map-support": { + "version": "0.8.1", + "resolved": "https://registry.npmjs.org/@cspotcode/source-map-support/-/source-map-support-0.8.1.tgz", + "integrity": "sha512-IchNf6dN4tHoMFIn/7OE8LWZ19Y6q/67Bmf6vnGREv8RSbBVb9LPJxEcnwrcwX6ixSvaiGoomAUvu4YSxXrVgw==", + "dev": true, + "license": "MIT", + "dependencies": { + "@jridgewell/trace-mapping": "0.3.9" + }, + "engines": { + "node": ">=12" + } + }, + "node_modules/@emnapi/runtime": { + "version": "1.11.2", + "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.11.2.tgz", + "integrity": "sha512-kyOl3X0DuTiT1h2ft8r2fYO8JYtU9a9Xis/zBSiGArNaagCOWx90N1k2wxp18czFDH+OgcWGb5ZP/XMt3dcyPA==", + "dev": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@esbuild/aix-ppc64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.28.1.tgz", + "integrity": "sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "aix" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-arm": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/android-arm/-/android-arm-0.28.1.tgz", + "integrity": "sha512-0k2F129Xdio1TdJfzJ8sy1Q47vUD2NnwdhiAf7drUN1EBTfPf4hsFCtmMgu/6m8JSzsBrlmVjudMBQqOfG8usQ==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/android-arm64/-/android-arm64-0.28.1.tgz", + "integrity": "sha512-34EGEbCIAgosYz6goLcopX6Mo7NyGv9tfwEM2/7Ce2VcVRk568iSvniGWcUXIy7wEDR1wzolcxcriFVrWYcwBg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/android-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/android-x64/-/android-x64-0.28.1.tgz", + "integrity": "sha512-dbwY7ltSMDWsRatcRpCnES4F+im88OCUgGZjy52shC7GqHRE/cYlxNbB4Z4UpJswpcc4Qxd2oE/ufM0p61IKng==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/darwin-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/darwin-arm64/-/darwin-arm64-0.28.1.tgz", + "integrity": "sha512-TZbWkQY7kvTAXbXUT7uVACR5cMHsDiSz9z7ZKAX/RTq/WJEk3QyRr0wZpNhBDX+/0CtdqUIJlOiodQcta6tY3Q==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/darwin-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/darwin-x64/-/darwin-x64-0.28.1.tgz", + "integrity": "sha512-zfdzgK9ACBNZLI/CyHTOx81SyNbM6YXn7rxSgX97VjyiPl9W1i4Ka4fgKECEoFCKGpvBj5qArWIGgQjOwkgskQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/freebsd-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/freebsd-arm64/-/freebsd-arm64-0.28.1.tgz", + "integrity": "sha512-wG2EA8ENdEI0qhkSZMjfqrdY+ziCYCPMmtZjjIwOmXFjmyzEHn+UUxk5of+SYsjtfs3VpnlC7QLzSI5hY/rOAw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/freebsd-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/freebsd-x64/-/freebsd-x64-0.28.1.tgz", + "integrity": "sha512-i7dZ9vQgnvSCzi/rYCXNgtF/U+eKZNJBzu3eTQbRgHnM7tNSizLOkRFAl3qzVc/Op/u5YkHHa4pf/3DOYHthLQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-arm": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-arm/-/linux-arm-0.28.1.tgz", + "integrity": "sha512-qVXBOHQS+d5Y722GwJzJUtOLlX7km3CraOaGormF1pDtPd2C/l1SHRPgjLunLGe51Sh5YYWKMFDyV4SxgMQYTQ==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-arm64/-/linux-arm64-0.28.1.tgz", + "integrity": "sha512-yHs+0uc8+nvEAfAfxrWQKK5peSNzBc4PegcMO0EJ2hT71uA7vB8Ihg2e77R2P7SG5uYjPbHlLLmve4LLLRCf0g==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-ia32": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-ia32/-/linux-ia32-0.28.1.tgz", + "integrity": "sha512-d1z4ZuP0ajrfz/FhGT4vv278rX8KnPPJx8i5+AtK7TYbx9Le9F1hyzurZpkEyjkGa9dUGhQow4C1NmeGvqxN2w==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-loong64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-loong64/-/linux-loong64-0.28.1.tgz", + "integrity": "sha512-M5sRjUVZrkm1OAPR3dlOYzNmN+loZKGVi1VUQGrwuqLcbR6qeAz+famMhjASeH3YVKvZz+zT1jlh/keC3Rj/lg==", + "cpu": [ + "loong64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-mips64el": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-mips64el/-/linux-mips64el-0.28.1.tgz", + "integrity": "sha512-mRObBZeHh2OxcBFPWE/FjylkRgZdYuiTR3vaTozquCGOH14iP9oN4x4Ge81CoIDYQrXmIxpFumJBu5MtZpnQJQ==", + "cpu": [ + "mips64el" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-ppc64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-ppc64/-/linux-ppc64-0.28.1.tgz", + "integrity": "sha512-slScBsMAb3GFDcdrCgLwZtPYRoH2H/youv10QiZyRjmsP48fznoveWytSgCI/R0ZcUgpc0ZhIUEx6LHts8yrfQ==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-riscv64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-riscv64/-/linux-riscv64-0.28.1.tgz", + "integrity": "sha512-kw0owk1o0GFETUJyW0jc0G4Yzs0BHZn0JDZ8JRT088vjJYX777BAs1fDGxAC+q831qOs2DTC96mNsG2opdfyyQ==", + "cpu": [ + "riscv64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-s390x": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-s390x/-/linux-s390x-0.28.1.tgz", + "integrity": "sha512-/lAIjX8aYFRByhh6L5rYtPEDRqa9de/4V/juOXcta5frjvzXO4/sqEtyytse0g3zZFuWu5cDN0MkLz2qRDD2Ag==", + "cpu": [ + "s390x" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/linux-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.28.1.tgz", + "integrity": "sha512-u/anNYF2mmVOEDwLtnQ1wOr3EZ9sTNGLWrsYGYwHWzGA3Si84IOkHXlbWTD1NB+9/1lcnweYKO54uhxZydNzfA==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/netbsd-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/netbsd-arm64/-/netbsd-arm64-0.28.1.tgz", + "integrity": "sha512-oks0DYbLwWMmaakTsCb+zL4E+aHRVLom9IJZOAthMQEPiQmydXHkziYEsGYRx0uNV/IjEKGAV941JzH02pflqw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "netbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/netbsd-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/netbsd-x64/-/netbsd-x64-0.28.1.tgz", + "integrity": "sha512-aeL6lAnN89Hz43Mlh1G8ARasbuoYvSITDEx0tHh5b7jJnHcssqgjy9Yx430GDpmCa6OyrKoS0aNRjKundRizGg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "netbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openbsd-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/openbsd-arm64/-/openbsd-arm64-0.28.1.tgz", + "integrity": "sha512-MEFJe5C3R8pwXdZ5Y21oo6m7ePiS0d9pWucn99O/wvyJZChoIQKrQDxKrGeW8F5+T0okTHesAmDeiHDTIq0V/Q==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openbsd-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/openbsd-x64/-/openbsd-x64-0.28.1.tgz", + "integrity": "sha512-i/ZLIOafE0Z8cI/XANJAixoJL/uRAoS2xOA3rb0xN+KK0K177cMAsQYkzHtBrtMXAKuAc7HGgcWiZ/sRC1Nxgw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openbsd" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/openharmony-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/openharmony-arm64/-/openharmony-arm64-0.28.1.tgz", + "integrity": "sha512-ge+Z7EXFNt2BO1oAMsVpiQ8EwndV9i1xXerAeTIK7AtPs3bKFXQM7nlRxDSIUIMeueR1CNXxqztLzdNeReKBJg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/sunos-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/sunos-x64/-/sunos-x64-0.28.1.tgz", + "integrity": "sha512-BEjgtECkL3vY+SaSQ6nzVfiALUeFxpawyp8Jmf5PtYhf1Ug40N1h/hxlhts+f1FvSvarEigdxS3BlSMI2PJLcQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "sunos" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-arm64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/win32-arm64/-/win32-arm64-0.28.1.tgz", + "integrity": "sha512-lCv9eK/H6ZJWbE7bh2nw54CZ9M2nupBxJcTsdk/QQnWkdSjKGuxmmH8/GWrlT1eMmZfn4dGcCjRte397WqfQXA==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-ia32": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/win32-ia32/-/win32-ia32-0.28.1.tgz", + "integrity": "sha512-zvb/mB2bSCoJOpoCBgYKKpX6YM6mJBlBUVUtVj41DlZJVEB6/0CKlRYxP5wWl1C1ILiCoAU5wZZ4q1P3qeS6Eg==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@esbuild/win32-x64": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.28.1.tgz", + "integrity": "sha512-bm4Mowrv+GXMlpWX++EcXw/iLyd1o3+bJkC2DkWXYVvgZCqD/bSj9ctZeAMC3cIxgjRVR2Dufaiu4YPxr5gW1A==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">=18" + } + }, + "node_modules/@img/colour": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@img/colour/-/colour-1.1.0.tgz", + "integrity": "sha512-Td76q7j57o/tLVdgS746cYARfSyxk8iEfRxewL9h4OMzYhbW4TAcppl0mT4eyqXddh6L/jwoM75mo7ixa/pCeQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + } + }, + "node_modules/@img/sharp-darwin-arm64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-darwin-arm64/-/sharp-darwin-arm64-0.34.5.tgz", + "integrity": "sha512-imtQ3WMJXbMY4fxb/Ndp6HBTNVtWCUI0WdobyheGf5+ad6xX8VIDO8u2xE4qc/fr08CKG/7dDseFtn6M6g/r3w==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-darwin-arm64": "1.2.4" + } + }, + "node_modules/@img/sharp-darwin-x64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-darwin-x64/-/sharp-darwin-x64-0.34.5.tgz", + "integrity": "sha512-YNEFAF/4KQ/PeW0N+r+aVVsoIY0/qxxikF2SWdp+NRkmMB7y9LBZAVqQ4yhGCm/H3H270OSykqmQMKLBhBJDEw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-darwin-x64": "1.2.4" + } + }, + "node_modules/@img/sharp-libvips-darwin-arm64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-arm64/-/sharp-libvips-darwin-arm64-1.2.4.tgz", + "integrity": "sha512-zqjjo7RatFfFoP0MkQ51jfuFZBnVE2pRiaydKJ1G/rHZvnsrHAOcQALIi9sA5co5xenQdTugCvtb1cuf78Vf4g==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "darwin" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-darwin-x64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-darwin-x64/-/sharp-libvips-darwin-x64-1.2.4.tgz", + "integrity": "sha512-1IOd5xfVhlGwX+zXv2N93k0yMONvUlANylbJw1eTah8K/Jtpi15KC+WSiaX/nBmbm2HxRM1gZ0nSdjSsrZbGKg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "darwin" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-arm": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm/-/sharp-libvips-linux-arm-1.2.4.tgz", + "integrity": "sha512-bFI7xcKFELdiNCVov8e44Ia4u2byA+l3XtsAj+Q8tfCwO6BQ8iDojYdvoPMqsKDkuoOo+X6HZA0s0q11ANMQ8A==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-arm64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-arm64/-/sharp-libvips-linux-arm64-1.2.4.tgz", + "integrity": "sha512-excjX8DfsIcJ10x1Kzr4RcWe1edC9PquDRRPx3YVCvQv+U5p7Yin2s32ftzikXojb1PIFc/9Mt28/y+iRklkrw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-ppc64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-ppc64/-/sharp-libvips-linux-ppc64-1.2.4.tgz", + "integrity": "sha512-FMuvGijLDYG6lW+b/UvyilUWu5Ayu+3r2d1S8notiGCIyYU/76eig1UfMmkZ7vwgOrzKzlQbFSuQfgm7GYUPpA==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-riscv64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-riscv64/-/sharp-libvips-linux-riscv64-1.2.4.tgz", + "integrity": "sha512-oVDbcR4zUC0ce82teubSm+x6ETixtKZBh/qbREIOcI3cULzDyb18Sr/Wcyx7NRQeQzOiHTNbZFF1UwPS2scyGA==", + "cpu": [ + "riscv64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-s390x": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-s390x/-/sharp-libvips-linux-s390x-1.2.4.tgz", + "integrity": "sha512-qmp9VrzgPgMoGZyPvrQHqk02uyjA0/QrTO26Tqk6l4ZV0MPWIW6LTkqOIov+J1yEu7MbFQaDpwdwJKhbJvuRxQ==", + "cpu": [ + "s390x" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linux-x64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linux-x64/-/sharp-libvips-linux-x64-1.2.4.tgz", + "integrity": "sha512-tJxiiLsmHc9Ax1bz3oaOYBURTXGIRDODBqhveVHonrHJ9/+k89qbLl0bcJns+e4t4rvaNBxaEZsFtSfAdquPrw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linuxmusl-arm64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-arm64/-/sharp-libvips-linuxmusl-arm64-1.2.4.tgz", + "integrity": "sha512-FVQHuwx1IIuNow9QAbYUzJ+En8KcVm9Lk5+uGUQJHaZmMECZmOlix9HnH7n1TRkXMS0pGxIJokIVB9SuqZGGXw==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-libvips-linuxmusl-x64": { + "version": "1.2.4", + "resolved": "https://registry.npmjs.org/@img/sharp-libvips-linuxmusl-x64/-/sharp-libvips-linuxmusl-x64-1.2.4.tgz", + "integrity": "sha512-+LpyBk7L44ZIXwz/VYfglaX/okxezESc6UxDSoyo2Ks6Jxc4Y7sGjpgU9s4PMgqgjj1gZCylTieNamqA1MF7Dg==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "LGPL-3.0-or-later", + "optional": true, + "os": [ + "linux" + ], + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-linux-arm": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm/-/sharp-linux-arm-0.34.5.tgz", + "integrity": "sha512-9dLqsvwtg1uuXBGZKsxem9595+ujv0sJ6Vi8wcTANSFpwV/GONat5eCkzQo/1O6zRIkh0m/8+5BjrRr7jDUSZw==", + "cpu": [ + "arm" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-arm": "1.2.4" + } + }, + "node_modules/@img/sharp-linux-arm64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-arm64/-/sharp-linux-arm64-0.34.5.tgz", + "integrity": "sha512-bKQzaJRY/bkPOXyKx5EVup7qkaojECG6NLYswgktOZjaXecSAeCWiZwwiFf3/Y+O1HrauiE3FVsGxFg8c24rZg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-arm64": "1.2.4" + } + }, + "node_modules/@img/sharp-linux-ppc64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-ppc64/-/sharp-linux-ppc64-0.34.5.tgz", + "integrity": "sha512-7zznwNaqW6YtsfrGGDA6BRkISKAAE1Jo0QdpNYXNMHu2+0dTrPflTLNkpc8l7MUP5M16ZJcUvysVWWrMefZquA==", + "cpu": [ + "ppc64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-ppc64": "1.2.4" + } + }, + "node_modules/@img/sharp-linux-riscv64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-riscv64/-/sharp-linux-riscv64-0.34.5.tgz", + "integrity": "sha512-51gJuLPTKa7piYPaVs8GmByo7/U7/7TZOq+cnXJIHZKavIRHAP77e3N2HEl3dgiqdD/w0yUfiJnII77PuDDFdw==", + "cpu": [ + "riscv64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-riscv64": "1.2.4" + } + }, + "node_modules/@img/sharp-linux-s390x": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-s390x/-/sharp-linux-s390x-0.34.5.tgz", + "integrity": "sha512-nQtCk0PdKfho3eC5MrbQoigJ2gd1CgddUMkabUj+rBevs8tZ2cULOx46E7oyX+04WGfABgIwmMC0VqieTiR4jg==", + "cpu": [ + "s390x" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-s390x": "1.2.4" + } + }, + "node_modules/@img/sharp-linux-x64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linux-x64/-/sharp-linux-x64-0.34.5.tgz", + "integrity": "sha512-MEzd8HPKxVxVenwAa+JRPwEC7QFjoPWuS5NZnBt6B3pu7EG2Ge0id1oLHZpPJdn3OQK+BQDiw9zStiHBTJQQQQ==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linux-x64": "1.2.4" + } + }, + "node_modules/@img/sharp-linuxmusl-arm64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-arm64/-/sharp-linuxmusl-arm64-0.34.5.tgz", + "integrity": "sha512-fprJR6GtRsMt6Kyfq44IsChVZeGN97gTD331weR1ex1c1rypDEABN6Tm2xa1wE6lYb5DdEnk03NZPqA7Id21yg==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linuxmusl-arm64": "1.2.4" + } + }, + "node_modules/@img/sharp-linuxmusl-x64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-linuxmusl-x64/-/sharp-linuxmusl-x64-0.34.5.tgz", + "integrity": "sha512-Jg8wNT1MUzIvhBFxViqrEhWDGzqymo3sV7z7ZsaWbZNDLXRJZoRGrjulp60YYtV4wfY8VIKcWidjojlLcWrd8Q==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-libvips-linuxmusl-x64": "1.2.4" + } + }, + "node_modules/@img/sharp-wasm32": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-wasm32/-/sharp-wasm32-0.34.5.tgz", + "integrity": "sha512-OdWTEiVkY2PHwqkbBI8frFxQQFekHaSSkUIJkwzclWZe64O1X4UlUjqqqLaPbUpMOQk6FBu/HtlGXNblIs0huw==", + "cpu": [ + "wasm32" + ], + "dev": true, + "license": "Apache-2.0 AND LGPL-3.0-or-later AND MIT", + "optional": true, + "dependencies": { + "@emnapi/runtime": "^1.7.0" + }, + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-win32-arm64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-arm64/-/sharp-win32-arm64-0.34.5.tgz", + "integrity": "sha512-WQ3AgWCWYSb2yt+IG8mnC6Jdk9Whs7O0gxphblsLvdhSpSTtmu69ZG1Gkb6NuvxsNACwiPV6cNSZNzt0KPsw7g==", + "cpu": [ + "arm64" + ], + "dev": true, + "license": "Apache-2.0 AND LGPL-3.0-or-later", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-win32-ia32": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-ia32/-/sharp-win32-ia32-0.34.5.tgz", + "integrity": "sha512-FV9m/7NmeCmSHDD5j4+4pNI8Cp3aW+JvLoXcTUo0IqyjSfAZJ8dIUmijx1qaJsIiU+Hosw6xM5KijAWRJCSgNg==", + "cpu": [ + "ia32" + ], + "dev": true, + "license": "Apache-2.0 AND LGPL-3.0-or-later", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@img/sharp-win32-x64": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/@img/sharp-win32-x64/-/sharp-win32-x64-0.34.5.tgz", + "integrity": "sha512-+29YMsqY2/9eFEiW93eqWnuLcWcufowXewwSNIT6UwZdUUCrM3oFjMWH/Z6/TMmb4hlFenmfAVbpWeup2jryCw==", + "cpu": [ + "x64" + ], + "dev": true, + "license": "Apache-2.0 AND LGPL-3.0-or-later", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + } + }, + "node_modules/@jridgewell/resolve-uri": { + "version": "3.1.2", + "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", + "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=6.0.0" + } + }, + "node_modules/@jridgewell/sourcemap-codec": { + "version": "1.5.5", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", + "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", + "dev": true, + "license": "MIT" + }, + "node_modules/@jridgewell/trace-mapping": { + "version": "0.3.9", + "resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.9.tgz", + "integrity": "sha512-3Belt6tdc8bPgAtbcmdtNJlirVoTmEb5e2gC94PnkwEW9jI6CAHUeoG85tjWP5WquqfavoMtMwiG4P926ZKKuQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@jridgewell/resolve-uri": "^3.0.3", + "@jridgewell/sourcemap-codec": "^1.4.10" + } + }, + "node_modules/@poppinss/colors": { + "version": "4.1.6", + "resolved": "https://registry.npmjs.org/@poppinss/colors/-/colors-4.1.6.tgz", + "integrity": "sha512-H9xkIdFswbS8n1d6vmRd8+c10t2Qe+rZITbbDHHkQixH5+2x1FDGmi/0K+WgWiqQFKPSlIYB7jlH6Kpfn6Fleg==", + "dev": true, + "license": "MIT", + "dependencies": { + "kleur": "^4.1.5" + } + }, + "node_modules/@poppinss/dumper": { + "version": "0.6.5", + "resolved": "https://registry.npmjs.org/@poppinss/dumper/-/dumper-0.6.5.tgz", + "integrity": "sha512-NBdYIb90J7LfOI32dOewKI1r7wnkiH6m920puQ3qHUeZkxNkQiFnXVWoE6YtFSv6QOiPPf7ys6i+HWWecDz7sw==", + "dev": true, + "license": "MIT", + "dependencies": { + "@poppinss/colors": "^4.1.5", + "@sindresorhus/is": "^7.0.2", + "supports-color": "^10.0.0" + } + }, + "node_modules/@poppinss/exception": { + "version": "1.2.3", + "resolved": "https://registry.npmjs.org/@poppinss/exception/-/exception-1.2.3.tgz", + "integrity": "sha512-dCED+QRChTVatE9ibtoaxc+WkdzOSjYTKi/+uacHWIsfodVfpsueo3+DKpgU5Px8qXjgmXkSvhXvSCz3fnP9lw==", + "dev": true, + "license": "MIT" + }, + "node_modules/@sindresorhus/is": { + "version": "7.2.0", + "resolved": "https://registry.npmjs.org/@sindresorhus/is/-/is-7.2.0.tgz", + "integrity": "sha512-P1Cz1dWaFfR4IR+U13mqqiGsLFf1KbayybWwdd2vfctdV6hDpUkgCY0nKOLLTMSoRd/jJNjtbqzf13K8DCCXQw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sindresorhus/is?sponsor=1" + } + }, + "node_modules/@speed-highlight/core": { + "version": "1.2.17", + "resolved": "https://registry.npmjs.org/@speed-highlight/core/-/core-1.2.17.tgz", + "integrity": "sha512-Z92FwKpCtfaW1V0jTU/fh3QzYEZN8wDwrzRIBoADCJfn4mJCNcJN/XegifX7BDrQ8/h9Xh/JnbyMchL0FqXrkg==", + "dev": true, + "license": "CC0-1.0" + }, + "node_modules/blake3-wasm": { + "version": "2.1.5", + "resolved": "https://registry.npmjs.org/blake3-wasm/-/blake3-wasm-2.1.5.tgz", + "integrity": "sha512-F1+K8EbfOZE49dtoPtmxUQrpXaBIl3ICvasLh+nJta0xkz+9kF/7uet9fLnwKqhDrmj6g+6K3Tw9yQPUg2ka5g==", + "dev": true, + "license": "MIT" + }, + "node_modules/cookie": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/cookie/-/cookie-1.1.1.tgz", + "integrity": "sha512-ei8Aos7ja0weRpFzJnEA9UHJ/7XQmqglbRwnf2ATjcB9Wq874VKH9kfjjirM6UhU2/E5fFYadylyhFldcqSidQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/detect-libc": { + "version": "2.1.2", + "resolved": "https://registry.npmjs.org/detect-libc/-/detect-libc-2.1.2.tgz", + "integrity": "sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==", + "dev": true, + "license": "Apache-2.0", + "engines": { + "node": ">=8" + } + }, + "node_modules/error-stack-parser-es": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/error-stack-parser-es/-/error-stack-parser-es-1.0.5.tgz", + "integrity": "sha512-5qucVt2XcuGMcEGgWI7i+yZpmpByQ8J1lHhcL7PwqCwu9FPP3VUXzT4ltHe5i2z9dePwEHcDVOAfSnHsOlCXRA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/antfu" + } + }, + "node_modules/esbuild": { + "version": "0.28.1", + "resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.28.1.tgz", + "integrity": "sha512-HrJrvZv5ayxBzPfwphOoNzkzOIIlifzk0KJrGK2c8R4+LKpMtpYLQeUdjnwjWv/LZlkH2laZk+4w78pi99D4Vw==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "bin": { + "esbuild": "bin/esbuild" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "@esbuild/aix-ppc64": "0.28.1", + "@esbuild/android-arm": "0.28.1", + "@esbuild/android-arm64": "0.28.1", + "@esbuild/android-x64": "0.28.1", + "@esbuild/darwin-arm64": "0.28.1", + "@esbuild/darwin-x64": "0.28.1", + "@esbuild/freebsd-arm64": "0.28.1", + "@esbuild/freebsd-x64": "0.28.1", + "@esbuild/linux-arm": "0.28.1", + "@esbuild/linux-arm64": "0.28.1", + "@esbuild/linux-ia32": "0.28.1", + "@esbuild/linux-loong64": "0.28.1", + "@esbuild/linux-mips64el": "0.28.1", + "@esbuild/linux-ppc64": "0.28.1", + "@esbuild/linux-riscv64": "0.28.1", + "@esbuild/linux-s390x": "0.28.1", + "@esbuild/linux-x64": "0.28.1", + "@esbuild/netbsd-arm64": "0.28.1", + "@esbuild/netbsd-x64": "0.28.1", + "@esbuild/openbsd-arm64": "0.28.1", + "@esbuild/openbsd-x64": "0.28.1", + "@esbuild/openharmony-arm64": "0.28.1", + "@esbuild/sunos-x64": "0.28.1", + "@esbuild/win32-arm64": "0.28.1", + "@esbuild/win32-ia32": "0.28.1", + "@esbuild/win32-x64": "0.28.1" + } + }, + "node_modules/fsevents": { + "version": "2.3.3", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", + "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, + "node_modules/kleur": { + "version": "4.1.5", + "resolved": "https://registry.npmjs.org/kleur/-/kleur-4.1.5.tgz", + "integrity": "sha512-o+NO+8WrRiQEE4/7nwRJhN1HWpVmJm511pBHUxPLtp0BUISzlBplORYSmTclCnJvQq2tKu/sgl3xVpkc7ZWuQQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=6" + } + }, + "node_modules/miniflare": { + "version": "4.20260708.1", + "resolved": "https://registry.npmjs.org/miniflare/-/miniflare-4.20260708.1.tgz", + "integrity": "sha512-c94O9zRDISdqO18EHt6l0iF/fWgWt8p18PJvRsA/L/NJZ9Cfke3s/F5Blg1XXF7WDutVRzWVWy8Vy4LaT5ifsA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@cspotcode/source-map-support": "0.8.1", + "sharp": "0.34.5", + "undici": "7.28.0", + "workerd": "1.20260708.1", + "ws": "8.21.0", + "youch": "4.1.0-beta.10" + }, + "bin": { + "miniflare": "bootstrap.js" + }, + "engines": { + "node": ">=22.0.0" + } + }, + "node_modules/path-to-regexp": { + "version": "6.3.0", + "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-6.3.0.tgz", + "integrity": "sha512-Yhpw4T9C6hPpgPeA28us07OJeqZ5EzQTkbfwuhsUg0c237RomFoETJgmp2sa3F/41gfLE6G5cqcYwznmeEeOlQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/pathe": { + "version": "2.0.3", + "resolved": "https://registry.npmjs.org/pathe/-/pathe-2.0.3.tgz", + "integrity": "sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==", + "dev": true, + "license": "MIT" + }, + "node_modules/semver": { + "version": "7.8.5", + "resolved": "https://registry.npmjs.org/semver/-/semver-7.8.5.tgz", + "integrity": "sha512-Y7/KDsb8LjooZpwaqGyulO6DQlksgCncchHGk+sZIY4SBvUocMBEFH5Ur1fI4dV+Jvl0w6cjvucaIi40puRioA==", + "dev": true, + "license": "ISC", + "bin": { + "semver": "bin/semver.js" + }, + "engines": { + "node": ">=10" + } + }, + "node_modules/sharp": { + "version": "0.34.5", + "resolved": "https://registry.npmjs.org/sharp/-/sharp-0.34.5.tgz", + "integrity": "sha512-Ou9I5Ft9WNcCbXrU9cMgPBcCK8LiwLqcbywW3t4oDV37n1pzpuNLsYiAV8eODnjbtQlSDwZ2cUEeQz4E54Hltg==", + "dev": true, + "hasInstallScript": true, + "license": "Apache-2.0", + "dependencies": { + "@img/colour": "^1.0.0", + "detect-libc": "^2.1.2", + "semver": "^7.7.3" + }, + "engines": { + "node": "^18.17.0 || ^20.3.0 || >=21.0.0" + }, + "funding": { + "url": "https://opencollective.com/libvips" + }, + "optionalDependencies": { + "@img/sharp-darwin-arm64": "0.34.5", + "@img/sharp-darwin-x64": "0.34.5", + "@img/sharp-libvips-darwin-arm64": "1.2.4", + "@img/sharp-libvips-darwin-x64": "1.2.4", + "@img/sharp-libvips-linux-arm": "1.2.4", + "@img/sharp-libvips-linux-arm64": "1.2.4", + "@img/sharp-libvips-linux-ppc64": "1.2.4", + "@img/sharp-libvips-linux-riscv64": "1.2.4", + "@img/sharp-libvips-linux-s390x": "1.2.4", + "@img/sharp-libvips-linux-x64": "1.2.4", + "@img/sharp-libvips-linuxmusl-arm64": "1.2.4", + "@img/sharp-libvips-linuxmusl-x64": "1.2.4", + "@img/sharp-linux-arm": "0.34.5", + "@img/sharp-linux-arm64": "0.34.5", + "@img/sharp-linux-ppc64": "0.34.5", + "@img/sharp-linux-riscv64": "0.34.5", + "@img/sharp-linux-s390x": "0.34.5", + "@img/sharp-linux-x64": "0.34.5", + "@img/sharp-linuxmusl-arm64": "0.34.5", + "@img/sharp-linuxmusl-x64": "0.34.5", + "@img/sharp-wasm32": "0.34.5", + "@img/sharp-win32-arm64": "0.34.5", + "@img/sharp-win32-ia32": "0.34.5", + "@img/sharp-win32-x64": "0.34.5" + } + }, + "node_modules/supports-color": { + "version": "10.2.2", + "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-10.2.2.tgz", + "integrity": "sha512-SS+jx45GF1QjgEXQx4NJZV9ImqmO2NPz5FNsIHrsDjh2YsHnawpan7SNQ1o8NuhrbHZy9AZhIoCUiCeaW/C80g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/chalk/supports-color?sponsor=1" + } + }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "dev": true, + "license": "0BSD", + "optional": true + }, + "node_modules/undici": { + "version": "7.28.0", + "resolved": "https://registry.npmjs.org/undici/-/undici-7.28.0.tgz", + "integrity": "sha512-cRZYrTDwWznlnRiPjggAGxZXanty6M8RV1ff8Wm4LWXBp7/IG8v5DnOm74DtUBp9OONpK75YlPnIjQqX0dBDtA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=20.18.1" + } + }, + "node_modules/unenv": { + "version": "2.0.0-rc.24", + "resolved": "https://registry.npmjs.org/unenv/-/unenv-2.0.0-rc.24.tgz", + "integrity": "sha512-i7qRCmY42zmCwnYlh9H2SvLEypEFGye5iRmEMKjcGi7zk9UquigRjFtTLz0TYqr0ZGLZhaMHl/foy1bZR+Cwlw==", + "dev": true, + "license": "MIT", + "dependencies": { + "pathe": "^2.0.3" + } + }, + "node_modules/workerd": { + "version": "1.20260708.1", + "resolved": "https://registry.npmjs.org/workerd/-/workerd-1.20260708.1.tgz", + "integrity": "sha512-WAK+Kt/VVCSldH2qSr8lx46XCJ4Q+bdlHNaFqUtOHthBEIB8C1N8HVW+VOLrxDoTCk0NGNv0zajnBeQK4JOB9w==", + "dev": true, + "hasInstallScript": true, + "license": "Apache-2.0", + "bin": { + "workerd": "bin/workerd" + }, + "engines": { + "node": ">=16" + }, + "optionalDependencies": { + "@cloudflare/workerd-darwin-64": "1.20260708.1", + "@cloudflare/workerd-darwin-arm64": "1.20260708.1", + "@cloudflare/workerd-linux-64": "1.20260708.1", + "@cloudflare/workerd-linux-arm64": "1.20260708.1", + "@cloudflare/workerd-windows-64": "1.20260708.1" + } + }, + "node_modules/wrangler": { + "version": "4.110.0", + "resolved": "https://registry.npmjs.org/wrangler/-/wrangler-4.110.0.tgz", + "integrity": "sha512-xZeXKYi7hxQRF5anL+v77RkufJNpF9f3Eqeyqq2QBsETpLZgh0Agj0jJ6JPtkbgn6ukZdh8OK5egsGPWIditgg==", + "dev": true, + "license": "MIT OR Apache-2.0", + "dependencies": { + "@cloudflare/kv-asset-handler": "0.5.0", + "@cloudflare/unenv-preset": "2.16.1", + "blake3-wasm": "2.1.5", + "esbuild": "0.28.1", + "miniflare": "4.20260708.1", + "path-to-regexp": "6.3.0", + "unenv": "2.0.0-rc.24", + "workerd": "1.20260708.1" + }, + "bin": { + "cf-wrangler": "bin/cf-wrangler.js", + "wrangler": "bin/wrangler.js", + "wrangler2": "bin/wrangler.js" + }, + "engines": { + "node": ">=22.0.0" + }, + "optionalDependencies": { + "fsevents": "2.3.3" + }, + "peerDependencies": { + "@cloudflare/workers-types": "^5.20260708.1" + }, + "peerDependenciesMeta": { + "@cloudflare/workers-types": { + "optional": true + } + } + }, + "node_modules/ws": { + "version": "8.21.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", + "integrity": "sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, + "node_modules/youch": { + "version": "4.1.0-beta.10", + "resolved": "https://registry.npmjs.org/youch/-/youch-4.1.0-beta.10.tgz", + "integrity": "sha512-rLfVLB4FgQneDr0dv1oddCVZmKjcJ6yX6mS4pU82Mq/Dt9a3cLZQ62pDBL4AUO+uVrCvtWz3ZFUL2HFAFJ/BXQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@poppinss/colors": "^4.1.5", + "@poppinss/dumper": "^0.6.4", + "@speed-highlight/core": "^1.2.7", + "cookie": "^1.0.2", + "youch-core": "^0.3.3" + } + }, + "node_modules/youch-core": { + "version": "0.3.3", + "resolved": "https://registry.npmjs.org/youch-core/-/youch-core-0.3.3.tgz", + "integrity": "sha512-ho7XuGjLaJ2hWHoK8yFnsUGy2Y5uDpqSTq1FkHLK4/oqKtyUU1AFbOOxY4IpC9f0fTLjwYbslUz0Po5BpD1wrA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@poppinss/exception": "^1.2.2", + "error-stack-parser-es": "^1.0.5" + } + } + } +} diff --git a/deploy/cloudflare-worker/package.json b/deploy/cloudflare-worker/package.json new file mode 100644 index 00000000..82c5c8c5 --- /dev/null +++ b/deploy/cloudflare-worker/package.json @@ -0,0 +1,11 @@ +{ + "name": "kakeya-inference-network-worker", + "private": true, + "scripts": { + "deploy": "wrangler deploy", + "check": "wrangler deploy --dry-run" + }, + "devDependencies": { + "wrangler": "^4.0.0" + } +} diff --git a/deploy/cloudflare-worker/src/index.js b/deploy/cloudflare-worker/src/index.js new file mode 100644 index 00000000..4e93c848 --- /dev/null +++ b/deploy/cloudflare-worker/src/index.js @@ -0,0 +1,33 @@ +const ORIGIN = "https://agent.kakeya.ai"; + +export default { + async fetch(request) { + const incoming = new URL(request.url); + let path = incoming.pathname; + + if (path === "/" || path === "/index.html") { + path = "/network"; + } else if (path === "/healthz") { + path = "/v1/network/summary"; + } else if ( + path !== "/network" && + !path.startsWith("/v1/network/") + ) { + path = "/network"; + } + + const target = new URL(path + incoming.search, ORIGIN); + const upstreamRequest = new Request(target, request); + const response = await fetch(upstreamRequest); + const headers = new Headers(response.headers); + headers.set("X-Kakeya-Surface", "inference-network"); + headers.set("X-Content-Type-Options", "nosniff"); + headers.set("Referrer-Policy", "same-origin"); + + return new Response(response.body, { + status: response.status, + statusText: response.statusText, + headers, + }); + }, +}; diff --git a/deploy/cloudflare-worker/wrangler.jsonc b/deploy/cloudflare-worker/wrangler.jsonc new file mode 100644 index 00000000..f5677776 --- /dev/null +++ b/deploy/cloudflare-worker/wrangler.jsonc @@ -0,0 +1,13 @@ +{ + "$schema": "node_modules/wrangler/config-schema.json", + "name": "kakeya-inference-network", + "main": "src/index.js", + "compatibility_date": "2026-07-11", + "account_id": "c54b05200639ec14333bba0643416723", + "routes": [ + { + "pattern": "kakeya.ai/*", + "zone_name": "kakeya.ai" + } + ] +} diff --git a/deploy/install_prefill_network_launchd.sh b/deploy/install_prefill_network_launchd.sh new file mode 100755 index 00000000..d400d2ba --- /dev/null +++ b/deploy/install_prefill_network_launchd.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +HOME_DIR="${HOME}" +VENV="${VENV:-$HOME/.venv-distwan}" +KEY_FILE="${KEY_FILE:-$HOME/.kakeya/network_api_key}" +TEMPLATE="$ROOT/deploy/launchd/ai.kakeya.prefill-network-head.plist" +TARGET="$HOME/Library/LaunchAgents/ai.kakeya.prefill-network.plist" + +[ -x "$VENV/bin/python" ] || { echo "missing $VENV/bin/python" >&2; exit 2; } +[ -s "$KEY_FILE" ] || { echo "missing network API key: $KEY_FILE" >&2; exit 2; } +mkdir -p "$HOME/.kakeya" "$HOME/Library/LaunchAgents" + +API_KEY="$(tr -d '\n' < "$KEY_FILE")" +python3 - "$TEMPLATE" "$TARGET" "$ROOT" "$HOME_DIR" "$VENV" "$API_KEY" <<'PY' +from pathlib import Path +import sys +source, target, repo, home, venv, key = sys.argv[1:] +text = Path(source).read_text() +for old, new in { + "__REPO__": repo, + "__HOME__": home, + "__VENV__": venv, + "__PYTHON__": f"{venv}/bin/python", + "__API_KEY__": key, +}.items(): + text = text.replace(old, new) +Path(target).write_text(text) +PY +chmod 600 "$TARGET" + +launchctl bootout "gui/$(id -u)/ai.kakeya.prefill-network" 2>/dev/null || true +pkill -f start_prefill_cache_node.py 2>/dev/null || true +launchctl bootstrap "gui/$(id -u)" "$TARGET" +launchctl kickstart -k "gui/$(id -u)/ai.kakeya.prefill-network" +echo "installed ai.kakeya.prefill-network" diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist new file mode 100644 index 00000000..c1237a60 --- /dev/null +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -0,0 +1,49 @@ + + + + Labelai.kakeya.grpc-runtime-prefill + ProgramArguments + /usr/bin/caffeinate-dimsu + /Users/fluffy314/.venv-distwan/bin/python + /Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-prefill-kv-network/scripts/start_grpc_runtime_server.py + --backendmlx + --verifier-id/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit + --bind127.0.0.1:51051 + --capacity1 + --sink4 + --window2048 + --max-concurrent-rpcs4 + --skip-cache-check + --enable-prefill-cache + --prefill-cache-gb1 + --cache-peer169.254.27.104:52051 + --cache-model-idgemma-4-26B-A4B-it-mlx-4bit + --model-revisionlocal-4bit-v1 + --tokenizer-revisiongemma4-v1 + --cache-quantization4bit-mlx + --cache-kv-dtypebfloat16 + --cache-block-tokens64 + --node-idhead-runtime + --advertise169.254.187.239:51051 + --cache-advertise169.254.187.239:51051 + --network-labelthunderbolt + --network-priority100 + --measured-rtt-ms0.55 + --network-telemetry-urlhttp://127.0.0.1:8090/v1/network/telemetry/tokens + --network-telemetry-api-key__NETWORK_KEY__ + --log-levelINFO + + WorkingDirectory/Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-prefill-kv-network + EnvironmentVariables + PATH/Users/fluffy314/.venv-distwan/bin:/usr/bin:/bin:/usr/sbin:/sbin + PYTHONPATH/Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-prefill-kv-network:/Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-prefill-kv-network/sdks/python + HF_HUB_OFFLINE0 + TRANSFORMERS_OFFLINE0 + + RunAtLoad + KeepAlive + ProcessTypeInteractive + StandardOutPath/Users/fluffy314/.kakeya/grpc-runtime-prefill.log + StandardErrorPath/Users/fluffy314/.kakeya/grpc-runtime-prefill.log + diff --git a/deploy/launchd/ai.kakeya.prefill-network-head.plist b/deploy/launchd/ai.kakeya.prefill-network-head.plist new file mode 100644 index 00000000..e110f6d3 --- /dev/null +++ b/deploy/launchd/ai.kakeya.prefill-network-head.plist @@ -0,0 +1,35 @@ + + + + Labelai.kakeya.prefill-network + ProgramArguments + __PYTHON__ + __REPO__/scripts/start_prefill_cache_node.py + --node-idhead-mini + --bind0.0.0.0:52051 + --advertise169.254.187.239:52051 + --peer169.254.27.104:52051 + --model-idgemma-4-26B-A4B-it-mlx-4bit + --model-revisionlocal-4bit-v1 + --tokenizer-revisiongemma4-v1 + --quantization4bit-mlx + --cache-gb2 + --networkthunderbolt + --priority100 + --rtt-ms0.55 + --http-port8090 + --state-path__HOME__/.kakeya/inference_network_prod.json + --api-key__API_KEY__ + + WorkingDirectory__REPO__ + EnvironmentVariables + PATH__VENV__/bin:/usr/bin:/bin:/usr/sbin:/sbin + PYTHONPATH__REPO__ + + RunAtLoad + KeepAlive + ProcessTypeInteractive + StandardOutPath__HOME__/.kakeya/prefill-network.log + StandardErrorPath__HOME__/.kakeya/prefill-network.log + diff --git a/deploy/launchd/ai.kakeya.prefill-network-peer.plist b/deploy/launchd/ai.kakeya.prefill-network-peer.plist new file mode 100644 index 00000000..e6ad1bd1 --- /dev/null +++ b/deploy/launchd/ai.kakeya.prefill-network-peer.plist @@ -0,0 +1,33 @@ + + + + Labelai.kakeya.prefill-network + ProgramArguments + /Users/allen/kakeya-prefill-venv/bin/python + /Users/allen/kakeya-prefill-network/start_prefill_cache_node.py + --node-idallens-mini + --bind169.254.27.104:52051 + --advertise169.254.27.104:52051 + --model-idgemma-4-26B-A4B-it-mlx-4bit + --model-revisionlocal-4bit-v1 + --tokenizer-revisiongemma4-v1 + --quantization4bit-mlx + --layer-geometry-hash93d9585b0f06b60bac8e1cadf50b29df1adbf086c862e61720b6127d22c30e2b + --cache-gb2 + --networkthunderbolt + --priority100 + --rtt-ms0.55 + --state-path/Users/allen/.kakeya/inference_network_peer.json + + WorkingDirectory/Users/allen/kakeya-prefill-network + EnvironmentVariables + PATH/Users/allen/kakeya-prefill-venv/bin:/usr/bin:/bin:/usr/sbin:/sbin + PYTHONPATH/Users/allen/kakeya-prefill-network + + RunAtLoad + KeepAlive + ProcessTypeInteractive + StandardOutPath/Users/allen/.kakeya/prefill-cache.log + StandardErrorPath/Users/allen/.kakeya/prefill-cache.log + diff --git a/dllm/__init__.py b/dllm/__init__.py new file mode 100644 index 00000000..f51242b9 --- /dev/null +++ b/dllm/__init__.py @@ -0,0 +1,7 @@ +"""Compatibility namespace for legacy dllm-hub remote model imports. + +The legacy Qwen diffusion checkpoint declares ``dllm`` in its Transformers +auto-map source even though Kakeya does not import runtime symbols from that +package. Keeping this empty namespace on PYTHONPATH satisfies the static +dependency check consistently in CI and source checkouts. +""" diff --git a/docs/adr/0016-distributed-prefill-kv-cache-network.md b/docs/adr/0016-distributed-prefill-kv-cache-network.md new file mode 100644 index 00000000..7f208df6 --- /dev/null +++ b/docs/adr/0016-distributed-prefill-kv-cache-network.md @@ -0,0 +1,211 @@ +# ADR 0016 — Distributed Prefill KV Cache Network + +- **Status:** Proposed / MVP implementation +- **Date:** 2026-07-11 +- **Relates to:** ADR 0008 (session runtime), ADR 0009 (capability gossip), + ADR 0013 (distributed topology), ADR 0015 (engine substrate) + +## Context + +Kakeya's bounded-memory runtime controls resident KV growth, but long prompt +prefill remains expensive. Multiple Mac minis often run compatible model +revisions and have idle unified memory that cannot become a coherent shared +address space over Thunderbolt. The useful resource is therefore not remote +RAM itself, but immutable prefill state that another node can import once. + +The initial rejected interpretation was a hot-path remote KV server queried +during every decoded token. That would add layer-by-layer network latency and +turn remote memory into distributed attention. This ADR instead keeps decode +local and places all network work before it. + +## Decision + +Kakeya nodes form a symmetric P2P fleet. Every node can be an inference head, +a prefill-cache requester, and a prefill-cache provider. The stored object is a +restorable prefill K/V snapshot; the reuse policy is chained longest-prefix +matching. + +Remote cache access happens only before decode: + +1. Tokenize the request and compute chained fixed-size block hashes. +2. Query the local cache and compatible live peers concurrently. +3. Select the longest contiguous prefix whose transfer/import cost is lower + than local prefill recomputation. +4. Transfer one immutable snapshot at the selected prefix boundary. +5. Import it, compute the missing suffix locally, then decode entirely locally. +6. Publish newly computed prefix-boundary snapshots asynchronously. + +There are no per-token remote reads and no coherent shared-memory illusion. + +The product name is **Distributed Prefill KV Cache**. “Prefill” names the +stored artifact; “prefix cache” names the chained longest-prefix lookup policy. + +## Three layers + +### Inference layer + +- `PrefixCacheStore` is an immutable, content-addressed, memory-bounded LRU. +- Each chained prefix hash maps to a complete restorable bounded-cache + checkpoint at that boundary. +- The MLX adapter serializes per-layer K/V, logical position, cached token + sequence, and next-token logits without pickle. +- `AppendTokensCoordinator` accepts an optional prefill-cache hook for cold + sessions. Cache failure is always a local-prefill fallback. + +### Network layer + +- Existing `CapabilityService` gossip remains the decentralized control plane. +- `NodeCapability` advertises interface-specific endpoints and exact cache + compatibility cards. +- `PrefillCacheService.LookupPrefix` is metadata-only. +- `PrefillCacheService.FetchBlocks` is a point-to-point streaming gRPC data + plane. +- Thunderbolt endpoints receive higher priority than LAN/Tailscale endpoints. +- Cards expire by TTL. Cache entries are immutable; reads require no + distributed lock. + +### UI and product layer + +The inference-network dashboard exposes: + +- node registration and expiring pairing tokens; +- online node inventory and coarse region distribution; +- cache discovery and pairing state; +- inference groups; +- completed and KV-assisted token totals; +- cache capacity, hit ratio, and transfer telemetry. + +Raw prompts, raw prefix hashes, exact IPs, and hardware serials are not public +UI data. + +## Compatibility tuple + +Remote K/V is accepted only when all fields match: + +- model id and exact weights revision; +- tokenizer/chat-template revision; +- quantization and K/V dtype; +- cache schema version; +- RoPE/position configuration; +- layer geometry; +- token block size. + +A mismatch is a cache miss, never a best-effort import. + +## Failure model + +- peer unavailable / lookup timeout → local prefill; +- stale card → TTL eviction; +- lease expired → local prefill or another lookup; +- incomplete stream / checksum mismatch → reject and local prefill; +- slow transfer → caller may cancel and recompute; +- node restart → cache epoch changes; stale leases are invalid. + +Remote cache availability must never determine request correctness. + +## Security + +The MVP assumes a trusted private fleet. Production pairing must bind `node_id` +to mTLS credentials or a fleet PSK. Prompt-derived block hashes should be HMACed +when membership attacks are in scope. + +## Observability + +Every node reports: + +- capability-card freshness and endpoint RTT; +- cache bytes used/free, entry count and epoch; +- local/remote lookup hit/miss counts; +- snapshot publish/fetch bytes and checksum failures; +- completed inference tokens and KV-assisted prompt tokens; +- fallback reason and local recompute count. + +The public dashboard shows aggregate/coarse data. Exact addresses, prompt +hashes and payload metadata remain administrator-only. + +## Consequences + +Positive: + +- expensive prompt prefill is reused across processes, restarts and nodes; +- the remote peer's memory becomes a useful cache tier without changing decode; +- node failure cannot change output correctness; +- the existing capability gossip and tensor codec remain the shared substrate; +- cache-only peers need no loaded model when they receive replicated snapshots. + +Costs: + +- snapshots are large and must be chunked; +- each model/tokenizer/cache revision creates a separate namespace; +- cache snapshots duplicate state at block boundaries in the MVP; +- peer memory is volatile and cold after restart; +- trusted-LAN deployment precedes production identity/auth hardening. + +## Alternatives considered + +### Per-token remote KV reads + +Rejected. Attention would require remote communication in every transformer +layer, adding latency to the autoregressive critical path. + +### Coherent shared memory over Thunderbolt + +Rejected. Thunderbolt Bridge exposes IP networking, not a cache-coherent +CPU/GPU address space or GPU-direct remote memory. + +### Exact full-prompt cache only + +Rejected as the sole policy. It is simpler but misses common-system-prompt and +shared-prefix reuse. Chained block hashes preserve causal correctness while +allowing suffix-only prefill. + +### Arbitrary block/hole reuse + +Rejected. Later K/V depends on the complete preceding token sequence and +positions. Only the longest contiguous prefix is safe. + +### SMB/NFS snapshot files + +Rejected for the serving data plane. Files remain useful for offline +checkpoints, but gRPC provides explicit framing, checksums, leases, +backpressure and cancellation. + +### Central registry + +Rejected. Existing symmetric CapabilityService gossip already provides +eventually-consistent discovery and TTL expiry without a new coordinator. + +## Rollout + +1. Enable cache services on a private two-Mac Thunderbolt fleet. +2. Run shadow lookup/publish while still computing prefill locally. +3. Compare imported continuation logits against local prefill. +4. Enable remote import with mandatory compatibility/checksum validation. +5. Add launchd supervision, token telemetry and public read-only dashboard. +6. Require mTLS/PSK before expanding beyond trusted private nodes. + +The cache feature is disabled by omitting `--enable-prefill-cache`. + +## Rollback + +Stop cache/dashboard launchd services and restart RuntimeService without the +cache flag. No data migration is needed because all entries are immutable, +volatile optimizations. The direct gateway remains the fallback public origin. + +## Evidence + +The two-Mac report is: +[`docs/reports/distributed-prefill-kv-mac-thunderbolt.md`](../reports/distributed-prefill-kv-mac-thunderbolt.md). + +On Apple M4 Macs over Thunderbolt, a 93-token Gemma 26B prompt measured 5.926 s +cold prefill and 0.061 s after remote snapshot import (approximately 97× for +that prompt). The peer served 36.4 MB across two snapshots with SHA-256 +validation. + +## Non-goals + +- combining physical RAM into one address space; +- remote attention in the per-token loop; +- arbitrary-hole K/V reuse; +- cross-model or cross-tokenizer cache conversion; +- using gossip to carry tensor payloads. diff --git a/docs/adr/README.md b/docs/adr/README.md index dfe4b723..bd7a7410 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -45,6 +45,7 @@ reader what was *not* chosen. | 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 | 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 new file mode 100644 index 00000000..50c69535 --- /dev/null +++ b/docs/ops/distributed-prefill-kv-network.md @@ -0,0 +1,172 @@ +# 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. + +## Production surfaces + +- dashboard: `https://kakeya.ai/` +- health: `https://kakeya.ai/healthz` +- nodes: `https://kakeya.ai/v1/network/nodes` +- groups: `https://kakeya.ai/v1/network/groups` +- token counters: `https://kakeya.ai/v1/network/tokens` + +Public reads expose aliases, coarse regions and aggregate metrics only. Writes +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 | + +## Compatibility lock + +All peers in one inference group must match: + +```text +model_id +model_revision +tokenizer_revision / chat template +quantization +KV dtype +cache format version +RoPE hash +layer geometry hash +block size +``` + +Changing any field creates a new cache namespace. Never convert or import a +best-effort mismatch. + +## Head runtime + +The release launchd asset is: + +```text +deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +``` + +It runs Gemma 26B MLX on `51051`, queries the peer on `52051`, asynchronously +publishes new snapshots, and reports generated/reused token telemetry to the +network API. + +Check: + +```bash +launchctl list | grep ai.kakeya.grpc-runtime-prefill +lsof -nP -iTCP:51051 -sTCP:LISTEN +tail -f ~/.kakeya/grpc-runtime-prefill.log +``` + +## Head dashboard/control node + +Install: + +```bash +openssl rand -hex 24 > ~/.kakeya/network_api_key +chmod 600 ~/.kakeya/network_api_key +bash deploy/install_prefill_network_launchd.sh +``` + +Check: + +```bash +launchctl list | grep ai.kakeya.prefill-network +curl -fsS http://127.0.0.1:8090/healthz +curl -fsS http://127.0.0.1:8090/v1/network/nodes +``` + +## Cache peer + +Use an isolated venv and copy/sync the repository package. The peer plist is: + +```text +deploy/launchd/ai.kakeya.prefill-network-peer.plist +``` + +Check from the head over Thunderbolt: + +```bash +ping -c 3 169.254.27.104 +nc -vz 169.254.27.104 52051 +``` + +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. + +## Node registration and groups + +Create a registration: + +```bash +KEY="$(cat ~/.kakeya/network_api_key)" +curl -fsS -X POST http://127.0.0.1:8090/v1/network/nodes/register \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $KEY" \ + -d '{"alias":"peer-mini","address":"169.254.27.104:52051","region":"Private","role":"cache"}' +``` + +Create a paired group: + +```bash +curl -fsS -X POST http://127.0.0.1:8090/v1/network/groups \ + -H "Content-Type: application/json" \ + -H "X-API-Key: $KEY" \ + -d '{"name":"Thunderbolt Pair","node_ids":["head-mini","allens-mini"]}' +``` + +## Health and acceptance + +Expected invariants: + +- both node cards appear within two gossip intervals; +- stale nodes disappear after TTL; +- remote lookup returns only the longest contiguous prefix; +- imported snapshot checksum and compatibility fingerprint match; +- remote failure falls back to local prefill; +- no remote RPC occurs in autoregressive decode; +- completed and KV-assisted token counters increase after live calls. + +Minimal acceptance: + +```bash +curl -fsS https://kakeya.ai/healthz +curl -fsS https://kakeya.ai/v1/network/summary +curl -fsS https://kakeya.ai/v1/network/tokens +``` + +## Rollback + +The cache is an optimization; inference correctness does not depend on it. + +1. Stop the cache services: + + ```bash + launchctl bootout "gui/$(id -u)/ai.kakeya.prefill-network" + launchctl bootout "gui/$(id -u)/ai.kakeya.grpc-runtime-prefill" + ``` + +2. Restart the previous RuntimeService without `--enable-prefill-cache`. +3. Roll back the Cloudflare Worker deployment with Wrangler versions/deployments. +4. `agent.kakeya.ai` remains available as the direct gateway origin. + +In-memory cache entries require no migration or cleanup after rollback. + +## Security before untrusted fleets + +The live MVP assumes trusted private Macs. Before accepting third-party nodes: + +- require mTLS or fleet-PSK authentication; +- bind signed node identity to `node_id`; +- HMAC prompt-derived block hashes; +- rate-limit registration, lookup and publish; +- cap block size and stream bytes before allocation; +- maintain revocation and audit logs; +- never expose raw prompts, hashes, IPs or cache payloads in the public UI. diff --git a/docs/quickstart.md b/docs/quickstart.md index b69cb9de..90b8606f 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -78,15 +78,11 @@ Same steps as Mac, plus a `torch` build matching your CUDA toolkit. ### Common pitfalls - **`ModuleNotFoundError: dllm`**: the legacy diffusion proposer (v0.2) - references a `dllm` package that v0.3 doesn't need but transformers' - static imports flag. Fix: - ```bash - python3 -c "import site, os; \ - p = os.path.join(site.getusersitepackages(), 'dllm'); \ - os.makedirs(p, exist_ok=True); \ - open(os.path.join(p, '__init__.py'), 'a').close()" - ``` - `setup_mac.sh` and `setup_cuda.sh` do this automatically. + declares a `dllm` package that Kakeya does not call at runtime but + Transformers validates statically. Current source checkouts include the empty + compatibility namespace at `dllm/__init__.py`; ensure the repository root is + on `PYTHONPATH`. Older tags can use `setup_mac.sh` / `setup_cuda.sh`, which + create the same namespace in site-packages. - **Connection refused to `huggingface.co`**: set `HF_ENDPOINT` before setup (see above) or pre-warm offline. diff --git a/docs/reports/distributed-prefill-kv-mac-thunderbolt.md b/docs/reports/distributed-prefill-kv-mac-thunderbolt.md new file mode 100644 index 00000000..1de69de6 --- /dev/null +++ b/docs/reports/distributed-prefill-kv-mac-thunderbolt.md @@ -0,0 +1,48 @@ +# Distributed Prefill KV — two Mac mini live report + +Date: 2026-07-11 + +## Topology + +- head Mac mini: Apple M4, 24 GB, `169.254.187.239` +- cache peer Mac mini: Apple M4, 16 GB, `169.254.27.104` +- link: Thunderbolt Bridge, 40 Gb/s physical, measured RTT ≈ 0.55 ms +- runtime model: local Gemma 26B-A4B MLX 4-bit +- cache block boundary: 64 tokens +- peer cache allocation: 2 GiB + +## Live result + +Prompt length: 93 tokens. + +| Run | Local cache | Peer cache | AppendTokens prefill | +|---|---:|---:|---:| +| cold | empty | empty | 5.926 s | +| after runtime restart | empty | 2 remote snapshots (36.4 MB) | 0.061 s | + +Observed prefill acceleration: approximately **97×** for this prompt. + +The remote hit imported 20,951,040 live KV bytes into the head runtime. The +peer's telemetry reported 93 tokens served. SHA-256 validation passed on +point-to-point publish/lookup/fetch tests. + +## Product surfaces + +- public dashboard: `https://kakeya.ai/` +- public summary API: `https://kakeya.ai/v1/network/summary` +- node/group API: `/v1/network/nodes`, `/v1/network/groups` +- token accounting: `/v1/network/tokens` + +Both cache services and the head Gemma runtime run under launchd with KeepAlive. +Cloudflare Worker `kakeya-inference-network` owns `kakeya.ai/*`; deployed +version at validation time: `e45f67be-721a-413d-804e-33f7e28e80d8`. + +## Known constraints + +- The allens Miniconda Python process lacks macOS Local Network permission for + outbound sockets. Head→peer lookup/publish/fetch works over Thunderbolt; + reverse peer→head gossip is disabled until that permission is granted. +- `agent.kakeya.ai` remains the direct gateway origin and rollback path. +- Cache entries are in-memory and intentionally disappear on peer restart. +- The MVP trusts the private Thunderbolt fleet. Production requires mTLS/PSK + identity binding before accepting remote tensor payloads. diff --git a/inference_engine/backends/mlx/prefill_snapshot.py b/inference_engine/backends/mlx/prefill_snapshot.py new file mode 100644 index 00000000..969b9afd --- /dev/null +++ b/inference_engine/backends/mlx/prefill_snapshot.py @@ -0,0 +1,166 @@ +"""Portable snapshot adapter for MLX prefill caches. + +Snapshots are immutable checkpoints at token-block boundaries. They contain +the current bounded per-layer K/V tensors, logical position, cached token +sequence, and optional next-token logits. The wire container is JSON metadata +plus raw tensor buffers; no pickle is used. +""" + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass +from typing import Any, Sequence + +import numpy as np + +from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.prefill_cache import compatibility_fingerprint +from inference_engine.distributed.tensor_codec import ( + WireTensor, + from_proto_fields, + mlx_to_wire, + to_proto_fields, + torch_to_wire, + wire_to_mlx, + wire_to_torch, +) + +_MAGIC = b"KPKV1" + + +@dataclass(frozen=True) +class ImportedPrefillSnapshot: + token_count: int + cached_token_ids: tuple[int, ...] + next_token_logits: Any | None + + +def export_mlx_prefill_snapshot( + cache: Sequence[Any], + *, + token_count: int, + cached_token_ids: Sequence[int], + compatibility: CacheCompatibility, + next_token_logits: Any | None = None, +) -> bytes: + """Serialize current MLX cache state at one prefix boundary.""" + if token_count <= 0: + raise ValueError("token_count must be > 0") + tensors: list[tuple[str, WireTensor, str]] = [] + for index, layer in enumerate(cache): + keys = getattr(layer, "keys", None) + values = getattr(layer, "values", None) + if keys is None or values is None: + raise ValueError(f"cache layer {index} is empty") + tensors.append((f"layer.{index}.k", mlx_to_wire(keys), "mlx")) + tensors.append((f"layer.{index}.v", mlx_to_wire(values), "mlx")) + if next_token_logits is not None: + if hasattr(next_token_logits, "detach"): + tensors.append(("next_token_logits", torch_to_wire(next_token_logits), "torch")) + else: + tensors.append(("next_token_logits", mlx_to_wire(next_token_logits), "mlx")) + return _pack( + tensors, + metadata={ + "compatibility_sha256": compatibility_fingerprint(compatibility).hex(), + "token_count": int(token_count), + "cached_token_ids": [int(token) for token in cached_token_ids], + "layer_count": len(cache), + }, + ) + + +def import_mlx_prefill_snapshot( + payload: bytes, + cache: Sequence[Any], + *, + compatibility: CacheCompatibility, +) -> ImportedPrefillSnapshot: + """Restore a snapshot into an allocated MLX cache list.""" + metadata, tensors = _unpack(payload) + expected = compatibility_fingerprint(compatibility).hex() + if metadata.get("compatibility_sha256") != expected: + raise ValueError("prefill snapshot compatibility fingerprint mismatch") + layer_count = int(metadata.get("layer_count", -1)) + if layer_count != len(cache): + raise ValueError( + f"snapshot layer_count {layer_count} != allocated cache {len(cache)}", + ) + token_count = int(metadata["token_count"]) + for index, layer in enumerate(cache): + key_wire, _ = tensors[f"layer.{index}.k"] + value_wire, _ = tensors[f"layer.{index}.v"] + keys = wire_to_mlx(key_wire) + values = wire_to_mlx(value_wire) + if hasattr(layer, "state"): + layer.state = (keys, values) + else: + layer.keys, layer.values = keys, values + if hasattr(layer, "offset"): + layer.offset = token_count + next_logits = None + if "next_token_logits" in tensors: + wire, framework = tensors["next_token_logits"] + next_logits = wire_to_torch(wire) if framework == "torch" else wire_to_mlx(wire) + return ImportedPrefillSnapshot( + token_count=token_count, + cached_token_ids=tuple(int(t) for t in metadata["cached_token_ids"]), + next_token_logits=next_logits, + ) + + +def _pack( + tensors: Sequence[tuple[str, WireTensor, str]], + *, + metadata: dict, +) -> bytes: + raw_parts: list[bytes] = [] + tensor_meta: list[dict] = [] + offset = 0 + for name, wire, framework in tensors: + dtype, shape, data = to_proto_fields(wire) + tensor_meta.append({ + "name": name, + "dtype": dtype, + "shape": shape, + "offset": offset, + "length": len(data), + "framework": framework, + }) + raw_parts.append(data) + offset += len(data) + header = json.dumps( + {**metadata, "tensors": tensor_meta}, + sort_keys=True, + separators=(",", ":"), + ).encode() + return _MAGIC + struct.pack(" tuple[dict, dict[str, tuple[WireTensor, str]]]: + if not payload.startswith(_MAGIC) or len(payload) < len(_MAGIC) + 4: + raise ValueError("invalid prefill snapshot magic") + header_len = struct.unpack(" len(payload): + raise ValueError("truncated prefill snapshot header") + metadata = json.loads(payload[header_start:header_end]) + raw = memoryview(payload)[header_end:] + tensors: dict[str, tuple[WireTensor, str]] = {} + for item in metadata.pop("tensors"): + start = int(item["offset"]) + end = start + int(item["length"]) + if start < 0 or end > len(raw): + raise ValueError("truncated prefill snapshot tensor") + tensors[item["name"]] = ( + from_proto_fields( + item["dtype"], + item["shape"], + bytes(raw[start:end]), + ), + item["framework"], + ) + return metadata, tensors diff --git a/inference_engine/backends/mlx/verifier.py b/inference_engine/backends/mlx/verifier.py index 49dd3125..4268fee0 100644 --- a/inference_engine/backends/mlx/verifier.py +++ b/inference_engine/backends/mlx/verifier.py @@ -107,20 +107,41 @@ def __init__(self, config: Optional[VerifierConfig] = None) -> None: # ``kv_live_bytes`` accessor. Mirrors the CPU verifier; # reads dims from the wrapped HF config so GQA / MQA via # ``num_key_value_heads`` is honored. - cfg = self.model.config if hasattr(self.model, "config") else self.model - num_layers = int(getattr(cfg, "num_hidden_layers")) - num_kv_heads = int( - getattr(cfg, "num_key_value_heads", None) - or getattr(cfg, "num_attention_heads") - ) - head_dim = int( - getattr(cfg, "head_dim", None) - or (cfg.hidden_size // cfg.num_attention_heads) + cfg = ( + getattr(self.model, "config", None) + or getattr(self.model, "args", None) + or self.model ) + cfg = getattr(cfg, "text_config", None) or cfg itemsize = torch.tensor([], dtype=self.config.dtype).element_size() - self._bytes_per_kv_token = ( - num_layers * num_kv_heads * head_dim * itemsize * 2 - ) + try: + num_layers = int(getattr(cfg, "num_hidden_layers")) + num_kv_heads = int( + getattr(cfg, "num_key_value_heads", None) + or getattr(cfg, "num_attention_heads") + ) + head_dim = int( + getattr(cfg, "head_dim", None) + or (cfg.hidden_size // cfg.num_attention_heads) + ) + self._bytes_per_kv_token = ( + num_layers * num_kv_heads * head_dim * itemsize * 2 + ) + except Exception: + from inference_engine.backends.mlx.cross_model_dlm_verifier import ( + per_layer_kv_geometry, + resolve_mlx_text_model, + ) + geometry = per_layer_kv_geometry(resolve_mlx_text_model(self.model)) + if not geometry or any( + num_kv_heads <= 0 or head_dim <= 0 + for num_kv_heads, head_dim, _layer_type in geometry + ): + raise + self._bytes_per_kv_token = sum( + num_kv_heads * head_dim * itemsize * 2 + for num_kv_heads, head_dim, _layer_type in geometry + ) # ---------------------------- public API ---------------------------- # diff --git a/inference_engine/distributed/capability.py b/inference_engine/distributed/capability.py index 1ae05cfc..627396ea 100644 --- a/inference_engine/distributed/capability.py +++ b/inference_engine/distributed/capability.py @@ -48,6 +48,7 @@ class CapabilityRole(enum.IntEnum): PROPOSER = 2 EMBEDDER = 3 TOOL = 4 + PREFILL_CACHE = 5 @dataclass(frozen=True) @@ -77,6 +78,119 @@ def from_proto(cls, msg: distributed_pb2.ModelCapability) -> "ModelCapability": ) +@dataclass(frozen=True) +class NodeEndpoint: + """One interface-specific address advertised by a node.""" + + address: str + network: str = "" + priority: int = 0 + measured_rtt_ms: float = 0.0 + + def to_proto(self) -> distributed_pb2.NodeEndpoint: + return distributed_pb2.NodeEndpoint( + address=self.address, + network=self.network, + priority=self.priority, + measured_rtt_ms=self.measured_rtt_ms, + ) + + @classmethod + def from_proto(cls, msg: distributed_pb2.NodeEndpoint) -> "NodeEndpoint": + return cls( + address=msg.address, + network=msg.network, + priority=msg.priority, + measured_rtt_ms=msg.measured_rtt_ms, + ) + + +@dataclass(frozen=True) +class CacheCompatibility: + """Exact compatibility tuple for reusable prefill K/V blocks.""" + + model_id: str + model_revision: str = "" + tokenizer_revision: str = "" + cache_format_version: str = "kv-v1" + quantization: str = "" + rope_hash: str = "" + layer_geometry_hash: str = "" + kv_dtype: str = "" + block_size_tokens: int = 64 + + def to_proto(self) -> distributed_pb2.CacheCompatibility: + return distributed_pb2.CacheCompatibility( + model_id=self.model_id, + model_revision=self.model_revision, + tokenizer_revision=self.tokenizer_revision, + cache_format_version=self.cache_format_version, + quantization=self.quantization, + rope_hash=self.rope_hash, + layer_geometry_hash=self.layer_geometry_hash, + kv_dtype=self.kv_dtype, + block_size_tokens=self.block_size_tokens, + ) + + @classmethod + def from_proto( + cls, msg: distributed_pb2.CacheCompatibility, + ) -> "CacheCompatibility": + return cls( + model_id=msg.model_id, + model_revision=msg.model_revision, + tokenizer_revision=msg.tokenizer_revision, + cache_format_version=msg.cache_format_version, + quantization=msg.quantization, + rope_hash=msg.rope_hash, + layer_geometry_hash=msg.layer_geometry_hash, + kv_dtype=msg.kv_dtype, + block_size_tokens=msg.block_size_tokens, + ) + + +@dataclass(frozen=True) +class CacheCapability: + """One compatible prefill-cache offering on a node.""" + + compatibility: CacheCompatibility + cache_address: str = "" + cache_bytes_used: int = 0 + cache_bytes_free: int = 0 + entry_count: int = 0 + cache_epoch: int = 0 + load: float = 0.0 + tokens_served: int = 0 + bloom_filter: bytes = b"" + + def to_proto(self) -> distributed_pb2.CacheCapability: + return distributed_pb2.CacheCapability( + compatibility=self.compatibility.to_proto(), + cache_address=self.cache_address, + cache_bytes_used=self.cache_bytes_used, + cache_bytes_free=self.cache_bytes_free, + entry_count=self.entry_count, + cache_epoch=self.cache_epoch, + load=self.load, + tokens_served=self.tokens_served, + bloom_filter=self.bloom_filter, + ) + + @classmethod + def from_proto(cls, msg: distributed_pb2.CacheCapability) -> "CacheCapability": + return cls( + compatibility=CacheCompatibility.from_proto(msg.compatibility), + cache_address=msg.cache_address, + cache_bytes_used=msg.cache_bytes_used, + cache_bytes_free=msg.cache_bytes_free, + entry_count=msg.entry_count, + cache_epoch=msg.cache_epoch, + load=msg.load, + tokens_served=msg.tokens_served, + bloom_filter=msg.bloom_filter, + ) + + @dataclass(frozen=True) class NodeCapability: """One node's capability card. See distributed.proto for field docs.""" @@ -90,6 +204,8 @@ class NodeCapability: announced_at_unix: float = 0.0 ttl_seconds: float = DEFAULT_TTL_SECONDS ring_address: str = "" + caches: Tuple[CacheCapability, ...] = () + endpoints: Tuple[NodeEndpoint, ...] = () def __post_init__(self) -> None: if not self.node_id: @@ -121,6 +237,8 @@ def to_proto(self) -> distributed_pb2.NodeCapability: announced_at_unix=self.announced_at_unix, ttl_seconds=self.ttl_seconds, ring_address=self.ring_address, + caches=[c.to_proto() for c in self.caches], + endpoints=[e.to_proto() for e in self.endpoints], ) @classmethod @@ -135,6 +253,8 @@ def from_proto(cls, msg: distributed_pb2.NodeCapability) -> "NodeCapability": announced_at_unix=msg.announced_at_unix, ttl_seconds=msg.ttl_seconds, 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), ) diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py new file mode 100644 index 00000000..d931d924 --- /dev/null +++ b/inference_engine/distributed/prefill_cache.py @@ -0,0 +1,276 @@ +"""Immutable, content-addressed distributed prefill K/V cache. + +The cache stores an opaque restorable snapshot at selected token-block +boundaries. Model-specific adapters own serialization/import; this module owns +deterministic prefix hashing, exact compatibility matching, +longest-contiguous-prefix lookup, leases, accounting, and memory-pressure +eviction. A hit transfers only the snapshot at the longest matched boundary. + +Decode never reads this store. A requester imports a hit once, computes the +missing suffix locally, and keeps the autoregressive loop local. +""" + +from __future__ import annotations + +import hashlib +import json +import secrets +import threading +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Iterable, Sequence + +from inference_engine.distributed.capability import CacheCompatibility + +DEFAULT_LEASE_SECONDS = 30.0 + + +def compatibility_fingerprint(compatibility: CacheCompatibility) -> bytes: + """Stable SHA-256 of every field that affects K/V interpretation.""" + payload = { + "block_size_tokens": compatibility.block_size_tokens, + "cache_format_version": compatibility.cache_format_version, + "kv_dtype": compatibility.kv_dtype, + "layer_geometry_hash": compatibility.layer_geometry_hash, + "model_id": compatibility.model_id, + "model_revision": compatibility.model_revision, + "quantization": compatibility.quantization, + "rope_hash": compatibility.rope_hash, + "tokenizer_revision": compatibility.tokenizer_revision, + } + return hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode(), + ).digest() + + +def chained_block_hashes( + token_ids: Sequence[int], + compatibility: CacheCompatibility, +) -> list[bytes]: + """Hash fixed-size token blocks, chaining each hash to its predecessor. + + Chaining means block N is only reusable after blocks 0..N-1 matched, which + enforces the causal longest-prefix rule and prevents arbitrary-hole reuse. + """ + size = int(compatibility.block_size_tokens) + if size <= 0: + raise ValueError("block_size_tokens must be > 0") + namespace = compatibility_fingerprint(compatibility) + previous = bytes(32) + hashes: list[bytes] = [] + 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() + hashes.append(previous) + return hashes + + +@dataclass(frozen=True) +class CacheBlock: + block_hash: bytes + token_count: int + payload: bytes + payload_sha256: bytes + + @classmethod + def create(cls, block_hash: bytes, token_count: int, payload: bytes) -> "CacheBlock": + if len(block_hash) != 32: + raise ValueError("block_hash must be SHA-256 (32 bytes)") + if token_count <= 0: + raise ValueError("token_count must be > 0") + data = bytes(payload) + return cls( + block_hash=bytes(block_hash), + token_count=int(token_count), + payload=data, + payload_sha256=hashlib.sha256(data).digest(), + ) + + @property + def nbytes(self) -> int: + return len(self.payload) + + +@dataclass(frozen=True) +class PrefixLease: + lease_id: str + block_hashes: tuple[bytes, ...] + hit_block_count: int + hit_token_count: int + transfer_bytes: int + cache_epoch: int + expires_at_unix: float + payload_sha256: bytes + + +@dataclass(frozen=True) +class CacheStats: + bytes_used: int + max_bytes: int + entry_count: int + cache_epoch: int + lookup_hits: int + lookup_misses: int + tokens_served: int + bytes_served: int + + +class PrefixCacheStore: + """Thread-safe in-memory LRU of immutable K/V block payloads.""" + + def __init__( + self, + compatibility: CacheCompatibility, + *, + max_bytes: int, + node_id: str, + ) -> None: + if max_bytes <= 0: + raise ValueError("max_bytes must be > 0") + if not node_id: + raise ValueError("node_id must be non-empty") + self.compatibility = compatibility + self.max_bytes = int(max_bytes) + self.node_id = node_id + self._blocks: OrderedDict[bytes, CacheBlock] = OrderedDict() + self._leases: dict[str, PrefixLease] = {} + self._bytes_used = 0 + self._epoch = 1 + self._lookup_hits = 0 + self._lookup_misses = 0 + self._tokens_served = 0 + self._bytes_served = 0 + self._lock = threading.RLock() + + def put(self, block: CacheBlock) -> bool: + """Publish one immutable block. Returns False for an identical hit.""" + if block.nbytes > self.max_bytes: + raise ValueError("block payload exceeds cache capacity") + with self._lock: + existing = self._blocks.get(block.block_hash) + if existing is not None: + if existing.payload_sha256 != block.payload_sha256: + raise ValueError("content-address collision with different payload") + self._blocks.move_to_end(block.block_hash) + return False + self._blocks[block.block_hash] = block + self._bytes_used += block.nbytes + self._epoch += 1 + self._evict_to_budget() + return True + + def put_prefix( + self, + token_ids: Sequence[int], + payloads: Sequence[bytes], + ) -> list[bytes]: + hashes = chained_block_hashes(token_ids, self.compatibility) + if len(hashes) != len(payloads): + raise ValueError("one payload is required for every token block") + size = self.compatibility.block_size_tokens + for index, (block_hash, payload) in enumerate(zip(hashes, payloads)): + prefix_count = min((index + 1) * size, len(token_ids)) + self.put(CacheBlock.create(block_hash, prefix_count, payload)) + return hashes + + def lookup( + self, + block_hashes: Sequence[bytes], + *, + lease_seconds: float = DEFAULT_LEASE_SECONDS, + now: float | None = None, + ) -> PrefixLease: + """Lease the longest contiguous prefix held by this store.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be > 0") + now = time.time() if now is None else now + with self._lock: + self._expire_leases(now) + matched: list[CacheBlock] = [] + for raw_hash in block_hashes: + block_hash = bytes(raw_hash) + block = self._blocks.get(block_hash) + if block is None: + break + matched.append(block) + self._blocks.move_to_end(block_hash) + if not matched: + self._lookup_misses += 1 + return PrefixLease("", (), 0, 0, 0, self._epoch, now, bytes(32)) + self._lookup_hits += 1 + lease_id = secrets.token_urlsafe(18) + snapshot = matched[-1] + lease = PrefixLease( + lease_id=lease_id, + block_hashes=(snapshot.block_hash,), + hit_block_count=len(matched), + hit_token_count=snapshot.token_count, + transfer_bytes=snapshot.nbytes, + cache_epoch=self._epoch, + expires_at_unix=now + lease_seconds, + payload_sha256=snapshot.payload_sha256, + ) + self._leases[lease_id] = lease + return lease + + def fetch(self, lease_id: str, *, now: float | None = None) -> tuple[CacheBlock, ...]: + now = time.time() if now is None else now + with self._lock: + self._expire_leases(now) + lease = self._leases.get(lease_id) + if lease is None: + raise KeyError("unknown or expired cache lease") + blocks: list[CacheBlock] = [] + for block_hash in lease.block_hashes: + block = self._blocks.get(block_hash) + if block is None: + raise KeyError("leased block was evicted") + blocks.append(block) + self._tokens_served += lease.hit_token_count + self._bytes_served += lease.transfer_bytes + return tuple(blocks) + + def stats(self) -> CacheStats: + with self._lock: + return CacheStats( + bytes_used=self._bytes_used, + max_bytes=self.max_bytes, + entry_count=len(self._blocks), + cache_epoch=self._epoch, + lookup_hits=self._lookup_hits, + lookup_misses=self._lookup_misses, + tokens_served=self._tokens_served, + bytes_served=self._bytes_served, + ) + + def block_hashes(self) -> tuple[bytes, ...]: + with self._lock: + return tuple(self._blocks) + + def _expire_leases(self, now: float) -> None: + for lease_id, lease in list(self._leases.items()): + if now > lease.expires_at_unix: + del self._leases[lease_id] + + def _pinned_hashes(self) -> set[bytes]: + return { + block_hash + for lease in self._leases.values() + for block_hash in lease.block_hashes + } + + def _evict_to_budget(self) -> None: + pinned = self._pinned_hashes() + while self._bytes_used > self.max_bytes and self._blocks: + victim = next((h for h in self._blocks if h not in pinned), None) + if victim is None: + break + block = self._blocks.pop(victim) + self._bytes_used -= block.nbytes + self._epoch += 1 + + +def total_payload_bytes(blocks: Iterable[CacheBlock]) -> int: + return sum(block.nbytes for block in blocks) diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py new file mode 100644 index 00000000..0b2c5484 --- /dev/null +++ b/inference_engine/distributed/prefill_cache_runtime.py @@ -0,0 +1,259 @@ +"""Synchronous runtime hook that applies distributed prefill-cache hits. + +The gRPC RuntimeService uses a synchronous verifier underneath its asyncio +handlers. This hook keeps that contract: peer lookups run concurrently in a +small thread pool, a winning snapshot is imported once, and missing suffix +blocks are prefetched locally before decode begins. +""" + +from __future__ import annotations + +import hashlib +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Any, Sequence + +import grpc + +from inference_engine.backends.mlx.prefill_snapshot import ( + export_mlx_prefill_snapshot, + import_mlx_prefill_snapshot, +) +from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.prefill_cache import ( + CacheBlock, + PrefixCacheStore, + chained_block_hashes, +) +from inference_engine.server.proto_gen.kakeya.v1 import ( + distributed_pb2, + distributed_pb2_grpc, +) + + +@dataclass +class PrefillReuseStats: + local_hits: int = 0 + remote_hits: int = 0 + misses: int = 0 + tokens_reused: int = 0 + tokens_computed: int = 0 + bytes_received: int = 0 + + +@dataclass(frozen=True) +class _Hit: + source: str + lease_id: str + hit_blocks: int + hit_tokens: int + transfer_bytes: int + payload: bytes | None = None + + +class DistributedPrefillCacheHook: + """Prepare a cold verifier using local/remote longest-prefix snapshots.""" + + def __init__( + self, + local_store: PrefixCacheStore, + *, + peers: Sequence[str] = (), + lookup_timeout_s: float = 2.0, + fetch_timeout_s: float = 30.0, + on_reuse=None, + ) -> None: + self.local_store = local_store + self.compatibility = local_store.compatibility + self.peers = tuple(dict.fromkeys(peer for peer in peers if peer)) + self.lookup_timeout_s = float(lookup_timeout_s) + self.fetch_timeout_s = float(fetch_timeout_s) + self.stats = PrefillReuseStats() + self._on_reuse = on_reuse + self._publisher = ThreadPoolExecutor( + max_workers=max(1, min(4, len(self.peers))), + thread_name_prefix="prefill-kv-publish", + ) + + def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int: + """Restore the longest prefix, compute suffix, publish all new boundaries. + + Returns the number of tokens reused from cache. + """ + tokens = [int(token) for token in token_ids] + if not tokens: + return 0 + hashes = chained_block_hashes(tokens, self.compatibility) + 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) + verifier.reset() + imported = import_mlx_prefill_snapshot( + payload, + verifier.cache, + compatibility=self.compatibility, + ) + reused = min(imported.token_count, len(tokens)) + verifier.cached_token_sequence = list(imported.cached_token_ids) + verifier.next_global_position = reused + if imported.next_token_logits is not None: + verifier.next_token_logits = imported.next_token_logits + self.stats.tokens_reused += reused + if self._on_reuse is not None: + self._on_reuse(reused) + if hit.source == "local": + self.stats.local_hits += 1 + else: + self.stats.remote_hits += 1 + else: + self.stats.misses += 1 + + self._compute_and_publish(verifier, tokens, hashes, reused) + return reused + + def _compute_and_publish( + self, + verifier: Any, + tokens: list[int], + hashes: list[bytes], + 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 + logits = verifier.forward_block(block_tokens) + verifier.commit_or_truncate( + forwarded=len(block_tokens), + accepted=len(block_tokens), + ) + verifier.next_token_logits = logits[-1].clone() + self.stats.tokens_computed += len(block_tokens) + self._publish_boundary(verifier, tokens, hashes, block_index, end) + + def _publish_boundary( + self, + verifier: Any, + tokens: list[int], + hashes: list[bytes], + block_index: int, + prefix_end: int, + ) -> None: + 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 = CacheBlock.create(hashes[block_index], prefix_end, payload) + self.local_store.put(block) + if self.peers: + from inference_engine.distributed.prefill_cache_service import ( + publish_block_sync, + ) + for peer in self.peers: + self._publisher.submit( + publish_block_sync, + peer, + self.compatibility, + block, + timeout_s=self.fetch_timeout_s, + ) + + def close(self) -> None: + self._publisher.shutdown(wait=False, cancel_futures=True) + + def _best_hit(self, hashes: Sequence[bytes]) -> _Hit | None: + candidates: list[_Hit] = [] + local = self.local_store.lookup(hashes) + if local.lease_id: + blocks = self.local_store.fetch(local.lease_id) + candidates.append(_Hit( + source="local", + lease_id=local.lease_id, + hit_blocks=local.hit_block_count, + hit_tokens=local.hit_token_count, + transfer_bytes=local.transfer_bytes, + payload=blocks[-1].payload, + )) + if self.peers: + with ThreadPoolExecutor(max_workers=min(8, len(self.peers))) as pool: + futures = { + pool.submit(self._lookup_peer, peer, hashes): peer + for peer in self.peers + } + for future in as_completed(futures): + hit = future.result() + if hit is not None: + candidates.append(hit) + if not candidates: + return None + return max( + candidates, + key=lambda hit: ( + hit.hit_tokens, + hit.source == "local", + -hit.transfer_bytes, + ), + ) + + 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) + response = stub.LookupPrefix( + distributed_pb2.LookupPrefixRequest( + compatibility=self.compatibility.to_proto(), + block_hashes=hashes, + ), + timeout=self.lookup_timeout_s, + ) + except grpc.RpcError: + return None + if not response.lease_id or response.hit_block_count == 0: + 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, + ) + + def _fetch_remote(self, hit: _Hit) -> bytes: + parts: dict[int, bytes] = {} + expected_chunks = 0 + expected_sha = b"" + try: + with grpc.insecure_channel(hit.source) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + for chunk in stub.FetchBlocks( + distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id), + timeout=self.fetch_timeout_s, + ): + parts[chunk.chunk_index] = bytes(chunk.data) + expected_chunks = chunk.total_chunks + expected_sha = bytes(chunk.block_sha256) + 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") + 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") + self.stats.bytes_received += len(payload) + return payload diff --git a/inference_engine/distributed/prefill_cache_service.py b/inference_engine/distributed/prefill_cache_service.py new file mode 100644 index 00000000..ec70056c --- /dev/null +++ b/inference_engine/distributed/prefill_cache_service.py @@ -0,0 +1,376 @@ +"""gRPC service/client for distributed immutable prefill K/V blocks.""" + +from __future__ import annotations + +import asyncio +import hashlib +from dataclasses import dataclass +from typing import Sequence + +import grpc + +from inference_engine.distributed.capability import ( + CacheCapability, + CacheCompatibility, + NodeCapability, +) +from inference_engine.distributed.prefill_cache import ( + CacheBlock, + PrefixCacheStore, + PrefixLease, +) +from inference_engine.server.proto_gen.kakeya.v1 import ( + distributed_pb2, + distributed_pb2_grpc, +) + +DEFAULT_CHUNK_BYTES = 2 * 1024 * 1024 + + +def cache_capability( + store: PrefixCacheStore, + *, + cache_address: str, + load: float = 0.0, +) -> CacheCapability: + stats = store.stats() + return CacheCapability( + compatibility=store.compatibility, + cache_address=cache_address, + cache_bytes_used=stats.bytes_used, + cache_bytes_free=max(0, stats.max_bytes - stats.bytes_used), + entry_count=stats.entry_count, + cache_epoch=stats.cache_epoch, + load=load, + tokens_served=stats.tokens_served, + ) + + +class PrefillCacheServiceServicer( + distributed_pb2_grpc.PrefillCacheServiceServicer, +): + def __init__( + self, + store: PrefixCacheStore, + *, + cache_address: str, + chunk_bytes: int = DEFAULT_CHUNK_BYTES, + ) -> 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) + + async def GetCacheSummary( # noqa: N802 + self, + request: distributed_pb2.GetCacheSummaryRequest, + context: grpc.aio.ServicerContext, + ) -> distributed_pb2.GetCacheSummaryResponse: + requested = CacheCompatibility.from_proto(request.compatibility) + caches = [] + if requested == self.store.compatibility: + caches.append( + cache_capability( + self.store, + cache_address=self.cache_address, + ).to_proto(), + ) + return distributed_pb2.GetCacheSummaryResponse( + node_id=self.store.node_id, + caches=caches, + ) + + async def LookupPrefix( # noqa: N802 + self, + request: distributed_pb2.LookupPrefixRequest, + context: grpc.aio.ServicerContext, + ) -> distributed_pb2.LookupPrefixResponse: + requested = CacheCompatibility.from_proto(request.compatibility) + if requested != self.store.compatibility: + return distributed_pb2.LookupPrefixResponse( + node_id=self.store.node_id, + cache_epoch=self.store.stats().cache_epoch, + ) + lease = self.store.lookup(request.block_hashes) + return _lease_to_proto(self.store.node_id, lease) + + async def FetchBlocks( # noqa: N802 + self, + request: distributed_pb2.FetchBlocksRequest, + context: grpc.aio.ServicerContext, + ): + try: + blocks = self.store.fetch(request.lease_id) + except KeyError as exc: + await context.abort(grpc.StatusCode.NOT_FOUND, str(exc)) + return # pragma: no cover - grpc abort raises + epoch = self.store.stats().cache_epoch + for block_index, block in enumerate(blocks): + total_chunks = max( + 1, (len(block.payload) + self.chunk_bytes - 1) // self.chunk_bytes, + ) + for chunk_index in range(total_chunks): + start = chunk_index * self.chunk_bytes + yield distributed_pb2.FetchBlocksResponse( + block_hash=block.block_hash, + block_index=block_index, + token_count=block.token_count, + chunk_index=chunk_index, + total_chunks=total_chunks, + data=block.payload[start:start + self.chunk_bytes], + block_sha256=block.payload_sha256, + cache_epoch=epoch, + ) + + async def PublishBlock( # noqa: N802 + self, + request_iterator, + context: grpc.aio.ServicerContext, + ) -> distributed_pb2.PublishBlockResponse: + parts: dict[int, bytes] = {} + first = None + async for chunk in request_iterator: + if first is None: + first = chunk + elif ( + chunk.block_hash != first.block_hash + or chunk.total_chunks != first.total_chunks + or chunk.block_sha256 != first.block_sha256 + ): + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "inconsistent publish chunk metadata", + ) + parts[chunk.chunk_index] = bytes(chunk.data) + if first is None: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "empty publish stream", + ) + requested = CacheCompatibility.from_proto(first.compatibility) + if requested != self.store.compatibility: + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "prefill cache compatibility mismatch", + ) + if len(parts) != first.total_chunks: + await context.abort( + grpc.StatusCode.DATA_LOSS, + "incomplete publish stream", + ) + payload = b"".join(parts[index] for index in range(first.total_chunks)) + if hashlib.sha256(payload).digest() != bytes(first.block_sha256): + await context.abort( + grpc.StatusCode.DATA_LOSS, + "prefill cache payload checksum mismatch", + ) + try: + stored = self.store.put(CacheBlock.create( + bytes(first.block_hash), + first.token_count, + payload, + )) + except ValueError as exc: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(exc)) + return distributed_pb2.PublishBlockResponse( + stored=stored, + cache_epoch=self.store.stats().cache_epoch, + ) + + +def add_prefill_cache_service( + server: grpc.aio.Server, + store: PrefixCacheStore, + *, + cache_address: str, + chunk_bytes: int = DEFAULT_CHUNK_BYTES, +) -> PrefillCacheServiceServicer: + servicer = PrefillCacheServiceServicer( + store, + cache_address=cache_address, + chunk_bytes=chunk_bytes, + ) + distributed_pb2_grpc.add_PrefillCacheServiceServicer_to_server( + servicer, server, + ) + return servicer + + +@dataclass(frozen=True) +class RemotePrefixHit: + address: str + node_id: str + lease_id: str + hit_block_count: int + hit_token_count: int + transfer_bytes: int + cache_epoch: int + expires_at_unix: float + payload_sha256: bytes + + @property + def found(self) -> bool: + return self.hit_block_count > 0 and bool(self.lease_id) + + +async def lookup_peer( + address: str, + compatibility: CacheCompatibility, + block_hashes: Sequence[bytes], + *, + timeout_s: float = 3.0, +) -> RemotePrefixHit: + try: + async with grpc.aio.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + response = await stub.LookupPrefix( + distributed_pb2.LookupPrefixRequest( + compatibility=compatibility.to_proto(), + block_hashes=block_hashes, + ), + timeout=timeout_s, + ) + except grpc.aio.AioRpcError: + return RemotePrefixHit(address, "", "", 0, 0, 0, 0, 0.0, b"") + return RemotePrefixHit( + address=address, + node_id=response.node_id, + lease_id=response.lease_id, + hit_block_count=response.hit_block_count, + hit_token_count=response.hit_token_count, + transfer_bytes=response.transfer_bytes, + cache_epoch=response.cache_epoch, + expires_at_unix=response.lease_expires_at_unix, + payload_sha256=response.payload_sha256, + ) + + +async def lookup_best_peer( + peers: Sequence[str], + compatibility: CacheCompatibility, + block_hashes: Sequence[bytes], + *, + timeout_s: float = 3.0, +) -> RemotePrefixHit | None: + """Fan out concurrently and choose longest hit, then smallest transfer.""" + if not peers: + return None + hits = await asyncio.gather(*( + lookup_peer( + peer, + compatibility, + block_hashes, + timeout_s=timeout_s, + ) + for peer in peers + )) + candidates = [hit for hit in hits if hit.found] + if not candidates: + return None + return max( + candidates, + key=lambda hit: ( + hit.hit_block_count, + -hit.transfer_bytes, + hit.expires_at_unix, + ), + ) + + +async def fetch_remote_blocks( + hit: RemotePrefixHit, + *, + timeout_s: float = 30.0, +) -> list[distributed_pb2.FetchBlocksResponse]: + async with grpc.aio.insecure_channel(hit.address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + stream = stub.FetchBlocks( + distributed_pb2.FetchBlocksRequest(lease_id=hit.lease_id), + timeout=timeout_s, + ) + return [chunk async for chunk in stream] + + +def publish_block_sync( + address: str, + compatibility: CacheCompatibility, + block: CacheBlock, + *, + timeout_s: float = 30.0, + chunk_bytes: int = DEFAULT_CHUNK_BYTES, +) -> bool: + """Publish one immutable snapshot to a peer (used by background workers).""" + try: + total_chunks = max( + 1, (block.nbytes + chunk_bytes - 1) // chunk_bytes, + ) + + 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(), + ) + + with grpc.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + response = stub.PublishBlock( + chunks(), + timeout=timeout_s, + ) + except grpc.RpcError: + return False + return bool(response.stored) + + +def compatible_cache_peers( + cards: Sequence[NodeCapability], + compatibility: CacheCompatibility, +) -> list[str]: + """Choose each compatible card's highest-priority reachable endpoint.""" + peers: list[str] = [] + for card in cards: + matching = [ + cache + for cache in card.caches + if cache.compatibility == compatibility + ] + if not matching: + continue + cache_address = matching[0].cache_address + if cache_address: + peers.append(cache_address) + continue + endpoints = sorted( + card.endpoints, + key=lambda endpoint: ( + endpoint.priority, + -endpoint.measured_rtt_ms, + ), + reverse=True, + ) + peers.append(endpoints[0].address if endpoints else card.grpc_address) + return peers + + +def _lease_to_proto( + node_id: str, + lease: PrefixLease, +) -> distributed_pb2.LookupPrefixResponse: + return distributed_pb2.LookupPrefixResponse( + node_id=node_id, + hit_block_count=lease.hit_block_count, + hit_token_count=lease.hit_token_count, + transfer_bytes=lease.transfer_bytes, + cache_epoch=lease.cache_epoch, + lease_id=lease.lease_id, + lease_expires_at_unix=lease.expires_at_unix, + payload_sha256=lease.payload_sha256, + ) diff --git a/inference_engine/network/__init__.py b/inference_engine/network/__init__.py new file mode 100644 index 00000000..ae05e4a1 --- /dev/null +++ b/inference_engine/network/__init__.py @@ -0,0 +1 @@ +"""Kakeya inference-network management API and dashboard.""" diff --git a/inference_engine/network/api.py b/inference_engine/network/api.py new file mode 100644 index 00000000..79d8eb26 --- /dev/null +++ b/inference_engine/network/api.py @@ -0,0 +1,131 @@ +"""FastAPI management/telemetry surface for the Kakeya inference network.""" + +from __future__ import annotations + +import asyncio +import json +import os +from typing import Optional + +from fastapi import Depends, FastAPI, Header, HTTPException +from fastapi.responses import HTMLResponse, StreamingResponse +from pydantic import BaseModel, Field + +from inference_engine.network.dashboard import dashboard_html +from inference_engine.network.state import NetworkState + + +class RegisterNodeRequest(BaseModel): + alias: str = Field(min_length=1, max_length=100) + address: str = Field(min_length=3, max_length=255) + region: str = Field(default="Private", max_length=100) + role: str = Field(default="hybrid", pattern="^(head|cache|hybrid|inference)$") + + +class CreateGroupRequest(BaseModel): + name: str = Field(min_length=1, max_length=100) + node_ids: list[str] = Field(min_length=1) + + +class TokenTelemetryRequest(BaseModel): + node_id: str = Field(min_length=1, max_length=100) + completed: int = Field(ge=0) + kv_assisted: int = Field(default=0, ge=0) + + +def create_network_app( + state: NetworkState, + *, + api_key: Optional[str] = None, +) -> FastAPI: + key = (api_key if api_key is not None else os.environ.get( + "KAKEYA_NETWORK_API_KEY", "", + )).strip() + app = FastAPI(title="Kakeya Inference Network", version="0.1.0") + + def require_key( + x_api_key: Optional[str] = Header(default=None), + ) -> None: + if key and x_api_key != key: + raise HTTPException(status_code=401, detail="invalid X-API-Key") + + @app.get("/", response_class=HTMLResponse) + @app.get("/network", response_class=HTMLResponse) + def dashboard() -> str: + return dashboard_html() + + @app.get("/healthz") + def healthz(): + summary = state.summary() + return { + "status": "ok", + "online_nodes": summary["online_nodes"], + "cache_bytes_used": summary["cache_bytes_used"], + } + + @app.get("/v1/network/summary") + def summary(): + return state.summary() + + @app.get("/v1/network/nodes") + def nodes(): + return state.nodes() + + @app.post("/v1/network/nodes/register", dependencies=[Depends(require_key)]) + def register(request: RegisterNodeRequest): + try: + return state.register_node(**request.model_dump()) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.get("/v1/network/groups") + def groups(): + return state.groups() + + @app.post("/v1/network/groups", dependencies=[Depends(require_key)]) + def create_group(request: CreateGroupRequest): + try: + return state.create_group(**request.model_dump()) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.get("/v1/network/topology") + def topology(): + return state.topology() + + @app.get("/v1/network/tokens") + def tokens(): + summary = state.summary() + return { + "completed": summary["completed_tokens"], + "kv_assisted": summary["kv_assisted_tokens"], + "hit_rate": summary["kv_hit_rate"], + } + + @app.post( + "/v1/network/telemetry/tokens", + dependencies=[Depends(require_key)], + ) + def record_tokens(request: TokenTelemetryRequest): + try: + state.record_tokens(**request.model_dump()) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + return {"status": "accepted"} + + @app.get("/v1/network/events") + async def events(once: bool = False): + async def stream(): + while True: + payload = json.dumps({ + "type": "summary", + "data": state.summary(), + }, separators=(",", ":")) + yield f"event: summary\ndata: {payload}\n\n" + if once: + return + await asyncio.sleep(5) # pragma: no cover - persistent SSE loop + + return StreamingResponse(stream(), media_type="text/event-stream") + + return app diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py new file mode 100644 index 00000000..2802e920 --- /dev/null +++ b/inference_engine/network/dashboard.py @@ -0,0 +1,58 @@ +"""Self-contained kakeya.ai inference-network dashboard.""" + +from __future__ import annotations + + +def dashboard_html() -> str: + return r""" + + + +Kakeya Inference Network + + +
+

Kakeya Inference Network

P2P Prefill KV sharing across trusted inference nodes
+
+ +
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
+
+

Online node distribution

Live KV discovery

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

Cache capacity

0 / 0 GB
+
+ + +

Exact IPs, raw prompt hashes and cache keys are administrator-only. Region is operator-selected and coarse.

+
+""" diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py new file mode 100644 index 00000000..82a1be2f --- /dev/null +++ b/inference_engine/network/state.py @@ -0,0 +1,237 @@ +"""Persistent product state projected from the P2P capability plane.""" + +from __future__ import annotations + +import json +import secrets +import threading +import time +from pathlib import Path +from typing import Any + +from inference_engine.distributed.capability import CapabilityRegistry +from inference_engine.distributed.prefill_cache import PrefixCacheStore + + +class NetworkState: + def __init__( + self, + registry: CapabilityRegistry, + cache_store: PrefixCacheStore, + *, + state_path: str | Path, + ) -> None: + self.registry = registry + self.cache_store = cache_store + self.state_path = Path(state_path).expanduser() + self._lock = threading.RLock() + self._data = self._load() + + def register_node( + self, + *, + alias: str, + address: str, + region: str, + role: str = "hybrid", + ) -> dict[str, Any]: + if not alias or not address: + raise ValueError("alias and address are required") + now = time.time() + item = { + "id": secrets.token_hex(8), + "alias": alias, + "address": address, + "region": region or "Private", + "role": role, + "status": "pending", + "pairing_token": "kn_pair_" + secrets.token_urlsafe(18), + "expires_at": now + 600, + "created_at": now, + } + with self._lock: + self._data["registrations"].append(item) + self._save() + return dict(item) + + def create_group(self, *, name: str, node_ids: list[str]) -> dict[str, Any]: + if not name or not node_ids: + raise ValueError("name and node_ids are required") + group = { + "id": secrets.token_hex(6), + "name": name, + "node_ids": list(dict.fromkeys(node_ids)), + "created_at": time.time(), + } + with self._lock: + self._data["groups"].append(group) + self._save() + return dict(group) + + def record_tokens( + self, + *, + node_id: str, + completed: int, + kv_assisted: int = 0, + ) -> None: + if completed < 0 or kv_assisted < 0 or kv_assisted > completed: + raise ValueError("invalid token counters") + with self._lock: + counters = self._data["tokens"].setdefault( + node_id, + {"completed": 0, "kv_assisted": 0}, + ) + counters["completed"] += int(completed) + counters["kv_assisted"] += int(kv_assisted) + self._save() + + def nodes(self) -> list[dict[str, Any]]: + registrations = { + item["alias"]: item + for item in self._data["registrations"] + } + output: 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 + endpoint = sorted( + card.endpoints, + key=lambda item: item.priority, + reverse=True, + ) + output.append({ + "id": card.node_id, + "alias": card.node_id, + "region": registration.get("region", "Private"), + "role": registration.get( + "role", + "hybrid" if cache else "inference", + ), + "status": "online", + "platform": card.platform, + "memory_bytes": card.unified_memory_bytes, + "models": [ + { + "model_id": model.model_id, + "role": model.role.name.lower(), + "quantization": model.quantization, + "tokens_per_second": model.tokens_per_second, + } + for model in card.models + ], + "cache": ( + { + "bytes_used": cache.cache_bytes_used, + "bytes_free": cache.cache_bytes_free, + "entry_count": cache.entry_count, + "epoch": cache.cache_epoch, + "tokens_served": cache.tokens_served, + "format": cache.compatibility.cache_format_version, + "model_id": cache.compatibility.model_id, + } + if cache else None + ), + "endpoint": ( + { + "address": endpoint[0].address, + "network": endpoint[0].network, + "priority": endpoint[0].priority, + "rtt_ms": endpoint[0].measured_rtt_ms, + } + if endpoint else { + "address": card.grpc_address, + "network": "default", + "priority": 0, + "rtt_ms": 0, + } + ), + }) + live_ids = {item["id"] for item in output} + for registration in self._data["registrations"]: + if registration["alias"] not in live_ids: + output.append({ + "id": registration["alias"], + "alias": registration["alias"], + "region": registration["region"], + "role": registration["role"], + "status": registration["status"], + "platform": "", + "memory_bytes": 0, + "models": [], + "cache": None, + "endpoint": { + "address": registration["address"], + "network": "pending", + "priority": 0, + "rtt_ms": 0, + }, + }) + return output + + def groups(self) -> list[dict[str, Any]]: + nodes = {node["id"]: node for node in self.nodes()} + groups = [] + for group in self._data["groups"]: + members = [nodes[node_id] for node_id in group["node_ids"] if node_id in nodes] + groups.append({ + **group, + "members": members, + "online": sum(member["status"] == "online" for member in members), + }) + return groups + + def summary(self) -> dict[str, Any]: + nodes = self.nodes() + counters = list(self._data["tokens"].values()) + cache_stats = self.cache_store.stats() + completed = sum(item["completed"] for item in counters) + assisted = sum(item["kv_assisted"] for item in counters) + return { + "online_nodes": sum(node["status"] == "online" for node in nodes), + "registered_nodes": len(nodes), + "groups": len(self._data["groups"]), + "completed_tokens": completed, + "kv_assisted_tokens": assisted, + "kv_hit_rate": (assisted / completed if completed else 0.0), + "cache_bytes_used": sum( + (node["cache"] or {}).get("bytes_used", 0) for node in nodes + ), + "cache_bytes_free": sum( + (node["cache"] or {}).get("bytes_free", 0) for node in nodes + ), + "local_lookup_hits": cache_stats.lookup_hits, + "local_lookup_misses": cache_stats.lookup_misses, + "local_tokens_served": cache_stats.tokens_served, + } + + def topology(self) -> dict[str, Any]: + nodes = self.nodes() + edges = [] + for group in self.groups(): + ids = group["node_ids"] + if len(ids) > 1: + edges.extend({ + "source": ids[0], + "target": target, + "group_id": group["id"], + } for target in ids[1:]) + return {"nodes": nodes, "edges": edges} + + def _load(self) -> dict[str, Any]: + if self.state_path.exists(): + try: + data = json.loads(self.state_path.read_text()) + data.setdefault("registrations", []) + data.setdefault("groups", []) + data.setdefault("tokens", {}) + return data + except (OSError, ValueError): + pass + return {"registrations": [], "groups": [], "tokens": {}} + + def _save(self) -> None: + self.state_path.parent.mkdir(parents=True, exist_ok=True) + temp = self.state_path.with_suffix(".tmp") + temp.write_text(json.dumps(self._data, indent=2, sort_keys=True)) + temp.replace(self.state_path) diff --git a/inference_engine/server/__init__.py b/inference_engine/server/__init__.py index dfc95f20..25cc4c1e 100644 --- a/inference_engine/server/__init__.py +++ b/inference_engine/server/__init__.py @@ -26,10 +26,6 @@ decoder from the user's chosen verifier/proposer pair. """ -from .config import ServerConfig -from .engine import Engine, EngineResult, SpeculativeEngine -from .tokenizer import Tokenizer - __all__ = [ "ServerConfig", "Engine", @@ -37,3 +33,21 @@ "SpeculativeEngine", "Tokenizer", ] + + +def __getattr__(name): + """Lazy public exports keep proto-only/cache-only nodes lightweight.""" + if name == "ServerConfig": + from .config import ServerConfig + return ServerConfig + if name in {"Engine", "EngineResult", "SpeculativeEngine"}: + from .engine import Engine, EngineResult, SpeculativeEngine + return { + "Engine": Engine, + "EngineResult": EngineResult, + "SpeculativeEngine": SpeculativeEngine, + }[name] + if name == "Tokenizer": + from .tokenizer import Tokenizer + return Tokenizer + raise AttributeError(name) diff --git a/inference_engine/server/grpc_app.py b/inference_engine/server/grpc_app.py index e1541de2..02a093ad 100644 --- a/inference_engine/server/grpc_app.py +++ b/inference_engine/server/grpc_app.py @@ -380,6 +380,8 @@ def create_grpc_server( capability_registry: Optional[object] = None, proposers: Optional[object] = None, default_proposer_model_id: str = "", + prefill_cache_store: Optional[object] = None, + prefill_cache_address: str = "", ) -> grpc.aio.Server: """Build, but do not start, a configured gRPC asyncio server. @@ -444,6 +446,20 @@ def create_grpc_server( _logger.info( "gRPC ProposerService enabled for models: %s", sorted(proposers), ) + if prefill_cache_store is not None: + from inference_engine.distributed.prefill_cache_service import ( + add_prefill_cache_service, + ) + + add_prefill_cache_service( + server, + prefill_cache_store, + cache_address=prefill_cache_address or config.bind_address, + ) + _logger.info( + "gRPC PrefillCacheService enabled at %s", + prefill_cache_address or config.bind_address, + ) server.add_insecure_port(config.bind_address) _logger.info("gRPC RuntimeService bound to %s", config.bind_address) return server 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 d6225aff..699e21ec 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.py @@ -24,59 +24,83 @@ -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\"\xee\x01\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\"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\"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*\xa5\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\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\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\"\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') _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=1874 - _globals['_CAPABILITYROLE']._serialized_end=2039 + _globals['_CAPABILITYROLE']._serialized_start=3547 + _globals['_CAPABILITYROLE']._serialized_end=3747 _globals['_MODELCAPABILITY']._serialized_start=42 _globals['_MODELCAPABILITY']._serialized_end=167 _globals['_NODECAPABILITY']._serialized_start=170 - _globals['_NODECAPABILITY']._serialized_end=408 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_start=410 - _globals['_EXCHANGECAPABILITIESREQUEST']._serialized_end=487 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_start=489 - _globals['_EXCHANGECAPABILITIESRESPONSE']._serialized_end=567 - _globals['_GETNODECAPABILITYREQUEST']._serialized_start=569 - _globals['_GETNODECAPABILITYREQUEST']._serialized_end=595 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_start=597 - _globals['_GETNODECAPABILITYRESPONSE']._serialized_end=665 - _globals['_PROPOSEBLOCKREQUEST']._serialized_start=667 - _globals['_PROPOSEBLOCKREQUEST']._serialized_end=774 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_start=776 - _globals['_PROPOSEBLOCKRESPONSE']._serialized_end=897 - _globals['_TENSOR']._serialized_start=899 - _globals['_TENSOR']._serialized_end=951 - _globals['_LAYERKV']._serialized_start=953 - _globals['_LAYERKV']._serialized_end=1037 - _globals['_RESTOREREQUEST']._serialized_start=1040 - _globals['_RESTOREREQUEST']._serialized_end=1172 - _globals['_RESTORERESPONSE']._serialized_start=1174 - _globals['_RESTORERESPONSE']._serialized_end=1276 - _globals['_SEEDCONTEXTREQUEST']._serialized_start=1278 - _globals['_SEEDCONTEXTREQUEST']._serialized_end=1369 - _globals['_SEEDCONTEXTRESPONSE']._serialized_start=1371 - _globals['_SEEDCONTEXTRESPONSE']._serialized_end=1413 - _globals['_DRAFTBLOCKREQUEST']._serialized_start=1415 - _globals['_DRAFTBLOCKREQUEST']._serialized_end=1519 - _globals['_DRAFTBLOCKRESPONSE']._serialized_start=1521 - _globals['_DRAFTBLOCKRESPONSE']._serialized_end=1621 - _globals['_EXTENDCONTEXTREQUEST']._serialized_start=1623 - _globals['_EXTENDCONTEXTREQUEST']._serialized_end=1716 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_start=1718 - _globals['_EXTENDCONTEXTRESPONSE']._serialized_end=1762 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_start=1764 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONREQUEST']._serialized_end=1826 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_start=1828 - _globals['_DFLASHPROPOSERSERVICECLOSESESSIONRESPONSE']._serialized_end=1871 - _globals['_CAPABILITYSERVICE']._serialized_start=2042 - _globals['_CAPABILITYSERVICE']._serialized_end=2262 - _globals['_PROPOSERSERVICE']._serialized_start=2264 - _globals['_PROPOSERSERVICE']._serialized_end=2362 - _globals['_DFLASHPROPOSERSERVICE']._serialized_start=2365 - _globals['_DFLASHPROPOSERSERVICE']._serialized_end=2814 + _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 # @@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 66ad84c0..8141a68c 100644 --- a/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi +++ b/inference_engine/server/proto_gen/kakeya/v1/distributed_pb2.pyi @@ -14,11 +14,13 @@ class CapabilityRole(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): CAPABILITY_ROLE_PROPOSER: _ClassVar[CapabilityRole] CAPABILITY_ROLE_EMBEDDER: _ClassVar[CapabilityRole] CAPABILITY_ROLE_TOOL: _ClassVar[CapabilityRole] + CAPABILITY_ROLE_PREFILL_CACHE: _ClassVar[CapabilityRole] 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 class ModelCapability(_message.Message): __slots__ = ("model_id", "role", "quantization", "tokens_per_second") @@ -33,7 +35,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") + __slots__ = ("node_id", "grpc_address", "platform", "unified_memory_bytes", "mlx_version", "models", "announced_at_unix", "ttl_seconds", "ring_address", "caches", "endpoints") NODE_ID_FIELD_NUMBER: _ClassVar[int] GRPC_ADDRESS_FIELD_NUMBER: _ClassVar[int] PLATFORM_FIELD_NUMBER: _ClassVar[int] @@ -43,6 +45,8 @@ class NodeCapability(_message.Message): ANNOUNCED_AT_UNIX_FIELD_NUMBER: _ClassVar[int] TTL_SECONDS_FIELD_NUMBER: _ClassVar[int] RING_ADDRESS_FIELD_NUMBER: _ClassVar[int] + CACHES_FIELD_NUMBER: _ClassVar[int] + ENDPOINTS_FIELD_NUMBER: _ClassVar[int] node_id: str grpc_address: str platform: str @@ -52,7 +56,65 @@ class NodeCapability(_message.Message): announced_at_unix: float ttl_seconds: float ring_address: str - 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] = ...) -> None: ... + 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: ... + +class NodeEndpoint(_message.Message): + __slots__ = ("address", "network", "priority", "measured_rtt_ms") + ADDRESS_FIELD_NUMBER: _ClassVar[int] + NETWORK_FIELD_NUMBER: _ClassVar[int] + PRIORITY_FIELD_NUMBER: _ClassVar[int] + MEASURED_RTT_MS_FIELD_NUMBER: _ClassVar[int] + address: str + network: str + priority: int + measured_rtt_ms: float + 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") + MODEL_ID_FIELD_NUMBER: _ClassVar[int] + MODEL_REVISION_FIELD_NUMBER: _ClassVar[int] + TOKENIZER_REVISION_FIELD_NUMBER: _ClassVar[int] + CACHE_FORMAT_VERSION_FIELD_NUMBER: _ClassVar[int] + QUANTIZATION_FIELD_NUMBER: _ClassVar[int] + ROPE_HASH_FIELD_NUMBER: _ClassVar[int] + LAYER_GEOMETRY_HASH_FIELD_NUMBER: _ClassVar[int] + KV_DTYPE_FIELD_NUMBER: _ClassVar[int] + BLOCK_SIZE_TOKENS_FIELD_NUMBER: _ClassVar[int] + model_id: str + model_revision: str + tokenizer_revision: str + cache_format_version: str + quantization: str + rope_hash: str + 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: ... + +class CacheCapability(_message.Message): + __slots__ = ("compatibility", "cache_address", "cache_bytes_used", "cache_bytes_free", "entry_count", "cache_epoch", "load", "tokens_served", "bloom_filter") + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + CACHE_ADDRESS_FIELD_NUMBER: _ClassVar[int] + CACHE_BYTES_USED_FIELD_NUMBER: _ClassVar[int] + CACHE_BYTES_FREE_FIELD_NUMBER: _ClassVar[int] + ENTRY_COUNT_FIELD_NUMBER: _ClassVar[int] + CACHE_EPOCH_FIELD_NUMBER: _ClassVar[int] + LOAD_FIELD_NUMBER: _ClassVar[int] + TOKENS_SERVED_FIELD_NUMBER: _ClassVar[int] + BLOOM_FILTER_FIELD_NUMBER: _ClassVar[int] + compatibility: CacheCompatibility + cache_address: str + cache_bytes_used: int + cache_bytes_free: int + entry_count: int + cache_epoch: int + 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: ... class ExchangeCapabilitiesRequest(_message.Message): __slots__ = ("known_nodes",) @@ -76,6 +138,104 @@ class GetNodeCapabilityResponse(_message.Message): node: NodeCapability def __init__(self, node: _Optional[_Union[NodeCapability, _Mapping]] = ...) -> None: ... +class GetCacheSummaryRequest(_message.Message): + __slots__ = ("compatibility",) + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + compatibility: CacheCompatibility + def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ...) -> None: ... + +class GetCacheSummaryResponse(_message.Message): + __slots__ = ("node_id", "caches") + NODE_ID_FIELD_NUMBER: _ClassVar[int] + CACHES_FIELD_NUMBER: _ClassVar[int] + node_id: str + caches: _containers.RepeatedCompositeFieldContainer[CacheCapability] + def __init__(self, node_id: _Optional[str] = ..., caches: _Optional[_Iterable[_Union[CacheCapability, _Mapping]]] = ...) -> None: ... + +class LookupPrefixRequest(_message.Message): + __slots__ = ("compatibility", "block_hashes") + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + BLOCK_HASHES_FIELD_NUMBER: _ClassVar[int] + compatibility: CacheCompatibility + block_hashes: _containers.RepeatedScalarFieldContainer[bytes] + def __init__(self, compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ..., block_hashes: _Optional[_Iterable[bytes]] = ...) -> None: ... + +class LookupPrefixResponse(_message.Message): + __slots__ = ("node_id", "hit_block_count", "hit_token_count", "transfer_bytes", "cache_epoch", "lease_id", "lease_expires_at_unix", "payload_sha256") + NODE_ID_FIELD_NUMBER: _ClassVar[int] + HIT_BLOCK_COUNT_FIELD_NUMBER: _ClassVar[int] + HIT_TOKEN_COUNT_FIELD_NUMBER: _ClassVar[int] + TRANSFER_BYTES_FIELD_NUMBER: _ClassVar[int] + CACHE_EPOCH_FIELD_NUMBER: _ClassVar[int] + LEASE_ID_FIELD_NUMBER: _ClassVar[int] + LEASE_EXPIRES_AT_UNIX_FIELD_NUMBER: _ClassVar[int] + PAYLOAD_SHA256_FIELD_NUMBER: _ClassVar[int] + node_id: str + hit_block_count: int + hit_token_count: int + transfer_bytes: int + cache_epoch: int + lease_id: str + lease_expires_at_unix: float + payload_sha256: bytes + def __init__(self, node_id: _Optional[str] = ..., hit_block_count: _Optional[int] = ..., hit_token_count: _Optional[int] = ..., transfer_bytes: _Optional[int] = ..., cache_epoch: _Optional[int] = ..., lease_id: _Optional[str] = ..., lease_expires_at_unix: _Optional[float] = ..., payload_sha256: _Optional[bytes] = ...) -> None: ... + +class FetchBlocksRequest(_message.Message): + __slots__ = ("lease_id",) + LEASE_ID_FIELD_NUMBER: _ClassVar[int] + lease_id: str + def __init__(self, lease_id: _Optional[str] = ...) -> None: ... + +class FetchBlocksResponse(_message.Message): + __slots__ = ("block_hash", "block_index", "token_count", "chunk_index", "total_chunks", "data", "block_sha256", "cache_epoch") + BLOCK_HASH_FIELD_NUMBER: _ClassVar[int] + BLOCK_INDEX_FIELD_NUMBER: _ClassVar[int] + TOKEN_COUNT_FIELD_NUMBER: _ClassVar[int] + CHUNK_INDEX_FIELD_NUMBER: _ClassVar[int] + TOTAL_CHUNKS_FIELD_NUMBER: _ClassVar[int] + DATA_FIELD_NUMBER: _ClassVar[int] + BLOCK_SHA256_FIELD_NUMBER: _ClassVar[int] + CACHE_EPOCH_FIELD_NUMBER: _ClassVar[int] + block_hash: bytes + block_index: int + token_count: int + chunk_index: int + total_chunks: int + data: bytes + block_sha256: bytes + cache_epoch: int + def __init__(self, block_hash: _Optional[bytes] = ..., block_index: _Optional[int] = ..., token_count: _Optional[int] = ..., chunk_index: _Optional[int] = ..., total_chunks: _Optional[int] = ..., data: _Optional[bytes] = ..., block_sha256: _Optional[bytes] = ..., cache_epoch: _Optional[int] = ...) -> None: ... + +class PublishBlockRequest(_message.Message): + __slots__ = ("block_hash", "block_index", "token_count", "chunk_index", "total_chunks", "data", "block_sha256", "cache_epoch", "compatibility") + BLOCK_HASH_FIELD_NUMBER: _ClassVar[int] + BLOCK_INDEX_FIELD_NUMBER: _ClassVar[int] + TOKEN_COUNT_FIELD_NUMBER: _ClassVar[int] + CHUNK_INDEX_FIELD_NUMBER: _ClassVar[int] + TOTAL_CHUNKS_FIELD_NUMBER: _ClassVar[int] + DATA_FIELD_NUMBER: _ClassVar[int] + BLOCK_SHA256_FIELD_NUMBER: _ClassVar[int] + CACHE_EPOCH_FIELD_NUMBER: _ClassVar[int] + COMPATIBILITY_FIELD_NUMBER: _ClassVar[int] + block_hash: bytes + block_index: int + token_count: int + chunk_index: int + total_chunks: int + data: bytes + block_sha256: bytes + cache_epoch: int + compatibility: CacheCompatibility + def __init__(self, block_hash: _Optional[bytes] = ..., block_index: _Optional[int] = ..., token_count: _Optional[int] = ..., chunk_index: _Optional[int] = ..., total_chunks: _Optional[int] = ..., data: _Optional[bytes] = ..., block_sha256: _Optional[bytes] = ..., cache_epoch: _Optional[int] = ..., compatibility: _Optional[_Union[CacheCompatibility, _Mapping]] = ...) -> None: ... + +class PublishBlockResponse(_message.Message): + __slots__ = ("stored", "cache_epoch") + STORED_FIELD_NUMBER: _ClassVar[int] + CACHE_EPOCH_FIELD_NUMBER: _ClassVar[int] + stored: bool + cache_epoch: int + def __init__(self, stored: _Optional[bool] = ..., cache_epoch: _Optional[int] = ...) -> 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 3ec69370..8885f712 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 @@ -255,6 +255,219 @@ def ProposeBlock(request, _registered_method=True) +class PrefillCacheServiceStub: + """PrefillCacheService exposes immutable, content-addressed prefill K/V blocks. + Lookup is metadata-only; FetchBlocks is the point-to-point bulk data plane. + Decode never calls this service: a requester imports a hit once, computes the + missing suffix locally, and keeps the autoregressive loop local. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.GetCacheSummary = channel.unary_unary( + '/kakeya.v1.PrefillCacheService/GetCacheSummary', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryResponse.FromString, + _registered_method=True) + self.LookupPrefix = channel.unary_unary( + '/kakeya.v1.PrefillCacheService/LookupPrefix', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.LookupPrefixRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.LookupPrefixResponse.FromString, + _registered_method=True) + self.FetchBlocks = channel.unary_stream( + '/kakeya.v1.PrefillCacheService/FetchBlocks', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.FetchBlocksRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.FetchBlocksResponse.FromString, + _registered_method=True) + self.PublishBlock = channel.stream_unary( + '/kakeya.v1.PrefillCacheService/PublishBlock', + request_serializer=kakeya_dot_v1_dot_distributed__pb2.PublishBlockRequest.SerializeToString, + response_deserializer=kakeya_dot_v1_dot_distributed__pb2.PublishBlockResponse.FromString, + _registered_method=True) + + +class PrefillCacheServiceServicer: + """PrefillCacheService exposes immutable, content-addressed prefill K/V blocks. + Lookup is metadata-only; FetchBlocks is the point-to-point bulk data plane. + Decode never calls this service: a requester imports a hit once, computes the + missing suffix locally, and keeps the autoregressive loop local. + """ + + def GetCacheSummary(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 LookupPrefix(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 FetchBlocks(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 PublishBlock(self, request_iterator, 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_PrefillCacheServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'GetCacheSummary': grpc.unary_unary_rpc_method_handler( + servicer.GetCacheSummary, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryResponse.SerializeToString, + ), + 'LookupPrefix': grpc.unary_unary_rpc_method_handler( + servicer.LookupPrefix, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.LookupPrefixRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.LookupPrefixResponse.SerializeToString, + ), + 'FetchBlocks': grpc.unary_stream_rpc_method_handler( + servicer.FetchBlocks, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.FetchBlocksRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.FetchBlocksResponse.SerializeToString, + ), + 'PublishBlock': grpc.stream_unary_rpc_method_handler( + servicer.PublishBlock, + request_deserializer=kakeya_dot_v1_dot_distributed__pb2.PublishBlockRequest.FromString, + response_serializer=kakeya_dot_v1_dot_distributed__pb2.PublishBlockResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'kakeya.v1.PrefillCacheService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('kakeya.v1.PrefillCacheService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class PrefillCacheService: + """PrefillCacheService exposes immutable, content-addressed prefill K/V blocks. + Lookup is metadata-only; FetchBlocks is the point-to-point bulk data plane. + Decode never calls this service: a requester imports a hit once, computes the + missing suffix locally, and keeps the autoregressive loop local. + """ + + @staticmethod + def GetCacheSummary(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.PrefillCacheService/GetCacheSummary', + kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.GetCacheSummaryResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def LookupPrefix(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.PrefillCacheService/LookupPrefix', + kakeya_dot_v1_dot_distributed__pb2.LookupPrefixRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.LookupPrefixResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def FetchBlocks(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_stream( + request, + target, + '/kakeya.v1.PrefillCacheService/FetchBlocks', + kakeya_dot_v1_dot_distributed__pb2.FetchBlocksRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.FetchBlocksResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def PublishBlock(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_unary( + request_iterator, + target, + '/kakeya.v1.PrefillCacheService/PublishBlock', + kakeya_dot_v1_dot_distributed__pb2.PublishBlockRequest.SerializeToString, + kakeya_dot_v1_dot_distributed__pb2.PublishBlockResponse.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/inference_engine/session/coordinator.py b/inference_engine/session/coordinator.py index a9c7e0b0..5bdc8128 100644 --- a/inference_engine/session/coordinator.py +++ b/inference_engine/session/coordinator.py @@ -48,7 +48,7 @@ from __future__ import annotations -from typing import Iterable, List, Protocol +from typing import Any, Iterable, List, Protocol import torch @@ -100,6 +100,13 @@ def kv_live_bytes(self, session: Session) -> int: ... # pragma: no cover - Protocol body, never executed +class PrefillCacheHookProtocol(Protocol): + """Optional cold-prefill accelerator used by distributed cache nodes.""" + + def prepare(self, verifier: Any, token_ids: List[int]) -> int: + ... # pragma: no cover + + def _sync_slab_bytes(session: Session, verifier: "VerifierProtocol") -> None: """Mirror the verifier's current KV byte count onto the session's slab placeholder (PR-E1c). @@ -137,6 +144,7 @@ def __init__( store: SessionStore, verifier: VerifierProtocol, resolver=None, + prefill_cache: PrefillCacheHookProtocol | None = None, ) -> None: self._store = store self._verifier = verifier @@ -144,6 +152,7 @@ def __init__( # binding (multi-tenant). When None, the single ``verifier`` is used # for every session (v0.3 single-tenant behaviour, unchanged). self._resolver = resolver + self._prefill_cache = prefill_cache def _verifier_for(self, session_id: str) -> "VerifierProtocol": return self._resolver(session_id) if self._resolver else self._verifier @@ -193,7 +202,10 @@ def append_tokens( # - next_global_position = sum of all tokens ever appended # - next_token_logits predicts position == next_global_position if session.next_global_position == 0: - verifier.prefill(token_list) + if self._prefill_cache is not None: + self._prefill_cache.prepare(verifier, token_list) + else: + verifier.prefill(token_list) else: block_logits = verifier.forward_block(token_list) verifier.commit_or_truncate( diff --git a/inference_engine/session/generator.py b/inference_engine/session/generator.py index a3dc1097..6abf39ed 100644 --- a/inference_engine/session/generator.py +++ b/inference_engine/session/generator.py @@ -47,7 +47,7 @@ import time from dataclasses import dataclass -from typing import Iterator, Optional, Union +from typing import Callable, Iterator, Optional, Union import torch @@ -115,11 +115,13 @@ def __init__( store: SessionStore, verifier: VerifierProtocol, resolver=None, + on_tokens: Optional[Callable[[int], None]] = None, ) -> None: self._store = store self._verifier = verifier # PR-A3c: optional per-session verifier resolver (multi-tenant). self._resolver = resolver + self._on_tokens = on_tokens def _verifier_for(self, session_id: str) -> "VerifierProtocol": return self._resolver(session_id) if self._resolver else self._verifier @@ -242,6 +244,8 @@ def generate( # capacity, this value plateaus and the caller can # observe the architectural KV bound empirically. _sync_slab_bytes(session, verifier) + if self._on_tokens is not None: + self._on_tokens(generated_count) yield DoneEvent( stop_reason=STOP_REASON_EOS, generated_token_count=generated_count, @@ -251,6 +255,8 @@ def generate( return _sync_slab_bytes(session, verifier) + if self._on_tokens is not None: + self._on_tokens(generated_count) yield DoneEvent( stop_reason=STOP_REASON_MAX_TOKENS, generated_token_count=generated_count, diff --git a/kv_cache_proposer/proposer.py b/kv_cache_proposer/proposer.py index 4883d10b..5e599002 100644 --- a/kv_cache_proposer/proposer.py +++ b/kv_cache_proposer/proposer.py @@ -25,6 +25,7 @@ from __future__ import annotations +from contextlib import contextmanager from dataclasses import dataclass from typing import List, Optional @@ -33,6 +34,34 @@ from transformers import AutoModelForMaskedLM, AutoTokenizer +@contextmanager +def _legacy_dllm_registration_compat(): + """Bridge the dllm-hub Qwen remote code to strict Transformers 4.x. + + The checkpoint registers ``A2DQwen3LMHeadModel`` against + ``A2DQwen3Config`` but inherits Qwen3's default ``config_class``. Newer + Transformers 4.x validates that pair. Limit the correction to these exact + remote class names and restore the registry method after model import. + """ + from transformers.models.auto.auto_factory import _BaseAutoModelClass + + original = _BaseAutoModelClass.register.__func__ + + def register(cls, config_class, model_class, exist_ok=False): + if ( + config_class.__name__ == "A2DQwen3Config" + and model_class.__name__ == "A2DQwen3LMHeadModel" + ): + model_class.config_class = config_class + return original(cls, config_class, model_class, exist_ok=exist_ok) + + _BaseAutoModelClass.register = classmethod(register) + try: + yield + finally: + _BaseAutoModelClass.register = classmethod(original) + + @dataclass class ProposerConfig: model_id: str = "dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1" @@ -71,11 +100,12 @@ def __init__(self, config: Optional[ProposerConfig] = None) -> None: self.config = config or ProposerConfig() self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_id) # The dLLM checkpoint uses a custom `A2DQwen3LMHeadModel` head. - self.model = AutoModelForMaskedLM.from_pretrained( - self.config.model_id, - dtype=self.config.dtype, - trust_remote_code=True, - ) + with _legacy_dllm_registration_compat(): + self.model = AutoModelForMaskedLM.from_pretrained( + self.config.model_id, + dtype=self.config.dtype, + trust_remote_code=True, + ) self.model.to(self.config.device).eval() self.mask_id: int = self.tokenizer.mask_token_id if self.mask_id is None: diff --git a/proto/kakeya/v1/distributed.proto b/proto/kakeya/v1/distributed.proto index afa50169..8e51fa6b 100644 --- a/proto/kakeya/v1/distributed.proto +++ b/proto/kakeya/v1/distributed.proto @@ -61,6 +61,17 @@ service ProposerService { rpc ProposeBlock(ProposeBlockRequest) returns (ProposeBlockResponse); } +// PrefillCacheService exposes immutable, content-addressed prefill K/V blocks. +// Lookup is metadata-only; FetchBlocks is the point-to-point bulk data plane. +// Decode never calls this service: a requester imports a hit once, computes the +// missing suffix locally, and keeps the autoregressive loop local. +service PrefillCacheService { + rpc GetCacheSummary(GetCacheSummaryRequest) returns (GetCacheSummaryResponse); + rpc LookupPrefix(LookupPrefixRequest) returns (LookupPrefixResponse); + rpc FetchBlocks(FetchBlocksRequest) returns (stream FetchBlocksResponse); + rpc PublishBlock(stream PublishBlockRequest) returns (PublishBlockResponse); +} + // ----------------------------------------------------------------------------- // Capability messages // ----------------------------------------------------------------------------- @@ -81,6 +92,10 @@ enum CapabilityRole { // Reserved for future exchanged capabilities (ADR 0009 §4). CAPABILITY_ROLE_EMBEDDER = 3; CAPABILITY_ROLE_TOOL = 4; + + // The node can answer PrefillCacheService lookups and stream compatible + // immutable prefill K/V blocks. + CAPABILITY_ROLE_PREFILL_CACHE = 5; } // ModelCapability is one (model, role) a node offers. @@ -141,6 +156,46 @@ message NodeCapability { // Optional mlx.distributed ring endpoint for bulk-tensor flows // (ADR 0009 §4 item 4). Empty when the node is not part of a ring. string ring_address = 9; + + // Optional prefill-cache offerings. Kept separate from ModelCapability so + // placement can reject incompatible cache formats before sending a lookup. + repeated CacheCapability caches = 10; + + // Reachable interfaces ordered by operator preference (Thunderbolt before + // LAN/Tailscale). grpc_address remains the compatibility default. + repeated NodeEndpoint endpoints = 11; +} + +message NodeEndpoint { + string address = 1; + string network = 2; // thunderbolt|lan|tailscale|public + uint32 priority = 3; + double measured_rtt_ms = 4; +} + +message CacheCompatibility { + string model_id = 1; + string model_revision = 2; + string tokenizer_revision = 3; + string cache_format_version = 4; + string quantization = 5; + string rope_hash = 6; + string layer_geometry_hash = 7; + string kv_dtype = 8; + uint32 block_size_tokens = 9; +} + +message CacheCapability { + CacheCompatibility compatibility = 1; + string cache_address = 2; + uint64 cache_bytes_used = 3; + uint64 cache_bytes_free = 4; + uint64 entry_count = 5; + uint64 cache_epoch = 6; + double load = 7; + uint64 tokens_served = 8; + // Compact probabilistic summary. Empty means "query me directly". + bytes bloom_filter = 9; } message ExchangeCapabilitiesRequest { @@ -160,6 +215,68 @@ message GetNodeCapabilityResponse { NodeCapability node = 1; } +// ----------------------------------------------------------------------------- +// Distributed prefill K/V cache +// ----------------------------------------------------------------------------- + +message GetCacheSummaryRequest { + CacheCompatibility compatibility = 1; +} + +message GetCacheSummaryResponse { + string node_id = 1; + repeated CacheCapability caches = 2; +} + +message LookupPrefixRequest { + CacheCompatibility compatibility = 1; + // Ordered chained hashes from token block zero onward. + repeated bytes block_hashes = 2; +} + +message LookupPrefixResponse { + string node_id = 1; + uint32 hit_block_count = 2; + uint64 hit_token_count = 3; + uint64 transfer_bytes = 4; + uint64 cache_epoch = 5; + string lease_id = 6; + double lease_expires_at_unix = 7; + bytes payload_sha256 = 8; +} + +message FetchBlocksRequest { + string lease_id = 1; +} + +message FetchBlocksResponse { + bytes block_hash = 1; + uint32 block_index = 2; + uint32 token_count = 3; + uint32 chunk_index = 4; + uint32 total_chunks = 5; + bytes data = 6; + bytes block_sha256 = 7; + uint64 cache_epoch = 8; +} + +message PublishBlockRequest { + bytes block_hash = 1; + uint32 block_index = 2; + uint32 token_count = 3; + uint32 chunk_index = 4; + uint32 total_chunks = 5; + bytes data = 6; + bytes block_sha256 = 7; + uint64 cache_epoch = 8; + CacheCompatibility compatibility = 9; +} + +message PublishBlockResponse { + bool stored = 1; + uint64 cache_epoch = 2; +} + // ----------------------------------------------------------------------------- // Remote proposal messages // ----------------------------------------------------------------------------- diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index e464a51d..d5529ade 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -67,21 +67,39 @@ def _resolve_kv_dims(verifier) -> Tuple[int, int, int]: verifier's HF config means the per-session byte numbers reported over gRPC match what the verifier is actually holding. """ - cfg = verifier.model.config - # Gemma 4 is multimodal: decoder dims live under config.text_config. - cfg = getattr(cfg, "text_config", None) or cfg - num_layers = int(getattr(cfg, "num_hidden_layers")) - # Qwen3 / Gemma / DeepSeek all support GQA — kv-heads is the - # dimension that matters for KV cache size, not attention-heads. - num_kv_heads = int( - getattr(cfg, "num_key_value_heads", None) - or getattr(cfg, "num_attention_heads") - ) - head_dim = int( - getattr(cfg, "head_dim", None) - or (cfg.hidden_size // cfg.num_attention_heads) - ) - return num_layers, num_kv_heads, head_dim + try: + cfg = ( + getattr(verifier.model, "config", None) + or getattr(verifier.model, "args", None) + ) + if cfg is None: + raise AttributeError("model exposes neither config nor args") + if hasattr(cfg, "get_text_config"): + cfg = cfg.get_text_config() + cfg = getattr(cfg, "text_config", None) or cfg + num_layers = int(getattr(cfg, "num_hidden_layers")) + num_kv_heads = int( + getattr(cfg, "num_key_value_heads", None) + or getattr(cfg, "num_attention_heads") + ) + head_dim = int( + getattr(cfg, "head_dim", None) + or (cfg.hidden_size // cfg.num_attention_heads) + ) + return num_layers, num_kv_heads, head_dim + except AttributeError: + from inference_engine.backends.mlx.cross_model_dlm_verifier import ( + per_layer_kv_geometry, + resolve_mlx_text_model, + ) + geometry = per_layer_kv_geometry(resolve_mlx_text_model(verifier.model)) + if not geometry: + raise + return ( + len(geometry), + max(item[0] for item in geometry), + max(item[1] for item in geometry), + ) def _build_verifier( @@ -143,7 +161,12 @@ def _total_memory_bytes() -> int: return 0 -def _build_capability_registry(args: argparse.Namespace, *, backend: str): +def _build_capability_registry( + args: argparse.Namespace, + *, + backend: str, + cache_store=None, +): """Build this node's CapabilityRegistry from CLI flags + probes.""" import platform as _platform import time @@ -155,6 +178,7 @@ def _build_capability_registry(args: argparse.Namespace, *, backend: str): CapabilityRole, ModelCapability, NodeCapability, + NodeEndpoint, ) from inference_engine.distributed.mlx_ring import probe_ring_environment @@ -173,6 +197,24 @@ def _build_capability_registry(args: argparse.Namespace, *, backend: str): quantization="none", ), ) + caches = () + if cache_store is not None: + from inference_engine.distributed.prefill_cache_service import ( + cache_capability, + ) + models.append( + ModelCapability( + model_id=args.verifier_id, + role=CapabilityRole.PREFILL_CACHE, + quantization=args.cache_quantization, + ), + ) + caches = ( + cache_capability( + cache_store, + cache_address=args.cache_advertise or args.advertise or args.bind, + ), + ) mlx_env = probe_environment() ring_env = probe_ring_environment() @@ -187,6 +229,15 @@ def _build_capability_registry(args: argparse.Namespace, *, backend: str): announced_at_unix=time.time(), ttl_seconds=args.capability_ttl_s, ring_address=ring_env.ring_address(hostname), + caches=caches, + endpoints=( + NodeEndpoint( + address=args.advertise or args.bind, + network=args.network_label, + priority=args.network_priority, + measured_rtt_ms=args.measured_rtt_ms, + ), + ), ) _LOG.info("capability card: %s @ %s ring=%r models=%s", self_card.node_id, self_card.grpc_address, @@ -195,11 +246,32 @@ def _build_capability_registry(args: argparse.Namespace, *, backend: str): return CapabilityRegistry(self_card=self_card) -async def _exchange_loop(registry, peers, interval_s: float) -> None: +async def _exchange_loop( + registry, + peers, + interval_s: float, + *, + cache_store=None, + cache_address: str = "", +) -> None: """Periodic gossip with seed peers until the task is cancelled.""" from inference_engine.distributed.exchange import exchange_once while True: + if cache_store is not None: + from dataclasses import replace + from inference_engine.distributed.prefill_cache_service import ( + cache_capability, + ) + registry.self_card = replace( + registry.self_card, + caches=( + cache_capability( + cache_store, + cache_address=cache_address, + ), + ), + ) report = await exchange_once(registry, peers) if report.errors: _LOG.warning("gossip errors: %s", report.errors) @@ -211,6 +283,42 @@ async def _exchange_loop(registry, peers, interval_s: float) -> None: await asyncio.sleep(interval_s) +def _build_token_telemetry_callback(args): + if not args.network_telemetry_url: + return None + + def report(count: int, kv_assisted: int = 0) -> None: + import json + import threading + import urllib.request + + def send() -> None: + body = json.dumps({ + "node_id": args.node_id or "runtime", + "completed": int(count), + "kv_assisted": int(kv_assisted), + }).encode() + headers = {"Content-Type": "application/json"} + if args.network_telemetry_api_key: + headers["X-API-Key"] = args.network_telemetry_api_key + try: + with urllib.request.urlopen( + urllib.request.Request( + args.network_telemetry_url, + data=body, + headers=headers, + ), + timeout=3, + ): + pass + except OSError: + _LOG.warning("failed to publish token telemetry") + + threading.Thread(target=send, daemon=True).start() + + return report + + async def _serve(args: argparse.Namespace) -> int: logging.basicConfig( level=getattr(logging, args.log_level.upper(), logging.INFO), @@ -259,6 +367,53 @@ async def _serve(args: argparse.Namespace) -> int: ) pool = SlabPool(num_slabs=args.capacity, slab_config=slab_cfg) + prefill_store = None + prefill_hook = 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.prefill_cache import PrefixCacheStore + from inference_engine.distributed.prefill_cache_runtime import ( + DistributedPrefillCacheHook, + ) + + geometry = f"{num_layers}:{num_kv_heads}:{head_dim}" + compatibility = CacheCompatibility( + model_id=args.cache_model_id or args.verifier_id, + model_revision=args.model_revision, + tokenizer_revision=args.tokenizer_revision, + cache_format_version=args.cache_format_version, + quantization=args.cache_quantization, + rope_hash=args.rope_hash, + layer_geometry_hash=hashlib.sha256(geometry.encode()).hexdigest(), + kv_dtype=args.cache_kv_dtype, + block_size_tokens=args.cache_block_tokens, + ) + prefill_store = PrefixCacheStore( + compatibility, + max_bytes=int(args.prefill_cache_gb * (1 << 30)), + node_id=args.node_id or (__import__("platform").node() or "localhost"), + ) + telemetry_callback = _build_token_telemetry_callback(args) + prefill_hook = DistributedPrefillCacheHook( + prefill_store, + peers=args.cache_peer, + lookup_timeout_s=args.cache_lookup_timeout_s, + fetch_timeout_s=args.cache_fetch_timeout_s, + on_reuse=( + (lambda count: telemetry_callback(count, count)) + if telemetry_callback is not None else None + ), + ) + _LOG.info( + "distributed prefill cache enabled: %.2f GiB, block=%d, peers=%s", + args.prefill_cache_gb, + args.cache_block_tokens, + args.cache_peer, + ) + # PR-A3c: per-session binding for true multi-tenant serving. Each session # gets its own verifier adapter (isolated KV) sharing the model weights via # the adapter's spawn(); the registry doubles as cache_inspector + resolver. @@ -282,19 +437,36 @@ async def _serve(args: argparse.Namespace) -> int: slab_pool=pool, ) resolver = registry.get if registry is not None else None - append_coord = AppendTokensCoordinator(store, verifier, resolver=resolver) - gen_coord = GenerationCoordinator(store, verifier, resolver=resolver) + append_coord = AppendTokensCoordinator( + store, + verifier, + resolver=resolver, + prefill_cache=prefill_hook, + ) + gen_coord = GenerationCoordinator( + store, + verifier, + resolver=resolver, + on_tokens=_build_token_telemetry_callback(args), + ) # Multi-host capability plane (ADR 0009): only constructed when # the operator opts in via --enable-capability-exchange (implied # by --peer / --serve-ngram-proposer). distributed_enabled = bool( - args.enable_capability_exchange or args.peer or args.serve_ngram_proposer + args.enable_capability_exchange + or args.peer + or args.serve_ngram_proposer + or args.enable_prefill_cache ) registry = None proposers = None if distributed_enabled: - registry = _build_capability_registry(args, backend=args.backend) + registry = _build_capability_registry( + args, + backend=args.backend, + cache_store=prefill_store, + ) if args.serve_ngram_proposer: from inference_engine.distributed.capability import NGRAM_MODEL_ID from inference_engine.distributed.ngram import NGramProposer @@ -312,15 +484,55 @@ async def _serve(args: argparse.Namespace) -> int: on_session_close=on_session_close, capability_registry=registry, proposers=proposers, + prefill_cache_store=prefill_store, + prefill_cache_address=args.cache_advertise or args.advertise or args.bind, ) await server.start() _LOG.info("kakeya gRPC RuntimeService listening on %s", args.bind) + http_server = None + http_task = None + if args.network_http_port: + if registry is None or prefill_store is None: + raise SystemExit( + "--network-http-port requires --enable-prefill-cache", + ) + import uvicorn + from inference_engine.network.api import create_network_app + from inference_engine.network.state import NetworkState + + network_state = NetworkState( + registry, + prefill_store, + state_path=args.network_state, + ) + http_server = uvicorn.Server(uvicorn.Config( + create_network_app( + network_state, + api_key=args.network_api_key, + ), + host=args.network_http_host, + port=args.network_http_port, + log_level=args.log_level.lower(), + )) + http_task = asyncio.create_task(http_server.serve()) + _LOG.info( + "inference-network dashboard listening on http://%s:%d/network", + args.network_http_host, + args.network_http_port, + ) + exchange_task = None if registry is not None and args.peer: exchange_task = asyncio.create_task( - _exchange_loop(registry, list(args.peer), args.exchange_interval_s), + _exchange_loop( + registry, + list(args.peer), + args.exchange_interval_s, + cache_store=prefill_store, + cache_address=args.cache_advertise or args.advertise or args.bind, + ), ) _LOG.info( "capability gossip every %.1fs with peers: %s", @@ -347,6 +559,10 @@ def _on_signal(sig: int) -> None: await exchange_task except asyncio.CancelledError: pass + if http_server is not None: + http_server.should_exit = True + if http_task is not None: + await http_task await server.stop(grace=args.shutdown_grace_s) _LOG.info("kakeya gRPC RuntimeService stopped cleanly") return 0 @@ -419,6 +635,48 @@ def main() -> int: ap.add_argument("--capability-ttl-s", type=float, default=120.0, help="TTL of this node's capability card; peers drop " "it if not refreshed within this window.") + # --- Distributed prefill K/V cache -------------------------------- + ap.add_argument("--enable-prefill-cache", action="store_true", + help="Enable immutable longest-prefix prefill K/V reuse " + "and serve PrefillCacheService.") + ap.add_argument("--prefill-cache-gb", type=float, default=4.0, + help="Maximum local in-memory snapshot cache size.") + ap.add_argument("--cache-peer", action="append", default=[], + help="Peer PrefillCacheService host:port. Repeatable.") + ap.add_argument("--cache-advertise", default="", + help="Peer-reachable PrefillCacheService address. " + "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-model-id", default="", + help="Logical cache model id; defaults to --verifier-id. " + "Use this when verifier-id is a host-specific path.") + ap.add_argument("--cache-quantization", default="", + help="Exact model quantization label used in compatibility.") + ap.add_argument("--cache-kv-dtype", default="bfloat16") + ap.add_argument("--model-revision", default="", + help="Exact weights revision/hash for cache compatibility.") + ap.add_argument("--tokenizer-revision", default="", + help="Exact tokenizer/chat-template revision.") + ap.add_argument("--rope-hash", default="", + 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("--network-label", default="lan", + help="Advertised interface: thunderbolt|lan|tailscale|public.") + ap.add_argument("--network-priority", type=int, default=50) + ap.add_argument("--measured-rtt-ms", type=float, default=0.0) + ap.add_argument("--network-http-host", default="127.0.0.1") + ap.add_argument("--network-http-port", type=int, default=0, + help="Serve the inference-network API/dashboard; 0 disables.") + ap.add_argument("--network-state", + default="~/.kakeya/inference_network.json") + ap.add_argument("--network-api-key", default="", + help="X-API-Key required for registration/group/telemetry writes.") + ap.add_argument("--network-telemetry-url", default="", + help="Optional POST endpoint receiving completed token counters.") + ap.add_argument("--network-telemetry-api-key", default="") ap.add_argument("--skip-cache-check", action="store_true", help="Skip the HF-cache pre-flight assertion. By " "default the server fails fast if the verifier " diff --git a/scripts/start_prefill_cache_node.py b/scripts/start_prefill_cache_node.py new file mode 100644 index 00000000..4a55f3ae --- /dev/null +++ b/scripts/start_prefill_cache_node.py @@ -0,0 +1,196 @@ +"""Start a standalone Kakeya P2P prefill-cache node and optional dashboard. + +This process is useful for cache-only peers and control-plane deployment. A +full inference node should instead pass --enable-prefill-cache to +start_grpc_runtime_server.py so the same store is wired into cold prefill. +""" + +from __future__ import annotations + +import argparse +import asyncio +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 uvicorn + +from inference_engine.distributed.capability import ( + CacheCompatibility, + CapabilityRegistry, + CapabilityRole, + ModelCapability, + NodeCapability, + NodeEndpoint, +) +from inference_engine.distributed.exchange import ( + add_capability_service, + exchange_once, +) +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.network.api import create_network_app +from inference_engine.network.state import NetworkState + + +def physical_memory_bytes() -> int: + try: + return int(__import__("os").sysconf("SC_PAGE_SIZE") + * __import__("os").sysconf("SC_PHYS_PAGES")) + except (ValueError, OSError, AttributeError): + return 0 + + +async def serve(args) -> None: + compatibility = CacheCompatibility( + model_id=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, + ) + store = PrefixCacheStore( + compatibility, + max_bytes=int(args.cache_gb * (1 << 30)), + node_id=args.node_id, + ) + card = NodeCapability( + node_id=args.node_id, + grpc_address=args.advertise, + platform=args.platform or f"{platform.system()}-{platform.machine()}", + unified_memory_bytes=args.memory_bytes or physical_memory_bytes(), + models=( + ModelCapability( + args.model_id, + CapabilityRole.PREFILL_CACHE, + args.quantization, + ), + ), + announced_at_unix=time.time(), + ttl_seconds=args.ttl_seconds, + caches=(cache_capability(store, cache_address=args.advertise),), + endpoints=( + NodeEndpoint( + args.advertise, + args.network, + args.priority, + args.rtt_ms, + ), + ), + ) + registry = CapabilityRegistry(card) + grpc_server = grpc.aio.server() + add_capability_service(grpc_server, registry) + add_prefill_cache_service( + grpc_server, + store, + cache_address=args.advertise, + ) + grpc_server.add_insecure_port(args.bind) + await grpc_server.start() + + network_state = NetworkState( + registry, + store, + state_path=args.state_path, + ) + http_server = None + http_task = None + if args.http_port: + http_server = uvicorn.Server(uvicorn.Config( + create_network_app(network_state, api_key=args.api_key), + host=args.http_host, + port=args.http_port, + log_level="info", + )) + http_task = asyncio.create_task(http_server.serve()) + + stop = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, stop.set) + + async def gossip(): + while not stop.is_set(): + registry.self_card = replace( + registry.self_card, + caches=(cache_capability(store, cache_address=args.advertise),), + ) + await exchange_once( + registry, + args.peer, + timeout_s=args.gossip_timeout_s, + ) + try: + await asyncio.wait_for(stop.wait(), args.gossip_interval_s) + except asyncio.TimeoutError: + pass + + gossip_task = asyncio.create_task(gossip()) + print( + f"[prefill-cache] {args.node_id} grpc={args.bind} " + f"advertise={args.advertise} peers={args.peer}", + flush=True, + ) + if args.http_port: + print( + f"[prefill-cache] dashboard=http://{args.http_host}:{args.http_port}/network", + flush=True, + ) + await stop.wait() + gossip_task.cancel() + if http_server is not None: + http_server.should_exit = True + if http_task is not None: + await http_task + await grpc_server.stop(3) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--node-id", required=True) + ap.add_argument("--bind", default="0.0.0.0:52051") + ap.add_argument("--advertise", required=True) + ap.add_argument("--peer", action="append", default=[]) + 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("--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("--cache-gb", type=float, default=4) + ap.add_argument("--platform", default="") + ap.add_argument("--memory-bytes", type=int, default=0) + ap.add_argument("--network", default="lan") + ap.add_argument("--priority", type=int, default=50) + ap.add_argument("--rtt-ms", type=float, default=0) + ap.add_argument("--ttl-seconds", type=float, default=120) + ap.add_argument("--gossip-interval-s", type=float, default=10) + ap.add_argument("--gossip-timeout-s", type=float, default=3) + ap.add_argument("--http-host", default="127.0.0.1") + ap.add_argument("--http-port", type=int, default=0) + ap.add_argument("--state-path", default="~/.kakeya/inference_network.json") + ap.add_argument("--api-key", default="") + args = ap.parse_args() + 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 318f9cb0..e35e7130 100644 --- a/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts +++ b/sdks/typescript/src/proto_gen/kakeya/v1/distributed.ts @@ -11,7 +11,11 @@ import { type ChannelCredentials, Client, type ClientOptions, + type ClientReadableStream, type ClientUnaryCall, + type ClientWritableStream, + type handleClientStreamingCall, + type handleServerStreamingCall, type handleUnaryCall, makeGenericClientConstructor, type Metadata, @@ -38,6 +42,11 @@ export enum CapabilityRole { /** EMBEDDER - Reserved for future exchanged capabilities (ADR 0009 §4). */ EMBEDDER = 3, TOOL = 4, + /** + * PREFILL_CACHE - The node can answer PrefillCacheService lookups and stream compatible + * immutable prefill K/V blocks. + */ + PREFILL_CACHE = 5, UNRECOGNIZED = -1, } @@ -58,6 +67,9 @@ export function capabilityRoleFromJSON(object: any): CapabilityRole { case 4: case "CAPABILITY_ROLE_TOOL": return CapabilityRole.TOOL; + case 5: + case "CAPABILITY_ROLE_PREFILL_CACHE": + return CapabilityRole.PREFILL_CACHE; case -1: case "UNRECOGNIZED": default: @@ -77,6 +89,8 @@ export function capabilityRoleToJSON(object: CapabilityRole): string { return "CAPABILITY_ROLE_EMBEDDER"; case CapabilityRole.TOOL: return "CAPABILITY_ROLE_TOOL"; + case CapabilityRole.PREFILL_CACHE: + return "CAPABILITY_ROLE_PREFILL_CACHE"; case CapabilityRole.UNRECOGNIZED: default: return "UNRECOGNIZED"; @@ -148,6 +162,49 @@ export interface NodeCapability { * (ADR 0009 §4 item 4). Empty when the node is not part of a ring. */ ringAddress: string; + /** + * Optional prefill-cache offerings. Kept separate from ModelCapability so + * placement can reject incompatible cache formats before sending a lookup. + */ + caches: CacheCapability[]; + /** + * Reachable interfaces ordered by operator preference (Thunderbolt before + * LAN/Tailscale). grpc_address remains the compatibility default. + */ + endpoints: NodeEndpoint[]; +} + +export interface NodeEndpoint { + address: string; + /** thunderbolt|lan|tailscale|public */ + network: string; + priority: number; + measuredRttMs: number; +} + +export interface CacheCompatibility { + modelId: string; + modelRevision: string; + tokenizerRevision: string; + cacheFormatVersion: string; + quantization: string; + ropeHash: string; + layerGeometryHash: string; + kvDtype: string; + blockSizeTokens: number; +} + +export interface CacheCapability { + compatibility?: CacheCompatibility | undefined; + cacheAddress: string; + cacheBytesUsed: string; + cacheBytesFree: string; + entryCount: string; + cacheEpoch: string; + load: number; + tokensServed: string; + /** Compact probabilistic summary. Empty means "query me directly". */ + bloomFilter: Uint8Array; } export interface ExchangeCapabilitiesRequest { @@ -168,6 +225,66 @@ export interface GetNodeCapabilityResponse { node?: NodeCapability | undefined; } +export interface GetCacheSummaryRequest { + compatibility?: CacheCompatibility | undefined; +} + +export interface GetCacheSummaryResponse { + nodeId: string; + caches: CacheCapability[]; +} + +export interface LookupPrefixRequest { + compatibility?: + | CacheCompatibility + | undefined; + /** Ordered chained hashes from token block zero onward. */ + blockHashes: Uint8Array[]; +} + +export interface LookupPrefixResponse { + nodeId: string; + hitBlockCount: number; + hitTokenCount: string; + transferBytes: string; + cacheEpoch: string; + leaseId: string; + leaseExpiresAtUnix: number; + payloadSha256: Uint8Array; +} + +export interface FetchBlocksRequest { + leaseId: string; +} + +export interface FetchBlocksResponse { + blockHash: Uint8Array; + blockIndex: number; + tokenCount: number; + chunkIndex: number; + totalChunks: number; + data: Uint8Array; + blockSha256: Uint8Array; + cacheEpoch: string; +} + +export interface PublishBlockRequest { + blockHash: Uint8Array; + blockIndex: number; + tokenCount: number; + chunkIndex: number; + totalChunks: number; + data: Uint8Array; + blockSha256: Uint8Array; + cacheEpoch: string; + compatibility?: CacheCompatibility | undefined; +} + +export interface PublishBlockResponse { + stored: boolean; + cacheEpoch: string; +} + export interface ProposeBlockRequest { /** * The committed prefix (prompt + accepted tokens), raw token ids in @@ -423,6 +540,8 @@ function createBaseNodeCapability(): NodeCapability { announcedAtUnix: 0, ttlSeconds: 0, ringAddress: "", + caches: [], + endpoints: [], }; } @@ -455,6 +574,12 @@ export const NodeCapability: MessageFns = { if (message.ringAddress !== "") { writer.uint32(74).string(message.ringAddress); } + for (const v of message.caches) { + CacheCapability.encode(v!, writer.uint32(82).fork()).join(); + } + for (const v of message.endpoints) { + NodeEndpoint.encode(v!, writer.uint32(90).fork()).join(); + } return writer; }, @@ -537,6 +662,22 @@ export const NodeCapability: MessageFns = { message.ringAddress = reader.string(); continue; } + case 10: { + if (tag !== 82) { + break; + } + + message.caches.push(CacheCapability.decode(reader, reader.uint32())); + continue; + } + case 11: { + if (tag !== 90) { + break; + } + + message.endpoints.push(NodeEndpoint.decode(reader, reader.uint32())); + continue; + } } if ((tag & 7) === 4 || tag === 0) { break; @@ -587,6 +728,12 @@ export const NodeCapability: MessageFns = { : isSet(object.ring_address) ? globalThis.String(object.ring_address) : "", + caches: globalThis.Array.isArray(object?.caches) + ? object.caches.map((e: any) => CacheCapability.fromJSON(e)) + : [], + endpoints: globalThis.Array.isArray(object?.endpoints) + ? object.endpoints.map((e: any) => NodeEndpoint.fromJSON(e)) + : [], }; }, @@ -619,6 +766,12 @@ export const NodeCapability: MessageFns = { if (message.ringAddress !== "") { obj.ringAddress = message.ringAddress; } + if (message.caches?.length) { + obj.caches = message.caches.map((e) => CacheCapability.toJSON(e)); + } + if (message.endpoints?.length) { + obj.endpoints = message.endpoints.map((e) => NodeEndpoint.toJSON(e)); + } return obj; }, @@ -636,26 +789,174 @@ export const NodeCapability: MessageFns = { message.announcedAtUnix = object.announcedAtUnix ?? 0; message.ttlSeconds = object.ttlSeconds ?? 0; message.ringAddress = object.ringAddress ?? ""; + message.caches = object.caches?.map((e) => CacheCapability.fromPartial(e)) || []; + message.endpoints = object.endpoints?.map((e) => NodeEndpoint.fromPartial(e)) || []; return message; }, }; -function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest { - return { knownNodes: [] }; +function createBaseNodeEndpoint(): NodeEndpoint { + return { address: "", network: "", priority: 0, measuredRttMs: 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 NodeEndpoint: MessageFns = { + encode(message: NodeEndpoint, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.address !== "") { + writer.uint32(10).string(message.address); + } + if (message.network !== "") { + writer.uint32(18).string(message.network); + } + if (message.priority !== 0) { + writer.uint32(24).uint32(message.priority); + } + if (message.measuredRttMs !== 0) { + writer.uint32(33).double(message.measuredRttMs); } return writer; }, - decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesRequest { + decode(input: BinaryReader | Uint8Array, length?: number): NodeEndpoint { const reader = input instanceof BinaryReader ? input : new BinaryReader(input); const end = length === undefined ? reader.len : reader.pos + length; - const message = createBaseExchangeCapabilitiesRequest(); + const message = createBaseNodeEndpoint(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.address = reader.string(); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.network = reader.string(); + continue; + } + case 3: { + if (tag !== 24) { + break; + } + + message.priority = reader.uint32(); + continue; + } + case 4: { + if (tag !== 33) { + break; + } + + message.measuredRttMs = reader.double(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): NodeEndpoint { + return { + address: isSet(object.address) ? globalThis.String(object.address) : "", + network: isSet(object.network) ? globalThis.String(object.network) : "", + priority: isSet(object.priority) ? globalThis.Number(object.priority) : 0, + measuredRttMs: isSet(object.measuredRttMs) + ? globalThis.Number(object.measuredRttMs) + : isSet(object.measured_rtt_ms) + ? globalThis.Number(object.measured_rtt_ms) + : 0, + }; + }, + + toJSON(message: NodeEndpoint): unknown { + const obj: any = {}; + if (message.address !== "") { + obj.address = message.address; + } + if (message.network !== "") { + obj.network = message.network; + } + if (message.priority !== 0) { + obj.priority = Math.round(message.priority); + } + if (message.measuredRttMs !== 0) { + obj.measuredRttMs = message.measuredRttMs; + } + return obj; + }, + + create, I>>(base?: I): NodeEndpoint { + return NodeEndpoint.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): NodeEndpoint { + const message = createBaseNodeEndpoint(); + message.address = object.address ?? ""; + message.network = object.network ?? ""; + message.priority = object.priority ?? 0; + message.measuredRttMs = object.measuredRttMs ?? 0; + return message; + }, +}; + +function createBaseCacheCompatibility(): CacheCompatibility { + return { + modelId: "", + modelRevision: "", + tokenizerRevision: "", + cacheFormatVersion: "", + quantization: "", + ropeHash: "", + layerGeometryHash: "", + kvDtype: "", + blockSizeTokens: 0, + }; +} + +export const CacheCompatibility: MessageFns = { + encode(message: CacheCompatibility, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.modelId !== "") { + writer.uint32(10).string(message.modelId); + } + if (message.modelRevision !== "") { + writer.uint32(18).string(message.modelRevision); + } + if (message.tokenizerRevision !== "") { + writer.uint32(26).string(message.tokenizerRevision); + } + if (message.cacheFormatVersion !== "") { + writer.uint32(34).string(message.cacheFormatVersion); + } + if (message.quantization !== "") { + writer.uint32(42).string(message.quantization); + } + if (message.ropeHash !== "") { + writer.uint32(50).string(message.ropeHash); + } + if (message.layerGeometryHash !== "") { + writer.uint32(58).string(message.layerGeometryHash); + } + if (message.kvDtype !== "") { + writer.uint32(66).string(message.kvDtype); + } + if (message.blockSizeTokens !== 0) { + writer.uint32(72).uint32(message.blockSizeTokens); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): CacheCompatibility { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseCacheCompatibility(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { @@ -664,7 +965,1251 @@ export const ExchangeCapabilitiesRequest: MessageFns, I>>(base?: I): CacheCompatibility { + return CacheCompatibility.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): CacheCompatibility { + const message = createBaseCacheCompatibility(); + message.modelId = object.modelId ?? ""; + message.modelRevision = object.modelRevision ?? ""; + message.tokenizerRevision = object.tokenizerRevision ?? ""; + message.cacheFormatVersion = object.cacheFormatVersion ?? ""; + message.quantization = object.quantization ?? ""; + message.ropeHash = object.ropeHash ?? ""; + message.layerGeometryHash = object.layerGeometryHash ?? ""; + message.kvDtype = object.kvDtype ?? ""; + message.blockSizeTokens = object.blockSizeTokens ?? 0; + return message; + }, +}; + +function createBaseCacheCapability(): CacheCapability { + return { + compatibility: undefined, + cacheAddress: "", + cacheBytesUsed: "0", + cacheBytesFree: "0", + entryCount: "0", + cacheEpoch: "0", + load: 0, + tokensServed: "0", + bloomFilter: new Uint8Array(0), + }; +} + +export const CacheCapability: MessageFns = { + encode(message: CacheCapability, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(10).fork()).join(); + } + if (message.cacheAddress !== "") { + writer.uint32(18).string(message.cacheAddress); + } + if (message.cacheBytesUsed !== "0") { + writer.uint32(24).uint64(message.cacheBytesUsed); + } + if (message.cacheBytesFree !== "0") { + writer.uint32(32).uint64(message.cacheBytesFree); + } + if (message.entryCount !== "0") { + writer.uint32(40).uint64(message.entryCount); + } + if (message.cacheEpoch !== "0") { + writer.uint32(48).uint64(message.cacheEpoch); + } + if (message.load !== 0) { + writer.uint32(57).double(message.load); + } + if (message.tokensServed !== "0") { + writer.uint32(64).uint64(message.tokensServed); + } + if (message.bloomFilter.length !== 0) { + writer.uint32(74).bytes(message.bloomFilter); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): CacheCapability { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseCacheCapability(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.compatibility = CacheCompatibility.decode(reader, reader.uint32()); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.cacheAddress = reader.string(); + continue; + } + case 3: { + if (tag !== 24) { + break; + } + + message.cacheBytesUsed = reader.uint64().toString(); + continue; + } + case 4: { + if (tag !== 32) { + break; + } + + message.cacheBytesFree = reader.uint64().toString(); + continue; + } + case 5: { + if (tag !== 40) { + break; + } + + message.entryCount = reader.uint64().toString(); + continue; + } + case 6: { + if (tag !== 48) { + break; + } + + message.cacheEpoch = reader.uint64().toString(); + continue; + } + case 7: { + if (tag !== 57) { + break; + } + + message.load = reader.double(); + continue; + } + case 8: { + if (tag !== 64) { + break; + } + + message.tokensServed = reader.uint64().toString(); + continue; + } + case 9: { + if (tag !== 74) { + break; + } + + message.bloomFilter = reader.bytes(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): CacheCapability { + return { + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + cacheAddress: isSet(object.cacheAddress) + ? globalThis.String(object.cacheAddress) + : isSet(object.cache_address) + ? globalThis.String(object.cache_address) + : "", + cacheBytesUsed: isSet(object.cacheBytesUsed) + ? globalThis.String(object.cacheBytesUsed) + : isSet(object.cache_bytes_used) + ? globalThis.String(object.cache_bytes_used) + : "0", + cacheBytesFree: isSet(object.cacheBytesFree) + ? globalThis.String(object.cacheBytesFree) + : isSet(object.cache_bytes_free) + ? globalThis.String(object.cache_bytes_free) + : "0", + entryCount: isSet(object.entryCount) + ? globalThis.String(object.entryCount) + : isSet(object.entry_count) + ? globalThis.String(object.entry_count) + : "0", + cacheEpoch: isSet(object.cacheEpoch) + ? globalThis.String(object.cacheEpoch) + : isSet(object.cache_epoch) + ? globalThis.String(object.cache_epoch) + : "0", + load: isSet(object.load) ? globalThis.Number(object.load) : 0, + tokensServed: isSet(object.tokensServed) + ? globalThis.String(object.tokensServed) + : isSet(object.tokens_served) + ? globalThis.String(object.tokens_served) + : "0", + bloomFilter: isSet(object.bloomFilter) + ? bytesFromBase64(object.bloomFilter) + : isSet(object.bloom_filter) + ? bytesFromBase64(object.bloom_filter) + : new Uint8Array(0), + }; + }, + + toJSON(message: CacheCapability): unknown { + const obj: any = {}; + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + if (message.cacheAddress !== "") { + obj.cacheAddress = message.cacheAddress; + } + if (message.cacheBytesUsed !== "0") { + obj.cacheBytesUsed = message.cacheBytesUsed; + } + if (message.cacheBytesFree !== "0") { + obj.cacheBytesFree = message.cacheBytesFree; + } + if (message.entryCount !== "0") { + obj.entryCount = message.entryCount; + } + if (message.cacheEpoch !== "0") { + obj.cacheEpoch = message.cacheEpoch; + } + if (message.load !== 0) { + obj.load = message.load; + } + if (message.tokensServed !== "0") { + obj.tokensServed = message.tokensServed; + } + if (message.bloomFilter.length !== 0) { + obj.bloomFilter = base64FromBytes(message.bloomFilter); + } + return obj; + }, + + create, I>>(base?: I): CacheCapability { + return CacheCapability.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): CacheCapability { + const message = createBaseCacheCapability(); + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + message.cacheAddress = object.cacheAddress ?? ""; + message.cacheBytesUsed = object.cacheBytesUsed ?? "0"; + message.cacheBytesFree = object.cacheBytesFree ?? "0"; + message.entryCount = object.entryCount ?? "0"; + message.cacheEpoch = object.cacheEpoch ?? "0"; + message.load = object.load ?? 0; + message.tokensServed = object.tokensServed ?? "0"; + message.bloomFilter = object.bloomFilter ?? new Uint8Array(0); + return message; + }, +}; + +function createBaseExchangeCapabilitiesRequest(): ExchangeCapabilitiesRequest { + return { knownNodes: [] }; +} + +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(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): ExchangeCapabilitiesResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseExchangeCapabilitiesResponse(); + 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): ExchangeCapabilitiesResponse { + 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: ExchangeCapabilitiesResponse): unknown { + const obj: any = {}; + if (message.knownNodes?.length) { + obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e)); + } + return obj; + }, + + create, I>>(base?: I): ExchangeCapabilitiesResponse { + return ExchangeCapabilitiesResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): ExchangeCapabilitiesResponse { + const message = createBaseExchangeCapabilitiesResponse(); + message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || []; + return message; + }, +}; + +function createBaseGetNodeCapabilityRequest(): GetNodeCapabilityRequest { + return {}; +} + +export const GetNodeCapabilityRequest: MessageFns = { + encode(_: GetNodeCapabilityRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): GetNodeCapabilityRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseGetNodeCapabilityRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(_: any): GetNodeCapabilityRequest { + return {}; + }, + + toJSON(_: GetNodeCapabilityRequest): unknown { + const obj: any = {}; + return obj; + }, + + create, I>>(base?: I): GetNodeCapabilityRequest { + return GetNodeCapabilityRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(_: I): GetNodeCapabilityRequest { + const message = createBaseGetNodeCapabilityRequest(); + return message; + }, +}; + +function createBaseGetNodeCapabilityResponse(): GetNodeCapabilityResponse { + return { node: undefined }; +} + +export const GetNodeCapabilityResponse: MessageFns = { + encode(message: GetNodeCapabilityResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.node !== undefined) { + NodeCapability.encode(message.node, writer.uint32(10).fork()).join(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): GetNodeCapabilityResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseGetNodeCapabilityResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.node = NodeCapability.decode(reader, reader.uint32()); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): GetNodeCapabilityResponse { + return { node: isSet(object.node) ? NodeCapability.fromJSON(object.node) : undefined }; + }, + + toJSON(message: GetNodeCapabilityResponse): unknown { + const obj: any = {}; + if (message.node !== undefined) { + obj.node = NodeCapability.toJSON(message.node); + } + return obj; + }, + + create, I>>(base?: I): GetNodeCapabilityResponse { + return GetNodeCapabilityResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): GetNodeCapabilityResponse { + const message = createBaseGetNodeCapabilityResponse(); + message.node = (object.node !== undefined && object.node !== null) + ? NodeCapability.fromPartial(object.node) + : undefined; + return message; + }, +}; + +function createBaseGetCacheSummaryRequest(): GetCacheSummaryRequest { + return { compatibility: undefined }; +} + +export const GetCacheSummaryRequest: MessageFns = { + encode(message: GetCacheSummaryRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(10).fork()).join(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): GetCacheSummaryRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseGetCacheSummaryRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + 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): GetCacheSummaryRequest { + return { + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + }; + }, + + toJSON(message: GetCacheSummaryRequest): unknown { + const obj: any = {}; + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + return obj; + }, + + create, I>>(base?: I): GetCacheSummaryRequest { + return GetCacheSummaryRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): GetCacheSummaryRequest { + const message = createBaseGetCacheSummaryRequest(); + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + return message; + }, +}; + +function createBaseGetCacheSummaryResponse(): GetCacheSummaryResponse { + return { nodeId: "", caches: [] }; +} + +export const GetCacheSummaryResponse: MessageFns = { + encode(message: GetCacheSummaryResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.nodeId !== "") { + writer.uint32(10).string(message.nodeId); + } + for (const v of message.caches) { + CacheCapability.encode(v!, writer.uint32(18).fork()).join(); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): GetCacheSummaryResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseGetCacheSummaryResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.nodeId = reader.string(); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.caches.push(CacheCapability.decode(reader, reader.uint32())); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): GetCacheSummaryResponse { + return { + nodeId: isSet(object.nodeId) + ? globalThis.String(object.nodeId) + : isSet(object.node_id) + ? globalThis.String(object.node_id) + : "", + caches: globalThis.Array.isArray(object?.caches) + ? object.caches.map((e: any) => CacheCapability.fromJSON(e)) + : [], + }; + }, + + toJSON(message: GetCacheSummaryResponse): unknown { + const obj: any = {}; + if (message.nodeId !== "") { + obj.nodeId = message.nodeId; + } + if (message.caches?.length) { + obj.caches = message.caches.map((e) => CacheCapability.toJSON(e)); + } + return obj; + }, + + create, I>>(base?: I): GetCacheSummaryResponse { + return GetCacheSummaryResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): GetCacheSummaryResponse { + const message = createBaseGetCacheSummaryResponse(); + message.nodeId = object.nodeId ?? ""; + message.caches = object.caches?.map((e) => CacheCapability.fromPartial(e)) || []; + return message; + }, +}; + +function createBaseLookupPrefixRequest(): LookupPrefixRequest { + return { compatibility: undefined, blockHashes: [] }; +} + +export const LookupPrefixRequest: MessageFns = { + encode(message: LookupPrefixRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.compatibility !== undefined) { + CacheCompatibility.encode(message.compatibility, writer.uint32(10).fork()).join(); + } + for (const v of message.blockHashes) { + writer.uint32(18).bytes(v!); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): LookupPrefixRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseLookupPrefixRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.compatibility = CacheCompatibility.decode(reader, reader.uint32()); + continue; + } + case 2: { + if (tag !== 18) { + break; + } + + message.blockHashes.push(reader.bytes()); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): LookupPrefixRequest { + return { + compatibility: isSet(object.compatibility) ? CacheCompatibility.fromJSON(object.compatibility) : undefined, + 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)) + : [], + }; + }, + + toJSON(message: LookupPrefixRequest): unknown { + const obj: any = {}; + if (message.compatibility !== undefined) { + obj.compatibility = CacheCompatibility.toJSON(message.compatibility); + } + if (message.blockHashes?.length) { + obj.blockHashes = message.blockHashes.map((e) => base64FromBytes(e)); + } + return obj; + }, + + create, I>>(base?: I): LookupPrefixRequest { + return LookupPrefixRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): LookupPrefixRequest { + const message = createBaseLookupPrefixRequest(); + message.compatibility = (object.compatibility !== undefined && object.compatibility !== null) + ? CacheCompatibility.fromPartial(object.compatibility) + : undefined; + message.blockHashes = object.blockHashes?.map((e) => e) || []; + return message; + }, +}; + +function createBaseLookupPrefixResponse(): LookupPrefixResponse { + return { + nodeId: "", + hitBlockCount: 0, + hitTokenCount: "0", + transferBytes: "0", + cacheEpoch: "0", + leaseId: "", + leaseExpiresAtUnix: 0, + payloadSha256: new Uint8Array(0), + }; +} + +export const LookupPrefixResponse: MessageFns = { + encode(message: LookupPrefixResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.nodeId !== "") { + writer.uint32(10).string(message.nodeId); + } + if (message.hitBlockCount !== 0) { + writer.uint32(16).uint32(message.hitBlockCount); + } + if (message.hitTokenCount !== "0") { + writer.uint32(24).uint64(message.hitTokenCount); + } + if (message.transferBytes !== "0") { + writer.uint32(32).uint64(message.transferBytes); + } + if (message.cacheEpoch !== "0") { + writer.uint32(40).uint64(message.cacheEpoch); + } + if (message.leaseId !== "") { + writer.uint32(50).string(message.leaseId); + } + if (message.leaseExpiresAtUnix !== 0) { + writer.uint32(57).double(message.leaseExpiresAtUnix); + } + if (message.payloadSha256.length !== 0) { + writer.uint32(66).bytes(message.payloadSha256); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): LookupPrefixResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseLookupPrefixResponse(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.nodeId = reader.string(); + continue; + } + case 2: { + if (tag !== 16) { + break; + } + + message.hitBlockCount = reader.uint32(); + continue; + } + case 3: { + if (tag !== 24) { + break; + } + + message.hitTokenCount = reader.uint64().toString(); + continue; + } + case 4: { + if (tag !== 32) { + break; + } + + message.transferBytes = reader.uint64().toString(); + continue; + } + case 5: { + if (tag !== 40) { + break; + } + + message.cacheEpoch = reader.uint64().toString(); + continue; + } + case 6: { + if (tag !== 50) { + break; + } + + message.leaseId = reader.string(); + continue; + } + case 7: { + if (tag !== 57) { + break; + } + + message.leaseExpiresAtUnix = reader.double(); + continue; + } + case 8: { + if (tag !== 66) { + break; + } + + message.payloadSha256 = reader.bytes(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): LookupPrefixResponse { + return { + nodeId: isSet(object.nodeId) + ? globalThis.String(object.nodeId) + : isSet(object.node_id) + ? globalThis.String(object.node_id) + : "", + hitBlockCount: isSet(object.hitBlockCount) + ? globalThis.Number(object.hitBlockCount) + : isSet(object.hit_block_count) + ? globalThis.Number(object.hit_block_count) + : 0, + hitTokenCount: isSet(object.hitTokenCount) + ? globalThis.String(object.hitTokenCount) + : isSet(object.hit_token_count) + ? globalThis.String(object.hit_token_count) + : "0", + transferBytes: isSet(object.transferBytes) + ? globalThis.String(object.transferBytes) + : isSet(object.transfer_bytes) + ? globalThis.String(object.transfer_bytes) + : "0", + cacheEpoch: isSet(object.cacheEpoch) + ? globalThis.String(object.cacheEpoch) + : isSet(object.cache_epoch) + ? globalThis.String(object.cache_epoch) + : "0", + leaseId: isSet(object.leaseId) + ? globalThis.String(object.leaseId) + : isSet(object.lease_id) + ? globalThis.String(object.lease_id) + : "", + leaseExpiresAtUnix: isSet(object.leaseExpiresAtUnix) + ? globalThis.Number(object.leaseExpiresAtUnix) + : isSet(object.lease_expires_at_unix) + ? globalThis.Number(object.lease_expires_at_unix) + : 0, + payloadSha256: isSet(object.payloadSha256) + ? bytesFromBase64(object.payloadSha256) + : isSet(object.payload_sha256) + ? bytesFromBase64(object.payload_sha256) + : new Uint8Array(0), + }; + }, + + toJSON(message: LookupPrefixResponse): unknown { + const obj: any = {}; + if (message.nodeId !== "") { + obj.nodeId = message.nodeId; + } + if (message.hitBlockCount !== 0) { + obj.hitBlockCount = Math.round(message.hitBlockCount); + } + if (message.hitTokenCount !== "0") { + obj.hitTokenCount = message.hitTokenCount; + } + if (message.transferBytes !== "0") { + obj.transferBytes = message.transferBytes; + } + if (message.cacheEpoch !== "0") { + obj.cacheEpoch = message.cacheEpoch; + } + if (message.leaseId !== "") { + obj.leaseId = message.leaseId; + } + if (message.leaseExpiresAtUnix !== 0) { + obj.leaseExpiresAtUnix = message.leaseExpiresAtUnix; + } + if (message.payloadSha256.length !== 0) { + obj.payloadSha256 = base64FromBytes(message.payloadSha256); + } + return obj; + }, + + create, I>>(base?: I): LookupPrefixResponse { + return LookupPrefixResponse.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): LookupPrefixResponse { + const message = createBaseLookupPrefixResponse(); + message.nodeId = object.nodeId ?? ""; + message.hitBlockCount = object.hitBlockCount ?? 0; + message.hitTokenCount = object.hitTokenCount ?? "0"; + message.transferBytes = object.transferBytes ?? "0"; + message.cacheEpoch = object.cacheEpoch ?? "0"; + message.leaseId = object.leaseId ?? ""; + message.leaseExpiresAtUnix = object.leaseExpiresAtUnix ?? 0; + message.payloadSha256 = object.payloadSha256 ?? new Uint8Array(0); + return message; + }, +}; + +function createBaseFetchBlocksRequest(): FetchBlocksRequest { + return { leaseId: "" }; +} + +export const FetchBlocksRequest: MessageFns = { + encode(message: FetchBlocksRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.leaseId !== "") { + writer.uint32(10).string(message.leaseId); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): FetchBlocksRequest { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseFetchBlocksRequest(); + while (reader.pos < end) { + const tag = reader.uint32(); + switch (tag >>> 3) { + case 1: { + if (tag !== 10) { + break; + } + + message.leaseId = reader.string(); + continue; + } + } + if ((tag & 7) === 4 || tag === 0) { + break; + } + reader.skip(tag & 7); + } + return message; + }, + + fromJSON(object: any): FetchBlocksRequest { + return { + leaseId: isSet(object.leaseId) + ? globalThis.String(object.leaseId) + : isSet(object.lease_id) + ? globalThis.String(object.lease_id) + : "", + }; + }, + + toJSON(message: FetchBlocksRequest): unknown { + const obj: any = {}; + if (message.leaseId !== "") { + obj.leaseId = message.leaseId; + } + return obj; + }, + + create, I>>(base?: I): FetchBlocksRequest { + return FetchBlocksRequest.fromPartial(base ?? ({} as any)); + }, + fromPartial, I>>(object: I): FetchBlocksRequest { + const message = createBaseFetchBlocksRequest(); + message.leaseId = object.leaseId ?? ""; + return message; + }, +}; + +function createBaseFetchBlocksResponse(): FetchBlocksResponse { + return { + blockHash: new Uint8Array(0), + blockIndex: 0, + tokenCount: 0, + chunkIndex: 0, + totalChunks: 0, + data: new Uint8Array(0), + blockSha256: new Uint8Array(0), + cacheEpoch: "0", + }; +} + +export const FetchBlocksResponse: MessageFns = { + encode(message: FetchBlocksResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.blockHash.length !== 0) { + writer.uint32(10).bytes(message.blockHash); + } + if (message.blockIndex !== 0) { + writer.uint32(16).uint32(message.blockIndex); + } + if (message.tokenCount !== 0) { + writer.uint32(24).uint32(message.tokenCount); + } + if (message.chunkIndex !== 0) { + writer.uint32(32).uint32(message.chunkIndex); + } + if (message.totalChunks !== 0) { + writer.uint32(40).uint32(message.totalChunks); + } + 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); + } + return writer; + }, + + decode(input: BinaryReader | Uint8Array, length?: number): FetchBlocksResponse { + const reader = input instanceof BinaryReader ? input : new BinaryReader(input); + const end = length === undefined ? reader.len : reader.pos + length; + const message = createBaseFetchBlocksResponse(); + 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; } } @@ -676,50 +2221,143 @@ export const ExchangeCapabilitiesRequest: MessageFns NodeCapability.fromJSON(e)) - : globalThis.Array.isArray(object?.known_nodes) - ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e)) - : [], + 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", }; }, - toJSON(message: ExchangeCapabilitiesRequest): unknown { + toJSON(message: FetchBlocksResponse): unknown { const obj: any = {}; - if (message.knownNodes?.length) { - obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e)); + 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; } return obj; }, - create, I>>(base?: I): ExchangeCapabilitiesRequest { - return ExchangeCapabilitiesRequest.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): FetchBlocksResponse { + return FetchBlocksResponse.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): FetchBlocksResponse { + const message = createBaseFetchBlocksResponse(); + 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"; return message; }, }; -function createBaseExchangeCapabilitiesResponse(): ExchangeCapabilitiesResponse { - return { knownNodes: [] }; +function createBasePublishBlockRequest(): PublishBlockRequest { + return { + blockHash: new Uint8Array(0), + blockIndex: 0, + tokenCount: 0, + chunkIndex: 0, + totalChunks: 0, + data: new Uint8Array(0), + blockSha256: new Uint8Array(0), + cacheEpoch: "0", + compatibility: undefined, + }; } -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(); +export const PublishBlockRequest: MessageFns = { + encode(message: PublishBlockRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { + if (message.blockHash.length !== 0) { + writer.uint32(10).bytes(message.blockHash); + } + if (message.blockIndex !== 0) { + writer.uint32(16).uint32(message.blockIndex); + } + if (message.tokenCount !== 0) { + writer.uint32(24).uint32(message.tokenCount); + } + if (message.chunkIndex !== 0) { + writer.uint32(32).uint32(message.chunkIndex); + } + if (message.totalChunks !== 0) { + writer.uint32(40).uint32(message.totalChunks); + } + 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): ExchangeCapabilitiesResponse { + 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 = createBaseExchangeCapabilitiesResponse(); + const message = createBasePublishBlockRequest(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { @@ -728,62 +2366,73 @@ export const ExchangeCapabilitiesResponse: MessageFns NodeCapability.fromJSON(e)) - : globalThis.Array.isArray(object?.known_nodes) - ? object.known_nodes.map((e: any) => NodeCapability.fromJSON(e)) - : [], - }; - }, + message.blockIndex = reader.uint32(); + continue; + } + case 3: { + if (tag !== 24) { + break; + } - toJSON(message: ExchangeCapabilitiesResponse): unknown { - const obj: any = {}; - if (message.knownNodes?.length) { - obj.knownNodes = message.knownNodes.map((e) => NodeCapability.toJSON(e)); - } - return obj; - }, + message.tokenCount = reader.uint32(); + continue; + } + case 4: { + if (tag !== 32) { + break; + } - create, I>>(base?: I): ExchangeCapabilitiesResponse { - return ExchangeCapabilitiesResponse.fromPartial(base ?? ({} as any)); - }, - fromPartial, I>>(object: I): ExchangeCapabilitiesResponse { - const message = createBaseExchangeCapabilitiesResponse(); - message.knownNodes = object.knownNodes?.map((e) => NodeCapability.fromPartial(e)) || []; - return message; - }, -}; + message.chunkIndex = reader.uint32(); + continue; + } + case 5: { + if (tag !== 40) { + break; + } -function createBaseGetNodeCapabilityRequest(): GetNodeCapabilityRequest { - return {}; -} + message.totalChunks = reader.uint32(); + continue; + } + case 6: { + if (tag !== 50) { + break; + } -export const GetNodeCapabilityRequest: MessageFns = { - encode(_: GetNodeCapabilityRequest, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { - return writer; - }, + message.data = reader.bytes(); + continue; + } + case 7: { + if (tag !== 58) { + break; + } - decode(input: BinaryReader | Uint8Array, length?: number): GetNodeCapabilityRequest { - const reader = input instanceof BinaryReader ? input : new BinaryReader(input); - const end = length === undefined ? reader.len : reader.pos + length; - const message = createBaseGetNodeCapabilityRequest(); - while (reader.pos < end) { - const tag = reader.uint32(); - switch (tag >>> 3) { + 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; @@ -793,49 +2442,136 @@ export const GetNodeCapabilityRequest: MessageFns = { return message; }, - fromJSON(_: any): GetNodeCapabilityRequest { - return {}; + 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(_: GetNodeCapabilityRequest): unknown { + 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): GetNodeCapabilityRequest { - return GetNodeCapabilityRequest.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): PublishBlockRequest { + return PublishBlockRequest.fromPartial(base ?? ({} as any)); }, - fromPartial, I>>(_: I): GetNodeCapabilityRequest { - const message = createBaseGetNodeCapabilityRequest(); + 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 createBaseGetNodeCapabilityResponse(): GetNodeCapabilityResponse { - return { node: undefined }; +function createBasePublishBlockResponse(): PublishBlockResponse { + return { stored: false, cacheEpoch: "0" }; } -export const GetNodeCapabilityResponse: MessageFns = { - encode(message: GetNodeCapabilityResponse, writer: BinaryWriter = new BinaryWriter()): BinaryWriter { - if (message.node !== undefined) { - NodeCapability.encode(message.node, writer.uint32(10).fork()).join(); +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): GetNodeCapabilityResponse { + 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 = createBaseGetNodeCapabilityResponse(); + const message = createBasePublishBlockResponse(); while (reader.pos < end) { const tag = reader.uint32(); switch (tag >>> 3) { case 1: { - if (tag !== 10) { + if (tag !== 8) { break; } - message.node = NodeCapability.decode(reader, reader.uint32()); + message.stored = reader.bool(); + continue; + } + case 2: { + if (tag !== 16) { + break; + } + + message.cacheEpoch = reader.uint64().toString(); continue; } } @@ -847,26 +2583,35 @@ export const GetNodeCapabilityResponse: MessageFns = return message; }, - fromJSON(object: any): GetNodeCapabilityResponse { - return { node: isSet(object.node) ? NodeCapability.fromJSON(object.node) : undefined }; + 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: GetNodeCapabilityResponse): unknown { + toJSON(message: PublishBlockResponse): unknown { const obj: any = {}; - if (message.node !== undefined) { - obj.node = NodeCapability.toJSON(message.node); + if (message.stored !== false) { + obj.stored = message.stored; + } + if (message.cacheEpoch !== "0") { + obj.cacheEpoch = message.cacheEpoch; } return obj; }, - create, I>>(base?: I): GetNodeCapabilityResponse { - return GetNodeCapabilityResponse.fromPartial(base ?? ({} as any)); + create, I>>(base?: I): PublishBlockResponse { + return PublishBlockResponse.fromPartial(base ?? ({} as any)); }, - fromPartial, I>>(object: I): GetNodeCapabilityResponse { - const message = createBaseGetNodeCapabilityResponse(); - message.node = (object.node !== undefined && object.node !== null) - ? NodeCapability.fromPartial(object.node) - : undefined; + fromPartial, I>>(object: I): PublishBlockResponse { + const message = createBasePublishBlockResponse(); + message.stored = object.stored ?? false; + message.cacheEpoch = object.cacheEpoch ?? "0"; return message; }, }; @@ -2507,6 +4252,127 @@ export const ProposerServiceClient = makeGenericClientConstructor( serviceName: string; }; +/** + * PrefillCacheService exposes immutable, content-addressed prefill K/V blocks. + * Lookup is metadata-only; FetchBlocks is the point-to-point bulk data plane. + * Decode never calls this service: a requester imports a hit once, computes the + * missing suffix locally, and keeps the autoregressive loop local. + */ +export type PrefillCacheServiceService = typeof PrefillCacheServiceService; +export const PrefillCacheServiceService = { + getCacheSummary: { + path: "/kakeya.v1.PrefillCacheService/GetCacheSummary" as const, + requestStream: false as const, + responseStream: false as const, + requestSerialize: (value: GetCacheSummaryRequest): Buffer => + Buffer.from(GetCacheSummaryRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): GetCacheSummaryRequest => GetCacheSummaryRequest.decode(value), + responseSerialize: (value: GetCacheSummaryResponse): Buffer => + Buffer.from(GetCacheSummaryResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): GetCacheSummaryResponse => GetCacheSummaryResponse.decode(value), + }, + lookupPrefix: { + path: "/kakeya.v1.PrefillCacheService/LookupPrefix" as const, + requestStream: false as const, + responseStream: false as const, + requestSerialize: (value: LookupPrefixRequest): Buffer => Buffer.from(LookupPrefixRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): LookupPrefixRequest => LookupPrefixRequest.decode(value), + responseSerialize: (value: LookupPrefixResponse): Buffer => + Buffer.from(LookupPrefixResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): LookupPrefixResponse => LookupPrefixResponse.decode(value), + }, + fetchBlocks: { + path: "/kakeya.v1.PrefillCacheService/FetchBlocks" as const, + requestStream: false as const, + responseStream: true as const, + requestSerialize: (value: FetchBlocksRequest): Buffer => Buffer.from(FetchBlocksRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): FetchBlocksRequest => FetchBlocksRequest.decode(value), + responseSerialize: (value: FetchBlocksResponse): Buffer => Buffer.from(FetchBlocksResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): FetchBlocksResponse => FetchBlocksResponse.decode(value), + }, + publishBlock: { + path: "/kakeya.v1.PrefillCacheService/PublishBlock" as const, + requestStream: true as const, + responseStream: false as const, + requestSerialize: (value: PublishBlockRequest): Buffer => Buffer.from(PublishBlockRequest.encode(value).finish()), + requestDeserialize: (value: Buffer): PublishBlockRequest => PublishBlockRequest.decode(value), + responseSerialize: (value: PublishBlockResponse): Buffer => + Buffer.from(PublishBlockResponse.encode(value).finish()), + responseDeserialize: (value: Buffer): PublishBlockResponse => PublishBlockResponse.decode(value), + }, +} as const; + +export interface PrefillCacheServiceServer extends UntypedServiceImplementation { + getCacheSummary: handleUnaryCall; + lookupPrefix: handleUnaryCall; + fetchBlocks: handleServerStreamingCall; + publishBlock: handleClientStreamingCall; +} + +export interface PrefillCacheServiceClient extends Client { + getCacheSummary( + request: GetCacheSummaryRequest, + callback: (error: ServiceError | null, response: GetCacheSummaryResponse) => void, + ): ClientUnaryCall; + getCacheSummary( + request: GetCacheSummaryRequest, + metadata: Metadata, + callback: (error: ServiceError | null, response: GetCacheSummaryResponse) => void, + ): ClientUnaryCall; + getCacheSummary( + request: GetCacheSummaryRequest, + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: GetCacheSummaryResponse) => void, + ): ClientUnaryCall; + lookupPrefix( + request: LookupPrefixRequest, + callback: (error: ServiceError | null, response: LookupPrefixResponse) => void, + ): ClientUnaryCall; + lookupPrefix( + request: LookupPrefixRequest, + metadata: Metadata, + callback: (error: ServiceError | null, response: LookupPrefixResponse) => void, + ): ClientUnaryCall; + lookupPrefix( + request: LookupPrefixRequest, + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: LookupPrefixResponse) => void, + ): ClientUnaryCall; + fetchBlocks(request: FetchBlocksRequest, options?: Partial): ClientReadableStream; + fetchBlocks( + request: FetchBlocksRequest, + metadata?: Metadata, + options?: Partial, + ): ClientReadableStream; + publishBlock( + callback: (error: ServiceError | null, response: PublishBlockResponse) => void, + ): ClientWritableStream; + publishBlock( + metadata: Metadata, + callback: (error: ServiceError | null, response: PublishBlockResponse) => void, + ): ClientWritableStream; + publishBlock( + options: Partial, + callback: (error: ServiceError | null, response: PublishBlockResponse) => void, + ): ClientWritableStream; + publishBlock( + metadata: Metadata, + options: Partial, + callback: (error: ServiceError | null, response: PublishBlockResponse) => void, + ): ClientWritableStream; +} + +export const PrefillCacheServiceClient = makeGenericClientConstructor( + PrefillCacheServiceService, + "kakeya.v1.PrefillCacheService", +) as unknown as { + new (address: string, credentials: ChannelCredentials, options?: Partial): PrefillCacheServiceClient; + service: typeof PrefillCacheServiceService; + serviceName: string; +}; + /** * DFlashProposerService: stateful remote DFlash drafter + f_θ restoration. * Per turn: Restore (prompt -> f_θ-projected verifier K/V) then SeedContext diff --git a/tests/backends/mlx/test_prefill_snapshot.py b/tests/backends/mlx/test_prefill_snapshot.py new file mode 100644 index 00000000..558c623b --- /dev/null +++ b/tests/backends/mlx/test_prefill_snapshot.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") +torch = pytest.importorskip("torch") + +from inference_engine.backends.mlx.prefill_snapshot import ( # noqa: E402 + export_mlx_prefill_snapshot, + import_mlx_prefill_snapshot, +) +from inference_engine.distributed.capability import CacheCompatibility # noqa: E402 + + +class Layer: + def __init__(self, value: float = 0.0): + self.keys = mx.full((1, 2, 3, 4), value) + self.values = mx.full((1, 2, 3, 4), value + 1) + self.offset = 3 + + @property + def state(self): + return self.keys, self.values + + @state.setter + def state(self, value): + self.keys, self.values = value + + +def test_snapshot_round_trip_and_compatibility(): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=3) + source = [Layer(1), Layer(2)] + payload = export_mlx_prefill_snapshot( + source, + token_count=3, + cached_token_ids=[1, 2, 3], + compatibility=compatibility, + next_token_logits=torch.tensor([1.0, 2.0]), + ) + target = [Layer(9), Layer(9)] + imported = import_mlx_prefill_snapshot( + payload, + target, + compatibility=compatibility, + ) + assert imported.token_count == 3 + assert imported.cached_token_ids == (1, 2, 3) + assert torch.equal(imported.next_token_logits, torch.tensor([1.0, 2.0])) + assert bool(mx.all(target[0].keys == source[0].keys)) + assert target[0].offset == 3 + with pytest.raises(ValueError, match="compatibility"): + import_mlx_prefill_snapshot( + payload, + target, + compatibility=CacheCompatibility(model_id="other"), + ) + + +def test_snapshot_rejects_empty_layer_and_bad_payload(): + compatibility = CacheCompatibility(model_id="m") + layer = Layer() + layer.keys = None + with pytest.raises(ValueError, match="empty"): + export_mlx_prefill_snapshot( + [layer], + token_count=1, + cached_token_ids=[1], + compatibility=compatibility, + ) + with pytest.raises(ValueError, match="magic"): + import_mlx_prefill_snapshot( + b"bad", + [Layer()], + compatibility=compatibility, + ) diff --git a/tests/backends/mlx/test_verifier.py b/tests/backends/mlx/test_verifier.py index d3f6646a..90b7b893 100644 --- a/tests/backends/mlx/test_verifier.py +++ b/tests/backends/mlx/test_verifier.py @@ -460,26 +460,16 @@ def test_mlx_kv_live_bytes_zero_before_prefill() -> None: def test_mlx_kv_live_bytes_equals_k_seq_length_times_per_token() -> None: - """kv_live_bytes = k_seq_length × per-token bytes, computed from - the wrapped HF config the same way the verifier does.""" + """kv_live_bytes = k_seq_length × resolved per-token bytes. + + The verifier may resolve geometry from an HF config or directly from a + multimodal mlx-lm text-model wrapper (Gemma 4). + """ v = _build_mlx_verifier(sink=2, window=8) v.prefill([10, 20, 30, 40, 50]) k_len = v.k_seq_length(session=None) assert k_len == 5 - cfg = v.model.config if hasattr(v.model, "config") else v.model - num_layers = int(cfg.num_hidden_layers) - num_kv_heads = int( - getattr(cfg, "num_key_value_heads", None) - or cfg.num_attention_heads - ) - head_dim = int( - getattr(cfg, "head_dim", None) - or (cfg.hidden_size // cfg.num_attention_heads) - ) - bytes_per_token = ( - num_layers * num_kv_heads * head_dim - * v.config.dtype.itemsize * 2 - ) + bytes_per_token = v._bytes_per_kv_token expected = k_len * bytes_per_token assert v.kv_live_bytes(session=None) == expected assert expected > 0 diff --git a/tests/inference_engine/distributed/test_capability.py b/tests/inference_engine/distributed/test_capability.py index c75da6b0..c8978ef8 100644 --- a/tests/inference_engine/distributed/test_capability.py +++ b/tests/inference_engine/distributed/test_capability.py @@ -16,8 +16,11 @@ NGRAM_MODEL_ID, CapabilityRegistry, CapabilityRole, + CacheCapability, + CacheCompatibility, ModelCapability, NodeCapability, + NodeEndpoint, ) T0 = 1_000_000.0 @@ -101,6 +104,47 @@ def test_model_capability_proto_round_trip(): assert ModelCapability.from_proto(model.to_proto()) == model +def test_cache_capability_and_endpoints_proto_round_trip(): + compatibility = CacheCompatibility( + model_id="gemma", + model_revision="abc", + tokenizer_revision="tok", + cache_format_version="kv-v1", + quantization="4bit", + rope_hash="rope", + layer_geometry_hash="geometry", + kv_dtype="bfloat16", + block_size_tokens=64, + ) + card = NodeCapability( + node_id="cache-peer", + grpc_address="peer:50051", + caches=( + CacheCapability( + compatibility, + cache_address="169.254.27.104:52051", + cache_bytes_used=10, + cache_bytes_free=20, + entry_count=3, + cache_epoch=4, + load=0.5, + tokens_served=100, + bloom_filter=b"filter", + ), + ), + endpoints=( + NodeEndpoint( + "169.254.27.104:52051", + "thunderbolt", + 100, + 0.45, + ), + ), + ) + assert NodeCapability.from_proto(card.to_proto()) == card + assert CapabilityRole.PREFILL_CACHE.value == 5 + + # --------------------------------------------------------------------------- # CapabilityRegistry merge semantics # --------------------------------------------------------------------------- diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py new file mode 100644 index 00000000..5295080b --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import pytest + +from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.prefill_cache import ( + CacheBlock, + PrefixCacheStore, + chained_block_hashes, + compatibility_fingerprint, + total_payload_bytes, +) + + +def _compat(block_size: int = 2) -> CacheCompatibility: + return CacheCompatibility( + model_id="gemma", + model_revision="weights-1", + tokenizer_revision="tok-1", + cache_format_version="kv-v1", + quantization="4bit", + rope_hash="rope", + layer_geometry_hash="geometry", + kv_dtype="bfloat16", + block_size_tokens=block_size, + ) + + +def test_compatibility_fingerprint_is_stable_and_sensitive(): + a = _compat() + assert compatibility_fingerprint(a) == compatibility_fingerprint(a) + assert compatibility_fingerprint(a) != compatibility_fingerprint(_compat(4)) + + +def test_chained_hashes_require_longest_contiguous_prefix(): + hashes = chained_block_hashes([1, 2, 3, 4, 5], _compat()) + changed = chained_block_hashes([1, 9, 3, 4, 5], _compat()) + assert len(hashes) == 3 + assert hashes[0] != changed[0] + assert hashes[1] != changed[1] + with pytest.raises(ValueError, match="block_size"): + chained_block_hashes([1], _compat(0)) + + +def test_store_returns_longest_snapshot_only(): + store = PrefixCacheStore(_compat(), max_bytes=100, node_id="peer") + hashes = store.put_prefix([1, 2, 3, 4, 5], [b"a", b"bb", b"ccc"]) + lease = store.lookup(hashes + [bytes(32)], now=10.0) + assert lease.hit_block_count == 3 + assert lease.hit_token_count == 5 + assert lease.transfer_bytes == 3 + assert store.fetch(lease.lease_id, now=10.0)[0].payload == b"ccc" + + +def test_store_miss_expiry_collision_and_lru(): + store = PrefixCacheStore(_compat(), max_bytes=4, node_id="peer") + hashes = chained_block_hashes([1, 2, 3, 4], _compat()) + store.put(CacheBlock.create(hashes[0], 2, b"aa")) + assert not store.put(CacheBlock.create(hashes[0], 2, b"aa")) + with pytest.raises(ValueError, match="collision"): + store.put(CacheBlock.create(hashes[0], 2, b"zz")) + store.put(CacheBlock.create(hashes[1], 4, b"bbb")) + assert hashes[0] not in store.block_hashes() + miss = store.lookup([hashes[0]], now=20.0) + assert not miss.lease_id + lease = store.lookup([hashes[1]], lease_seconds=1, now=20.0) + with pytest.raises(KeyError, match="expired"): + store.fetch(lease.lease_id, now=22.0) + + +def test_validation_and_stats(): + with pytest.raises(ValueError, match="max_bytes"): + PrefixCacheStore(_compat(), max_bytes=0, node_id="x") + with pytest.raises(ValueError, match="node_id"): + PrefixCacheStore(_compat(), max_bytes=1, node_id="") + with pytest.raises(ValueError, match="SHA"): + CacheBlock.create(b"x", 1, b"") + with pytest.raises(ValueError, match="token_count"): + CacheBlock.create(bytes(32), 0, b"") + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + with pytest.raises(ValueError, match="capacity"): + store.put(CacheBlock.create(bytes(32), 1, b"x" * 11)) + stats = store.stats() + assert stats.entry_count == 0 + assert stats.max_bytes == 10 + with pytest.raises(ValueError, match="one payload"): + store.put_prefix([1, 2, 3], [b"only-one"]) + with pytest.raises(ValueError, match="lease_seconds"): + store.lookup([], lease_seconds=0) + assert total_payload_bytes([ + CacheBlock.create(bytes(32), 1, b"12"), + CacheBlock.create(bytes.fromhex("01" * 32), 1, b"345"), + ]) == 5 + + +def test_pinned_eviction_and_missing_leased_block_guards(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + block = CacheBlock.create(bytes(32), 1, b"12345") + store.put(block) + lease = store.lookup([block.block_hash], now=1) + store.max_bytes = 1 + store._evict_to_budget() + assert store.block_hashes() == (block.block_hash,) + store._blocks.pop(block.block_hash) + with pytest.raises(KeyError, match="evicted"): + store.fetch(lease.lease_id, now=1) diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime.py b/tests/inference_engine/distributed/test_prefill_cache_runtime.py new file mode 100644 index 00000000..9c84b0ae --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") +torch = pytest.importorskip("torch") + +from inference_engine.distributed.capability import CacheCompatibility # noqa: E402 +from inference_engine.distributed.prefill_cache import PrefixCacheStore # noqa: E402 +from inference_engine.distributed.prefill_cache_runtime import ( # noqa: E402 + DistributedPrefillCacheHook, +) + + +class Layer: + def __init__(self): + self.keys = None + self.values = None + self.offset = 0 + + @property + def state(self): + return self.keys, self.values + + @state.setter + def state(self, value): + self.keys, self.values = value + + +class Verifier: + def __init__(self): + self.cache = None + self.cached_token_sequence = [] + self.next_global_position = 0 + self.next_token_logits = torch.zeros(2) + self.prefill_calls = 0 + self.forwarded = 0 + + def reset(self): + self.cache = [Layer()] + self.cached_token_sequence = [] + self.next_global_position = 0 + + def prefill(self, tokens): + self.reset() + self.prefill_calls += 1 + self._append(tokens) + + def forward_block(self, tokens): + self.forwarded += len(tokens) + self._append(tokens) + return torch.stack([torch.tensor([float(t), 0.0]) for t in tokens]) + + def commit_or_truncate(self, *, forwarded, accepted): + assert forwarded == accepted + + def _append(self, tokens): + values = mx.array(tokens, dtype=mx.float32).reshape(1, 1, -1, 1) + layer = self.cache[0] + layer.keys = values if layer.keys is None else mx.concatenate([layer.keys, values], axis=2) + layer.values = layer.keys + 1 + layer.offset += len(tokens) + self.cached_token_sequence.extend(tokens) + self.next_global_position += len(tokens) + self.next_token_logits = torch.tensor([float(tokens[-1]), 1.0]) + + +def test_local_snapshot_hit_skips_prefill_and_computes_suffix(): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1 << 20, node_id="head") + reused_events = [] + hook = DistributedPrefillCacheHook(store, on_reuse=reused_events.append) + + first = Verifier() + assert hook.prepare(first, [1, 2, 3, 4]) == 0 + assert first.prefill_calls == 1 + assert store.stats().entry_count == 2 + + exact = Verifier() + assert hook.prepare(exact, [1, 2, 3, 4]) == 4 + assert exact.prefill_calls == 0 + assert exact.forwarded == 0 + assert exact.cached_token_sequence == [1, 2, 3, 4] + + suffix = Verifier() + assert hook.prepare(suffix, [1, 2, 3, 4, 5, 6]) == 4 + assert suffix.prefill_calls == 0 + assert suffix.forwarded == 2 + assert suffix.cached_token_sequence == [1, 2, 3, 4, 5, 6] + assert hook.stats.local_hits == 2 + assert hook.stats.tokens_reused == 8 + assert reused_events == [4, 4] diff --git a/tests/inference_engine/distributed/test_prefill_cache_service.py b/tests/inference_engine/distributed/test_prefill_cache_service.py new file mode 100644 index 00000000..e4f4cabd --- /dev/null +++ b/tests/inference_engine/distributed/test_prefill_cache_service.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import asyncio + +import grpc +import pytest + +from inference_engine.distributed.capability import ( + CacheCapability, + CacheCompatibility, + NodeCapability, + NodeEndpoint, +) +from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.distributed.prefill_cache import CacheBlock +from inference_engine.distributed.prefill_cache_service import ( + PrefillCacheServiceServicer, + add_prefill_cache_service, + compatible_cache_peers, + fetch_remote_blocks, + lookup_best_peer, + publish_block_sync, +) +from inference_engine.server.proto_gen.kakeya.v1 import ( # noqa: E402 + distributed_pb2, + distributed_pb2_grpc, +) + + +def _compat() -> CacheCompatibility: + return CacheCompatibility(model_id="m", block_size_tokens=2) + + +@pytest.mark.asyncio +async def test_lookup_and_fetch_over_real_grpc(): + store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer") + hashes = store.put_prefix([1, 2, 3], [b"first", b"latest-snapshot"]) + 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=4) + await server.start() + try: + hit = await lookup_best_peer([address], _compat(), hashes) + assert hit is not None + assert hit.node_id == "peer" + assert hit.hit_block_count == 2 + assert hit.hit_token_count == 3 + chunks = await fetch_remote_blocks(hit) + assert b"".join(chunk.data for chunk in chunks) == b"latest-snapshot" + assert len(chunks) > 1 + async with grpc.aio.insecure_channel(address) as channel: + stub = distributed_pb2_grpc.PrefillCacheServiceStub(channel) + summary = await stub.GetCacheSummary( + distributed_pb2.GetCacheSummaryRequest( + compatibility=_compat().to_proto(), + ), + ) + assert summary.node_id == "peer" + with pytest.raises(grpc.aio.AioRpcError) as error: + async for _ in stub.FetchBlocks( + distributed_pb2.FetchBlocksRequest(lease_id="missing"), + ): + pass + assert error.value.code() == grpc.StatusCode.NOT_FOUND + 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") + hashes = store.put_prefix([1, 2], [b"payload"]) + 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) + await server.start() + try: + wrong = CacheCompatibility(model_id="other", block_size_tokens=2) + assert await lookup_best_peer([], wrong, hashes) is None + assert await lookup_best_peer([address], wrong, hashes) is None + assert await lookup_best_peer(["127.0.0.1:1"], _compat(), hashes) is None + finally: + await server.stop(0) + + +@pytest.mark.asyncio +async def test_publish_block_replication(): + 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) + await server.start() + try: + block = CacheBlock.create(bytes(32), 2, b"snapshot") + stored = await asyncio.to_thread( + publish_block_sync, + address, + _compat(), + block, + ) + assert stored + assert store.stats().entry_count == 1 + assert not await asyncio.to_thread( + publish_block_sync, + address, + _compat(), + block, + ) + finally: + await server.stop(0) + + +@pytest.mark.asyncio +async def test_publish_rejects_malformed_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) + 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)) + + base = dict( + block_hash=bytes(32), + token_count=2, + total_chunks=1, + block_sha256=__import__("hashlib").sha256(b"x").digest(), + compatibility=_compat().to_proto(), + ) + try: + with pytest.raises(grpc.aio.AioRpcError): + await call([]) + with pytest.raises(grpc.aio.AioRpcError): + await call([ + distributed_pb2.PublishBlockRequest(**base, chunk_index=0, data=b"x"), + distributed_pb2.PublishBlockRequest( + **{**base, "block_hash": bytes.fromhex("01" * 32)}, + chunk_index=1, + data=b"x", + ), + ]) + with pytest.raises(grpc.aio.AioRpcError): + await call([ + distributed_pb2.PublishBlockRequest( + **{**base, "block_hash": b"short"}, + chunk_index=0, + data=b"x", + ), + ]) + with pytest.raises(grpc.aio.AioRpcError): + await call([ + distributed_pb2.PublishBlockRequest( + **{**base, "compatibility": CacheCompatibility( + model_id="wrong", + ).to_proto()}, + chunk_index=0, + data=b"x", + ), + ]) + with pytest.raises(grpc.aio.AioRpcError): + await call([ + distributed_pb2.PublishBlockRequest( + **{**base, "total_chunks": 2}, + chunk_index=0, + data=b"x", + ), + ]) + with pytest.raises(grpc.aio.AioRpcError): + await call([ + distributed_pb2.PublishBlockRequest( + **{**base, "block_sha256": bytes(32)}, + chunk_index=0, + data=b"x", + ), + ]) + finally: + await server.stop(0) + + +def test_compatible_peer_selection_prefers_cache_address_and_endpoint(): + compatibility = _compat() + with_cache_address = NodeCapability( + node_id="a", + grpc_address="a:1", + caches=(CacheCapability(compatibility, cache_address="a:2"),), + ) + with_endpoint = NodeCapability( + node_id="b", + grpc_address="b:1", + caches=(CacheCapability(compatibility),), + endpoints=( + NodeEndpoint("b-lan:2", "lan", 50, 2.0), + NodeEndpoint("b-tb:2", "thunderbolt", 100, 0.4), + ), + ) + incompatible = NodeCapability( + node_id="c", + grpc_address="c:1", + caches=(CacheCapability(CacheCompatibility(model_id="x")),), + ) + fallback = NodeCapability( + node_id="d", + grpc_address="d:1", + caches=(CacheCapability(compatibility),), + ) + assert compatible_cache_peers( + [with_cache_address, with_endpoint, incompatible, fallback], + compatibility, + ) == ["a:2", "b-tb:2", "d:1"] + + +def test_service_validation_and_dead_publish(): + store = PrefixCacheStore(_compat(), max_bytes=1024, node_id="peer") + with pytest.raises(ValueError, match="chunk_bytes"): + PrefillCacheServiceServicer( + store, + cache_address="peer:1", + chunk_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) diff --git a/tests/inference_engine/network/test_network_api.py b/tests/inference_engine/network/test_network_api.py new file mode 100644 index 00000000..97957ba6 --- /dev/null +++ b/tests/inference_engine/network/test_network_api.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from fastapi.testclient import TestClient + +from inference_engine.distributed.capability import ( + CacheCompatibility, + CapabilityRegistry, + NodeCapability, +) +from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.network.api import create_network_app +from inference_engine.network.state import NetworkState + + +def _client(tmp_path): + compatibility = CacheCompatibility(model_id="m") + state = NetworkState( + CapabilityRegistry(NodeCapability(node_id="head", grpc_address="head:1")), + PrefixCacheStore(compatibility, max_bytes=100, node_id="head"), + state_path=tmp_path / "state.json", + ) + client = TestClient(create_network_app(state, api_key="secret")) + client.network_state = state + return client + + +def test_dashboard_health_and_read_apis(tmp_path): + client = _client(tmp_path) + assert client.get("/").status_code == 200 + assert "Kakeya Inference Network" in client.get("/network").text + assert client.get("/healthz").json()["status"] == "ok" + assert client.get("/v1/network/summary").json()["online_nodes"] == 1 + assert len(client.get("/v1/network/nodes").json()) == 1 + assert client.get("/v1/network/groups").json() == [] + assert "nodes" in client.get("/v1/network/topology").json() + assert client.get("/v1/network/tokens").json()["completed"] == 0 + events = client.get("/v1/network/events?once=true") + assert events.status_code == 200 + assert "event: summary" in events.text + + +def test_write_apis_require_key_and_update_state(tmp_path): + client = _client(tmp_path) + body = {"alias": "peer", "address": "peer:2", "region": "HK"} + assert client.post("/v1/network/nodes/register", json=body).status_code == 401 + registered = client.post( + "/v1/network/nodes/register", + json=body, + headers={"X-API-Key": "secret"}, + ) + assert registered.status_code == 200 + assert registered.json()["pairing_token"].startswith("kn_pair_") + group = client.post( + "/v1/network/groups", + json={"name": "g", "node_ids": ["head", "peer"]}, + headers={"X-API-Key": "secret"}, + ) + assert group.status_code == 200 + telemetry = client.post( + "/v1/network/telemetry/tokens", + json={"node_id": "head", "completed": 10, "kv_assisted": 7}, + headers={"X-API-Key": "secret"}, + ) + assert telemetry.json()["status"] == "accepted" + assert client.get("/v1/network/tokens").json()["kv_assisted"] == 7 + + +def test_write_error_mapping_and_event_stream(tmp_path, monkeypatch): + client = _client(tmp_path) + monkeypatch.setattr( + client.network_state, + "register_node", + lambda **_: (_ for _ in ()).throw(ValueError("bad registration")), + ) + response = client.post( + "/v1/network/nodes/register", + json={"alias": "a", "address": "a:1"}, + headers={"X-API-Key": "secret"}, + ) + assert response.status_code == 400 + monkeypatch.setattr( + client.network_state, + "create_group", + lambda **_: (_ for _ in ()).throw(ValueError("bad group")), + ) + assert client.post( + "/v1/network/groups", + json={"name": "g", "node_ids": ["a"]}, + headers={"X-API-Key": "secret"}, + ).status_code == 400 + monkeypatch.setattr( + client.network_state, + "record_tokens", + lambda **_: (_ for _ in ()).throw(ValueError("bad tokens")), + ) + assert client.post( + "/v1/network/telemetry/tokens", + json={"node_id": "a", "completed": 1}, + headers={"X-API-Key": "secret"}, + ).status_code == 400 diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py new file mode 100644 index 00000000..c92cc96d --- /dev/null +++ b/tests/inference_engine/network/test_network_state.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from inference_engine.distributed.capability import ( + CacheCapability, + CacheCompatibility, + CapabilityRegistry, + NodeCapability, + NodeEndpoint, +) +from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.network.state import NetworkState + + +def _state(tmp_path): + compatibility = CacheCompatibility(model_id="gemma") + store = PrefixCacheStore(compatibility, max_bytes=1000, node_id="head") + card = NodeCapability( + node_id="head", + grpc_address="head:1", + platform="mac-m4", + unified_memory_bytes=24 << 30, + caches=( + CacheCapability( + compatibility, + cache_address="head:2", + cache_bytes_free=1000, + ), + ), + endpoints=(NodeEndpoint("head:2", "thunderbolt", 100, 0.4),), + ) + return NetworkState( + CapabilityRegistry(self_card=card), + store, + state_path=tmp_path / "network.json", + ) + + +def test_registration_groups_tokens_and_persistence(tmp_path): + state = _state(tmp_path) + registration = state.register_node( + alias="peer", + address="peer:2", + region="Hong Kong", + ) + assert registration["pairing_token"].startswith("kn_pair_") + group = state.create_group(name="Studio", node_ids=["head", "peer"]) + state.record_tokens(node_id="head", completed=100, kv_assisted=70) + summary = state.summary() + assert summary["online_nodes"] == 1 + assert summary["registered_nodes"] == 2 + assert summary["completed_tokens"] == 100 + assert summary["kv_hit_rate"] == 0.7 + assert state.groups()[0]["id"] == group["id"] + assert state.topology()["edges"][0]["target"] == "peer" + + reloaded = _state(tmp_path) + assert reloaded.summary()["completed_tokens"] == 100 + assert reloaded.groups()[0]["name"] == "Studio" + + +def test_state_validates_inputs(tmp_path): + state = _state(tmp_path) + for kwargs in ( + {"alias": "", "address": "x", "region": "r"}, + {"alias": "x", "address": "", "region": "r"}, + ): + try: + state.register_node(**kwargs) + except ValueError: + pass + else: + raise AssertionError("expected registration validation") + try: + state.create_group(name="", node_ids=[]) + except ValueError: + pass + else: + raise AssertionError("expected group validation") + try: + state.record_tokens(node_id="x", completed=1, kv_assisted=2) + except ValueError: + pass + else: + raise AssertionError("expected token validation") + + +def test_invalid_persisted_state_falls_back_to_empty(tmp_path): + path = tmp_path / "network.json" + path.write_text("{not-json") + compatibility = CacheCompatibility(model_id="m") + state = NetworkState( + CapabilityRegistry(NodeCapability(node_id="head", grpc_address="h:1")), + PrefixCacheStore(compatibility, max_bytes=10, node_id="head"), + state_path=path, + ) + assert state.groups() == [] + assert state.summary()["completed_tokens"] == 0 diff --git a/tests/inference_engine/server/test_grpc_app.py b/tests/inference_engine/server/test_grpc_app.py index 81e88b8f..8d141a2b 100644 --- a/tests/inference_engine/server/test_grpc_app.py +++ b/tests/inference_engine/server/test_grpc_app.py @@ -758,6 +758,28 @@ def generate(self, session_id, *, max_tokens, **kw): await server.stop(grace=0.1) +async def test_factory_wires_prefill_cache_service(): + from inference_engine.distributed.capability import CacheCompatibility + from inference_engine.distributed.prefill_cache import PrefixCacheStore + store = SessionStore(capacity=1) + cache = PrefixCacheStore( + CacheCompatibility(model_id="m"), + max_bytes=100, + node_id="cache-node", + ) + server = create_grpc_server( + session_store=store, + config=GrpcServerConfig(bind_address="127.0.0.1:0"), + prefill_cache_store=cache, + prefill_cache_address="cache:1", + ) + await server.start() + # grpc.aio does not expose the selected ephemeral port on the wrapper, + # so this test pins factory wiring by construction/logical coverage; the + # real wire path is covered in test_prefill_cache_service. + await server.stop(grace=0.1) + + # --------------------------------------------------------------------------- # Generate (PR-B3) — verifier-independent paths only #