Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d453d5deb8
|
||
|
|
01c604195d
|
||
|
|
ea1e7f1997
|
||
|
|
e94633c3a3
|
||
|
|
7cc37ebf37
|
||
|
|
98d63cc3fd
|
||
|
|
a26bd94f13
|
||
|
|
0fe7aa38b0
|
||
|
|
d1ea8d28d3
|
||
|
|
b5e77148df
|
@@ -20,6 +20,11 @@ pub struct MessagesRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Anthropic's Messages API carries `top_k` natively, so it maps
|
||||
/// straight onto the OpenAI-side field rather than being dropped in
|
||||
/// translation (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(flatten)]
|
||||
|
||||
@@ -18,6 +18,31 @@ pub struct ChatCompletionRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Truncate sampling to the `k` most likely tokens (#272).
|
||||
///
|
||||
/// Not in the OpenAI schema but universal among serving stacks
|
||||
/// (vLLM, SGLang, llama.cpp, Ollama), and required by models that
|
||||
/// publish it — Qwen3 specifies `top_k = 20`. A named field rather
|
||||
/// than `extra`, because a sampling parameter that lands in `extra`
|
||||
/// is accepted and silently discarded.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
/// Pin the sampler's RNG for reproducible output (OpenAI core).
|
||||
///
|
||||
/// Previously ignored on every text path, which called
|
||||
/// `unix_subsec_nanos()` regardless — a caller asking for
|
||||
/// determinism got fresh randomness and a `200`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Penalty applied to recently-generated tokens; `1.0` disables.
|
||||
/// Defaults to 1.1 when absent (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repetition_penalty: Option<f32>,
|
||||
/// How many recent tokens `repetition_penalty` considers. Defaults
|
||||
/// to 64 — fine for chat, short for a long reasoning block, which is
|
||||
/// why it is a knob rather than a constant (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repeat_last_n: Option<usize>,
|
||||
/// The deprecated spelling of the output cap. Still what most
|
||||
/// clients send; see [`ChatCompletionRequest::effective_max_tokens`].
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -149,8 +174,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 +208,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,
|
||||
}
|
||||
|
||||
|
||||
@@ -55,6 +55,21 @@ pub struct ResponsesRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Truncate sampling to the `k` most likely tokens (#272). Not in
|
||||
/// the Responses schema, but a sampling parameter that lands in
|
||||
/// `extra` is accepted and silently discarded, which is the defect
|
||||
/// this closes.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
/// Pin the sampler's RNG for reproducible output (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Penalty on recently-generated tokens; `1.0` disables (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repetition_penalty: Option<f32>,
|
||||
/// Window `repetition_penalty` considers (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repeat_last_n: Option<usize>,
|
||||
/// Chained-conversation identifier. We don't store responses
|
||||
/// server-side yet; if this is `Some`, the handler returns 400.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -249,8 +264,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 +280,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,
|
||||
}
|
||||
|
||||
@@ -289,6 +304,16 @@ pub mod events {
|
||||
pub const OUTPUT_TEXT_DONE: &str = "response.output_text.done";
|
||||
pub const CONTENT_PART_DONE: &str = "response.content_part.done";
|
||||
pub const OUTPUT_ITEM_DONE: &str = "response.output_item.done";
|
||||
/// Reasoning-item event family. A reasoning model's think block is
|
||||
/// its own output item, emitted ahead of the message item, with a
|
||||
/// summary part carrying the text. Streaming these is what keeps the
|
||||
/// connection observably alive while a model thinks — clients time
|
||||
/// out on silence, and SSE comment keep-alives don't count because
|
||||
/// idle timers reset on parsed events, not comments.
|
||||
pub const REASONING_SUMMARY_PART_ADDED: &str = "response.reasoning_summary_part.added";
|
||||
pub const REASONING_SUMMARY_TEXT_DELTA: &str = "response.reasoning_summary_text.delta";
|
||||
pub const REASONING_SUMMARY_TEXT_DONE: &str = "response.reasoning_summary_text.done";
|
||||
pub const REASONING_SUMMARY_PART_DONE: &str = "response.reasoning_summary_part.done";
|
||||
pub const FUNCTION_CALL_ARGUMENTS_DELTA: &str = "response.function_call_arguments.delta";
|
||||
pub const FUNCTION_CALL_ARGUMENTS_DONE: &str = "response.function_call_arguments.done";
|
||||
pub const COMPLETED: &str = "response.completed";
|
||||
|
||||
@@ -101,6 +101,13 @@ pub fn anthropic_to_openai(req: MessagesRequest) -> ChatCompletionRequest {
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
// Anthropic models top_k too, so it survives the hop (#272).
|
||||
top_k: req.top_k,
|
||||
// Anthropic's Messages API has no seed or penalty knobs; absent
|
||||
// means the model's own generation_config decides.
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
max_tokens: Some(req.max_tokens),
|
||||
// Anthropic has one spelling; it maps onto the legacy field.
|
||||
max_completion_tokens: None,
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
//! in Prometheus text format.
|
||||
|
||||
use anyhow::Result;
|
||||
use metrics_exporter_prometheus::PrometheusBuilder;
|
||||
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder};
|
||||
use std::net::SocketAddr;
|
||||
|
||||
/// Install the Prometheus metrics recorder and return a handle.
|
||||
@@ -12,8 +12,7 @@ use std::net::SocketAddr;
|
||||
pub fn install(listen: &str) -> Result<()> {
|
||||
let addr: SocketAddr = listen.parse()?;
|
||||
|
||||
PrometheusBuilder::new()
|
||||
.with_http_listener(addr)
|
||||
with_buckets(PrometheusBuilder::new().with_http_listener(addr))?
|
||||
.install()
|
||||
.map_err(|e| anyhow::anyhow!("failed to install Prometheus exporter: {e}"))?;
|
||||
|
||||
@@ -25,13 +24,77 @@ pub fn install(listen: &str) -> Result<()> {
|
||||
/// Install a recorder for testing (no HTTP listener). Returns a handle
|
||||
/// that can render the current metrics as Prometheus text.
|
||||
pub fn install_test_recorder() -> Result<metrics_exporter_prometheus::PrometheusHandle> {
|
||||
let handle = PrometheusBuilder::new()
|
||||
let handle = with_buckets(PrometheusBuilder::new())?
|
||||
.install_recorder()
|
||||
.map_err(|e| anyhow::anyhow!("failed to install test recorder: {e}"))?;
|
||||
describe_metrics();
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
/// Give every histogram explicit buckets, so it is exported as a
|
||||
/// Prometheus **histogram** rather than a summary.
|
||||
///
|
||||
/// Without this, `metrics-exporter-prometheus` renders every
|
||||
/// `histogram!` as a summary: `{quantile="0.95"}` series and no
|
||||
/// `_bucket` series at all. Three things break as a result, and the
|
||||
/// third is what made it visible.
|
||||
///
|
||||
/// 1. `histogram_quantile()` has nothing to read, so any dashboard
|
||||
/// panel written the idiomatic way returns *No data* forever. The
|
||||
/// fleet dashboard's TTFT panel did exactly that while cortex was
|
||||
/// serving thousands of requests.
|
||||
/// 2. Summary quantiles are computed per process over a rolling
|
||||
/// window, so they **cannot be aggregated**. Averaging p95 across
|
||||
/// two gateways is not p95 of anything.
|
||||
/// 3. They decay: with no samples in the window every quantile reads
|
||||
/// `0`, which is indistinguishable from "genuinely instant" on a
|
||||
/// graph.
|
||||
///
|
||||
/// Buckets are per-metric because the quantities differ by orders of
|
||||
/// magnitude — a request lasting minutes and a TTFT of milliseconds
|
||||
/// have no useful shared scale. Ranges are chosen from observed fleet
|
||||
/// behaviour rather than round numbers: decode runs for minutes on a
|
||||
/// long turn, prefill is seconds on a long prompt, and decode
|
||||
/// throughput sits in the tens of tokens/sec.
|
||||
fn with_buckets(builder: PrometheusBuilder) -> Result<PrometheusBuilder> {
|
||||
let seconds_short = &[
|
||||
0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0, 120.0,
|
||||
];
|
||||
let seconds_long = &[
|
||||
0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0, 1200.0,
|
||||
];
|
||||
let tokens_per_second = &[
|
||||
1.0, 2.5, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 120.0, 200.0,
|
||||
];
|
||||
builder
|
||||
// Whole-request latency: a long agentic turn legitimately runs
|
||||
// for many minutes, so the tail has to reach there.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_request_duration_seconds".into()),
|
||||
seconds_long,
|
||||
)?
|
||||
// Prefill: sub-second for a short prompt, seconds for a long
|
||||
// one. Anything past a minute is pathological and belongs in
|
||||
// the overflow bucket.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_time_to_first_token_seconds".into()),
|
||||
seconds_short,
|
||||
)?
|
||||
// Image generation is inherently slower than a text turn and
|
||||
// scales with resolution and step count.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_images_generation_seconds".into()),
|
||||
seconds_long,
|
||||
)?
|
||||
// Not a duration — decode throughput, tens of tokens/sec on
|
||||
// this fleet.
|
||||
.set_buckets_for_metric(
|
||||
Matcher::Full("cortex_tokens_per_second".into()),
|
||||
tokens_per_second,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("failed to configure histogram buckets: {e}"))
|
||||
}
|
||||
|
||||
fn describe_metrics() {
|
||||
metrics::describe_histogram!(
|
||||
"cortex_request_duration_seconds",
|
||||
@@ -58,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",
|
||||
@@ -130,3 +197,43 @@ fn describe_metrics() {
|
||||
"Live prefill throughput per neuron:model, tokens/sec EMA (#137)"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod bucket_tests {
|
||||
/// Histograms must export as Prometheus **histograms**, with
|
||||
/// `_bucket` series — not as summaries.
|
||||
///
|
||||
/// This asserts the exported *shape*, not the recording call,
|
||||
/// because the recording was never the problem: cortex measured
|
||||
/// TTFT correctly for months while the fleet dashboard's panel sat
|
||||
/// on "No data", since `histogram_quantile()` reads `_bucket`
|
||||
/// series and a summary has none. A metric that is collected but
|
||||
/// cannot be queried is indistinguishable from one that was never
|
||||
/// collected — from the graph, and from the operator's chair.
|
||||
#[test]
|
||||
fn histograms_export_buckets_not_summary_quantiles() {
|
||||
let handle = match super::install_test_recorder() {
|
||||
Ok(h) => h,
|
||||
// Another test in this binary owns the global recorder;
|
||||
// it installs the same buckets, so skipping is honest.
|
||||
Err(_) => return,
|
||||
};
|
||||
metrics::histogram!("cortex_time_to_first_token_seconds", "node" => "n", "model" => "m")
|
||||
.record(0.42);
|
||||
let rendered = handle.render();
|
||||
|
||||
assert!(
|
||||
rendered.contains("cortex_time_to_first_token_seconds_bucket"),
|
||||
"no _bucket series — histogram_quantile() cannot read this:\n{rendered}"
|
||||
);
|
||||
assert!(
|
||||
rendered.contains("# TYPE cortex_time_to_first_token_seconds histogram"),
|
||||
"exported as the wrong metric type:\n{rendered}"
|
||||
);
|
||||
assert!(
|
||||
!rendered
|
||||
.contains(r#"cortex_time_to_first_token_seconds{node="n",model="m",quantile="#),
|
||||
"still exporting summary quantiles, which cannot be aggregated across gateways"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
+159
-132
@@ -13,7 +13,7 @@ use async_trait::async_trait;
|
||||
use candle_core::quantized::gguf_file;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
use candle_transformers::generation::{LogitsProcessor, Sampling};
|
||||
use candle_transformers::generation::LogitsProcessor;
|
||||
use candle_transformers::models::llama as llama_dense;
|
||||
use candle_transformers::models::quantized_llama::ModelWeights as QuantizedLlamaWeights;
|
||||
use candle_transformers::models::quantized_qwen3::ModelWeights as QuantizedQwen3Weights;
|
||||
@@ -568,8 +568,30 @@ impl LoadedImageModel {
|
||||
/// A loaded model with its tokenizer, device placement, and architecture-
|
||||
/// specific weights. The `arch` is `Arc<Mutex<>>` so the lock guard can be
|
||||
/// moved into `spawn_blocking` for synchronous candle forward passes.
|
||||
/// Collect the sampling knobs a request expressed (#272).
|
||||
///
|
||||
/// Everything optional — absent means "no opinion", which defers to the
|
||||
/// model's `generation_config.json` and then to the built-in fallback.
|
||||
/// Kept in one place so a new knob is added once rather than at each of
|
||||
/// the four request-handling paths.
|
||||
fn requested_sampling(request: &ChatCompletionRequest) -> super::sampling::RequestedSampling {
|
||||
super::sampling::RequestedSampling {
|
||||
temperature: request.temperature,
|
||||
top_p: request.top_p,
|
||||
top_k: request.top_k,
|
||||
seed: request.seed,
|
||||
repeat_penalty: request.repetition_penalty,
|
||||
repeat_last_n: request.repeat_last_n,
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LoadedModel {
|
||||
pub model_id: String,
|
||||
/// What the model published in `generation_config.json` (#272).
|
||||
/// Defaults for any sampling field the request leaves unset — so a
|
||||
/// caller that says nothing gets what the model's authors intended
|
||||
/// rather than an untruncated distribution.
|
||||
pub generation_defaults: super::sampling::ModelGenerationDefaults,
|
||||
/// Local (async-side) handle to the model architecture. `Some`
|
||||
/// only when the model loaded onto the CPU device (no CUDA
|
||||
/// available); the inference path then takes this mutex via
|
||||
@@ -742,6 +764,8 @@ impl LoadedModel {
|
||||
#[cfg(feature = "cuda")]
|
||||
pub struct TpLoadedModel {
|
||||
pub model_id: String,
|
||||
/// What the model published in `generation_config.json` (#272).
|
||||
pub generation_defaults: super::sampling::ModelGenerationDefaults,
|
||||
pub tokenizer: Tokenizer,
|
||||
pub devices: Vec<u32>,
|
||||
/// One end-to-end gate: the pool's RPC stream to the subprocess
|
||||
@@ -2815,10 +2839,14 @@ impl CandleHarness {
|
||||
// inference closure takes ownership of `prompt_tokens`.
|
||||
let prompt_opened_reasoning =
|
||||
prompt_opens_reasoning(&prompt_tokens, loaded.reasoning_tokens.as_ref());
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
// Resolve sampling once: request → the model's
|
||||
// generation_config.json → built-in fallback (#272).
|
||||
let sampling = super::sampling::SamplingParams::resolve(
|
||||
&requested_sampling(&request),
|
||||
&loaded.generation_defaults,
|
||||
unix_subsec_nanos(),
|
||||
);
|
||||
let max_new = request.max_tokens.unwrap_or(8192) as usize;
|
||||
let seed = unix_subsec_nanos();
|
||||
|
||||
let eos_id = loaded
|
||||
.tokenizer
|
||||
@@ -2829,8 +2857,10 @@ impl CandleHarness {
|
||||
tracing::info!(
|
||||
prompt_len,
|
||||
max_new,
|
||||
temperature,
|
||||
?top_p,
|
||||
temperature = sampling.temperature,
|
||||
top_p = ?sampling.top_p,
|
||||
top_k = ?sampling.top_k,
|
||||
seed = sampling.seed,
|
||||
?eos_id,
|
||||
vram_free_mb,
|
||||
vram_total_mb,
|
||||
@@ -2874,9 +2904,7 @@ impl CandleHarness {
|
||||
images.clone(),
|
||||
*image_token_id,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling.clone(),
|
||||
eos_id,
|
||||
)
|
||||
.await
|
||||
@@ -2889,9 +2917,7 @@ impl CandleHarness {
|
||||
loaded.prefix_cache.as_deref(),
|
||||
loaded.tokenizer.token_to_id("<|im_start|>"),
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling.clone(),
|
||||
eos_id,
|
||||
)
|
||||
.await
|
||||
@@ -2941,9 +2967,7 @@ impl CandleHarness {
|
||||
loaded_for_cache.prefix_cache.as_deref(),
|
||||
im_start_id,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling.clone(),
|
||||
eos_id,
|
||||
)
|
||||
})
|
||||
@@ -3273,10 +3297,14 @@ impl CandleHarness {
|
||||
None
|
||||
};
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
// Resolve sampling once: request → the model's
|
||||
// generation_config.json → built-in fallback (#272).
|
||||
let sampling = super::sampling::SamplingParams::resolve(
|
||||
&requested_sampling(&request),
|
||||
&loaded.generation_defaults,
|
||||
unix_subsec_nanos(),
|
||||
);
|
||||
let max_new = request.max_tokens.unwrap_or(8192) as usize;
|
||||
let seed = unix_subsec_nanos();
|
||||
|
||||
let eos_id = loaded
|
||||
.tokenizer
|
||||
@@ -3333,8 +3361,10 @@ impl CandleHarness {
|
||||
tracing::info!(
|
||||
prompt_len,
|
||||
max_new,
|
||||
temperature,
|
||||
?top_p,
|
||||
temperature = sampling.temperature,
|
||||
top_p = ?sampling.top_p,
|
||||
top_k = ?sampling.top_k,
|
||||
seed = sampling.seed,
|
||||
?eos_id,
|
||||
vram_free_mb,
|
||||
vram_total_mb,
|
||||
@@ -3392,9 +3422,7 @@ impl CandleHarness {
|
||||
.submit(super::engine::EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling: sampling.clone(),
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
@@ -3423,9 +3451,7 @@ impl CandleHarness {
|
||||
loaded_for_task.prefix_cache.as_deref(),
|
||||
&loaded_for_task.prefill_rate,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling.clone(),
|
||||
eos_id,
|
||||
reasoning_tokens_inner,
|
||||
tool_call_tokens_inner,
|
||||
@@ -3487,9 +3513,7 @@ impl CandleHarness {
|
||||
&prompt_tokens,
|
||||
loaded_for_task.prefix_cache.as_deref(),
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling.clone(),
|
||||
eos_id,
|
||||
reasoning_tokens_inner.as_ref(),
|
||||
tool_call_tokens_inner.as_ref(),
|
||||
@@ -4047,8 +4071,17 @@ impl Harness for CandleHarness {
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
// The model's published sampling, read from the snapshot dir
|
||||
// beside the tokenizer (#272). Absent is fine — it falls back to
|
||||
// the built-in defaults, which is what every model did before.
|
||||
let generation_defaults = super::sampling::ModelGenerationDefaults::load_from_dir(
|
||||
tokenizer_path
|
||||
.parent()
|
||||
.unwrap_or_else(|| std::path::Path::new(".")),
|
||||
);
|
||||
let loaded = Arc::new(LoadedModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
generation_defaults,
|
||||
arch: arch_local,
|
||||
tokenizer,
|
||||
device,
|
||||
@@ -4680,8 +4713,16 @@ impl CandleHarness {
|
||||
);
|
||||
}
|
||||
|
||||
// Same source as the single-GPU path (#272): the model's own
|
||||
// generation_config.json, read from the snapshot dir.
|
||||
let generation_defaults = super::sampling::ModelGenerationDefaults::load_from_dir(
|
||||
tokenizer_path
|
||||
.parent()
|
||||
.unwrap_or_else(|| std::path::Path::new(".")),
|
||||
);
|
||||
let tp_loaded = StdArc::new(TpLoadedModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
generation_defaults,
|
||||
tokenizer,
|
||||
devices: devices.clone(),
|
||||
pool: StdArc::new(TMutex::new(pool)),
|
||||
@@ -4991,10 +5032,14 @@ impl CandleHarness {
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
// Resolve sampling once: request → the model's
|
||||
// generation_config.json → built-in fallback (#272).
|
||||
let sampling = super::sampling::SamplingParams::resolve(
|
||||
&requested_sampling(&request),
|
||||
&tp.generation_defaults,
|
||||
unix_subsec_nanos(),
|
||||
);
|
||||
let max_new = request.max_tokens.unwrap_or(8192) as usize;
|
||||
let seed = unix_subsec_nanos();
|
||||
|
||||
let eos_id = tp
|
||||
.tokenizer
|
||||
@@ -5047,8 +5092,10 @@ impl CandleHarness {
|
||||
parent: &span,
|
||||
prompt_len,
|
||||
max_new,
|
||||
temperature,
|
||||
?top_p,
|
||||
temperature = sampling.temperature,
|
||||
top_p = ?sampling.top_p,
|
||||
top_k = ?sampling.top_k,
|
||||
seed = sampling.seed,
|
||||
?eos_id,
|
||||
vram_free_mb,
|
||||
vram_total_mb,
|
||||
@@ -5094,9 +5141,7 @@ impl CandleHarness {
|
||||
.submit(super::engine::EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling: sampling.clone(),
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
@@ -5160,6 +5205,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,17 +5230,10 @@ impl CandleHarness {
|
||||
}
|
||||
}
|
||||
};
|
||||
reused_for_task = reused as u32;
|
||||
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling())
|
||||
};
|
||||
|
||||
// Chunked prefill — see `chunked_prefill_tp`. Each
|
||||
@@ -5599,12 +5641,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
|
||||
@@ -5710,10 +5769,14 @@ async fn chat_completion_tp_inner(
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
// Resolve sampling once: request → the model's
|
||||
// generation_config.json → built-in fallback (#272).
|
||||
let sampling = super::sampling::SamplingParams::resolve(
|
||||
&requested_sampling(&request),
|
||||
&tp.generation_defaults,
|
||||
unix_subsec_nanos(),
|
||||
);
|
||||
let max_new = request.max_tokens.unwrap_or(8192) as usize;
|
||||
let seed = unix_subsec_nanos();
|
||||
|
||||
let eos_id = tp
|
||||
.tokenizer
|
||||
@@ -5725,8 +5788,10 @@ async fn chat_completion_tp_inner(
|
||||
model = %model_id,
|
||||
prompt_len,
|
||||
max_new,
|
||||
temperature,
|
||||
?top_p,
|
||||
temperature = sampling.temperature,
|
||||
top_p = ?sampling.top_p,
|
||||
top_k = ?sampling.top_k,
|
||||
seed = sampling.seed,
|
||||
?eos_id,
|
||||
vram_free_mb,
|
||||
vram_total_mb,
|
||||
@@ -5786,17 +5851,8 @@ async fn chat_completion_tp_inner(
|
||||
"TP chat_completion: kv cache ready"
|
||||
);
|
||||
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut generated: Vec<u32> = Vec::new();
|
||||
let mut finish_reason = "length".to_string();
|
||||
@@ -6859,22 +6915,11 @@ async fn run_inference_with_images_via_worker(
|
||||
images: Vec<super::device_worker::jobs::ImageInput>,
|
||||
image_token_id: u32,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
seed: u64,
|
||||
sampling: super::sampling::SamplingParams,
|
||||
eos_id: Option<u32>,
|
||||
) -> Result<(Vec<u32>, String)> {
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut generated: Vec<u32> = Vec::new();
|
||||
let prompt_len = prompt_tokens.len();
|
||||
@@ -7223,22 +7268,11 @@ async fn run_inference_via_worker(
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
im_start_id: Option<u32>,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
seed: u64,
|
||||
sampling: super::sampling::SamplingParams,
|
||||
eos_id: Option<u32>,
|
||||
) -> Result<(Vec<u32>, String)> {
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut generated: Vec<u32> = Vec::new();
|
||||
let prompt_len = prompt_tokens.len();
|
||||
@@ -7350,9 +7384,7 @@ async fn stream_inference_via_worker(
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
prefill_rate: &super::context_limit::PrefillRateEma,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
seed: u64,
|
||||
sampling: super::sampling::SamplingParams,
|
||||
eos_id: Option<u32>,
|
||||
reasoning_tokens: Option<ReasoningTokenPair>,
|
||||
tool_call_tokens: Option<ToolCallTokenPair>,
|
||||
@@ -7363,17 +7395,8 @@ async fn stream_inference_via_worker(
|
||||
// identity would be unsound for vision requests — they bypass the
|
||||
// prefix cache entirely (no restore, no snapshot).
|
||||
let prefix_cache = if images.is_some() { None } else { prefix_cache };
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut all_tokens: Vec<u32> = Vec::new();
|
||||
// Incremental detokenizer. Replaces the old "decode cumulative
|
||||
@@ -7416,6 +7439,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 +7456,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 +7640,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,
|
||||
@@ -7638,22 +7677,11 @@ fn run_inference(
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
im_start_id: Option<u32>,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
seed: u64,
|
||||
sampling: super::sampling::SamplingParams,
|
||||
eos_id: Option<u32>,
|
||||
) -> Result<(Vec<u32>, String)> {
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut generated: Vec<u32> = Vec::new();
|
||||
|
||||
@@ -7709,26 +7737,15 @@ fn run_inference_streaming(
|
||||
prompt_tokens: &[u32],
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
seed: u64,
|
||||
sampling: super::sampling::SamplingParams,
|
||||
eos_id: Option<u32>,
|
||||
reasoning_tokens: Option<&ReasoningTokenPair>,
|
||||
tool_call_tokens: Option<&ToolCallTokenPair>,
|
||||
tool_schemas: ToolSchemas,
|
||||
tx: &mpsc::Sender<InferenceEvent>,
|
||||
) -> Result<()> {
|
||||
let mut logits_processor = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut logits_processor =
|
||||
{ LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling()) };
|
||||
|
||||
let mut all_tokens: Vec<u32> = Vec::new();
|
||||
// Incremental detokenizer. See `stream_inference_via_worker` for
|
||||
@@ -7885,11 +7902,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
|
||||
@@ -397,9 +406,10 @@ async fn op_prefill(
|
||||
pub struct EngineRequest {
|
||||
pub prompt_tokens: Vec<u32>,
|
||||
pub max_new: usize,
|
||||
pub temperature: f64,
|
||||
pub top_p: Option<f64>,
|
||||
pub seed: u64,
|
||||
/// Sampling resolved once from request → model `generation_config`
|
||||
/// → fallback (#272), rather than loose scalars each caller could
|
||||
/// default differently.
|
||||
pub sampling: crate::harness::sampling::SamplingParams,
|
||||
pub eos_id: Option<u32>,
|
||||
pub tool_schemas: ToolSchemas,
|
||||
pub tx: mpsc::Sender<InferenceEvent>,
|
||||
@@ -439,6 +449,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 +474,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 +642,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,
|
||||
@@ -759,14 +775,10 @@ async fn prefill_join(
|
||||
session: &mut ActiveSession,
|
||||
req: EngineRequest,
|
||||
) -> Result<Option<(Slot, u64)>> {
|
||||
use candle_transformers::generation::Sampling;
|
||||
|
||||
let EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling,
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
@@ -774,21 +786,11 @@ async fn prefill_join(
|
||||
span,
|
||||
} = req;
|
||||
|
||||
let mut lp = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut lp = LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling());
|
||||
|
||||
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 +829,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 +893,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 +984,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 +1028,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 +1056,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.
|
||||
@@ -1046,9 +1175,17 @@ mod tests {
|
||||
.submit(EngineRequest {
|
||||
prompt_tokens: prompt,
|
||||
max_new,
|
||||
temperature: 0.0, // greedy — deterministic
|
||||
top_p: None,
|
||||
seed: 0,
|
||||
// Greedy, deterministic: the fixture asserts exact
|
||||
// token sequences.
|
||||
sampling: crate::harness::sampling::SamplingParams::resolve(
|
||||
&crate::harness::sampling::RequestedSampling {
|
||||
temperature: Some(0.0),
|
||||
seed: Some(0),
|
||||
..Default::default()
|
||||
},
|
||||
&crate::harness::sampling::ModelGenerationDefaults::default(),
|
||||
0,
|
||||
),
|
||||
eos_id: None,
|
||||
tool_schemas: ToolSchemas::new(),
|
||||
tx,
|
||||
|
||||
@@ -12,6 +12,7 @@ pub mod image;
|
||||
pub mod prefix_cache;
|
||||
pub mod preflight;
|
||||
pub mod preprocess;
|
||||
pub mod sampling;
|
||||
pub mod speculative;
|
||||
pub mod tp;
|
||||
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
//! Sampling parameter resolution (#272).
|
||||
//!
|
||||
//! Three sources, in priority order:
|
||||
//!
|
||||
//! 1. what the **request** asked for,
|
||||
//! 2. what the **model** published in `generation_config.json`,
|
||||
//! 3. a built-in fallback.
|
||||
//!
|
||||
//! Before this, (2) did not exist and (3) was a hardcoded
|
||||
//! `temperature = 0.7` with `top_p = None` — and `None` selects
|
||||
//! `Sampling::All`, meaning no truncation at all. So a caller that
|
||||
//! omitted sampling fields (the common case: the DeepSeek Harness sends
|
||||
//! only `temperature`) got untruncated sampling across the full ~250k
|
||||
//! vocabulary on a model whose authors published `top_k = 20,
|
||||
//! top_p = 0.95, temperature = 1.0`.
|
||||
//!
|
||||
//! Picking "no truncation" for an absent field is not a neutral default,
|
||||
//! it is the widest possible one. The neutral default is what the model
|
||||
//! shipped with.
|
||||
//!
|
||||
//! Resolved **once** per request and threaded whole. The eight
|
||||
//! `LogitsProcessor` construction sites this replaces could each have
|
||||
//! disagreed about a default — the same shape of bug as #252, where two
|
||||
//! config files described one model and drifted.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
/// Fallback temperature when neither the request nor the model says.
|
||||
/// Retained from the previous hardcoded value so a model without a
|
||||
/// `generation_config.json` behaves exactly as it did before.
|
||||
const FALLBACK_TEMPERATURE: f64 = 0.7;
|
||||
|
||||
/// Repetition penalty applied to recently-generated tokens before
|
||||
/// sampling. 1.0 disables it; >1.0 makes recently-emitted tokens less
|
||||
/// likely. mistral.rs and llama.cpp default to 1.1, which is enough to
|
||||
/// stop small quantized models from degenerating into "Wait, no, no..."
|
||||
/// loops without distorting normal output.
|
||||
pub const DEFAULT_REPEAT_PENALTY: f32 = 1.1;
|
||||
|
||||
/// Number of recently-generated tokens fed into the repetition penalty.
|
||||
/// Matches the candle quantized-qwen3 example default.
|
||||
///
|
||||
/// Defensible for chat and poorly suited to a 24k-token reasoning block,
|
||||
/// where a derivation repeated 500 tokens later is well outside the
|
||||
/// window — which is why it is now a per-request knob rather than a
|
||||
/// constant.
|
||||
pub const DEFAULT_REPEAT_LAST_N: usize = 64;
|
||||
|
||||
/// What a model published in its `generation_config.json`.
|
||||
///
|
||||
/// Every field optional: the file is a HuggingFace convention, not a
|
||||
/// schema, and models populate different subsets of it.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct ModelGenerationDefaults {
|
||||
pub temperature: Option<f64>,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
}
|
||||
|
||||
impl ModelGenerationDefaults {
|
||||
/// Read `generation_config.json` from a model snapshot directory.
|
||||
///
|
||||
/// Absent or unparseable is not an error — we fall back exactly as
|
||||
/// before — but it is logged, because silently guessing the model's
|
||||
/// sampling is how the untruncated default went unnoticed.
|
||||
pub fn load_from_dir(dir: &Path) -> Self {
|
||||
let path = dir.join("generation_config.json");
|
||||
let text = match std::fs::read_to_string(&path) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"sampling: no generation_config.json; using built-in fallbacks"
|
||||
);
|
||||
return Self::default();
|
||||
}
|
||||
};
|
||||
match serde_json::from_str::<Self>(&text) {
|
||||
Ok(d) => {
|
||||
tracing::info!(
|
||||
path = %path.display(),
|
||||
temperature = ?d.temperature,
|
||||
top_p = ?d.top_p,
|
||||
top_k = ?d.top_k,
|
||||
"sampling: model generation defaults loaded"
|
||||
);
|
||||
d
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"sampling: generation_config.json did not parse; using built-in fallbacks"
|
||||
);
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the caller asked for. All optional — absent means "no opinion",
|
||||
/// which defers to the model's published value.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RequestedSampling {
|
||||
pub temperature: Option<f64>,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
pub seed: Option<u64>,
|
||||
pub repeat_penalty: Option<f32>,
|
||||
pub repeat_last_n: Option<usize>,
|
||||
}
|
||||
|
||||
/// Fully-resolved sampling for one request.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct SamplingParams {
|
||||
pub temperature: f64,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
pub seed: u64,
|
||||
pub repeat_penalty: f32,
|
||||
pub repeat_last_n: usize,
|
||||
}
|
||||
|
||||
impl SamplingParams {
|
||||
/// Resolve request over model defaults over fallback.
|
||||
///
|
||||
/// `random_seed` is the per-request seed used when the caller did
|
||||
/// not pin one; it is passed in rather than generated here so this
|
||||
/// stays a pure function and the "same seed reproduces" property is
|
||||
/// testable without a clock.
|
||||
pub fn resolve(
|
||||
requested: &RequestedSampling,
|
||||
model: &ModelGenerationDefaults,
|
||||
random_seed: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
temperature: requested
|
||||
.temperature
|
||||
.or(model.temperature)
|
||||
.unwrap_or(FALLBACK_TEMPERATURE),
|
||||
// No `unwrap_or` here on purpose: `None` still means "do not
|
||||
// truncate", but it is now only reachable when neither the
|
||||
// caller nor the model expressed a preference, rather than
|
||||
// whenever the caller stayed quiet.
|
||||
top_p: requested.top_p.or(model.top_p),
|
||||
top_k: requested.top_k.or(model.top_k),
|
||||
seed: requested.seed.unwrap_or(random_seed),
|
||||
repeat_penalty: requested.repeat_penalty.unwrap_or(DEFAULT_REPEAT_PENALTY),
|
||||
repeat_last_n: requested.repeat_last_n.unwrap_or(DEFAULT_REPEAT_LAST_N),
|
||||
}
|
||||
}
|
||||
|
||||
/// The candle sampling strategy these parameters select.
|
||||
///
|
||||
/// `temperature <= 0` is greedy regardless of the rest — a caller
|
||||
/// asking for determinism that way should not have it silently
|
||||
/// softened by a `top_k` the model happened to publish.
|
||||
pub fn to_sampling(&self) -> candle_transformers::generation::Sampling {
|
||||
use candle_transformers::generation::Sampling;
|
||||
if self.temperature <= 0.0 {
|
||||
return Sampling::ArgMax;
|
||||
}
|
||||
let temperature = self.temperature;
|
||||
match (self.top_k, self.top_p) {
|
||||
(Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature },
|
||||
(Some(k), None) => Sampling::TopK { k, temperature },
|
||||
(None, Some(p)) => Sampling::TopP { p, temperature },
|
||||
(None, None) => Sampling::All { temperature },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_transformers::generation::Sampling;
|
||||
|
||||
fn qwen3_defaults() -> ModelGenerationDefaults {
|
||||
ModelGenerationDefaults {
|
||||
temperature: Some(1.0),
|
||||
top_p: Some(0.95),
|
||||
top_k: Some(20),
|
||||
}
|
||||
}
|
||||
|
||||
/// The case that motivated #271: a caller that names no sampling
|
||||
/// parameters must get the model's published ones, not untruncated
|
||||
/// sampling at a hardcoded 0.7.
|
||||
#[test]
|
||||
fn a_silent_request_inherits_the_models_published_sampling() {
|
||||
let p = SamplingParams::resolve(&RequestedSampling::default(), &qwen3_defaults(), 42);
|
||||
assert_eq!(p.temperature, 1.0);
|
||||
assert_eq!(p.top_p, Some(0.95));
|
||||
assert_eq!(p.top_k, Some(20));
|
||||
assert_eq!(
|
||||
p.to_sampling(),
|
||||
Sampling::TopKThenTopP {
|
||||
k: 20,
|
||||
p: 0.95,
|
||||
temperature: 1.0
|
||||
},
|
||||
"the combination Qwen3 asks for must be expressible"
|
||||
);
|
||||
}
|
||||
|
||||
/// The caller still wins. A harness that deliberately sets a low
|
||||
/// temperature must not have it overridden by the model's file.
|
||||
#[test]
|
||||
fn the_request_overrides_the_model() {
|
||||
let requested = RequestedSampling {
|
||||
temperature: Some(0.2),
|
||||
top_k: Some(5),
|
||||
..Default::default()
|
||||
};
|
||||
let p = SamplingParams::resolve(&requested, &qwen3_defaults(), 42);
|
||||
assert_eq!(p.temperature, 0.2);
|
||||
assert_eq!(p.top_k, Some(5));
|
||||
// Untouched fields still fall through to the model.
|
||||
assert_eq!(p.top_p, Some(0.95));
|
||||
}
|
||||
|
||||
/// A model with no `generation_config.json` must behave exactly as
|
||||
/// it did before this landed, so nothing regresses on the models
|
||||
/// that do not ship one.
|
||||
#[test]
|
||||
fn without_model_defaults_the_old_behaviour_is_preserved() {
|
||||
let p = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(p.temperature, FALLBACK_TEMPERATURE);
|
||||
assert_eq!(p.top_p, None);
|
||||
assert_eq!(p.top_k, None);
|
||||
assert_eq!(p.to_sampling(), Sampling::All { temperature: 0.7 });
|
||||
}
|
||||
|
||||
/// A pinned seed is honoured; an absent one takes the per-request
|
||||
/// random value. Reproducibility is the whole point — the text paths
|
||||
/// previously called `unix_subsec_nanos()` unconditionally, so
|
||||
/// `seed` was accepted and discarded.
|
||||
#[test]
|
||||
fn seed_is_pinned_when_given_and_random_otherwise() {
|
||||
let pinned = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
seed: Some(7),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(pinned.seed, 7);
|
||||
|
||||
let unpinned = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(unpinned.seed, 42);
|
||||
}
|
||||
|
||||
/// Greedy beats every truncation setting: a caller asking for
|
||||
/// `temperature = 0` wants determinism, not a softened version of it.
|
||||
#[test]
|
||||
fn zero_temperature_is_greedy_even_with_model_top_k() {
|
||||
let requested = RequestedSampling {
|
||||
temperature: Some(0.0),
|
||||
..Default::default()
|
||||
};
|
||||
let p = SamplingParams::resolve(&requested, &qwen3_defaults(), 42);
|
||||
assert_eq!(p.to_sampling(), Sampling::ArgMax);
|
||||
}
|
||||
|
||||
/// Each truncation knob alone selects the matching strategy.
|
||||
#[test]
|
||||
fn single_knobs_select_their_own_strategy() {
|
||||
let only_k = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
top_k: Some(20),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(
|
||||
only_k.to_sampling(),
|
||||
Sampling::TopK {
|
||||
k: 20,
|
||||
temperature: 0.7
|
||||
}
|
||||
);
|
||||
|
||||
let only_p = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
top_p: Some(0.9),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(
|
||||
only_p.to_sampling(),
|
||||
Sampling::TopP {
|
||||
p: 0.9,
|
||||
temperature: 0.7
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// Repetition penalty and its window are now per-request, defaulting
|
||||
/// to the values that were previously constants.
|
||||
#[test]
|
||||
fn repetition_knobs_default_to_the_former_constants() {
|
||||
let d = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(d.repeat_penalty, DEFAULT_REPEAT_PENALTY);
|
||||
assert_eq!(d.repeat_last_n, DEFAULT_REPEAT_LAST_N);
|
||||
|
||||
let tuned = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
repeat_penalty: Some(1.05),
|
||||
repeat_last_n: Some(512),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(tuned.repeat_penalty, 1.05);
|
||||
assert_eq!(tuned.repeat_last_n, 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_generation_config_is_not_an_error() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert!(d.temperature.is_none() && d.top_p.is_none() && d.top_k.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_real_generation_config_parses() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
// Verbatim from Qwen/Qwen3.8-27B, extra fields included — the
|
||||
// reader must tolerate the ones it does not model.
|
||||
std::fs::write(
|
||||
dir.path().join("generation_config.json"),
|
||||
r#"{"bos_token_id":248044,"do_sample":true,"eos_token_id":[248046,248044],
|
||||
"pad_token_id":248044,"temperature":1.0,"top_k":20,"top_p":0.95}"#,
|
||||
)
|
||||
.expect("write");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert_eq!(d.temperature, Some(1.0));
|
||||
assert_eq!(d.top_p, Some(0.95));
|
||||
assert_eq!(d.top_k, Some(20));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_generation_config_falls_back_rather_than_failing() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
std::fs::write(dir.path().join("generation_config.json"), "{not json").expect("write");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert!(d.temperature.is_none());
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -156,6 +156,10 @@ pub fn request_to_chat(req: ResponsesRequest) -> Result<ChatCompletionRequest, T
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
top_k: req.top_k,
|
||||
seed: req.seed,
|
||||
repetition_penalty: req.repetition_penalty,
|
||||
repeat_last_n: req.repeat_last_n,
|
||||
max_tokens: req.max_output_tokens,
|
||||
// Responses spells the cap `max_output_tokens`, already
|
||||
// resolved onto the legacy field above.
|
||||
@@ -365,6 +369,12 @@ async fn run_projection(
|
||||
// final `response.completed` payload so clients that only read
|
||||
// the terminal response still see the calls.
|
||||
let mut tool_items: Vec<Value> = Vec::new();
|
||||
// Reasoning item state (#267) and whether the message item's
|
||||
// opening frames have gone out. The message is opened lazily,
|
||||
// because its `output_index` depends on whether a reasoning item
|
||||
// claimed index 0 first.
|
||||
let mut reasoning = ReasoningTracker::default();
|
||||
let mut message_open = false;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
@@ -378,12 +388,22 @@ async fn run_projection(
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Thinking is over the moment visible text starts.
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
if !message_open {
|
||||
if !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
message_open = true;
|
||||
}
|
||||
accumulated.push_str(&text);
|
||||
let frame = ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_TEXT_DELTA,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": reasoning.message_index(),
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}),
|
||||
@@ -392,10 +412,45 @@ async fn run_projection(
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::ReasoningDelta(_) => {
|
||||
// No representation in our Responses model yet.
|
||||
// Stage where it'd land: a `response.reasoning_*`
|
||||
// event family alongside `response.output_text.*`.
|
||||
InferenceEvent::ReasoningDelta(text) => {
|
||||
// Stream the think block as its own output item (#267).
|
||||
// This used to be dropped on the floor, which meant a
|
||||
// model that reasoned for minutes emitted no events at
|
||||
// all and clients closed the stream as dead — dsh
|
||||
// reports `pi-ai stream idle timeout after 300000ms`.
|
||||
// Its watchdog resets on any parsed event, so surfacing
|
||||
// reasoning is what keeps a thinking model alive on the
|
||||
// wire; SSE comment keep-alives cannot, because idle
|
||||
// timers count events, not comments.
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if !reasoning.open {
|
||||
if reasoning.item.is_some() {
|
||||
// Reasoning resumed after visible text. Our
|
||||
// model has one reasoning item per response, so
|
||||
// fold it into the existing summary rather than
|
||||
// opening a second item out of order.
|
||||
continue;
|
||||
}
|
||||
if !emit_reasoning_open_frames(&tx, &meta).await {
|
||||
return;
|
||||
}
|
||||
reasoning.open = true;
|
||||
}
|
||||
reasoning.text.push_str(&text);
|
||||
let frame = ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
data: json!({
|
||||
"item_id": reasoning_item_id(&meta),
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"delta": text,
|
||||
}),
|
||||
};
|
||||
if tx.send(frame).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::ToolCall {
|
||||
index,
|
||||
@@ -403,9 +458,19 @@ async fn run_projection(
|
||||
name,
|
||||
arguments,
|
||||
} => {
|
||||
// The message item opened by the start frames holds
|
||||
// output_index 0; function_call items follow it.
|
||||
let output_index = 1 + index as u64;
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
if !message_open {
|
||||
if !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
message_open = true;
|
||||
}
|
||||
// function_call items follow the message item, which
|
||||
// itself follows any reasoning item — so the base is the
|
||||
// message's index, not a hardcoded 0.
|
||||
let output_index = reasoning.message_index() + 1 + index as u64;
|
||||
let item_id = format!(
|
||||
"fc_{}_{index}",
|
||||
meta.response_id.trim_start_matches("resp_")
|
||||
@@ -479,6 +544,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: _,
|
||||
@@ -495,7 +561,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,
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -508,18 +579,158 @@ async fn run_projection(
|
||||
return;
|
||||
}
|
||||
|
||||
// A stream can end while still thinking — a reasoning model that
|
||||
// exhausts max_output_tokens mid-block never emits visible text.
|
||||
// Close the item so the client sees a finished response rather than
|
||||
// a reasoning item left open forever.
|
||||
if !close_reasoning(&tx, &meta, &mut reasoning).await {
|
||||
return;
|
||||
}
|
||||
// The message item is opened lazily, so a response that produced no
|
||||
// text (tool-call-only turns, or the case above) has not announced
|
||||
// it yet. The finish frames reference it, so it must exist first.
|
||||
if !message_open && !emit_message_open_frames(&tx, &meta, reasoning.message_index()).await {
|
||||
return;
|
||||
}
|
||||
|
||||
let reason = finish.unwrap_or(FinishReason::Stop);
|
||||
let _ = emit_finish_frames(
|
||||
&tx,
|
||||
&meta,
|
||||
&accumulated,
|
||||
reason,
|
||||
usage.as_ref(),
|
||||
&tool_items,
|
||||
FinishContext {
|
||||
full_text: &accumulated,
|
||||
reason,
|
||||
usage: usage.as_ref(),
|
||||
tool_items: &tool_items,
|
||||
message_index: reasoning.message_index(),
|
||||
reasoning_item: reasoning.item.as_ref(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Streaming state for a reasoning model's think block (#267).
|
||||
///
|
||||
/// The block is its own output item, ahead of the message, so the
|
||||
/// message's `output_index` depends on whether reasoning happened —
|
||||
/// which is only known once the first token arrives. This tracks that
|
||||
/// decision and the accumulated summary text.
|
||||
#[derive(Default)]
|
||||
struct ReasoningTracker {
|
||||
/// Frames opening the reasoning item have been emitted.
|
||||
open: bool,
|
||||
/// Everything the model has thought so far, for the `*.done` frames
|
||||
/// and the terminal `completed` payload.
|
||||
text: String,
|
||||
/// The finished reasoning item, replayed inside `response.completed`
|
||||
/// so a client that only reads the terminal response still sees it.
|
||||
item: Option<Value>,
|
||||
}
|
||||
|
||||
impl ReasoningTracker {
|
||||
/// `output_index` the message item takes: 1 when a reasoning item
|
||||
/// occupies 0, else 0.
|
||||
fn message_index(&self) -> u64 {
|
||||
u64::from(self.open || self.item.is_some())
|
||||
}
|
||||
}
|
||||
|
||||
fn reasoning_item_id(meta: &ResponseMeta) -> String {
|
||||
format!("rs_{}", meta.response_id.trim_start_matches("resp_"))
|
||||
}
|
||||
|
||||
/// Open the reasoning item and its summary part, once.
|
||||
async fn emit_reasoning_open_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
) -> bool {
|
||||
let item_id = reasoning_item_id(meta);
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_ADDED,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": {
|
||||
"type": "reasoning",
|
||||
"id": item_id,
|
||||
"summary": [],
|
||||
"status": "in_progress",
|
||||
},
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_PART_ADDED,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"part": { "type": "summary_text", "text": "" },
|
||||
}),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Close the reasoning item, if one was opened, and record it for the
|
||||
/// terminal payload. Idempotent — every path that emits a non-reasoning
|
||||
/// item calls this first, and only the first call does anything.
|
||||
async fn close_reasoning(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
reasoning: &mut ReasoningTracker,
|
||||
) -> bool {
|
||||
if !reasoning.open {
|
||||
return true;
|
||||
}
|
||||
reasoning.open = false;
|
||||
let item_id = reasoning_item_id(meta);
|
||||
let full_item = json!({
|
||||
"type": "reasoning",
|
||||
"id": item_id,
|
||||
"summary": [{ "type": "summary_text", "text": reasoning.text }],
|
||||
"status": "completed",
|
||||
});
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_TEXT_DONE,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"text": reasoning.text,
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::REASONING_SUMMARY_PART_DONE,
|
||||
data: json!({
|
||||
"item_id": item_id,
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"part": { "type": "summary_text", "text": reasoning.text },
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_DONE,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": full_item.clone(),
|
||||
}),
|
||||
},
|
||||
];
|
||||
reasoning.item = Some(full_item);
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &ResponseMeta) -> bool {
|
||||
let shell = response_shell(meta, "in_progress", &[], None);
|
||||
let frames = [
|
||||
@@ -540,10 +751,31 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
event_name: events::IN_PROGRESS,
|
||||
data: json!({ "response": shell }),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Open the message item, once, at `output_index`.
|
||||
///
|
||||
/// Deferred rather than emitted with the start frames because the
|
||||
/// message is no longer guaranteed to be `output[0]`: a reasoning model
|
||||
/// puts its think block in an item ahead of it, and an item's index has
|
||||
/// to be right the first time it is announced.
|
||||
async fn emit_message_open_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
output_index: u64,
|
||||
) -> bool {
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_ADDED,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"output_index": output_index,
|
||||
"item": empty_message_item(&meta.message_item_id),
|
||||
}),
|
||||
},
|
||||
@@ -551,7 +783,7 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
event_name: events::CONTENT_PART_ADDED,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": output_index,
|
||||
"content_index": 0,
|
||||
"part": { "type": "output_text", "text": "", "annotations": [] },
|
||||
}),
|
||||
@@ -565,14 +797,34 @@ async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &Respon
|
||||
true
|
||||
}
|
||||
|
||||
/// Everything the terminal frames need that isn't stream identity.
|
||||
/// Grouped rather than passed loose: the reasoning item (#267) pushed
|
||||
/// the argument list past the point where positional `&str`/`u64`
|
||||
/// parameters are safe to read at a call site.
|
||||
struct FinishContext<'a> {
|
||||
full_text: &'a str,
|
||||
reason: FinishReason,
|
||||
usage: Option<&'a ResponsesUsage>,
|
||||
tool_items: &'a [Value],
|
||||
/// Where the message item sits — 1 when a reasoning item took 0.
|
||||
message_index: u64,
|
||||
/// The completed reasoning item, replayed inside `completed`.
|
||||
reasoning_item: Option<&'a Value>,
|
||||
}
|
||||
|
||||
async fn emit_finish_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
full_text: &str,
|
||||
reason: FinishReason,
|
||||
usage: Option<&ResponsesUsage>,
|
||||
tool_items: &[Value],
|
||||
ctx: FinishContext<'_>,
|
||||
) -> bool {
|
||||
let FinishContext {
|
||||
full_text,
|
||||
reason,
|
||||
usage,
|
||||
tool_items,
|
||||
message_index,
|
||||
reasoning_item,
|
||||
} = ctx;
|
||||
let status = finish_to_status(reason);
|
||||
let full_part = json!({
|
||||
"type": "output_text",
|
||||
@@ -586,10 +838,14 @@ async fn emit_finish_frames(
|
||||
"content": [full_part.clone()],
|
||||
"status": status,
|
||||
});
|
||||
// Terminal output array: the message item (always present — its
|
||||
// open frames were emitted eagerly at stream start) followed by
|
||||
// any completed function_call items.
|
||||
let mut output_items = Vec::with_capacity(1 + tool_items.len());
|
||||
// Terminal output array, in wire order: the reasoning item (when
|
||||
// the model thought), then the message, then any completed
|
||||
// function_call items. Order matters — a client reconstructing the
|
||||
// response from `completed` alone should see what the stream showed.
|
||||
let mut output_items = Vec::with_capacity(2 + tool_items.len());
|
||||
if let Some(item) = reasoning_item {
|
||||
output_items.push(item.clone());
|
||||
}
|
||||
output_items.push(full_item.clone());
|
||||
output_items.extend(tool_items.iter().cloned());
|
||||
let frames = [
|
||||
@@ -597,7 +853,7 @@ async fn emit_finish_frames(
|
||||
event_name: events::OUTPUT_TEXT_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"content_index": 0,
|
||||
"text": full_text,
|
||||
}),
|
||||
@@ -606,7 +862,7 @@ async fn emit_finish_frames(
|
||||
event_name: events::CONTENT_PART_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"content_index": 0,
|
||||
"part": full_part,
|
||||
}),
|
||||
@@ -614,7 +870,7 @@ async fn emit_finish_frames(
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_DONE,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"output_index": message_index,
|
||||
"item": full_item.clone(),
|
||||
}),
|
||||
},
|
||||
@@ -756,6 +1012,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -778,6 +1038,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -801,6 +1065,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: Some("resp_prev".into()),
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -833,6 +1101,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -863,6 +1135,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -908,6 +1184,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -945,6 +1225,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -971,6 +1255,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1095,6 +1383,208 @@ mod tests {
|
||||
out
|
||||
}
|
||||
|
||||
/// The regression #267 was filed for: a reasoning model must put
|
||||
/// events on the wire while it thinks.
|
||||
///
|
||||
/// These deltas used to be dropped, so a model that reasoned for
|
||||
/// minutes emitted nothing at all and clients closed the stream as
|
||||
/// dead — dsh reports `pi-ai stream idle timeout after 300000ms`
|
||||
/// and retries, so a task whose thinking outlasts the client's idle
|
||||
/// 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);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("think".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("ing".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("answer".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 2,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert_eq!(
|
||||
names,
|
||||
vec![
|
||||
events::CREATED,
|
||||
events::IN_PROGRESS,
|
||||
events::OUTPUT_ITEM_ADDED,
|
||||
events::REASONING_SUMMARY_PART_ADDED,
|
||||
events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
events::REASONING_SUMMARY_TEXT_DELTA,
|
||||
events::REASONING_SUMMARY_TEXT_DONE,
|
||||
events::REASONING_SUMMARY_PART_DONE,
|
||||
events::OUTPUT_ITEM_DONE,
|
||||
events::OUTPUT_ITEM_ADDED,
|
||||
events::CONTENT_PART_ADDED,
|
||||
events::OUTPUT_TEXT_DELTA,
|
||||
events::OUTPUT_TEXT_DONE,
|
||||
events::CONTENT_PART_DONE,
|
||||
events::OUTPUT_ITEM_DONE,
|
||||
events::COMPLETED,
|
||||
]
|
||||
);
|
||||
|
||||
// The reasoning item owns output_index 0, so the message shifts
|
||||
// to 1 — an item's index has to be right when it is announced,
|
||||
// which is why the message is opened lazily.
|
||||
let reasoning_added = &frames[2];
|
||||
assert_eq!(reasoning_added.data["output_index"], 0);
|
||||
assert_eq!(reasoning_added.data["item"]["type"], "reasoning");
|
||||
let message_added = &frames[9];
|
||||
assert_eq!(message_added.data["output_index"], 1);
|
||||
assert_eq!(message_added.data["item"]["type"], "message");
|
||||
assert_eq!(frames[11].data["output_index"], 1, "text delta index");
|
||||
|
||||
// Accumulated thinking, not just the last fragment.
|
||||
assert_eq!(frames[6].data["text"], "thinking");
|
||||
|
||||
// The terminal payload replays both items, in wire order, so a
|
||||
// client reading only `completed` sees what the stream showed.
|
||||
let output = frames.last().unwrap().data["response"]["output"]
|
||||
.as_array()
|
||||
.expect("output array");
|
||||
assert_eq!(output.len(), 2);
|
||||
assert_eq!(output[0]["type"], "reasoning");
|
||||
assert_eq!(output[0]["summary"][0]["text"], "thinking");
|
||||
assert_eq!(output[1]["type"], "message");
|
||||
}
|
||||
|
||||
/// A model can exhaust its budget mid-thought and never produce
|
||||
/// visible text. The reasoning item must still be closed and the
|
||||
/// message still opened, or the client waits on a response that
|
||||
/// never resolves.
|
||||
#[tokio::test]
|
||||
async fn reasoning_only_stream_still_closes_cleanly() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("pondering".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Length,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert!(names.contains(&events::REASONING_SUMMARY_PART_DONE));
|
||||
assert!(
|
||||
names.contains(&events::CONTENT_PART_ADDED),
|
||||
"message item must still be opened so the finish frames refer to something"
|
||||
);
|
||||
assert_eq!(names.last(), Some(&events::COMPLETED));
|
||||
let last = frames.last().unwrap();
|
||||
assert_eq!(last.data["response"]["status"], "incomplete");
|
||||
}
|
||||
|
||||
/// Tool calls sit after the message, which sits after any reasoning
|
||||
/// item — so their index is relative to the message, not a
|
||||
/// hardcoded 0.
|
||||
#[tokio::test]
|
||||
async fn tool_call_index_shifts_past_a_reasoning_item() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("plan".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::ToolCall {
|
||||
index: 0,
|
||||
id: "call_1".into(),
|
||||
name: "bash".into(),
|
||||
arguments: "{}".into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::ToolCalls,
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let call_added = frames
|
||||
.iter()
|
||||
.find(|f| {
|
||||
f.event_name == events::OUTPUT_ITEM_ADDED
|
||||
&& f.data["item"]["type"] == "function_call"
|
||||
})
|
||||
.expect("function_call item announced");
|
||||
// reasoning=0, message=1, first call=2.
|
||||
assert_eq!(call_added.data["output_index"], 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_stream_emits_expected_event_sequence() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
@@ -1112,6 +1602,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1173,6 +1664,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1245,6 +1737,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1274,6 +1767,7 @@ mod tests {
|
||||
prompt_tokens: 30,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 4,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1305,6 +1799,7 @@ mod tests {
|
||||
prompt_tokens: 8,
|
||||
completion_tokens: 3,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1330,6 +1825,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1377,6 +1873,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