"tests/vscode:/vscode.git/clone" did not exist on "e3a20fef428ba870bf906caf9218945b10b3c0d1"
server.rs 71.9 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
    GenerateResponse, GrammarType, HubModelInfo, HubProcessorConfig, HubTokenizerConfig, Info,
16
17
    Message, MessageChunk, MessageContent, 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
    let CompletionRequest {
600
        model,
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
        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),
    };
616
617
618

    // if suffix is present throw an error
    if req.suffix.is_some() {
619
        metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
620
621
622
623
624
625
626
627
628
629
        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(),
            }),
        ));
    }

630
    if req.prompt.0.len() > info.max_client_batch_size {
631
        metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
632
633
634
635
636
637
638
639
640
641
642
643
644
645
        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
646
        .0
647
648
649
650
651
        .iter()
        .map(|prompt| GenerateRequest {
            inputs: prompt.to_string(),
            parameters: GenerateParameters {
                best_of: None,
652
                temperature,
653
654
655
656
657
                repetition_penalty: req.repetition_penalty,
                frequency_penalty: req.frequency_penalty,
                top_k: None,
                top_p: req.top_p,
                typical_p: None,
658
                do_sample,
659
660
                max_new_tokens,
                return_full_text: None,
661
                stop: stop.clone(),
662
663
664
665
666
667
668
                truncate: None,
                watermark: false,
                details: true,
                decoder_input_details: !stream,
                seed,
                top_n_tokens: None,
                grammar: None,
669
                adapter_id: model.as_ref().filter(|m| *m != "tgi").map(String::from),
670
671
672
673
674
675
676
            },
        })
        .collect();

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

    if stream {
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
        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
699
                        .json_data(Completion::Chunk(Chunk {
700
701
702
703
704
705
706
707
708
709
710
711
                            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(),
712
                        }))
713
                        .unwrap_or_else(|_e| Event::default())
714
715
716
717
718
719
720
721
722
723
724
725
726
727
                };

                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;
728

729
730
                    // send and dont wait for response
                    let _ = header_tx.send(header_map);
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
760
                    // 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(),
                    }),
761
                )
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
            })?;
            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);
        }
780

781
782
783
784
785
786
787
788
        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());
        }
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
814
        // 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());
815
816
817
818
819
820
821
        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();

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
923
        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)))?;
924

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

942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
        // 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());
        }
958
959
960
961
        Ok((headers, Json(response)).into_response())
    }
}

962
963
/// Generate tokens
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
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"})),
)
)]
984
#[instrument(
OlivierDehaene's avatar
OlivierDehaene committed
985
986
987
988
989
990
991
992
993
994
995
skip_all,
fields(
// parameters = ? req.parameters,
total_time,
validation_time,
queue_time,
inference_time,
time_per_token,
seed,
)
)]
996
997
async fn chat_completions(
    Extension(infer): Extension<Infer>,
998
    Extension(compute_type): Extension<ComputeType>,
999
1000
1001
    Extension(info): Extension<Info>,
    Json(req): Json<ChatRequest>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
1002
    let span = tracing::Span::current();
1003
    metrics::counter!("tgi_request_count").increment(1);
1004
    let ChatRequest {
1005
        model,
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
        logprobs,
        max_tokens,
        messages,
        presence_penalty,
        seed,
        stop,
        stream,
        tools,
        tool_choice,
        tool_prompt,
1016
        temperature,
drbh's avatar
drbh committed
1017
        response_format,
1018
1019
1020
1021
1022
1023
1024
1025
        ..
    } = 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();
1026
1027
1028
1029
1030
    // enable greedy only when temperature is 0
    let (do_sample, temperature) = match temperature {
        Some(temperature) if temperature == 0.0 => (false, None),
        other => (true, other),
    };
1031

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

1044
1045
1046
    // extract tool grammar if present
    let tool_grammar = match ToolGrammar::apply(tools, tool_choice) {
        Ok(grammar) => grammar,
1047
        Err(err) => {
1048
            metrics::counter!("tgi_request_failure", "err" => "validation").increment(1);
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
            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
1060
1061
    // determine the appropriate arguments for apply_chat_template
    let tools_grammar_prompt = tool_grammar
1062
1063
        .as_ref()
        .map(|t| (GrammarType::Json(serde_json::json!(t)), tool_prompt));
drbh's avatar
drbh committed
1064

drbh's avatar
drbh committed
1065
1066
1067
1068
1069
1070
1071
    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
1072

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

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

    // 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();

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

drbh's avatar
drbh committed
1134
1135
1136
1137
            // 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 {
1138
1139
1140
1141
1142
1143
1144
                let content = if !stream_token.token.special {
                    Some(stream_token.token.text)
                } else {
                    None
                };

                (content, None)
drbh's avatar
drbh committed
1145
1146
            };

1147
            event
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
                .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()),
                    ),
1158
                ))
1159
1160
1161
1162
                .unwrap_or_else(|e| {
                    println!("Failed to serialize ChatCompletionChunk: {:?}", e);
                    Event::default()
                })
1163
1164
        };

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

        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
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
        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(),
                        }),
                    )
                })?;
1196
            let tool_calls = vec![ToolCall {
1197
                id: "0".to_string(),
drbh's avatar
drbh committed
1198
1199
1200
                r#type: "function".to_string(),
                function: FunctionDefinition {
                    description: None,
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
                    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
1218
                },
1219
1220
            }];
            (Some(tool_calls), None)
drbh's avatar
drbh committed
1221
1222
1223
        } else {
            (None, Some(generation.generated_text))
        };
1224
        // build the complete response object with the full text
1225
        let response = CompletionType::ChatCompletion(ChatCompletion::new(
1226
1227
            model_id,
            system_fingerprint,
drbh's avatar
drbh committed
1228
            output,
1229
1230
1231
            current_time,
            generation.details.unwrap(),
            logprobs,
drbh's avatar
drbh committed
1232
            tool_calls,
1233
        ));
1234
1235
1236
1237

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

drbh's avatar
drbh committed
1240
1241
/// Generate tokens from Vertex request
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
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
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
#[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>)> {
1274
    let span = tracing::Span::current();
1275
    metrics::counter!("tgi_request_count").increment(1);
drbh's avatar
drbh committed
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
1304
1305

    // 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 {
1306
                generate_internal(
drbh's avatar
drbh committed
1307
                    Extension(infer.clone()),
1308
                    compute_type.clone(),
drbh's avatar
drbh committed
1309
                    Json(generate_request),
1310
                    span.clone(),
drbh's avatar
drbh committed
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
                )
                .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())
}

1333
1334
/// Tokenize inputs
#[utoipa::path(
OlivierDehaene's avatar
OlivierDehaene committed
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
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"})),
)
)]
1345
1346
1347
1348
#[instrument(skip_all)]
async fn tokenize(
    Extension(infer): Extension<Infer>,
    Json(req): Json<GenerateRequest>,
1349
) -> Result<Json<TokenizeResponse>, (StatusCode, Json<ErrorResponse>)> {
1350
1351
1352
1353
1354
1355
1356
1357
    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))| {
1358
1359
                let text: String =
                    String::from_utf8_lossy(&input.as_bytes()[start..stop]).to_string();
1360
1361
1362
1363
1364
1365
1366
1367
                SimpleToken {
                    id,
                    text,
                    start,
                    stop,
                }
            })
            .collect();
1368
        Ok(Json(TokenizeResponse(tokens)))
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
    } 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(),
            }),
        ))
    }
}

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

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

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

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1500
    // Create state
1501
1502
1503
1504
1505
1506
    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
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
1638
1639
1640
1641

    // 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}");

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

1655
    let infer = Infer::new(
OlivierDehaene's avatar
OlivierDehaene committed
1656
        scheduler,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1657
        validation,
1658
        max_concurrent_requests,
1659
        tokenizer_config,
drbh's avatar
drbh committed
1660
        processor_config,
1661
    );
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1662

1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
    // 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
1677
        .map(|x| (max_input_tokens as f64 / 100.0) * (x + 1) as f64)
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
        .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"));
1691
    let batch_size_buckets: Vec<f64> = (0..1024).map(|x| (x + 1) as f64).collect();
OlivierDehaene's avatar
OlivierDehaene committed
1692
1693
1694
    // 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();
1695

1696
    // Prometheus handler
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
    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
1707
1708
        .unwrap()
        .set_buckets_for_metric(skipped_matcher, &skipped_buckets)
1709
        .unwrap();
1710
1711
1712
1713
    let prom_handle = builder
        .install_recorder()
        .expect("failed to install metrics recorder");

1714
1715
1716
1717
1718
1719
1720
    // 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);

1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
    // 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
1731
        max_input_tokens,
1732
1733
1734
1735
        max_total_tokens,
        waiting_served_ratio,
        max_batch_total_tokens,
        max_waiting_tokens,
1736
        max_batch_size,
1737
        validation_workers,
1738
        max_client_batch_size,
1739
        router: env!("CARGO_PKG_NAME"),
1740
1741
        version: env!("CARGO_PKG_VERSION"),
        sha: option_env!("VERGEN_GIT_SHA"),
1742
        docker_label: option_env!("DOCKER_LABEL"),
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
1778
1779
1780
1781
    #[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,
1782
                kserve_model_infer,
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
            ),
            components(schemas(
                InferenceOutput,
                InferenceRequest,
                LiveResponse,
                MetadataServerResponse,
                OutputChunk,
                ReadyResponse,
            ))
        )]
        struct KServeApiDoc;

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

1798
    // Configure Swagger UI
drbh's avatar
drbh committed
1799
    let swagger_ui = SwaggerUi::new("/docs").url("/api-doc/openapi.json", doc);
1800
1801
1802

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

    // Conditional AWS Sagemaker route
1817
    let aws_sagemaker_route = if messages_api_enabled {
1818
1819
1820
1821
1822
        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
    };

1823
1824
    let compute_type =
        ComputeType(std::env::var("COMPUTE_TYPE").unwrap_or("gpu+optimized".to_string()));
1825

1826
    // Combine routes and layers
drbh's avatar
drbh committed
1827
    let mut app = Router::new()
1828
1829
        .merge(swagger_ui)
        .merge(base_routes)
drbh's avatar
drbh committed
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
        .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));
        }
    }

1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
    #[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
1867
1868
    // add layers after routes
    app = app
1869
        .layer(Extension(info))
1870
        .layer(Extension(health_ext.clone()))
1871
1872
        .layer(Extension(compat_return_full_text))
        .layer(Extension(infer))
1873
        .layer(Extension(compute_type))
1874
        .layer(Extension(prom_handle.clone()))
Nicolas Patry's avatar
Nicolas Patry committed
1875
        .layer(OtelAxumLayer::default())
1876
        .layer(cors_layer);
Olivier Dehaene's avatar
Olivier Dehaene committed
1877

OlivierDehaene's avatar
OlivierDehaene committed
1878
1879
    tracing::info!("Connected");

1880
1881
1882
    if ngrok {
        #[cfg(feature = "ngrok")]
        {
1883
            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.");
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897

            // 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
1898
1899
1900

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

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

/// 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,
1944
            InferError::TemplateError(_) => StatusCode::UNPROCESSABLE_ENTITY,
1945
            InferError::ToolError(_) => StatusCode::UNPROCESSABLE_ENTITY,
1946
1947
1948
1949
1950
1951
        };

        (
            status_code,
            Json(ErrorResponse {
                error: err.to_string(),
1952
                error_type: err.error_type().to_string(),
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
            }),
        )
    }
}

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

#[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),
}