Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ea1e7f1997
|
||
|
|
e94633c3a3
|
||
|
|
7cc37ebf37
|
||
|
|
98d63cc3fd
|
@@ -149,8 +149,9 @@ pub struct Usage {
|
||||
/// `completion_tokens`, never added into `total_tokens`).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub completion_tokens_details: Option<CompletionTokensDetails>,
|
||||
/// OpenAI-standard breakdown of `prompt_tokens`. Populated once
|
||||
/// prompt caching lands (#11); `None` until then.
|
||||
/// OpenAI-standard breakdown of `prompt_tokens`, carrying
|
||||
/// `cached_tokens` (#269). Omitted when nothing was reused, so
|
||||
/// "absent" stays distinguishable from "measured as zero".
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
||||
/// helexa extension (non-OpenAI): server-measured prefill/decode
|
||||
@@ -182,8 +183,9 @@ pub struct CompletionTokensDetails {
|
||||
/// Sub-counts of `Usage::prompt_tokens`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PromptTokensDetails {
|
||||
/// Prompt tokens served from cache (cache-read rate). Populated
|
||||
/// once prompt caching lands (#11).
|
||||
/// Prompt tokens served from neuron's prefix KV cache (#11), so a
|
||||
/// caller can see the saving rather than assume none (#269). A
|
||||
/// sub-count of `prompt_tokens`, never added into `total_tokens`.
|
||||
pub cached_tokens: u64,
|
||||
}
|
||||
|
||||
|
||||
@@ -249,8 +249,8 @@ pub struct ResponsesUsage {
|
||||
/// sub-count of `output_tokens`, never added into `total_tokens`).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub output_tokens_details: Option<OutputTokensDetails>,
|
||||
/// OpenAI-standard breakdown of `input_tokens`. Populated once
|
||||
/// prompt caching lands (#11); `None` until then.
|
||||
/// OpenAI-standard breakdown of `input_tokens`, carrying
|
||||
/// `cached_tokens` (#269). Omitted when nothing was reused.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub input_tokens_details: Option<InputTokensDetails>,
|
||||
}
|
||||
@@ -265,8 +265,8 @@ pub struct OutputTokensDetails {
|
||||
/// Sub-counts of `ResponsesUsage::input_tokens`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InputTokensDetails {
|
||||
/// Input tokens served from cache (cache-read rate). Populated
|
||||
/// once prompt caching lands (#11).
|
||||
/// Input tokens served from neuron's prefix KV cache (#11), so a
|
||||
/// caller can see the saving rather than assume none (#269).
|
||||
pub cached_tokens: u64,
|
||||
}
|
||||
|
||||
|
||||
@@ -121,6 +121,10 @@ fn describe_metrics() {
|
||||
"cortex_request_errors_total",
|
||||
"Total number of failed proxy requests"
|
||||
);
|
||||
metrics::describe_counter!(
|
||||
"cortex_cached_prompt_tokens_total",
|
||||
"Prompt tokens served from neuron's prefix KV cache; over cortex_prompt_tokens_total this is the cache-hit rate (#269)"
|
||||
);
|
||||
metrics::describe_counter!("cortex_evictions_total", "Total number of model evictions");
|
||||
metrics::describe_counter!(
|
||||
"cortex_cold_starts_total",
|
||||
|
||||
@@ -171,6 +171,9 @@ impl IntoResponse for ProxyError {
|
||||
// over the decode window (first→last chunk); falls back to the
|
||||
// full request duration for single-chunk (non-streaming) bodies
|
||||
// cortex_prompt_tokens_total / cortex_completion_tokens_total (counters)
|
||||
// cortex_cached_prompt_tokens_total (counter) — prompt tokens served
|
||||
// from neuron's prefix cache; over prompt_tokens_total this is the
|
||||
// fleet's cache-hit rate (#269)
|
||||
|
||||
/// Cap on the retained body tail. The usage object rides on the final
|
||||
/// chunk, so a generous tail is plenty; the cap bounds memory on huge
|
||||
@@ -243,6 +246,15 @@ impl ChunkObserver for CortexMetrics {
|
||||
if let Some(prompt) = prompt {
|
||||
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
|
||||
}
|
||||
// Prefix-cache reuse (#269), read from the same usage object.
|
||||
// Against `cortex_prompt_tokens_total` this gives fleet-wide
|
||||
// cache-hit rate as a ratio, which was previously only
|
||||
// obtainable by grepping `reused=` out of a neuron journal —
|
||||
// per-host, unaggregated, and gone on log rotation.
|
||||
if let Some(cached) = last_count_for(self.tail.as_str(), "cached_tokens") {
|
||||
metrics::counter!("cortex_cached_prompt_tokens_total", &self.labels)
|
||||
.increment(cached);
|
||||
}
|
||||
if let Some(completion) = completion.filter(|c| *c > 0) {
|
||||
metrics::counter!("cortex_completion_tokens_total", &self.labels)
|
||||
.increment(completion);
|
||||
|
||||
@@ -5160,6 +5160,10 @@ impl CandleHarness {
|
||||
let mut prefill_ms_measured: u32 = 0;
|
||||
let mut decode_start: Option<std::time::Instant> = None;
|
||||
|
||||
// Hoisted out of `'work` so the completion path below
|
||||
// can report it (#269) — the reuse count is decided
|
||||
// inside the block but consumed after it.
|
||||
let mut reused_for_task: u32 = 0;
|
||||
'work: {
|
||||
// Prefix-cache decision (#11): vision requests
|
||||
// clear as before; text requests restore the
|
||||
@@ -5181,6 +5185,7 @@ impl CandleHarness {
|
||||
}
|
||||
}
|
||||
};
|
||||
reused_for_task = reused as u32;
|
||||
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
@@ -5599,12 +5604,29 @@ impl CandleHarness {
|
||||
.prefill_rate
|
||||
.record_decode(all_tokens.len(), d.elapsed());
|
||||
}
|
||||
// Close the journal trail (#268) — the serialized
|
||||
// streaming path, like the engine one, previously
|
||||
// ended without saying how. Mirrors the field set
|
||||
// the non-streaming `chat_completion: done` emits.
|
||||
tracing::info!(
|
||||
prompt_tokens = prompt_len,
|
||||
completion_tokens = all_tokens.len(),
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
cached_tokens = reused_for_task,
|
||||
finish_reason = ?finish_reason,
|
||||
prefill_ms = prefill_ms_measured,
|
||||
decode_ms = decode_start
|
||||
.map(|d| d.elapsed().as_millis() as u32)
|
||||
.unwrap_or(0),
|
||||
"chat_completion (stream): done"
|
||||
);
|
||||
let _ = tx
|
||||
.send(InferenceEvent::Finish {
|
||||
reason: finish_reason,
|
||||
prompt_tokens: prompt_len as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
cached_tokens: reused_for_task,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_ms_measured,
|
||||
decode_ms: decode_start
|
||||
@@ -7416,6 +7438,9 @@ async fn stream_inference_via_worker(
|
||||
// effective rate — and the throughput ceiling it feeds — rises.
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let prefill_prompt_len = prompt_tokens.len();
|
||||
// Hoisted out of the match so the Finish event below can report it
|
||||
// (#269). Vision requests clear the cache outright, so they keep 0.
|
||||
let mut cached_tokens: u32 = 0;
|
||||
let logits_vec = match images {
|
||||
Some((imgs, image_token_id)) => {
|
||||
worker
|
||||
@@ -7430,6 +7455,7 @@ async fn stream_inference_via_worker(
|
||||
None => {
|
||||
let reused =
|
||||
restore_or_clear_via_worker(&worker, handle, prefix_cache, &prompt_tokens).await?;
|
||||
cached_tokens = reused as u32;
|
||||
// Two-stage prefill around the retokenization-stable
|
||||
// snapshot boundary — see `run_inference_via_worker`.
|
||||
let cut = if prefix_cache.is_some() {
|
||||
@@ -7613,12 +7639,24 @@ async fn stream_inference_via_worker(
|
||||
}
|
||||
// Fold this request's decode throughput into the model tracker (#137).
|
||||
prefill_rate.record_decode(all_tokens.len(), decode_start.elapsed());
|
||||
// Close the journal trail (#268).
|
||||
tracing::info!(
|
||||
prompt_tokens = prompt_tokens.len(),
|
||||
completion_tokens = all_tokens.len(),
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
cached_tokens,
|
||||
finish_reason = ?finish_reason,
|
||||
prefill_ms = prefill_elapsed.as_millis(),
|
||||
decode_ms = decode_start.elapsed().as_millis(),
|
||||
"chat_completion (stream): done"
|
||||
);
|
||||
let _ = tx
|
||||
.send(InferenceEvent::Finish {
|
||||
reason: finish_reason,
|
||||
prompt_tokens: prompt_tokens.len() as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
cached_tokens,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_elapsed.as_millis() as u32,
|
||||
decode_ms: decode_start.elapsed().as_millis() as u32,
|
||||
@@ -7885,11 +7923,21 @@ fn run_inference_streaming(
|
||||
if emitted_tool_call && finish_reason == FinishReason::Stop {
|
||||
finish_reason = FinishReason::ToolCalls;
|
||||
}
|
||||
// Close the journal trail (#268).
|
||||
tracing::info!(
|
||||
prompt_tokens = prompt_tokens.len(),
|
||||
completion_tokens = all_tokens.len(),
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
cached_tokens = reused,
|
||||
finish_reason = ?finish_reason,
|
||||
"chat_completion (stream): done"
|
||||
);
|
||||
let _ = tx.blocking_send(InferenceEvent::Finish {
|
||||
reason: finish_reason,
|
||||
prompt_tokens: prompt_tokens.len() as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
cached_tokens: reused as u32,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_elapsed.as_millis() as u32,
|
||||
decode_ms: decode_start.elapsed().as_millis() as u32,
|
||||
|
||||
@@ -297,14 +297,23 @@ async fn op_drop_snap(cfg: &EngineConfig, session: &mut ActiveSession, id: u64)
|
||||
|
||||
/// Prefill one joining request at B=1 through the backend's existing
|
||||
/// chunked-prefill + prefix-cache paths. Returns the prefill logits.
|
||||
/// Prefill, returning the logits **and** how many leading tokens the
|
||||
/// prefix cache let us skip (#269).
|
||||
///
|
||||
/// The reuse count is computed here and was previously only logged, so
|
||||
/// no caller could report it — every client saw a 0% cache hit rate on
|
||||
/// a cache that was working. Returning it is what lets `usage` carry
|
||||
/// `cached_tokens`.
|
||||
async fn op_prefill(
|
||||
cfg: &EngineConfig,
|
||||
session: &mut ActiveSession,
|
||||
prompt_tokens: &[u32],
|
||||
) -> Result<Vec<f32>> {
|
||||
) -> Result<(Vec<f32>, usize)> {
|
||||
let prompt_len = prompt_tokens.len();
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let logits = match (&cfg.backend, session) {
|
||||
// Each arm yields its logits plus how many leading tokens the
|
||||
// prefix cache served, so neither can be forgotten.
|
||||
let (logits, cached_tokens) = match (&cfg.backend, session) {
|
||||
(
|
||||
BackendConfig::Single {
|
||||
worker,
|
||||
@@ -340,7 +349,7 @@ async fn op_prefill(
|
||||
None => chunked_prefill_via_worker(worker, *handle, prompt_tokens, reused).await?,
|
||||
};
|
||||
prefill_rate.record(prompt_len, prefill_start.elapsed());
|
||||
logits
|
||||
(logits, reused)
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
(BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
|
||||
@@ -384,12 +393,12 @@ async fn op_prefill(
|
||||
}
|
||||
};
|
||||
tp.prefill_rate.record(prompt_len, prefill_start.elapsed());
|
||||
logits
|
||||
(logits, reused)
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
_ => anyhow::bail!("engine backend/session mismatch"),
|
||||
};
|
||||
Ok(logits)
|
||||
Ok((logits, cached_tokens))
|
||||
}
|
||||
|
||||
/// One queued request. Admission has already been passed — the permit
|
||||
@@ -439,6 +448,8 @@ enum RouterMsg {
|
||||
reason: FinishReason,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
/// Leading prompt tokens the prefix cache let us skip (#269).
|
||||
cached_tokens: u32,
|
||||
timing: FinishTiming,
|
||||
},
|
||||
}
|
||||
@@ -462,6 +473,9 @@ struct Slot {
|
||||
finished: Option<FinishReason>,
|
||||
prefill_ms: u32,
|
||||
prefill_tokens: u32,
|
||||
/// Prompt tokens served from the prefix cache (#269), reported to
|
||||
/// the caller as `cached_tokens` so a working cache is visible.
|
||||
cached_tokens: u32,
|
||||
decode_start: std::time::Instant,
|
||||
_admit: AdmissionPermit,
|
||||
}
|
||||
@@ -627,6 +641,7 @@ async fn finish_slot(slot: &mut Slot, reason: FinishReason, rate: &PrefillRateEm
|
||||
reason,
|
||||
prompt_tokens: slot.prompt_len as u32,
|
||||
completion_tokens: slot.generated.len() as u32,
|
||||
cached_tokens: slot.cached_tokens,
|
||||
timing: FinishTiming {
|
||||
prefill_ms: slot.prefill_ms,
|
||||
decode_ms: slot.decode_start.elapsed().as_millis() as u32,
|
||||
@@ -788,7 +803,7 @@ async fn prefill_join(
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let logits_vec = op_prefill(cfg, session, &prompt_tokens).await?;
|
||||
let (logits_vec, cached_tokens) = op_prefill(cfg, session, &prompt_tokens).await?;
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
|
||||
// First token from the prefill logits.
|
||||
@@ -827,6 +842,7 @@ async fn prefill_join(
|
||||
let mut slot = Slot {
|
||||
prefix_len: prompt_len,
|
||||
prompt_len,
|
||||
cached_tokens: cached_tokens as u32,
|
||||
generated,
|
||||
next_token: first,
|
||||
max_new,
|
||||
@@ -890,6 +906,9 @@ async fn run_router(
|
||||
let mut tool_call_idx: usize = 0;
|
||||
let mut emitted_tool_call = false;
|
||||
let mut consumer_alive = true;
|
||||
// Whether a Finish was seen, so the fall-through below can tell a
|
||||
// clean end from a dropped slot.
|
||||
let mut finished = false;
|
||||
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
@@ -978,17 +997,43 @@ async fn run_router(
|
||||
mut reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cached_tokens,
|
||||
timing,
|
||||
} => {
|
||||
if emitted_tool_call && reason == FinishReason::Stop {
|
||||
reason = FinishReason::ToolCalls;
|
||||
}
|
||||
// Close the journal trail for this turn (#268). The
|
||||
// non-streaming path has always logged this; streamed
|
||||
// turns logged nothing at all, so whether one finished,
|
||||
// hit its output cap, or was abandoned mid-flight was
|
||||
// unanswerable afterwards — and an absence of lines was
|
||||
// repeatedly misread as an absence of work.
|
||||
//
|
||||
// `run_router` is instrumented with the request's span,
|
||||
// so req_id and model come along without being passed.
|
||||
tracing::info!(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
cached_tokens,
|
||||
finish_reason = ?reason,
|
||||
prefill_ms = timing.prefill_ms,
|
||||
decode_ms = timing.decode_ms,
|
||||
// The turn ran to completion with nobody reading it:
|
||||
// it held its admission permit and the pool lock the
|
||||
// whole time. Invisible until now.
|
||||
client_gone = !consumer_alive,
|
||||
"chat_completion (stream): done"
|
||||
);
|
||||
finished = true;
|
||||
let _ = tx
|
||||
.send(InferenceEvent::Finish {
|
||||
reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
cached_tokens,
|
||||
timing: Some(timing),
|
||||
})
|
||||
.await;
|
||||
@@ -996,6 +1041,17 @@ async fn run_router(
|
||||
}
|
||||
}
|
||||
}
|
||||
// Fell out of the loop without a Finish: the engine dropped this
|
||||
// slot's channel (fatal worker error, shutdown). The client's stream
|
||||
// ends either way, but silently — so say so rather than leaving the
|
||||
// turn's last trace to be a prefill line (#268).
|
||||
if !finished {
|
||||
tracing::warn!(
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
client_gone = !consumer_alive,
|
||||
"chat_completion (stream): ended without a finish event — engine dropped the slot"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// `tracing::Instrument` without importing the trait at every use
|
||||
@@ -1013,6 +1069,92 @@ mod tests {
|
||||
use crate::config::AdmissionConfig;
|
||||
use crate::harness::admission::AdmissionController;
|
||||
|
||||
/// A streamed turn must close its own journal trail (#268).
|
||||
///
|
||||
/// Streaming used to end silently: no token counts, no
|
||||
/// `finish_reason`, no duration. That made a finished turn, a turn
|
||||
/// that hit its output cap, and a turn abandoned by a disconnected
|
||||
/// client indistinguishable after the fact — and during the #267
|
||||
/// investigation an absence of log lines was twice mistaken for an
|
||||
/// absence of work on a daemon that was healthy and decoding.
|
||||
///
|
||||
/// Asserts the emitted line rather than the struct, because a
|
||||
/// `Finish` event that reaches the client but leaves no server-side
|
||||
/// record is exactly the state this fixes.
|
||||
#[tokio::test]
|
||||
async fn a_streamed_turn_logs_how_it_ended() {
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct Capture(Arc<Mutex<Vec<u8>>>);
|
||||
impl std::io::Write for Capture {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.0
|
||||
.lock()
|
||||
.expect("capture poisoned")
|
||||
.extend_from_slice(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Capture {
|
||||
type Writer = Self;
|
||||
fn make_writer(&'a self) -> Self::Writer {
|
||||
self.clone()
|
||||
}
|
||||
}
|
||||
|
||||
let capture = Capture::default();
|
||||
let subscriber = tracing_subscriber::registry()
|
||||
.with(tracing_subscriber::fmt::layer().with_writer(capture.clone()));
|
||||
let _guard = tracing::subscriber::set_default(subscriber);
|
||||
|
||||
let (router_tx, router_rx) = mpsc::channel::<RouterMsg>(8);
|
||||
let (tx, mut rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let handle = tokio::spawn(run_router(
|
||||
tiny_tokenizer(8),
|
||||
None,
|
||||
None,
|
||||
ToolSchemas::default(),
|
||||
false,
|
||||
tx,
|
||||
Arc::new(AtomicBool::new(false)),
|
||||
router_rx,
|
||||
));
|
||||
|
||||
router_tx
|
||||
.send(RouterMsg::Finish {
|
||||
reason: FinishReason::Length,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 24576,
|
||||
cached_tokens: 2068,
|
||||
timing: FinishTiming {
|
||||
prefill_ms: 6000,
|
||||
decode_ms: 900_000,
|
||||
prefill_tokens: 7413,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.expect("router accepts finish");
|
||||
drop(router_tx);
|
||||
handle.await.expect("router task joins");
|
||||
while rx.recv().await.is_some() {}
|
||||
|
||||
let logged = String::from_utf8(capture.0.lock().expect("capture poisoned").clone())
|
||||
.expect("utf8 log output");
|
||||
assert!(
|
||||
logged.contains("chat_completion (stream): done"),
|
||||
"streamed turn left no completion line:\n{logged}"
|
||||
);
|
||||
// The counts are the point — "it ended" without them cannot
|
||||
// distinguish a complete answer from an exhausted budget.
|
||||
assert!(logged.contains("24576"), "no completion_tokens:\n{logged}");
|
||||
assert!(logged.contains("Length"), "no finish_reason:\n{logged}");
|
||||
}
|
||||
|
||||
/// A WordLevel tokenizer whose vocab covers the whole fixture
|
||||
/// vocab (`w0`..`w511`), so every decoded token maps to a unique
|
||||
/// word and the emitted text uniquely encodes the token sequence.
|
||||
|
||||
@@ -84,6 +84,13 @@ pub enum InferenceEvent {
|
||||
/// `output_tokens_details.reasoning_tokens` (responses).
|
||||
/// Zero for non-reasoning models.
|
||||
reasoning_tokens: u32,
|
||||
/// Leading prompt tokens served from the prefix cache (#269) — a
|
||||
/// sub-count of `prompt_tokens`, never added to `total_tokens`.
|
||||
/// Surfaced as `prompt_tokens_details.cached_tokens` (chat) /
|
||||
/// `input_tokens_details.cached_tokens` (responses). The cache
|
||||
/// has worked since #11; without this the saving was invisible
|
||||
/// and every client reported a 0% hit rate.
|
||||
cached_tokens: u32,
|
||||
/// Server-measured prefill/decode timing for the request, or
|
||||
/// `None` on paths that don't measure it (CPU fallback that
|
||||
/// doesn't instrument, tests). Streaming projectors surface
|
||||
|
||||
@@ -27,7 +27,8 @@
|
||||
//! propagate without us writing any logic.
|
||||
|
||||
use cortex_core::openai::{
|
||||
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, Usage,
|
||||
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, PromptTokensDetails,
|
||||
Usage,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc;
|
||||
@@ -219,6 +220,7 @@ pub fn project_chat_stream_with(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
timing,
|
||||
} => {
|
||||
// The finish_reason chunk, then an OpenAI-style
|
||||
@@ -233,9 +235,12 @@ pub fn project_chat_stream_with(
|
||||
&id,
|
||||
created,
|
||||
&model_id,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
UsageCounts {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
},
|
||||
timing,
|
||||
),
|
||||
]
|
||||
@@ -385,15 +390,32 @@ fn final_chunk(
|
||||
/// ignore the empty-choices chunk; clients that do (opencode, and
|
||||
/// cortex's Anthropic translator) get the token counts they need to
|
||||
/// track context.
|
||||
/// The token counts a finished turn reports. Grouped rather than passed
|
||||
/// loose: adding `cached_tokens` (#269) put the list past the point
|
||||
/// where four adjacent bare `u32`s are safe to read at a call site.
|
||||
struct UsageCounts {
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
/// Sub-count of `completion_tokens` — tokens spent reasoning.
|
||||
reasoning_tokens: u32,
|
||||
/// Sub-count of `prompt_tokens` — tokens served from the prefix
|
||||
/// cache. Neither sub-count is added into `total_tokens`.
|
||||
cached_tokens: u32,
|
||||
}
|
||||
|
||||
fn usage_chunk(
|
||||
id: &str,
|
||||
created: u64,
|
||||
model_id: &str,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
reasoning_tokens: u32,
|
||||
counts: UsageCounts,
|
||||
timing: Option<FinishTiming>,
|
||||
) -> ChatCompletionChunk {
|
||||
let UsageCounts {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
} = counts;
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
@@ -409,7 +431,12 @@ fn usage_chunk(
|
||||
completion_tokens_details: (reasoning_tokens > 0).then_some(CompletionTokensDetails {
|
||||
reasoning_tokens: reasoning_tokens as u64,
|
||||
}),
|
||||
prompt_tokens_details: None,
|
||||
// Prefix-cache reuse (#269). Omitted at zero so a client
|
||||
// sees unchanged JSON when nothing was cached, and so
|
||||
// "absent" stays distinguishable from "measured as none".
|
||||
prompt_tokens_details: (cached_tokens > 0).then_some(PromptTokensDetails {
|
||||
cached_tokens: cached_tokens as u64,
|
||||
}),
|
||||
// helexa extension (#85): server-measured prefill/decode
|
||||
// timing for the bench harness. Omitted on paths that don't
|
||||
// measure it so standard clients see unchanged JSON.
|
||||
@@ -458,6 +485,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -496,6 +524,7 @@ mod tests {
|
||||
prompt_tokens: 128,
|
||||
completion_tokens: 64,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: 200,
|
||||
decode_ms: 1500,
|
||||
@@ -520,6 +549,73 @@ mod tests {
|
||||
assert_eq!(timing.prefill_tokens, 128);
|
||||
}
|
||||
|
||||
/// A cache hit must reach the caller (#269).
|
||||
///
|
||||
/// Prefix caching has worked since #11 — neuron logs `reused=2068`
|
||||
/// on prefills — but nothing carried the number onto the wire, so
|
||||
/// every client reported a 0% hit rate on a cache that was saving
|
||||
/// real prefill. dsh showed `Cache hit 0%` against a 42.1K-token
|
||||
/// context, and an agentic harness that believes nothing is cached
|
||||
/// may restructure the very prefix that was being reused.
|
||||
#[tokio::test]
|
||||
async fn cached_tokens_reach_the_usage_chunk() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 2068,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let out = collect(out_rx).await;
|
||||
let usage = out
|
||||
.iter()
|
||||
.find_map(|c| c.usage.as_ref())
|
||||
.expect("usage chunk present");
|
||||
let details = usage
|
||||
.prompt_tokens_details
|
||||
.as_ref()
|
||||
.expect("prompt_tokens_details populated when the cache was hit");
|
||||
assert_eq!(details.cached_tokens, 2068);
|
||||
// A sub-count of prompt_tokens, never added into the total —
|
||||
// otherwise a cache hit would inflate what a caller is billed.
|
||||
assert_eq!(usage.prompt_tokens, 7413);
|
||||
assert_eq!(usage.total_tokens, 7425);
|
||||
}
|
||||
|
||||
/// No reuse means no `prompt_tokens_details` at all, so a client
|
||||
/// sees byte-identical JSON to before this landed — and "absent"
|
||||
/// stays distinguishable from "measured as zero".
|
||||
#[tokio::test]
|
||||
async fn no_cache_hit_omits_the_details_object() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 1,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let usage = collect(out_rx)
|
||||
.await
|
||||
.iter()
|
||||
.find_map(|c| c.usage.clone())
|
||||
.expect("usage chunk present");
|
||||
assert!(usage.prompt_tokens_details.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_text_delta_is_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
@@ -541,6 +637,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -637,6 +734,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -689,6 +787,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -735,6 +834,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -781,6 +881,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -809,6 +910,7 @@ mod tests {
|
||||
prompt_tokens: 42,
|
||||
completion_tokens: 5,
|
||||
reasoning_tokens: 2,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -843,6 +945,7 @@ mod tests {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 7,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
|
||||
@@ -33,9 +33,9 @@
|
||||
|
||||
use cortex_core::openai::{ChatCompletionRequest, ChatMessage, MessageContent};
|
||||
use cortex_core::responses::{
|
||||
OutputTokensDetails, ResponsesContentPart, ResponsesInput, ResponsesInputElement,
|
||||
ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent, ResponsesOutputItem,
|
||||
ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
|
||||
InputTokensDetails, OutputTokensDetails, ResponsesContentPart, ResponsesInput,
|
||||
ResponsesInputElement, ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent,
|
||||
ResponsesOutputItem, ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
@@ -540,6 +540,7 @@ async fn run_projection(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
// Responses-side `helexa_timing` surfacing not wired yet;
|
||||
// the bench harness reads timing off the chat path (#85).
|
||||
timing: _,
|
||||
@@ -556,7 +557,12 @@ async fn run_projection(
|
||||
output_tokens_details: (reasoning_tokens > 0).then_some(OutputTokensDetails {
|
||||
reasoning_tokens: reasoning_tokens as u64,
|
||||
}),
|
||||
input_tokens_details: None,
|
||||
// Prefix-cache reuse (#269). Omitted at zero so a
|
||||
// client sees unchanged JSON when nothing was
|
||||
// cached.
|
||||
input_tokens_details: (cached_tokens > 0).then_some(InputTokensDetails {
|
||||
cached_tokens: cached_tokens as u64,
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -1351,6 +1357,41 @@ mod tests {
|
||||
/// timeout could never finish. Its watchdog resets on any parsed
|
||||
/// event, which is why emitting these fixes it and why SSE comment
|
||||
/// keep-alives could not: idle timers count events, not comments.
|
||||
/// The Responses surface reports cache reuse too (#269), in its own
|
||||
/// spelling: `input_tokens_details.cached_tokens`.
|
||||
///
|
||||
/// This is the surface dsh uses, and the one that showed
|
||||
/// `Cache hit 0%` against a 42.1K-token context while neuron was
|
||||
/// reusing thousands of tokens per prefill.
|
||||
#[tokio::test]
|
||||
async fn cached_tokens_reach_the_completed_usage() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("ok".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 2068,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let usage = &frames.last().unwrap().data["response"]["usage"];
|
||||
assert_eq!(usage["input_tokens_details"]["cached_tokens"], 2068);
|
||||
assert_eq!(usage["input_tokens"], 7413);
|
||||
// Sub-count, not additive — a cache hit must not inflate the total.
|
||||
assert_eq!(usage["total_tokens"], 7425);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_projects_a_full_item_lifecycle_before_the_message() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
@@ -1371,6 +1412,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 2,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1444,6 +1486,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1487,6 +1530,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1522,6 +1566,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1583,6 +1628,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1655,6 +1701,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1684,6 +1731,7 @@ mod tests {
|
||||
prompt_tokens: 30,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 4,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1715,6 +1763,7 @@ mod tests {
|
||||
prompt_tokens: 8,
|
||||
completion_tokens: 3,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1740,6 +1789,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1787,6 +1837,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
|
||||
Reference in New Issue
Block a user