aprender-serve 0.64.0

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

//! The handler half of aprender#2376(3): turn an axum client-disconnect into a
//! cancellation signal the decode loops can actually observe.
//!
//! # Why "axum drops the future" was not already enough
//!
//! Every generate handler in this crate calls its backend *synchronously* from
//! inside its `async fn`:
//!
//! ```ignore
//! pub async fn generate_handler(...) -> ... {
//!     let generated = model.generate_with_cache(&prompt_ids, &q_config)?; // no .await
//!     ...
//! }
//! ```
//!
//! Axum cancels an abandoned request by **dropping the response future**, and a
//! future can only be dropped while it is suspended at an `.await`. There is no
//! `.await` anywhere inside a synchronous decode loop, so the task never yielded,
//! the drop never landed, and generation ran all the way to `max_tokens` for a
//! client that had already hung up. That is the ~250%-CPU-with-zero-open-
//! connections symptom in aprender#2376(3).
//!
//! # The shape that works
//!
//! [`cancel_on_disconnect`] is a `tower` layer over the whole router. Per request
//! it:
//!
//! 1. mints a [`CancelToken`] and puts a clone in the request extensions, so any
//!    handler can pull it out with `Extension<CancelToken>` and install it on its
//!    generation config;
//! 2. holds a [`CancelOnDrop`](crate::generate::CancelOnDrop) guard for the whole
//!    middleware future — this is the piece axum drops on disconnect; and
//! 3. runs the inner handler in a **separate task** ([`tokio::spawn`]) and awaits
//!    its `JoinHandle`.
//!
//! Step 3 is not decoration. Awaiting is what gives this future a suspension point
//! to be dropped at, and running the handler in its own task is what lets the
//! decode loop still be alive — and therefore still able to observe the flag —
//! *after* the drop. Dropping a `JoinHandle` detaches its task rather than killing
//! it, and nothing can preempt a synchronous loop anyway, which is exactly why the
//! loop has to cooperate by polling.
//!
//! Take any one of the three away and the defect returns:
//!
//! | Missing | Result |
//! |---------|--------|
//! | 1 | The loop polls a token nobody shares — always false. |
//! | 2 | Nothing ever sets the flag. |
//! | 3 | No await point; the drop never happens mid-generation. |
//!
//! # Why the guard is disarmed on completion (aprender#2375(1))
//!
//! Step 2's guard originally fired on BOTH exits — abandonment and normal
//! completion — because firing late was assumed harmless. It is harmless only
//! for a handler that finishes its generation before returning. The streaming
//! chat backends do the opposite: they hand the decode loop to
//! `spawn_blocking` and return the SSE response at once, so this future
//! completes while the loop is still in prefill. The guard then cancelled it at
//! the very first poll, and `POST /v1/chat/completions` with `"stream":true`
//! returned a well-formed event stream containing the opening chunk, the
//! terminal chunk, and **zero content deltas** — every streamed reply empty.
//!
//! So a completed handler disarms the guard. Abandonment (drop) and a panicking
//! handler still cancel. An abandoned *stream* is still stopped, by the
//! mechanism that has always covered it: hyper drops the response body → the
//! SSE receiver drops → the generator's `on_token` send fails → the loop breaks.
//!
//! # Panics are preserved
//!
//! A panicking handler is re-raised with [`std::panic::resume_unwind`] so hyper
//! still sees a panic rather than this layer converting it into a 500. Interposing
//! a task must not change what a completed — or a failing — request returns.
//!
//! # Contract
//!
//! `contracts/apr-serve-cancellation-v1.yaml`.

use axum::{
    body::Body,
    http::{Request, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
    Json,
};

use crate::generate::CancelToken;

use super::ErrorResponse;

/// Per-request cancellation layer. See the module docs for why all three steps
/// are load-bearing.
pub(crate) async fn cancel_on_disconnect(mut request: Request<Body>, next: Next) -> Response {
    let token = CancelToken::new();
    request.extensions_mut().insert(token.clone());

    // (2) Lives in THIS future. Dropped when axum abandons the request because
    // the client went away — and DISARMED first when the handler returns, since
    // a returned response may still be streaming from a background decode loop.
    let mut disconnect_guard = token.cancel_on_drop();

    // (3) The handler outlives this future's drop, so its decode loop is still
    // running to see the flag the guard just set.
    let handle = tokio::spawn(async move { next.run(request).await });

    match handle.await {
        Ok(response) => {
            // The handler produced a response, so this request was not
            // abandoned. Cancelling here would stop a streaming body that has
            // not been written yet — see the module docs (#2375(1)).
            disconnect_guard.disarm();
            response
        },
        Err(join_err) if join_err.is_panic() => std::panic::resume_unwind(join_err.into_panic()),
        Err(join_err) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse {
                error: format!("Request task did not complete: {join_err}"),
            }),
        )
            .into_response(),
    }
}

/// The token a handler should install on its generation config.
///
/// Handlers take `Extension<CancelToken>`; this exists for the paths that hold a
/// `Request` rather than running as an extractor-based handler, and for tests that
/// invoke a handler without going through [`cancel_on_disconnect`].
#[must_use]
pub fn request_cancel_token(request: &Request<Body>) -> CancelToken {
    request
        .extensions()
        .get::<CancelToken>()
        .cloned()
        .unwrap_or_else(CancelToken::never)
}

#[cfg(test)]
#[path = "tests/cancel_scope_2376.rs"]
mod cancel_scope_2376;