[sgl-router] refactor - chat_completions() into modules (#39848)

2026-09-20 13:51:25 +08:00
parent dd83b54611
commit 99d53fe0c2
3 changed files with 2039 additions and 2810 deletions
@@ -0,0 +1,337 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Plain and PD chat forwarding, including load tracking and streaming metrics.
+
+use super::preparation::{generate_room_id, BootstrapFields, PreparedChatRequest};
+use crate::discovery::WorkerMode;
+use crate::policies::active_load::ActiveLoadGuard;
+use crate::proxy::sse::StreamEnd;
+use crate::server::app_context::AppContext;
+use crate::server::error::ApiError;
+use crate::server::metrics::{
+    classify_stream_end, outcome_from_status, MetricsRegistry, RequestLogContext, RequestOutcome,
+    StaleRequestOutcome, WorkerModeLabel,
+};
+use crate::workers::{LoadGuard, Worker};
+use axum::body::Body;
+use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
+use axum::response::IntoResponse;
+use bytes::Bytes;
+use std::sync::Arc;
+use std::time::Instant;
+
+const CHAT_PATH: &str = "/v1/chat/completions";
+// Expose the selected decode worker to both PD workers and the client.
+const X_SGL_DECODE_URL: HeaderName = HeaderName::from_static("x-sgl-decode-url");
+type LoadGuards = (LoadGuard, ActiveLoadGuard);
+
+/// A plain worker, or a prefill worker paired with a decode worker for PD.
+pub(super) struct SelectedWorkers {
+    pub(super) prefill: Arc<Worker>,
+    pub(super) decode: Option<Arc<Worker>>,
+    pub(super) track_dispatch_timestamps: bool,
+}
+
+pub(super) async fn forward_chat_request(
+    ctx: &AppContext,
+    request: PreparedChatRequest,
+    workers: SelectedWorkers,
+    mut headers: HeaderMap,
+    request_started_at: Instant,
+) -> Result<Response<Body>, ApiError> {
+    let SelectedWorkers {
+        prefill,
+        decode,
+        track_dispatch_timestamps,
+    } = workers;
+    let decode_url_header = decode
+        .as_ref()
+        .and_then(|worker| parse_decode_url_header(&worker.url));
+    if let Some(hint) = &decode_url_header {
+        headers.insert(X_SGL_DECODE_URL, hint.clone());
+    }
+
+    // Track worker occupancy and the prompt's contribution to active load.
+    let worker_load_guard = if track_dispatch_timestamps {
+        prefill.timestamped_load_guard()
+    } else {
+        prefill.load_guard()
+    };
+    let active_request_guard = ctx.active_load.register(
+        prefill.id.clone(),
+        prefill.url.clone(),
+        request.input_token_count,
+        0,
+    );
+    // PD requests keep using the prefill expiration token after dispatching decode.
+    let expiration_token = active_request_guard.cancel_token().clone();
+    // Attribute the outcome to the worker supplying the client-visible response.
+    let metrics = DispatchMetrics::new(
+        ctx,
+        &request,
+        decode.as_deref().unwrap_or(&prefill),
+        request_started_at,
+    );
+    // Both PD workers receive the same bootstrap room to coordinate KV transfer.
+    let pd = decode.map(|decode| {
+        let bootstrap = BootstrapFields {
+            host: prefill.bootstrap_host().to_string(),
+            port: prefill.bootstrap_port(),
+            room: generate_room_id(),
+        };
+        (decode, bootstrap)
+    });
+    let body = request.into_outgoing_body(ctx, pd.as_ref().map(|(_, bootstrap)| bootstrap))?;
+    let prefill_load_guards = (worker_load_guard, active_request_guard);
+
+    // In PD mode, prefill runs independently and decode supplies the client response.
+    let (response_worker, response_load_guards) = if let Some((decode, bootstrap)) = pd {
+        spawn_prefill_request(
+            ctx,
+            prefill,
+            headers.clone(),
+            body.clone(),
+            prefill_load_guards,
+            bootstrap.room,
+        );
+        let decode_load_guards = (
+            decode.load_guard(),
+            ctx.active_load
+                .register(decode.id.clone(), decode.url.clone(), 0, 1),
+        );
+        (decode, decode_load_guards)
+    } else {
+        (prefill, prefill_load_guards)
+    };
+
+    let response_future = forward_to_response_worker(
+        ctx,
+        &response_worker,
+        &headers,
+        body,
+        response_load_guards,
+        &metrics,
+    );
+    // A ready response wins if request expiration fires in the same poll.
+    let result = tokio::select! {
+        biased;
+        result = response_future => result,
+        _ = expiration_token.cancelled() => Err(ApiError::StaleRequestExpired {
+            model: metrics.model.clone(),
+        }),
+    };
+    let log_context = metrics.record_dispatch_result(&result);
+    // Materialize dispatch errors here so the access log retains the selected worker.
+    let mut response = match result {
+        Ok(mut response) => {
+            if let Some(hint) = decode_url_header {
+                response.headers_mut().insert(X_SGL_DECODE_URL, hint);
+            }
+            response
+        }
+        Err(error) => error.into_response(),
+    };
+    response.extensions_mut().insert(log_context);
+    Ok(response)
+}
+
+fn parse_decode_url_header(decode_url: &str) -> Option<HeaderValue> {
+    HeaderValue::from_str(decode_url)
+        .map_err(|error| {
+            tracing::warn!(
+                decode_url = %decode_url,
+                %error,
+                "decode worker URL rejected by header parser; sending request without decode hint",
+            );
+        })
+        .ok()
+}
+
+fn spawn_prefill_request(
+    ctx: &AppContext,
+    prefill_worker: Arc<Worker>,
+    headers: HeaderMap,
+    body: Bytes,
+    load_guards: LoadGuards,
+    bootstrap_room: u64,
+) {
+    let proxy = Arc::clone(&ctx.proxy);
+    // Let prefill finish KV transfer after client cancellation; router shutdown still cancels it.
+    tokio::spawn(async move {
+        let _load_guards = load_guards;
+        match proxy
+            .forward_json_to(
+                &prefill_worker.url,
+                prefill_worker.protocol(),
+                &prefill_worker.breaker,
+                CHAT_PATH,
+                &headers,
+                body,
+            )
+            .await
+        {
+            Ok(_) => tracing::debug!(
+                prefill_url = %prefill_worker.url, bootstrap_room, "prefill side completed",
+            ),
+            // Prefill failures surface to the client through decode's bootstrap timeout.
+            Err(error) => tracing::warn!(
+                prefill_url = %prefill_worker.url,
+                bootstrap_room,
+                %error,
+                "prefill request failed; decode will time out on bootstrap_room",
+            ),
+        }
+    });
+}
+
+async fn forward_to_response_worker(
+    ctx: &AppContext,
+    worker: &Worker,
+    headers: &HeaderMap,
+    body: Bytes,
+    load_guards: LoadGuards,
+    metrics: &DispatchMetrics,
+) -> Result<Response<Body>, ApiError> {
+    if metrics.streaming {
+        // Load and duration guards live until the SSE pump ends, not just until headers arrive.
+        let stream_guards: Box<dyn Send + 'static> =
+            Box::new((load_guards, metrics.stream_duration_guard()));
+        ctx.proxy
+            .forward_streaming_to(
+                &worker.url,
+                worker.protocol(),
+                &worker.breaker,
+                CHAT_PATH,
+                headers,
+                body,
+                Some(stream_guards),
+                Some(metrics.first_byte_callback()),
+                Some(metrics.stream_end_callback(worker.url.clone())),
+            )
+            .await
+    } else {
+        // JSON forwarding reads the full response body before releasing load guards.
+        let _load_guards = load_guards;
+        ctx.proxy
+            .forward_json_to(
+                &worker.url,
+                worker.protocol(),
+                &worker.breaker,
+                CHAT_PATH,
+                headers,
+                body,
+            )
+            .await
+    }
+}
+
+struct DispatchMetrics {
+    registry: Arc<MetricsRegistry>,
+    model: String,
+    worker_url: String,
+    mode: WorkerModeLabel,
+    streaming: bool,
+    request_started_at: Instant,
+}
+
+impl DispatchMetrics {
+    fn new(
+        ctx: &AppContext,
+        request: &PreparedChatRequest,
+        response_worker: &Worker,
+        request_started_at: Instant,
+    ) -> Self {
+        Self {
+            registry: Arc::clone(&ctx.metrics),
+            model: request.model.0.clone(),
+            worker_url: response_worker.url.clone(),
+            mode: match response_worker.mode() {
+                WorkerMode::Prefill => WorkerModeLabel::Prefill,
+                WorkerMode::Decode => WorkerModeLabel::Decode,
+                WorkerMode::Plain => WorkerModeLabel::Plain,
+            },
+            streaming: request.streaming,
+            request_started_at,
+        }
+    }
+
+    // TTFT uses the first successful stream chunk, measured from request arrival.
+    fn first_byte_callback(&self) -> Box<dyn FnOnce() + Send + 'static> {
+        let metrics = Arc::clone(&self.registry);
+        let model = self.model.clone();
+        let request_started_at = self.request_started_at;
+        Box::new(move || metrics.observe_ttft(&model, request_started_at.elapsed().as_secs_f64()))
+    }
+
+    fn stream_duration_guard(&self) -> StreamDurationGuard {
+        StreamDurationGuard {
+            metrics: Arc::clone(&self.registry),
+            model: self.model.clone(),
+            request_started_at: self.request_started_at,
+        }
+    }
+
+    fn stream_end_callback(
+        &self,
+        response_worker_url: String,
+    ) -> Box<dyn FnOnce(StreamEnd) + Send + 'static> {
+        let metrics = Arc::clone(&self.registry);
+        let model = self.model.clone();
+        Box::new(move |end| {
+            metrics.record_stream_outcome(&response_worker_url, &model, classify_stream_end(end));
+        })
+    }
+
+    // HTTP status determines the outcome; router cancellations and dispatch failures stay distinct.
+    fn record_dispatch_result(
+        &self,
+        result: &Result<Response<Body>, ApiError>,
+    ) -> RequestLogContext {
+        let http_status = match result {
+            Ok(response) => response.status().as_u16(),
+            Err(error) => error.status_code().as_u16(),
+        };
+        let outcome = match result {
+            Err(ApiError::StaleRequestExpired { .. }) => {
+                self.registry
+                    .record_stale_request(StaleRequestOutcome::Expired);
+                RequestOutcome::Cancelled
+            }
+            // These 503s come from the router, not worker backpressure.
+            Err(ApiError::BreakerOpen { .. } | ApiError::WorkerMisconfigured { .. }) => {
+                RequestOutcome::Error
+            }
+            _ => outcome_from_status(http_status),
+        };
+        self.registry
+            .record_worker_request(&self.worker_url, &self.model, self.mode, outcome);
+        if !self.streaming {
+            self.registry.observe_request_duration(
+                &self.model,
+                self.request_started_at.elapsed().as_secs_f64(),
+            );
+        }
+        // The app middleware emits the access log and edge counters exactly once.
+        RequestLogContext {
+            worker_url: self.worker_url.clone(),
+            model_id: self.model.clone(),
+            streaming: self.streaming,
+            outcome,
+        }
+    }
+}
+
+/// Record total request duration when streaming ends or setup fails.
+struct StreamDurationGuard {
+    metrics: Arc<MetricsRegistry>,
+    model: String,
+    request_started_at: Instant,
+}
+
+impl Drop for StreamDurationGuard {
+    fn drop(&mut self) {
+        self.metrics
+            .observe_request_duration(&self.model, self.request_started_at.elapsed().as_secs_f64());
+    }
+}