Skip to content

Feat: Add KV Cache pool snapshot history - #485

Closed
kanaad-lims wants to merge 5 commits into
ovg-project:mainfrom
kanaad-lims:feature/pool-snapshot-history
Closed

kanaad-lims wants to merge 5 commits into
ovg-project:mainfrom
kanaad-lims:feature/pool-snapshot-history

Conversation

@kanaad-lims

@kanaad-lims kanaad-lims commented Sep 8, 2026 •

Copy link
Copy Markdown

Summary

Fixes #483
Add bounded in-memory history for KVCachePoolSnapshot observability data.

Changes

  • Maintain snapshot history independently for each (id(manager), group_id) pair to avoid collisions between pools sharing the same group_id.
  • Automatically record a timestamp + KVCachePoolSnapshot whenever build_kv_cache_pool_snapshot() is called.
  • Bound each history to 120 snapshots.
  • Add thread-safe access to the shared in-memory history.
  • Clean up a pool's history automatically when its manager is garbage-collected.
  • Add get_kv_cache_pool_snapshot_history(manager, group_id) to retrieve historical snapshots.
  • Store history entries as a timestamp plus a nested snapshot object to avoid timestamp field collisions.
  • Add clear_kv_cache_pool_history() for tests and manual resets.
  • Advertise support through get_capabilities()["features"]["pool_snapshot_history"].
  • Document that history records snapshot/poll calls rather than only state changes, and that history is shared among consumers observing the same pool.

Scope

This remains within the existing read-only, poll-only, record-only observability model. History is kept only in process memory and is lost when the serving process exits.

Testing

  • Added coverage for history accumulation.
  • Added timestamp and snapshot data validation.
  • Added isolation coverage for different managers sharing the same group_id.
  • Added bounded-history behavior.
  • Added history clearing and unknown-group behavior.
  • Added manager garbage-collection cleanup coverage.
  • Existing observability tests continue to pass.

pytest tests/test_observability.py -v — 16 passed.

@rishabhsinha17 rishabhsinha17 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ran tests/test_observability.py on de239db locally, 15 passed as the body says, and read the change against main and the two in-flight observability PRs. The bounded deque is the right primitive, the per-group isolation and unknown-group tests check real behavior, and storing snapshot objects rather than dicts means fields added later flow into history automatically, which keeps this compatible with #474's capability record and the lifecycle_phase field #478 adds. The findings below need work before this merges, the last one mechanical.

  1. Keying history by group_id conflates pools. group_id defaults to 0 and identifies a KV cache group inside one model (kv_cache_manager.py:70,86), not a pool. pool_registry keys by id(manager) with weakref cleanup because one process can run several pools (pool_registry.py:12-33), and get_registered_kv_cache_pool_snapshots sorts by integration then pool_name then group_id because group_id alone does not identify a pool (observability.py:312). Registered pools that both leave group_id at 0 interleave their timelines in one deque at :251. The isolation test only exercises distinct group ids (3 and 4), so it cannot catch this. The key needs pool identity, pool_name plus group_id at minimum.

  2. A reader can crash mid-poll. get_kv_cache_pool_snapshot_history iterates the live deque without a lock (:263-268) while build appends from whatever thread snapshots that pool, and CPython raises RuntimeError, deque mutated during iteration, on that interleaving. A full deque widens the window because append then also drops the left end. I reproduced that exception standalone with a plain deque. observability_snapshot's @synchronized (kv_cache_manager.py:716) serializes one manager only, and the duck-typed direct path (:293 here) does not take a lock, so nothing protects the module global. pool_registry guards its dict with a module lock for the same access pattern (_registered_pools_lock, pool_registry.py:13). Copying under such a lock and building dicts outside it closes this. Related retention gap, nothing ever drops a history key, a garbage-collected manager leaves its deque behind, and clear_kv_cache_pool_history only clears everything (:271-272) while get is per group. Processes that cycle models accumulate deques for pools that no longer exist. The registry's weakref cleanup is the in-repo precedent for both.

  3. The flattened entry has a silent-shadowing hazard. {"timestamp": t, **snapshot.to_dict()} (:266-268) puts the timestamp at the same level as record fields, and an entry then no longer matches the advertised pool_snapshot_fields (:135). Because the unpacking comes second, any future record field named timestamp would overwrite the history timestamp with no error, and #474's compatibility rule adds new optional record keys without a schema_version bump, which keeps that collision reachable. Nesting, {"timestamp": t, "snapshot": snapshot.to_dict()}, keeps the record exact and removes the hazard.

  4. The recording semantics deserve a sentence in the docstrings. History records poll calls, not state changes. Every consumer that snapshots shares the 120-entry budget, the registered-pools view appends too since it flows through the same builder (:293), and clear wipes other consumers' history. That is a workable v1 contract, but build_kv_cache_pool_snapshot's docstring does not mention the new side effect and the two new functions have no docstrings, in the module whose stated point is a documented stable surface. Worth also deciding on _HISTORY_MAXLEN (:21). The issue said configurable, the PR hardcodes 120. An env knob next to the ones in kvcached/utils.py would match convention, or keep it fixed and say so in the docstring.

  5. Mechanical. No CI has run yet, first-PR workflows sit awaiting approval, and when they run, pre-commit will fail isort on both files. kvcached/observability.py needs import time before from collections import deque with a blank line after the future import, and in the test file isort alphabetizes the two new names inside the existing import block. tests/test_observability.py also lacks its trailing newline. The pre-commit config comments out end-of-file-fixer, and the missing newline will not block, but worth fixing while in there. Clearing history at the start of each new test works. An autouse fixture that clears would also cover the pre-existing tests in this file, which now all append to the global as a side effect.

Happy to take another look after the rework.

@kanaad-lims

Copy link
Copy Markdown
Author

Hey @rishabhsinha17
Will take a look and make the changes.
Thank you!

@rishabhsinha17 rishabhsinha17 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed at a6832b2. This round resolves the keying, the unlocked reader, and the schema findings from the first one, and I checked each rather than reading the diff alone.

Keying is now (id(manager), group_id) with a weakref finalizer dropping a dead manager's entries, so two pools sharing group 0 no longer share history, and your isolation and GC tests pin both halves. The reader copies the deque under the module RLock. A two-second tight loop with two writer and two reader threads at a6832b2 raised nothing, while the same loop at de239db raised on roughly a third of reads. History entries nest the snapshot dict under a separate timestamp key, so nothing shadows snapshot fields and the record-keys contract coming in #474 stays satisfiable. Ran tests/test_observability.py at head, 16 passed, and the CI-pinned mypy over kvcached and tests matches main.

Two nits, neither blocking. The trailing-whitespace hook will fail CI on the whitespace-only lines at observability.py:296 and :301. clear_kv_cache_pool_history() skips the lock the other paths take, and while dict.clear() cannot crash a reader here, taking the lock there too would keep every access uniform. CI has never run on this branch and needs a maintainer to approve the workflow run.

@shipiyouniao shipiyouniao left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I rechecked a6832b2 and its clean merge with current main. The existing observability suite passes in both cases (16 passed), and the manager identity, nested snapshot record, bounded history and locked reader address the earlier review.

There is one additional compatibility regression in the unconditional weakref registration: the previously usable build_kv_cache_pool_snapshot(manager-like) path now raises for non-weak-referenceable adapters. I reproduced this on the same SimpleNamespace adapter: current main returns the snapshot, this head raises TypeError. Details below.

The trailing whitespace at observability.py:296 and :301 reported earlier is also still present (git diff --check fails). CI has not run on this head, so the required checks still need to complete before merge.

Comment thread kvcached/observability.py Outdated

with _pool_snapshot_history_lock:
if manager_id not in _pool_snapshot_manager_refs:
reference = weakref.ref(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep optional history from breaking the existing manager-like snapshot API. This builder previously accepted any adapter exposing the required fields/methods, including SimpleNamespace and slotted adapters without weakref. After successfully constructing the snapshot, the unconditional weakref.ref(manager, ...) now raises TypeError for those objects, so callers cannot even obtain the current snapshot. Reproduction using the test fixture: m = FakeManager(); adapter = types.SimpleNamespace(**{k: getattr(m, k) for k in dir(m) if not k.startswith("__")}); build_kv_cache_pool_snapshot(adapter). It succeeds on current main and fails here with cannot create weak reference to 'types.SimpleNamespace' object. Please preserve snapshot construction for these adapters and define a bounded/no-retention history fallback, rather than making weakref support a new implicit requirement. Add a contract regression test for this case.

@kanaad-lims

Copy link
Copy Markdown
Author

Hello @rishabhsinha17 and @shipiyouniao
Thank you for the review. I have made the requested changes as per the feedback.

  • Made weakref conditional so the objects that are not referable (e.g SimpleNamespace) can still build snapshots and have their history recorded.
  • Added regression test covering the non-weak-referenceable adapter case (based on @shipiyouniao test example)
  • Added the history lock to clear_kv_cache_pool_history().
  • Removed the trailing whitespaces. git diff --check returns empty.

pytest tests/test_observability.py -v passes all 17 test cases.

For non-weak-referenceable adapters, weakref-based GC cleanup is not available, so their history cannot be automatically removed when the adapter is garbage-collected. The history itself remains bounded by the existing 120-entry limit and does not retain a strong reference to the adapter.

@shipiyouniao shipiyouniao left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked 2ced054. The TypeError regression is fixed, the eight history-specific tests pass locally, and the whitespace issue is gone. There is still a lifecycle problem in the fallback, so I cannot approve yet.

When weakref.ref(manager) raises TypeError, the code continues into _pool_snapshot_history.setdefault((id(manager), group_id), ...). No weakref or other owner is registered for that entry. I reproduced this with the same SimpleNamespace adapter used in the new test: after deleting the adapter and running gc.collect(), its history entry remains in the global dictionary. The 120-entry limit bounds each deque, not the number of abandoned pools. Also, if Python later reuses that object ID, a different adapter can inherit the old history; I did not reproduce ID reuse in this run, but the lookup has no identity check to prevent it.

The simplest fallback is to return the snapshot without recording history when the object cannot be weak-referenced, and document that history requires a weak-referenceable manager. If history for these adapters is required, it needs an explicit ownership/cleanup mechanism instead. Please add a deletion/cleanup regression test rather than only checking the initial read. CI checks are still not reported for this head.

@shipiyouniao shipiyouniao left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked 006a613. Non-weak-referenceable adapters now return their snapshot without entering the history table, and the regression test checks that no entry remains after deletion. The new docstring also makes this limitation explicit. This resolves my outstanding lifecycle concern.

All 17 observability tests passed locally at 110d8a0; the only subsequent change in 006a613 is that explanatory docstring. A two-second concurrent check with two writers, two readers and a history-clear thread also completed without errors. No remaining blocking findings from my review.

@RixinLiu this is ready for your final double-check from my side. GitHub still reports no CI checks, so please ensure the required workflows run and pass before merging.

@kanaad-lims

Copy link
Copy Markdown
Author

Hello @shipiyouniao
I have addressed the lifecycle issue for non-weak referenceable manager adapters. Snapshot construction continues to work normally, but the history is no longer recorded when the manager cannot be weak-referenced.

Let me know if any additional changes are required.
Thank you!

@rishabhsinha17 rishabhsinha17 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed at 006a613 against my a6832b2 approval. The interdiff touches only the weakref registration, the clear path and the new adapter test, and neither earlier finding regresses. Keying stays (id(manager), group_id) with the finalizer, and the TypeError path returns before touching the history table, which the new adapter test pins. The reader still copies under the lock, and clear_kv_cache_pool_history now takes it too, closing the one unlocked mutation that remained. 17 tests pass locally at this head and a two second loop with two writers, two readers and a clear thread raised nothing. Still approved.

@RixinLiu

Copy link
Copy Markdown
Collaborator

Really appreciate you trying to contribute to kvcached. But i have to say there is no need to retain polling history inside kvcached. This should only done by those consumers who asked for these information.
Sorry for the previous rounds of reviews, i will talk about this later to @shipiyouniao and @rishabhsinha17

@RixinLiu RixinLiu closed this Sep 21, 2026
@kanaad-lims

Copy link
Copy Markdown
Author

Thank you for the clarification @RixinLiu. I appreciate you and other maintainers @rishabhsinha17 and @shipiyouniao for taking the time to review the PR and providing feedback throughout the process. I’m still excited to contribute to kvcached and work on other issues where I can be helpful.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Feature: Add in-memory history tracking to KVCachePoolSnapshot observability.

4 participants