From b4811bd3cb4b150f31782ef532ea3a1afbef709b Mon Sep 17 00:00:00 2001 From: 7774596 <2238674546@qq.com> Date: Thu, 28 May 2026 16:06:39 +0800 Subject: [PATCH 1/4] docs(optimizations): plan multi-document chunk reuse production rollout Consolidate the production-grade plan for user-managed multi-document chunk reuse into a single optimization record under ``docs/optimizations/``. The document captures positioning, SLA targets, connector capability boundary, staged roadmap (observability, DocSet management, hierarchical cache, accuracy safety net), data model extensions, CLI flags, test matrix, risks, rollback plan, and PR breakdown. No source code changes; later PRs land each stage behind independent feature flags. --- docs/optimizations/5_multi_doc_production.md | 516 +++++++++++++++++++ 1 file changed, 516 insertions(+) create mode 100644 docs/optimizations/5_multi_doc_production.md diff --git a/docs/optimizations/5_multi_doc_production.md b/docs/optimizations/5_multi_doc_production.md new file mode 100644 index 0000000..37295f5 --- /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={lru, lfu, pinned_attempt} | + | `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 | +| `--metrics-port` | `9090` | A | Prometheus metrics port | + +## 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. From 99302dc06bb9223e6db772f9f20367177ab1b590 Mon Sep 17 00:00:00 2001 From: 7774596 <2238674546@qq.com> Date: Thu, 28 May 2026 17:28:47 +0800 Subject: [PATCH 2/4] feat(server): track chunk access count and time on lookup hits Lay the metadata foundation for production observability and the later LFU+LRU hybrid eviction policy. - ChunkMeta gains ``access_count`` and ``last_access_time``; both default to zero / ``created_at`` so existing call sites are unaffected. - ``MetadataStore.load`` now reads chunks through a ``_chunk_meta_from_payload`` helper that drops unknown payload keys, keeping msgpack records both forward- and backward-compatible as later production stages add more fields. - ``MetadataStore.touch`` updates the stats and ``ServerCore.lookup`` calls it for every matched chunk, so hit traffic alone is enough to drive eviction and metrics later. - Unit tests cover the new defaults, the touch contract, both msgpack compatibility directions, and the existing lookup paths. --- daser/server/core.py | 6 ++ daser/server/metadata_store.py | 63 +++++++++++++++- tests/server/test_core.py | 44 +++++++++++ tests/server/test_metadata_store.py | 112 ++++++++++++++++++++++++++++ 4 files changed, 222 insertions(+), 3 deletions(-) diff --git a/daser/server/core.py b/daser/server/core.py index fca10c7..6f2c948 100644 --- a/daser/server/core.py +++ b/daser/server/core.py @@ -245,6 +245,10 @@ async def rebuild_retrieval_index(self) -> None: async def lookup(self, tokens: list[int], model_id: str) -> list[ChunkInfo]: """Look up cached chunks for token IDs. + Hits update per-chunk access statistics (``access_count`` and + ``last_access_time``) on the underlying ``MetadataStore`` so later + eviction policies and observability endpoints can read them. + Args: tokens: prompt token IDs. model_id: model identifier. @@ -259,6 +263,8 @@ async def lookup(self, tokens: list[int], model_id: str) -> list[ChunkInfo]: 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( diff --git a/daser/server/metadata_store.py b/daser/server/metadata_store.py index dc50e1a..51403b0 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,39 @@ 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 +186,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 +283,8 @@ 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/tests/server/test_core.py b/tests/server/test_core.py index ca5484e..7203406 100644 --- a/tests/server/test_core.py +++ b/tests/server/test_core.py @@ -104,6 +104,50 @@ 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 = _hash_tokens(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_miss_leaves_access_stats_untouched() -> None: + core = make_core() + tokens = [1, 2, 3, 4] + key = _hash_tokens(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_metadata_store.py b/tests/server/test_metadata_store.py index 62e5f60..115e9cc 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,113 @@ 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 From 09df282e28441e3b8cb36d7c596ed1ba9c27efb2 Mon Sep 17 00:00:00 2001 From: 7774596 <2238674546@qq.com> Date: Mon, 8 Jun 2026 12:06:44 +0800 Subject: [PATCH 3/4] test(server): use rolling prefix keys in access stat tests Align Stage A lookup tests with master prefix index helpers after rebase. --- tests/server/test_core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/server/test_core.py b/tests/server/test_core.py index 7203406..5a273f5 100644 --- a/tests/server/test_core.py +++ b/tests/server/test_core.py @@ -108,7 +108,7 @@ async def test_wait_for_committed_chunks_times_out() -> None: async def test_lookup_hit_updates_chunk_access_stats() -> None: core = make_core() tokens = [1, 2, 3, 4] - key = _hash_tokens(tokens) + key = first_rolling_key(tokens) await core.alloc_chunk(key, token_count=len(tokens), model_id="m") await core.commit_chunk(key) @@ -131,7 +131,7 @@ async def test_lookup_hit_updates_chunk_access_stats() -> None: async def test_lookup_miss_leaves_access_stats_untouched() -> None: core = make_core() tokens = [1, 2, 3, 4] - key = _hash_tokens(tokens) + key = first_rolling_key(tokens) await core.alloc_chunk(key, token_count=len(tokens), model_id="m") await core.commit_chunk(key) From 483d8cd2e3c95ba7533e80156855e844c9833fd8 Mon Sep 17 00:00:00 2001 From: 7774596 <2238674546@qq.com> Date: Tue, 9 Jun 2026 16:38:52 +0800 Subject: [PATCH 4/4] feat(server): add Stage A observability endpoints Expose opt-in Prometheus metrics and diagnostic cache explanations for the HTTP service, while keeping diagnostic lookups read-only so access counters reflect production demand. Add a baseline measurement script and tests for metrics, diagnostics, CLI wiring, and lookup access accounting. --- daser/server/__main__.py | 13 +- daser/server/chunk_manager.py | 5 + daser/server/core.py | 51 ++- daser/server/http/app.py | 179 +++++++++- daser/server/http/metrics.py | 313 ++++++++++++++++++ daser/server/metadata_store.py | 7 +- docs/optimizations/5_multi_doc_production.md | 4 +- examples/baseline_measure/measure.py | 240 ++++++++++++++ pyproject.toml | 1 + tests/server/test_core.py | 17 + tests/server/test_http_server.py | 328 +++++++++++++++++++ tests/server/test_main_cli.py | 22 ++ tests/server/test_metadata_store.py | 8 +- 13 files changed, 1157 insertions(+), 31 deletions(-) create mode 100644 daser/server/http/metrics.py create mode 100644 examples/baseline_measure/measure.py 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 6f2c948..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,16 +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 (``access_count`` and - ``last_access_time``) on the underlying ``MetadataStore`` so later - eviction policies and observability endpoints can read them. + 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. @@ -260,11 +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 - for match in matches: - self._cm.store.touch(match.meta.chunk_key) + 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( @@ -438,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. @@ -505,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]) @@ -642,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, @@ -659,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: @@ -687,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 51403b0..d38a8b1 100644 --- a/daser/server/metadata_store.py +++ b/daser/server/metadata_store.py @@ -54,9 +54,7 @@ def __post_init__(self) -> None: self.last_access_time = self.created_at -_CHUNK_META_FIELD_NAMES: frozenset[str] = frozenset( - f.name for f in fields(ChunkMeta) -) +_CHUNK_META_FIELD_NAMES: frozenset[str] = frozenset(f.name for f in fields(ChunkMeta)) def _chunk_meta_from_payload(payload: dict[str, Any]) -> ChunkMeta: @@ -283,8 +281,7 @@ def load(self, path: str) -> None: self._total_slots = payload["total_slots"] self._chunk_index = { - k: _chunk_meta_from_payload(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 index 37295f5..cd6cf22 100644 --- a/docs/optimizations/5_multi_doc_production.md +++ b/docs/optimizations/5_multi_doc_production.md @@ -256,7 +256,7 @@ Subparts: | `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={lru, lfu, pinned_attempt} | + | `daser_eviction_total` | reason={ring, explicit, document_delete} | | `daser_l1_l2_l3_size_bytes` | level | | `daser_accuracy_drift_estimate` | docset_name | @@ -424,7 +424,7 @@ service is byte-compatible with master when no new flag is set). | `--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 | -| `--metrics-port` | `9090` | A | Prometheus metrics port | +| `--enable-metrics` | `false` | A | Serve Prometheus metrics at `/metrics` on the main HTTP server | ## 8. Test Matrix 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 5a273f5..f9d3f23 100644 --- a/tests/server/test_core.py +++ b/tests/server/test_core.py @@ -127,6 +127,23 @@ async def test_lookup_hit_updates_chunk_access_stats() -> None: 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() 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 115e9cc..4efb00a 100644 --- a/tests/server/test_metadata_store.py +++ b/tests/server/test_metadata_store.py @@ -137,9 +137,7 @@ def test_load_backward_compatible_with_legacy_payload(tmp_path: Path) -> None: {"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) - ], + + [{"kind": "cont", "chunk_key": None, "num_slots": 0} for _ in range(6)], } path = str(tmp_path / "legacy.index") with open(path, "wb") as f: @@ -165,9 +163,7 @@ def test_load_forward_compatible_with_future_fields(tmp_path: Path) -> None: {"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) - ], + + [{"kind": "cont", "chunk_key": None, "num_slots": 0} for _ in range(6)], } path = str(tmp_path / "future.index") with open(path, "wb") as f: