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
141
142
143
144
145
146
147
148
149
150
151
152
//! The six handles a decode runs against, taken once and carried
//! together.
//!
//! Every generating route has to do the same thing before it can call
//! [`crate::run_generation`]: pin the model handle so a mid-flight
//! `/admin/models/load` cannot splice two checkpoints into one answer,
//! then clone the KV pool, the paged pool, the prefix cache, the
//! continuous batcher and the context ceiling off `AppState` and the
//! active model, then hand all six to an eight-argument function in the
//! right order.
//!
//! That block was written out four times -- `/v1/chat/completions`,
//! `/v1/completions`, `/v1/messages` and `/v1/responses` -- and the
//! *argument order* is the part that made it dangerous: `kv_pool` and
//! `paged_kv` are both `Option<&_>`, so transposing them is not a type
//! error. This repo's recorded lesson is that a copied path diverges
//! and nothing notices, so the block is one struct here and the routes
//! name their protocol instead.
//!
//! Two methods, matching the two shapes a route needs:
//!
//! - [`DecodeHandles::run`], the buffered one, wrapped by
//! [`buffered`] which also does the `spawn_blocking` handoff and the
//! two error mappings.
//! - [`DecodeHandles::run_emit`], which streams chunks to a sink. The
//! sink is where the protocols genuinely differ (reasoning splits,
//! tool-call parsing, plain text), so that part stays with each
//! route; the handles do not.
use std::sync::{Arc, Mutex};
use frink_models::PrefixCache;
use crate::generate::GenerationParams;
use crate::{
budget, decode_error_response, generate, join_error_response, serving, ActiveModel, ApiError,
AppState, Model,
};
/// Everything a decode needs from the server, pinned to one model.
pub(crate) struct DecodeHandles {
model: Arc<Model>,
kv_pool: Option<generate::KvPoolConfig>,
paged_kv: Option<generate::PagedKvConfig>,
prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
batcher: Option<serving::batch::ContinuousBatcher>,
ceiling: Option<Arc<budget::ContextCeiling>>,
metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
}
impl DecodeHandles {
/// Take the handles for one request.
///
/// `active` is the caller's own `require_active()` result rather
/// than a fresh read of `AppState`, so the model, batcher and
/// ceiling all come from the same swap generation -- a ceiling
/// derived for a checkpoint that is no longer loaded prices the
/// wrong KV geometry.
///
/// Fails when `active` is an encoder-only checkpoint: there is no
/// decode to take handles for. Every generating route passes
/// through here, so this one `?` is what keeps an encoder off all
/// five of them.
pub(crate) fn take(state: &AppState, active: &ActiveModel) -> Result<Self, ApiError> {
Ok(DecodeHandles {
model: Arc::clone(active.generative()?),
kv_pool: state.kv_pool.clone(),
paged_kv: state.paged_kv.clone(),
prefix_cache: state.prefix_cache.clone(),
batcher: active.batcher.clone(),
ceiling: active.ceiling.clone(),
metal_private_decode_gate: state.metal_private_decode_gate.clone(),
})
}
/// The model this decode is pinned to. Read for its name or its
/// tokenizer; the handle itself never leaves.
pub(crate) fn model(&self) -> &Model {
&self.model
}
/// Whether a prefix cache is configured, which is the only
/// circumstance under which prompt KV is reused across requests.
///
/// Read by `/completion`, whose `cache_prompt: false` is a
/// *requirement* not to reuse rather than a permission to.
pub(crate) fn has_prefix_cache(&self) -> bool {
self.prefix_cache.is_some()
}
/// Generate, collecting the whole answer. Blocking: call it from
/// `spawn_blocking` (or through [`buffered`], which does that).
pub(crate) fn run(
&self,
prompt: &str,
params: &GenerationParams,
// Per choice, choice 0 first. See `crate::run_generation_emit`.
) -> Result<generate::Generated, generate::DecodeError> {
crate::run_generation(
&self.model,
prompt,
params,
self.kv_pool.as_ref(),
self.paged_kv.as_ref(),
self.prefix_cache.as_deref(),
self.batcher.as_ref(),
self.ceiling.as_deref(),
self.metal_private_decode_gate.as_deref(),
)
}
/// Generate, handing every decoded chunk to `emit` as it arrives.
/// Blocking, for the same reason.
pub(crate) fn run_emit(
&self,
prompt: &str,
params: &GenerationParams,
emit: impl FnMut(usize, &str),
// Per choice, choice 0 first. See `crate::run_generation_emit`.
) -> Result<generate::Generated, generate::DecodeError> {
crate::run_generation_emit(
&self.model,
prompt,
params,
self.kv_pool.as_ref(),
self.paged_kv.as_ref(),
self.prefix_cache.as_deref(),
self.batcher.as_ref(),
self.ceiling.as_deref(),
self.metal_private_decode_gate.as_deref(),
emit,
)
}
}
/// One buffered generation, off the request thread.
///
/// Generation is CPU-bound and would otherwise block a Tokio worker for
/// the length of a completion, so it runs on `spawn_blocking`; a panic
/// in it becomes a 500 through `join_error_response` rather than a
/// hung request, and a decode error becomes the status
/// `decode_error_response` names.
pub(crate) async fn buffered(
handles: DecodeHandles,
prompt: String,
params: GenerationParams,
) -> Result<generate::Generated, ApiError> {
tokio::task::spawn_blocking(move || handles.run(&prompt, ¶ms))
.await
.map_err(join_error_response)?
.map_err(decode_error_response)
}