infer.rs 33.5 KB
Newer Older
1
2
/// Batching and inference logic
use crate::validation::{Validation, ValidationError};
3
4
5
6
use crate::{
    ChatTemplateInputs, Entry, GenerateRequest, GenerateStreamResponse, HubTokenizerConfig,
    Message, PrefillToken, Queue, Token,
};
7
use futures::future::try_join_all;
8
use minijinja::{Environment, ErrorKind, Template};
9
use nohash_hasher::IntMap;
10
11
12
13
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
14
use text_generation_client::{
Nicolas Patry's avatar
Nicolas Patry committed
15
    Batch, CachedBatch, ClientError, GeneratedText, Generation, ShardedClient, Tokens,
16
17
};
use thiserror::Error;
OlivierDehaene's avatar
OlivierDehaene committed
18
use tokio::sync::mpsc::error::SendError;
19
use tokio::sync::{mpsc, Notify, Semaphore, TryAcquireError};
20
use tokio::time::Instant;
OlivierDehaene's avatar
OlivierDehaene committed
21
22
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_stream::StreamExt;
23
use tracing::{info_span, instrument, Instrument, Span};
24
25
26
27
28
29

/// Inference struct
#[derive(Clone)]
pub struct Infer {
    /// Validation
    validation: Validation,
30
31
    /// Request queue
    queue: Queue,
32
33
34
35
    /// Shared state
    shared: Arc<Shared>,
    /// Inference limit
    limit_concurrent_requests: Arc<Semaphore>,
36
37
38
39
40
41
    /// Chat template (template, bos_token, eos_token)
    template: (
        Option<Template<'static, 'static>>,
        Option<String>,
        Option<String>,
    ),
42
43
44
45
46
47
48
49
}

/// Infer shared state
struct Shared {
    /// Batching background Tokio task notifier
    batching_task: Notify,
}

50
51
52
53
54
/// Raise a exception (custom function) used in the chat templates
fn raise_exception(err_text: String) -> Result<String, minijinja::Error> {
    Err(minijinja::Error::new(ErrorKind::SyntaxError, err_text))
}

55
impl Infer {
56
    #[allow(clippy::too_many_arguments)]
57
58
59
    pub(crate) fn new(
        client: ShardedClient,
        validation: Validation,
60
        waiting_served_ratio: f32,
61
        max_batch_prefill_tokens: u32,
62
        max_batch_total_tokens: u32,
63
64
        max_waiting_tokens: usize,
        max_concurrent_requests: usize,
65
        requires_padding: bool,
66
        window_size: Option<u32>,
Nicolas Patry's avatar
Nicolas Patry committed
67
        speculate: u32,
68
        generation_health: Arc<AtomicBool>,
69
        tokenizer_config: HubTokenizerConfig,
70
71
    ) -> Self {
        // Infer shared state
Nicolas Patry's avatar
Nicolas Patry committed
72
        let queue = Queue::new(requires_padding, 16, window_size, speculate);
73
74
75
76
77
78
79
        let shared = Arc::new(Shared {
            batching_task: Notify::new(),
        });

        // Spawn batching background task that contains all the inference logic
        tokio::spawn(batching_task(
            client,
80
            waiting_served_ratio,
81
            max_batch_prefill_tokens,
82
            max_batch_total_tokens,
83
            max_waiting_tokens,
84
            queue.clone(),
85
            shared.clone(),
86
            generation_health,
87
88
89
90
91
        ));

        // Inference limit with a semaphore
        let semaphore = Arc::new(Semaphore::new(max_concurrent_requests));

92
        let template = tokenizer_config.chat_template.map(|t| {
93
            let mut env = Box::new(Environment::new());
94
            let template_str = t.into_boxed_str();
95
            env.add_function("raise_exception", raise_exception);
96
97
98
99
100
            // leaking env and template_str as read-only, static resources for performance.
            Box::leak(env)
                .template_from_str(Box::leak(template_str))
                .unwrap()
        });
101
102
103
104
105
106
107
108
        let eos_token = tokenizer_config
            .eos_token
            .map_or_else(String::new, |t| t)
            .into();
        let bos_token = tokenizer_config
            .bos_token
            .map_or_else(String::new, |t| t)
            .into();
109
110
        Self {
            validation,
111
            queue,
112
113
            shared,
            limit_concurrent_requests: semaphore,
114
            template: (template, eos_token, bos_token),
115
116
117
        }
    }

118
    /// Add a new request to the queue and return a stream of InferStreamResponse
119
    #[instrument(skip_all)]
120
121
122
    pub(crate) async fn generate_stream(
        &self,
        request: GenerateRequest,
123
    ) -> Result<GenerateStreamResponse, InferError> {
124
        // Limit concurrent requests by acquiring a permit from the semaphore
125
126
127
128
129
        let permit = self
            .clone()
            .limit_concurrent_requests
            .try_acquire_owned()
            .map_err(|err| {
130
                metrics::increment_counter!("tgi_request_failure", "err" => "overloaded");
131
132
133
                tracing::error!("{err}");
                err
            })?;
134
135

        // Validate request
136
137
138
139
140
        let valid_request = self.validation.validate(request).await.map_err(|err| {
            metrics::increment_counter!("tgi_request_failure", "err" => "validation");
            tracing::error!("{err}");
            err
        })?;
141
142

        // MPSC channel to communicate with the background batching task
OlivierDehaene's avatar
OlivierDehaene committed
143
        let (response_tx, response_rx) = mpsc::unbounded_channel();
144
        let input_length = valid_request.input_length;
145

146
147
        // Append the request to the queue
        self.queue.append(Entry {
148
149
            request: valid_request,
            response_tx,
150
151
152
            span: Span::current(),
            temp_span: None,
            queue_time: Instant::now(),
153
154
155
            batch_time: None,
        });

156
        // Notify the background task that we have a new entry in the queue that needs
157
158
159
160
        // to be batched
        self.shared.batching_task.notify_one();

        // Return stream
161
162
163
164
165
        Ok((
            permit,
            input_length,
            UnboundedReceiverStream::new(response_rx),
        ))
166
167
    }

168
169
    /// Apply the chat template to the chat request
    #[instrument(skip_all)]
170
171
172
    pub(crate) fn apply_chat_template(&self, messages: Vec<Message>) -> Result<String, InferError> {
        let (template, bos_token, eos_token) = &self.template;
        template
173
174
            .as_ref()
            .ok_or_else(|| InferError::TemplateError(ErrorKind::TemplateNotFound.into()))?
175
176
177
178
179
            .render(ChatTemplateInputs {
                messages,
                eos_token: eos_token.as_deref(),
                bos_token: bos_token.as_deref(),
            })
180
181
182
183
184
185
186
            .map_err(|e| {
                metrics::increment_counter!("tgi_request_failure", "err" => "template");
                tracing::error!("{e}");
                InferError::TemplateError(e)
            })
    }

187
    /// Add a new request to the queue and return a InferResponse
188
    #[instrument(skip_all)]
189
190
191
192
    pub(crate) async fn generate(
        &self,
        request: GenerateRequest,
    ) -> Result<InferResponse, InferError> {
Nicolas Patry's avatar
Nicolas Patry committed
193
194
        let use_top_tokens = request.parameters.top_n_tokens.is_some_and(|x| x > 0);

195
        // Create stream and keep semaphore permit as long as generate lives
196
        let (_permit, _input_length, mut stream) = self.generate_stream(request).await?;
197
198
199
200

        // Return values
        let mut result_prefill = Vec::new();
        let mut result_tokens = Vec::new();
Nicolas Patry's avatar
Nicolas Patry committed
201
        let mut result_top_tokens = Vec::new();
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        let mut result_generated_text = None;
        let mut result_start = None;
        let mut result_queued = None;

        // Iterate on stream
        while let Some(response) = stream.next().await {
            match response? {
                // Add prefill tokens
                InferStreamResponse::Prefill(tokens) => {
                    // Create Token objects
                    // We do that here instead of in the Python code as Rust for loops are faster
                    result_prefill = tokens
                        .ids
                        .into_iter()
                        .zip(tokens.logprobs.into_iter())
                        .zip(tokens.texts.into_iter())
218
                        .map(|((id, logprob), text)| PrefillToken { id, text, logprob })
219
220
221
                        .collect();
                }
                // Push last token
Nicolas Patry's avatar
Nicolas Patry committed
222
223
224
225
                InferStreamResponse::Intermediate { token, top_tokens } => {
                    result_tokens.push(token);
                    result_top_tokens.push(top_tokens);
                }
226
227
228
229
230
231
232
                // Final message
                // Set return values
                InferStreamResponse::End {
                    token,
                    generated_text,
                    start,
                    queued,
Nicolas Patry's avatar
Nicolas Patry committed
233
                    top_tokens,
234
235
                } => {
                    result_tokens.push(token);
Nicolas Patry's avatar
Nicolas Patry committed
236
                    result_top_tokens.push(top_tokens);
237
238
239
240
241
242
243
244
245
246
247
248
249
                    result_generated_text = Some(generated_text);
                    result_start = Some(start);
                    result_queued = Some(queued)
                }
            }
        }

        // Check that we received a `InferStreamResponse::End` message
        if let (Some(generated_text), Some(queued), Some(start)) =
            (result_generated_text, result_queued, result_start)
        {
            Ok(InferResponse {
                prefill: result_prefill,
250
                _input_length,
251
252
253
254
                tokens: result_tokens,
                generated_text,
                queued,
                start,
Nicolas Patry's avatar
Nicolas Patry committed
255
256
257
258
259
                top_tokens: if use_top_tokens {
                    result_top_tokens
                } else {
                    Vec::new()
                },
260
261
            })
        } else {
262
            let err = InferError::IncompleteGeneration;
263
            metrics::increment_counter!("tgi_request_failure", "err" => "incomplete");
264
265
            tracing::error!("{err}");
            Err(err)
266
267
        }
    }
268
269
    /// Add best_of new requests to the queue and return a InferResponse of the sequence with
    /// the highest log probability per token
270
    #[instrument(skip(self, request))]
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
    pub(crate) async fn generate_best_of(
        &self,
        request: GenerateRequest,
        best_of: usize,
    ) -> Result<(InferResponse, Vec<InferResponse>), InferError> {
        // validate  best_of parameter separately
        let best_of = self.validation.validate_best_of(best_of)?;

        // create multiple generate requests
        let mut infer_responses: Vec<InferResponse> =
            try_join_all((0..best_of).map(|_| self.generate(request.clone()))).await?;

        // get the sequence with the highest log probability per token
        let mut max_index = 0;
        let mut max_logprob: f32 = f32::MIN;

        for (i, response) in infer_responses.iter().enumerate() {
            // mean logprobs of the generated tokens
            let sequence_logprob = response
                .tokens
                .iter()
                .map(|token| token.logprob)
                .sum::<f32>()
                / response.tokens.len() as f32;

            // set best sequence
            if sequence_logprob > max_logprob {
                max_index = i;
                max_logprob = sequence_logprob;
            }
        }
        let best_response = infer_responses.remove(max_index);
        Ok((best_response, infer_responses))
    }
305
306
307
308
309
310
}

/// Batching logic
/// Will be launched in a background Tokio task
///
/// Batches requests and sends them to the inference server
311
#[allow(clippy::too_many_arguments)]
312
313
async fn batching_task(
    mut client: ShardedClient,
314
    waiting_served_ratio: f32,
315
    max_batch_prefill_tokens: u32,
316
    max_batch_total_tokens: u32,
317
    max_waiting_tokens: usize,
318
    queue: Queue,
319
    shared: Arc<Shared>,
320
    generation_health: Arc<AtomicBool>,
321
322
323
324
325
326
) {
    // Infinite loop
    loop {
        // Wait for a notification from the Infer struct
        shared.batching_task.notified().await;

327
        // Get the next batch from the queue
328
        // This batch might be smaller than the maximum batch size if there are not enough requests
329
        // waiting in the queue
330
331
332
        while let Some((mut entries, batch, span)) = queue
            .next_batch(None, max_batch_prefill_tokens, max_batch_total_tokens)
            .await
333
        {
334
            let mut cached_batch = prefill(&mut client, batch, &mut entries, &generation_health)
335
336
                .instrument(span)
                .await;
337
338
339
340
341
342
343
            let mut waiting_tokens = 1;

            // We loop until we do not receive any cached batch from the inference server (== until
            // all requests have met their stopping criteria)
            while let Some(batch) = cached_batch {
                // Get current batch info
                let batch_size = batch.size;
344
                let batch_max_tokens = batch.max_tokens;
345
                let mut batches = vec![batch];
346
                metrics::gauge!("tgi_batch_current_size", batch_size as f64);
347
348
349
350
351
352
353
354
355
356
357
                metrics::gauge!("tgi_batch_current_max_tokens", batch_max_tokens as f64);

                let min_size = if waiting_tokens >= max_waiting_tokens {
                    // If we didn't onboard any new requests since >= max_waiting_tokens, we try
                    // to add a new batch even though its size might be small
                    None
                } else {
                    // Minimum batch size
                    Some((batch_size as f32 * waiting_served_ratio).floor() as usize)
                };

358
                let token_budget = max_batch_total_tokens.saturating_sub(batch_max_tokens);
359
360

                // Try to get a new batch
361
362
363
                if let Some((mut new_entries, new_batch, span)) = queue
                    .next_batch(min_size, max_batch_prefill_tokens, token_budget)
                    .await
364
365
366
367
368
369
370
                {
                    // Tracking metrics
                    if min_size.is_some() {
                        metrics::increment_counter!("tgi_batch_concat", "reason" => "backpressure");
                    } else {
                        metrics::increment_counter!("tgi_batch_concat", "reason" => "wait_exceeded");
                    }
371

372
373
374
375
376
377
378
379
380
381
382
383
                    entries.iter_mut().for_each(|(_, entry)| {
                        // Create a new span to add the info that this entry is waiting
                        // because a new batch is being computed
                        let entry_waiting_span = info_span!(parent: &entry.span, "waiting");
                        // Add relationships
                        span.follows_from(&entry_waiting_span);
                        entry_waiting_span.follows_from(&span);
                        // Update entry
                        entry.temp_span = Some(entry_waiting_span);
                    });

                    // Generate one token for this new batch to have the attention past in cache
384
385
386
387
                    let new_cached_batch =
                        prefill(&mut client, new_batch, &mut new_entries, &generation_health)
                            .instrument(span)
                            .await;
388
389
390
391
392
393
                    // Reset waiting counter
                    waiting_tokens = 1;
                    // Extend current batch with the new batch
                    if let Some(new_cached_batch) = new_cached_batch {
                        entries.extend(new_entries);
                        batches.push(new_cached_batch);
394
395
                    }
                }
396

397
398
399
400
401
402
                // Create span for this batch to add context to inference calls
                let next_batch_size = entries.len();
                let next_batch_span =
                    info_span!(parent: None, "batch", batch_size = next_batch_size);
                entries.iter_mut().for_each(|(_, entry)| {
                    // Create a new span to link the batch back to this entry
403
                    let entry_batch_span = info_span!(parent: &entry.span, "infer");
404
405
                    // Add relationships
                    next_batch_span.follows_from(&entry_batch_span);
406
407
408
409
                    entry_batch_span.follows_from(&next_batch_span);
                    // Update entry
                    entry.temp_span = Some(entry_batch_span);
                });
410

411
                cached_batch = decode(&mut client, batches, &mut entries, &generation_health)
412
413
                    .instrument(next_batch_span)
                    .await;
414
415
                waiting_tokens += 1;
            }
416
            metrics::gauge!("tgi_batch_current_size", 0.0);
417
            metrics::gauge!("tgi_batch_current_max_tokens", 0.0);
418
419
420
421
        }
    }
}

422
#[instrument(skip_all)]
423
424
425
async fn prefill(
    client: &mut ShardedClient,
    batch: Batch,
426
    entries: &mut IntMap<u64, Entry>,
427
    generation_health: &Arc<AtomicBool>,
428
) -> Option<CachedBatch> {
429
    let start_time = Instant::now();
430
    let batch_id = batch.id;
431
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "prefill");
432
433

    match client.prefill(batch).await {
434
        Ok((generations, next_batch, timings)) => {
435
436
            // Update health
            generation_health.store(true, Ordering::SeqCst);
437
438

            let start_filtering_time = Instant::now();
439
            // Send generated tokens and filter stopped entries
440
441
442
            filter_send_generations(generations, entries);

            // Filter next batch and remove requests that were stopped
443
            let next_batch = filter_batch(client, next_batch, entries).await;
444

445
446
447
            metrics::histogram!("tgi_batch_forward_duration", timings.forward.as_secs_f64(), "method" => "prefill");
            metrics::histogram!("tgi_batch_decode_duration", timings.decode.as_secs_f64(), "method" => "prefill");
            metrics::histogram!("tgi_batch_filter_duration", start_filtering_time.elapsed().as_secs_f64(), "method" => "prefill");
448
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "prefill");
449
450
451
452
453
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "prefill");
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
454
455
            // Update health
            generation_health.store(false, Ordering::SeqCst);
456
            let _ = client.clear_cache(Some(batch_id)).await;
457
458
459
460
461
462
463
464
465
466
            send_errors(err, entries);
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "prefill");
            None
        }
    }
}

#[instrument(skip_all)]
async fn decode(
    client: &mut ShardedClient,
467
    batches: Vec<CachedBatch>,
468
    entries: &mut IntMap<u64, Entry>,
469
    generation_health: &Arc<AtomicBool>,
470
) -> Option<CachedBatch> {
471
    let start_time = Instant::now();
472
    let batch_ids: Vec<u64> = batches.iter().map(|b| b.id).collect();
473
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "decode");
474
475

    match client.decode(batches).await {
476
        Ok((generations, next_batch, timings)) => {
477
478
            // Update health
            generation_health.store(true, Ordering::SeqCst);
479
480

            let start_filtering_time = Instant::now();
481
            // Send generated tokens and filter stopped entries
482
483
484
            filter_send_generations(generations, entries);

            // Filter next batch and remove requests that were stopped
485
            let next_batch = filter_batch(client, next_batch, entries).await;
486

487
488
489
490
491
492
            if let Some(concat_duration) = timings.concat {
                metrics::histogram!("tgi_batch_concat_duration", concat_duration.as_secs_f64(), "method" => "decode");
            }
            metrics::histogram!("tgi_batch_forward_duration", timings.forward.as_secs_f64(), "method" => "decode");
            metrics::histogram!("tgi_batch_decode_duration", timings.decode.as_secs_f64(), "method" => "decode");
            metrics::histogram!("tgi_batch_filter_duration", start_filtering_time.elapsed().as_secs_f64(), "method" => "decode");
493
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "decode");
494
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "decode");
495
496
497
498
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
499
            generation_health.store(false, Ordering::SeqCst);
500
501
502
            for id in batch_ids {
                let _ = client.clear_cache(Some(id)).await;
            }
503
            send_errors(err, entries);
504
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "decode");
505
506
507
508
509
            None
        }
    }
}

510
511
/// Filter a `batch` and remove all requests not present in `entries`
#[instrument(skip_all)]
512
513
async fn filter_batch(
    client: &mut ShardedClient,
514
    next_batch: Option<CachedBatch>,
515
    entries: &IntMap<u64, Entry>,
516
) -> Option<CachedBatch> {
517
518
519
520
521
522
523
524
525
526
    let mut batch = next_batch?;

    // No need to filter
    if batch.size as usize == entries.len() {
        return Some(batch);
    }

    let id = batch.id;

    // Retain only requests that are still in entries
527
    batch.request_ids.retain(|id| entries.contains_key(id));
528

529
    if batch.request_ids.is_empty() {
530
531
532
533
534
535
536
537
538
        // All requests have been filtered out
        // Next batch is now empty
        // Clear it from the Python shards cache
        // We unwrap here as we need to panic since we cannot recover if this method fails
        client.clear_cache(Some(id)).await.unwrap();
        None
    } else {
        // Filter Python shard cache
        // We unwrap here as we need to panic since we cannot recover if this method fails
539
        client.filter_batch(id, batch.request_ids).await.unwrap()
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
    }
}

/// Send one or multiple `InferStreamResponse` to Infer for all `entries`
/// and filter entries
#[instrument(skip_all)]
fn filter_send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entry>) {
    generations.into_iter().for_each(|generation| {
        let id = generation.request_id;
        // Get entry
        // We can `expect` here as the request id should always be in the entries
        let entry = entries
            .get(&id)
            .expect("ID not found in entries. This is a bug.");

        // Create and enter a span to link this function back to the entry
        let _span = info_span!(parent: entry.temp_span.as_ref().expect("batch_span is None. This is a bug."), "send_generation", generation = ?generation).entered();
        // Send generation responses back to the infer task
        // If the receive an error from the Flume channel, it means that the client dropped the
        // request and we need to stop generating hence why we unwrap_or(true)
        let stopped = send_responses(generation, entry).map_err(|err| {
OlivierDehaene's avatar
OlivierDehaene committed
561
            tracing::error!("Entry response channel error.");
562
563
564
565
566
567
568
569
570
571
572
573
574
            metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
            err
        }).unwrap_or(true);
        if stopped {
            entries.remove(&id).expect("ID not found in entries. This is a bug.");
        }
    });
}

/// Send responses through the `entry` response channel
fn send_responses(
    generation: Generation,
    entry: &Entry,
OlivierDehaene's avatar
OlivierDehaene committed
575
) -> Result<bool, Box<SendError<Result<InferStreamResponse, InferError>>>> {
576
    // Return directly if the channel is disconnected
OlivierDehaene's avatar
OlivierDehaene committed
577
578
    if entry.response_tx.is_closed() {
        metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
579
580
581
        return Ok(true);
    }

582
583
584
585
    let mut stopped = false;

    if let Some(prefill_tokens) = generation.prefill_tokens {
        // Send message
OlivierDehaene's avatar
OlivierDehaene committed
586
587
588
        entry
            .response_tx
            .send(Ok(InferStreamResponse::Prefill(prefill_tokens)))?;
589
590
591
    }

    // Create last Token
Nicolas Patry's avatar
Nicolas Patry committed
592
593
594
595
596
597
    let tokens_ = generation.tokens.expect("Non empty tokens in generation");
    let n = tokens_.ids.len();
    metrics::histogram!("tgi_request_skipped_tokens", (n - 1) as f64);
    let mut iterator = tokens_
        .ids
        .into_iter()
598
599
600
        .zip(tokens_.logprobs)
        .zip(tokens_.texts)
        .zip(tokens_.is_special)
Nicolas Patry's avatar
Nicolas Patry committed
601
602
603
604
605
606
607
608
609
610
        .enumerate()
        .peekable();
    while let Some((i, (((id, logprob), text), special))) = iterator.next() {
        let token = Token {
            id,
            text,
            logprob,
            special,
        };
        let top_tokens = if let Some(top_tokens_) = generation.top_tokens.get(i) {
Nicolas Patry's avatar
Nicolas Patry committed
611
612
            top_tokens_
                .ids
Nicolas Patry's avatar
Nicolas Patry committed
613
614
615
616
617
                .iter()
                .zip(top_tokens_.logprobs.iter())
                .zip(top_tokens_.texts.iter())
                .zip(top_tokens_.is_special.iter())
                .map(|(((&id, &logprob), text), &special)| Token {
Nicolas Patry's avatar
Nicolas Patry committed
618
                    id,
Nicolas Patry's avatar
Nicolas Patry committed
619
                    text: text.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
620
621
                    logprob,
                    special,
Nicolas Patry's avatar
Nicolas Patry committed
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
                })
                .collect()
        } else {
            vec![]
        };
        match (&generation.generated_text, iterator.peek()) {
            (Some(generated_text), None) => {
                // Generation has ended
                stopped = true;
                // Send message
                entry.response_tx.send(Ok(InferStreamResponse::End {
                    token,
                    top_tokens,
                    generated_text: generated_text.clone(),
                    queued: entry.queue_time,
                    start: entry.batch_time.unwrap(),
                }))?;
            }
            _ => {
                // Send message
                entry
                    .response_tx
                    .send(Ok(InferStreamResponse::Intermediate { token, top_tokens }))?;
            }
        }
Nicolas Patry's avatar
Nicolas Patry committed
647
648
    }

649
650
651
    Ok(stopped)
}

652
/// Send errors to Infer for all `entries`
653
654
#[instrument(skip_all)]
fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
655
    entries.drain().for_each(|(_, entry)| {
656
657
658
        // Create and enter a span to link this function back to the entry
        let _send_error_span = info_span!(parent: entry.temp_span.as_ref().expect("batch_span is None. This is a bug."), "send_error").entered();
        let err = InferError::GenerationError(error.to_string());
659
        metrics::increment_counter!("tgi_request_failure", "err" => "generation");
660
661
        tracing::error!("{err}");

662
663
664
        // unwrap_or is valid here as we don't care if the receiver is gone.
        entry
            .response_tx
OlivierDehaene's avatar
OlivierDehaene committed
665
            .send(Err(err))
666
667
668
669
670
671
672
            .unwrap_or(());
    });
}

#[derive(Debug)]
pub(crate) enum InferStreamResponse {
    // Optional first message
Nicolas Patry's avatar
Nicolas Patry committed
673
    Prefill(Tokens),
674
    // Intermediate messages
Nicolas Patry's avatar
Nicolas Patry committed
675
676
677
678
    Intermediate {
        token: Token,
        top_tokens: Vec<Token>,
    },
679
680
681
    // Last message
    End {
        token: Token,
Nicolas Patry's avatar
Nicolas Patry committed
682
        top_tokens: Vec<Token>,
683
684
685
686
687
688
689
690
        generated_text: GeneratedText,
        start: Instant,
        queued: Instant,
    },
}

#[derive(Debug)]
pub(crate) struct InferResponse {
691
692
693
694
    /// input_length is the input as perceived by the rust tokenizer in the
    /// validation pathway. It is redundant with prefill.len() but prefill
    /// has data only if the user asked for it. This will always be filled.
    pub(crate) _input_length: u32,
695
    pub(crate) prefill: Vec<PrefillToken>,
696
697
698
699
    pub(crate) tokens: Vec<Token>,
    pub(crate) generated_text: GeneratedText,
    pub(crate) queued: Instant,
    pub(crate) start: Instant,
Nicolas Patry's avatar
Nicolas Patry committed
700
    pub(crate) top_tokens: Vec<Vec<Token>>,
701
702
703
704
705
706
707
708
709
710
711
712
}

#[derive(Debug, Error)]
pub enum InferError {
    #[error("Request failed during generation: {0}")]
    GenerationError(String),
    #[error("Model is overloaded")]
    Overloaded(#[from] TryAcquireError),
    #[error("Input validation error: {0}")]
    ValidationError(#[from] ValidationError),
    #[error("Incomplete generation")]
    IncompleteGeneration,
713
714
    #[error("Template error: {0}")]
    TemplateError(#[from] minijinja::Error),
715
}
716
717
718
719
720
721
722
723

impl InferError {
    pub(crate) fn error_type(&self) -> &str {
        match self {
            InferError::GenerationError(_) => "generation",
            InferError::Overloaded(_) => "overloaded",
            InferError::ValidationError(_) => "validation",
            InferError::IncompleteGeneration => "incomplete_generation",
724
            InferError::TemplateError(_) => "template_error",
725
726
727
        }
    }
}
728
729
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
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
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
814
815
816
817
818
819
820
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
923
924
925
926
927
928
929

// tests
#[cfg(test)]
mod tests {
    use crate::infer::raise_exception;
    use crate::ChatTemplateInputs;
    use crate::Message;
    use minijinja::Environment;

    #[test]
    fn test_chat_template() {
        let env = Environment::new();

        let source = r#"
        {% for message in messages %}
            {% if message['role'] == 'system' %}
                {% if message['content']%}
                    {{'### System:\n' + message['content']+'\n\n'}}
                {% endif %}
            {% elif message['role'] == 'user' %}
                {{'### User:\n' + message['content']+'\n\n'}}
            {% elif message['role'] == 'assistant' %}
                {{'### Assistant:\n'  + message['content']}}
            {% endif %}
            {% if loop.last and add_generation_prompt %}
                {{ '### Assistant:\n' }}
            {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
                Message {
                    role: "user".to_string(),
                    content: "Hi!".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();

        assert_eq!(
            result,
            r#"### User:
Hi!

### Assistant:
Hello how can I help?### User:
What is Deep Learning?

### Assistant:
magic!"#
        );
    }

    #[test]
    fn test_chat_template_invalid_with_raise() {
        let mut env = Environment::new();
        env.add_function("raise_exception", raise_exception);

        let source = r#"
        {{ bos_token }}
        {% for message in messages %}
        {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
        {% endif %}
        {% if message['role'] == 'user' %}
        {{ '[INST] ' + message['content'] + ' [/INST]' }}
        {% elif message['role'] == 'assistant' %}
        {{ message['content'] + eos_token}}
        {% else %}
        {{ raise_exception('Only user and assistant roles are supported!') }}
        {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
                Message {
                    role: "user".to_string(),
                    content: "Hi!".to_string(),
                },
                Message {
                    role: "user".to_string(),
                    content: "Hi again!".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
        };

        let result = tmpl.unwrap().render(chat_template_inputs); //.err().unwrap();

        match result {
            Ok(_) => panic!("Should have failed"),
            Err(e) => {
                assert_eq!(
                    e.detail().unwrap(),
                    "Conversation roles must alternate user/assistant/user/assistant/..."
                );
            }
        }
    }

    #[test]
    fn test_chat_template_valid_with_raise() {
        let mut env = Environment::new();
        env.add_function("raise_exception", raise_exception);

        let source = r#"
        {{ bos_token }}
        {% for message in messages %}
        {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
        {% endif %}
        {% if message['role'] == 'user' %}
        {{ '[INST] ' + message['content'] + ' [/INST]' }}
        {% elif message['role'] == 'assistant' %}
        {{ message['content'] + eos_token}}
        {% else %}
        {{ raise_exception('Only user and assistant roles are supported!') }}
        {% endif %}
        {% endfor %}"#;

        // trim all the whitespace
        let source = source
            .lines()
            .map(|line| line.trim())
            .collect::<Vec<&str>>()
            .join("");

        let tmpl = env.template_from_str(&source);

        let chat_template_inputs = ChatTemplateInputs {
            messages: vec![
                Message {
                    role: "user".to_string(),
                    content: "Hi!".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
        assert_eq!(result, "[BOS][INST] Hi! [/INST]Hello how can I help?[EOS][INST] What is Deep Learning? [/INST]magic![EOS]");
    }
}