server.rs 71.8 KB
Newer Older
1
/// HTTP Server logic
OlivierDehaene's avatar
OlivierDehaene committed
2
3
4
5
6
use crate::config::Config;
use crate::infer::v2::SchedulerV2;
use crate::infer::v3::SchedulerV3;
use crate::infer::{HealthCheck, Scheduler};
use crate::infer::{Infer, InferError, InferResponse, InferStreamResponse, ToolGrammar};
7
8
9
10
11
#[cfg(feature = "kserve")]
use crate::kserve::{
    kerve_server_metadata, kserve_health_live, kserve_health_ready, kserve_model_infer,
    kserve_model_metadata, kserve_model_metadata_ready,
};
12
use crate::validation::ValidationError;
13
use crate::{
14
    BestOfSequence, Details, ErrorResponse, FinishReason, GenerateParameters, GenerateRequest,
15
16
17
    GenerateResponse, GrammarType, HubModelInfo, HubProcessorConfig, HubTokenizerConfig, Info,
    Message, PrefillToken, SimpleToken, StreamDetails, StreamResponse, Token, TokenizeResponse,
    Usage, Validation,
18
19
20
21
};
use crate::{
    ChatCompletion, ChatCompletionChoice, ChatCompletionChunk, ChatCompletionComplete,
    ChatCompletionDelta, ChatCompletionLogprob, ChatCompletionLogprobs, ChatCompletionTopLogprob,
22
23
    ChatRequest, Chunk, CompatGenerateRequest, Completion, CompletionComplete, CompletionFinal,
    CompletionRequest, CompletionType, DeltaToolCall, Function, Prompt, Tool, VertexRequest,
24
    VertexResponse,
25
};
26
use crate::{FunctionDefinition, HubPreprocessorConfig, ToolCall, ToolType};
27
use async_stream::__private::AsyncStream;
Olivier Dehaene's avatar
Olivier Dehaene committed
28
use axum::extract::Extension;
29
use axum::http::{HeaderMap, Method, StatusCode};
30
use axum::response::sse::{Event, KeepAlive, Sse};
31
use axum::response::{IntoResponse, Response};
Olivier Dehaene's avatar
Olivier Dehaene committed
32
use axum::routing::{get, post};
33
use axum::{http, Json, Router};
Nicolas Patry's avatar
Nicolas Patry committed
34
use axum_tracing_opentelemetry::middleware::OtelAxumLayer;
35
use futures::stream::StreamExt;
36
use futures::stream::{FuturesOrdered, FuturesUnordered};
37
use futures::Stream;
drbh's avatar
drbh committed
38
use futures::TryStreamExt;
39
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
drbh's avatar
drbh committed
40
use serde_json::Value;
41
use std::convert::Infallible;
Olivier Dehaene's avatar
Olivier Dehaene committed
42
use std::net::SocketAddr;
43
44
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
OlivierDehaene's avatar
OlivierDehaene committed
45
46
use text_generation_client::{v2, v3, ClientError, ShardInfo};
use thiserror::Error;
Olivier Dehaene's avatar
Olivier Dehaene committed
47
use tokenizers::Tokenizer;
48
use tokio::select;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
49
use tokio::signal;
50
use tokio::sync::oneshot;
Olivier Dehaene's avatar
Olivier Dehaene committed
51
use tokio::time::Instant;
52
use tower_http::cors::{AllowOrigin, CorsLayer};
53
use tracing::{info_span, instrument, Instrument};
54
55
use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi;
Olivier Dehaene's avatar
Olivier Dehaene committed
56

57
58
/// Generate tokens if `stream == false` or a stream of token if `stream == true`
#[utoipa::path(
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
post,
tag = "Text Generation Inference",
path = "/",
request_body = CompatGenerateRequest,
responses(
(status = 200, description = "Generated Text",
content(
("application/json" = GenerateResponse),
("text/event-stream" = StreamResponse),
)),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"})),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"})),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"})),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"})),
)
78
)]
79
#[instrument(skip(infer, req))]
80
async fn compat_generate(
81
    Extension(default_return_full_text): Extension<bool>,
82
    infer: Extension<Infer>,
83
    compute_type: Extension<ComputeType>,
84
    Json(mut req): Json<CompatGenerateRequest>,
85
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
86
87
    // default return_full_text given the pipeline_tag
    if req.parameters.return_full_text.is_none() {
88
        req.parameters.return_full_text = Some(default_return_full_text)
89
90
    }

91
92
    // switch on stream
    if req.stream {
93
        Ok(generate_stream(infer, compute_type, Json(req.into()))
94
95
96
            .await
            .into_response())
    } else {
97
        let (headers, Json(generation)) = generate(infer, compute_type, Json(req.into())).await?;
98
        // wrap generation inside a Vec to match api-inference
99
        Ok((headers, Json(vec![generation])).into_response())
100
101
102
    }
}

103
104
/// Text Generation Inference endpoint info
#[utoipa::path(
105
106
107
108
get,
tag = "Text Generation Inference",
path = "/info",
responses((status = 200, description = "Served model info", body = Info))
109
110
)]
#[instrument]
111
112
async fn get_model_info(info: Extension<Info>) -> Json<Info> {
    Json(info.0)
113
114
}

115
#[utoipa::path(
116
117
118
119
120
121
122
123
get,
tag = "Text Generation Inference",
path = "/health",
responses(
(status = 200, description = "Everything is working fine"),
(status = 503, description = "Text generation inference is down", body = ErrorResponse,
example = json ! ({"error": "unhealthy", "error_type": "healthcheck"})),
)
124
125
)]
#[instrument(skip(health))]
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
126
/// Health check method
OlivierDehaene's avatar
OlivierDehaene committed
127
128
129
async fn health(
    mut health: Extension<HealthCheck>,
) -> Result<(), (StatusCode, Json<ErrorResponse>)> {
130
131
132
133
134
135
136
137
138
139
    match health.check().await {
        true => Ok(()),
        false => Err((
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ErrorResponse {
                error: "unhealthy".to_string(),
                error_type: "healthcheck".to_string(),
            }),
        )),
    }
Olivier Dehaene's avatar
Olivier Dehaene committed
140
141
}

142
143
/// Generate tokens
#[utoipa::path(
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
post,
tag = "Text Generation Inference",
path = "/generate",
request_body = GenerateRequest,
responses(
(status = 200, description = "Generated Text", body = GenerateResponse),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"})),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"})),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"})),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"})),
)
159
)]
160
#[instrument(
161
162
skip_all,
fields(
163
parameters = ? req.parameters,
164
165
166
167
168
169
170
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
171
)]
Olivier Dehaene's avatar
Olivier Dehaene committed
172
async fn generate(
173
    infer: Extension<Infer>,
174
    Extension(ComputeType(compute_type)): Extension<ComputeType>,
175
    Json(req): Json<GenerateRequest>,
176
) -> Result<(HeaderMap, Json<GenerateResponse>), (StatusCode, Json<ErrorResponse>)> {
177
    let span = tracing::Span::current();
178
179
180
    generate_internal(infer, ComputeType(compute_type), Json(req), span).await
}

181
pub(crate) async fn generate_internal(
182
183
184
185
186
    infer: Extension<Infer>,
    ComputeType(compute_type): ComputeType,
    Json(req): Json<GenerateRequest>,
    span: tracing::Span,
) -> Result<(HeaderMap, Json<GenerateResponse>), (StatusCode, Json<ErrorResponse>)> {
187
    let start_time = Instant::now();
188
    metrics::counter!("tgi_request_count").increment(1);
189

190
191
    // Do not long ultra long inputs, like image payloads.
    tracing::debug!("Input: {}", &req.inputs[..1000.min(req.inputs.len())]);
192

193
    let compute_characters = req.inputs.chars().count();
194
    let mut add_prompt = None;
195
196
    if req.parameters.return_full_text.unwrap_or(false) {
        add_prompt = Some(req.inputs.clone());
197
198
    }

Nicolas Patry's avatar
Nicolas Patry committed
199
    let details: bool = req.parameters.details || req.parameters.decoder_input_details;
200
201

    // Inference
202
    let (response, best_of_responses) = match req.parameters.best_of {
203
        Some(best_of) if best_of > 1 => {
204
            let (response, best_of_responses) = infer.generate_best_of(req, best_of).await?;
205
206
            (response, Some(best_of_responses))
        }
207
        _ => (infer.generate(req).await?, None),
208
    };
Olivier Dehaene's avatar
Olivier Dehaene committed
209

OlivierDehaene's avatar
OlivierDehaene committed
210
    // Token details
211
    let input_length = response._input_length;
OlivierDehaene's avatar
OlivierDehaene committed
212
    let details = match details {
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        true => {
            // convert best_of_responses
            let best_of_sequences = best_of_responses.map(|responses: Vec<InferResponse>| {
                responses
                    .into_iter()
                    .map(|response: InferResponse| {
                        // Add prompt if return_full_text
                        let mut output_text = response.generated_text.text;
                        if let Some(prompt) = &add_prompt {
                            output_text = prompt.clone() + &output_text;
                        }

                        BestOfSequence {
                            generated_text: output_text,
OlivierDehaene's avatar
OlivierDehaene committed
227
                            finish_reason: response.generated_text.finish_reason,
228
229
230
                            generated_tokens: response.generated_text.generated_tokens,
                            prefill: response.prefill,
                            tokens: response.tokens,
Nicolas Patry's avatar
Nicolas Patry committed
231
                            top_tokens: response.top_tokens,
232
233
234
235
236
237
238
                            seed: response.generated_text.seed,
                        }
                    })
                    .collect()
            });

            Some(Details {
OlivierDehaene's avatar
OlivierDehaene committed
239
                finish_reason: response.generated_text.finish_reason,
240
241
242
243
244
                generated_tokens: response.generated_text.generated_tokens,
                prefill: response.prefill,
                tokens: response.tokens,
                seed: response.generated_text.seed,
                best_of_sequences,
Nicolas Patry's avatar
Nicolas Patry committed
245
                top_tokens: response.top_tokens,
246
247
            })
        }
OlivierDehaene's avatar
OlivierDehaene committed
248
249
250
        false => None,
    };

251
252
253
254
    // Timings
    let total_time = start_time.elapsed();
    let validation_time = response.queued - start_time;
    let queue_time = response.start - response.queued;
255
256
    let inference_time = Instant::now() - response.start;
    let time_per_token = inference_time / response.generated_text.generated_tokens;
257

258
259
260
261
262
263
264
265
    // Tracing metadata
    span.record("total_time", format!("{total_time:?}"));
    span.record("validation_time", format!("{validation_time:?}"));
    span.record("queue_time", format!("{queue_time:?}"));
    span.record("inference_time", format!("{inference_time:?}"));
    span.record("time_per_token", format!("{time_per_token:?}"));
    span.record("seed", format!("{:?}", response.generated_text.seed));

266
267
    // Headers
    let mut headers = HeaderMap::new();
268
    headers.insert("x-compute-type", compute_type.parse().unwrap());
269
270
    headers.insert(
        "x-compute-time",
Nicolas Patry's avatar
Nicolas Patry committed
271
        total_time.as_secs_f64().to_string().parse().unwrap(),
272
273
274
275
276
    );
    headers.insert(
        "x-compute-characters",
        compute_characters.to_string().parse().unwrap(),
    );
277
278
279
280
281
282
283
284
285
286
287
    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
288
    );
289
290
291
292
293
294
295
296
    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(),
    );
297
298
299
300
301
    headers.insert("x-prompt-tokens", input_length.into());
    headers.insert(
        "x-generated-tokens",
        response.generated_text.generated_tokens.into(),
    );
302

303
    // Metrics
304
305
306
307
308
309
310
311
312
    metrics::counter!("tgi_request_success").increment(1);
    metrics::histogram!("tgi_request_duration").record(total_time.as_secs_f64());
    metrics::histogram!("tgi_request_validation_duration").record(validation_time.as_secs_f64());
    metrics::histogram!("tgi_request_queue_duration").record(queue_time.as_secs_f64());
    metrics::histogram!("tgi_request_inference_duration").record(inference_time.as_secs_f64());
    metrics::histogram!("tgi_request_mean_time_per_token_duration")
        .record(time_per_token.as_secs_f64());
    metrics::histogram!("tgi_request_generated_tokens")
        .record(response.generated_text.generated_tokens as f64);
313

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
314
    // Send response
315
316
317
318
319
    let mut output_text = response.generated_text.text;
    if let Some(prompt) = add_prompt {
        output_text = prompt + &output_text;
    }

320
321
    tracing::debug!("Output: {}", output_text);
    tracing::info!("Success");
322

323
    let response = GenerateResponse {
324
        generated_text: output_text,
OlivierDehaene's avatar
OlivierDehaene committed
325
        details,
326
    };
327
    Ok((headers, Json(response)))
Olivier Dehaene's avatar
Olivier Dehaene committed
328
329
}

Yannic Kilcher's avatar
Yannic Kilcher committed
330
/// Generate a stream of token using Server-Sent Events
331
#[utoipa::path(
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
post,
tag = "Text Generation Inference",
path = "/generate_stream",
request_body = GenerateRequest,
responses(
(status = 200, description = "Generated Text", body = StreamResponse,
content_type = "text/event-stream"),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"}),
content_type = "text/event-stream"),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"}),
content_type = "text/event-stream"),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"}),
content_type = "text/event-stream"),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"}),
content_type = "text/event-stream"),
)
352
)]
353
#[instrument(
354
355
skip_all,
fields(
356
parameters = ? req.parameters,
357
358
359
360
361
362
363
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
364
365
)]
async fn generate_stream(
366
    Extension(infer): Extension<Infer>,
367
    Extension(compute_type): Extension<ComputeType>,
368
    Json(req): Json<GenerateRequest>,
369
370
371
372
) -> (
    HeaderMap,
    Sse<impl Stream<Item = Result<Event, Infallible>>>,
) {
373
    let span = tracing::Span::current();
374
375
376
377
378
    let on_message_callback = |stream_token: StreamResponse| {
        let event = Event::default();
        event.json_data(stream_token).unwrap()
    };
    let (headers, response_stream) =
379
        generate_stream_internal(infer, compute_type, Json(req), on_message_callback, span).await;
380
381
382
383
384
385
    let sse = Sse::new(response_stream).keep_alive(KeepAlive::default());
    (headers, sse)
}

async fn generate_stream_internal(
    infer: Infer,
386
    ComputeType(compute_type): ComputeType,
387
388
    Json(req): Json<GenerateRequest>,
    on_message_callback: impl Fn(StreamResponse) -> Event,
389
    span: tracing::Span,
390
) -> (HeaderMap, impl Stream<Item = Result<Event, Infallible>>) {
391
    let start_time = Instant::now();
392
    metrics::counter!("tgi_request_count").increment(1);
393

394
    tracing::debug!("Input: {}", req.inputs);
395

396
    let compute_characters = req.inputs.chars().count();
397
398

    let mut headers = HeaderMap::new();
399
    headers.insert("x-compute-type", compute_type.parse().unwrap());
400
401
402
403
    headers.insert(
        "x-compute-characters",
        compute_characters.to_string().parse().unwrap(),
    );
404
    headers.insert("X-Accel-Buffering", "no".parse().unwrap());
405

406
407
408
409
    let stream = async_stream::stream! {
        // Inference
        let mut end_reached = false;
        let mut error = false;
410
411

        let mut add_prompt = None;
412
413
        if req.parameters.return_full_text.unwrap_or(false) {
            add_prompt = Some(req.inputs.clone());
414
        }
415
        let details = req.parameters.details;
416

417
        let best_of = req.parameters.best_of.unwrap_or(1);
418
419
        if best_of != 1 {
            let err = InferError::from(ValidationError::BestOfStream);
420
            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
421
422
            tracing::error!("{err}");
            yield Ok(Event::from(err));
423
        } else if req.parameters.decoder_input_details {
424
            let err = InferError::from(ValidationError::PrefillDetailsStream);
425
            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
426
427
428
            tracing::error!("{err}");
            yield Ok(Event::from(err));
        } else {
429
            match infer.generate_stream(req).instrument(info_span!(parent: &span, "async_stream")).await {
430
                // Keep permit as long as generate_stream lives
431
                Ok((_permit, _input_length, mut response_stream)) => {
432
                    let mut index = 0;
433
434
                    // Server-Sent Event stream
                    while let Some(response) = response_stream.next().await {
435
                        index += 1;
436
437
438
439
440
441
                        match response {
                            Ok(response) => {
                                match response {
                                    // Prefill is ignored
                                    InferStreamResponse::Prefill(_) => {}
                                    // Yield event for every new token
Nicolas Patry's avatar
Nicolas Patry committed
442
443
444
445
                                    InferStreamResponse::Intermediate{
                                        token,
                                        top_tokens,
                                    } => {
446
447
                                        tracing::debug!(parent: &span, "Token: {:?}", token);

448
449
                                        // StreamResponse
                                        let stream_token = StreamResponse {
450
                                            index,
451
                                            token,
Nicolas Patry's avatar
Nicolas Patry committed
452
                                            top_tokens,
453
454
455
                                            generated_text: None,
                                            details: None,
                                        };
456
457
                                        let event = on_message_callback(stream_token);
                                        yield Ok(event);
458
                                    }
459
460
                                    // Yield event for last token and compute timings
                                    InferStreamResponse::End {
461
                                        token,
462
463
464
                                        generated_text,
                                        start,
                                        queued,
Nicolas Patry's avatar
Nicolas Patry committed
465
                                        top_tokens,
466
467
468
469
                                    } => {
                                        // Token details
                                        let details = match details {
                                            true => Some(StreamDetails {
OlivierDehaene's avatar
OlivierDehaene committed
470
                                                finish_reason: generated_text.finish_reason,
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
                                                generated_tokens: generated_text.generated_tokens,
                                                seed: generated_text.seed,
                                            }),
                                            false => None,
                                        };

                                        // Timings
                                        let total_time = start_time.elapsed();
                                        let validation_time = queued - start_time;
                                        let queue_time = start - queued;
                                        let inference_time = Instant::now() - start;
                                        let time_per_token = inference_time / generated_text.generated_tokens;

                                        // Tracing metadata
                                        span.record("total_time", format!("{total_time:?}"));
                                        span.record("validation_time", format!("{validation_time:?}"));
                                        span.record("queue_time", format!("{queue_time:?}"));
                                        span.record("inference_time", format!("{inference_time:?}"));
                                        span.record("time_per_token", format!("{time_per_token:?}"));
                                        span.record("seed", format!("{:?}", generated_text.seed));

                                        // Metrics
493
494
495
496
497
498
499
                                        metrics::counter!("tgi_request_success").increment(1);
                                        metrics::histogram!("tgi_request_duration").record(total_time.as_secs_f64());
                                        metrics::histogram!("tgi_request_validation_duration").record(validation_time.as_secs_f64());
                                        metrics::histogram!("tgi_request_queue_duration").record(queue_time.as_secs_f64());
                                        metrics::histogram!("tgi_request_inference_duration").record(inference_time.as_secs_f64());
                                        metrics::histogram!("tgi_request_mean_time_per_token_duration").record(time_per_token.as_secs_f64());
                                        metrics::histogram!("tgi_request_generated_tokens").record(generated_text.generated_tokens as f64);
500
501
502
503
504
505
506
507
508

                                        // StreamResponse
                                        end_reached = true;

                                        let mut output_text = generated_text.text;
                                        if let Some(prompt) = add_prompt {
                                            output_text = prompt + &output_text;
                                        }

509
510
                                        tracing::debug!(parent: &span, "Output: {}", output_text);
                                        tracing::info!(parent: &span, "Success");
511

512
                                        let stream_token = StreamResponse {
513
                                            index,
514
                                            token,
Nicolas Patry's avatar
Nicolas Patry committed
515
                                            top_tokens,
516
517
518
519
                                            generated_text: Some(output_text),
                                            details
                                        };

520
521
522

                                        let event = on_message_callback(stream_token);
                                        yield Ok(event);
523
524
                                        break;
                                    }
525
526
                                }
                            }
527
528
529
530
531
532
                            // yield error
                            Err(err) => {
                                error = true;
                                yield Ok(Event::from(err));
                                break;
                            }
533
534
                        }
                    }
535
536
537
538
539
                },
                // yield error
                Err(err) => {
                    error = true;
                    yield Ok(Event::from(err));
540
                }
541
542
543
544
545
            }
            // Check if generation reached the end
            // Skip if we already sent an error
            if !end_reached && !error {
                let err = InferError::IncompleteGeneration;
546
                metrics::counter!("tgi_request_failure", "err" => "incomplete").increment(1);
547
                tracing::error!("{err}");
548
                yield Ok(Event::from(err));
549
550
551
552
            }
        }
    };

553
554
555
    (headers, stream)
}

556
557
/// Generate tokens
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
post,
tag = "Text Generation Inference",
path = "/v1/completions",
request_body = CompletionRequest,
responses(
(status = 200, description = "Generated Chat Completion",
content(
("application/json" = Completion),
("text/event-stream" = CompletionCompleteChunk),
)),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"})),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"})),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"})),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"})),
)
)]
578
#[instrument(
OlivierDehaene's avatar
OlivierDehaene committed
579
580
581
582
583
584
585
586
587
588
589
skip_all,
fields(
// parameters = ? req.parameters,
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
)]
590
591
592
593
594
595
async fn completions(
    Extension(infer): Extension<Infer>,
    Extension(compute_type): Extension<ComputeType>,
    Extension(info): Extension<Info>,
    Json(req): Json<CompletionRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
596
    let span = tracing::Span::current();
597
    metrics::counter!("tgi_request_count").increment(1);
598

599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
    let CompletionRequest {
        max_tokens,
        seed,
        stop,
        stream,
        temperature,
        ..
    } = req;

    let max_new_tokens = max_tokens.or(Some(100));
    let stop = stop.unwrap_or_default();
    // enable greedy only when temperature is 0
    let (do_sample, temperature) = match temperature {
        Some(temperature) if temperature == 0.0 => (false, None),
        other => (true, other),
    };
615
616
617

    // if suffix is present throw an error
    if req.suffix.is_some() {
618
        metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
619
620
621
622
623
624
625
626
627
628
        return Err((
            StatusCode::UNPROCESSABLE_ENTITY,
            Json(ErrorResponse {
                error: "Suffix is not supported and can be achieved by preprocessing the prompt."
                    .to_string(),
                error_type: "suffix not supported".to_string(),
            }),
        ));
    }

629
    if req.prompt.0.len() > info.max_client_batch_size {
630
        metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
631
632
633
634
635
636
637
638
639
640
641
642
643
644
        return Err((
            StatusCode::UNPROCESSABLE_ENTITY,
            Json(ErrorResponse {
                error: format!(
                    "Number of prompts exceeds the maximum allowed batch size of {}",
                    info.max_client_batch_size
                ),
                error_type: "batch size exceeded".to_string(),
            }),
        ));
    }

    let generate_requests: Vec<GenerateRequest> = req
        .prompt
645
        .0
646
647
648
649
650
        .iter()
        .map(|prompt| GenerateRequest {
            inputs: prompt.to_string(),
            parameters: GenerateParameters {
                best_of: None,
651
                temperature,
652
653
654
655
656
                repetition_penalty: req.repetition_penalty,
                frequency_penalty: req.frequency_penalty,
                top_k: None,
                top_p: req.top_p,
                typical_p: None,
657
                do_sample,
658
659
                max_new_tokens,
                return_full_text: None,
660
                stop: stop.clone(),
661
662
663
664
665
666
667
                truncate: None,
                watermark: false,
                details: true,
                decoder_input_details: !stream,
                seed,
                top_n_tokens: None,
                grammar: None,
drbh's avatar
drbh committed
668
                ..Default::default()
669
670
671
672
673
674
675
            },
        })
        .collect();

    let mut x_compute_type = None;
    let mut x_compute_characters = 0u32;
    let mut x_accel_buffering = None;
676
677

    if stream {
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
        let mut response_streams = FuturesOrdered::new();
        for (index, generate_request) in generate_requests.into_iter().enumerate() {
            let model_id = info.model_id.clone();
            let system_fingerprint =
                format!("{}-{}", info.version, info.docker_label.unwrap_or("native"));
            let infer_clone = infer.clone();
            let compute_type_clone = compute_type.clone();
            let span_clone = span.clone();

            // Create a future for each generate_stream_internal call.
            let generate_future = async move {
                let on_message_callback = move |stream_token: StreamResponse| {
                    let event = Event::default();

                    let current_time = std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap_or_else(|_| std::time::Duration::from_secs(0))
                        .as_secs();

                    event
698
                        .json_data(Completion::Chunk(Chunk {
699
700
701
702
703
704
705
706
707
708
709
710
                            id: "".to_string(),
                            created: current_time,

                            choices: vec![CompletionComplete {
                                finish_reason: "".to_string(),
                                index: index as u32,
                                logprobs: None,
                                text: stream_token.token.text,
                            }],

                            model: model_id.clone(),
                            system_fingerprint: system_fingerprint.clone(),
711
                        }))
712
                        .unwrap_or_else(|_e| Event::default())
713
714
715
716
717
718
719
720
721
722
723
724
725
726
                };

                let (header_tx, header_rx) = oneshot::channel();
                let (sse_tx, sse_rx) = tokio::sync::mpsc::unbounded_channel();

                tokio::spawn(async move {
                    let (header_map, sse) = generate_stream_internal(
                        infer_clone.clone(),
                        compute_type_clone.clone(),
                        Json(generate_request),
                        on_message_callback,
                        span_clone.clone(),
                    )
                    .await;
727

728
729
                    // send and dont wait for response
                    let _ = header_tx.send(header_map);
730

731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
                    // pin an emit messages to the sse_tx
                    let mut sse = Box::pin(sse);
                    while let Some(event) = sse.next().await {
                        if sse_tx.send(event).is_err() {
                            tracing::error!("Failed to send event. Receiver dropped.");
                            break;
                        }
                    }
                });

                (header_rx, sse_rx)
            };
            response_streams.push_back(generate_future);
        }

        let mut all_rxs = vec![];

        while let Some((header_rx, sse_rx)) = response_streams.next().await {
            all_rxs.push(sse_rx);

            // get the headers from the first response of each stream
            let headers = header_rx.await.map_err(|e| {
                tracing::error!("Failed to get headers: {:?}", e);
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(ErrorResponse {
                        error: "Failed to get headers".to_string(),
                        error_type: "headers".to_string(),
                    }),
760
                )
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
            })?;
            if x_compute_type.is_none() {
                x_compute_type = headers
                    .get("x-compute-type")
                    .and_then(|v| v.to_str().ok())
                    .map(|v| v.to_string());

                x_accel_buffering = headers
                    .get("x-accel-buffering")
                    .and_then(|v| v.to_str().ok())
                    .map(|v| v.to_string());
            }
            x_compute_characters += headers
                .get("x-compute-characters")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| v.parse().ok())
                .unwrap_or(0);
        }
779

780
781
782
783
784
785
786
787
        let mut headers = HeaderMap::new();
        if let Some(x_compute_type) = x_compute_type {
            headers.insert("x-compute-type", x_compute_type.parse().unwrap());
        }
        headers.insert("x-compute-characters", x_compute_characters.into());
        if let Some(x_accel_buffering) = x_accel_buffering {
            headers.insert("x-accel-buffering", x_accel_buffering.parse().unwrap());
        }
788

789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
        // now sink the sse streams into a single stream and remove the ones that are done
        let stream: AsyncStream<Result<Event, Infallible>, _> = async_stream::stream! {
            loop {
                let mut i = 0;
                while i < all_rxs.len() {
                    let rx = &mut all_rxs[i];
                    select! {
                        Some(event) = rx.recv() => {
                            yield event;
                        }
                        else => {
                            all_rxs.remove(i);
                            continue; // skip the increment to handle the next element at the same index
                        }
                    }
                    i += 1; // only increment when no element was removed
                }

                if all_rxs.is_empty() {
                    break;
                }
            }
        };

        let sse = Sse::new(stream).keep_alive(KeepAlive::default());
814
815
816
817
818
819
820
        Ok((headers, sse).into_response())
    } else {
        let current_time = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_else(|_| std::time::Duration::from_secs(0))
            .as_secs();

821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
        let responses = FuturesUnordered::new();
        for (index, generate_request) in generate_requests.into_iter().enumerate() {
            let infer_clone = infer.clone();
            let compute_type_clone = compute_type.clone();
            let span_clone = span.clone();
            let response_future = async move {
                let result = generate_internal(
                    Extension(infer_clone),
                    compute_type_clone,
                    Json(generate_request),
                    span_clone,
                )
                .await;
                result.map(|(headers, generation)| (index, headers, generation))
            };
            responses.push(response_future);
        }
        let generate_responses = responses.try_collect::<Vec<_>>().await?;

        let mut prompt_tokens = 0u32;
        let mut completion_tokens = 0u32;
        let mut total_tokens = 0u32;

        let mut x_compute_time = 0u32;
        let mut x_total_time = 0u32;
        let mut x_validation_time = 0u32;
        let mut x_queue_time = 0u32;
        let mut x_inference_time = 0u32;
        let mut x_time_per_token = 0u32;
        let mut x_prompt_tokens = 0u32;
        let mut x_generated_tokens = 0u32;

        let choices = generate_responses
            .into_iter()
            .map(|(index, headers, Json(generation))| {
                let details = generation.details.ok_or((
                    // this should never happen but handle if details are missing unexpectedly
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(ErrorResponse {
                        error: "No details in generation".to_string(),
                        error_type: "no details".to_string(),
                    }),
                ))?;

                if x_compute_type.is_none() {
                    x_compute_type = headers
                        .get("x-compute-type")
                        .and_then(|v| v.to_str().ok())
                        .map(|v| v.to_string());
                }

                // accumulate headers and usage from each response
                x_compute_time += headers
                    .get("x-compute-time")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_compute_characters += headers
                    .get("x-compute-characters")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_total_time += headers
                    .get("x-total-time")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_validation_time += headers
                    .get("x-validation-time")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_queue_time += headers
                    .get("x-queue-time")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_inference_time += headers
                    .get("x-inference-time")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_time_per_token += headers
                    .get("x-time-per-token")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_prompt_tokens += headers
                    .get("x-prompt-tokens")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);
                x_generated_tokens += headers
                    .get("x-generated-tokens")
                    .and_then(|v| v.to_str().ok()?.parse().ok())
                    .unwrap_or(0);

                prompt_tokens += details.prefill.len() as u32;
                completion_tokens += details.generated_tokens;
                total_tokens += details.prefill.len() as u32 + details.generated_tokens;

                Ok(CompletionComplete {
                    finish_reason: details.finish_reason.to_string(),
                    index: index as u32,
                    logprobs: None,
                    text: generation.generated_text,
                })
            })
            .collect::<Result<Vec<_>, _>>()
            .map_err(|(status, Json(err))| (status, Json(err)))?;
923

924
        let response = Completion::Final(CompletionFinal {
925
926
927
928
929
930
931
932
            id: "".to_string(),
            created: current_time,
            model: info.model_id.clone(),
            system_fingerprint: format!(
                "{}-{}",
                info.version,
                info.docker_label.unwrap_or("native")
            ),
933
            choices,
934
            usage: Usage {
935
936
937
                prompt_tokens,
                completion_tokens,
                total_tokens,
938
            },
939
        });
940

941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
        // headers similar to `generate` but aggregated
        let mut headers = HeaderMap::new();
        if let Some(x_compute_type) = x_compute_type {
            headers.insert("x-compute-type", x_compute_type.parse().unwrap());
        }
        headers.insert("x-compute-characters", x_compute_characters.into());
        headers.insert("x-total-time", x_total_time.into());
        headers.insert("x-validation-time", x_validation_time.into());
        headers.insert("x-queue-time", x_queue_time.into());
        headers.insert("x-inference-time", x_inference_time.into());
        headers.insert("x-time-per-token", x_time_per_token.into());
        headers.insert("x-prompt-tokens", x_prompt_tokens.into());
        headers.insert("x-generated-tokens", x_generated_tokens.into());
        if let Some(x_accel_buffering) = x_accel_buffering {
            headers.insert("x-accel-buffering", x_accel_buffering.parse().unwrap());
        }
957
958
959
960
        Ok((headers, Json(response)).into_response())
    }
}

961
962
/// Generate tokens
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
post,
tag = "Text Generation Inference",
path = "/v1/chat/completions",
request_body = ChatRequest,
responses(
(status = 200, description = "Generated Chat Completion",
content(
("application/json" = ChatCompletion),
("text/event-stream" = ChatCompletionChunk),
)),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"})),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"})),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"})),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"})),
)
)]
983
#[instrument(
OlivierDehaene's avatar
OlivierDehaene committed
984
985
986
987
988
989
990
991
992
993
994
skip_all,
fields(
// parameters = ? req.parameters,
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
)]
995
996
async fn chat_completions(
    Extension(infer): Extension<Infer>,
997
    Extension(compute_type): Extension<ComputeType>,
998
999
1000
    Extension(info): Extension<Info>,
    Json(req): Json<ChatRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
1001
    let span = tracing::Span::current();
1002
    metrics::counter!("tgi_request_count").increment(1);
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
    let ChatRequest {
        logprobs,
        max_tokens,
        messages,
        presence_penalty,
        seed,
        stop,
        stream,
        tools,
        tool_choice,
        tool_prompt,
1014
        temperature,
drbh's avatar
drbh committed
1015
        response_format,
1016
1017
1018
1019
1020
1021
1022
1023
        ..
    } = req;

    let repetition_penalty = presence_penalty.map(|x| x + 2.0);
    let max_new_tokens = max_tokens.or(Some(100));
    let logprobs = logprobs.unwrap_or(false);
    let tool_prompt = tool_prompt.unwrap_or_default();
    let stop = stop.unwrap_or_default();
1024
1025
1026
1027
1028
    // enable greedy only when temperature is 0
    let (do_sample, temperature) = match temperature {
        Some(temperature) if temperature == 0.0 => (false, None),
        other => (true, other),
    };
1029

drbh's avatar
drbh committed
1030
1031
    // response_format and tools are mutually exclusive
    if response_format.is_some() && tools.as_ref().is_some() {
1032
        metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
drbh's avatar
drbh committed
1033
1034
1035
1036
1037
1038
1039
1040
1041
        return Err((
            StatusCode::UNPROCESSABLE_ENTITY,
            Json(ErrorResponse {
                error: "Grammar and tools are mutually exclusive".to_string(),
                error_type: "grammar and tools".to_string(),
            }),
        ));
    }

1042
1043
1044
    // extract tool grammar if present
    let tool_grammar = match ToolGrammar::apply(tools, tool_choice) {
        Ok(grammar) => grammar,
1045
        Err(err) => {
1046
            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
            tracing::error!("{err}");
            return Err((
                StatusCode::UNPROCESSABLE_ENTITY,
                Json(ErrorResponse {
                    error: err.to_string(),
                    error_type: err.error_type().to_string(),
                }),
            ));
        }
    };

drbh's avatar
drbh committed
1058
1059
    // determine the appropriate arguments for apply_chat_template
    let tools_grammar_prompt = tool_grammar
1060
1061
        .as_ref()
        .map(|t| (GrammarType::Json(serde_json::json!(t)), tool_prompt));
drbh's avatar
drbh committed
1062

drbh's avatar
drbh committed
1063
1064
1065
1066
1067
1068
1069
    let (tools_grammar_prompt, grammar) = match response_format {
        Some(response_format) => (None, Some(response_format)),
        None => (
            tools_grammar_prompt.clone(),
            tools_grammar_prompt.map(|(grammar, _)| grammar.clone()),
        ),
    };
drbh's avatar
drbh committed
1070

1071
    // apply chat template to flatten the request into a single input
drbh's avatar
drbh committed
1072
    let inputs = match infer.apply_chat_template(messages, tools_grammar_prompt) {
1073
1074
        Ok(inputs) => inputs,
        Err(err) => {
1075
            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
1076
1077
            tracing::error!("{err}");
            return Err((
drbh's avatar
drbh committed
1078
1079
                StatusCode::UNPROCESSABLE_ENTITY,
                Json(ErrorResponse {
1080
1081
                    error: err.to_string(),
                    error_type: err.error_type().to_string(),
drbh's avatar
drbh committed
1082
                }),
1083
1084
            ));
        }
drbh's avatar
drbh committed
1085
1086
    };

1087
1088
1089
1090
1091
    // build the request passing some parameters
    let generate_request = GenerateRequest {
        inputs: inputs.to_string(),
        parameters: GenerateParameters {
            best_of: None,
1092
            temperature,
1093
            repetition_penalty,
1094
            frequency_penalty: req.frequency_penalty,
1095
            top_k: None,
1096
            top_p: req.top_p,
1097
            typical_p: None,
1098
            do_sample,
1099
1100
            max_new_tokens,
            return_full_text: None,
1101
            stop,
1102
1103
1104
            truncate: None,
            watermark: false,
            details: true,
1105
            decoder_input_details: !stream,
1106
            seed,
1107
            top_n_tokens: req.top_logprobs,
drbh's avatar
drbh committed
1108
            grammar,
drbh's avatar
drbh committed
1109
            ..Default::default()
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
        },
    };

    // static values that will be returned in all cases
    let model_id = info.model_id.clone();
    let system_fingerprint = format!("{}-{}", info.version, info.docker_label.unwrap_or("native"));

    // switch on stream
    if stream {
        // pass this callback to the stream generation and build the required event structure
        let on_message_callback = move |stream_token: StreamResponse| {
            let event = Event::default();

            let current_time = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_else(|_| std::time::Duration::from_secs(0))
                .as_secs();

1128
1129
1130
1131
            let logprobs = logprobs.then(|| {
                ChatCompletionLogprobs::from((stream_token.token.clone(), stream_token.top_tokens))
            });

drbh's avatar
drbh committed
1132
1133
1134
1135
            // replace the content with the tool calls if grammar is present
            let (content, tool_calls) = if tool_grammar.is_some() {
                (None, Some(vec![stream_token.token.text]))
            } else {
1136
1137
1138
1139
1140
1141
1142
                let content = if !stream_token.token.special {
                    Some(stream_token.token.text)
                } else {
                    None
                };

                (content, None)
drbh's avatar
drbh committed
1143
1144
            };

1145
            event
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
                .json_data(CompletionType::ChatCompletionChunk(
                    ChatCompletionChunk::new(
                        model_id.clone(),
                        system_fingerprint.clone(),
                        content,
                        tool_calls,
                        current_time,
                        logprobs,
                        stream_token.details.map(|d| d.finish_reason.to_string()),
                    ),
1156
                ))
1157
1158
1159
1160
                .unwrap_or_else(|e| {
                    println!("Failed to serialize ChatCompletionChunk: {:?}", e);
                    Event::default()
                })
1161
1162
        };

1163
1164
1165
1166
1167
        let (headers, response_stream) = generate_stream_internal(
            infer,
            compute_type,
            Json(generate_request),
            on_message_callback,
1168
            span,
1169
1170
        )
        .await;
1171
1172
1173
        let sse = Sse::new(response_stream).keep_alive(KeepAlive::default());
        Ok((headers, sse).into_response())
    } else {
1174
1175
        let (headers, Json(generation)) =
            generate_internal(Extension(infer), compute_type, Json(generate_request), span).await?;
1176
1177
1178
1179
1180
1181

        let current_time = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_else(|_| std::time::Duration::from_secs(0))
            .as_secs();

drbh's avatar
drbh committed
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
        let (tool_calls, output) = if tool_grammar.is_some() {
            // gen_text should be valid json
            let gen_text_value: Value =
                serde_json::from_str(&generation.generated_text).map_err(|e| {
                    (
                        StatusCode::UNPROCESSABLE_ENTITY,
                        Json(ErrorResponse {
                            error: e.to_string(),
                            error_type: "Input validation error".to_string(),
                        }),
                    )
                })?;
1194
            let tool_calls = vec![ToolCall {
1195
                id: "0".to_string(),
drbh's avatar
drbh committed
1196
1197
1198
                r#type: "function".to_string(),
                function: FunctionDefinition {
                    description: None,
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
                    name: gen_text_value
                        .get("function")
                        .and_then(|f| f.get("_name"))
                        .and_then(|name| name.as_str())
                        .unwrap_or("default_function_name")
                        .to_string(),
                    // Serialize the JSON object obtained from "function" to an escaped JSON string
                    arguments: gen_text_value
                        .get("function")
                        .map(|f| {
                            let mut f_cloned = f.clone();
                            if let Value::Object(ref mut props) = f_cloned {
                                props.remove("_name");
                            }
                            f_cloned
                        })
                        .unwrap_or_default(),
drbh's avatar
drbh committed
1216
                },
1217
1218
            }];
            (Some(tool_calls), None)
drbh's avatar
drbh committed
1219
1220
1221
        } else {
            (None, Some(generation.generated_text))
        };
1222
        // build the complete response object with the full text
1223
        let response = CompletionType::ChatCompletion(ChatCompletion::new(
1224
1225
            model_id,
            system_fingerprint,
drbh's avatar
drbh committed
1226
            output,
1227
1228
1229
            current_time,
            generation.details.unwrap(),
            logprobs,
drbh's avatar
drbh committed
1230
            tool_calls,
1231
        ));
1232
1233
1234
1235

        // wrap generation inside a Vec to match api-inference
        Ok((headers, Json(response)).into_response())
    }
1236
1237
}

drbh's avatar
drbh committed
1238
1239
/// Generate tokens from Vertex request
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
post,
tag = "Text Generation Inference",
path = "/vertex",
request_body = VertexRequest,
responses(
(status = 200, description = "Generated Text", body = VertexResponse),
(status = 424, description = "Generation Error", body = ErrorResponse,
example = json ! ({"error": "Request failed during generation"})),
(status = 429, description = "Model is overloaded", body = ErrorResponse,
example = json ! ({"error": "Model is overloaded"})),
(status = 422, description = "Input validation error", body = ErrorResponse,
example = json ! ({"error": "Input validation error"})),
(status = 500, description = "Incomplete generation", body = ErrorResponse,
example = json ! ({"error": "Incomplete generation"})),
)
)]
drbh's avatar
drbh committed
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
#[instrument(
    skip_all,
    fields(
        total_time,
        validation_time,
        queue_time,
        inference_time,
        time_per_token,
        seed,
    )
)]
async fn vertex_compatibility(
    Extension(infer): Extension<Infer>,
    Extension(compute_type): Extension<ComputeType>,
    Json(req): Json<VertexRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
1272
    let span = tracing::Span::current();
1273
    metrics::counter!("tgi_request_count").increment(1);
drbh's avatar
drbh committed
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303

    // check that theres at least one instance
    if req.instances.is_empty() {
        return Err((
            StatusCode::UNPROCESSABLE_ENTITY,
            Json(ErrorResponse {
                error: "Input validation error".to_string(),
                error_type: "Input validation error".to_string(),
            }),
        ));
    }

    // Process all instances
    let predictions = req
        .instances
        .iter()
        .map(|instance| {
            let generate_request = GenerateRequest {
                inputs: instance.inputs.clone(),
                parameters: GenerateParameters {
                    do_sample: true,
                    max_new_tokens: instance.parameters.as_ref().and_then(|p| p.max_new_tokens),
                    seed: instance.parameters.as_ref().and_then(|p| p.seed),
                    details: true,
                    decoder_input_details: true,
                    ..Default::default()
                },
            };

            async {
1304
                generate_internal(
drbh's avatar
drbh committed
1305
                    Extension(infer.clone()),
1306
                    compute_type.clone(),
drbh's avatar
drbh committed
1307
                    Json(generate_request),
1308
                    span.clone(),
drbh's avatar
drbh committed
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
                )
                .await
                .map(|(_, Json(generation))| generation.generated_text)
                .map_err(|_| {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        Json(ErrorResponse {
                            error: "Incomplete generation".into(),
                            error_type: "Incomplete generation".into(),
                        }),
                    )
                })
            }
        })
        .collect::<FuturesUnordered<_>>()
        .try_collect::<Vec<_>>()
        .await?;

    let response = VertexResponse { predictions };
    Ok((HeaderMap::new(), Json(response)).into_response())
}

1331
1332
/// Tokenize inputs
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
post,
tag = "Text Generation Inference",
path = "/tokenize",
request_body = GenerateRequest,
responses(
(status = 200, description = "Tokenized ids", body = TokenizeResponse),
(status = 404, description = "No tokenizer found", body = ErrorResponse,
example = json ! ({"error": "No fast tokenizer available"})),
)
)]
1343
1344
1345
1346
#[instrument(skip_all)]
async fn tokenize(
    Extension(infer): Extension<Infer>,
    Json(req): Json<GenerateRequest>,
1347
) -> Result<Json<TokenizeResponse>, (StatusCode, Json<ErrorResponse>)> {
1348
1349
1350
1351
1352
1353
1354
1355
    let input = req.inputs.clone();
    let encoding = infer.tokenize(req).await?;
    if let Some(encoding) = encoding {
        let tokens: Vec<SimpleToken> = encoding
            .get_ids()
            .iter()
            .zip(encoding.get_offsets())
            .map(|(&id, &(start, stop))| {
1356
1357
                let text: String =
                    String::from_utf8_lossy(&input.as_bytes()[start..stop]).to_string();
1358
1359
1360
1361
1362
1363
1364
1365
                SimpleToken {
                    id,
                    text,
                    start,
                    stop,
                }
            })
            .collect();
1366
        Ok(Json(TokenizeResponse(tokens)))
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
    } else {
        Err((
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "No fast tokenizer or tokenizer.json for this model".to_string(),
                error_type: "no fast tokenizer".to_string(),
            }),
        ))
    }
}

1378
1379
/// Prometheus metrics scrape endpoint
#[utoipa::path(
1380
1381
1382
1383
    get,
    tag = "Text Generation Inference",
    path = "/metrics",
    responses((status = 200, description = "Prometheus Metrics", body = String))
1384
1385
1386
1387
1388
)]
async fn metrics(prom_handle: Extension<PrometheusHandle>) -> String {
    prom_handle.render()
}

1389
1390
1391
#[derive(Clone, Debug)]
pub(crate) struct ComputeType(String);

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1392
1393
1394
/// Serving method
#[allow(clippy::too_many_arguments)]
pub async fn run(
OlivierDehaene's avatar
OlivierDehaene committed
1395
    master_shard_uds_path: String,
1396
    model_info: HubModelInfo,
1397
    compat_return_full_text: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1398
    max_concurrent_requests: usize,
1399
    max_best_of: usize,
1400
    max_stop_sequences: usize,
Nicolas Patry's avatar
Nicolas Patry committed
1401
    max_top_n_tokens: u32,
OlivierDehaene's avatar
OlivierDehaene committed
1402
    max_input_tokens: usize,
1403
    max_total_tokens: usize,
1404
    waiting_served_ratio: f32,
1405
    max_batch_prefill_tokens: u32,
OlivierDehaene's avatar
OlivierDehaene committed
1406
    max_batch_total_tokens: Option<u32>,
1407
    max_waiting_tokens: usize,
1408
    max_batch_size: Option<usize>,
1409
    tokenizer: Option<Tokenizer>,
1410
    config: Option<Config>,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1411
1412
    validation_workers: usize,
    addr: SocketAddr,
1413
    allow_origin: Option<AllowOrigin>,
1414
    ngrok: bool,
1415
1416
    _ngrok_authtoken: Option<String>,
    _ngrok_edge: Option<String>,
1417
    tokenizer_config: HubTokenizerConfig,
1418
    preprocessor_config: Option<HubPreprocessorConfig>,
drbh's avatar
drbh committed
1419
    processor_config: HubProcessorConfig,
1420
    messages_api_enabled: bool,
drbh's avatar
drbh committed
1421
    grammar_support: bool,
1422
    max_client_batch_size: usize,
1423
    print_schema_command: bool,
OlivierDehaene's avatar
OlivierDehaene committed
1424
) -> Result<(), WebServerError> {
1425
1426
1427
    // OpenAPI documentation
    #[derive(OpenApi)]
    #[openapi(
1428
1429
1430
1431
1432
1433
    paths(
    health,
    get_model_info,
    compat_generate,
    generate,
    generate_stream,
1434
    chat_completions,
1435
    completions,
1436
    tokenize,
1437
1438
1439
1440
1441
1442
1443
    metrics,
    ),
    components(
    schemas(
    Info,
    CompatGenerateRequest,
    GenerateRequest,
1444
    GrammarType,
1445
1446
    ChatRequest,
    Message,
1447
    ChatCompletionComplete,
1448
1449
1450
    ChatCompletionChoice,
    ChatCompletionDelta,
    ChatCompletionChunk,
1451
1452
1453
    ChatCompletionLogprob,
    ChatCompletionLogprobs,
    ChatCompletionTopLogprob,
1454
    ChatCompletion,
1455
1456
    CompletionRequest,
    CompletionComplete,
1457
1458
1459
1460
    Chunk,
    Completion,
    CompletionFinal,
    Prompt,
1461
1462
1463
1464
    GenerateParameters,
    PrefillToken,
    Token,
    GenerateResponse,
1465
1466
    TokenizeResponse,
    SimpleToken,
1467
1468
1469
1470
1471
1472
    BestOfSequence,
    Details,
    FinishReason,
    StreamResponse,
    StreamDetails,
    ErrorResponse,
drbh's avatar
drbh committed
1473
    GrammarType,
1474
    Usage,
OlivierDehaene's avatar
OlivierDehaene committed
1475
1476
1477
1478
1479
1480
    DeltaToolCall,
    ToolType,
    Tool,
    ToolCall,
    Function,
    FunctionDefinition,
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
    )
    ),
    tags(
    (name = "Text Generation Inference", description = "Hugging Face Text Generation Inference API")
    ),
    info(
    title = "Text Generation Inference",
    license(
    name = "Apache 2.0",
    url = "https://www.apache.org/licenses/LICENSE-2.0"
    )
    )
1493
1494
1495
    )]
    struct ApiDoc;

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1496
    // Create state
1497
1498
1499
1500
1501
1502
    if print_schema_command {
        let api_doc = ApiDoc::openapi();
        let api_doc = serde_json::to_string_pretty(&api_doc).unwrap();
        println!("{}", api_doc);
        std::process::exit(0);
    }
OlivierDehaene's avatar
OlivierDehaene committed
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637

    // Open connection, get model info and warmup
    let (scheduler, health_ext, shard_info, max_batch_total_tokens): (
        Arc<dyn Scheduler + Send + Sync>,
        HealthCheck,
        ShardInfo,
        u32,
    ) = {
        // Helper function to check both v2 and v3
        let check_max_batch_total_tokens = |max_supported_batch_total_tokens: Option<u32>| {
            match max_supported_batch_total_tokens {
                // Older models do not support automatic max-batch-total-tokens
                None => {
                    let max_batch_total_tokens = max_batch_total_tokens.unwrap_or(
                        16000.max((max_total_tokens as u32).max(max_batch_prefill_tokens)),
                    );
                    tracing::warn!("Model does not support automatic max batch total tokens");
                    Ok(max_batch_total_tokens)
                }
                // Flash attention models return their max supported total tokens
                Some(max_supported_batch_total_tokens) => {
                    // Warn if user added his own max-batch-total-tokens as we will ignore it
                    if max_batch_total_tokens.is_some() {
                        tracing::warn!(
                            "`--max-batch-total-tokens` is deprecated for Flash \
                        Attention models."
                        );
                        tracing::warn!(
                            "Inferred max batch total tokens: {max_supported_batch_total_tokens}"
                        );
                    }
                    if max_total_tokens as u32 > max_supported_batch_total_tokens {
                        return Err(WebServerError::NotEnoughMemory(max_total_tokens));
                    }

                    Ok(max_supported_batch_total_tokens)
                }
            }
        };

        let generation_health = Arc::new(AtomicBool::new(false));

        match v3::ShardedClient::connect_uds(master_shard_uds_path.clone()).await {
            Ok(mut sharded_client) => {
                // server is running on v3
                // Clear the cache; useful if the webserver rebooted
                sharded_client
                    .clear_cache(None)
                    .await
                    .map_err(WebServerError::Cache)?;
                // Get info from the shard
                let shard_info = sharded_client.info().await.map_err(WebServerError::Info)?;

                // Warmup model
                tracing::info!("Warming up model");
                let max_batch_total_tokens = check_max_batch_total_tokens(
                    sharded_client
                        .warmup(
                            max_input_tokens as u32,
                            max_batch_prefill_tokens,
                            max_total_tokens as u32,
                            max_batch_size,
                        )
                        .await
                        .map_err(WebServerError::Warmup)?,
                )?;

                let health_ext =
                    HealthCheck::new(Arc::new(sharded_client.clone()), generation_health.clone());
                let scheduler = Arc::new(SchedulerV3::new(
                    sharded_client,
                    waiting_served_ratio,
                    max_batch_prefill_tokens,
                    max_batch_total_tokens,
                    max_waiting_tokens,
                    max_batch_size,
                    shard_info.requires_padding,
                    shard_info.window_size,
                    shard_info.speculate,
                    generation_health,
                ));
                tracing::info!("Using scheduler V3");

                (scheduler, health_ext, shard_info, max_batch_total_tokens)
            }
            Err(_) => {
                let mut sharded_client = v2::ShardedClient::connect_uds(master_shard_uds_path)
                    .await
                    .map_err(WebServerError::Connection)?;

                // server is running on v2
                // Clear the cache; useful if the webserver rebooted
                sharded_client
                    .clear_cache(None)
                    .await
                    .map_err(WebServerError::Cache)?;
                // Get info from the shard
                let shard_info = sharded_client.info().await.map_err(WebServerError::Info)?;

                // Warmup model
                tracing::info!("Warming up model");
                let max_batch_total_tokens = check_max_batch_total_tokens(
                    sharded_client
                        .warmup(
                            max_input_tokens as u32,
                            max_batch_prefill_tokens,
                            max_total_tokens as u32,
                            max_batch_size,
                        )
                        .await
                        .map_err(WebServerError::Warmup)?,
                )?;

                let health_ext =
                    HealthCheck::new(Arc::new(sharded_client.clone()), generation_health.clone());
                let scheduler = Arc::new(SchedulerV2::new(
                    sharded_client,
                    waiting_served_ratio,
                    max_batch_prefill_tokens,
                    max_batch_total_tokens,
                    max_waiting_tokens,
                    max_batch_size,
                    shard_info.requires_padding,
                    shard_info.window_size,
                    shard_info.speculate,
                    generation_health,
                ));
                tracing::info!("Using scheduler V2");

                (scheduler, health_ext, shard_info, max_batch_total_tokens)
            }
        }
    };
    tracing::info!("Setting max batch total tokens to {max_batch_total_tokens}");

1638
1639
1640
    let validation = Validation::new(
        validation_workers,
        tokenizer,
1641
        config,
1642
        preprocessor_config,
1643
        max_best_of,
1644
        max_stop_sequences,
Nicolas Patry's avatar
Nicolas Patry committed
1645
        max_top_n_tokens,
OlivierDehaene's avatar
OlivierDehaene committed
1646
        max_input_tokens,
1647
        max_total_tokens,
drbh's avatar
drbh committed
1648
        grammar_support,
1649
    );
OlivierDehaene's avatar
OlivierDehaene committed
1650

1651
    let infer = Infer::new(
OlivierDehaene's avatar
OlivierDehaene committed
1652
        scheduler,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1653
        validation,
1654
        max_concurrent_requests,
1655
        tokenizer_config,
drbh's avatar
drbh committed
1656
        processor_config,
1657
    );
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1658

1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
    // Duration buckets
    let duration_matcher = Matcher::Suffix(String::from("duration"));
    let n_duration_buckets = 35;
    let mut duration_buckets = Vec::with_capacity(n_duration_buckets);
    // Minimum duration in seconds
    let mut value = 0.0001;
    for _ in 0..n_duration_buckets {
        // geometric sequence
        value *= 1.5;
        duration_buckets.push(value);
    }
    // Input Length buckets
    let input_length_matcher = Matcher::Full(String::from("tgi_request_input_length"));
    let input_length_buckets: Vec<f64> = (0..100)
OlivierDehaene's avatar
OlivierDehaene committed
1673
        .map(|x| (max_input_tokens as f64 / 100.0) * (x + 1) as f64)
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
        .collect();
    // Generated tokens buckets
    let generated_tokens_matcher = Matcher::Full(String::from("tgi_request_generated_tokens"));
    let generated_tokens_buckets: Vec<f64> = (0..100)
        .map(|x| (max_total_tokens as f64 / 100.0) * (x + 1) as f64)
        .collect();
    // Input Length buckets
    let max_new_tokens_matcher = Matcher::Full(String::from("tgi_request_max_new_tokens"));
    let max_new_tokens_buckets: Vec<f64> = (0..100)
        .map(|x| (max_total_tokens as f64 / 100.0) * (x + 1) as f64)
        .collect();
    // Batch size buckets
    let batch_size_matcher = Matcher::Full(String::from("tgi_batch_next_size"));
1687
    let batch_size_buckets: Vec<f64> = (0..1024).map(|x| (x + 1) as f64).collect();
OlivierDehaene's avatar
OlivierDehaene committed
1688
1689
1690
    // Speculated tokens buckets
    let skipped_matcher = Matcher::Full(String::from("tgi_request_skipped_tokens"));
    let skipped_buckets: Vec<f64> = (0..shard_info.speculate + 1).map(|x| x as f64).collect();
1691

1692
    // Prometheus handler
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
    let builder = PrometheusBuilder::new()
        .set_buckets_for_metric(duration_matcher, &duration_buckets)
        .unwrap()
        .set_buckets_for_metric(input_length_matcher, &input_length_buckets)
        .unwrap()
        .set_buckets_for_metric(generated_tokens_matcher, &generated_tokens_buckets)
        .unwrap()
        .set_buckets_for_metric(max_new_tokens_matcher, &max_new_tokens_buckets)
        .unwrap()
        .set_buckets_for_metric(batch_size_matcher, &batch_size_buckets)
OlivierDehaene's avatar
OlivierDehaene committed
1703
1704
        .unwrap()
        .set_buckets_for_metric(skipped_matcher, &skipped_buckets)
1705
        .unwrap();
1706
1707
1708
1709
    let prom_handle = builder
        .install_recorder()
        .expect("failed to install metrics recorder");

1710
1711
1712
1713
1714
1715
1716
    // CORS layer
    let allow_origin = allow_origin.unwrap_or(AllowOrigin::any());
    let cors_layer = CorsLayer::new()
        .allow_methods([Method::GET, Method::POST])
        .allow_headers([http::header::CONTENT_TYPE])
        .allow_origin(allow_origin);

1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
    // Endpoint info
    let info = Info {
        model_id: model_info.model_id,
        model_sha: model_info.sha,
        model_dtype: shard_info.dtype,
        model_device_type: shard_info.device_type,
        model_pipeline_tag: model_info.pipeline_tag,
        max_concurrent_requests,
        max_best_of,
        max_stop_sequences,
OlivierDehaene's avatar
OlivierDehaene committed
1727
        max_input_tokens,
1728
1729
1730
1731
        max_total_tokens,
        waiting_served_ratio,
        max_batch_total_tokens,
        max_waiting_tokens,
1732
        max_batch_size,
1733
        validation_workers,
1734
        max_client_batch_size,
1735
        router: env!("CARGO_PKG_NAME"),
1736
1737
        version: env!("CARGO_PKG_VERSION"),
        sha: option_env!("VERGEN_GIT_SHA"),
1738
        docker_label: option_env!("DOCKER_LABEL"),
1739
1740
    };

1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
    #[allow(unused_mut)] // mut is needed for conditional compilation
    let mut doc = ApiDoc::openapi();

    #[cfg(feature = "google")]
    {
        use crate::VertexInstance;

        #[derive(OpenApi)]
        #[openapi(
            paths(vertex_compatibility),
            components(schemas(VertexInstance, VertexRequest, VertexResponse))
        )]
        struct VertexApiDoc;

        doc.merge(VertexApiDoc::openapi());
    }

    #[cfg(feature = "kserve")]
    {
        use crate::kserve::{
            InferenceOutput, InferenceRequest, LiveResponse, MetadataServerResponse, OutputChunk,
            ReadyResponse,
        };
        use crate::kserve::{
            __path_kerve_server_metadata, __path_kserve_health_live, __path_kserve_health_ready,
            __path_kserve_model_infer, __path_kserve_model_metadata,
            __path_kserve_model_metadata_ready,
        };

        #[derive(OpenApi)]
        #[openapi(
            paths(
                kserve_health_live,
                kserve_health_ready,
                kerve_server_metadata,
                kserve_model_metadata,
                kserve_model_metadata_ready,
1778
                kserve_model_infer,
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
            ),
            components(schemas(
                InferenceOutput,
                InferenceRequest,
                LiveResponse,
                MetadataServerResponse,
                OutputChunk,
                ReadyResponse,
            ))
        )]
        struct KServeApiDoc;

        doc.merge(KServeApiDoc::openapi());
    }
drbh's avatar
drbh committed
1793

1794
    // Configure Swagger UI
drbh's avatar
drbh committed
1795
    let swagger_ui = SwaggerUi::new("/docs").url("/api-doc/openapi.json", doc);
1796
1797
1798

    // Define base and health routes
    let base_routes = Router::new()
1799
        .route("/", post(compat_generate))
1800
        .route("/", get(health))
1801
        .route("/info", get(get_model_info))
Olivier Dehaene's avatar
Olivier Dehaene committed
1802
        .route("/generate", post(generate))
1803
        .route("/generate_stream", post(generate_stream))
1804
        .route("/v1/chat/completions", post(chat_completions))
1805
        .route("/v1/completions", post(completions))
drbh's avatar
drbh committed
1806
        .route("/vertex", post(vertex_compatibility))
1807
        .route("/tokenize", post(tokenize))
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1808
        .route("/health", get(health))
1809
        .route("/ping", get(health))
1810
1811
1812
        .route("/metrics", get(metrics));

    // Conditional AWS Sagemaker route
1813
    let aws_sagemaker_route = if messages_api_enabled {
1814
1815
1816
1817
1818
        Router::new().route("/invocations", post(chat_completions)) // Use 'chat_completions' for OAI_ENABLED
    } else {
        Router::new().route("/invocations", post(compat_generate)) // Use 'compat_generate' otherwise
    };

1819
1820
    let compute_type =
        ComputeType(std::env::var("COMPUTE_TYPE").unwrap_or("gpu+optimized".to_string()));
1821

1822
    // Combine routes and layers
drbh's avatar
drbh committed
1823
    let mut app = Router::new()
1824
1825
        .merge(swagger_ui)
        .merge(base_routes)
drbh's avatar
drbh committed
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
        .merge(aws_sagemaker_route);

    #[cfg(feature = "google")]
    {
        tracing::info!("Built with `google` feature");
        tracing::info!(
            "Environment variables `AIP_PREDICT_ROUTE` and `AIP_HEALTH_ROUTE` will be respected."
        );
        if let Ok(env_predict_route) = std::env::var("AIP_PREDICT_ROUTE") {
            app = app.route(&env_predict_route, post(vertex_compatibility));
        }
        if let Ok(env_health_route) = std::env::var("AIP_HEALTH_ROUTE") {
            app = app.route(&env_health_route, get(health));
        }
    }

1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
    #[cfg(feature = "kserve")]
    {
        tracing::info!("Built with `kserve` feature");
        app = app
            .route(
                "/v2/models/:model_name/versions/:model_version/infer",
                post(kserve_model_infer),
            )
            .route(
                "/v2/models/:model_name/versions/:model_version",
                get(kserve_model_metadata),
            )
            .route("/v2/health/ready", get(kserve_health_ready))
            .route("/v2/health/live", get(kserve_health_live))
            .route("/v2", get(kerve_server_metadata))
            .route(
                "/v2/models/:model_name/versions/:model_version/ready",
                get(kserve_model_metadata_ready),
            );
    }

drbh's avatar
drbh committed
1863
1864
    // add layers after routes
    app = app
1865
        .layer(Extension(info))
1866
        .layer(Extension(health_ext.clone()))
1867
1868
        .layer(Extension(compat_return_full_text))
        .layer(Extension(infer))
1869
        .layer(Extension(compute_type))
1870
        .layer(Extension(prom_handle.clone()))
Nicolas Patry's avatar
Nicolas Patry committed
1871
        .layer(OtelAxumLayer::default())
1872
        .layer(cors_layer);
Olivier Dehaene's avatar
Olivier Dehaene committed
1873

OlivierDehaene's avatar
OlivierDehaene committed
1874
1875
    tracing::info!("Connected");

1876
1877
1878
    if ngrok {
        #[cfg(feature = "ngrok")]
        {
1879
            panic!("ngrok feature is not functional with axum=0.7 and hyper=1, waiting on https://github.com/ngrok/ngrok-rust/pull/137/files to re-enable.");
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893

            // Run server
        }
        #[cfg(not(feature = "ngrok"))]
        {
            let _ngrok_authtoken = ngrok_authtoken;
            let _ngrok_domain = ngrok_domain;
            let _ngrok_username = ngrok_username;
            let _ngrok_password = ngrok_password;

            panic!("`text-generation-router` was compiled without the `ngrok` feature");
        }
    } else {
        // Run server
1894
1895
1896

        let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
        axum::serve(listener, app)
1897
            .with_graceful_shutdown(shutdown_signal())
OlivierDehaene's avatar
OlivierDehaene committed
1898
1899
            .await
            .map_err(|err| WebServerError::Axum(Box::new(err)))?;
1900
    }
1901
    Ok(())
Olivier Dehaene's avatar
Olivier Dehaene committed
1902
}
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928

/// 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");
1929
    opentelemetry::global::shutdown_tracer_provider();
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1930
}
1931
1932
1933
1934
1935
1936
1937
1938
1939

/// Convert to Axum supported formats
impl From<InferError> for (StatusCode, Json<ErrorResponse>) {
    fn from(err: InferError) -> Self {
        let status_code = match err {
            InferError::GenerationError(_) => StatusCode::FAILED_DEPENDENCY,
            InferError::Overloaded(_) => StatusCode::TOO_MANY_REQUESTS,
            InferError::ValidationError(_) => StatusCode::UNPROCESSABLE_ENTITY,
            InferError::IncompleteGeneration => StatusCode::INTERNAL_SERVER_ERROR,
1940
            InferError::TemplateError(_) => StatusCode::UNPROCESSABLE_ENTITY,
1941
            InferError::ToolError(_) => StatusCode::UNPROCESSABLE_ENTITY,
1942
1943
1944
1945
1946
1947
        };

        (
            status_code,
            Json(ErrorResponse {
                error: err.to_string(),
1948
                error_type: err.error_type().to_string(),
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
            }),
        )
    }
}

impl From<InferError> for Event {
    fn from(err: InferError) -> Self {
        Event::default()
            .json_data(ErrorResponse {
                error: err.to_string(),
1959
                error_type: err.error_type().to_string(),
1960
1961
1962
1963
            })
            .unwrap()
    }
}
OlivierDehaene's avatar
OlivierDehaene committed
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979

#[derive(Debug, Error)]
pub enum WebServerError {
    #[error("Unable to connect to the Python model shards: {0}")]
    Connection(ClientError),
    #[error("Unable to clear the Python model shards cache: {0}")]
    Cache(ClientError),
    #[error("Unable to get the Python model shards info: {0}")]
    Info(ClientError),
    #[error("Unable to warmup the Python model shards: {0}")]
    Warmup(ClientError),
    #[error("Not enough memory to handle `max_total_tokens={0}`")]
    NotEnoughMemory(usize),
    #[error("Axum error: {0}")]
    Axum(#[from] axum::BoxError),
}