Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/webui/runtime.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@

`completed_requests_total`, `completion_tokens_total` and `generation_time_ms_total` come from successfully recorded route results. Their window is the loaded provider's lifetime, and they can reset after model reload. They exclude errors not recorded by the route and are not admission counters. Generation time has route-specific semantics, including elapsed time for some non-chat routes; it is not a decode-only denominator.

A chat completion counts the same whether it streams or not. A streamed one, which is what the WebUI chat sends, is recorded once when its generation finishes, with the token counts its usage chunk reports, whether or not the request asked for that chunk. A stream the client cancelled is not a completion and is not counted: a disconnect of a stream sent without `X-Conversation-Id`, `DELETE /v1/stream`, or replacement of the session by a new request with the same `X-Conversation-Id`. A resumable stream (`X-Conversation-Id`) whose client only disconnects keeps generating into its session, so it is counted when it finishes.

`decode_tokens_total` and `decode_time_us_total` read the existing completed-request `BatchObservability` counters. They remain unavailable until a completed decode timing sample has been published. A backend that never publishes them is not shown as a measured zero. Independent atomic reads must not be converted into a claimed exact per-request or instantaneous rate. TTFT (request send to first output token, with thinking/content distinctions) and decode rate remain unavailable on this runtime surface; request-aware chat instrumentation is a separate source.

## Slots and context
Expand Down
22 changes: 21 additions & 1 deletion src/server/routes/chat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
//! preparation and option merging to shared helpers, and streams chunk payloads
//! back through `server/streaming.rs`.

use std::sync::atomic::Ordering;

use axum::{
Json,
extract::State,
Expand Down Expand Up @@ -1664,7 +1666,7 @@ async fn stream_chat_completion(
prepared.videos,
prepared.media,
queue_reservation,
cancelled,
cancelled.clone(),
&live,
|token, lp_data| {
slot.on_token(&token);
Expand Down Expand Up @@ -1932,6 +1934,24 @@ async fn stream_chat_completion(
);
}

// Record the completion once (#1911), with the counts the usage chunk
// below carries, whether or not `stream_options.include_usage` asked
// for that chunk. A stream the client cancelled (a disconnect without
// `X-Conversation-Id`, `DELETE /v1/stream`, or session replacement) is
// not a completion. The worker still answers `Ok` with the partial
// count after an abort, so the flag is what tells the two apart. A
// resumable stream whose client only disconnected keeps generating,
// leaves the flag clear and is counted.
if let Ok(r) = &result
&& !cancelled.load(Ordering::Acquire)
{
state.metrics.record_request(
r.prompt_tokens,
r.completion_tokens,
r.generation_time_ms,
);
}

// Check for tool calls in accumulated output
let mut finish_reason = match &result {
Ok(r) => r.finish_reason.clone(),
Expand Down
4 changes: 2 additions & 2 deletions src/server/routes/metrics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -144,13 +144,13 @@ pub async fn metrics(State(state): State<AppState>) -> Response {

let pp_snapshot = state.pp_observability.snapshot();
let body = format!(
"# HELP mlxcel_requests_total Total number of generation requests\n\
"# HELP mlxcel_requests_total Total number of completed generation requests; chat completion streams cancelled by the client are excluded\n\
# TYPE mlxcel_requests_total counter\n\
mlxcel_requests_total {requests}\n\
# HELP mlxcel_prompt_tokens_total Total prompt tokens processed\n\
# TYPE mlxcel_prompt_tokens_total counter\n\
mlxcel_prompt_tokens_total {prompt_tokens}\n\
# HELP mlxcel_completion_tokens_total Total completion tokens generated\n\
# HELP mlxcel_completion_tokens_total Total completion tokens of the requests counted in mlxcel_requests_total\n\
# TYPE mlxcel_completion_tokens_total counter\n\
mlxcel_completion_tokens_total {completion_tokens}\n\
# HELP mlxcel_generation_time_seconds_total Total generation time in seconds\n\
Expand Down
241 changes: 241 additions & 0 deletions src/server/routes/stream_route_metrics_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
// Copyright 2025-2026 Lablup Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! `Metrics` recording for streamed chat completions (#1911).
//!
//! The WebUI chat always streams, and its Activity tiles read
//! `Metrics.requests_total` and `completion_tokens_total`, so a streamed
//! completion has to be recorded like a non-streamed one. The scripted
//! provider's result reports `prompt_tokens: 1`, `generation_time_ms: 1` and
//! one completion token per emitted step.

use std::sync::atomic::Ordering;

use axum::http::StatusCode;
use tower::ServiceExt;

use super::{
chat_request, collect_body, get_request, read_until, scripted_app_with_state, wait_until,
};
use crate::server::model_provider::ScriptedStreamHandle;
use crate::server::{AppState, ServerConfig};

/// The tokens a completed stream emits. Distinctive enough that reading one
/// back cannot match the completion id or the role chunk.
const TOKENS: [&str; 3] = ["alpha", "beta", "gamma"];

/// Drive one plain (no `X-Conversation-Id`) streaming chat to completion:
/// emit [`TOKENS`], finish, and read the body through `[DONE]`.
///
/// The route records the completion on the generation thread before it sends
/// `[DONE]`, and the SSE channel orders the two, so the counters are final
/// once this returns.
async fn complete_plain_stream(app: &axum::Router, handle: &ScriptedStreamHandle) {
let response = app
.clone()
.oneshot(chat_request(None, None, None))
.await
.expect("stream request");
assert_eq!(response.status(), StatusCode::OK);
let mut body = response.into_body();
for token in TOKENS {
handle.token(token);
}
handle.finish();
let mut bytes = Vec::new();
read_until(&mut body, &mut bytes, "[DONE]").await;
}

/// Wait until the streaming generation task has returned.
///
/// The task holds its completion-control registration until it returns,
/// which is after the point where it decides whether to record, so an empty
/// registry means that decision is made. Asserting a counter is still zero
/// any earlier would pass whether or not the route skips the record.
async fn wait_for_generation_task_exit(state: &AppState) {
wait_until("the streaming generation task exits", || {
state.completion_controls.len_for_tests() == 0
})
.await;
}

/// The acceptance case: a completed stream advances every counter by the
/// result's own numbers. The request sets no `stream_options`, so the record
/// does not depend on `include_usage`.
#[tokio::test]
async fn completed_stream_records_one_request_with_its_counts() {
let (app, handle, _options_rx, state) = scripted_app_with_state(ServerConfig::default());

complete_plain_stream(&app, &handle).await;

let m = &state.metrics;
assert_eq!(m.requests_total.load(Ordering::Relaxed), 1);
assert_eq!(
m.completion_tokens_total.load(Ordering::Relaxed),
TOKENS.len() as u64
);
assert_eq!(m.prompt_tokens_total.load(Ordering::Relaxed), 1);
assert_eq!(m.generation_time_ms_total.load(Ordering::Relaxed), 1);

// Recorded once: nothing after `[DONE]` adds a second record.
wait_for_generation_task_exit(&state).await;
assert_eq!(m.requests_total.load(Ordering::Relaxed), 1);
}

/// A plain stream the client disconnects from is cancelled, and a cancelled
/// stream is not a completion. The worker still answers `Ok` with the tokens
/// it emitted before the abort, so only the cancellation flag keeps this
/// request out of the counters.
#[tokio::test]
async fn client_cancelled_plain_stream_is_not_recorded() {
let (app, handle, _options_rx, state) = scripted_app_with_state(ServerConfig::default());

let response = app
.clone()
.oneshot(chat_request(None, None, None))
.await
.expect("stream request");
assert_eq!(response.status(), StatusCode::OK);
let mut body = response.into_body();
let mut bytes = Vec::new();
handle.token(TOKENS[0]);
read_until(&mut body, &mut bytes, TOKENS[0]).await;
drop(body);

let flag = handle.cancellation_flag(0).expect("one generation");
wait_until("disconnect cancels the plain stream", || {
handle.token(TOKENS[1]);
flag.load(Ordering::Acquire)
})
.await;
handle.finish();
wait_for_generation_task_exit(&state).await;

let m = &state.metrics;
assert_eq!(
m.requests_total.load(Ordering::Relaxed),
0,
"a client-cancelled stream must not count as a completed request"
);
assert_eq!(m.completion_tokens_total.load(Ordering::Relaxed), 0);
}

/// A resumable stream survives a client disconnect: the generation runs to
/// the end into its session, the cancellation flag stays clear, and the
/// finished completion counts with every token it produced.
#[tokio::test]
async fn resumable_stream_is_recorded_after_a_client_disconnect() {
let (app, handle, _options_rx, state) = scripted_app_with_state(ServerConfig::default());

let response = app
.clone()
.oneshot(chat_request(Some("conv-metrics"), None, None))
.await
.expect("stream request");
assert_eq!(response.status(), StatusCode::OK);
let mut body = response.into_body();
let mut bytes = Vec::new();
handle.token(TOKENS[0]);
read_until(&mut body, &mut bytes, TOKENS[0]).await;
drop(body);

for token in &TOKENS[1..] {
handle.token(token);
}
handle.finish();
wait_for_generation_task_exit(&state).await;

let flag = handle.cancellation_flag(0).expect("one generation");
assert!(
!flag.load(Ordering::Acquire),
"a client disconnect must not cancel a resumable stream"
);
let m = &state.metrics;
assert_eq!(m.requests_total.load(Ordering::Relaxed), 1);
assert_eq!(
m.completion_tokens_total.load(Ordering::Relaxed),
TOKENS.len() as u64
);
}

/// `/metrics` reports a completed stream. The lines are matched whole, so
/// `mlxcel_requests_total 11` cannot satisfy the first expectation.
#[tokio::test]
async fn metrics_endpoint_reports_a_completed_stream() {
let config = ServerConfig {
enable_metrics_endpoint: true,
..Default::default()
};
let (app, handle, _options_rx, _state) = scripted_app_with_state(config);

complete_plain_stream(&app, &handle).await;

let response = app
.clone()
.oneshot(get_request("/metrics", None))
.await
.expect("scrape");
assert_eq!(response.status(), StatusCode::OK);
let text = String::from_utf8(collect_body(response).await).expect("utf8 body");
let expected = [
"mlxcel_requests_total 1".to_string(),
format!("mlxcel_completion_tokens_total {}", TOKENS.len()),
];
for line in &expected {
assert!(
text.lines().any(|l| l == line.as_str()),
"/metrics must print {line:?}:\n{text}"
);
}
let help = text
.lines()
.find(|l| l.starts_with("# HELP mlxcel_requests_total "))
.expect("requests_total HELP line");
assert!(
help.contains("cancelled by the client are excluded"),
"the HELP line must state the cancelled-stream exclusion: {help}"
);
}

/// The WebUI Activity tiles read `completed_requests_total` and
/// `completion_tokens_total` from the runtime projection; a completed stream
/// moves both.
///
/// `runtime_snapshot` lives behind the `webui` feature (`server::webui` is
/// `#[cfg(feature = "webui")]` in `server/mod.rs`), so this test needs the
/// same gate; without it, `--no-default-features --features xla-diagnostics`
/// fails to compile this file.
#[cfg(feature = "webui")]
#[tokio::test]
async fn activity_projection_counts_a_completed_stream() {
let config = ServerConfig {
enable_metrics_endpoint: true,
..Default::default()
};
let (app, handle, _options_rx, state) = scripted_app_with_state(config);

complete_plain_stream(&app, &handle).await;

let snapshot = crate::server::webui::runtime::runtime_snapshot(
"srv-route-test".into(),
"mdl-route-test".into(),
1,
1,
&state.config,
Some(&state),
);
let value = |key: &str| snapshot.measurements.get(key).and_then(|m| m.value);
assert_eq!(value("completed_requests_total"), Some(1.0));
assert_eq!(value("completion_tokens_total"), Some(TOKENS.len() as f64));
}
22 changes: 21 additions & 1 deletion src/server/routes/stream_route_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,21 @@ fn scripted_app(
axum::Router,
ScriptedStreamHandle,
mpsc::Receiver<ServerGenerateOptions>,
) {
let (app, handle, options_rx, _state) = scripted_app_with_state(config);
(app, handle, options_rx)
}

/// [`scripted_app`], also handing back the `AppState` the router serves, so a
/// test can read what the routes publish into it (`state.metrics`, the
/// completion-control registry).
fn scripted_app_with_state(
config: ServerConfig,
) -> (
axum::Router,
ScriptedStreamHandle,
mpsc::Receiver<ServerGenerateOptions>,
AppState,
) {
let (options_tx, options_rx) = mpsc::channel();
let (provider, handle) = ModelProvider::scripted_streaming_for_route_tests(options_tx);
Expand All @@ -51,7 +66,7 @@ fn scripted_app(
PathBuf::from("route-test-model"),
batch_metrics,
);
(create_app(state), handle, options_rx)
(create_app(state.clone()), handle, options_rx, state)
}

fn chat_body(reasoning_control: Option<bool>) -> String {
Expand Down Expand Up @@ -696,3 +711,8 @@ async fn reasoning_end_on_an_unarmed_completion_reports_not_enabled() {
);
handle.finish();
}

/// What a streamed chat completion records in `Metrics` (#1911), driven
/// through the same router and helpers as the lifecycle tests above.
#[path = "stream_route_metrics_tests.rs"]
mod metrics_tests;
4 changes: 2 additions & 2 deletions src/server/webui/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,10 +187,10 @@ fn counter_source(key: &str) -> &'static str {
"BatchMetrics pending prefill / single-stream queue; point-in-time sample"
}
"completed_requests_total" => {
"Metrics successfully recorded route completions; loaded-provider lifetime, not all admissions or errors"
"Metrics successfully recorded route completions, streamed and non-streamed chat alike, excluding chat streams the client cancelled; loaded-provider lifetime, not all admissions or errors"
}
"completion_tokens_total" => {
"Metrics completion tokens from recorded route results; loaded-provider lifetime"
"Metrics completion tokens from recorded route results, streamed and non-streamed chat alike, excluding chat streams the client cancelled; loaded-provider lifetime"
}
"generation_time_ms_total" => {
"Metrics sum of route-reported generation/elapsed times; loaded-provider lifetime, not decode-only or TTFT"
Expand Down
4 changes: 2 additions & 2 deletions tests/fixtures/webui/scenarios/runtime.ready-shared.json
Original file line number Diff line number Diff line change
Expand Up @@ -94,14 +94,14 @@
"unit": "requests",
"scope": "model",
"measured_at": "2026-09-12T03:04:05Z",
"reason": "Metrics successfully recorded route completions; loaded-provider lifetime, not all admissions or errors"
"reason": "Metrics successfully recorded route completions, streamed and non-streamed chat alike, excluding chat streams the client cancelled; loaded-provider lifetime, not all admissions or errors"
},
"completion_tokens_total": {
"value": 0.0,
"unit": "tokens",
"scope": "model",
"measured_at": "2026-09-12T03:04:05Z",
"reason": "Metrics completion tokens from recorded route results; loaded-provider lifetime"
"reason": "Metrics completion tokens from recorded route results, streamed and non-streamed chat alike, excluding chat streams the client cancelled; loaded-provider lifetime"
},
"generation_time_ms_total": {
"value": 0.0,
Expand Down
Loading