Merge pull request #188 from FluffyAIcode/AgentMemory/precise-snapshot-reservation-0717

fix(prefill): preserve unrelated snapshots during jobs

Author: FluffyAIcode

inference_engine/distributed/prefill_worker.py

@@ -267,12 +267,12 @@ 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.
+            # Reserve a conservative estimate before model compute. Atomic
+            # publish protects the final snapshot without evicting unrelated
+            # content-addressed snapshots needed by later restore requests.
             self.cache_store.reserve(
                 job.job_id,
-                self.cache_store.max_bytes,
+                len(job.token_ids) * self.estimated_snapshot_bytes_per_token,
             )
             blocks = tuple(self._engine_for_current_thread().compute_prefill(
                 job.token_ids,

tests/inference_engine/distributed/test_prefill_worker.py

@@ -256,6 +256,32 @@ def test_job_rejects_snapshot_capacity_before_model_compute():
         jobs.close()
 
 
+def test_job_reservation_preserves_unrelated_restore_snapshot():
+    cache = PrefixCacheStore(COMPAT, max_bytes=100, node_id="shared")
+    old = CacheBlock.create(b"z" * 32, 2, b"old-snapshot")
+    cache.put(old)
+    jobs = PrefillJobStore(
+        _Engine(),
+        cache,
+        estimated_snapshot_bytes_per_token=10,
+    )
+    try:
+        job = jobs.submit(
+            request_id="new-prefix",
+            tenant_id="tenant",
+            token_ids=[1, 2],
+            block_hashes=[b"a" * 32],
+            compatibility=COMPAT,
+            compression=CompressionCodec.NONE,
+        )
+        job.future.result(timeout=1)
+        assert job.state == PrefillJobState.COMPLETED
+        assert old.block_hash in cache.block_hashes()
+        assert job.block_hash in cache.block_hashes()
+    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 = []