diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index 620298d..809e2fd 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/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, diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 3a86912..357b46f 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/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 = []