diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist index 560e9085..0bc8eb92 100644 --- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -17,8 +17,8 @@ --skip-cache-check --enable-prefill-cache --prefill-cache-gb1 - --peer169.254.27.104:52051 - --cache-peer169.254.27.104:52051 + --peer169.254.27.104:53051 + --cache-peer169.254.27.104:53051 --cache-model-idgemma-4-26B-A4B-it-mlx-4bit --model-revisionlocal-4bit-v1 --tokenizer-revisiongemma4-v1 @@ -37,6 +37,9 @@ --measured-rtt-ms0.55 --cache-link-mbps10000 --cache-default-rtt-ms0.55 + --remote-prefill-min-tokens0 + --prefill-worker-timeout-s300 + --prefill-policyremote-required --network-http-host127.0.0.1 --network-http-port8090 --network-api-key__NETWORK_KEY__ diff --git a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist new file mode 100644 index 00000000..cbe2ebb5 --- /dev/null +++ b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist @@ -0,0 +1,43 @@ + + + + Labelai.kakeya.prefill-worker + ProgramArguments + /Users/allen/.venv-distwan/bin/python + /Users/allen/Kakeya-LLM-Inference-engine/scripts/start_prefill_worker_node.py + --node-idallens-mini + --bind169.254.27.104:53051 + --advertise169.254.27.104:53051 + --model-id/Users/allen/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit + --cache-model-idgemma-4-26B-A4B-it-mlx-4bit + --model-revisionlocal-4bit-v1 + --tokenizer-revisiongemma4-v1 + --cache-format-versionkakeya-prefill-v3-kl-d4-q38 + --cache-compressionkakeyalattice-d4 + --quantization4bit-mlx + --layer-geometry-hash93d9585b0f06b60bac8e1cadf50b29df1adbf086c862e61720b6127d22c30e2b + --tenant-idprivate-fleet + --fleet-psk-file/Users/allen/.kakeya/fleet.psk + --sink4 + --window2048 + --block-size-tokens64 + --cache-gb0.25 + --prefill-tps1 + --max-concurrent-jobs1 + --networkthunderbolt + --priority100 + --rtt-ms0.55 + --log-levelINFO + + WorkingDirectory/Users/allen/Kakeya-LLM-Inference-engine + EnvironmentVariables + PATH/Users/allen/.venv-distwan/bin:/usr/bin:/bin:/usr/sbin:/sbin + PYTHONPATH/Users/allen/Kakeya-LLM-Inference-engine:/Users/allen/Kakeya-LLM-Inference-engine/sdks/python + + RunAtLoad + KeepAlive + ProcessTypeInteractive + StandardOutPath/Users/allen/.kakeya/prefill-worker.log + StandardErrorPath/Users/allen/.kakeya/prefill-worker.log + diff --git a/docs/adr/0017-prefill-compute-worker-orchestration.md b/docs/adr/0017-prefill-compute-worker-orchestration.md index 662c24e4..944c3590 100644 --- a/docs/adr/0017-prefill-compute-worker-orchestration.md +++ b/docs/adr/0017-prefill-compute-worker-orchestration.md @@ -49,6 +49,12 @@ Every remote error (lookup, job, lease, fetch, checksum, decompress, import) resets the verifier and falls back to full local prefill. Cache availability must never determine request correctness. +For deployments that require a strictly decode-only primary, +`--prefill-policy remote-required` changes this failure contract: only a +complete cache hit or completed remote worker job is accepted. Partial hits are +not extended on Primary, cost gating is bypassed, and worker failure returns +`UNAVAILABLE` instead of silently running local prefill. + ### Discovery and placement Capability gossip is the only membership source. Static `--cache-peer` and @@ -72,6 +78,9 @@ operator-configured defaults. - Decode KV remains local to the primary. - Peer memory is a pre-decode snapshot tier, not coherent remote attention RAM. +- Cache mounts are exposed as one content-addressed `kv://` namespace for + management. This virtualizes naming and location only; fetch/import still + copies the selected snapshot into Primary memory. - Snapshot payloads support zlib framing and retain SHA-256 of the uncompressed bytes. - Replication uses rendezvous hashing and a bounded replication factor instead diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 15f44407..6e28d42b 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -92,10 +92,9 @@ Use an isolated venv and copy/sync the repository package. The peer plist is: deploy/launchd/ai.kakeya.prefill-network-peer.plist ``` -In the two-Mac profile, `allens-mini` runs this role only. It does not load the -model or apply a chat template. Snapshots arrive from the primary's fallback -prefill today, and from separately deployed compute workers when the fleet has -additional machines. +The cache-only plist remains available for rollback or additional RAM-only +nodes. In the strict two-Mac decode/prefill profile, allens instead runs the +prefill-worker plist below with a co-located cache. Check from the head over Thunderbolt: @@ -114,10 +113,10 @@ The worker loads the exact same MLX model as the primary, accepts queued prefill-only jobs, writes immutable snapshots into its co-located RAM cache and never serves user decode. -This is an additional fleet role, not the `allens-mini` cache-only role in the -two-Mac profile. Deploy it on Worker A/B/C addresses when those machines exist. -Workers receive canonical token IDs from the scheduler; they do not construct -their own chat template. +In the strict two-Mac profile, `allens-mini` runs this role and Primary uses +`--prefill-policy remote-required`. Workers receive canonical token IDs from +the scheduler; they do not construct their own chat template. The worker stores +the resulting snapshots in its co-located content-addressed cache. Create a fleet PSK once and copy it to every trusted node: @@ -233,6 +232,16 @@ separate Worker A/B/C path; that mode additionally requires `remote_jobs`. Decode throughput is reported separately because all autoregressive decode remains on the primary. +The logical cross-node KV namespace is available at: + +```bash +curl -fsS http://127.0.0.1:8090/v1/network/kvfs +``` + +The returned `kv://` URI and mount table virtualize naming and management only. +Payloads remain in each Mac's physical RAM and are copied into Primary once +before decode; `coherent_shared_memory` is always false. + ## Maintenance cache saturation Enable the bounded, memory-only first-append capture queue on the primary: diff --git a/inference_engine/distributed/kv_namespace.py b/inference_engine/distributed/kv_namespace.py new file mode 100644 index 00000000..41e24ce0 --- /dev/null +++ b/inference_engine/distributed/kv_namespace.py @@ -0,0 +1,53 @@ +"""Logical content-addressed namespace over physically separate KV stores.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + +from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.prefill_cache import compatibility_fingerprint + + +@dataclass(frozen=True) +class VirtualKVMount: + node_id: str + address: str + bytes_used: int + bytes_free: int + entry_count: int + network: str + + +class VirtualKVNamespace: + """Present cache-node RAM as one lookup namespace, never as coherent RAM.""" + + def __init__(self, compatibility: CacheCompatibility) -> None: + self.compatibility = compatibility + fingerprint = compatibility_fingerprint(compatibility).hex() + tenant = compatibility.tenant_namespace or "default" + self.uri = f"kv://{tenant}/{compatibility.model_id}/{fingerprint}" + + def describe(self, nodes: Sequence[dict[str, Any]]) -> dict[str, Any]: + mounts = [] + for node in nodes: + cache = node.get("cache") + if not cache or cache.get("model_id") != self.compatibility.model_id: + continue + endpoint = node.get("endpoint") or {} + mounts.append(VirtualKVMount( + node_id=node["id"], + address=endpoint.get("address", ""), + bytes_used=int(cache.get("bytes_used", 0)), + bytes_free=int(cache.get("bytes_free", 0)), + entry_count=int(cache.get("entry_count", 0)), + network=endpoint.get("network", "default"), + )) + return { + "uri": self.uri, + "access": "content-addressed-lookup-fetch-import", + "coherent_shared_memory": False, + "mounts": [mount.__dict__ for mount in mounts], + "bytes_used": sum(mount.bytes_used for mount in mounts), + "bytes_free": sum(mount.bytes_free for mount in mounts), + "entry_count": sum(mount.entry_count for mount in mounts), + } diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py index 79ad9da6..81c47f4e 100644 --- a/inference_engine/distributed/prefill_cache_runtime.py +++ b/inference_engine/distributed/prefill_cache_runtime.py @@ -43,6 +43,7 @@ from inference_engine.distributed.prefill_scheduler import ( PrefillCostConfig, choose_prefill_worker, + compatible_prefill_workers, remote_import_wins, select_cache_replicas, ) @@ -80,6 +81,10 @@ class PrefillReuseStats: last_publish_error: str = "" +class RemotePrefillRequiredError(RuntimeError): + """Raised when decode-only primary policy cannot obtain a complete KV.""" + + @dataclass(frozen=True) class _Hit: source: str @@ -114,6 +119,7 @@ def __init__( cost_config: PrefillCostConfig | None = None, auth: FleetAuthConfig | None = None, on_reuse=None, + require_remote_compute: bool = False, ) -> None: if min( lookup_timeout_s, @@ -143,6 +149,7 @@ def __init__( self.replication_factor = int(replication_factor) self.cost_config = cost_config or PrefillCostConfig() self.auth = auth + self.require_remote_compute = bool(require_remote_compute) self._hash_key = auth.tenant_hash_key() if auth is not None else b"" self.stats = PrefillReuseStats() self._stats_lock = threading.Lock() @@ -167,14 +174,29 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int: ) hit = self._best_hit(hashes) reused = 0 + if ( + self.require_remote_compute + and hit is not None + and hit.hit_tokens != len(tokens) + ): + hit = None if hit is not None: reused = self._try_import(verifier, tokens, hit) - elif len(tokens) >= self.remote_compute_min_tokens: + if reused == 0 and ( + self.require_remote_compute + or len(tokens) >= self.remote_compute_min_tokens + ): remote_hit = self._compute_remote(tokens, hashes) if remote_hit is not None: reused = self._try_import(verifier, tokens, remote_hit) if reused == 0: self.stats.misses += 1 + if self.require_remote_compute and reused != len(tokens): + verifier.reset() + reason = self.stats.last_fallback_reason or ( + "no compatible remote prefill worker completed the request" + ) + raise RemotePrefillRequiredError(reason) self._compute_and_publish(verifier, tokens, hashes, reused) return reused @@ -267,15 +289,28 @@ def _compute_remote( card for card in self._cards() if card.node_id != self.local_store.node_id ) - target = choose_prefill_worker( - cards, - self.compatibility, - prompt_tokens=len(tokens), - estimated_snapshot_bytes=( - len(tokens) * self.estimated_snapshot_bytes_per_token - ), - config=self.cost_config, - ) + if self.require_remote_compute: + candidates = compatible_prefill_workers(cards, self.compatibility) + target = min( + candidates, + key=lambda item: ( + item.capability.load, + item.capability.queued_tokens, + -item.capability.tokens_per_second_prefill, + item.node_id, + ), + default=None, + ) + else: + target = choose_prefill_worker( + cards, + self.compatibility, + prompt_tokens=len(tokens), + estimated_snapshot_bytes=( + len(tokens) * self.estimated_snapshot_bytes_per_token + ), + config=self.cost_config, + ) if target is None: return None request = distributed_pb2.SubmitPrefillJobRequest( diff --git a/inference_engine/network/api.py b/inference_engine/network/api.py index 2e84ef72..78e2173a 100644 --- a/inference_engine/network/api.py +++ b/inference_engine/network/api.py @@ -107,6 +107,10 @@ def create_group(request: CreateGroupRequest): def topology(): return state.topology() + @app.get("/v1/network/kvfs") + def virtual_kv_file(): + return state.virtual_kv_file() + @app.get("/v1/network/tokens") def tokens(): summary = state.summary() diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 0d7482f0..18133437 100644 --- a/inference_engine/network/state.py +++ b/inference_engine/network/state.py @@ -11,6 +11,7 @@ from typing import Any, Callable from inference_engine.distributed.capability import CapabilityRegistry +from inference_engine.distributed.kv_namespace import VirtualKVNamespace from inference_engine.distributed.prefill_cache import PrefixCacheStore @@ -25,6 +26,7 @@ def __init__( ) -> None: self.registry = registry self.cache_store = cache_store + self.kv_namespace = VirtualKVNamespace(cache_store.compatibility) self.state_path = Path(state_path).expanduser() self.prefill_stats_provider = prefill_stats_provider self._lock = threading.RLock() @@ -257,6 +259,9 @@ def topology(self) -> dict[str, Any]: } for target in ids[1:]) return {"nodes": nodes, "edges": edges} + def virtual_kv_file(self) -> dict[str, Any]: + return self.kv_namespace.describe(self.nodes()) + def _load(self) -> dict[str, Any]: if self.state_path.exists(): try: diff --git a/inference_engine/server/grpc_app.py b/inference_engine/server/grpc_app.py index 3a7f45af..24539317 100644 --- a/inference_engine/server/grpc_app.py +++ b/inference_engine/server/grpc_app.py @@ -36,6 +36,9 @@ import grpc from inference_engine.memory.pool import PoolExhausted +from inference_engine.distributed.prefill_cache_runtime import ( + RemotePrefillRequiredError, +) from inference_engine.server.proto_gen.kakeya.v1 import ( runtime_pb2, runtime_pb2_grpc, @@ -210,6 +213,8 @@ async def AppendTokens( # noqa: N802 — gRPC-generated method casing ) except SessionNotFoundError as exc: await context.abort(grpc.StatusCode.NOT_FOUND, str(exc)) + except RemotePrefillRequiredError as exc: + await context.abort(grpc.StatusCode.UNAVAILABLE, str(exc)) except ValueError as exc: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(exc)) except InvariantViolation as exc: diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index dccc740f..732fa428 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -453,6 +453,7 @@ async def _serve(args: argparse.Namespace) -> int: primary_compute_penalty_ms=args.primary_prefill_penalty_ms, ), auth=prefill_auth, + require_remote_compute=(args.prefill_policy == "remote-required"), on_reuse=( (lambda count: telemetry_callback(count, count)) if telemetry_callback is not None else None @@ -767,6 +768,13 @@ def main() -> int: "with prefill; drives work to compute peers.") ap.add_argument("--remote-prefill-min-tokens", type=int, default=128) ap.add_argument("--prefill-worker-timeout-s", type=float, default=120.0) + ap.add_argument( + "--prefill-policy", + choices=["local-fallback", "remote-required"], + default="local-fallback", + help="Use remote workers opportunistically, or require complete remote " + "prefill so the primary remains decode-only.", + ) ap.add_argument("--network-label", default="lan", help="Advertised interface: thunderbolt|lan|tailscale|public.") ap.add_argument("--network-priority", type=int, default=50) diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index b1dce440..4d37b571 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -5,6 +5,7 @@ INSTALLER = ROOT / "deploy" / "install_prefill_worker_launchd.sh" HEAD_PLIST = ROOT / "deploy" / "launchd" / "ai.kakeya.grpc-runtime-prefill.plist" PEER_PLIST = ROOT / "deploy" / "launchd" / "ai.kakeya.prefill-network-peer.plist" +WORKER_PLIST = ROOT / "deploy" / "launchd" / "ai.kakeya.prefill-worker-peer.plist" def test_worker_installer_emits_full_cache_compatibility_contract(): @@ -31,17 +32,20 @@ def test_worker_installer_emits_full_cache_compatibility_contract(): assert 'launchctl kickstart -k "$DOMAIN/$LABEL"' in source -def test_two_mac_deployment_uses_allens_as_cache_only(): +def test_two_mac_deployment_uses_allens_as_prefill_only(): plist = HEAD_PLIST.read_text() assert ( - "--peer169.254.27.104:52051" + "--peer169.254.27.104:53051" in plist ) assert ( - "--cache-peer169.254.27.104:52051" + "--cache-peer169.254.27.104:53051" + in plist + ) + assert ( + "--prefill-policyremote-required" in plist ) - assert "--primary-prefill-penalty-ms" not in plist assert ( "--cache-tenant-idprivate-fleet" in plist @@ -52,8 +56,10 @@ def test_two_mac_deployment_uses_allens_as_cache_only(): "kakeyalattice-d4" in plist ) - peer = PEER_PLIST.read_text() - assert "scripts/start_prefill_cache_node.py" in peer - assert "scripts/start_prefill_worker_node.py" not in peer - assert "--cache-gb8" in peer - assert "--window2048" in peer + worker = WORKER_PLIST.read_text() + assert "scripts/start_prefill_worker_node.py" in worker + assert "scripts/start_prefill_cache_node.py" not in worker + assert "--cache-gb0.25" in worker + assert "--window2048" in worker + assert "--prefill-tps1" in worker + assert "scripts/start_prefill_cache_node.py" in PEER_PLIST.read_text() diff --git a/tests/inference_engine/distributed/test_kv_namespace.py b/tests/inference_engine/distributed/test_kv_namespace.py new file mode 100644 index 00000000..60059b7d --- /dev/null +++ b/tests/inference_engine/distributed/test_kv_namespace.py @@ -0,0 +1,56 @@ +from inference_engine.distributed.capability import CacheCompatibility +from inference_engine.distributed.kv_namespace import VirtualKVNamespace + + +def test_virtual_namespace_aggregates_matching_cache_mounts(): + namespace = VirtualKVNamespace(CacheCompatibility( + model_id="gemma", + tenant_namespace="private", + )) + result = namespace.describe([ + { + "id": "head", + "cache": { + "model_id": "gemma", + "bytes_used": 10, + "bytes_free": 20, + "entry_count": 2, + }, + "endpoint": {"address": "head:1", "network": "thunderbolt"}, + }, + { + "id": "peer", + "cache": { + "model_id": "other", + "bytes_used": 99, + "bytes_free": 1, + "entry_count": 9, + }, + "endpoint": None, + }, + {"id": "worker", "cache": None}, + ]) + assert result["uri"].startswith("kv://private/gemma/") + assert result["coherent_shared_memory"] is False + assert result["bytes_used"] == 10 + assert result["bytes_free"] == 20 + assert result["entry_count"] == 2 + assert result["mounts"] == [{ + "node_id": "head", + "address": "head:1", + "bytes_used": 10, + "bytes_free": 20, + "entry_count": 2, + "network": "thunderbolt", + }] + + +def test_virtual_namespace_defaults_tenant_and_endpoint(): + namespace = VirtualKVNamespace(CacheCompatibility(model_id="m")) + result = namespace.describe([{ + "id": "cache", + "cache": {"model_id": "m"}, + }]) + assert result["uri"].startswith("kv://default/m/") + assert result["mounts"][0]["address"] == "" + assert result["mounts"][0]["network"] == "default" diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py index 6755aea0..bc2a0d39 100644 --- a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py @@ -15,9 +15,11 @@ ) from inference_engine.distributed.prefill_cache_runtime import ( DistributedPrefillCacheHook, + RemotePrefillRequiredError, _Hit, ) from inference_engine.distributed.prefill_scheduler import PrefillCostConfig +from inference_engine.distributed.prefill_worker import PrefillJobState from inference_engine.server.proto_gen.kakeya.v1 import distributed_pb2 @@ -137,6 +139,98 @@ def test_runtime_validation_empty_and_provider_failure(): hook.close() +def test_remote_required_never_falls_back_to_primary_prefill(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") + hook = DistributedPrefillCacheHook( + store, + require_remote_compute=True, + remote_compute_min_tokens=0, + ) + monkeypatch.setattr(hook, "_best_hit", lambda hashes: None) + monkeypatch.setattr(hook, "_compute_remote", lambda tokens, hashes: None) + verifier = _Verifier() + with __import__("pytest").raises(RemotePrefillRequiredError): + hook.prepare(verifier, [1, 2]) + assert verifier.prefill_calls == 0 + assert hook.stats.tokens_computed == 0 + hook.close() + + +def test_remote_required_rejects_partial_cache_hit(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + hook = DistributedPrefillCacheHook( + PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"), + require_remote_compute=True, + ) + monkeypatch.setattr( + hook, + "_best_hit", + lambda hashes: _Hit("peer", "lease", 1, 2, 10), + ) + monkeypatch.setattr( + hook, + "_try_import", + lambda *args: (_ for _ in ()).throw(AssertionError("partial import")), + ) + monkeypatch.setattr(hook, "_compute_remote", lambda tokens, hashes: None) + with __import__("pytest").raises(RemotePrefillRequiredError): + hook.prepare(_Verifier(), [1, 2, 3, 4]) + hook.close() + + +def test_remote_required_forces_compatible_worker_despite_cost(monkeypatch): + compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) + hook = DistributedPrefillCacheHook( + PrefixCacheStore(compatibility, max_bytes=1024, node_id="head"), + require_remote_compute=True, + worker_poll_interval_s=0.001, + ) + capability = type("Capability", (), { + "load": 0.0, + "queued_tokens": 0, + "tokens_per_second_prefill": 0.01, + })() + target = type("Target", (), { + "node_id": "worker", + "address": "worker:1", + "rtt_ms": 1.0, + "capability": capability, + })() + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "compatible_prefill_workers", + lambda cards, compat: [target], + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "choose_prefill_worker", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("cost path")), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "submit_prefill_job_sync", + lambda *args, **kwargs: type("Response", (), {"job_id": "j"})(), + ) + monkeypatch.setattr( + "inference_engine.distributed.prefill_cache_runtime." + "get_prefill_job_sync", + lambda *args, **kwargs: type("Status", (), { + "status": int(PrefillJobState.COMPLETED), + "cache_address": "worker:1", + "lease_id": "lease", + "tokens_computed": 2, + "transfer_bytes": 10, + "payload_sha256": b"s" * 32, + })(), + ) + hit = hook._compute_remote([1, 2], [b"h" * 32]) + assert hit is not None, hook.stats.last_fallback_reason + assert hit.hit_tokens == 2 + assert hook.stats.remote_jobs == 1 + hook.close() + + def test_successful_local_import_suffix_and_on_reuse(monkeypatch): compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") diff --git a/tests/inference_engine/network/test_network_api.py b/tests/inference_engine/network/test_network_api.py index 0d87866a..5faead90 100644 --- a/tests/inference_engine/network/test_network_api.py +++ b/tests/inference_engine/network/test_network_api.py @@ -48,6 +48,7 @@ def test_dashboard_health_and_read_apis(tmp_path): 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/kvfs").json()["uri"].startswith("kv://") assert client.get("/v1/network/tokens").json()["completed"] == 0 assert client.get("/v1/network/prefill").json()["remote_jobs"] == 3 assert client.get("/v1/network/maintenance/capture").status_code == 401 diff --git a/tests/inference_engine/server/test_grpc_app.py b/tests/inference_engine/server/test_grpc_app.py index 8d141a2b..7036aba0 100644 --- a/tests/inference_engine/server/test_grpc_app.py +++ b/tests/inference_engine/server/test_grpc_app.py @@ -378,6 +378,43 @@ def append_tokens(self, session_id, token_ids): await server.stop(grace=0.1) +async def test_append_tokens_remote_required_error_returns_unavailable(): + from inference_engine.distributed.prefill_cache_runtime import ( + RemotePrefillRequiredError, + ) + from inference_engine.session import AppendTokensCoordinator + + class Coordinator(AppendTokensCoordinator): + def append_tokens(self, session_id, token_ids): + raise RemotePrefillRequiredError("remote worker unavailable") + + class Aborted(Exception): + pass + + class Context: + code = None + detail = "" + + async def abort(self, code, detail): + self.code = code + self.detail = detail + raise Aborted + + store = SessionStore(capacity=1) + servicer = RuntimeServiceServicer( + store, + append_coordinator=Coordinator(store, verifier=None), + ) + context = Context() + with pytest.raises(Aborted): + await servicer.AppendTokens( + runtime_pb2.AppendTokensRequest(session_id="s", token_ids=[1]), + context, + ) + assert context.code == grpc.StatusCode.UNAVAILABLE + assert "remote worker unavailable" in context.detail + + async def test_append_tokens_invariant_violation_returns_failed_precondition(): """InvariantViolation raised by the coordinator → FAILED_PRECONDITION on the wire. Verifier is never consulted on this path."""