diff --git a/daser/server/__main__.py b/daser/server/__main__.py index c0d2d51..233add4 100644 --- a/daser/server/__main__.py +++ b/daser/server/__main__.py @@ -24,6 +24,7 @@ from daser.server.core import ServerCore from daser.server.doc_registry import DocRegistry from daser.server.http import HTTPServerConfig, VLLMClient, build_http_app +from daser.server.http.metrics import MetricsCollector from daser.server.ipc import IPCServer from daser.server.metadata_store import MetadataStore @@ -196,6 +197,11 @@ def _parse_args() -> argparse.Namespace: help="L1 memory-tier capacity for --transfer-mode=iouring. Defaults " "to min(1GiB, --l2-size).", ) + parser.add_argument( + "--enable-metrics", + action="store_true", + help="Enable Prometheus metrics on the main HTTP server at /metrics.", + ) return parser.parse_args() @@ -481,7 +487,12 @@ async def run_server(args: argparse.Namespace) -> None: if cfg.transfer_mode != "gds": await ipc_server.initialize_transfer() - app = build_http_app(_build_http_config(args), core) + metrics_collector = MetricsCollector(enabled=True) if args.enable_metrics else None + app = build_http_app( + _build_http_config(args), + core, + metrics_collector=metrics_collector, + ) uvicorn_config = uvicorn.Config( app=app, host=args.host, diff --git a/daser/server/chunk_manager.py b/daser/server/chunk_manager.py index 9c1e31c..2540b87 100644 --- a/daser/server/chunk_manager.py +++ b/daser/server/chunk_manager.py @@ -57,6 +57,11 @@ def store(self) -> MetadataStore: """The underlying MetadataStore.""" return self._store + @property + def total_slots(self) -> int: + """Total slot capacity of the ring buffer.""" + return self._total_slots + @property def free_slots(self) -> int: """Number of slots currently available without eviction. diff --git a/daser/server/core.py b/daser/server/core.py index fca10c7..fe72c08 100644 --- a/daser/server/core.py +++ b/daser/server/core.py @@ -224,6 +224,7 @@ def __init__( self._late_evicted_commits = 0 self._lookup_requests = 0 self._lookup_hits = 0 + self._eviction_counts: dict[str, int] = {} self._committed_chunk_keys: set[str] = set() self._commit_waiters: dict[str, set[asyncio.Future[None]]] = {} @@ -242,12 +243,23 @@ async def rebuild_retrieval_index(self) -> None: await self._ri.insert(meta) self._committed_chunk_keys.add(meta.chunk_key) - async def lookup(self, tokens: list[int], model_id: str) -> list[ChunkInfo]: + async def lookup( + self, + tokens: list[int], + model_id: str, + record_access: bool = True, + ) -> list[ChunkInfo]: """Look up cached chunks for token IDs. + Hits update per-chunk access statistics by default. Read-only callers + such as diagnostics can pass ``record_access=False`` to avoid making an + inspection request look like production cache demand. + Args: tokens: prompt token IDs. model_id: model identifier. + record_access: whether hits should update access counters and + aggregate lookup hit/request statistics. Returns: List of matching chunks, possibly empty. @@ -256,9 +268,12 @@ async def lookup(self, tokens: list[int], model_id: str) -> list[ChunkInfo]: Performs no blocking I/O and should run on the server event loop. """ matches = await self._ri.lookup(tokens, model_id) - self._lookup_requests += 1 - if matches: - self._lookup_hits += 1 + if record_access: + self._lookup_requests += 1 + if matches: + self._lookup_hits += 1 + for match in matches: + self._cm.store.touch(match.meta.chunk_key) return [self._chunk_info(match) for match in matches] async def alloc_chunk( @@ -432,6 +447,17 @@ async def commit_stats(self) -> dict[str, int]: "lookup_hits": self._lookup_hits, } + def eviction_stats(self) -> dict[str, int]: + """Return cumulative chunk eviction counters by reason. + + Returns: + Mapping from eviction reason to cumulative count. + + Async/thread-safety: + Reads in-memory counters on the server event loop. + """ + return dict(self._eviction_counts) + async def live_allocations(self, allocations: list[dict[str, Any]]) -> list[str]: """Return chunk keys that still own their allocated slot ranges. @@ -499,6 +525,7 @@ async def evict_chunk(self, chunk_key: str) -> None: if meta is not None: self._mark_chunk_evicted_in_docs(meta) self._cm.store.remove(chunk_key) + self._record_eviction("explicit") self._committed_chunk_keys.discard(chunk_key) self._evicted_chunk_keys.add(chunk_key) logger.debug("[CORE] evict_chunk key=%s", chunk_key[:8]) @@ -636,6 +663,8 @@ async def delete_document(self, doc_id: str) -> DeleteDocumentResult: if self._detach_doc_from_chunk(doc_id, key): await self._ri.remove(key) chunks_evicted += 1 + if chunks_evicted: + self._record_eviction("document_delete", chunks_evicted) logger.info( "[CORE] delete_document doc_id=%s chunks_evicted=%d", doc_id, @@ -653,6 +682,7 @@ async def _drain_ring_evictions(self) -> None: await self._ri.remove(chunk_key) self._committed_chunk_keys.discard(chunk_key) self._evicted_chunk_keys.add(chunk_key) + self._record_eviction("ring") logger.debug("[CORE] removed auto-evicted chunk key=%s", chunk_key[:8]) def _notify_commit_waiters(self, chunk_key: str) -> None: @@ -681,6 +711,15 @@ def _mark_chunk_evicted_in_docs(self, meta: ChunkMeta) -> None: for doc_id in list(meta.doc_ids): registry.mark_chunk_evicted(doc_id, meta.chunk_key) + def _record_eviction(self, reason: str, count: int = 1) -> None: + """Increment the cumulative eviction counter. + + Args: + reason: eviction reason label. + count: number of chunks evicted for this reason. + """ + self._eviction_counts[reason] = self._eviction_counts.get(reason, 0) + count + def _detach_doc_from_chunk(self, doc_id: str, chunk_key: str) -> bool: """Detach a document reference and remove unreferenced chunks. diff --git a/daser/server/http/app.py b/daser/server/http/app.py index cbd1d9f..e4e821e 100644 --- a/daser/server/http/app.py +++ b/daser/server/http/app.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Standard +from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass from importlib import resources @@ -9,8 +10,8 @@ import uuid # Third Party -from fastapi import FastAPI, HTTPException -from fastapi.responses import FileResponse +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import FileResponse, Response from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Field @@ -19,6 +20,7 @@ from daser.server.core import ServerCore from daser.server.doc_registry import DocEntry from daser.server.http.chunker import Chunker, TokenChunk +from daser.server.http.metrics import MetricsCollector from daser.server.http.vllm_client import VLLMClient logger = init_logger(__name__) @@ -59,6 +61,13 @@ class HTTPServerConfig: align_document_chunks: bool = False +class DiagExplainRequest(BaseModel): + """Request body for ``POST /diag/explain``.""" + + doc_ids: list[str] = Field(..., description="Doc IDs to include in the prompt") + task: str = Field(default="", description="User task appended after documents") + + class UploadRequest(BaseModel): """Request body for ``POST /documents``.""" @@ -462,7 +471,7 @@ async def _prewarm_fixed_segments( chunk = chunker.single_chunk(segment_tokens, pad_token) if chunk.chunk_key in prewarmed_fixed_segments: continue - if await core.lookup(chunk.tokens, cfg.model): + if await core.lookup(chunk.tokens, cfg.model, record_access=False): prewarmed_fixed_segments.add(chunk.chunk_key) continue await _prefill_chunks(vllm, [chunk], label) @@ -474,6 +483,7 @@ def build_http_app( core: ServerCore, tokenizer: Any | None = None, vllm: VLLMClient | None = None, + metrics_collector: MetricsCollector | None = None, ) -> FastAPI: """Construct the HTTP server app. @@ -482,6 +492,8 @@ def build_http_app( core: shared server core. tokenizer: optional tokenizer override for tests. vllm: optional vLLM client override for tests. + metrics_collector: optional Prometheus collector. When provided, + ``/metrics`` is served and request handlers record counters. Returns: FastAPI instance ready for uvicorn. @@ -532,6 +544,36 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: ) chunker = Chunker(block_tokens=cfg.block_tokens) + @app.middleware("http") + async def record_http_metrics( + request: Request, + call_next: Callable[[Request], Awaitable[Response]], + ) -> Response: + """Record HTTP request counters when metrics are enabled. + + Args: + request: incoming FastAPI request. + call_next: downstream ASGI handler. + + Returns: + HTTP response from the downstream handler. + + Async/thread-safety: + Runs on the FastAPI event loop and only mutates the injected + ``MetricsCollector``. + """ + if metrics_collector is None: + return await call_next(request) + status = 500 + try: + response = await call_next(request) + status = response.status_code + return response + finally: + route = request.scope.get("route") + endpoint = str(getattr(route, "path", request.url.path)) + metrics_collector.record_http_request(request.method, endpoint, status) + @app.get("/", include_in_schema=False) async def web_ui() -> FileResponse: """Serve the built-in DaseR Web UI.""" @@ -546,7 +588,7 @@ async def _ensure_fixed_segment_cached(label: str, tokens: list[int]) -> None: chunk = chunker.single_chunk(tokens, pad_token) if chunk.chunk_key in prewarmed_fixed_segments: return - if await core.lookup(chunk.tokens, cfg.model): + if await core.lookup(chunk.tokens, cfg.model, record_access=False): prewarmed_fixed_segments.add(chunk.chunk_key) return await _prefill_chunks(vllm, [chunk], label) @@ -561,6 +603,27 @@ async def health() -> dict[str, Any]: "vllm": vllm_ok, } + @app.get("/metrics") + async def metrics() -> Response: + """Expose Prometheus text-format metrics. + + Returns: + Plain-text Prometheus metrics payload. Returns ``404`` when the + metrics collector is not enabled (backward compatible). + + Async/thread-safety: + Reads in-memory gauges from ``ServerCore`` and calls + ``generate_latest`` which is thread-safe through the + prometheus_client internal lock. + """ + if metrics_collector is None or not metrics_collector.enabled: + raise HTTPException(status_code=404, detail="metrics disabled") + metrics_collector.state_snapshot(core) + return Response( + content=metrics_collector.export_metrics(), + media_type="text/plain; version=0.0.4", + ) + @app.post("/documents", status_code=201) async def upload_document(req: UploadRequest) -> dict[str, Any]: """Upload a document, prefill chunk KV, and register it.""" @@ -587,6 +650,8 @@ async def upload_document(req: UploadRequest) -> dict[str, Any]: chunk_keys = await _prefill_chunks(vllm, chunks, "document") await _wait_for_committed_chunks(core, chunks) prefill_ms = (time.time() - t0) * 1000 + if metrics_collector is not None: + metrics_collector.record_document_prefill(prefill_ms) prompt_tokens = ( _tokens_from_chunks(chunks) if cfg.align_document_chunks else tokens ) @@ -649,6 +714,83 @@ async def delete_document(doc_id: str) -> dict[str, Any]: ) from exc return {"ok": True, "chunks_evicted": result.chunks_evicted} + @app.post("/diag/explain") + async def diag_explain(req: DiagExplainRequest) -> dict[str, Any]: + """Simulate prompt assembly and cache lookup without running inference. + + Returns chunk-level hit/miss information and the assembled prompt + token count so operators can inspect cache behaviour before sending + a live ``/infer`` request. + + Async/thread-safety: + Reads in-memory state on the FastAPI event loop through ``core``. + """ + if not req.doc_ids: + raise HTTPException(status_code=400, detail="doc_ids must not be empty") + + docs: list[DocEntry] = [] + missing_docs: list[str] = [] + for doc_id in req.doc_ids: + doc = await core.get_document(doc_id) + if doc is None: + missing_docs.append(doc_id) + continue + docs.append(doc) + + if missing_docs: + raise HTTPException( + status_code=404, + detail=f"documents not found: {', '.join(missing_docs)}", + ) + + prompt_segments, prompt_preview = _build_prompt_segments( + tokenizer, + cfg.system_prompt, + cfg.doc_separator, + req.task, + docs, + ) + prompt_tokens: list[int] = [] + for segment in prompt_segments: + prompt_tokens.extend( + _segment_tokens( + chunker, + segment.tokens, + pad_token, + cfg.align_document_chunks and segment.fixed, + ) + ) + + cache_hits = [ + chunk.to_dict() + for chunk in await core.lookup( + prompt_tokens, + cfg.model, + record_access=False, + ) + ] + hit_count = len(cache_hits) + hit_token_count = sum(int(hit.get("token_count", 0)) for hit in cache_hits) + coverage_ratio = hit_token_count / len(prompt_tokens) if prompt_tokens else 0.0 + hit_tier = "L1" if hit_count else "miss" + routing_decision = "chunk_reuse" if hit_count else "full_prefill" + + return { + "doc_ids": req.doc_ids, + "prompt_tokens": len(prompt_tokens), + "prompt_preview": prompt_preview, + "cache_hits": cache_hits, + "hit_count": hit_count, + "hit_token_count": hit_token_count, + "coverage_ratio": round(coverage_ratio, 4), + "routing_decision": routing_decision, + "hit_tier": hit_tier, + "expected_ttft_ms": None, + "accuracy_net_enabled": False, + "cache_reuse_mode": cfg.cache_reuse_mode, + "align_document_chunks": cfg.align_document_chunks, + } + @app.post("/infer") async def infer(req: InferRequest) -> dict[str, Any]: """Run inference on a chat-template prompt with cached documents.""" @@ -687,11 +829,29 @@ async def infer(req: InferRequest) -> dict[str, Any]: ) ) - cache_hits: list[dict[str, Any]] = [] - if req.use_kv_cache and req.trace_cache: - cache_hits = [ - chunk.to_dict() for chunk in await core.lookup(prompt_tokens, cfg.model) + observed_cache_hits: list[dict[str, Any]] = [] + if req.use_kv_cache and (req.trace_cache or metrics_collector is not None): + observed_cache_hits = [ + chunk.to_dict() + for chunk in await core.lookup( + prompt_tokens, + cfg.model, + record_access=False, + ) ] + cache_hits = observed_cache_hits if req.trace_cache else [] + + if metrics_collector is not None and req.use_kv_cache: + hit = len(observed_cache_hits) > 0 + metrics_collector.record_lookup_result(hit) + hit_chunk_keys = { + str(hit_entry.get("chunk_key", "")) for hit_entry in observed_cache_hits + } + for doc in docs: + if any(key in hit_chunk_keys for key in doc.chunk_keys): + metrics_collector.record_lookup_hit(doc.doc_id) + else: + metrics_collector.record_lookup_miss(doc.doc_id) # Tell the connector to skip persisting this request's KV. The # /infer prompt is system + doc tokens + task; doc chunks are @@ -715,6 +875,9 @@ async def infer(req: InferRequest) -> dict[str, Any]: ) from exc elapsed_ms = (time.time() - t0) * 1000 + if metrics_collector is not None: + metrics_collector.record_inference(ttft_ms, elapsed_ms) + text = "" if result.get("choices"): text = result["choices"][0].get("text", "") diff --git a/daser/server/http/metrics.py b/daser/server/http/metrics.py new file mode 100644 index 0000000..a154db0 --- /dev/null +++ b/daser/server/http/metrics.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prometheus metrics collector for DaseR Stage A observability. + +The ``MetricsCollector`` owns counters, gauges, and histograms registered +with the ``prometheus_client`` library. It is injected into ``build_http_app`` +as an optional parameter so existing tests and deployments that do not enable +metrics remain unaffected. +""" + +# Standard +from typing import TYPE_CHECKING + +# Third Party +from prometheus_client import ( + CollectorRegistry, + Counter, + Gauge, + Histogram, + generate_latest, +) + +if TYPE_CHECKING: + from daser.server.core import ServerCore + + +class MetricsCollector: + """Prometheus metrics collector for the DaseR HTTP server. + + Each collector instance owns an independent ``CollectorRegistry`` so + multiple inspectors or test fixtures can co-exist without duplicate + metric errors. + + Args: + enabled: when False, ``export_metrics`` and route hooks are no-ops. + + Async/thread-safety: + Each instance uses its own ``CollectorRegistry``. Prometheus metrics + use internal locks and are safe for concurrent calls from FastAPI + endpoints. + """ + + def __init__(self, enabled: bool = True) -> None: + self._enabled = enabled + reg = CollectorRegistry(auto_describe=True) + + histogram_buckets = ( + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + float("inf"), + ) + + self.http_requests_total = Counter( + "daser_http_requests_total", + "Total HTTP requests served", + ["method", "endpoint", "status"], + registry=reg, + ) + self.chunk_hits_total = Counter( + "daser_chunk_hits_total", + "Chunk cache hits by document", + ["doc_id"], + registry=reg, + ) + self.chunk_misses_total = Counter( + "daser_chunk_misses_total", + "Chunk cache misses by document", + ["doc_id"], + registry=reg, + ) + self.lookup_total = Counter( + "daser_lookup_total", + "Server-side cache lookups by result", + ["status"], + registry=reg, + ) + self.ttft_seconds = Histogram( + "daser_ttft_seconds", + "Time-to-first-token for inference requests", + buckets=histogram_buckets, + registry=reg, + ) + self.inference_latency_seconds = Histogram( + "daser_inference_latency_seconds", + "End-to-end inference wall latency", + buckets=histogram_buckets, + registry=reg, + ) + self.document_prefill_seconds = Histogram( + "daser_document_prefill_seconds", + "Document upload prefill latency", + buckets=histogram_buckets, + registry=reg, + ) + self.cache_chunks = Gauge( + "daser_cache_chunks", + "Number of chunks currently stored in the ring buffer", + registry=reg, + ) + self.cache_free_slots = Gauge( + "daser_cache_free_slots", + "Number of free slots in the ring buffer", + registry=reg, + ) + self.cache_total_slots = Gauge( + "daser_cache_total_slots", + "Total slot capacity of the ring buffer", + registry=reg, + ) + self.eviction_total = Counter( + "daser_eviction_total", + "Chunk evictions by reason", + ["reason"], + registry=reg, + ) + self.documents_total = Gauge( + "daser_documents_total", + "Number of registered documents", + registry=reg, + ) + + self._registry = reg + self._seen_eviction_counts: dict[str, int] = {} + + @property + def enabled(self) -> bool: + """Return whether metrics collection is active. + + Returns: + ``True`` when route hooks should record metrics. + + Async/thread-safety: + Reads immutable configuration. + """ + return self._enabled + + def export_metrics(self) -> bytes: + """Return the current Prometheus text-format metrics payload. + + Returns: + Bytes suitable for a ``text/plain; version=0.0.4`` response body. + When the collector is disabled the payload contains only a disabled + marker comment. + + Async/thread-safety: + Delegates to prometheus_client, which guards registry collection + with its own locks. + """ + if not self._enabled: + return b"# metrics disabled\n" + return generate_latest(self._registry) + + def state_snapshot(self, core: "ServerCore") -> None: + """Sync point-in-time gauges from the shared server core. + + Args: + core: active ``ServerCore`` instance whose ``ChunkManager`` and + ``DocRegistry`` provide the current cache and document state. + + Returns: + None. + + Async/thread-safety: + Reads in-memory state on the FastAPI event loop. The same + event-loop ownership used by ``ServerCore`` guarantees exclusive + access. + """ + if not self._enabled: + return + cm = core.chunk_manager + self.cache_chunks.set(len(cm.store)) + self.cache_free_slots.set(cm.free_slots) + self.cache_total_slots.set(cm.total_slots) + registry = getattr(cm, "doc_registry", None) + if registry is not None: + self.documents_total.set(len(registry)) + else: + self.documents_total.set(0) + for reason, count in core.eviction_stats().items(): + previous = self._seen_eviction_counts.get(reason, 0) + if count > previous: + self.eviction_total.labels(reason=reason).inc(count - previous) + self._seen_eviction_counts[reason] = count + + def record_lookup_hit(self, doc_id: str) -> None: + """Record a chunk cache hit for one document. + + Args: + doc_id: document identifier that owns the cached chunk. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client counter guarded by its internal lock. + """ + if not self._enabled: + return + self.chunk_hits_total.labels(doc_id=doc_id).inc() + + def record_lookup_miss(self, doc_id: str) -> None: + """Record a chunk cache miss for one document. + + Args: + doc_id: document identifier referenced by the lookup. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client counter guarded by its internal lock. + """ + if not self._enabled: + return + self.chunk_misses_total.labels(doc_id=doc_id).inc() + + def record_http_request(self, method: str, endpoint: str, status: int) -> None: + """Record one HTTP request. + + Args: + method: HTTP method name. + endpoint: route template, for example ``"/documents/{doc_id}"``. + status: HTTP response status code. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client counter guarded by its internal lock. + """ + if not self._enabled: + return + self.http_requests_total.labels( + method=method, + endpoint=endpoint, + status=str(status), + ).inc() + + def record_eviction(self, reason: str) -> None: + """Record a chunk eviction. + + Args: + reason: human-readable cause, for example ``"ring"`` or + ``"document_delete"`` in Stage A. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client counter guarded by its internal lock. + """ + if not self._enabled: + return + self.eviction_total.labels(reason=reason).inc() + + def record_lookup_result(self, hit: bool) -> None: + """Record a server-side cache lookup result. + + Args: + hit: True when at least one chunk matched the lookup prompt. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client counter guarded by its internal lock. + """ + if not self._enabled: + return + status = "hit" if hit else "miss" + self.lookup_total.labels(status=status).inc() + + def record_inference(self, ttft_ms: float, latency_ms: float) -> None: + """Record inference latency for a single completion request. + + Args: + ttft_ms: time-to-first-token in milliseconds. + latency_ms: end-to-end wall latency in milliseconds. + + Returns: + None. + + Async/thread-safety: + Updates prometheus_client histograms guarded by internal locks. + """ + if not self._enabled: + return + self.ttft_seconds.observe(ttft_ms / 1000.0) + self.inference_latency_seconds.observe(latency_ms / 1000.0) + + def record_document_prefill(self, prefill_ms: float) -> None: + """Record document upload prefill latency. + + Args: + prefill_ms: time spent in prefill during upload, in milliseconds. + + Returns: + None. + + Async/thread-safety: + Updates a prometheus_client histogram guarded by its internal lock. + """ + if not self._enabled: + return + self.document_prefill_seconds.observe(prefill_ms / 1000.0) diff --git a/daser/server/metadata_store.py b/daser/server/metadata_store.py index dc50e1a..d38a8b1 100644 --- a/daser/server/metadata_store.py +++ b/daser/server/metadata_store.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # Standard -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, fields import time -from typing import Literal, Optional +from typing import Any, Literal, Optional # Third Party import msgpack @@ -29,6 +29,11 @@ class ChunkMeta: doc_ids: list of doc_ids that reference this chunk (empty when the chunk belongs to no registered document). Serves as the back-pointer used by cascading eviction. + access_count: number of times this chunk has been served from cache. + Updated by ``MetadataStore.touch`` on each lookup hit. Used by + later production stages (LFU/LRU hybrid eviction, metrics). + last_access_time: unix timestamp of the most recent ``touch`` call, + or ``created_at`` if the chunk has never been accessed. """ chunk_key: str @@ -39,10 +44,37 @@ class ChunkMeta: model_id: str created_at: float = 0.0 doc_ids: list[str] = field(default_factory=list) + access_count: int = 0 + last_access_time: float = 0.0 def __post_init__(self) -> None: if self.created_at == 0.0: self.created_at = time.time() + if self.last_access_time == 0.0: + self.last_access_time = self.created_at + + +_CHUNK_META_FIELD_NAMES: frozenset[str] = frozenset(f.name for f in fields(ChunkMeta)) + + +def _chunk_meta_from_payload(payload: dict[str, Any]) -> ChunkMeta: + """Build a ChunkMeta from a deserialized dict, dropping unknown keys. + + Args: + payload: dict from msgpack, possibly written by a newer schema + that added fields this build does not know about. + + Returns: + A ChunkMeta populated from the recognized subset of ``payload``. + Missing fields fall back to dataclass defaults so older on-disk + records remain readable. + + Async/thread-safety: + Pure function; safe to call from any thread. + """ + return ChunkMeta( + **{k: v for k, v in payload.items() if k in _CHUNK_META_FIELD_NAMES} + ) @dataclass @@ -152,6 +184,28 @@ def remove(self, chunk_key: str) -> None: del self._chunk_index[chunk_key] logger.debug("[INDEX] remove chunk_key=%s", chunk_key) + def touch(self, chunk_key: str, now: Optional[float] = None) -> None: + """Record one access for a stored chunk. + + Increments ``access_count`` and updates ``last_access_time`` on the + ChunkMeta for ``chunk_key``. Used on cache hits so later eviction + policies and observability metrics have per-chunk access stats. + + Args: + chunk_key: key of the chunk being accessed. + now: optional unix timestamp override (useful for tests). When + ``None``, ``time.time()`` is used. + + Async/thread-safety: + In-memory mutation; expected to run on the server event loop + together with other ``MetadataStore`` operations. + """ + meta = self._chunk_index.get(chunk_key) + if meta is None: + return + meta.access_count += 1 + meta.last_access_time = time.time() if now is None else now + # ------------------------------------------------------------------ # Query # ------------------------------------------------------------------ @@ -227,7 +281,7 @@ def load(self, path: str) -> None: self._total_slots = payload["total_slots"] self._chunk_index = { - k: ChunkMeta(**v) for k, v in payload["chunk_index"].items() + k: _chunk_meta_from_payload(v) for k, v in payload["chunk_index"].items() } self._slot_map = [ SlotEntry( diff --git a/docs/optimizations/5_multi_doc_production.md b/docs/optimizations/5_multi_doc_production.md new file mode 100644 index 0000000..cd6cf22 --- /dev/null +++ b/docs/optimizations/5_multi_doc_production.md @@ -0,0 +1,516 @@ +# Multi-Document Chunk Reuse Production Plan + +**Date:** 2026-05-28 +**Target:** turn user-managed multi-document chunk reuse into a production +service with SLA-grade cache hit ratio and bounded accuracy drift +**Scope:** `daser/server/`, `daser/connector/`, `examples/`, `docs/` + +This document consolidates the project positioning, SLA targets, connector +capability boundary, staged roadmap, data model extensions, CLI flags, test +matrix, and rollback plan. No vLLM source changes are required at any stage. + +## 1. Positioning + +### 1.1 Value Proposition + +DaseR is a KV cache service for **user-managed multi-document RAG**: + +- **User contract:** the client passes explicit `doc_ids` on `/infer`. DaseR + does not run a retrieval / semantic-similarity stage on its own. +- **Primary advantages:** cross-process / cross-restart persistence on NVMe, + multi-document chunk concatenation, predictable cache hit ratio, and bounded + seam-precision repair. +- **Primary metric:** **Cache Hit Ratio**, which directly determines TTFT, + GPU prefill cost, and the user-perceived latency floor. +- **Commitment:** doc-sets that the user declares as frequently used should + stay pinned in cache, achieve close to full hit ratio, and remain observable + at the production-SLA level. + +### 1.2 Differentiation + +| Capability | vLLM Prefix Cache | SGLang Radix | DaseR | +|-------------------------------------------|:-----------------:|:------------:|:---------:| +| Cross-process / restart persistence | no | no | yes (NVMe) | +| Multi-document concatenation (non-prefix) | no | partial | yes | +| User-explicit doc-set registration | no | no | yes | +| Hierarchical caching (L1 / L2 / L3) | no | no | yes | +| Seam-precision repair | no | no | yes | +| Production observability and SLA | basic | basic | yes | + +No open-source competitor combines explicit doc-set management with persistent +multi-document KV cache reuse. + +## 2. SLA Targets + +| Indicator | Definition | Target | +|----------------------------|--------------------------------------------------------|----------------:| +| Chunk Hit Ratio | `chunk_hits / chunks_requested` | >= 95% | +| DocSet Hit Ratio | `docset_full_hits / total_requests` (registered sets) | >= 80% | +| TTFT P50 (warm) | warm-hit TTFT median | <= 50 ms | +| TTFT P99 (warm) | warm-hit TTFT 99th percentile | <= 200 ms | +| Cold TTFT (with prefetch) | first request in a new session after prefetch | <= 200 ms | +| Accuracy Drift | chunk reuse vs full prefill task-level accuracy loss | <= 2% | +| Pinned Eviction Rate | pinned chunk evictions per 1000 requests | 0 | +| Prefetch Lead Time | session start to prefetch completion | <= 3 s | + +## 3. Connector Capability Boundary + +### 3.1 What the connector can do + +| Capability | Hook | Note | +|-----------------------------------------------|--------------------------------------------------|-----------------------------------------------------------------| +| Set external-token count | `get_num_new_matched_tokens` returns one int | N tokens starting from `num_computed_tokens`, must be contiguous | +| Skip load / save per request | `daser_skip_load` / `daser_skip_save` | Binary request-level flags | +| Mutate KV at load time | `_transform_loaded_staging_batch` | K/V scaling, RoPE delta, sink correction | +| Extract KV at save time | `_stage_store_batch` plus `wait_for_save` | Whole-chunk store keyed by `chunk_key` | +| Per-chunk metadata | `ChunkMeta` | Extensible | +| Multi-chunk assembly | `ReqLoadSpec.chunks` | Must be head-to-tail contiguous, no gap | +| Cross-request reuse | `chunk_key` index | Already implemented | + +### 3.2 What the connector cannot do + +- Read attention scores (QK^T or attention weights). +- Drive per-layer selective load / store. +- Express non-contiguous external supply (skip-head, reuse-middle, skip-tail) + inside one request. +- Modify attention computation (temperature scaling, in-kernel masking, etc.). +- Share the vLLM GPU KV block allocation across requests. + +### 3.3 Project constraints (CLAUDE.md) + +| Constraint | Plan landing point | +|---------------------------------------------|-------------------------------------------------------------------| +| Control plane in server | DocSet manager, pin policy, calibration, metrics in `daser/server/` | +| Data plane in connector | Sink correction in worker; scheduler only routes | +| Cross-layer access via ABC / IPC only | Scheduler reads new fields through IPC `lookup`, not server state | +| No vLLM source changes | All code lives in `daser/`, `examples/`, `docs/` | +| All I/O is asyncio | Prefetch and preheat run as server async tasks | +| Transfer mode fixed at startup | Independent of this plan | + +## 4. Architecture Overview + +```mermaid +flowchart TB + subgraph CALIB["Offline calibration (Stage D)"] + HF["HF Transformers
output_attentions=True"] + SIGMA["compute chunk sigma
and docset sigma"] + CAL[("calibration.json")] + HF --> SIGMA --> CAL + end + + subgraph MGMT["Management plane (Stage B)"] + DSAPI["/docsets API
user-declared doc combinations"] + PREHEAT["DocSet preheat scheduler
off-peak full-prefill into L2/L3"] + DSAPI --> PREHEAT + end + + subgraph UPLOAD["Upload path · POST /documents (Stage D injection)"] + DOCAPI["FastAPI /documents"] + CHUNK["chunker -> tokens"] + VLLMPRE["vLLM prefill (max_tokens=1)"] + SAVE["Worker._stage_store_batch
+ compute sink_stat"] + IPCSTORE["IPC store + sink_stat"] + SVRREG["Server.register_chunk
ChunkMeta += sigma, sink_stat, is_real_sink"] + DOCAPI --> CHUNK --> VLLMPRE --> SAVE --> IPCSTORE --> SVRREG + CAL -. sigma lookup .-> SVRREG + end + + subgraph STORE["Storage (data plane)"] + L1[("L1 single-chunk KV
ring buffer + pinned memory")] + L2[("L2 doc-pair seam KV")] + L3[("L3 full-prefix KV per docset")] + end + + subgraph INFER["Inference path · POST /infer (Stage A/C/D)"] + INFAPI["FastAPI /infer"] + SESSION["Session prefetch (Stage C)
async pull chunks into L1"] + LOOKUP["Server.lookup
L3 -> L2 -> L1 priority"] + ROUTE{"hit level + sigma"} + L3HIT["L3 hit:
accuracy lossless"] + L2HIT["L2 hit:
seam precomputed"] + L1FIX["L1 only + sigma <= sigma_high:
accuracy safety net (Stage D)"] + L1STD["L1 only + sigma > sigma_high:
standard chunk reuse"] + WSTART["Worker.start_load_kv"] + XFORM["_transform_loaded_staging_batch
sink drift correction"] + VLLM[(vLLM KV cache)] + + INFAPI --> SESSION + INFAPI --> LOOKUP + LOOKUP --> STORE + LOOKUP --> ROUTE + ROUTE --> L3HIT & L2HIT & L1FIX & L1STD + L3HIT --> WSTART + L2HIT --> WSTART + L1FIX --> WSTART + L1STD --> WSTART + SESSION -. async pull .-> STORE + WSTART --> XFORM --> VLLM + end + + subgraph OBS["Observability (Stage A)"] + METRICS["/metrics · Prometheus"] + DIAG["/diag/explain · route explanation"] + end + + MGMT -. ChunkMeta + pin .-> STORE + UPLOAD -. ChunkMeta .-> STORE + INFER -.-> OBS +``` + +## 5. Data Model Extensions + +All new fields are `Optional` or have safe defaults so msgpack records remain +backward and forward compatible across stages. + +### 5.1 `ChunkMeta` (`daser/server/metadata_store.py`) + +```python +@dataclass +class ChunkMeta: + # existing + chunk_key: str + start_slot: int + num_slots: int + token_count: int + pos_offset: int + model_id: str + created_at: float + doc_ids: list[str] + + # Stage A: observability and access tracking + access_count: int = 0 + last_access_time: float = 0.0 + + # Stage B: pinning policy + pinned: bool = False + docset_names: list[str] = field(default_factory=list) + + # Stage C: hierarchical cache + cache_level: Literal["L1", "L2", "L3"] = "L1" + + # Stage D: accuracy safety net + self_contained_score: Optional[float] = None # sigma in [0, 1] + sink_stat_bytes: Optional[bytes] = None # bf16 tensor [layers, kv_heads, head_dim] + sink_k_fs: int = 0 # 0 disables drift correction + is_real_sink: bool = False # skip drift correction when True +``` + +### 5.2 `DocSetMeta` (`daser/server/docset/`, new package) + +```python +@dataclass +class DocSetMeta: + name: str + doc_ids: list[str] + chunk_keys: list[str] # ordered chunk sequence + pinned: bool = True # registration pins by default + docset_sigma: Optional[float] = None # doc-set level sigma + l2_chunk_keys: list[str] = field(default_factory=list) # doc-pair seam KV + l3_chunk_key: Optional[str] = None # full-prefix KV + created_at: float = 0.0 + last_used_at: float = 0.0 +``` + +### 5.3 `ChunkInfo` (`daser/server/core.py`) + +The IPC `lookup` payload gains the same new fields. Older connectors ignore +unknown keys, which keeps the rollout out of lockstep. + +### 5.4 `ReqLoadSpec` (`daser/connector/metadata.py`) + +```python +@dataclass +class ReqLoadSpec: + # existing + ... + + cache_level: Literal["L1", "L2", "L3"] = "L1" + sink_stat_bytes: Optional[bytes] = None + sink_k_fs: int = 0 + is_real_sink: bool = False +``` + +## 6. Staged Roadmap + +| Stage | Span | Theme | Direct goal | Main subparts | +|:-----:|:-------:|----------------------------------------|------------------------------------------------------------|---------------------------------------------------------------| +| A | 1-2 wk | Production observability | Quantify hit ratio, TTFT, accuracy drift | A.1 metrics, A.2 diag, A.3 baseline measurement | +| B | 2-3 wk | DocSet management and pinning | Registered doc-set hit ratio >= 80% | B.1 `/docsets` API, B.2 preheat, B.3 LFU+LRU, B.4 pin | +| C | 3-4 wk | Hierarchical cache and session prefetch | Warm TTFT P50 <= 50 ms | C.1 L2, C.2 L3, C.3 session, C.4 async load | +| D | 4-6 wk | Accuracy safety net | Accuracy loss on L1 fallback path <= 2% | D.1 sigma, D.2 sink_stat, D.3 routing, D.4 trim, D.5 drift, D.6 real-sink | +| E | ongoing | Long-term evolution | Approach full-prefill accuracy | E.1 multi-version, E.2 distillation, E.3 online calibration, E.4 pair affinity, E.5 overlap chunking | + +### Stage A · Production Observability + +Stage A is a prerequisite for every later stage: without baseline numbers, +nothing else can prove a regression-free improvement. + +Subparts: + +- **A.1 `/metrics` Prometheus endpoint.** Exposes: + + | Metric | Labels | + |-----------------------------------|----------------------------------------| + | `daser_chunk_hits_total` | doc_id | + | `daser_chunk_misses_total` | doc_id | + | `daser_docset_hit_ratio` | docset_name | + | `daser_cache_level_hits_total` | level={L1, L2, L3} | + | `daser_ttft_seconds` | mode={baseline, chunk, docset} histogram | + | `daser_eviction_total` | reason={ring, explicit, document_delete} | + | `daser_l1_l2_l3_size_bytes` | level | + | `daser_accuracy_drift_estimate` | docset_name | + +- **A.2 `/diag/explain` endpoint.** Given a prompt and a `doc_ids` list, returns + the routing decision, hit tier, expected TTFT, and accuracy-net flag. +- **A.3 Baseline measurement script** under `examples/baseline_measure/` to + collect per-doc and per-docset hit-ratio and TTFT distribution over a week + of production traffic. + +Landing points: + +- `daser/server/http/metrics.py` (new) +- `daser/server/http/diag.py` (new) +- `daser/server/core.py` access counters + +Rollback: disabling metrics only loses observability; no behavioral change. + +### Stage B · DocSet Management and Pinning + +Goal: let the user declare frequently used doc-sets and have the system +guarantee a hit-ratio SLA on them. + +Subparts: + +- **B.1 `/docsets` API.** + + ``` + POST /docsets register { name, doc_ids } + GET /docsets list registrations + DELETE /docsets/{name} unregister (un-pin) + POST /docsets/{name}/preheat trigger preheat manually + ``` + +- **B.2 Preheat scheduler.** After registration, an async task runs a full + prefill path for each doc-set during off-peak windows, pulls all chunks into + L1, and pins them. +- **B.3 LFU + LRU hybrid eviction.** Replaces the current LRU-only policy: + + ``` + evict_score(chunk) = alpha * (1 - normalized_access_count) + + (1 - alpha) * time_since_last_access + ``` + + Default `alpha = 0.7` weighs LFU higher to keep hot chunks stable. + +- **B.4 Pin mechanism.** Chunks with `pinned=True` are never evicted, even + when L1 is full; only non-pinned chunks are eligible. Doc-set registration + pins all member chunks; unregistration unpins. + +Landing points: + +- `daser/server/docset/` (new package) +- `daser/server/chunk_manager.py` eviction policy +- `daser/server/http/app.py` route +- `daser/server/metadata_store.py` field extensions + +Rollback: `--enable-docsets=false` returns 501 from `/docsets` and falls back +to current behavior. + +### Stage C · Hierarchical Cache and Session Prefetch + +Goal: drive warm TTFT down to cache-read magnitude and remove cold-start spikes. + +Subparts: + +- **C.1 L2 cache (doc-pair seam KV).** On doc-set registration, identify + adjacent doc pairs, run one full-prefill that materializes the seam KV + segment (`doc_a + doc_b`), and store the seam segment as an L2 chunk. On hit, + the L1 chunks plus the L2 seam chunk concatenate cleanly without asking + vLLM to re-prefill the boundary. +- **C.2 L3 cache (full-prefix KV).** One full-prefill over the entire + registered doc-set, stored as a single chunk. On hit, no chunk assembly + or seam prefill is needed. **Accuracy is lossless.** +- **C.3 Session state and prefetch.** HTTP header `X-Session-Id` marks the + session. The first request triggers `prefetch_to_l1(doc_ids)` as an async + task. Subsequent requests in the same session hit L1 almost entirely. + Sessions expire after 10 minutes. +- **C.4 Async load.** Integrates with issue #42 so chunk reads overlap with + vLLM scheduling instead of blocking the connector load path. + +Landing points: + +- `daser/server/cache/levels.py` (new): L2 / L3 indexing and lookup priority +- `daser/server/prefetch.py` (new): async prefetch scheduler +- `daser/connector/worker.py`: async prefetch hooks (shared with #42) + +Rollback: each subpart has an independent flag. L2 / L3 disabled falls back to +L1 chunk reuse; session prefetch disabled falls back to synchronous load. + +### Stage D · Accuracy Safety Net + +**Critical scope:** Stage D applies only on the **L1 fallback path**. L3 and +L2 hits are accuracy-lossless by construction and need no repair. + +**Routing change:** because `doc_ids` are user-supplied, Stage D removes the +"skip reuse" branch entirely. Routing decides only "repair vs not": + +- sigma > sigma_high: standard chunk reuse (no repair). +- sigma <= sigma_high: trigger boundary trim and sink drift correction. + +Subparts: + +- **D.1 Offline sigma calibration.** Lives in `daser/server/calibration/`. Runs + HF Transformers with `output_attentions=True` and computes chunk-level and + doc-set-level sigma, writing `calibration.json`. CLI: + `python -m daser.server.calibration --model --corpus --out `. +- **D.2 Online sink_stat computation.** Worker `_stage_store_batch` averages + the K of the first `k_fs` tokens of each chunk per layer and submits the + result alongside the IPC store request. **Math note:** RoPE's fixed-delta + rotation is a linear operator, so `RoPE_delta(mean_j K[j]) = mean_j(RoPE_delta(K[j]))`. + This means the cache-time-domain mean rotated by the same delta at load + time equals the load-time-domain mean. No NoPE inverse rotation is needed. +- **D.3 Routing decision (scheduler).** Aggregates doc-set sigma falling back + to per-chunk `min(sigma)` to decide whether to engage the repair path. +- **D.4 Boundary trim (scheduler).** Cedes the first `k_head` and last + `k_tail` tokens of the loaded prefix to vLLM prefill. Only whole-prefix + single-sided trim is feasible: `_contiguous_prefix_tokens` forbids gaps, + so per-chunk two-sided trim is **not** achievable at this layer. +- **D.5 Sink drift correction (worker).** Inside `_transform_loaded_staging_batch`, + after the existing RoPE delta rotation, apply `K[:k_fs] -= lambda * RoPE_delta(sink_stat)` + for non-real-sink chunks. +- **D.6 Real-Sink marker.** System-prompt chunks set `is_real_sink=True` so + D.5 skips them; they are the true attention sink and must remain unchanged. + +Landing points: + +- `daser/server/calibration/` (new) +- `daser/connector/worker.py::_stage_store_batch` +- `daser/connector/scheduler.py::get_num_new_matched_tokens` +- `daser/connector/staging.py::_transform_loaded_staging_batch` +- `daser/server/http/app.py`: `/documents` accepts an `is_real_sink` flag + +Rollback: `--enable-cross-attn-fix=false` short-circuits all D subparts. + +### Stage E · Long-Term Evolution + +| Id | Theme | Replaces / extends | Note | +|:---:|--------------------------------------|----------------------------|----------------------------------------------------------------------| +| E.1 | Context-conditioned multi-version | extends Stage D `chunk_key` | Cache different versions of the same chunk keyed by upstream-context hash | +| E.2 | KV residual distillation | replaces D.5 | Learn a small residual correction model with better accuracy than heuristic sink reduction | +| E.3 | Online streaming calibration | replaces D.1 | Sample shadow full-prefills at inference time and update sigma via EMA | +| E.4 | Chunk-pair affinity | strengthens D.3 | Use doc-pair co-occurrence statistics in routing | +| E.5 | Overlap-aware chunking | replaces D.4 | Chunker emits overlapping chunks so vLLM trim is not needed | + +E items are research bets, not committed scope. + +## 7. CLI Configuration + +`python -m daser.server` gains the following flags (all default off; the +service is byte-compatible with master when no new flag is set). + +| Flag | Default | Stage | Purpose | +|--------------------------------|--------------------------------------|:-----:|--------------------------------------------------| +| `--enable-docsets` | `false` | B | Enable the `/docsets` API | +| `--enable-cache-l2` | `false` | C | Enable doc-pair seam (L2) cache | +| `--enable-cache-l3` | `false` | C | Enable full-prefix (L3) cache | +| `--enable-session-prefetch` | `false` | C | Enable session-level prefetch | +| `--enable-cross-attn-fix` | `false` | D | Enable accuracy safety net | +| `--session-timeout-seconds` | `600` | C | Session-state expiry | +| `--lfu-weight` | `0.7` | B | LFU weight in the hybrid eviction score | +| `--cross-attn-calibration` | `${store-dir}/calibration.json` | D | Calibration file path | +| `--cross-attn-sigma-default` | `0.7` | D | Fallback sigma when calibration is missing | +| `--cross-attn-sigma-high` | `0.8` | D | High self-containment threshold | +| `--cross-attn-k-head` | `8` | D | Whole-prefix head trim length | +| `--cross-attn-k-tail` | `4` | D | Whole-prefix tail trim length | +| `--cross-attn-sink-k-fs` | `4` | D | Per-chunk head sink window | +| `--cross-attn-sink-lambda` | `0.4` | D | Drift correction coefficient | +| `--enable-metrics` | `false` | A | Serve Prometheus metrics at `/metrics` on the main HTTP server | + +## 8. Test Matrix + +### 8.1 Unit tests + +| Module | Cases | +|-------------------------------------------------|--------------------------------------------------------------------| +| `metadata_store` | Old/new ChunkMeta msgpack roundtrip; DocSetMeta serialization | +| `chunk_manager` | LFU+LRU eviction ordering; pinned chunks resist eviction | +| `docset` | Register / unregister / preheat lifecycle; concurrent registration | +| `cache.levels` | L3 -> L2 -> L1 lookup priority | +| `prefetch` | Session state management; timeout cleanup | +| `calibration` | Synthetic attention matrix -> known sigma | +| `worker._stage_store_batch` | sink_stat numerics correct; `is_real_sink` skips computation | +| `scheduler` | Sigma binary routing; trim leaves chunks contiguous; min(sigma) aggregation | +| `staging._transform_loaded_staging_batch` | Behavior preserved when sink_stat=None or lambda=0; numeric oracle | + +### 8.2 Integration tests + +| Scenario | Expected | +|----------------------------------------------------------|---------------------------------------------------------| +| All stage flags off | Byte-identical to master | +| Doc-set registered and preheated, then served | L3 hit, TTFT visibly lower | +| Repeated requests inside one session | First triggers prefetch; later requests are L1 hits | +| Pinned chunks under L1 pressure | Non-pinned chunks evicted first | +| L1 fallback with sigma <= sigma_high | Triggers trim and drift correction | +| Real-Sink chunk | K[:k_fs] matches the no-correction baseline | + +### 8.3 Performance and accuracy regression + +- **TTFT baseline:** master vs each cumulative stage, reporting P50 / P99. +- **Hit ratio:** master vs Stage B and C cumulatively. +- **Accuracy:** CMRC2018 and a LongBench subset, all-on vs all-off vs + full-prefill baseline; loss <= 2%. +- **Observability:** 100% metric coverage of new code paths; alert + thresholds verified. + +## 9. Risks and Rollback + +| Risk | Mitigation | +|---------------------------------------------------------------|--------------------------------------------------------------------------------------------------| +| Too many registered doc-sets exceed L1 capacity | `/docsets` registration checks `chunk_count * slot_size < l1_capacity * 0.8` | +| Session prefetch competes with foreground requests for GPU | Prefetch runs outside `--gpu-memory-utilization`; paused under load | +| L3 invalidation when a doc in a doc-set is updated | Doc-update events cascade-invalidate dependent L3 entries; next request re-preheats | +| Calibration drift vs production model weights | Sigma is a ratio of attention shares; model fingerprint stored alongside calibration | +| Drift correction degrades decode-time KV | End-to-end accuracy regression is the safety net; lambda default is conservative | +| Boundary trim splits a chunk into non block-aligned residue | Scheduler double-checks block alignment after trim; falls back to full reuse if misaligned | +| Old connector talks to new server or vice versa | All new fields are Optional with msgpack defaults; CI covers cross-version compatibility | + +**Global rollback:** every stage has an independent flag. Setting +`--enable-docsets=false`, `--enable-cache-l2=false`, `--enable-cache-l3=false`, +`--enable-session-prefetch=false`, and `--enable-cross-attn-fix=false` +returns DaseR to byte-compatible behavior with master. + +## 10. Out of Scope (Research Items) + +These mechanisms require vLLM source changes and are tracked as future research: + +| Mechanism | Required vLLM change | +|------------------------------------------|----------------------------------------------------------------------------------------| +| Token-grain dynamic selective recompute | Export per-layer attention weights; scheduler supports mixed prefill + KV reuse | +| Per-chunk two-sided boundary recompute | `get_num_new_matched_tokens` accepts `list[(start, length)]` non-contiguous external supply | +| Per-layer recompute ratio adaptation | Per-layer load / store interface | +| In-kernel sink correction / APE rescale | PagedAttention kernel instrumentation | +| NoPE-format KV with fused RoPE | Attention kernel applies RoPE at attention time instead of at store time | + +## 11. PR Breakdown + +| PR | Stage | Content | Estimate | +|:---:|:-----:|--------------------------------------------------------------------------|---------:| +| #1 | A | ChunkMeta access_count + last_access_time + msgpack compatibility tests | ~200 lines | +| #2 | A | `/metrics` Prometheus endpoint | ~300 lines | +| #3 | A | `/diag/explain` endpoint + baseline measurement script | ~250 lines | +| #4 | B | DocSetMeta + `/docsets` API + register / unregister | ~400 lines | +| #5 | B | LFU+LRU hybrid eviction + pin mechanism | ~300 lines | +| #6 | B | DocSet preheat scheduler | ~300 lines | +| #7 | C | L2 cache (doc-pair seam KV) | ~400 lines | +| #8 | C | L3 cache (full-prefix KV) | ~300 lines | +| #9 | C | Session state + prefetch | ~400 lines | +| #10 | D | ChunkMeta accuracy fields + ReqLoadSpec extension | ~200 lines | +| #11 | D | Offline calibration CLI | ~400 lines | +| #12 | D | sink_stat computation + IPC store field passthrough | ~300 lines | +| #13 | D | Scheduler sigma routing + boundary trim | ~250 lines | +| #14 | D | Staging drift correction | ~200 lines | +| #15 | all | End-to-end integration tests + accuracy regression + documentation | ~600 lines | + +Each PR ships independently under its own flag and stays disabled by default, +so any single merge keeps master behavior byte-identical for callers who do +not opt in. diff --git a/examples/baseline_measure/measure.py b/examples/baseline_measure/measure.py new file mode 100644 index 0000000..42c4a8c --- /dev/null +++ b/examples/baseline_measure/measure.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Baseline measurement script for DaseR Stage A observability. + +Runs a sequence of document uploads and inference requests against a +running DaseR HTTP server and reports cache-hit ratio, TTFT, and latency +distributions. This script is standalone and not imported by DaseR. + +Usage:: + + python examples/baseline_measure/measure.py \\ + --service-url http://127.0.0.1:2026 \\ + --samples 20 + +Requirements: ``httpx`` (already a DaseR dependency). +""" + +# Standard +import argparse +import json +import statistics +import time +from typing import Any + + +def _upload(client: Any, service_url: str, title: str, text: str) -> dict: + """Upload a single document. + + Args: + client: ``httpx.Client`` connected to the DaseR server. + service_url: base URL of the DaseR HTTP API. + title: display title for the document. + text: raw document text. + + Returns: + The JSON response body from ``POST /documents``. + + Raises: + RuntimeError: if the upload fails. + """ + resp = client.post( + f"{service_url}/documents", + json={"title": title, "text": text}, + timeout=600.0, + ) + if resp.status_code != 201: + raise RuntimeError(f"upload failed: {resp.status_code} {resp.text}") + return resp.json() + + +def _infer( + client: Any, + service_url: str, + doc_ids: list[str], + task: str, + use_kv_cache: bool = True, +) -> tuple[dict, float]: + """Run one inference request. + + Args: + client: ``httpx.Client`` connected to the DaseR server. + service_url: base URL of the DaseR HTTP API. + doc_ids: ordered document identifiers. + task: user task text. + use_kv_cache: when False, sets ``daser_skip_load`` to bypass KV cache. + + Returns: + ``(response_json, wall_seconds)``. + """ + t0 = time.time() + resp = client.post( + f"{service_url}/infer", + json={ + "doc_ids": doc_ids, + "task": task, + "use_kv_cache": use_kv_cache, + "trace_cache": True, + "gen_params": {"max_tokens": 80, "temperature": 0.0, "stop": ["\n\n"]}, + }, + timeout=600.0, + ) + elapsed = time.time() - t0 + if resp.status_code != 200: + raise RuntimeError(f"infer failed: {resp.status_code} {resp.text}") + return resp.json(), elapsed + + +def _run_baseline(args: argparse.Namespace) -> None: + """Upload two documents and compare KV-cached vs non-cached inference. + + Args: + args: parsed CLI arguments. + """ + import httpx + + client = httpx.Client(base_url=args.service_url, timeout=600.0) + + print(f"=== DaseR Baseline Measurement ({args.samples} samples) ===") + print(f" service: {args.service_url}") + + # Health check + try: + health = client.get("/health").json() + print(f" health: {json.dumps(health)}") + except Exception as exc: # noqa: BLE001 + print(f" health: UNREACHABLE ({exc})") + return + + # Upload documents + print("\n--- Upload ---") + doc_a = _upload( + client, + args.service_url, + "Document A", + "DaseR is a RAG-native KV cache service for large language model " + "inference. It integrates with vLLM through KVConnectorBase_V1 and " + "stores attention KV tensors on NVMe using NVIDIA cuFile (GDS) or " + "io_uring as a fallback. The NVMe ring buffer uses fixed-size slots " + "and chunk metadata for cache management.", + ) + doc_b = _upload( + client, + args.service_url, + "Document B", + "The DaseR HTTP RAG API handles document upload, listing, and " + "inference. Upload tokenizes text, creates block-aligned chunks, " + "prewarms each chunk through vLLM, commits chunk_keys in ServerCore, " + "and reuses existing ChunkMeta entries for duplicate documents.", + ) + print(f" doc_a: {doc_a['doc_id'][:8]}... ({doc_a['chunk_count']} chunks)") + print(f" doc_b: {doc_b['doc_id'][:8]}... ({doc_b['chunk_count']} chunks)") + + task = ( + "Write exactly two short sentences. " + "Sentence 1 must start with DaseR and summarize the cache service. " + "Sentence 2 must start with The HTTP RAG API and summarize the " + "upload/inference layer." + ) + + # Warm-up: one no-KV request to ensure vLLM is primed + print("\n--- Warm-up ---") + _infer( + client, + args.service_url, + [doc_a["doc_id"], doc_b["doc_id"]], + task, + use_kv_cache=False, + ) + print(" warm-up complete") + + # Collect samples + print(f"\n--- Sampling ({args.samples} rounds) ---") + kv_load_ttfts: list[float] = [] + kv_load_latencies: list[float] = [] + kv_load_hits: list[int] = [] + no_kv_ttfts: list[float] = [] + no_kv_latencies: list[float] = [] + + for i in range(args.samples): + # With KV cache + result, wall = _infer( + client, + args.service_url, + [doc_a["doc_id"], doc_b["doc_id"]], + task, + use_kv_cache=True, + ) + kv_load_ttfts.append(float(result.get("ttft_ms", 0.0))) + kv_load_latencies.append(float(result.get("latency_ms", 0.0))) + kv_load_hits.append(len(result.get("cache_hits", []))) + + # Without KV cache + result, wall = _infer( + client, + args.service_url, + [doc_a["doc_id"], doc_b["doc_id"]], + task, + use_kv_cache=False, + ) + no_kv_ttfts.append(float(result.get("ttft_ms", 0.0))) + no_kv_latencies.append(float(result.get("latency_ms", 0.0))) + + if (i + 1) % 5 == 0: + print(f" ... {i + 1}/{args.samples}") + + # Report + print("\n=== Results (ms) ===") + _print_row("KV-load TTFT", kv_load_ttfts) + _print_row("KV-load latency", kv_load_latencies) + _print_row("No-KV TTFT", no_kv_ttfts) + _print_row("No-KV latency", no_kv_latencies) + + avg_hits = statistics.mean(kv_load_hits) if kv_load_hits else 0.0 + print(f"\n Avg cache hits (KV-load): {avg_hits:.1f}") + if kv_load_hits: + hit_pct = sum(1 for h in kv_load_hits if h > 0) / len(kv_load_hits) * 100 + print(f" Hit rate: {hit_pct:.1f}%") + ttft_speedup = ( + statistics.mean(no_kv_ttfts) / statistics.mean(kv_load_ttfts) + if kv_load_ttfts + else 0.0 + ) + print(f" TTFT speedup: {ttft_speedup:.2f}x") + + +def _print_row(label: str, values: list[float]) -> None: + """Print a statistics row. + + Args: + label: metric name. + values: list of measured values in milliseconds. + """ + if not values: + print(f" {label:20s} (no data)") + return + sorted_vals = sorted(values) + p50 = sorted_vals[len(sorted_vals) // 2] + p99 = sorted_vals[min(len(sorted_vals) - 1, (len(sorted_vals) * 99) // 100)] + avg = statistics.mean(values) + std = statistics.stdev(values) if len(values) > 1 else 0.0 + print( + f" {label:20s} avg={avg:7.1f} p50={p50:7.1f} p99={p99:7.1f} std={std:6.1f}" + ) + + +def main() -> None: + """CLI entry point.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--service-url", default="http://127.0.0.1:2026") + parser.add_argument( + "--samples", + type=int, + default=20, + help="number of inference rounds (default: 20)", + ) + args = parser.parse_args() + _run_baseline(args) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 4ad2e65..d5ea034 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ license = { file = "LICENSE" } dependencies = [ "fastapi>=0.110", "httpx>=0.27", + "prometheus-client>=0.20", "msgpack>=1.0", "pydantic>=2.0", "transformers>=4.40", diff --git a/tests/server/test_core.py b/tests/server/test_core.py index ca5484e..f9d3f23 100644 --- a/tests/server/test_core.py +++ b/tests/server/test_core.py @@ -104,6 +104,67 @@ async def test_wait_for_committed_chunks_times_out() -> None: await core.wait_for_committed_chunks(["missing"], timeout_s=0.001) +@pytest.mark.asyncio +async def test_lookup_hit_updates_chunk_access_stats() -> None: + core = make_core() + tokens = [1, 2, 3, 4] + key = first_rolling_key(tokens) + + await core.alloc_chunk(key, token_count=len(tokens), model_id="m") + await core.commit_chunk(key) + + store = core.chunk_manager.store + meta_before = store.get(key) + assert meta_before is not None + assert meta_before.access_count == 0 + + await core.lookup(tokens, "m") + await core.lookup(tokens, "m") + + meta_after = store.get(key) + assert meta_after is not None + assert meta_after.access_count == 2 + assert meta_after.last_access_time >= meta_after.created_at + + +@pytest.mark.asyncio +async def test_lookup_can_skip_access_stats_for_diagnostics() -> None: + core = make_core() + tokens = [1, 2, 3, 4] + key = first_rolling_key(tokens) + + await core.alloc_chunk(key, token_count=len(tokens), model_id="m") + await core.commit_chunk(key) + + hits = await core.lookup(tokens, "m", record_access=False) + + meta = core.chunk_manager.store.get(key) + assert hits + assert meta is not None + assert meta.access_count == 0 + + +@pytest.mark.asyncio +async def test_lookup_miss_leaves_access_stats_untouched() -> None: + core = make_core() + tokens = [1, 2, 3, 4] + key = first_rolling_key(tokens) + + await core.alloc_chunk(key, token_count=len(tokens), model_id="m") + await core.commit_chunk(key) + + store = core.chunk_manager.store + baseline = store.get(key) + assert baseline is not None + baseline_count = baseline.access_count + + await core.lookup([99, 98, 97, 96], "m") + + after = store.get(key) + assert after is not None + assert after.access_count == baseline_count + + @pytest.mark.asyncio async def test_restored_orphan_committed_chunk_can_be_reused(tmp_path) -> None: tokens = [1, 2, 3, 4] diff --git a/tests/server/test_http_server.py b/tests/server/test_http_server.py index ffd1e6d..4fa8daa 100644 --- a/tests/server/test_http_server.py +++ b/tests/server/test_http_server.py @@ -17,6 +17,7 @@ from daser.server.core import ServerCore from daser.server.doc_registry import DocRegistry from daser.server.http import HTTPServerConfig, build_http_app +from daser.server.http.metrics import MetricsCollector from daser.server.metadata_store import MetadataStore SLOT_SIZE = 1024 @@ -950,3 +951,330 @@ def test_chunk_reuse_repeated_separator_keeps_hits_contiguous_before_suffix() -> ] assert hit_starts == [0, *hit_ends[:-1]] assert hit_ends[-1] == len(fake_vllm.completions[0][0]) - len(_chat_suffix()) + + +# ------------------------------------------------------------------ +# Stage A: /metrics +# ------------------------------------------------------------------ + + +def _make_client_with_metrics( + core: ServerCore, + vllm: FakeVLLMClient, + metrics: MetricsCollector | None = None, +) -> TestClient: + """Build a test client with a metrics collector attached.""" + return TestClient( + build_http_app( + HTTPServerConfig( + vllm_base_url="http://vllm", + model="m", + tokenizer="fake", + block_tokens=4, + system_prompt="S:", + doc_separator="|", + task_separator="? ", + answer_separator="! ", + ), + core, + tokenizer=FakeTokenizer(), + vllm=vllm, + metrics_collector=metrics, + ) + ) + + +def test_metrics_endpoint_returns_404_when_disabled() -> None: + """``/metrics`` returns 404 when no collector is configured.""" + core = make_core() + fake_vllm = FakeVLLMClient() + client = _make_client_with_metrics(core, fake_vllm, metrics=None) + + resp = client.get("/metrics") + + assert resp.status_code == 404 + assert "metrics disabled" in resp.json()["detail"] + + +def test_metrics_endpoint_returns_200_when_enabled() -> None: + """``/metrics`` returns Prometheus text when a collector is active.""" + core = make_core() + fake_vllm = FakeVLLMClient() + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + resp = client.get("/metrics") + + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/plain") + assert b"daser_cache_chunks" in resp.content + assert b"daser_cache_total_slots" in resp.content + + +def test_metrics_endpoint_disabled_collector_returns_404() -> None: + """``/metrics`` returns 404 when collector exists but is disabled.""" + core = make_core() + fake_vllm = FakeVLLMClient() + metrics = MetricsCollector(enabled=False) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + resp = client.get("/metrics") + + assert resp.status_code == 404 + + +def test_metrics_counter_increments_after_infer() -> None: + """Lookup and inference counters increment after a ``/infer`` request.""" + core = make_core() + fake_vllm = FakeVLLMClient(commit_core=core) + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + doc_a = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + client.post("/documents", json={"title": "b", "text": "efgh"}).json() + client.post( + "/infer", + json={ + "doc_ids": [doc_a["doc_id"]], + "task": "go", + "use_kv_cache": True, + "trace_cache": True, + }, + ) + + resp = client.get("/metrics") + + assert resp.status_code == 200 + # After inference the lookup_total counter should be present + assert b"daser_lookup_total" in resp.content + assert b"daser_ttft_seconds" in resp.content + assert b"daser_inference_latency_seconds" in resp.content + + +def test_metrics_lookup_does_not_depend_on_trace_cache() -> None: + """Metrics use read-only lookup even when response tracing is disabled.""" + metrics = MetricsCollector(enabled=True) + client, _, _ = _make_client_for_diag(metrics=metrics) + + doc_a = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + client.post( + "/infer", + json={ + "doc_ids": [doc_a["doc_id"]], + "task": "go", + "use_kv_cache": True, + "trace_cache": False, + }, + ) + + resp = client.get("/metrics") + body = resp.content.decode() + + assert resp.status_code == 200 + assert f'daser_chunk_hits_total{{doc_id="{doc_a["doc_id"]}"}} 1.0' in body + assert f'daser_chunk_misses_total{{doc_id="{doc_a["doc_id"]}"}}' not in body + + +def test_metrics_http_request_counter_records_route() -> None: + """HTTP request metrics use route templates instead of raw doc IDs.""" + core = make_core() + fake_vllm = FakeVLLMClient() + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + client.get("/documents/missing") + + resp = client.get("/metrics") + body = resp.content.decode() + + assert resp.status_code == 200 + assert ( + 'daser_http_requests_total{endpoint="/documents/{doc_id}",' + 'method="GET",status="404"} 1.0' + ) in body + + +def test_metrics_document_prefill_recorded() -> None: + """Upload metrics record document prefill latency.""" + core = make_core() + fake_vllm = FakeVLLMClient(commit_core=core) + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + client.post("/documents", json={"title": "a", "text": "abcd"}) + + resp = client.get("/metrics") + assert resp.status_code == 200 + assert b"daser_document_prefill_seconds" in resp.content + + +def test_metrics_state_snapshot_reflects_cache_state() -> None: + """Gauge metrics reflect the current ring-buffer state after upload.""" + core = make_core() + fake_vllm = FakeVLLMClient(commit_core=core) + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + + client.post("/documents", json={"title": "a", "text": "abcd"}) + + resp = client.get("/metrics") + body = resp.content.decode() + + assert resp.status_code == 200 + # After upload, at least one chunk should be in the cache + assert "daser_cache_chunks" in body + assert "daser_cache_free_slots" in body + assert "daser_documents_total" in body + + +def test_metrics_eviction_counter_reflects_document_delete() -> None: + """Eviction metrics include chunks removed when deleting a document.""" + core = make_core() + fake_vllm = FakeVLLMClient(commit_core=core) + metrics = MetricsCollector(enabled=True) + client = _make_client_with_metrics(core, fake_vllm, metrics=metrics) + doc = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + + delete_resp = client.delete(f"/documents/{doc['doc_id']}") + metrics_resp = client.get("/metrics") + body = metrics_resp.content.decode() + + assert delete_resp.status_code == 200 + assert metrics_resp.status_code == 200 + assert 'daser_eviction_total{reason="document_delete"} 1.0' in body + + +# ------------------------------------------------------------------ +# Stage A: /diag/explain +# ------------------------------------------------------------------ + + +def _make_client_for_diag( + metrics: MetricsCollector | None = None, +) -> tuple[TestClient, FakeVLLMClient, ServerCore]: + """Create a test client whose fake vLLM commits prefetched chunks. + + This mirrors ``_make_client`` but uses the chunk-reuse index so that + ``/diag/explain`` can report hits across the full prompt. + """ + core = make_core_with_index(ChunkReuseIndex(block_tokens=BLOCK_TOKENS)) + fake_vllm = FakeVLLMClient(commit_core=core) + app = build_http_app( + HTTPServerConfig( + vllm_base_url="http://vllm", + model="m", + tokenizer="fake", + block_tokens=4, + system_prompt="S:", + doc_separator="|", + task_separator="? ", + answer_separator="! ", + cache_reuse_mode="chunk", + align_document_chunks=True, + ), + core, + tokenizer=FakeTokenizer(), + vllm=fake_vllm, + metrics_collector=metrics, + ) + return TestClient(app), fake_vllm, core + + +def test_diag_explain_returns_hits() -> None: + """``/diag/explain`` returns cache-hit information for uploaded documents.""" + client, _, _ = _make_client_for_diag() + + doc_a = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + client.post("/documents", json={"title": "b", "text": "efgh"}).json() + + resp = client.post( + "/diag/explain", + json={ + "doc_ids": [doc_a["doc_id"]], + "task": "go", + }, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["doc_ids"] == [doc_a["doc_id"]] + assert body["prompt_tokens"] > 0 + assert "prompt_preview" in body + assert isinstance(body["cache_hits"], list) + assert isinstance(body["hit_count"], int) + assert isinstance(body["hit_token_count"], int) + assert 0.0 <= body["coverage_ratio"] <= 1.0 + assert body["routing_decision"] in ("chunk_reuse", "full_prefill") + assert body["hit_tier"] in ("L1", "miss") + assert body["expected_ttft_ms"] is None + assert body["accuracy_net_enabled"] is False + assert body["cache_reuse_mode"] == "chunk" + + +def test_diag_explain_empty_docs_returns_400() -> None: + """``/diag/explain`` rejects an empty doc_ids list.""" + core = make_core() + fake_vllm = FakeVLLMClient() + client = _make_client_with_metrics(core, fake_vllm) + + resp = client.post("/diag/explain", json={"doc_ids": []}) + + assert resp.status_code == 400 + assert "doc_ids" in resp.json()["detail"].lower() + + +def test_diag_explain_missing_doc_returns_404() -> None: + """``/diag/explain`` returns 404 when a doc_id is unknown.""" + core = make_core() + fake_vllm = FakeVLLMClient() + client = _make_client_with_metrics(core, fake_vllm) + + resp = client.post( + "/diag/explain", + json={"doc_ids": ["nonexistent"], "task": ""}, + ) + + assert resp.status_code == 404 + + +def test_diag_explain_coverage_partial() -> None: + """Coverage ratio is less than 1.0 when only some chunks hit.""" + client, _, _ = _make_client_for_diag() + + doc_a = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + + resp = client.post( + "/diag/explain", + json={ + "doc_ids": [doc_a["doc_id"]], + "task": "a task with extra tokens that are not cached", + }, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["coverage_ratio"] < 1.0 + + +def test_diag_explain_does_not_update_access_stats() -> None: + """Diagnostic lookup must not influence later LFU/LRU access counts.""" + client, _, core = _make_client_for_diag() + doc_a = client.post("/documents", json={"title": "a", "text": "abcd"}).json() + public_doc = client.get(f"/documents/{doc_a['doc_id']}").json() + chunk_key = public_doc["chunk_keys"][0] + meta_before = core.chunk_manager.store.get(chunk_key) + assert meta_before is not None + assert meta_before.access_count == 0 + + resp = client.post( + "/diag/explain", + json={ + "doc_ids": [doc_a["doc_id"]], + "task": "go", + }, + ) + + meta_after = core.chunk_manager.store.get(chunk_key) + assert resp.status_code == 200 + assert meta_after is not None + assert meta_after.access_count == 0 diff --git a/tests/server/test_main_cli.py b/tests/server/test_main_cli.py index 9f578ff..bb72c74 100644 --- a/tests/server/test_main_cli.py +++ b/tests/server/test_main_cli.py @@ -112,12 +112,34 @@ def test_documented_flags_populate_config(tmp_path: Path) -> None: assert http_cfg.tokenizer == str(model_path) assert http_cfg.align_document_chunks is True assert args.cache_reuse_mode == "chunk" + assert args.enable_metrics is False runtime = cfg.runtime_config() assert runtime["transfer_mode"] == "iouring" assert runtime["l1_size_bytes"] == 1000**3 assert runtime["l2_size_bytes"] == cfg.aligned_store_bytes +def test_enable_metrics_flag_is_opt_in(tmp_path: Path) -> None: + """Metrics are disabled by default and enabled only by an explicit flag.""" + model_path = tmp_path / "model" + store_dir = tmp_path / "store" + _write_model_config(model_path) + + args = _run_parse( + [ + "--model-path", + str(model_path), + "--store-dir", + str(store_dir), + "--vllm-base-url", + "http://127.0.0.1:8001", + "--enable-metrics", + ] + ) + + assert args.enable_metrics is True + + def test_default_transfer_mode_is_iouring(tmp_path: Path) -> None: """The server defaults to iouring unless a transfer mode is specified.""" model_path = tmp_path / "model" diff --git a/tests/server/test_metadata_store.py b/tests/server/test_metadata_store.py index 62e5f60..4efb00a 100644 --- a/tests/server/test_metadata_store.py +++ b/tests/server/test_metadata_store.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Standard +from dataclasses import asdict from pathlib import Path import time +import msgpack import pytest from daser.server.metadata_store import ChunkMeta, MetadataStore @@ -84,3 +86,109 @@ def test_save_and_load(tmp_path: Path) -> None: assert len(store2) == 2 assert store2.get_slot_entry(0).kind == "chunk" assert store2.get_slot_entry(1).kind == "cont" + + +def test_chunk_meta_defaults_initialize_access_stats() -> None: + meta = make_meta("abc", start=0, num=2) + assert meta.access_count == 0 + assert meta.last_access_time == meta.created_at + + +def test_touch_increments_access_count_and_time() -> None: + store = MetadataStore(total_slots=8) + store.insert(make_meta("abc", start=0, num=2)) + + store.touch("abc", now=100.0) + meta = store.get("abc") + assert meta is not None + assert meta.access_count == 1 + assert meta.last_access_time == 100.0 + + store.touch("abc", now=101.5) + meta_again = store.get("abc") + assert meta_again is not None + assert meta_again.access_count == 2 + assert meta_again.last_access_time == 101.5 + + +def test_touch_unknown_key_is_noop() -> None: + store = MetadataStore(total_slots=8) + store.touch("missing") + assert store.get("missing") is None + + +def test_load_backward_compatible_with_legacy_payload(tmp_path: Path) -> None: + """Records written before access_count was introduced must still load.""" + legacy_payload = { + "total_slots": 8, + "chunk_index": { + "legacy": { + "chunk_key": "legacy", + "start_slot": 0, + "num_slots": 2, + "token_count": 16, + "pos_offset": 0, + "model_id": "test-model", + "created_at": 1234.0, + "doc_ids": [], + } + }, + "slot_map": [ + {"kind": "chunk", "chunk_key": "legacy", "num_slots": 2}, + {"kind": "cont", "chunk_key": None, "num_slots": 0}, + ] + + [{"kind": "cont", "chunk_key": None, "num_slots": 0} for _ in range(6)], + } + path = str(tmp_path / "legacy.index") + with open(path, "wb") as f: + f.write(msgpack.packb(legacy_payload, use_bin_type=True)) + + store = MetadataStore(total_slots=8) + store.load(path) + meta = store.get("legacy") + assert meta is not None + assert meta.access_count == 0 + assert meta.last_access_time == meta.created_at == 1234.0 + + +def test_load_forward_compatible_with_future_fields(tmp_path: Path) -> None: + """Records carrying unknown future fields must load without errors.""" + base = asdict(make_meta("future", start=0, num=2)) + base["self_contained_score"] = 0.91 # not yet in ChunkMeta + base["future_extension_flag"] = True + payload = { + "total_slots": 8, + "chunk_index": {"future": base}, + "slot_map": [ + {"kind": "chunk", "chunk_key": "future", "num_slots": 2}, + {"kind": "cont", "chunk_key": None, "num_slots": 0}, + ] + + [{"kind": "cont", "chunk_key": None, "num_slots": 0} for _ in range(6)], + } + path = str(tmp_path / "future.index") + with open(path, "wb") as f: + f.write(msgpack.packb(payload, use_bin_type=True)) + + store = MetadataStore(total_slots=8) + store.load(path) + meta = store.get("future") + assert meta is not None + assert meta.chunk_key == "future" + assert not hasattr(meta, "self_contained_score") + + +def test_save_then_load_preserves_access_stats(tmp_path: Path) -> None: + store = MetadataStore(total_slots=8) + store.insert(make_meta("abc", start=0, num=2)) + store.touch("abc", now=200.0) + store.touch("abc", now=300.0) + + path = str(tmp_path / "daser.index") + store.save(path) + + store2 = MetadataStore(total_slots=8) + store2.load(path) + meta = store2.get("abc") + assert meta is not None + assert meta.access_count == 2 + assert meta.last_access_time == 300.0