From 3db2c5559d6b75e86f143f1a3d5639abf2f43227 Mon Sep 17 00:00:00 2001 From: RickZ <121943721+rickisba@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:08:44 +0800 Subject: [PATCH 1/4] observability: propagate traces to instance --- doc/architecture/observability-v1.md | 12 +- .../configuration-and-interfaces.md | 6 +- doc/developer-handbook/runtime-flows.md | 6 +- .../issue-141-unified-observability.md | 3 + instance/instance_api.py | 46 +++- instance/observability.py | 220 ++++++++++++++++++ proxy/queue/manager.py | 11 + .../test_instance_observability.py | 131 +++++++++++ .../test_scheduler_proxy_production_paths.py | 10 +- test/test_repository_governance.py | 6 +- 10 files changed, 431 insertions(+), 20 deletions(-) create mode 100644 instance/observability.py create mode 100644 test/observability/test_instance_observability.py diff --git a/doc/architecture/observability-v1.md b/doc/architecture/observability-v1.md index 64b6fc4..f5bb718 100644 --- a/doc/architecture/observability-v1.md +++ b/doc/architecture/observability-v1.md @@ -49,7 +49,7 @@ copied. Ambiguous or overwritten values are labeled `legacy_projected`, never upgraded to actual observations. The current `cacheroute_meta` has no `request_id`, and this foundation does not add one. -## Scheduler-to-Proxy propagation +## Internal propagation The Scheduler now creates the internal context and overwrites the complete reserved header set: `scheduler-request-id`, `x-cacheroute-trace-version`, @@ -57,8 +57,7 @@ reserved header set: `scheduler-request-id`, `x-cacheroute-trace-version`, `x-cacheroute-trace-sampled`, and `x-cacheroute-trace-created-at`. The request ID allocated by the Scheduler remains authoritative. Client values using these names are never trusted; Authorization forwarding and the serialized payload -remain unchanged. Propagation stops at the Proxy and no trace header or model -is sent to an Instance or returned to a client. +remain unchanged. Propagation now reaches the selected Instance through the same exact reserved header set. Proxy encodes its accepted Scheduler context, or its Proxy-local fallback context, immediately before forwarding to Instance and overwrites every reserved value for that internal call. No arbitrary client trace header becomes authoritative and no JSON/base64 trace payload is added to the request body. Instance validates the complete reserved set, accepts matching fresh context, or creates a request-local fallback context using its resolved startup profile and sample rate. The reserved headers stop at Instance: Instance-to-vLLM forwarding does not include CacheRoute trace headers, W3C Trace Context, or any canonical trace object, and no trace is returned to a client. Scheduler and Proxy each resolve `CACHEROUTE_RUNTIME_PROFILE` through an explicit lifespan startup helper, with `legacy` as the compatibility default. Because the current services do not have an implemented production v1 data path, `auto` resolves with `v1_available=False` and is stored as `legacy`; explicit `legacy`, `test/mock`, and `v1` remain valid metadata values, but no stored or propagated context can remain `auto`. A missing, malformed, stale, request-ID-mismatched, or profile-mismatched context causes a Proxy-local context to be created and never rejects an otherwise valid request. Profile metadata does not select runtime behavior. The propagation freshness rule retains the five-minute maximum age and also accepts only a bounded 30-second future clock skew between Scheduler and Proxy clocks. @@ -83,3 +82,10 @@ vLLM execution, prefill, decode, Gateway, or LMCache timings. Collection remains process-local and immutable. The Legacy trace mapping and client metadata are unchanged and are not copied into `RequestTrace`. There is no client export, external exporter, registry, debug endpoint, or persistence. + + +## Instance-observed stages + +Instance resolves `CACHEROUTE_RUNTIME_PROFILE` and `CACHEROUTE_TRACE_SAMPLE_RATE` once during FastAPI lifespan startup and stores the immutable result in application state. The compatibility default remains `legacy`; `auto` is resolved and never persisted; invalid sample rates fail closed to `0.0` with at most one bounded warning reason. These settings are correlation metadata only and do not select mock versus real vLLM mode, registration, heartbeat, topology discovery, routing, injection, fallback, retries, or timeouts. + +For sampled valid Instance requests, collection is request-local. Completion spans only the Instance-observed downstream boundary around the existing mock or real-vLLM helper. Streaming first response spans from downstream invocation to the first non-empty downstream chunk observed by Instance. Streaming decode spans from after that chunk until the downstream stream ends. Non-streaming requests skip first response and decode with `non_streaming_request`. Empty streams, failures, and cancellations finalize without running stages and use bounded static canonical errors. All Instance stages use `TraceComponent.INSTANCE`; this slice does not emit `VLLM_PREFILL`, does not claim vLLM or LMCache provenance, and does not provide authoritative internal vLLM prefill/decode timing, LMCache hit-token data, remote-read data, external export, persistence, registry, aggregation, or client-visible traces. Proxy and Instance may contain repeated stage names in separate process-local traces that share one `TraceContext`. diff --git a/doc/developer-handbook/configuration-and-interfaces.md b/doc/developer-handbook/configuration-and-interfaces.md index f47ae88..8b14704 100644 --- a/doc/developer-handbook/configuration-and-interfaces.md +++ b/doc/developer-handbook/configuration-and-interfaces.md @@ -6,8 +6,8 @@ Catalog only developer-facing surfaces. Source definitions remain authoritative; | Name | Owner | Type | Exact default | Allowed values | Scope | Effect | Invalid/fallback behavior | Source | Validation | |---|---|---|---|---|---|---|---|---|---| -| `CACHEROUTE_RUNTIME_PROFILE` | Scheduler/Proxy observability | string | unset -> service passes default legacy startup resolution | `auto`, `legacy`, `v1`, `test/mock` plus compatibility aliases accepted by normalizer | startup | Resolves trace runtime metadata | Current Scheduler and Proxy use `v1_available=False`; `auto` resolves to `legacy`; invalid values raise through runtime normalization | `scheduler/scheduler.py`, `proxy/proxy.py`, `src/cacheroute/observability/startup.py` | `test/observability` | -| `CACHEROUTE_TRACE_SAMPLE_RATE` | Scheduler/Proxy observability | float string | unset -> `0.0` | finite `0.0` through `1.0` | startup | Sets deterministic trace sampled flag | malformed/non-finite/out-of-range values fail closed to `0.0` and expose warning reason | `src/cacheroute/observability/propagation.py` | `test/observability` | +| `CACHEROUTE_RUNTIME_PROFILE` | Scheduler/Proxy/Instance observability | string | unset -> service passes default legacy startup resolution | `auto`, `legacy`, `v1`, `test/mock` plus compatibility aliases accepted by normalizer | startup | Resolves trace runtime metadata | Current Scheduler, Proxy, and Instance use `v1_available=False`; `auto` resolves to `legacy`; invalid values raise through runtime normalization | `scheduler/scheduler.py`, `proxy/proxy.py`, `src/cacheroute/observability/startup.py` | `test/observability` | +| `CACHEROUTE_TRACE_SAMPLE_RATE` | Scheduler/Proxy/Instance observability | float string | unset -> `0.0` | finite `0.0` through `1.0` | startup | Sets deterministic trace sampled flag | malformed/non-finite/out-of-range values fail closed to `0.0` and expose warning reason | `src/cacheroute/observability/propagation.py` | `test/observability` | | `SCHEDULER_LOG_FILE` | Scheduler | path/string | non-portable Transitional Legacy source default (repository-external absolute log path in `core/config.py`; override for portable deployments) | path-like string | startup | Scheduler log path/config value | falls back to source owner default | `core/config.py`, `scheduler/scheduler.py` | demos | | `SCHEDULER_VERBOSE_REQUEST_LOG` | Scheduler | int/string | `1` | intended `0` or `1` | request | Enables verbose request scheduling logs | compared to integer `1` in current source; string environment values do not match that equality | `core/config.py`, `scheduler/scheduler.py` | demos | | `SCHEDULER_MODEL_PATH` | Scheduler | path/string | non-portable Transitional Legacy source default from `core/config.py` (repository-external model path; override for portable deployments) | path/model identifier | startup | Model path used by scheduler tokenizer/model setup | falls back to source owner default | `scheduler/scheduler.py`, `core/config.py` | source checkout tests | @@ -117,7 +117,7 @@ Catalog only developer-facing surfaces. Source definitions remain authoritative; | Scheduler control plane | `/healthz`, `/v1/proxy/register`, `/v1/proxy/heartbeat`, `/v1/proxy/unregister`, `/v1/proxy/list`, `/debug/proxy_pool_resources`, `/v1/kdn/register`, `/v1/kdn/heartbeat`, `/v1/kdn/unregister`, `/v1/kdn/list` | Proxy/KDN registration, heartbeat, unregister, list params | JSON control-plane status/resources | Current Transitional | | Proxy data plane | `POST /v1/chat/completions`, `POST /v1/completions` | Scheduler-forwarded internal request model and reserved trace headers | Chat may stream `text/event-stream`; completions returns JSON; Proxy collects local prepare/ready/first response/decode/completion stages | Current | | Proxy control plane | `/healthz`, `/debug/status`, `/v1/instance/register`, `/v1/instance/heartbeat`, `/v1/instance/resource_snapshot`, `/v1/instance/unregister`, `/v1/instance/list`, `/debug/pool_resource`, `/debug/pool_resource_sources`, `/debug/instance_resources`, `/debug/instance_loads`, `/v1/topology/report`, `/v1/topology/kdn_links` | Instance registration/heartbeat/resource/topology payloads; list supports `include_dead` | JSON pool, resource, and topology snapshots | Current Transitional | -| Instance data plane | `POST /v1/chat/completions`, `POST /v1/completions` | Proxy-forwarded OpenAI-like body | Forwards to vLLM or mock; no canonical trace context propagated from Proxy to Instance | Current Transitional | +| Instance data plane | `POST /v1/chat/completions`, `POST /v1/completions` | Proxy-forwarded OpenAI-like body plus exact internal reserved trace headers | Forwards to vLLM or mock with unchanged public JSON/SSE; validates or locally replaces context; no trace header reaches vLLM and no canonical trace is client-visible | Current Transitional | | Instance control plane | `GET /healthz`, `POST /v1/kv/inject_ready` | health or KV ready payload | JSON health or KV injection signaling result | Current Transitional | | KDN | `/v1/topology/hello`, `/v1/topology/ping`, `/knowledge/snapshot`, `/knowledge/register_text`, `/knowledge/build_kv`, `/knowledge/search/text`, `/knowledge/delete`, `/knowledge/purge_all`, `/knowledge/inject_ready_kv`, `/knowledge/pool_status` | topology, knowledge, KV build/search/delete/inject payloads | JSON knowledge/KV/status payloads; network timing fields when simulated network is enabled | Current Transitional | | Client validation | not a service endpoint | CLI accepts full URL, `-H/--header`, `-d/--data/--data-raw`; validates OpenAI-like fields | Prints response and CacheRoute metrics; unsupported body keys are validation errors | Current | diff --git a/doc/developer-handbook/runtime-flows.md b/doc/developer-handbook/runtime-flows.md index 9458ccc..b70bac4 100644 --- a/doc/developer-handbook/runtime-flows.md +++ b/doc/developer-handbook/runtime-flows.md @@ -8,8 +8,8 @@ | Text versus KVCache injection | Scheduler, Proxy, KDN, Instance | `Injection_type`: `text`, `kvcache`, `hybrid`; knowledge IDs | Proxy strategies and request preparation | Proxy task-level timing/logging and UI state | If KV is unavailable, current compatibility may recompute text depending on path | Proxy demos | Unified policy package migration | | Cache artifact and cache-operation flow | `cacheroute.cache`, KDN contract models, legacy KDN paths | `artifact_`, `cacheop_`, endpoint compatibility IDs | KDN contract requests | Contract JSON models | Unsupported or unknown compatibility must remain explicit | contract foundation tests | Gateway/LMCache operation execution adapters | | Instance capability registration | Instance, Proxy | capability identity, compatibility profile, endpoint-like data | `instance/control_plane.py`, demo instance | Proxy Instance list and resource snapshots | Capability mismatch prevents unsafe reuse | instance/proxy demos | Canonical topology registry migration | -| Scheduler-to-Proxy observability propagation | Scheduler, Proxy | reserved trace headers, `trace_`, request ID, runtime profile | Scheduler forwarding path and Proxy request path | Process-local Proxy trace stages | Proxy validates headers or creates local fallback context | `test/observability` | Client-returned canonical traces, Proxy-to-Instance context | -| Proxy-local request-stage collection | Proxy | prepare queue, ready queue, first response, decode, completion | Proxy request execution | Process-local collector | These are transport/Proxy-observed intervals, not authoritative vLLM prefill/decode timings | `test/observability` | vLLM/Instance timing instrumentation | +| Scheduler-to-Proxy-to-Instance observability propagation | Scheduler, Proxy, Instance | reserved trace headers, `trace_`, request ID, runtime profile | Scheduler forwarding path, Proxy request path, Proxy-to-Instance forwarding, Instance handlers | Separate process-local Proxy and Instance trace stages | Proxy and Instance validate headers or create local fallback context; exact reserved set stops at Instance | `test/observability` | Client-returned canonical traces, cross-process aggregation | +| Proxy-local and Instance-local request-stage collection | Proxy, Instance | prepare queue, ready queue, first response, decode, completion | Proxy queue execution; Instance chat/completion handlers | Request-local collectors and immutable process-local traces | These are transport-observed intervals, not authoritative vLLM prefill/decode timings; no trace header reaches vLLM | `test/observability` | authoritative vLLM/LMCache timing instrumentation | | Legacy compatibility paths | Root packages and shims | legacy runtime profile, Redis/LMCache key projections | existing CLIs/demos | Existing debug/README outputs | Shims preserve imports and wire compatibility | governance and wheel tests | Removal after approved migration milestones | -Proxy-observed transport intervals are useful for experiments but are not authoritative vLLM prefill/decode execution timings. Treat vLLM prefill/decode as not currently instrumented by canonical trace collection. +Proxy- and Instance-observed transport intervals are useful for experiments but are not authoritative vLLM prefill/decode execution timings. Treat vLLM prefill/decode as not currently instrumented by canonical trace collection. diff --git a/doc/research/issue-141-unified-observability.md b/doc/research/issue-141-unified-observability.md index f3a6bad..27862a2 100644 --- a/doc/research/issue-141-unified-observability.md +++ b/doc/research/issue-141-unified-observability.md @@ -37,3 +37,6 @@ ready-queue, and downstream transport stages become a process-local immutable trace. The canonical trace still is not client metadata, Instance input, or an authoritative account of vLLM or LMCache execution. This increment therefore does not close the broader Issue #141 research program. + + +Issue #185 extends the internal propagation boundary from Proxy to Instance for the exact reserved header vocabulary only. Instance now records sampled request-local, process-local transport stages for completion, first non-empty streaming response, and stream interval, while preserving public request/response/SSE shapes and keeping Instance-to-vLLM requests free of trace headers. Cross-process aggregation, Scheduler stage collection, vLLM/LMCache provenance, authoritative prefill/decode, external export, persistence, and debug trace retrieval remain future work under the umbrella. diff --git a/instance/instance_api.py b/instance/instance_api.py index e28e84b..dca927b 100644 --- a/instance/instance_api.py +++ b/instance/instance_api.py @@ -34,6 +34,9 @@ from instance.pclient.proxy_client import ProxyControlClient from instance.capability_builder import build_instance_capability from core.instance_capability import capability_fingerprint +from cacheroute.observability import resolve_observability_startup +from cacheroute.observability.startup import ObservabilityStartupConfig +from instance.observability import collect_non_streaming, collect_streaming, start_instance_trace_session PROXY_CP_URL = os.environ.get("PROXY_CP_URL", config.PROXY_CP_URL).rstrip("/") @@ -169,6 +172,17 @@ async def lifespan(app: FastAPI): cp_host = os.environ.get("INSTANCE_CP_HOST", config.INSTANCE_CP_HOST) cp_port = int(os.environ.get("INSTANCE_CP_PORT", config.INSTANCE_CP_PORT)) logger = logging.getLogger("instance") + observability_config = resolve_observability_startup( + os.environ.get("CACHEROUTE_RUNTIME_PROFILE"), + os.environ.get("CACHEROUTE_TRACE_SAMPLE_RATE"), + v1_available=False, + ) + app.state._observability_config = observability_config # type: ignore + if observability_config.sample_rate_warning_reason is not None: + logger.warning( + "[Trace] instance startup warning reason=%s", + observability_config.sample_rate_warning_reason, + ) capabilities = build_instance_capability() local_capability_fingerprint = capability_fingerprint(capabilities) @@ -388,6 +402,23 @@ async def _vllm_text_completion(payload: Dict[str, Any]) -> Dict[str, Any]: return _safe_json_from_bytes(content_bytes) + +def _instance_observability_config(request: FastAPIRequest) -> ObservabilityStartupConfig: + return getattr( + request.app.state, + "_observability_config", + resolve_observability_startup(None, None, v1_available=False), + ) + +def _instance_trace_session(request: FastAPIRequest, endpoint_label: str): + clock = getattr(request.app.state, "_trace_clock", None) + return start_instance_trace_session( + request.headers, + _instance_observability_config(request), + clock=clock, + endpoint_label=endpoint_label, + ) + # ======================= Scheduler method routes ======================= @instance.post("/v1/chat/completions") async def instance_chat_completions(request: FastAPIRequest): @@ -410,14 +441,19 @@ async def instance_chat_completions(request: FastAPIRequest): print(f"[Instance] stream={stream}") # Streaming: always use SSE byte streams. + session = _instance_trace_session(request, "instance_chat_completions") + if stream: async def event_stream(): - async for chunk in _vllm_stream_chat(payload): + async for chunk in collect_streaming(session, _vllm_stream_chat(payload)): yield chunk - return StreamingResponse(event_stream(), media_type="text/event-stream") + response = StreamingResponse(event_stream(), media_type="text/event-stream") + request.state._cacheroute_trace_session = session + return response - resp_json = await _vllm_chat_completion(payload) + resp_json = await collect_non_streaming(session, lambda: _vllm_chat_completion(payload)) + request.state._cacheroute_trace_session = session return JSONResponse(content=resp_json) @@ -436,5 +472,7 @@ async def instance_completions(request: FastAPIRequest): content={"error": "invalid_json", "detail": str(e)}, ) - resp_json = await _vllm_text_completion(payload) + session = _instance_trace_session(request, "instance_completions") + resp_json = await collect_non_streaming(session, lambda: _vllm_text_completion(payload)) + request.state._cacheroute_trace_session = session return JSONResponse(content=resp_json) diff --git a/instance/observability.py b/instance/observability.py new file mode 100644 index 0000000..0925d62 --- /dev/null +++ b/instance/observability.py @@ -0,0 +1,220 @@ +"""Instance-local observability session helpers. + +This module adapts dependency-light canonical tracing primitives to the +transitional Instance FastAPI handlers without creating service-global trace +state or exposing traces on public responses. +""" +from __future__ import annotations + +import asyncio +from datetime import timedelta +import logging +from dataclasses import dataclass +from uuid import uuid4 +from collections.abc import Mapping +from typing import AsyncGenerator, Awaitable, Callable, Any + +from cacheroute.contracts.v1 import ContractErrorDetail, OutcomeCode +from cacheroute.observability import TraceCollector, create_trace_context, decode_trace_headers +from cacheroute.observability.clock import SystemTraceClock, TraceClock +from cacheroute.observability.propagation import RESERVED_TRACE_HEADERS, REQUEST_ID_HEADER, TracePropagationError +from cacheroute.observability.startup import ObservabilityStartupConfig +from cacheroute.observability.v1 import RequestTrace, TraceContext, TraceProvenance, TraceStageName +from cacheroute.observability.v1.enums import TraceComponent +from cacheroute.runtime import RuntimeProfile + +logger = logging.getLogger("instance") +_NON_STREAMING = "non_streaming_request" +_EMPTY_STREAM_REASON = "stream_ended_before_decode" +_DOWNSTREAM_FAILED = ContractErrorDetail(code=OutcomeCode.FAILED, message="instance downstream request failed") +_EMPTY_STREAM = ContractErrorDetail(code=OutcomeCode.FAILED, message="instance stream ended before first response") +_REQUEST_CANCELLED = ContractErrorDetail(code=OutcomeCode.CANCELLED, message="instance request cancelled") + + +def local_request_id() -> str: + return f"req_{uuid4().hex}" + + +@dataclass(frozen=True) +class InstanceTraceResult: + context: TraceContext + accepted_propagated: bool + fallback_reason: str | None + + +@dataclass +class InstanceTraceSession: + context: TraceContext + collector: TraceCollector | None + provenance: TraceProvenance | None + request_trace: RequestTrace | None = None + completion_stage_id: str | None = None + first_token_stage_id: str | None = None + decode_stage_id: str | None = None + + def start_completion(self) -> None: + if self.collector is None or self.provenance is None: + return + try: + self.completion_stage_id = self.collector.start_stage(TraceStageName.COMPLETION, self.provenance) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage start failed reason=collector_state_invalid") + + def start_first_response(self) -> None: + if self.collector is None or self.provenance is None: + return + try: + self.first_token_stage_id = self.collector.start_stage( + TraceStageName.FIRST_TOKEN, self.provenance, parent_stage_id=self.completion_stage_id + ) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage start failed reason=collector_state_invalid") + + def finish_first_response_and_start_decode(self) -> None: + self.finish_stage("first_token_stage_id", OutcomeCode.SUCCESS) + if self.collector is None or self.provenance is None: + return + try: + self.decode_stage_id = self.collector.start_stage( + TraceStageName.DECODE, self.provenance, parent_stage_id=self.completion_stage_id + ) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage start failed reason=collector_state_invalid") + + def skip_non_streaming(self) -> None: + if self.collector is None or self.provenance is None: + return + for name in (TraceStageName.FIRST_TOKEN, TraceStageName.DECODE): + try: + self.collector.skip_stage(name, self.provenance, reason=_NON_STREAMING, parent_stage_id=self.completion_stage_id) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage skip failed reason=collector_state_invalid") + + def skip_empty_decode(self) -> None: + if self.collector is None or self.provenance is None: + return + try: + self.collector.skip_stage(TraceStageName.DECODE, self.provenance, reason=_EMPTY_STREAM_REASON, parent_stage_id=self.completion_stage_id) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage skip failed reason=collector_state_invalid") + + def finish_stage(self, attr: str, outcome: OutcomeCode, error: ContractErrorDetail | None = None) -> None: + if self.collector is None: + return + stage_id = getattr(self, attr) + try: + self.collector.finish_stage(stage_id, outcome=outcome, error=error) + setattr(self, attr, None) + except (ValueError, TypeError): + logger.warning("[Trace] instance stage finalization failed reason=collector_state_invalid") + + def finalize(self, outcome: OutcomeCode, error: ContractErrorDetail | None = None) -> None: + if self.collector is None or self.request_trace is not None: + return + try: + self.request_trace = self.collector.export(outcome=outcome, error=error) + except (ValueError, TypeError): + logger.warning("[Trace] instance finalization failed reason=collector_state_invalid") + + +def resolve_instance_context( + headers: Mapping[str, str], config: ObservabilityStartupConfig, *, clock: TraceClock | None = None +) -> InstanceTraceResult: + active_clock = clock or SystemTraceClock() + fallback_request_id: str | None = None + try: + propagated = decode_trace_headers(headers, clock=active_clock) + fallback_request_id = propagated.request_id + if propagated.runtime_profile is not config.runtime_profile: + raise TracePropagationError("profile_mismatch") + return InstanceTraceResult(propagated, True, None) + except TracePropagationError as exc: + reason = exc.reason + try: + complete = decode_trace_headers(headers, clock=active_clock, max_age=timedelta.max) + fallback_request_id = complete.request_id + except TracePropagationError: + raw = headers.get(REQUEST_ID_HEADER) if hasattr(headers, "get") else None + fallback_request_id = str(raw) if raw and len(str(raw)) <= 128 else None + request_id = fallback_request_id or local_request_id() + context = create_trace_context( + request_id, RuntimeProfile.normalize(config.runtime_profile), + sample_rate=config.trace_sample_rate, clock=active_clock, + ) + logger.info("[Trace] instance context fallback reason=%s", reason) + return InstanceTraceResult(context, False, reason) + + +def start_instance_trace_session( + headers: Mapping[str, str], config: ObservabilityStartupConfig, *, clock: TraceClock | None = None, + endpoint_label: str = "instance_downstream", +) -> InstanceTraceSession: + active_clock = clock or SystemTraceClock() + result = resolve_instance_context(headers, config, clock=active_clock) + if not result.context.sampled: + return InstanceTraceSession(result.context, None, None) + provenance = TraceProvenance( + source_component=TraceComponent.INSTANCE, + runtime_profile=result.context.runtime_profile, + captured_at=active_clock.utc_now(), + source_endpoint=endpoint_label, + ) + return InstanceTraceSession( + result.context, + TraceCollector(result.context, clock=active_clock), + provenance, + ) + + +async def collect_non_streaming(session: InstanceTraceSession, call: Callable[[], Awaitable[Any]]) -> Any: + session.start_completion() + session.skip_non_streaming() + try: + result = await call() + except asyncio.CancelledError: + session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + raise + except Exception: + session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finalize(OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + raise + session.finish_stage("completion_stage_id", OutcomeCode.SUCCESS) + session.finalize(OutcomeCode.SUCCESS) + return result + + +async def collect_streaming(session: InstanceTraceSession, stream: AsyncGenerator[bytes, None]) -> AsyncGenerator[bytes, None]: + session.start_completion() + session.start_first_response() + seen_first = False + try: + async for chunk in stream: + if chunk and not seen_first: + seen_first = True + session.finish_first_response_and_start_decode() + yield chunk + except asyncio.CancelledError: + session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + raise + except Exception: + session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finish_stage("decode_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finalize(OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + raise + if seen_first: + session.finish_stage("decode_stage_id", OutcomeCode.SUCCESS) + session.finish_stage("completion_stage_id", OutcomeCode.SUCCESS) + session.finalize(OutcomeCode.SUCCESS) + else: + session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) + session.skip_empty_decode() + session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) + session.finalize(OutcomeCode.FAILED, _EMPTY_STREAM) + + +__all__ = ["InstanceTraceSession", "resolve_instance_context", "start_instance_trace_session", "collect_non_streaming", "collect_streaming", "local_request_id"] diff --git a/proxy/queue/manager.py b/proxy/queue/manager.py index fcc022d..dfd563b 100644 --- a/proxy/queue/manager.py +++ b/proxy/queue/manager.py @@ -13,6 +13,8 @@ from proxy.metrics.queue_predictor import queue_predictor, decode_tpot_predictor, predict_redis_pull_ms from proxy.resource import p_control_plane from cacheroute.contracts.v1 import ContractErrorDetail, OutcomeCode +from cacheroute.observability import encode_trace_headers +from cacheroute.observability.propagation import RESERVED_TRACE_HEADERS from cacheroute.observability.v1 import TraceStageName from .task import ProxyTask @@ -1257,10 +1259,19 @@ async def _ready_worker_loop(self, instance_id: str, worker_idx: int) -> None: ) seen_first_chunk = False + extra_headers = None + if task.trace_context is not None: + try: + encoded_headers = encode_trace_headers(task.trace_context) + extra_headers = {name: encoded_headers[name] for name in RESERVED_TRACE_HEADERS} + except (ValueError, TypeError): + logger.warning("[Trace] propagation encode failed reason=context_invalid") + async for chunk in forward_request( url=target_url, data=task.instance_body, use_chunked=use_chunked, + extra_headers=extra_headers, ): if chunk: if not seen_first_chunk: diff --git a/test/observability/test_instance_observability.py b/test/observability/test_instance_observability.py new file mode 100644 index 0000000..9989597 --- /dev/null +++ b/test/observability/test_instance_observability.py @@ -0,0 +1,131 @@ +from datetime import datetime, timedelta, timezone +import asyncio +import json +import re + +import pytest + +from cacheroute.contracts.v1 import OutcomeCode +from cacheroute.observability import ManualTraceClock, create_trace_context, encode_trace_headers +from cacheroute.observability.propagation import RESERVED_TRACE_HEADERS, TRACE_ID_HEADER +from cacheroute.observability.startup import resolve_observability_startup +from cacheroute.observability.v1 import TraceStageName +from cacheroute.observability.v1.enums import TraceComponent +from cacheroute.runtime import RuntimeProfile +from instance.observability import collect_non_streaming, collect_streaming, resolve_instance_context, start_instance_trace_session + +NOW = datetime(2026, 8, 5, tzinfo=timezone.utc) +TRACE_ID = "trace_" + "1" * 32 + + +def cfg(rate="1.0", profile="legacy"): + return resolve_observability_startup(profile, rate, v1_available=False) + + +def headers(): + return encode_trace_headers(create_trace_context("scheduler-1", RuntimeProfile.LEGACY, sample_rate=1.0, clock=ManualTraceClock(NOW), trace_id=TRACE_ID)) + + +def test_instance_context_accepts_valid_and_falls_back_for_missing_malformed_stale_and_mismatch(monkeypatch): + clock = ManualTraceClock(NOW) + accepted = resolve_instance_context(headers(), cfg(), clock=clock) + assert accepted.accepted_propagated is True + assert accepted.context.trace_id == TRACE_ID + assert accepted.context.request_id == "scheduler-1" + + missing = resolve_instance_context({}, cfg(), clock=clock) + assert missing.accepted_propagated is False + assert re.fullmatch(r"req_[0-9a-f]{32}", missing.context.request_id) + + malformed = dict(headers(), **{TRACE_ID_HEADER: "bad"}) + bad = resolve_instance_context(malformed, cfg(), clock=clock) + assert bad.fallback_reason == "trace_id_invalid" + + old = ManualTraceClock(NOW + timedelta(minutes=6)) + stale = resolve_instance_context(headers(), cfg(), clock=old) + assert stale.context.request_id == "scheduler-1" + assert stale.accepted_propagated is False + + mismatch = resolve_instance_context(headers(), cfg(profile="test_mock"), clock=clock) + assert mismatch.fallback_reason == "profile_mismatch" + assert mismatch.context.request_id == "scheduler-1" + + +def test_startup_resolution_defaults_and_invalid_sample_rate(): + default = resolve_observability_startup(None, None, v1_available=False) + assert default.runtime_profile is RuntimeProfile.LEGACY + invalid = resolve_observability_startup("auto", "bad", v1_available=False) + assert invalid.runtime_profile is RuntimeProfile.LEGACY + assert invalid.trace_sample_rate == 0.0 + assert invalid.sample_rate_warning_reason == "trace_sample_rate_malformed" + + +def test_unsampled_session_collects_no_trace(): + session = start_instance_trace_session(headers(), cfg(rate="0.0"), clock=ManualTraceClock(NOW)) + assert session.collector is None + + +@pytest.mark.parametrize("stream_chunks", [[b"data: one\n\n", b"data: [DONE]\n\n"], [b"", b"data: one\n\n"]]) +def test_streaming_timing_order_parent_provenance_and_shape(stream_chunks): + async def scenario(): + clock = ManualTraceClock(NOW) + session = start_instance_trace_session(headers(), cfg(), clock=clock) + async def stream(): + for chunk in stream_chunks: + clock.advance(nanoseconds=10) + yield chunk + clock.advance(nanoseconds=5) + out = [] + async for chunk in collect_streaming(session, stream()): + out.append(chunk) + assert out == stream_chunks + trace = session.request_trace + assert trace.outcome is OutcomeCode.SUCCESS + stages = trace.stages + assert [s.name for s in stages] == [TraceStageName.COMPLETION, TraceStageName.FIRST_TOKEN, TraceStageName.DECODE] + assert stages[1].parent_stage_id == stages[0].stage_id + assert stages[2].parent_stage_id == stages[0].stage_id + assert all(s.provenance.source_component is TraceComponent.INSTANCE for s in stages) + assert TraceStageName.VLLM_PREFILL not in {s.name for s in stages} + asyncio.run(scenario()) + + +def test_non_streaming_skips_and_completion_payload_unchanged(): + async def scenario(): + clock = ManualTraceClock(NOW) + session = start_instance_trace_session(headers(), cfg(), clock=clock) + payload = {"choices": [{"message": {"content": "secret generated"}}]} + result = await collect_non_streaming(session, lambda: asyncio.sleep(0, result=payload)) + assert result is payload + stages = session.request_trace.stages + assert [s.name for s in stages] == [TraceStageName.COMPLETION, TraceStageName.FIRST_TOKEN, TraceStageName.DECODE] + assert {s.skip_reason for s in stages if s.name is not TraceStageName.COMPLETION} == {"non_streaming_request"} + dumped = session.request_trace.model_dump_json() + assert "secret generated" not in dumped + asyncio.run(scenario()) + + +def test_failure_empty_and_cancelled_paths_have_bounded_errors(): + async def before_failure(): + clock = ManualTraceClock(NOW) + session = start_instance_trace_session(headers(), cfg(), clock=clock) + async def stream(): + raise RuntimeError("raw boom prompt Authorization") + yield b"" + with pytest.raises(RuntimeError): + async for _ in collect_streaming(session, stream()): + pass + assert session.request_trace.error.message == "instance downstream request failed" + assert "raw boom" not in session.request_trace.model_dump_json() + + async def empty(): + session = start_instance_trace_session(headers(), cfg(), clock=ManualTraceClock(NOW)) + async def stream(): + if False: + yield b"" + async for _ in collect_streaming(session, stream()): + pass + assert session.request_trace.error.message == "instance stream ended before first response" + + asyncio.run(before_failure()) + asyncio.run(empty()) diff --git a/test/observability/test_scheduler_proxy_production_paths.py b/test/observability/test_scheduler_proxy_production_paths.py index 7456a1b..5fd3a97 100644 --- a/test/observability/test_scheduler_proxy_production_paths.py +++ b/test/observability/test_scheduler_proxy_production_paths.py @@ -97,8 +97,8 @@ async def _run_real_queue(monkeypatch, task, chunks=(), *, fail_at=None, cancel_ manager._READY_DEQUEUE_INTERVAL_S = 0.0 task.trace["preexisting"] = 7 - async def fake_forward_request(url, data, use_chunked=False): - calls.append({"url": url, "data": data.copy(), "use_chunked": use_chunked}) + async def fake_forward_request(url, data, use_chunked=False, extra_headers=None): + calls.append({"url": url, "data": data.copy(), "use_chunked": use_chunked, "headers": dict(extra_headers or {})}) if fail_at == "before_first": raise RuntimeError("raw downstream boom should not enter trace") for index, chunk in enumerate(chunks): @@ -144,8 +144,8 @@ async def dispatch_wait(wait_task, instance_id): wait_task.has_started_forward = True monkeypatch.setattr(manager, "_wait_dispatch_turn", dispatch_wait) - async def timed_forward_request(url, data, use_chunked=False): - calls.append({"url": url, "data": data.copy(), "use_chunked": use_chunked}) + async def timed_forward_request(url, data, use_chunked=False, extra_headers=None): + calls.append({"url": url, "data": data.copy(), "use_chunked": use_chunked, "headers": dict(extra_headers or {})}) clock.advance(nanoseconds=11_000_000) yield b"data: token\n\n" clock.advance(nanoseconds=13_000_000) @@ -243,7 +243,7 @@ async def scenario(): assert calls and calls[0]["url"] == f"http://{task.instance_host}:{task.instance_port}{task.url_path}" assert calls[0]["data"] == original_body assert calls[0]["use_chunked"] is (endpoint == "chat/completions") - assert not (set(calls[0].get("headers", {})) & set(RESERVED_TRACE_HEADERS)) + assert set(calls[0].get("headers", {})) == set(RESERVED_TRACE_HEADERS) assert task.request_trace.outcome is expected_outcome assert task.request_trace.error == expected_error _assert_no_running_and_valid_refs(task.request_trace) diff --git a/test/test_repository_governance.py b/test/test_repository_governance.py index cb09981..499daab 100644 --- a/test/test_repository_governance.py +++ b/test/test_repository_governance.py @@ -296,12 +296,14 @@ def test_internal_trace_header_vocabulary_has_single_owner_and_services_import_i } for literal in header_literals: owners = {path.relative_to(ROOT) for path in _tracked_files() if path.suffix == ".py" and literal in path.read_text(encoding="utf-8")} - assert owners <= {owner.relative_to(ROOT), Path("test/observability/test_propagation.py"), Path("test/observability/test_scheduler_proxy_production_paths.py"), Path("test/test_repository_governance.py")} + assert owners <= {owner.relative_to(ROOT), Path("test/observability/test_propagation.py"), Path("test/observability/test_scheduler_proxy_production_paths.py"), Path("test/observability/test_instance_observability.py"), Path("test/test_repository_governance.py")} for relative in (Path("scheduler/scheduler.py"), Path("proxy/proxy.py")): text = (ROOT / relative).read_text(encoding="utf-8") assert "cacheroute.observability" in text assert "extra_headers=extra_headers" in (ROOT / "scheduler/scheduler.py").read_text(encoding="utf-8") - assert "extra_headers" not in (ROOT / "proxy/queue/manager.py").read_text(encoding="utf-8") + manager_text = (ROOT / "proxy/queue/manager.py").read_text(encoding="utf-8") + assert "encode_trace_headers" in manager_text + assert "extra_headers=extra_headers" in manager_text def test_observability_does_not_duplicate_canonical_model_classes(): From 248e96664b45c62d59cd7dc90613b2ad5994d132 Mon Sep 17 00:00:00 2001 From: RickZ <121943721+rickisba@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:33:27 +0800 Subject: [PATCH 2/4] observability: finalize Instance trace terminal paths --- doc/architecture/observability-v1.md | 2 +- .../configuration-and-interfaces.md | 4 +- instance/instance_api.py | 9 +- instance/observability.py | 15 +- .../test_instance_observability.py | 180 +++++++++++++++++- .../test_scheduler_proxy_production_paths.py | 44 ++++- 6 files changed, 245 insertions(+), 9 deletions(-) diff --git a/doc/architecture/observability-v1.md b/doc/architecture/observability-v1.md index f5bb718..2c46272 100644 --- a/doc/architecture/observability-v1.md +++ b/doc/architecture/observability-v1.md @@ -59,7 +59,7 @@ ID allocated by the Scheduler remains authoritative. Client values using these names are never trusted; Authorization forwarding and the serialized payload remain unchanged. Propagation now reaches the selected Instance through the same exact reserved header set. Proxy encodes its accepted Scheduler context, or its Proxy-local fallback context, immediately before forwarding to Instance and overwrites every reserved value for that internal call. No arbitrary client trace header becomes authoritative and no JSON/base64 trace payload is added to the request body. Instance validates the complete reserved set, accepts matching fresh context, or creates a request-local fallback context using its resolved startup profile and sample rate. The reserved headers stop at Instance: Instance-to-vLLM forwarding does not include CacheRoute trace headers, W3C Trace Context, or any canonical trace object, and no trace is returned to a client. -Scheduler and Proxy each resolve `CACHEROUTE_RUNTIME_PROFILE` through an explicit lifespan startup helper, with `legacy` as the compatibility default. Because the current services do not have an implemented production v1 data path, `auto` resolves with `v1_available=False` and is stored as `legacy`; explicit `legacy`, `test/mock`, and `v1` remain valid metadata values, but no stored or propagated context can remain `auto`. A missing, malformed, stale, request-ID-mismatched, or profile-mismatched context causes a Proxy-local context to be created and never rejects an otherwise valid request. Profile metadata does not select runtime behavior. The propagation freshness rule retains the five-minute maximum age and also accepts only a bounded 30-second future clock skew between Scheduler and Proxy clocks. +Scheduler and Proxy each resolve `CACHEROUTE_RUNTIME_PROFILE` through an explicit lifespan startup helper, with `legacy` as the compatibility default. Because the current services do not have an implemented production v1 data path, `auto` resolves with `v1_available=False` and is stored as `legacy`; explicit `legacy`, `test/mock`, and `v1` remain valid metadata values, but no stored or propagated context can remain `auto`. A missing, malformed, stale, request-ID-mismatched, or profile-mismatched context causes a Proxy-local context to be created and never rejects an otherwise valid request. Profile metadata does not select runtime behavior. The internal propagation freshness rule retains the five-minute maximum age and accepts only a bounded 30-second future clock skew among Scheduler, Proxy, and Instance clocks. `CACHEROUTE_TRACE_SAMPLE_RATE` defaults to `0.0`. Invalid configuration fails closed to that value. Rates of zero and one disable or enable collection for diff --git a/doc/developer-handbook/configuration-and-interfaces.md b/doc/developer-handbook/configuration-and-interfaces.md index 8b14704..2a6114e 100644 --- a/doc/developer-handbook/configuration-and-interfaces.md +++ b/doc/developer-handbook/configuration-and-interfaces.md @@ -6,8 +6,8 @@ Catalog only developer-facing surfaces. Source definitions remain authoritative; | Name | Owner | Type | Exact default | Allowed values | Scope | Effect | Invalid/fallback behavior | Source | Validation | |---|---|---|---|---|---|---|---|---|---| -| `CACHEROUTE_RUNTIME_PROFILE` | Scheduler/Proxy/Instance observability | string | unset -> service passes default legacy startup resolution | `auto`, `legacy`, `v1`, `test/mock` plus compatibility aliases accepted by normalizer | startup | Resolves trace runtime metadata | Current Scheduler, Proxy, and Instance use `v1_available=False`; `auto` resolves to `legacy`; invalid values raise through runtime normalization | `scheduler/scheduler.py`, `proxy/proxy.py`, `src/cacheroute/observability/startup.py` | `test/observability` | -| `CACHEROUTE_TRACE_SAMPLE_RATE` | Scheduler/Proxy/Instance observability | float string | unset -> `0.0` | finite `0.0` through `1.0` | startup | Sets deterministic trace sampled flag | malformed/non-finite/out-of-range values fail closed to `0.0` and expose warning reason | `src/cacheroute/observability/propagation.py` | `test/observability` | +| `CACHEROUTE_RUNTIME_PROFILE` | Scheduler/Proxy/Instance observability | string | unset -> service passes default legacy startup resolution | `auto`, `legacy`, `v1`, `test/mock` plus compatibility aliases accepted by normalizer | startup | Resolves trace runtime metadata | Current Scheduler, Proxy, and Instance use `v1_available=False`; `auto` resolves to `legacy`; invalid values raise through runtime normalization | `scheduler/scheduler.py`, `proxy/proxy.py`, `instance/instance_api.py`, `src/cacheroute/observability/startup.py` | `test/observability` | +| `CACHEROUTE_TRACE_SAMPLE_RATE` | Scheduler/Proxy/Instance observability | float string | unset -> `0.0` | finite `0.0` through `1.0` | startup | Sets deterministic trace sampled flag | malformed/non-finite/out-of-range values fail closed to `0.0` and expose warning reason | `instance/instance_api.py`, `src/cacheroute/observability/propagation.py` | `test/observability` | | `SCHEDULER_LOG_FILE` | Scheduler | path/string | non-portable Transitional Legacy source default (repository-external absolute log path in `core/config.py`; override for portable deployments) | path-like string | startup | Scheduler log path/config value | falls back to source owner default | `core/config.py`, `scheduler/scheduler.py` | demos | | `SCHEDULER_VERBOSE_REQUEST_LOG` | Scheduler | int/string | `1` | intended `0` or `1` | request | Enables verbose request scheduling logs | compared to integer `1` in current source; string environment values do not match that equality | `core/config.py`, `scheduler/scheduler.py` | demos | | `SCHEDULER_MODEL_PATH` | Scheduler | path/string | non-portable Transitional Legacy source default from `core/config.py` (repository-external model path; override for portable deployments) | path/model identifier | startup | Model path used by scheduler tokenizer/model setup | falls back to source owner default | `scheduler/scheduler.py`, `core/config.py` | source checkout tests | diff --git a/instance/instance_api.py b/instance/instance_api.py index dca927b..7faa506 100644 --- a/instance/instance_api.py +++ b/instance/instance_api.py @@ -403,11 +403,18 @@ async def _vllm_text_completion(payload: Dict[str, Any]) -> Dict[str, Any]: +_NO_LIFESPAN_OBSERVABILITY_CONFIG = resolve_observability_startup( + None, None, v1_available=False +) + + def _instance_observability_config(request: FastAPIRequest) -> ObservabilityStartupConfig: + # Production always reads the immutable lifespan-owned object. The module + # constant keeps direct handler tests compatible without resolving per call. return getattr( request.app.state, "_observability_config", - resolve_observability_startup(None, None, v1_available=False), + _NO_LIFESPAN_OBSERVABILITY_CONFIG, ) def _instance_trace_session(request: FastAPIRequest, endpoint_label: str): diff --git a/instance/observability.py b/instance/observability.py index 0925d62..0b42694 100644 --- a/instance/observability.py +++ b/instance/observability.py @@ -17,7 +17,7 @@ from cacheroute.contracts.v1 import ContractErrorDetail, OutcomeCode from cacheroute.observability import TraceCollector, create_trace_context, decode_trace_headers from cacheroute.observability.clock import SystemTraceClock, TraceClock -from cacheroute.observability.propagation import RESERVED_TRACE_HEADERS, REQUEST_ID_HEADER, TracePropagationError +from cacheroute.observability.propagation import TracePropagationError from cacheroute.observability.startup import ObservabilityStartupConfig from cacheroute.observability.v1 import RequestTrace, TraceContext, TraceProvenance, TraceStageName from cacheroute.observability.v1.enums import TraceComponent @@ -131,11 +131,12 @@ def resolve_instance_context( except TracePropagationError as exc: reason = exc.reason try: + # A second decode may relax freshness only. It must still validate + # the complete canonical header set before its request ID is kept. complete = decode_trace_headers(headers, clock=active_clock, max_age=timedelta.max) fallback_request_id = complete.request_id except TracePropagationError: - raw = headers.get(REQUEST_ID_HEADER) if hasattr(headers, "get") else None - fallback_request_id = str(raw) if raw and len(str(raw)) <= 128 else None + fallback_request_id = None request_id = fallback_request_id or local_request_id() context = create_trace_context( request_id, RuntimeProfile.normalize(config.runtime_profile), @@ -200,6 +201,14 @@ async def collect_streaming(session: InstanceTraceSession, stream: AsyncGenerato session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) raise + except GeneratorExit: + # ``aclose()`` injects GeneratorExit at the current yield. Finalize + # the request-local trace, then preserve normal generator closure. + session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + raise except Exception: session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) session.finish_stage("decode_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) diff --git a/test/observability/test_instance_observability.py b/test/observability/test_instance_observability.py index 9989597..ee654ae 100644 --- a/test/observability/test_instance_observability.py +++ b/test/observability/test_instance_observability.py @@ -7,7 +7,11 @@ from cacheroute.contracts.v1 import OutcomeCode from cacheroute.observability import ManualTraceClock, create_trace_context, encode_trace_headers -from cacheroute.observability.propagation import RESERVED_TRACE_HEADERS, TRACE_ID_HEADER +from cacheroute.observability.propagation import ( + RESERVED_TRACE_HEADERS, + REQUEST_ID_HEADER, + TRACE_ID_HEADER, +) from cacheroute.observability.startup import resolve_observability_startup from cacheroute.observability.v1 import TraceStageName from cacheroute.observability.v1.enums import TraceComponent @@ -40,6 +44,7 @@ def test_instance_context_accepts_valid_and_falls_back_for_missing_malformed_sta malformed = dict(headers(), **{TRACE_ID_HEADER: "bad"}) bad = resolve_instance_context(malformed, cfg(), clock=clock) assert bad.fallback_reason == "trace_id_invalid" + assert re.fullmatch(r"req_[0-9a-f]{32}", bad.context.request_id) old = ManualTraceClock(NOW + timedelta(minutes=6)) stale = resolve_instance_context(headers(), cfg(), clock=old) @@ -51,6 +56,36 @@ def test_instance_context_accepts_valid_and_falls_back_for_missing_malformed_sta assert mismatch.context.request_id == "scheduler-1" +class _ConflictingHeaders(dict): + def items(self): + values = list(super().items()) + values.append((TRACE_ID_HEADER.upper(), "trace_" + "2" * 32)) + return values + + +@pytest.mark.parametrize( + "invalid_headers,reason", + [ + ({REQUEST_ID_HEADER: "safe-looking"}, "headers_incomplete"), + ({REQUEST_ID_HEADER: "unsafe request id!"}, "headers_incomplete"), + (dict(headers(), **{TRACE_ID_HEADER: "bad", REQUEST_ID_HEADER: "safe-looking"}), "trace_id_invalid"), + (dict(headers(), **{"x-cacheroute-trace-unknown": "value"}), "header_unknown"), + (_ConflictingHeaders(headers()), "header_conflict"), + ], +) +def test_invalid_header_sets_never_retain_raw_request_id_or_fail_output(invalid_headers, reason): + async def scenario(): + session = start_instance_trace_session(invalid_headers, cfg(), clock=ManualTraceClock(NOW)) + assert re.fullmatch(r"req_[0-9a-f]{32}", session.context.request_id) + output = {"choices": []} + assert await collect_non_streaming(session, lambda: asyncio.sleep(0, result=output)) is output + assert session.request_trace.outcome is OutcomeCode.SUCCESS + + result = resolve_instance_context(invalid_headers, cfg(), clock=ManualTraceClock(NOW)) + assert result.fallback_reason == reason + asyncio.run(scenario()) + + def test_startup_resolution_defaults_and_invalid_sample_rate(): default = resolve_observability_startup(None, None, v1_available=False) assert default.runtime_profile is RuntimeProfile.LEGACY @@ -129,3 +164,146 @@ async def stream(): asyncio.run(before_failure()) asyncio.run(empty()) + + +@pytest.mark.parametrize("after_first", [False, True]) +def test_streaming_asyncio_cancellation_finalizes_every_stage(after_first): + async def scenario(): + session = start_instance_trace_session(headers(), cfg(), clock=ManualTraceClock(NOW)) + gate = asyncio.Event() + + async def stream(): + if after_first: + yield b"data: first\n\n" + await gate.wait() + yield b"never" + + collected = collect_streaming(session, stream()) + if after_first: + assert await anext(collected) == b"data: first\n\n" + pending = asyncio.create_task(anext(collected)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert session.request_trace.outcome is OutcomeCode.CANCELLED + assert session.request_trace.error.message == "instance request cancelled" + assert all(stage.state.value != "running" for stage in session.request_trace.stages) + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("after_first", [False, True]) +def test_streaming_aclose_finalizes_as_cancellation(after_first): + async def scenario(): + session = start_instance_trace_session(headers(), cfg(), clock=ManualTraceClock(NOW)) + + async def stream(): + yield b"data: first\n\n" if after_first else b"" + yield b"never" + + collected = collect_streaming(session, stream()) + first = await anext(collected) + assert first == (b"data: first\n\n" if after_first else b"") + await collected.aclose() + assert session.request_trace.outcome is OutcomeCode.CANCELLED + assert session.request_trace.error.message == "instance request cancelled" + assert all(stage.state.value != "running" for stage in session.request_trace.stages) + stage = next( + item for item in session.request_trace.stages + if item.name is (TraceStageName.DECODE if after_first else TraceStageName.FIRST_TOKEN) + ) + assert stage.outcome is OutcomeCode.CANCELLED + + asyncio.run(scenario()) + + +def test_streaming_failure_after_first_response_is_bounded(): + async def scenario(): + session = start_instance_trace_session(headers(), cfg(), clock=ManualTraceClock(NOW)) + + async def stream(): + yield b"data: first\n\n" + raise RuntimeError("private downstream exception") + + collected = collect_streaming(session, stream()) + assert await anext(collected) == b"data: first\n\n" + with pytest.raises(RuntimeError): + await anext(collected) + assert session.request_trace.outcome is OutcomeCode.FAILED + assert all(stage.state.value != "running" for stage in session.request_trace.stages) + assert "private downstream exception" not in session.request_trace.model_dump_json() + + asyncio.run(scenario()) + + +def test_startup_configuration_is_reused_without_reinvoking_resolver(monkeypatch): + from types import SimpleNamespace + from instance import instance_api + + configured = cfg() + request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(_observability_config=configured))) + monkeypatch.setattr( + instance_api, + "resolve_observability_startup", + lambda *args, **kwargs: pytest.fail("request path reran startup resolver"), + ) + assert instance_api._instance_observability_config(request) is configured + assert instance_api._instance_observability_config(request) is configured + + +def test_real_vllm_forwarding_preserves_payload_and_sends_no_trace_headers(monkeypatch): + from instance import instance_api + + async def scenario(): + payload = {"model": "unit", "messages": [{"role": "user", "content": "secret"}], "stream": True} + calls = [] + + async def forwarded(url, data, use_chunked=False, **kwargs): + calls.append((url, data, use_chunked, kwargs)) + yield b"data: unchanged\n\n" + + monkeypatch.setattr(instance_api, "use_mock", False) + monkeypatch.setattr(instance_api, "vllm_base_url", "http://vllm") + monkeypatch.setattr(instance_api, "forward_request", forwarded) + assert [chunk async for chunk in instance_api._vllm_stream_chat(payload)] == [b"data: unchanged\n\n"] + assert calls == [("http://vllm/v1/chat/completions", payload, True, {})] + assert all(name not in calls[0][3] for name in RESERVED_TRACE_HEADERS) + assert "traceparent" not in calls[0][3] and "tracestate" not in calls[0][3] + + asyncio.run(scenario()) + + +def test_mock_streaming_chat_and_nonstreaming_chat_and_text_shapes_are_unchanged(monkeypatch): + from instance import instance_api + + async def scenario(): + stream_payload = {"stream": True, "messages": []} + chat_payload = {"stream": False, "messages": []} + text_payload = {"prompt": "kept"} + stream_bytes = [b"data: mock\n\n", b"data: [DONE]\n\n"] + chat_json = {"choices": [{"message": {"content": "mock"}}]} + text_json = {"choices": [{"text": "mock"}]} + + async def mock_stream(payload): + assert payload is stream_payload + for chunk in stream_bytes: + yield chunk + + async def mock_chat(payload): + assert payload is chat_payload + return chat_json + + async def mock_text(payload): + assert payload is text_payload + return text_json + + monkeypatch.setattr(instance_api, "use_mock", True) + monkeypatch.setattr(instance_api, "mock_chat_stream", mock_stream) + monkeypatch.setattr(instance_api, "mock_chat_completion", mock_chat) + monkeypatch.setattr(instance_api, "mock_text_completion", mock_text) + assert [chunk async for chunk in instance_api._vllm_stream_chat(stream_payload)] == stream_bytes + assert await instance_api._vllm_chat_completion(chat_payload) is chat_json + assert await instance_api._vllm_text_completion(text_payload) is text_json + + asyncio.run(scenario()) diff --git a/test/observability/test_scheduler_proxy_production_paths.py b/test/observability/test_scheduler_proxy_production_paths.py index 5fd3a97..1f73108 100644 --- a/test/observability/test_scheduler_proxy_production_paths.py +++ b/test/observability/test_scheduler_proxy_production_paths.py @@ -284,7 +284,10 @@ async def scenario(): task = _task_from_request(req) chunks = [b"data: token\\n\\n"] if cancel_after_first else [] - async def fake_forward_request(url, data, use_chunked=False): + forwarded_headers = [] + + async def fake_forward_request(url, data, use_chunked=False, extra_headers=None): + forwarded_headers.append(extra_headers) if not cancel_after_first: await asyncio.sleep(10) yield b"" @@ -303,6 +306,8 @@ async def fake_forward_request(url, data, use_chunked=False): await asyncio.gather(*manager._worker_tasks.values(), return_exceptions=True) assert task.request_trace.outcome is OutcomeCode.CANCELLED assert task.request_trace.error == _REQUEST_CANCELLED + assert len(forwarded_headers) == 1 + assert set(forwarded_headers[0]) == set(RESERVED_TRACE_HEADERS) _assert_no_running_and_valid_refs(task.request_trace) asyncio.run(scenario()) @@ -325,6 +330,9 @@ async def scenario(): def test_cpu_only_scheduler_to_proxy_integration_with_real_workers(monkeypatch): async def scenario(): + from instance.observability import collect_streaming, start_instance_trace_session + from cacheroute.observability.startup import resolve_observability_startup + scheduler_clock = ManualTraceClock(NOW) req = _scheduler_request(123) def fixed_context(request_id, runtime_profile, *, sample_rate=0.0, clock=None, trace_id=None): @@ -350,6 +358,40 @@ def fixed_context(request_id, runtime_profile, *, sample_rate=0.0, clock=None, t assert task.request_trace.context.trace_id == trace_id assert calls[0]["data"] == task.instance_body assert not (set(calls[0]["data"]) & set(RESERVED_TRACE_HEADERS)) + assert set(calls[0]["headers"]) == set(RESERVED_TRACE_HEADERS) + assert calls[0]["headers"] == encode_trace_headers(context) + + instance_clock = ManualTraceClock(NOW) + instance_session = start_instance_trace_session( + calls[0]["headers"], + resolve_observability_startup("legacy", "1.0", v1_available=False), + clock=instance_clock, + ) + response_chunks = [b"data: token\\n\\n"] + + async def mocked_vllm(): + for chunk in response_chunks: + yield chunk + + assert [chunk async for chunk in collect_streaming(instance_session, mocked_vllm())] == response_chunks + assert instance_session.context.request_id == str(req.Request_ID) + assert instance_session.context.trace_id == trace_id + assert instance_session.request_trace is not task.request_trace + assert instance_session.request_trace.context == task.request_trace.context + assert all(stage.provenance.source_component is TraceComponent.PROXY for stage in task.request_trace.stages) + assert all(stage.provenance.source_component is TraceComponent.INSTANCE for stage in instance_session.request_trace.stages) + assert response_chunks == [b"data: token\\n\\n"] + + replacement = create_trace_context( + "local-request", RuntimeProfile.LEGACY, sample_rate=1.0, + clock=proxy_clock, trace_id="trace_" + "b" * 32, + ) + replacement_session = start_instance_trace_session( + encode_trace_headers(replacement), + resolve_observability_startup("legacy", "1.0", v1_available=False), + clock=instance_clock, + ) + assert replacement_session.context.trace_id == replacement.trace_id asyncio.run(scenario()) From 8ee6323af73bdb5e7ab746e4d598ab596fb56721 Mon Sep 17 00:00:00 2001 From: RickZ <121943721+rickisba@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:01:56 +0800 Subject: [PATCH 3/4] observability: close Instance streaming boundaries --- instance/instance_api.py | 25 ++-- instance/observability.py | 77 ++++++----- src/cacheroute/observability/propagation.py | 4 +- .../test_instance_observability.py | 123 +++++++++++++++++- test/test_repository_governance.py | 57 ++++++++ 5 files changed, 234 insertions(+), 52 deletions(-) diff --git a/instance/instance_api.py b/instance/instance_api.py index 7faa506..2c74d2b 100644 --- a/instance/instance_api.py +++ b/instance/instance_api.py @@ -39,6 +39,7 @@ from instance.observability import collect_non_streaming, collect_streaming, start_instance_trace_session +logger = logging.getLogger("instance") PROXY_CP_URL = os.environ.get("PROXY_CP_URL", config.PROXY_CP_URL).rstrip("/") INSTANCE_ADVERTISE_HOST = os.environ.get("INSTANCE_ADVERTISE_HOST", config.INSTANCE_HOST) INSTANCE_ADVERTISE_PORT = int(os.environ.get("INSTANCE_ADVERTISE_PORT", os.environ.get("INSTANCE_PORT", config.INSTANCE_PORT))) @@ -355,11 +356,18 @@ async def _vllm_stream_chat(payload: Dict[str, Any]) -> AsyncGenerator[bytes, No url = f"{vllm_base_url}/v1/chat/completions" # Forward the OpenAI-style body from Proxy directly. upstream_stream = forward_request(url, data=payload, use_chunked=True) # type: ignore - - async for chunk in upstream_stream: - # Do not parse or rewrite; pass through directly. - if chunk: - yield chunk + try: + async for chunk in upstream_stream: + # Do not parse or rewrite; pass through directly. + if chunk: + yield chunk + finally: + close = getattr(upstream_stream, "aclose", None) + if close is not None: + try: + await close() + except BaseException: + logger.warning("[Trace] instance stream cleanup failed reason=vllm_stream_close_failed") async def _vllm_chat_completion(payload: Dict[str, Any]) -> Dict[str, Any]: @@ -451,11 +459,8 @@ async def instance_chat_completions(request: FastAPIRequest): session = _instance_trace_session(request, "instance_chat_completions") if stream: - async def event_stream(): - async for chunk in collect_streaming(session, _vllm_stream_chat(payload)): - yield chunk - - response = StreamingResponse(event_stream(), media_type="text/event-stream") + collected_stream = collect_streaming(session, _vllm_stream_chat(payload)) + response = StreamingResponse(collected_stream, media_type="text/event-stream") request.state._cacheroute_trace_session = session return response diff --git a/instance/observability.py b/instance/observability.py index 0b42694..8ea5bb8 100644 --- a/instance/observability.py +++ b/instance/observability.py @@ -190,40 +190,49 @@ async def collect_streaming(session: InstanceTraceSession, stream: AsyncGenerato session.start_first_response() seen_first = False try: - async for chunk in stream: - if chunk and not seen_first: - seen_first = True - session.finish_first_response_and_start_decode() - yield chunk - except asyncio.CancelledError: - session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - raise - except GeneratorExit: - # ``aclose()`` injects GeneratorExit at the current yield. Finalize - # the request-local trace, then preserve normal generator closure. - session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) - raise - except Exception: - session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) - session.finish_stage("decode_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) - session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) - session.finalize(OutcomeCode.FAILED, _DOWNSTREAM_FAILED) - raise - if seen_first: - session.finish_stage("decode_stage_id", OutcomeCode.SUCCESS) - session.finish_stage("completion_stage_id", OutcomeCode.SUCCESS) - session.finalize(OutcomeCode.SUCCESS) - else: - session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) - session.skip_empty_decode() - session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) - session.finalize(OutcomeCode.FAILED, _EMPTY_STREAM) + try: + async for chunk in stream: + if chunk and not seen_first: + seen_first = True + session.finish_first_response_and_start_decode() + yield chunk + except asyncio.CancelledError: + session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + raise + except GeneratorExit: + # ``aclose()`` injects GeneratorExit at the current yield. Finalize + # the request-local trace, then preserve normal generator closure. + session.finish_stage("first_token_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("decode_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finish_stage("completion_stage_id", OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + session.finalize(OutcomeCode.CANCELLED, _REQUEST_CANCELLED) + raise + except Exception: + session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finish_stage("decode_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + session.finalize(OutcomeCode.FAILED, _DOWNSTREAM_FAILED) + raise + if seen_first: + session.finish_stage("decode_stage_id", OutcomeCode.SUCCESS) + session.finish_stage("completion_stage_id", OutcomeCode.SUCCESS) + session.finalize(OutcomeCode.SUCCESS) + else: + session.finish_stage("first_token_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) + session.skip_empty_decode() + session.finish_stage("completion_stage_id", OutcomeCode.FAILED, _EMPTY_STREAM) + session.finalize(OutcomeCode.FAILED, _EMPTY_STREAM) + finally: + close = getattr(stream, "aclose", None) + if close is not None: + try: + await close() + except BaseException: + # Cleanup is subordinate to the existing response or exception. + logger.warning("[Trace] instance stream cleanup failed reason=downstream_close_failed") __all__ = ["InstanceTraceSession", "resolve_instance_context", "start_instance_trace_session", "collect_non_streaming", "collect_streaming", "local_request_id"] diff --git a/src/cacheroute/observability/propagation.py b/src/cacheroute/observability/propagation.py index d590522..cb6ec68 100644 --- a/src/cacheroute/observability/propagation.py +++ b/src/cacheroute/observability/propagation.py @@ -1,4 +1,4 @@ -"""Canonical Scheduler-to-Proxy trace-context propagation. +"""Canonical internal CacheRoute trace-context propagation. This module deliberately contains no service configuration or I/O. Services resolve startup settings and pass the resulting values into these helpers. @@ -34,7 +34,7 @@ ) _TRACE_ID = re.compile(r"^trace_[0-9a-f]{32}$") _MAX_AGE = timedelta(minutes=5) -# Scheduler and Proxy clocks may differ slightly during startup or host sync. +# Scheduler, Proxy, and Instance clocks may differ during startup or host sync. # Accept only this bounded future skew and retain the maximum-age freshness rule. _FUTURE_SKEW_TOLERANCE = timedelta(seconds=30) diff --git a/test/observability/test_instance_observability.py b/test/observability/test_instance_observability.py index ee654ae..b0e85fd 100644 --- a/test/observability/test_instance_observability.py +++ b/test/observability/test_instance_observability.py @@ -2,6 +2,7 @@ import asyncio import json import re +from types import SimpleNamespace import pytest @@ -122,6 +123,10 @@ async def stream(): assert stages[2].parent_stage_id == stages[0].stage_id assert all(s.provenance.source_component is TraceComponent.INSTANCE for s in stages) assert TraceStageName.VLLM_PREFILL not in {s.name for s in stages} + expected_first = 10 if stream_chunks[0] else 20 + assert stages[1].elapsed_ns == expected_first + assert stages[2].elapsed_ns == 25 - expected_first + assert stages[0].elapsed_ns == 25 asyncio.run(scenario()) @@ -256,20 +261,126 @@ def test_real_vllm_forwarding_preserves_payload_and_sends_no_trace_headers(monke from instance import instance_api async def scenario(): - payload = {"model": "unit", "messages": [{"role": "user", "content": "secret"}], "stream": True} + stream_payload = {"model": "unit", "messages": [{"role": "user", "content": "secret"}], "stream": True} + chat_payload = {"model": "unit", "messages": [], "stream": False} + text_payload = {"model": "unit", "prompt": "secret", "stream": False} calls = [] async def forwarded(url, data, use_chunked=False, **kwargs): calls.append((url, data, use_chunked, kwargs)) - yield b"data: unchanged\n\n" + if use_chunked: + yield b"data: unchanged\n\n" + else: + yield json.dumps({"choices": []}).encode() monkeypatch.setattr(instance_api, "use_mock", False) monkeypatch.setattr(instance_api, "vllm_base_url", "http://vllm") monkeypatch.setattr(instance_api, "forward_request", forwarded) - assert [chunk async for chunk in instance_api._vllm_stream_chat(payload)] == [b"data: unchanged\n\n"] - assert calls == [("http://vllm/v1/chat/completions", payload, True, {})] - assert all(name not in calls[0][3] for name in RESERVED_TRACE_HEADERS) - assert "traceparent" not in calls[0][3] and "tracestate" not in calls[0][3] + stream_bytes = [chunk async for chunk in instance_api._vllm_stream_chat(stream_payload)] + chat_json = await instance_api._vllm_chat_completion(chat_payload) + text_json = await instance_api._vllm_text_completion(text_payload) + assert stream_bytes == [b"data: unchanged\n\n"] + assert chat_json == text_json == {"choices": []} + assert calls == [ + ("http://vllm/v1/chat/completions", stream_payload, True, {}), + ("http://vllm/v1/chat/completions", chat_payload, False, {}), + ("http://vllm/v1/completions", text_payload, False, {}), + ] + public_values = stream_bytes + [json.dumps(chat_json).encode(), json.dumps(text_json).encode()] + for value in public_values: + assert TRACE_ID.encode() not in value + assert b"RequestTrace" not in value + assert b"trace_id" not in value + assert all(name.encode() not in value for name in RESERVED_TRACE_HEADERS) + for _url, _payload, _chunked, kwargs in calls: + assert kwargs == {} + assert all(name not in kwargs for name in RESERVED_TRACE_HEADERS) + assert "traceparent" not in kwargs and "tracestate" not in kwargs + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("after_first", [False, True]) +def test_streaming_response_body_close_finalizes_and_closes_downstream(monkeypatch, after_first): + from instance import instance_api + + async def scenario(): + clock = ManualTraceClock(NOW) + app = SimpleNamespace(state=SimpleNamespace( + _observability_config=cfg(), + _trace_clock=clock, + )) + request = SimpleNamespace( + app=app, + headers=headers(), + state=SimpleNamespace(), + json=lambda: asyncio.sleep(0, result={"stream": True, "messages": []}), + ) + downstream_closed = False + emitted = b"data: first\n\n" if after_first else b"" + + async def downstream(_payload): + nonlocal downstream_closed + try: + yield emitted + yield b"data: never\n\n" + finally: + downstream_closed = True + + monkeypatch.setattr(instance_api, "_vllm_stream_chat", downstream) + response = await instance_api.instance_chat_completions(request) + assert await anext(response.body_iterator) == emitted + await response.body_iterator.aclose() + + trace = request.state._cacheroute_trace_session.request_trace + assert downstream_closed is True + assert trace.outcome is OutcomeCode.CANCELLED + assert trace.error.message == "instance request cancelled" + assert all(stage.state.value != "running" for stage in trace.stages) + terminal_name = TraceStageName.DECODE if after_first else TraceStageName.FIRST_TOKEN + assert next(stage for stage in trace.stages if stage.name is terminal_name).outcome is OutcomeCode.CANCELLED + assert next(stage for stage in trace.stages if stage.name is TraceStageName.COMPLETION).outcome is OutcomeCode.CANCELLED + assert emitted == (b"data: first\n\n" if after_first else b"") + assert TRACE_ID not in str(response.headers) + assert all(name not in response.headers for name in RESERVED_TRACE_HEADERS) + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("terminal", ["success", "failure", "cancel", "close"]) +def test_collect_streaming_always_closes_wrapped_iterator(monkeypatch, terminal): + async def scenario(): + session = start_instance_trace_session(headers(), cfg(), clock=ManualTraceClock(NOW)) + closed = False + + async def downstream(): + nonlocal closed + try: + if terminal == "failure": + raise RuntimeError("original failure") + yield b"data: first\n\n" + if terminal == "cancel": + await asyncio.Event().wait() + finally: + closed = True + + collected = collect_streaming(session, downstream()) + if terminal == "failure": + with pytest.raises(RuntimeError, match="original failure"): + await anext(collected) + elif terminal == "close": + await anext(collected) + await collected.aclose() + elif terminal == "cancel": + await anext(collected) + pending = asyncio.create_task(anext(collected)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + else: + assert [chunk async for chunk in collected] == [b"data: first\n\n"] + assert closed is True asyncio.run(scenario()) diff --git a/test/test_repository_governance.py b/test/test_repository_governance.py index 499daab..31f8015 100644 --- a/test/test_repository_governance.py +++ b/test/test_repository_governance.py @@ -306,6 +306,63 @@ def test_internal_trace_header_vocabulary_has_single_owner_and_services_import_i assert "extra_headers=extra_headers" in manager_text +def test_internal_trace_propagation_and_vllm_boundary_are_narrow(): + manager_tree = ast.parse((ROOT / "proxy/queue/manager.py").read_text(encoding="utf-8")) + instance_tree = ast.parse((ROOT / "instance/observability.py").read_text(encoding="utf-8")) + manager_calls = { + node.func.id for node in ast.walk(manager_tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + instance_calls = { + node.func.id for node in ast.walk(instance_tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "encode_trace_headers" in manager_calls + assert "decode_trace_headers" in instance_calls + + manager_forward_calls = [ + node for node in ast.walk(manager_tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + and node.func.id == "forward_request" + ] + assert len(manager_forward_calls) == 1 + forwarded = next( + keyword.value for keyword in manager_forward_calls[0].keywords + if keyword.arg == "extra_headers" + ) + assert isinstance(forwarded, ast.Name) and forwarded.id == "extra_headers" + + instance_api = ROOT / "instance/instance_api.py" + instance_api_text = instance_api.read_text(encoding="utf-8") + instance_api_tree = ast.parse(instance_api_text) + vllm_forward_calls = [ + node for node in ast.walk(instance_api_tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + and node.func.id == "forward_request" + ] + assert len(vllm_forward_calls) == 3 + assert all("extra_headers" not in {keyword.arg for keyword in call.keywords} for call in vllm_forward_calls) + instance_string_literals = { + node.value.casefold() for node in ast.walk(instance_api_tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + assert "traceparent" not in instance_string_literals + assert "tracestate" not in instance_string_literals + + canonical_classes = {"TraceContext": [], "RuntimeProfile": []} + for path in _tracked_files(): + if path.suffix != ".py": + continue + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name in {"TraceContext", "RuntimeProfile"}: + canonical_classes[node.name].append(path.relative_to(ROOT)) + assert canonical_classes == { + "TraceContext": [Path("src/cacheroute/observability/v1/models.py")], + "RuntimeProfile": [Path("src/cacheroute/runtime/profiles.py")], + } + + def test_observability_does_not_duplicate_canonical_model_classes(): prohibited = { "RuntimeProfile", "ContractModel", "OutcomeCode", "ContractErrorDetail", From 6b9a19437245ef8c3a3d40b5520e0af4c7299075 Mon Sep 17 00:00:00 2001 From: RickZ <121943721+rickisba@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:21:31 +0800 Subject: [PATCH 4/4] test: correct Instance observability sampling contracts --- .../test_instance_observability.py | 49 +++++++++++++++++-- 1 file changed, 46 insertions(+), 3 deletions(-) diff --git a/test/observability/test_instance_observability.py b/test/observability/test_instance_observability.py index b0e85fd..79c7cee 100644 --- a/test/observability/test_instance_observability.py +++ b/test/observability/test_instance_observability.py @@ -52,7 +52,9 @@ def test_instance_context_accepts_valid_and_falls_back_for_missing_malformed_sta assert stale.context.request_id == "scheduler-1" assert stale.accepted_propagated is False - mismatch = resolve_instance_context(headers(), cfg(profile="test_mock"), clock=clock) + mismatch = resolve_instance_context( + headers(), cfg(profile=RuntimeProfile.TEST_MOCK.value), clock=clock + ) assert mismatch.fallback_reason == "profile_mismatch" assert mismatch.context.request_id == "scheduler-1" @@ -96,11 +98,52 @@ def test_startup_resolution_defaults_and_invalid_sample_rate(): assert invalid.sample_rate_warning_reason == "trace_sample_rate_malformed" -def test_unsampled_session_collects_no_trace(): - session = start_instance_trace_session(headers(), cfg(rate="0.0"), clock=ManualTraceClock(NOW)) +def test_valid_propagated_unsampled_context_ignores_local_full_sample_rate(): + clock = ManualTraceClock(NOW) + propagated = create_trace_context( + "scheduler-1", + RuntimeProfile.LEGACY, + sample_rate=0.0, + clock=clock, + trace_id=TRACE_ID, + ) + session = start_instance_trace_session( + encode_trace_headers(propagated), cfg(rate="1.0"), clock=clock + ) + + assert session.context.sampled is False + assert session.collector is None + + +@pytest.mark.parametrize("incoming_headers", [{}, {REQUEST_ID_HEADER: "incomplete"}]) +def test_fallback_context_uses_local_zero_sample_rate(incoming_headers): + session = start_instance_trace_session( + incoming_headers, cfg(rate="0.0"), clock=ManualTraceClock(NOW) + ) + + assert re.fullmatch(r"req_[0-9a-f]{32}", session.context.request_id) + assert session.context.sampled is False assert session.collector is None +def test_valid_propagated_sampled_context_ignores_local_zero_sample_rate(): + clock = ManualTraceClock(NOW) + propagated = create_trace_context( + "scheduler-1", + RuntimeProfile.LEGACY, + sample_rate=1.0, + clock=clock, + trace_id=TRACE_ID, + ) + session = start_instance_trace_session( + encode_trace_headers(propagated), cfg(rate="0.0"), clock=clock + ) + + assert session.context.request_id == "scheduler-1" + assert session.context.sampled is True + assert session.collector is not None + + @pytest.mark.parametrize("stream_chunks", [[b"data: one\n\n", b"data: [DONE]\n\n"], [b"", b"data: one\n\n"]]) def test_streaming_timing_order_parent_provenance_and_shape(stream_chunks): async def scenario():