wehub-resource-sync
376 lines
18 KiB
Markdown
376 lines
18 KiB
Markdown
# Per-Request Root Span
|
||
|
||
This document describes the root span design for `MPServerTracingSubscriber`:
|
||
how a single `"request"` OTel span wraps all child operations for one request,
|
||
and how the span close is deferred correctly when GPU stores are still in flight.
|
||
|
||
## Problem
|
||
|
||
Before this change, `MPServerTracingSubscriber` emitted flat, orphaned child
|
||
spans (`mp.store`, `mp.retrieve`, `mp.lookup_prefetch`) with no parent context.
|
||
Traces in Tempo/Jaeger showed disconnected spans with no request-level view.
|
||
|
||
## Design
|
||
|
||
Each request gets one root `"request"` span that:
|
||
|
||
- Opens at `MP_REQUEST_START` — the first CPU-synchronous touch of a `request_id`
|
||
- Nests all child spans beneath it via OTel context propagation
|
||
- Closes at `MP_REQUEST_END`, deferred if async GPU stores are still in flight
|
||
|
||
### New Events
|
||
|
||
Four new `EventType` values, all CPU-synchronous:
|
||
|
||
| Event | Published from | Purpose |
|
||
|-------|---------------|---------|
|
||
| `MP_REQUEST_START` | `lookup_prefetch_start()`, top of method | Open root span at true request arrival |
|
||
| `MP_STORE_SUBMITTED` | `store()`, before `publish_on_stream(MP_STORE_START)` | Register a pending GPU store before it's enqueued |
|
||
| `MP_RETRIEVE_SUBMITTED` | `retrieve()`, before `publish_on_stream(MP_RETRIEVE_START)` | Register a pending GPU retrieve before it's enqueued |
|
||
| `MP_REQUEST_END` | `end_session()`, after `session_manager.remove()` | Signal that the session lifecycle is complete |
|
||
|
||
### Deferral Protocol
|
||
|
||
`end_session()` is CPU-synchronous; GPU store/retrieve callbacks (`MP_STORE_END`,
|
||
`MP_RETRIEVE_END`) fire later via CUDA host callbacks. Without coordination,
|
||
`MP_REQUEST_END` can arrive and close the root span before GPU work finishes —
|
||
producing orphaned child spans.
|
||
|
||
The fix: `MP_STORE_SUBMITTED` and `MP_RETRIEVE_SUBMITTED` are published *before*
|
||
the respective GPU work is enqueued, incrementing `_pending_store_count` and
|
||
`_pending_retrieve_count`. When `MP_REQUEST_END` arrives:
|
||
|
||
- If both counters are zero → close root immediately
|
||
- Otherwise → save the `REQUEST_END` timestamp; the last `MP_STORE_END` or
|
||
`MP_RETRIEVE_END` to decrement its counter to zero (when the other counter is
|
||
also zero) closes the root using that saved timestamp
|
||
|
||
Root end-time is always the `REQUEST_END` timestamp (the logical request end),
|
||
not the GPU callback timestamp.
|
||
|
||
**Why `MP_RETRIEVE_SUBMITTED` is needed**: vLLM's IPC completion event is
|
||
recorded on the CUDA stream between `MP_RETRIEVE_START` and `MP_RETRIEVE_END`.
|
||
When vLLM unblocks on that event, it can call `end_session()` before the GPU
|
||
callback for `MP_RETRIEVE_END` fires. EventBus queue becomes:
|
||
`→ MP_RETRIEVE_START → MP_REQUEST_END → MP_RETRIEVE_END`
|
||
Without `MP_RETRIEVE_SUBMITTED`, `_on_session_end` sees no in-flight work and
|
||
closes the root span before the retrieve child span ends.
|
||
|
||
## Root Span Attributes
|
||
|
||
In addition to `session_id`, the root `"request"` span carries three hit rate
|
||
attributes that are set when `MP_LOOKUP_PREFETCH_END` is processed:
|
||
|
||
| Attribute | OTel type | Value |
|
||
|-----------|-----------|-------|
|
||
| `hit_tokens` | `int` | tokens found in L1+L2 (numerator) |
|
||
| `requested_tokens` | `int` | chunk-aligned tokens submitted for lookup (denominator) |
|
||
| `hit_rate` | `float` | `hit_tokens / requested_tokens`; `0.0` when denominator is zero |
|
||
|
||
`hit_rate` is stored as a precomputed float because trace UIs (Tempo, Jaeger)
|
||
cannot derive it from two integer attributes at query time.
|
||
|
||
**Invariant:** these attributes are set at `MP_LOOKUP_PREFETCH_END` time, while
|
||
the root span is still open. `LP_END` always precedes `MP_REQUEST_END` in the
|
||
event stream, so the root span is guaranteed to be live in the registry when the
|
||
attributes are written.
|
||
|
||
**Store-only requests** (no `lookup_prefetch_start()` call) never emit
|
||
`MP_LOOKUP_PREFETCH_END`, so the root span will not carry these attributes.
|
||
|
||
### CB path — `cb.request` span
|
||
|
||
The same three attributes appear on the `"cb.request"` root span and are set
|
||
when `CB_LOOKUP_END` is processed by `BlendTracingSubscriber`.
|
||
|
||
`CB_LOOKUP_END` carries `hit_tokens` and `requested_tokens` in its metadata,
|
||
computed at the emit site in `lmcache/v1/multiprocess/modules/blend.py`:
|
||
|
||
| Field | Value |
|
||
|-------|-------|
|
||
| `hit_tokens` | `storage_hits * chunk_size` |
|
||
| `requested_tokens` | `(num_tokens // chunk_size) * chunk_size` (chunk-aligned) |
|
||
|
||
All three `CB_LOOKUP_END` emit sites (no-fingerprint-match, no-GPU-context,
|
||
happy path) populate these fields, so `hit_rate` is always present on the
|
||
`cb.request` span.
|
||
|
||
A fourth attribute is also set on `"cb.request"` at `CB_LOOKUP_END` time:
|
||
|
||
| Attribute | OTel type | Value |
|
||
|-----------|-----------|-------|
|
||
| `prefix_hits` | `int` | chunks found via the prefix probe (not fingerprint matching) |
|
||
|
||
#### Prefix probe
|
||
|
||
`cb_lookup_pre_computed` has two lookup paths:
|
||
|
||
1. **Fingerprint path** — `BlendTokenRangeMatcher.match_sub_sequence` finds
|
||
sub-sequence matches using polynomial rolling hashes. Covers arbitrary
|
||
(non-prefix) positions in the token sequence.
|
||
2. **Prefix probe** — a fallback that runs after the fingerprint path and fills
|
||
in chunks at contiguous prefix positions not already covered by fingerprint
|
||
results. It calls `token_hasher.compute_chunk_hashes(token_ids)` to derive
|
||
the same storage keys used by `cb_store_final` and `cb_store_pre_computed`,
|
||
then creates `CBMatchResult(old_st==cur_st)` candidates for uncovered slots.
|
||
These candidates flow through the same prefetch/poll/evict machinery as
|
||
fingerprint results.
|
||
|
||
The prefix probe closes the gap between the MP and CB storage paths: chunks
|
||
written by `cb_store_final` (which only registers fingerprints when
|
||
`worker_id in [0, None]`) and chunks written via the MP `store()` path (which
|
||
uses block hashes incompatible with fingerprint matching) are both visible to
|
||
`cb_lookup_pre_computed` through the prefix probe.
|
||
|
||
#### Lazy registration
|
||
|
||
When `cb_lookup_pre_computed` returns results that came *entirely* from the
|
||
prefix probe (i.e. `fingerprint_results` is empty) and the calling worker is
|
||
rank 0 or the driver (`worker_id in [0, None]`), the found prefix chunks are
|
||
registered into `BlendTokenRangeMatcher` so that future lookups can find them
|
||
via the faster fingerprint path. Registration is guarded by
|
||
`BlendTokenRangeMatcher.has_chunk(token_hash)` to prevent overwriting existing
|
||
compact-ID assignments when the range matcher already has entries for the same
|
||
token sequence.
|
||
|
||
`prefix_hits` counts the chunks found exclusively through the prefix probe
|
||
(after deduplication against fingerprint results). When `fingerprint_results`
|
||
is non-empty and prefix candidates fill in additional positions,
|
||
`prefix_hits` reflects only the prefix-probe portion of the total
|
||
`storage_hits`.
|
||
|
||
## Request Scenarios
|
||
|
||
### Scenario 1 — Full Cache Hit
|
||
|
||
Path: `lookup_prefetch → retrieve → store`
|
||
|
||
```
|
||
CPU ─[REQUEST_START]─[LP_START]─[LP_END]──[RETR_SUBMITTED]──[STORE_SUBMITTED]─[REQUEST_END]─►
|
||
GPU ──────────────────────────────[RETR_START]─[vLLM_IPC]─[RETR_END]──[STORE_START]─[STORE_END]─►
|
||
|
||
root "request" [═══════════════════════════════════════════════════════════════════════════════]
|
||
mp.lookup_prefetch [══════════]
|
||
mp.retrieve [══════════════════]
|
||
mp.store [══════════════════════]
|
||
```
|
||
|
||
Root closes at `REQUEST_END` (deferred until both retrieve and store complete).
|
||
|
||
---
|
||
|
||
### Scenario 2 — Cache Miss (no retrieve)
|
||
|
||
Path: `lookup_prefetch → store`, no retrieve
|
||
|
||
```
|
||
CPU ─[REQUEST_START]─[LP_START]─[LP_END]──────────[STORE_SUBMITTED]─[REQUEST_END]─►
|
||
GPU ───────────────────────────────────────────────────────[STORE_START]─[STORE_END]─►
|
||
|
||
root "request" [═══════════════════════════════════════════════════════════════════]
|
||
mp.lookup_prefetch [══════════]
|
||
mp.store [══════════════════════]
|
||
```
|
||
|
||
No retrieve occurred, so `mp.retrieve` is absent.
|
||
|
||
---
|
||
|
||
### Scenario 3 — Lookup Only
|
||
|
||
Path: `lookup_prefetch` only, no store
|
||
|
||
```
|
||
CPU ─[REQUEST_START]─[LP_START]─[LP_END]─[REQUEST_END]─►
|
||
|
||
root "request" [════════════════════════════════════════]
|
||
mp.lookup_prefetch [══════════]
|
||
```
|
||
|
||
Root closes immediately at `REQUEST_END`.
|
||
|
||
---
|
||
|
||
### Scenario 4 — Store Only (no lookup)
|
||
|
||
Path: `store` with no prior `lookup_prefetch_start()` call
|
||
|
||
```
|
||
CPU ─(no REQUEST_START)──────[STORE_SUBMITTED]─[REQUEST_END]─►
|
||
GPU ──────────────────────────────────[STORE_START]─[STORE_END]─►
|
||
|
||
root "request" (lazy, created at MP_STORE_START)
|
||
[═════════════════════════]
|
||
mp.store [══════════════]
|
||
```
|
||
|
||
`MP_REQUEST_START` is only emitted from `lookup_prefetch_start()`. If that path
|
||
was not taken, `_get_or_create_request_span()` is called lazily on the first child
|
||
`_on_start()`. Root start time equals `STORE_START` timestamp.
|
||
|
||
---
|
||
|
||
### Scenario 5 — REQUEST_END Races GPU Store
|
||
|
||
`end_session()` called before the GPU store callback fires.
|
||
|
||
```
|
||
CPU ─[REQUEST_START]─[LP_START]─[LP_END]─[STORE_SUBMITTED]─[REQUEST_END]────────────────────►
|
||
GPU ──────────────────────────────────────────────[STORE_START]──────────────[STORE_END]─────►
|
||
▲
|
||
REQUEST_END arrives here─┘ (before STORE_END)
|
||
|
||
root "request" [═══════════════════════════════════════════════════════════════════════════]
|
||
mp.lookup_prefetch [══════════]
|
||
mp.store [═══════════════════]
|
||
▲
|
||
STORE_SUBMITTED → count=1 │
|
||
REQUEST_END → count>0 → defer (save ts) │
|
||
STORE_END → count=0 → _close_request_span(deferred_ts) ────────────────┘
|
||
```
|
||
|
||
---
|
||
|
||
### Scenario 6 — Multiple Stores, Deferred Close
|
||
|
||
Two concurrent stores; root stays open until both complete.
|
||
|
||
```
|
||
CPU ─[REQUEST_START]─[LP_START]─[LP_END]─[SUBMITTED×2]─[REQUEST_END]──────────────────────────────────►
|
||
GPU ────────────────────────────────────────────────────[S1_START]─[S1_END]─[S2_START]─[S2_END]────────►
|
||
|
||
root "request" [═══════════════════════════════════════════════════════════════════════════════════════]
|
||
mp.lookup_prefetch [══════════]
|
||
mp.store (1) [══════════]
|
||
mp.store (2) [══════════]
|
||
▲
|
||
count=2 at REQUEST_END → defer │
|
||
S1_END → count=1 → still open │
|
||
S2_END → count=0 → _close_request_span(deferred_ts) ─────────────────────────────┘
|
||
```
|
||
|
||
## Summary
|
||
|
||
| Scenario | Root opens | Root closes |
|
||
|----------|-----------|-------------|
|
||
| Full hit | `MP_REQUEST_START` | last `MP_STORE_END` / `MP_RETRIEVE_END` (stamped at `REQUEST_END` time) |
|
||
| Cache miss | `MP_REQUEST_START` | last `MP_STORE_END` (stamped at `REQUEST_END` time) |
|
||
| Lookup only | `MP_REQUEST_START` | `REQUEST_END` (immediate) |
|
||
| Store only | `MP_STORE_START` (lazy) | `REQUEST_END` (immediate) |
|
||
| REQUEST_END races store | `MP_REQUEST_START` | last `MP_STORE_END` (stamped at `REQUEST_END` time) |
|
||
| REQUEST_END races retrieve | `MP_REQUEST_START` | last `MP_RETRIEVE_END` (stamped at `REQUEST_END` time) |
|
||
| Multiple stores | `MP_REQUEST_START` | last `MP_STORE_END` (stamped at `REQUEST_END` time) |
|
||
|
||
## Implementation
|
||
|
||
| File | Change |
|
||
|------|--------|
|
||
| `lmcache/v1/mp_observability/event.py` | Add `MP_REQUEST_START`, `MP_STORE_SUBMITTED`, `MP_RETRIEVE_SUBMITTED`, `MP_REQUEST_END` |
|
||
| `lmcache/v1/multiprocess/server.py` | Emit the 4 events at `lookup_prefetch_start()`, `store()`, `retrieve()`, `end_session()` |
|
||
| `lmcache/v1/mp_observability/subscribers/tracing/mp_server.py` | Root span logic: `_pending_store_count`, `_pending_retrieve_count`, `_deferred_session_end_ts`; handlers `_on_request_start`, `_on_store_submitted`, `_on_retrieve_submitted`, `_on_session_end`; helpers `_get_or_create_request_span`, `_close_request_span` |
|
||
| `lmcache/v1/mp_observability/subscribers/tracing/span_registry.py` | `SpanRegistry`: shared dict of open spans keyed by `(session_id, span_name)` for cross-subscriber parent lookup |
|
||
| `tests/v1/mp_observability/subscribers/tracing/test_mp_server.py` | Tests for all scenarios including retrieve deferral |
|
||
| `lmcache/v1/multiprocess/modules/blend.py` | Prefix probe in `cb_lookup_pre_computed`; lazy registration; `has_chunk` on `BlendTokenRangeMatcher`; `prefix_hits` in `CB_LOOKUP_END` metadata |
|
||
| `lmcache/v1/mp_observability/subscribers/tracing/cb_server.py` | Stamp `prefix_hits` on `"cb.request"` root span from `CB_LOOKUP_END` |
|
||
| `tests/v1/multiprocess/test_blend_server_v2.py` | `has_chunk` unit tests |
|
||
| `tests/v1/mp_observability/subscribers/tracing/test_cb_server.py` | `prefix_hits` attribute tests |
|
||
|
||
---
|
||
|
||
## Extending the Span Hierarchy
|
||
|
||
### How the registry works
|
||
|
||
`MPServerTracingSubscriber` writes every open span into a shared
|
||
`SpanRegistry` while it is live:
|
||
|
||
```
|
||
registry[(session_id, "request")] → (root_span, root_ctx) # open: REQUEST_START → REQUEST_END
|
||
registry[(session_id, "retrieve")] → (retrieve_span, ctx) # open: RETRIEVE_START → RETRIEVE_END
|
||
registry[(session_id, "store")] → (store_span, ctx) # open: STORE_START → STORE_END
|
||
registry[(session_id, "lookup_prefetch")] → (lp_span, ctx) # open: LP_START → LP_END
|
||
```
|
||
|
||
Any subscriber that receives the same `SpanRegistry` instance can call
|
||
`registry.get_context(session_id, "request")` (or any other name) to obtain
|
||
the OTel context needed to nest a new span.
|
||
|
||
---
|
||
|
||
### Example 1 — new span at the same level
|
||
|
||
To add an `l1.read` span nested directly under the root `"request"` span,
|
||
create a new subscriber file and register it with the shared registry.
|
||
No existing files need to change.
|
||
|
||
**`subscribers/tracing/l1.py`**:
|
||
|
||
```python
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
from opentelemetry import trace
|
||
|
||
from lmcache.v1.mp_observability.event import Event, EventType
|
||
from lmcache.v1.mp_observability.event_bus import EventCallback, EventSubscriber
|
||
from lmcache.v1.mp_observability.subscribers.tracing.span_registry import SpanRegistry
|
||
|
||
_tracer = trace.get_tracer("lmcache_mp.l1")
|
||
|
||
class L1TracingSubscriber(EventSubscriber):
|
||
def __init__(self, registry: SpanRegistry) -> None:
|
||
self._registry = registry
|
||
self._pending: dict[str, object] = {}
|
||
|
||
def get_subscriptions(self) -> dict[EventType, EventCallback]:
|
||
return {
|
||
EventType.L1_READ_RESERVED: self._on_start,
|
||
EventType.L1_READ_FINISHED: self._on_end,
|
||
}
|
||
|
||
def _on_start(self, event: Event) -> None:
|
||
parent_ctx = self._registry.get_context(event.session_id, "request")
|
||
span = _tracer.start_span(
|
||
"l1.read", context=parent_ctx, start_time=int(event.timestamp * 1e9)
|
||
)
|
||
self._pending[event.session_id] = span
|
||
|
||
def _on_end(self, event: Event) -> None:
|
||
span = self._pending.pop(event.session_id, None)
|
||
if span:
|
||
span.end(end_time=int(event.timestamp * 1e9))
|
||
```
|
||
|
||
**`config.py`** (the only change needed):
|
||
|
||
```python
|
||
registry = SpanRegistry()
|
||
bus.register_subscriber(MPServerTracingSubscriber(registry))
|
||
bus.register_subscriber(L1TracingSubscriber(registry)) # ← add this line
|
||
```
|
||
|
||
This produces: `request → l1.read` (alongside `mp.retrieve`, `mp.store`, etc.)
|
||
|
||
---
|
||
|
||
### Example 2 — sub-span nested under an existing child span
|
||
|
||
To nest a span *inside* `mp.retrieve` (e.g. an L2 disk load that happens
|
||
during a retrieve), look up `"retrieve"` as the parent instead of `"request"`.
|
||
The `"retrieve"` entry is live in the registry from `MP_RETRIEVE_START` to
|
||
`MP_RETRIEVE_END`.
|
||
|
||
```python
|
||
def _on_detail_start(self, event: Event) -> None:
|
||
sid = event.session_id
|
||
# Prefer the immediate parent; fall back to root if retrieve has ended.
|
||
parent_ctx = (
|
||
self._registry.get_context(sid, "retrieve")
|
||
or self._registry.get_context(sid, "request")
|
||
)
|
||
span = _tracer.start_span(
|
||
"l2.disk_load", context=parent_ctx, start_time=int(event.timestamp * 1e9)
|
||
)
|
||
self._pending[sid] = span
|
||
```
|
||
|
||
This produces a three-level trace: `request → mp.retrieve → l2.disk_load`.
|