server.rs 7 KB
Newer Older
1
use crate::{
OlivierDehaene's avatar
OlivierDehaene committed
2
    Batcher, Details, ErrorResponse, GenerateParameters, GenerateRequest, GeneratedText, Validation,
3
};
Olivier Dehaene's avatar
Olivier Dehaene committed
4
use axum::extract::Extension;
5
6
use axum::http::{HeaderMap, StatusCode};
use axum::response::IntoResponse;
Olivier Dehaene's avatar
Olivier Dehaene committed
7
use axum::routing::{get, post};
Olivier Dehaene's avatar
Olivier Dehaene committed
8
9
use axum::{Json, Router};
use std::net::SocketAddr;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
10
use std::sync::Arc;
11
use text_generation_client::ShardedClient;
Olivier Dehaene's avatar
Olivier Dehaene committed
12
use tokenizers::Tokenizer;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
13
14
use tokio::signal;
use tokio::sync::Semaphore;
Olivier Dehaene's avatar
Olivier Dehaene committed
15
16
17
use tokio::time::Instant;
use tracing::instrument;

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
18
19
20
21
22
23
// Server shared state
#[derive(Clone)]
struct ServerState {
    validation: Validation,
    batcher: Batcher,
    limit_concurrent_requests: Arc<Semaphore>,
Olivier Dehaene's avatar
Olivier Dehaene committed
24
25
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
26
/// Health check method
Olivier Dehaene's avatar
Olivier Dehaene committed
27
#[instrument(skip(state), fields(time, time_per_token))]
28
async fn health(state: Extension<ServerState>) -> Result<(), (StatusCode, Json<ErrorResponse>)> {
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
29
30
31
32
33
34
35
36
    // TODO: while this is the best health check we can do, it is a bit on the heavy side and might
    //       be a bit too slow for a health check.
    //       What we should do instead if check if the gRPC channels are still healthy.

    // Limit concurrent requests by acquiring a permit from the semaphore
    let _permit = state.limit_concurrent_requests.try_acquire().map_err(|_| {
        (
            StatusCode::TOO_MANY_REQUESTS,
37
38
39
            Json(ErrorResponse {
                error: "Model is overloaded".to_string(),
            }),
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
40
41
42
43
        )
    })?;

    // Send a small inference request
Olivier Dehaene's avatar
Olivier Dehaene committed
44
    state
Olivier Dehaene's avatar
Olivier Dehaene committed
45
        .batcher
Olivier Dehaene's avatar
Olivier Dehaene committed
46
47
48
49
50
51
52
53
54
55
        .infer(
            1,
            GenerateRequest {
                inputs: "liveness".to_string(),
                parameters: GenerateParameters {
                    temperature: 1.0,
                    top_k: 0,
                    top_p: 1.0,
                    do_sample: false,
                    max_new_tokens: 1,
56
                    stop: vec![],
OlivierDehaene's avatar
OlivierDehaene committed
57
                    details: false,
Olivier Dehaene's avatar
Olivier Dehaene committed
58
59
60
                },
            },
        )
Olivier Dehaene's avatar
Olivier Dehaene committed
61
62
        .await?;
    Ok(())
Olivier Dehaene's avatar
Olivier Dehaene committed
63
64
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
65
/// Generate method
66
67
68
69
70
71
72
73
74
75
#[instrument(
    skip(state),
    fields(
        total_time,
        validation_time,
        queue_time,
        inference_time,
        time_per_token
    )
)]
Olivier Dehaene's avatar
Olivier Dehaene committed
76
async fn generate(
Olivier Dehaene's avatar
Olivier Dehaene committed
77
    state: Extension<ServerState>,
Olivier Dehaene's avatar
Olivier Dehaene committed
78
    req: Json<GenerateRequest>,
79
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
80
    let start_time = Instant::now();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
81
82
    // Limit concurrent requests by acquiring a permit from the semaphore
    let _permit = state.limit_concurrent_requests.try_acquire().map_err(|_| {
83
        tracing::error!("Model is overloaded");
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
84
85
        (
            StatusCode::TOO_MANY_REQUESTS,
86
87
88
            Json(ErrorResponse {
                error: "Model is overloaded".to_string(),
            }),
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
89
90
        )
    })?;
Olivier Dehaene's avatar
Olivier Dehaene committed
91

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
92
    // Validate request
OlivierDehaene's avatar
OlivierDehaene committed
93
    let details = req.0.parameters.details;
94
95
    let (input_length, validated_request) =
        state.validation.validate(req.0).await.map_err(|err| {
96
97
98
            tracing::error!("{}", err.to_string());
            err
        })?;
Olivier Dehaene's avatar
Olivier Dehaene committed
99

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
100
    // Inference
101
102
103
104
105
106
107
108
    let response = state
        .batcher
        .infer(input_length, validated_request)
        .await
        .map_err(|err| {
            tracing::error!("{}", err.to_string());
            err
        })?;
Olivier Dehaene's avatar
Olivier Dehaene committed
109

OlivierDehaene's avatar
OlivierDehaene committed
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
    // Token details
    let details = match details {
        true => {
            let tokens = response
                .token_ids
                .into_iter()
                .zip(response.tokens.into_iter())
                .zip(response.logprobs.into_iter())
                .map(|((id, text), logprob)| (id, text, logprob))
                .collect();
            Some(Details {
                finish_reason: response.finish_reason,
                generated_tokens: response.generated_tokens,
                tokens,
            })
        }
        false => None,
    };

129
130
131
132
133
    // Timings
    let total_time = start_time.elapsed();
    let validation_time = response.queued - start_time;
    let queue_time = response.start - response.queued;
    let inference_time = response.end - response.start;
OlivierDehaene's avatar
OlivierDehaene committed
134
    let time_per_token = inference_time / response.generated_tokens;
135
136
137
138
139
140
141
142
143
144
145
146
147
148

    // Headers
    let mut headers = HeaderMap::new();
    headers.insert(
        "x-total-time",
        total_time.as_millis().to_string().parse().unwrap(),
    );
    headers.insert(
        "x-validation-time",
        validation_time.as_millis().to_string().parse().unwrap(),
    );
    headers.insert(
        "x-queue-time",
        queue_time.as_millis().to_string().parse().unwrap(),
Olivier Dehaene's avatar
Olivier Dehaene committed
149
    );
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    headers.insert(
        "x-inference-time",
        inference_time.as_millis().to_string().parse().unwrap(),
    );
    headers.insert(
        "x-time-per-token",
        time_per_token.as_millis().to_string().parse().unwrap(),
    );

    // Tracing metadata
    tracing::Span::current().record("total_time", format!("{:?}", total_time));
    tracing::Span::current().record("validation_time", format!("{:?}", validation_time));
    tracing::Span::current().record("queue_time", format!("{:?}", queue_time));
    tracing::Span::current().record("inference_time", format!("{:?}", inference_time));
    tracing::Span::current().record("time_per_token", format!("{:?}", time_per_token));
OlivierDehaene's avatar
OlivierDehaene committed
165
    tracing::info!("Output: {}", response.output_text);
Olivier Dehaene's avatar
Olivier Dehaene committed
166

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
167
    // Send response
168
    let response = vec![GeneratedText {
OlivierDehaene's avatar
OlivierDehaene committed
169
170
        generated_text: response.output_text,
        details,
171
172
    }];
    Ok((headers, Json(response)))
Olivier Dehaene's avatar
Olivier Dehaene committed
173
174
}

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
175
176
177
178
179
180
/// Serving method
#[allow(clippy::too_many_arguments)]
pub async fn run(
    max_concurrent_requests: usize,
    max_input_length: usize,
    max_batch_size: usize,
181
    max_waiting_tokens: usize,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
182
183
184
185
186
187
    client: ShardedClient,
    tokenizer: Tokenizer,
    validation_workers: usize,
    addr: SocketAddr,
) {
    // Create state
188
    let batcher = Batcher::new(client, max_batch_size, max_waiting_tokens);
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
189
190
191
192
193
194
195
196
    let validation = Validation::new(validation_workers, tokenizer, max_input_length);
    let shared_state = ServerState {
        validation,
        batcher,
        limit_concurrent_requests: Arc::new(Semaphore::new(max_concurrent_requests)),
    };

    // Create router
Olivier Dehaene's avatar
Olivier Dehaene committed
197
    let app = Router::new()
198
        .route("/", post(generate))
Olivier Dehaene's avatar
Olivier Dehaene committed
199
        .route("/generate", post(generate))
200
        .route("/", get(health))
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
201
        .route("/health", get(health))
Olivier Dehaene's avatar
Olivier Dehaene committed
202
        .layer(Extension(shared_state.clone()));
Olivier Dehaene's avatar
Olivier Dehaene committed
203

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
204
    // Run server
Olivier Dehaene's avatar
Olivier Dehaene committed
205
    axum::Server::bind(&addr)
Olivier Dehaene's avatar
Olivier Dehaene committed
206
        .serve(app.into_make_service())
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
207
208
        // Wait until all requests are finished to shut down
        .with_graceful_shutdown(shutdown_signal())
Olivier Dehaene's avatar
Olivier Dehaene committed
209
210
        .await
        .unwrap();
Olivier Dehaene's avatar
Olivier Dehaene committed
211
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238

/// Shutdown signal handler
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    tracing::info!("signal received, starting graceful shutdown");
}