Add read-only observability snapshots

kvcached/integration/sglang/interfaces.py

@@ -7,6 +7,7 @@
 import torch
 
 from kvcached.kv_cache_manager import KVCacheManager
+from kvcached.observability import build_runtime_snapshot
 from kvcached.tp_ipc_util import start_worker_listener_thread
 from kvcached.utils import CONTIGUOUS_LAYOUT, PAGE_SIZE, get_kvcached_logger, normalize_gpu_device
 from kvcached.vmm_ops import (
@@ -63,6 +64,24 @@ def shutdown_kvcached() -> None:
     _async_sched = False
 
 
+def observability_snapshot():
+    """Return a read-only snapshot of the SGLang integration state."""
+    return build_runtime_snapshot(
+        engine="sglang",
+        initialized=_kvcached_initialized,
+        device=_kvcached_device,
+        world_size=_world_size,
+        pp_rank=_pp_rank,
+        async_sched=_async_sched,
+        contiguous_layout=_contiguous_layout,
+    )
+
+
+def observability_snapshot_dict() -> Dict[str, Any]:
+    """Return a JSON-serializable snapshot of the SGLang integration state."""
+    return observability_snapshot().to_dict()
+
+
 def alloc_kv_cache(
     kvcache_shape: Tuple[int, ...],
     dtype: torch.dtype,

kvcached/integration/vllm/interfaces.py

@@ -2,11 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 
 from kvcached.kv_cache_manager import KVCacheManager
+from kvcached.observability import build_runtime_snapshot
 from kvcached.tp_ipc_util import start_worker_listener_thread
 from kvcached.utils import CONTIGUOUS_LAYOUT, PAGE_SIZE, get_kvcached_logger, normalize_gpu_device
 from kvcached.vmm_ops import (
@@ -85,6 +86,25 @@ def shutdown_kvcached() -> None:
     _async_sched = False
 
 
+def observability_snapshot():
+    """Return a read-only snapshot of the vLLM integration state."""
+    return build_runtime_snapshot(
+        engine="vllm",
+        initialized=_kvcached_initialized,
+        device=_kvcached_device,
+        world_size=_world_size,
+        pp_rank=_pp_rank,
+        async_sched=_async_sched,
+        contiguous_layout=_contiguous_layout,
+        is_worker=_is_worker,
+    )
+
+
+def observability_snapshot_dict() -> Dict[str, Any]:
+    """Return a JSON-serializable snapshot of the vLLM integration state."""
+    return observability_snapshot().to_dict()
+
+
 def alloc_kv_cache(
     kvcache_shape: Tuple[int, ...],
     block_size: int,

kvcached/kv_cache_manager.py

@@ -440,6 +440,24 @@ def get_mapped_memory_size(self, unit='bytes') -> float:
         else:
             raise ValueError(f"Unknown unit: {unit}")
 
+    @synchronized
+    def observability_snapshot(self, *, integration=None, pool_name=None):
+        """Return a read-only snapshot of this KV cache pool."""
+        from kvcached.observability import build_kv_cache_pool_snapshot
+        return build_kv_cache_pool_snapshot(
+            self,
+            integration=integration,
+            pool_name=pool_name,
+        )
+
+    @synchronized
+    def observability_snapshot_dict(self, *, integration=None, pool_name=None):
+        """Return a JSON-serializable read-only snapshot of this KV cache pool."""
+        return self.observability_snapshot(
+            integration=integration,
+            pool_name=pool_name,
+        ).to_dict()
+
     @synchronized
     def clear(self):
         """

kvcached/observability.py

@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: Copyright contributors to the kvcached project
+# SPDX-License-Identifier: Apache-2.0
+
+"""Read-only observability snapshots for kvcached integrations.
+
+The structures in this module are intentionally policy-free.  They expose
+allocator and runtime state so monitoring systems or external control planes
+can consume kvcached status without depending on private patch details.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional
+
+
+SCHEMA_VERSION = "kvcached.observability.v1"
+
+
+def _call_int(obj: Any, name: str) -> Optional[int]:
+    method = getattr(obj, name, None)
+    if method is None:
+        return None
+    return int(method())
+
+
+def _int_attr(obj: Any, name: str) -> Optional[int]:
+    value = getattr(obj, name, None)
+    if value is None:
+        return None
+    return int(value)
+
+
+@dataclass(frozen=True)
+class RuntimeSnapshot:
+    """Runtime integration state for an engine shim."""
+
+    schema_version: str
+    engine: str
+    initialized: bool
+    device: Optional[str]
+    world_size: int
+    pp_rank: int
+    async_sched: bool
+    contiguous_layout: bool
+    is_worker: Optional[bool] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class KVCachePoolSnapshot:
+    """Read-only state of one kvcached-backed KV pool."""
+
+    schema_version: str
+    pool_type: str
+    integration: Optional[str]
+    pool_name: Optional[str]
+    group_id: int
+    num_layers: int
+    num_kv_buffers: int
+    page_size_bytes: int
+    block_size_bytes: int
+    total_blocks: int
+    available_blocks: int
+    allocated_blocks: int
+    reserved_blocks: int
+    null_block_reserved: bool
+    virtual_per_layer_bytes: int
+    virtual_total_bytes: int
+    mapped_bytes: int
+    total_pages: Optional[int]
+    free_pages: Optional[int]
+    inuse_pages: Optional[int]
+    reserved_pages: Optional[int]
+    available_physical_pages: Optional[int]
+    effective_free_pages: Optional[int]
+    in_shrink: bool
+    shrink_target_blocks: Optional[int]
+    resize_target_bytes: Optional[int]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+def get_capabilities() -> Dict[str, Any]:
+    """Return the stable observability surface currently exposed by kvcached."""
+
+    return {
+        "schema_version": SCHEMA_VERSION,
+        "features": {
+            "runtime_snapshot": True,
+            "kv_cache_pool_snapshot": True,
+            "read_only": True,
+            "policy_control": False,
+        },
+        "pool_snapshot_fields": list(KVCachePoolSnapshot.__dataclass_fields__.keys()),
+        "runtime_snapshot_fields": list(RuntimeSnapshot.__dataclass_fields__.keys()),
+    }
+
+
+def build_runtime_snapshot(
+    *,
+    engine: str,
+    initialized: bool,
+    device: Optional[str],
+    world_size: int,
+    pp_rank: int,
+    async_sched: bool,
+    contiguous_layout: bool,
+    is_worker: Optional[bool] = None,
+) -> RuntimeSnapshot:
+    return RuntimeSnapshot(
+        schema_version=SCHEMA_VERSION,
+        engine=engine,
+        initialized=initialized,
+        device=device,
+        world_size=world_size,
+        pp_rank=pp_rank,
+        async_sched=async_sched,
+        contiguous_layout=contiguous_layout,
+        is_worker=is_worker,
+    )
+
+
+def build_kv_cache_pool_snapshot(
+    manager: Any,
+    *,
+    integration: Optional[str] = None,
+    pool_name: Optional[str] = None,
+) -> KVCachePoolSnapshot:
+    """Build a read-only snapshot from a ``KVCacheManager``-like object."""
+
+    allocator = manager.page_allocator
+    free_pages = _call_int(allocator, "get_num_free_pages")
+    reserved_pages = _call_int(allocator, "get_num_reserved_pages")
+    available_physical_pages = _call_int(allocator, "get_avail_physical_pages")
+    if free_pages is None or reserved_pages is None or available_physical_pages is None:
+        effective_free_pages = None
+    else:
+        effective_free_pages = min(free_pages, available_physical_pages + reserved_pages)
+
+    mapped_bytes = int(manager.get_mapped_memory_size("bytes"))
+    virtual_per_layer_bytes = _int_attr(manager, "mem_size") or 0
+    num_layers = _int_attr(manager, "num_layers") or 0
+    num_kv_buffers = _int_attr(manager, "num_kv_buffers") or 0
+
+    return KVCachePoolSnapshot(
+        schema_version=SCHEMA_VERSION,
+        pool_type="kv_cache",
+        integration=integration,
+        pool_name=pool_name,
+        group_id=_int_attr(manager, "group_id") or 0,
+        num_layers=num_layers,
+        num_kv_buffers=num_kv_buffers,
+        page_size_bytes=_int_attr(manager, "page_size") or 0,
+        block_size_bytes=_int_attr(manager, "block_mem_size") or 0,
+        total_blocks=_int_attr(manager, "num_blocks") or 0,
+        available_blocks=int(manager.available_size()),
+        allocated_blocks=int(manager._get_num_alloced_blocks()),
+        reserved_blocks=len(getattr(manager, "reserved_blocks", [])),
+        null_block_reserved=getattr(manager, "null_block", None) is not None,
+        virtual_per_layer_bytes=virtual_per_layer_bytes,
+        virtual_total_bytes=virtual_per_layer_bytes * num_layers * num_kv_buffers,
+        mapped_bytes=mapped_bytes,
+        total_pages=_call_int(allocator, "get_num_total_pages"),
+        free_pages=free_pages,
+        inuse_pages=_call_int(allocator, "get_num_inuse_pages"),
+        reserved_pages=reserved_pages,
+        available_physical_pages=available_physical_pages,
+        effective_free_pages=effective_free_pages,
+        in_shrink=bool(getattr(manager, "in_shrink", False)),
+        shrink_target_blocks=getattr(manager, "target_num_blocks", None),
+        resize_target_bytes=_call_int(allocator, "get_resize_target"),
+    )

tests/test_observability.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright contributors to the kvcached project
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import importlib.util
+import sys
+import types
+
+if importlib.util.find_spec("torch") is None:
+    sys.modules.setdefault("torch", types.ModuleType("torch"))
+
+from kvcached.observability import (  # noqa: E402
+    build_kv_cache_pool_snapshot,
+    build_runtime_snapshot,
+    get_capabilities,
+)
+
+
+class FakePageAllocator:
+    def get_num_free_pages(self):
+        return 10
+
+    def get_num_inuse_pages(self):
+        return 6
+
+    def get_num_total_pages(self):
+        return 20
+
+    def get_num_reserved_pages(self):
+        return 2
+
+    def get_avail_physical_pages(self):
+        return 4
+
+    def get_resize_target(self):
+        return 0
+
+
+class FakeManager:
+    num_blocks = 128
+    block_mem_size = 4096
+    num_layers = 8
+    num_kv_buffers = 2
+    group_id = 3
+    page_size = 2 * 1024 * 1024
+    mem_size = num_blocks * block_mem_size
+    reserved_blocks = [0, 7]
+    null_block = [0]
+    in_shrink = False
+    target_num_blocks = None
+    page_allocator = FakePageAllocator()
+
+    def available_size(self):
+        return 64
+
+    def _get_num_alloced_blocks(self):
+        return 16
+
+    def get_mapped_memory_size(self, unit="bytes"):
+        assert unit == "bytes"
+        return 6 * self.num_layers * self.page_size * self.num_kv_buffers
+
+
+def test_capabilities_are_json_serializable():
+    capabilities = get_capabilities()
+
+    assert capabilities["schema_version"] == "kvcached.observability.v1"
+    assert capabilities["features"]["read_only"] is True
+    assert capabilities["features"]["policy_control"] is False
+    json.dumps(capabilities)
+
+
+def test_runtime_snapshot_dict():
+    snapshot = build_runtime_snapshot(
+        engine="vllm",
+        initialized=True,
+        device="cuda:0",
+        world_size=2,
+        pp_rank=1,
+        async_sched=True,
+        contiguous_layout=False,
+        is_worker=True,
+    )
+
+    assert snapshot.to_dict() == {
+        "schema_version": "kvcached.observability.v1",
+        "engine": "vllm",
+        "initialized": True,
+        "device": "cuda:0",
+        "world_size": 2,
+        "pp_rank": 1,
+        "async_sched": True,
+        "contiguous_layout": False,
+        "is_worker": True,
+    }
+
+
+def test_kv_cache_pool_snapshot_from_manager_like_object():
+    snapshot = build_kv_cache_pool_snapshot(
+        FakeManager(),
+        integration="sglang",
+        pool_name="full_attention",
+    )
+    data = snapshot.to_dict()
+
+    assert data["pool_type"] == "kv_cache"
+    assert data["integration"] == "sglang"
+    assert data["pool_name"] == "full_attention"
+    assert data["group_id"] == 3
+    assert data["available_blocks"] == 64
+    assert data["allocated_blocks"] == 16
+    assert data["reserved_blocks"] == 2
+    assert data["null_block_reserved"] is True
+    assert data["virtual_total_bytes"] == 128 * 4096 * 8 * 2
+    assert data["mapped_bytes"] == 6 * 8 * (2 * 1024 * 1024) * 2
+    assert data["total_pages"] == 20
+    assert data["free_pages"] == 10
+    assert data["inuse_pages"] == 6
+    assert data["reserved_pages"] == 2
+    assert data["available_physical_pages"] == 4
+    assert data["effective_free_pages"] == 6
+    assert data["resize_target_bytes"] == 0
+    json.dumps(data)