Repository navigation
Add read-only observability snapshots - #385
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a small, read-only observability surface for kvcached integrations and KV cache pools, exposing structured “snapshot” dataclasses and helper methods so external monitoring/control-plane systems can introspect state without relying on private patch details.
Changes:
- Added
kvcached.observabilitymodule withRuntimeSnapshot/KVCachePoolSnapshot, plus capability discovery and snapshot builders. - Added
KVCacheManager.observability_snapshot()/observability_snapshot_dict()for pool-level snapshot export. - Added vLLM/SGLang integration runtime snapshot helpers and a new unit test covering JSON-serializability and snapshot field values.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/test_observability.py | Adds unit tests for capability discovery and snapshot serialization/field correctness. |
| kvcached/observability.py | Introduces stable schema/versioned snapshot dataclasses and snapshot builder helpers. |
| kvcached/kv_cache_manager.py | Exposes synchronized pool snapshot helper methods on KVCacheManager. |
| kvcached/integration/vllm/interfaces.py | Adds vLLM runtime integration snapshot + dict helper. |
| kvcached/integration/sglang/interfaces.py | Adds SGLang runtime integration snapshot + dict helper. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
cc5e037 to
287bf60
Compare
ae8b707 to
ccba645
Compare
0ef56da to
beb0ec1
Compare
beb0ec1 to
fb368a7
Compare
Two accounting problems surface when the snapshot fields are checked against a real KVCacheManager on GPU rather than a mock. 1. allocated_blocks was inflated by len(reserved_blocks). try_to_reserve() obtains its blocks through alloc(), so they have already left their pages and are counted by the full/avail page terms. _get_num_alloced_blocks() added len(self.reserved_blocks) on top, counting them twice: reserving 512 blocks moved allocated_blocks by 1024 instead of 512. The extra term also made the count *decrease* when alloc() drained the reserve ledger, even though nothing was freed. The term is dropped. The only in-tree caller is the lazy-shrink check in free(), where resize() has already asserted the ledger is empty, so shrink behaviour is unchanged. 2. available_blocks/available_bytes could be exported negative. available_size() derives from PageAllocator::get_num_free_pages(), which reads a mutable counter without holding the allocator mutex while every mutation takes it. A polling snapshot observed -127 available blocks. Clamp the exported gauges at zero; the unlocked read itself is a separate fix in csrc/page_allocator.cpp. Also documents the block-accounting fields: allocated/reserved/available overlap and must not be summed, and reserved_blocks (ledger) is unrelated to reserved_pages (prealloc thread).
KVCacheManager.observability_snapshot() is @synchronized, so it reads every field under a single hold of the manager lock -- the nested per-accessor acquisitions are reentrant. get_registered_kv_cache_pool_snapshots(), which is the path a monitor actually calls, bypassed that wrapper and invoked build_kv_cache_pool_snapshot() directly. That takes the lock once per accessor rather than once per snapshot, so available_size(), _get_num_alloced_blocks() and get_mapped_memory_size() are each internally consistent but come from different instants, and reserved_blocks / in_shrink / target_num_blocks are read under no lock at all. The exported snapshot can therefore describe a state the pool was never in. Measured with a writer thread mutating under the manager lock (async_sched=True, so _lock is a real RLock), counting writes that commit between a snapshot's first and last field read: registry path before: 35/40 snapshots torn (median 5, max 375) after: 0/40 manager path 0/40 in both cases Routing through the manager entry point also drops the call from ~119us to ~88us by acquiring the lock once instead of once per accessor. Manager-likes that do not implement observability_snapshot() still fall back to the module-level builder, so the duck-typed callers that function is documented to accept keep working. Note this only serializes under async_sched=True; with sync scheduling _lock is a NoOpLock by design. The unlocked read of num_free_pages_ in csrc/page_allocator.cpp remains a separate fix.
Every SGLang pool is created with a pool_name ("mha", "mla", "mamba"), but
the vLLM integration never passed one, so every vLLM pool reported
pool_name: None. That field is the only human-readable pool identifier in
the snapshot and part of its sort key, which left the feature effectively
SGLang-only.
vLLM has a single unified pool, injected by ElasticBlockPoolPatch over
vllm.v1.core.block_pool, so name it "block_pool" to match vLLM's own
terminology.
Verified through the patched ElasticBlockPool rather than the factory
directly: registered vLLM pools report pool_name='block_pool' where they
previously reported None.
… buffer
virtual_per_layer_bytes was set straight from manager.mem_size, which is the
virtual reservation for ONE KV buffer of ONE layer -- K (or V) for MHA, the
single combined buffer for MLA. That is why virtual_total_bytes had to
multiply by both num_layers and num_kv_buffers. So for MHA the field was
half the bytes actually reserved for a layer, and the name only held for
MLA:
num_layers=16, num_kv_buffers=2
virtual_total_bytes = 32768 MB
virtual_total / layers = 2048 MB <- bytes really reserved per layer
virtual_per_layer_bytes = 1024 MB <- what was reported
Scale by num_kv_buffers so the field means what its name says, and derive
the total as virtual_per_layer_bytes * num_layers.
virtual_total_bytes is unchanged for every configuration (the same three
factors, regrouped), and MLA pools are unchanged entirely since
num_kv_buffers is 1 there. Only virtual_per_layer_bytes moves, and only for
multi-buffer pools -- worth doing before the v1 schema ships, since the
downstream SGLang exporter turns this field name directly into the
kvcached_kv_cache_pool_virtual_per_layer_bytes gauge.
|
@shipiyouniao If the changes look good to you, we can merge this PR. |
Signed-off-by: 石皮幼鸟 <2960474346@qq.com>
|
Thanks, the latest changes look good to me. I added regression coverage for the corrected per-layer semantics, negative block gauges, the synchronized registry path, and reserved-block accounting; the targeted tests and all current CI checks pass. From my side, #385 is ready to merge. One follow-up is #436. Since #385 introduces the snapshot reader that can run concurrently with the allocator, and downstream monitoring/exporters may poll these fields continuously, the unlocked C++ counter reads should be fixed soon. I suggest landing #436 before #410/#421 and before enabling production polling. Clamping |
Sounds good to me! One thing i want to understand a bit more is that if we use lock to remove the data race, would that introduce too much overhead for kvcached because of the data polling thread could hold the lock and the kvcached process have to wait for the lock. |
|
Good question. I ran a focused benchmark:
At a normal 1-second polling interval, the polling thread itself introduced no measurable additional contention. The main overhead comes from acquiring the mutex on every foreground getter call. More importantly, adding I therefore think #436 should use relaxed atomic counters for the hot getters, plus a single locked snapshot method if monitoring requires a coherent multi-field view. Thanks for raising this concern. |
Extends the get_capabilities() record from ovg-project#385 with backend and per-integration capabilities, additive feature keys for landed and planned surfaces, and counter enumeration. No schema bump: new optional keys are additive and consumers treat a missing key as unsupported. Read-only, no behavior change. Implements item (3) of ovg-project#375.
Summary
Adds a read-only observability surface for kvcached runtime integrations and KV cache pools.
kvcached owns the allocator measurements and byte accounting. External monitoring systems can periodically collect and forward structured snapshots without depending on private allocator or engine-patch details.
What is included
kvcached.observabilitydataclasses for runtime state, KV pool state, and capability discoveryKVCacheManagerlifetimeKVCacheManager.observability_snapshot()andobservability_snapshot_dict()Design boundaries
Validation
Local validation:
GitHub Actions is green for pre-commit and all five MyPy jobs.
GPU validation used the SGLang v0.5.13 image on one explicitly isolated RTX 4090. A real CUDA VMM-backed MHA pool was created, sampled, allocated by 32 blocks, sampled again, freed, and sampled a third time: