fix: roll back partial allocation when alloc_page() fails in KVCacheManager.alloc()

[rishabhsinha17]

Summary

Roll back partial allocations when PageAllocator.alloc_page() fails mid-alloc, and return an allocation miss instead of leaking blocks.

Fixes #364.

Root cause

KVCacheManager._alloc() consumes reserved blocks and blocks from existing pages before it needs a new physical page. Under multi-instance pressure, available_size() can report enough logical capacity while another instance drains the shared physical pool, so alloc_page() raises RuntimeError after the allocation has already been partially served. Those blocks were removed from reserved_blocks and page free lists with no owner — leaked.

Changes

On RuntimeError from the page-consuming loop:

• blocks taken from pages are returned through the regular free() path — safe to call re-entrantly, since _lock is a threading.RLock/NoOpLock and try_to_reserve() already nests synchronized calls the same way. This also releases a page back to the shared physical pool if the rollback empties it, which is exactly the resource the competing instance is starved of
• blocks taken off the reservation ledger are prepended back onto reserved_blocks, so a try_to_reserve() caller does not silently lose its reservation
• alloc() returns None, so the caller's existing allocation-miss handling applies

The success path is unchanged: the loop body is identical, only indented into the try.

Validation

kv_cache_manager.py hard-imports the compiled kvcached.vmm_ops extension, so the added tests/test_alloc_rollback.py stubs it with pure-Python FakePage/FakePageAllocator in sys.modules (only when the real extension is unavailable) and constructs the manager without __init__ — the same approach as the regression test merged in #407. The fake allocator reports ample capacity but fails alloc_page() after N pages, reproducing the race in the issue. Covered: clean miss with no partial state, page blocks rolled back and immediately reusable, reserved blocks restored, mixed reserved + page blocks with the emptied page returned to the physical pool, allocation succeeding after the pool recovers, and a success-path regression guard.

The tests fail on the unfixed code (5 of 6; the passing one is the success-path guard):

git stash push kvcached/kv_cache_manager.py && python -m pytest tests/test_alloc_rollback.py -q  # 5 failed, 1 passed
git stash pop && python -m pytest tests/test_alloc_rollback.py -q                                # 6 passed

The full CPU-only suite passes with the stub in place. ruff check, isort --check-only, and the mypy 3.9–3.13 matrix pass. I have not run the GPU-backed tests/test_kvcache_manager.py, as this machine has no CUDA device.

Related: #365 also addresses this issue with a different rollback strategy (it keeps emptied pages held by the instance and does not restore the reservation ledger; this PR releases emptied pages to the shared pool and restores reservations).

[RixinLiu]

LGTM, thanks for the contribution @rishabhsinha17 !

#365 fixes the same issue with a different approach, so these two can't both land. I'm inclined to take this one, but I'd like @shipiyouniao to check before I merge.

Goal. A rollback should leave the manager in the state it was in immediately before the alloc() call that failed.

case	state before the call	reachable?	#365	#430
A. this call creates a page, then alloc_page() fails	page does not exist	yes	page kept, num_avail 0 → 4 ✗	page released, num_avail back to 0 ✓
B. reservation + partially-free pre-existing page	reserved=[0,1], page 0 free [2,3]	yes	reserved=[], page 0 → [0,1,2,3], num_avail 2 → 4 ✗	all three identical ✓
C. pre-existing page that is already fully free	page fully free yet still in avail_pages	no	✓	✗

Case B matters most to me: try_to_reserve() promises the caller that n blocks are held for it. Sending them back through the free path returns them to the general pool and removes them from reserved_blocks, so the promise is cancelled with no signal.

Case C is the one your test_alloc_page_failure_rolls_back_blocks_from_existing_page asserts on, and #430 fails it. I believe that state is unreachable: avail_pages is written in exactly two places, kv_cache_manager.py:308, right after page.alloc(num_from_page) with num_from_page >= 1, and :394, in the explicit else of if page.empty(). So every page in avail_pages has at least one block outstanding.

If that holds, free() only sees if page.empty(), but the rule we want is "release iff the page did not exist before the call". Since a pre-existing page always has another holder's block still out (never empty) while a page created by this call has none (always empty). So no separate record of which pages the call created is needed, and that equivalence breaks only in case C, which release_empty_pages=False is itself the only code path that can produce.

What I'd like @shipiyouniao to check: is there a path I've missed that leaves a fully-free page in avail_pages? If not, I'll merge #430 and close this.

[shipiyouniao]

Please narrow this RuntimeError handler to the page_allocator.alloc_page() call that this recovery path is designed for. As written, it also catches InternalPage.alloc()'s "Not enough free blocks" invariant failure. By then _pick_avail_page() may already have removed the current page from avail_pages, while that page's blocks are not yet in ret_index, so _rollback_partial_alloc() cannot restore the page and returning None would leave the manager running with lost state. That invariant failure should remain fail-loud; a regression test where page.alloc() raises would pin down the distinction.

[shipiyouniao]

I checked the avail_pages reachability question. The production write paths preserve the invariant Rixin described:

• the allocation path inserts a page only after allocating at least one block from it, so it is not fully free;
• free() returns an empty page to PageAllocator and reinserts it only in the non-empty branch;
• the manager lock prevents another operation from invalidating that reasoning during rollback.

Therefore case C is not reachable through the manager API. For a failed allocation, blocks from a pre-existing page restore that page to its prior partially occupied state, while a page created by the failed call becomes empty and is returned to the shared pool. Restoring consumed reservation entries separately is also the correct contract.

I ran the new rollback tests together with the related reservation-order and page-selection tests: 16 passed. I left one inline request to narrow the caught RuntimeError to alloc_page() so an unrelated InternalPage.alloc() invariant failure cannot be downgraded into a miss with an untracked page. With that scope tightened and covered, this looks ready to merge over #365.