From 9ab66521d54abe6c1b9e6db9d9e38ae6cf5ccecf Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 13:53:34 +0800 Subject: [PATCH 1/2] fix(prefill): atomically reserve and publish snapshots Reserve worker cache capacity before expensive compute, protect it from adaptive shrink, and atomically lease the final snapshot so completed full-context Prefill cannot be evicted before discovery. Co-authored-by: Cursor --- deploy/install_prefill_worker_launchd.sh | 6 +- .../ai.kakeya.grpc-runtime-prefill.plist | 2 +- .../ai.kakeya.prefill-worker-peer.plist | 5 +- docs/ops/distributed-prefill-kv-network.md | 5 + inference_engine/distributed/prefill_cache.py | 100 +++++++++++++++++- .../distributed/prefill_worker.py | 35 ++++-- scripts/start_prefill_worker_node.py | 8 ++ .../bridge/test_prefill_worker_launchd.py | 10 +- .../distributed/test_prefill_cache.py | 33 ++++++ .../distributed/test_prefill_worker.py | 53 +++++++--- 10 files changed, 229 insertions(+), 28 deletions(-) diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh index 82df7d14..3a3a6810 100755 --- a/deploy/install_prefill_worker_launchd.sh +++ b/deploy/install_prefill_worker_launchd.sh @@ -12,8 +12,9 @@ set -euo pipefail BIND="${KAKEYA_WORKER_BIND:-0.0.0.0:53051}" TENANT="${KAKEYA_TENANT_ID:-default}" CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}" -CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-0.25}" -MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-2}" +CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-1}" +MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-0.5}" +SNAPSHOT_BYTES_PER_TOKEN="${KAKEYA_SNAPSHOT_BYTES_PER_TOKEN:-400000}" ADAPTIVE_CACHE="${KAKEYA_WORKER_ADAPTIVE_CACHE:-0}" PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}" CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}" @@ -76,6 +77,7 @@ cat > "$PLIST" <--cache-gb$CACHE_GB --cache-min-gb$CACHE_MIN_GB --memory-reserve-gb$MEMORY_RESERVE_GB + --estimated-snapshot-bytes-per-token$SNAPSHOT_BYTES_PER_TOKEN $adaptive_xml --sink$SINK --window$WINDOW diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist index 77e0ff78..aa57e623 100644 --- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -9,7 +9,7 @@ /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 + --bind0.0.0.0:51051 --capacity1 --sink4 --window2048 diff --git a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist index 7bbe0dd3..a92dd5f6 100644 --- a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist +++ b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist @@ -23,9 +23,10 @@ --window2048 --block-size-tokens64 --cache-gb8 - --cache-min-gb0.25 + --cache-min-gb1 --adaptive-cache - --memory-reserve-gb2 + --memory-reserve-gb0.5 + --estimated-snapshot-bytes-per-token400000 --prefill-tps1 --max-concurrent-jobs1 --networkthunderbolt diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index ed0feed8..57bea782 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -349,6 +349,11 @@ and semantic fallback are forbidden. A global Critic score is valid only when `critic_omitted_tokens=0`. Long Prefill operations emit a heartbeat every 30 seconds; on the 16GB allens worker, full-context Critic Prefill may take 15–25 minutes. +The worker reserves estimated final-snapshot capacity before model compute, +prevents adaptive shrink from consuming active reservations, then atomically +publishes and leases the final snapshot before adding optional intermediate +boundaries. The 16GB allens deployment uses a 1 GiB cache floor and 0.5 GiB +memory reserve. Capacity failures are rejected before Prefill starts. Interactive prompt templates are deterministic and contain no per-run nonce, so repeating the same task can reuse allens cold-tier and Primary hot-tier KV. diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index fa97f6d3..15b6c075 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/inference_engine/distributed/prefill_cache.py @@ -159,6 +159,7 @@ def __init__( self._evictions = 0 self._bytes_evicted = 0 self._put_failures = 0 + self._reservations: dict[str, int] = {} self._lock = threading.RLock() def put(self, block: CacheBlock) -> bool: @@ -263,6 +264,80 @@ def fetch(self, lease_id: str, *, now: float | None = None) -> tuple[CacheBlock, self._bytes_served += lease.transfer_bytes return tuple(blocks) + def reserve(self, reservation_id: str, byte_count: int) -> None: + """Reserve cache capacity before an expensive Prefill job starts.""" + if not reservation_id or byte_count <= 0: + raise ValueError("reservation id and byte count must be positive") + with self._lock: + if reservation_id in self._reservations: + raise ValueError("duplicate cache reservation") + requested = int(byte_count) + reserved = sum(self._reservations.values()) + requested + if reserved > self.max_bytes: + raise ValueError( + f"snapshot reservation {requested} exceeds available " + f"cache budget {self.max_bytes - sum(self._reservations.values())}", + ) + self._expire_leases(time.time()) + self._evict_to_limit(self.max_bytes - reserved) + if self._bytes_used > self.max_bytes - reserved: + raise ValueError("cache capacity is pinned by active leases") + self._reservations[reservation_id] = requested + + def release_reservation(self, reservation_id: str) -> None: + with self._lock: + self._reservations.pop(reservation_id, None) + + def publish_and_lease( + self, + blocks: Sequence[CacheBlock], + block_hashes: Sequence[bytes], + *, + reservation_id: str, + lease_seconds: float = DEFAULT_LEASE_SECONDS, + ) -> PrefixLease: + """Atomically publish and pin the final computed snapshot.""" + if not blocks or len(blocks) != len(block_hashes): + raise ValueError("one computed snapshot is required per block hash") + if lease_seconds <= 0: + raise ValueError("lease_seconds must be > 0") + final = blocks[-1] + if final.block_hash != bytes(block_hashes[-1]): + raise ValueError("final snapshot hash does not match request") + now = time.time() + with self._lock: + reserved = self._reservations.get(reservation_id) + if reserved is None: + raise ValueError("unknown cache reservation") + if final.nbytes > reserved: + raise ValueError( + f"final snapshot {final.nbytes} exceeds reservation {reserved}", + ) + self._expire_leases(now) + self._evict_to_limit(self.max_bytes - final.nbytes) + if self._bytes_used > self.max_bytes - final.nbytes: + raise ValueError("cache capacity is pinned by active leases") + self._put_locked(final) + lease_id = secrets.token_urlsafe(18) + lease = PrefixLease( + lease_id=lease_id, + block_hashes=(final.block_hash,), + hit_block_count=len(block_hashes), + hit_token_count=final.token_count, + transfer_bytes=final.nbytes, + cache_epoch=self._epoch, + expires_at_unix=now + lease_seconds, + payload_sha256=final.payload_sha256, + ) + self._leases[lease_id] = lease + del self._reservations[reservation_id] + # Preserve longest useful boundaries when spare capacity remains. + for block in reversed(blocks[:-1]): + if block.nbytes + self._bytes_used > self.max_bytes: + continue + self._put_locked(block) + return lease + def stats(self) -> CacheStats: with self._lock: return CacheStats( @@ -285,9 +360,12 @@ def resize(self, max_bytes: int) -> bool: raise ValueError("max_bytes must be > 0") with self._lock: previous = self.max_bytes + reserved = sum(self._reservations.values()) + if int(max_bytes) < reserved: + return False self.max_bytes = int(max_bytes) - self._evict_to_budget() - if self._bytes_used > self.max_bytes: + self._evict_to_limit(self.max_bytes - reserved) + if self._bytes_used > self.max_bytes - reserved: self.max_bytes = max(previous, self._bytes_used) return False return True @@ -322,8 +400,11 @@ def _pinned_hashes(self) -> set[bytes]: } def _evict_to_budget(self) -> None: + self._evict_to_limit(self.max_bytes) + + def _evict_to_limit(self, limit: int) -> None: pinned = self._pinned_hashes() - while self._bytes_used > self.max_bytes and self._blocks: + while self._bytes_used > limit and self._blocks: victim = next((h for h in self._blocks if h not in pinned), None) if victim is None: break @@ -333,6 +414,19 @@ def _evict_to_budget(self) -> None: self._bytes_evicted += block.nbytes self._epoch += 1 + def _put_locked(self, block: CacheBlock) -> bool: + existing = self._blocks.get(block.block_hash) + if existing is not None: + if existing.payload_sha256 != block.payload_sha256: + self._put_failures += 1 + raise ValueError("content-address collision with different payload") + self._blocks.move_to_end(block.block_hash) + return False + self._blocks[block.block_hash] = block + self._bytes_used += block.nbytes + self._epoch += 1 + return True + def total_payload_bytes(blocks: Iterable[CacheBlock]) -> int: return sum(block.nbytes for block in blocks) diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index c30f6032..620298d0 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/inference_engine/distributed/prefill_worker.py @@ -93,12 +93,14 @@ def __init__( max_jobs: int = 128, completed_ttl_s: float = 600.0, max_prompt_tokens: int = 131_072, + estimated_snapshot_bytes_per_token: int = 16, ) -> None: if min( max_concurrent_jobs, max_jobs, completed_ttl_s, max_prompt_tokens, + estimated_snapshot_bytes_per_token, ) <= 0: raise ValueError("worker limits must be > 0") if (engine is None) == (engine_factory is None): @@ -110,6 +112,9 @@ def __init__( self.max_jobs = int(max_jobs) self.completed_ttl_s = float(completed_ttl_s) self.max_prompt_tokens = int(max_prompt_tokens) + self.estimated_snapshot_bytes_per_token = int( + estimated_snapshot_bytes_per_token, + ) self._jobs: dict[str, PrefillJob] = {} self._requests: dict[tuple[str, str], str] = {} self._lock = threading.RLock() @@ -190,6 +195,14 @@ def submit( if deadline_ms > 0 else 0.0 ), ) + estimated_bytes = ( + len(job.token_ids) * self.estimated_snapshot_bytes_per_token + ) + if estimated_bytes > self.cache_store.max_bytes: + raise ValueError( + f"estimated final snapshot {estimated_bytes} exceeds " + f"cache capacity {self.cache_store.max_bytes}", + ) self._jobs[job.job_id] = job self._requests[request_key] = job.job_id job.future = self._executor.submit(self._run, job.job_id) @@ -216,6 +229,7 @@ def cancel(self, job_id: str, tenant_id: str) -> bool: if job.future is not None and job.future.cancel(): job.state = PrefillJobState.CANCELLED job.finished_at = time.time() + self.cache_store.release_reservation(job.job_id) return True def stats(self) -> tuple[int, int, float, int]: @@ -239,6 +253,7 @@ def _run(self, job_id: str) -> None: if job.cancelled.is_set(): job.state = PrefillJobState.CANCELLED job.finished_at = time.time() + self.cache_store.release_reservation(job.job_id) return job.state = PrefillJobState.RUNNING started = time.perf_counter() @@ -252,6 +267,13 @@ def _run(self, job_id: str) -> None: timer.daemon = True timer.start() try: + # MLX workers are single-job. Reserve the full current budget + # before model compute so adaptive resizing and unrelated + # boundaries cannot evict the final snapshot before leasing. + self.cache_store.reserve( + job.job_id, + self.cache_store.max_bytes, + ) blocks = tuple(self._engine_for_current_thread().compute_prefill( job.token_ids, job.block_hashes, @@ -264,13 +286,11 @@ def _run(self, job_id: str) -> None: raise RuntimeError( "prefill engine must return one snapshot per block hash", ) - for block in blocks: - if job.cancelled.is_set(): - raise InterruptedError("prefill job cancelled") - self.cache_store.put(block) - lease = self.cache_store.lookup(job.block_hashes) - if not lease.lease_id: - raise RuntimeError("computed snapshot was not discoverable") + lease = self.cache_store.publish_and_lease( + blocks, + job.block_hashes, + reservation_id=job.job_id, + ) if job.cancelled.is_set(): raise InterruptedError("prefill job cancelled") with self._lock: @@ -292,6 +312,7 @@ def _run(self, job_id: str) -> None: job.state = PrefillJobState.FAILED job.failure_reason = f"{type(exc).__name__}: {exc}" finally: + self.cache_store.release_reservation(job.job_id) if timer is not None: timer.cancel() with self._lock: diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index 37462a8c..5924c301 100644 --- a/scripts/start_prefill_worker_node.py +++ b/scripts/start_prefill_worker_node.py @@ -136,6 +136,9 @@ def engine_factory() -> MLXPrefillComputeEngine: max_jobs=args.max_jobs, completed_ttl_s=args.job_ttl_s, max_prompt_tokens=args.max_prompt_tokens, + estimated_snapshot_bytes_per_token=( + args.estimated_snapshot_bytes_per_token + ), ) jobs.warmup() @@ -290,6 +293,11 @@ def main() -> None: parser.add_argument("--max-concurrent-jobs", type=int, default=1) parser.add_argument("--max-jobs", type=int, default=128) parser.add_argument("--max-prompt-tokens", type=int, default=131072) + parser.add_argument( + "--estimated-snapshot-bytes-per-token", + type=int, + default=400_000, + ) parser.add_argument("--job-ttl-s", type=float, default=600.0) parser.add_argument("--prefill-tps", type=float, default=20.0) parser.add_argument("--max-concurrent-rpcs", type=int, default=32) diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 18479d2e..32fb872f 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -24,6 +24,7 @@ def test_worker_installer_emits_full_cache_compatibility_contract(): "--max-prompt-tokens", "--cache-min-gb", "--memory-reserve-gb", + "--estimated-snapshot-bytes-per-token", ): assert f"{flag}" in source assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source @@ -48,6 +49,7 @@ def test_two_mac_deployment_uses_allens_as_prefill_only(): "--prefill-policyremote-required" in plist ) + assert "--bind0.0.0.0:51051" in plist assert ( "--prefill-worker-timeout-s3600" in plist @@ -66,7 +68,13 @@ def test_two_mac_deployment_uses_allens_as_prefill_only(): assert "scripts/start_prefill_worker_node.py" in worker assert "scripts/start_prefill_cache_node.py" not in worker assert "--cache-gb8" in worker - assert "--cache-min-gb0.25" in worker + assert "--cache-min-gb1" in worker + assert "--memory-reserve-gb0.5" in worker + assert ( + "--estimated-snapshot-bytes-per-token" + "400000" + in worker + ) assert "--adaptive-cache" in worker assert "--window2048" in worker assert "--prefill-tps1" in worker diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index 0895eca3..10499d7f 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -154,3 +154,36 @@ def test_resize_evicts_cold_blocks_and_preserves_pinned_budget(): pass else: raise AssertionError("expected resize validation") + + +def test_reservation_blocks_adaptive_shrink_and_rejects_upfront(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + store.reserve("job", 6) + assert not store.resize(5) + assert store.stats().max_bytes == 10 + with pytest.raises(ValueError, match="available cache budget"): + store.reserve("too-large", 5) + store.release_reservation("job") + assert store.resize(5) + + +def test_publish_and_lease_atomically_pins_final_snapshot(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + old = CacheBlock.create(bytes.fromhex("03" * 32), 1, b"old!") + store.put(old) + hashes = chained_block_hashes([1, 2, 3, 4], _compat()) + blocks = ( + CacheBlock.create(hashes[0], 2, b"mid"), + CacheBlock.create(hashes[1], 4, b"final!"), + ) + store.reserve("job", 6) + lease = store.publish_and_lease( + blocks, + hashes, + reservation_id="job", + ) + assert lease.hit_block_count == 2 + assert lease.hit_token_count == 4 + assert store.fetch(lease.lease_id) == (blocks[-1],) + assert hashes[-1] in store.block_hashes() + assert not store.resize(1) diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 83d880ca..3a86912b 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -234,6 +234,28 @@ def test_job_store_validation_queue_stats_and_gc(): completed_jobs.close() +def test_job_rejects_snapshot_capacity_before_model_compute(): + engine = _Engine() + jobs = PrefillJobStore( + engine, + PrefixCacheStore(COMPAT, max_bytes=1024, node_id="small"), + estimated_snapshot_bytes_per_token=600, + ) + try: + with pytest.raises(ValueError, match="estimated final snapshot"): + jobs.submit( + request_id="too-large", + tenant_id="tenant", + token_ids=[1, 2], + block_hashes=[b"a" * 32], + compatibility=COMPAT, + compression=CompressionCodec.NONE, + ) + assert engine.calls == 0 + finally: + jobs.close() + + def test_factory_engine_is_warmed_and_used_on_same_compute_thread(): cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w") created_on = [] @@ -365,7 +387,7 @@ def compute_prefill(self, token_ids, block_hashes, **kwargs): missing = missing_jobs.submit(request_id="missing", **common) missing.future.result(timeout=1) assert missing.state == PrefillJobState.FAILED - assert "not discoverable" in missing.failure_reason + assert "final snapshot hash does not match request" in missing.failure_reason missing_jobs.close() deadline_engine = _Engine() @@ -396,13 +418,16 @@ def compute_prefill(self, token_ids, block_hashes, *, cancelled, **kwargs): engine = Engine() - class CancelOnPutStore(PrefixCacheStore): - def put(self, block): - result = super().put(block) + class CancelBeforeAtomicPublishStore(PrefixCacheStore): + def publish_and_lease(self, *args, **kwargs): engine.event.set() - return result + return super().publish_and_lease(*args, **kwargs) - store = CancelOnPutStore(COMPAT, max_bytes=4096, node_id="put") + store = CancelBeforeAtomicPublishStore( + COMPAT, + max_bytes=4096, + node_id="put", + ) jobs = PrefillJobStore(engine, store) job = jobs.submit( request_id="put", @@ -418,13 +443,17 @@ def put(self, block): engine2 = Engine() - class CancelOnLookupStore(PrefixCacheStore): - def lookup(self, *args, **kwargs): - lease = super().lookup(*args, **kwargs) + class CancelAfterAtomicPublishStore(PrefixCacheStore): + def publish_and_lease(self, *args, **kwargs): + lease = super().publish_and_lease(*args, **kwargs) engine2.event.set() return lease - store2 = CancelOnLookupStore(COMPAT, max_bytes=4096, node_id="lookup") + store2 = CancelAfterAtomicPublishStore( + COMPAT, + max_bytes=4096, + node_id="lookup", + ) jobs2 = PrefillJobStore(engine2, store2) job2 = jobs2.submit( request_id="lookup", @@ -444,10 +473,10 @@ def __init__(self): def is_set(self): self.calls += 1 - return self.calls >= 5 + return self.calls >= 4 def set(self): - self.calls = 5 + self.calls = 4 final_store = PrefixCacheStore(COMPAT, max_bytes=4096, node_id="final") final_jobs = PrefillJobStore(_Engine(), final_store) From 25ee62553225adfa4248d609f3bab25c6bad61e8 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 13:59:22 +0800 Subject: [PATCH 2/2] test(prefill): cover reservation failure boundaries Exercise every atomic publish and capacity guard while making public cache writes respect active reservations. Co-authored-by: Cursor --- inference_engine/distributed/prefill_cache.py | 10 +-- .../distributed/test_prefill_cache.py | 64 +++++++++++++++++++ 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index 15b6c075..0b8029d3 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/inference_engine/distributed/prefill_cache.py @@ -164,11 +164,11 @@ def __init__( def put(self, block: CacheBlock) -> bool: """Publish one immutable block. Returns False for an identical hit.""" - if block.nbytes > self.max_bytes: - with self._lock: - self._put_failures += 1 - raise ValueError("block payload exceeds cache capacity") with self._lock: + available = self.max_bytes - sum(self._reservations.values()) + if block.nbytes > available: + self._put_failures += 1 + raise ValueError("block payload exceeds cache capacity") self._expire_leases(time.time()) existing = self._blocks.get(block.block_hash) if existing is not None: @@ -180,7 +180,7 @@ def put(self, block: CacheBlock) -> bool: self._blocks[block.block_hash] = block self._bytes_used += block.nbytes self._epoch += 1 - self._evict_to_budget() + self._evict_to_limit(available) if block.block_hash not in self._blocks: self._put_failures += 1 raise ValueError( diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index 10499d7f..c5ad2d62 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -187,3 +187,67 @@ def test_publish_and_lease_atomically_pins_final_snapshot(): assert store.fetch(lease.lease_id) == (blocks[-1],) assert hashes[-1] in store.block_hashes() assert not store.resize(1) + + +def test_reservation_and_atomic_publish_validation_guards(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + with pytest.raises(ValueError, match="reservation id"): + store.reserve("", 1) + store.reserve("job", 2) + with pytest.raises(ValueError, match="duplicate"): + store.reserve("job", 1) + block = CacheBlock.create(bytes(32), 1, b"abc") + with pytest.raises(ValueError, match="one computed snapshot"): + store.publish_and_lease([], [], reservation_id="job") + with pytest.raises(ValueError, match="lease_seconds"): + store.publish_and_lease( + [block], + [block.block_hash], + reservation_id="job", + lease_seconds=0, + ) + with pytest.raises(ValueError, match="unknown cache reservation"): + store.publish_and_lease( + [block], + [block.block_hash], + reservation_id="missing", + ) + with pytest.raises(ValueError, match="exceeds reservation"): + store.publish_and_lease( + [block], + [block.block_hash], + reservation_id="job", + ) + store.release_reservation("job") + + +def test_reservation_and_publish_reject_pinned_capacity(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + pinned = CacheBlock.create(bytes.fromhex("04" * 32), 1, b"12345") + store.put(pinned) + store.lookup([pinned.block_hash]) + with pytest.raises(ValueError, match="pinned"): + store.reserve("blocked", 6) + + store2 = PrefixCacheStore(_compat(), max_bytes=10, node_id="y") + store2.reserve("job", 6) + # Simulate a concurrently pinned service block to exercise the atomic + # publish guard independently of the reservation-aware public put path. + store2._put_locked(pinned) + store2.lookup([pinned.block_hash]) + final = CacheBlock.create(bytes.fromhex("05" * 32), 2, b"final!") + with pytest.raises(ValueError, match="pinned"): + store2.publish_and_lease( + [final], + [final.block_hash], + reservation_id="job", + ) + + +def test_internal_put_is_idempotent_and_rejects_collision(): + store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x") + block = CacheBlock.create(bytes(32), 1, b"abc") + assert store._put_locked(block) + assert not store._put_locked(block) + with pytest.raises(ValueError, match="collision"): + store._put_locked(CacheBlock.create(bytes(32), 1, b"xyz"))