infer.rs 36.9 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
        max_waiting_tokens: usize,
64
        max_batch_size: Option<usize>,
65
        max_concurrent_requests: usize,
66
        requires_padding: bool,
67
        window_size: Option<u32>,
Nicolas Patry's avatar
Nicolas Patry committed
68
        speculate: u32,
69
        generation_health: Arc<AtomicBool>,
70
        tokenizer_config: HubTokenizerConfig,
71
72
    ) -> Self {
        // Infer shared state
Nicolas Patry's avatar
Nicolas Patry committed
73
        let queue = Queue::new(requires_padding, 16, window_size, speculate);
74
75
76
77
78
79
80
        let shared = Arc::new(Shared {
            batching_task: Notify::new(),
        });

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

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

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

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

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

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

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

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

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

170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
    /// Tokenizer the input
    #[instrument(skip_all)]
    pub(crate) async fn tokenize(
        &self,
        request: GenerateRequest,
    ) -> Result<Option<tokenizers::Encoding>, InferError> {
        // Tokenize request
        let inputs = request.inputs;
        let truncate = request.parameters.truncate;
        let encoding = self
            .validation
            .tokenize(inputs, truncate)
            .await
            .map_err(|err| {
                tracing::error!("Tokenization {err}");
                err
            })?;

        // Return Encoding
        Ok(encoding.map(|(encoding, _)| encoding))
    }

192
193
    /// Apply the chat template to the chat request
    #[instrument(skip_all)]
194
195
196
    pub(crate) fn apply_chat_template(&self, messages: Vec<Message>) -> Result<String, InferError> {
        let (template, bos_token, eos_token) = &self.template;
        template
197
198
            .as_ref()
            .ok_or_else(|| InferError::TemplateError(ErrorKind::TemplateNotFound.into()))?
199
200
201
202
            .render(ChatTemplateInputs {
                messages,
                eos_token: eos_token.as_deref(),
                bos_token: bos_token.as_deref(),
203
                add_generation_prompt: true,
204
            })
205
206
207
208
209
210
211
            .map_err(|e| {
                metrics::increment_counter!("tgi_request_failure", "err" => "template");
                tracing::error!("{e}");
                InferError::TemplateError(e)
            })
    }

212
    /// Add a new request to the queue and return a InferResponse
213
    #[instrument(skip_all)]
214
215
216
217
    pub(crate) async fn generate(
        &self,
        request: GenerateRequest,
    ) -> Result<InferResponse, InferError> {
Nicolas Patry's avatar
Nicolas Patry committed
218
219
        let use_top_tokens = request.parameters.top_n_tokens.is_some_and(|x| x > 0);

220
        // Create stream and keep semaphore permit as long as generate lives
221
        let (_permit, _input_length, mut stream) = self.generate_stream(request).await?;
222
223
224
225

        // Return values
        let mut result_prefill = Vec::new();
        let mut result_tokens = Vec::new();
Nicolas Patry's avatar
Nicolas Patry committed
226
        let mut result_top_tokens = Vec::new();
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        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())
243
                        .map(|((id, logprob), text)| PrefillToken { id, text, logprob })
244
245
246
                        .collect();
                }
                // Push last token
Nicolas Patry's avatar
Nicolas Patry committed
247
248
249
250
                InferStreamResponse::Intermediate { token, top_tokens } => {
                    result_tokens.push(token);
                    result_top_tokens.push(top_tokens);
                }
251
252
253
254
255
256
257
                // Final message
                // Set return values
                InferStreamResponse::End {
                    token,
                    generated_text,
                    start,
                    queued,
Nicolas Patry's avatar
Nicolas Patry committed
258
                    top_tokens,
259
260
                } => {
                    result_tokens.push(token);
Nicolas Patry's avatar
Nicolas Patry committed
261
                    result_top_tokens.push(top_tokens);
262
263
264
265
266
267
268
269
270
271
272
273
274
                    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,
275
                _input_length,
276
277
278
279
                tokens: result_tokens,
                generated_text,
                queued,
                start,
Nicolas Patry's avatar
Nicolas Patry committed
280
281
282
283
284
                top_tokens: if use_top_tokens {
                    result_top_tokens
                } else {
                    Vec::new()
                },
285
286
            })
        } else {
287
            let err = InferError::IncompleteGeneration;
288
            metrics::increment_counter!("tgi_request_failure", "err" => "incomplete");
289
290
            tracing::error!("{err}");
            Err(err)
291
292
        }
    }
293
294
    /// Add best_of new requests to the queue and return a InferResponse of the sequence with
    /// the highest log probability per token
295
    #[instrument(skip(self, request))]
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
    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))
    }
330
331
332
333
334
335
}

/// Batching logic
/// Will be launched in a background Tokio task
///
/// Batches requests and sends them to the inference server
336
#[allow(clippy::too_many_arguments)]
337
338
async fn batching_task(
    mut client: ShardedClient,
339
    waiting_served_ratio: f32,
340
    max_batch_prefill_tokens: u32,
341
    max_batch_total_tokens: u32,
342
    max_waiting_tokens: usize,
343
    max_batch_size: Option<usize>,
344
    queue: Queue,
345
    shared: Arc<Shared>,
346
    generation_health: Arc<AtomicBool>,
347
348
349
350
351
352
) {
    // Infinite loop
    loop {
        // Wait for a notification from the Infer struct
        shared.batching_task.notified().await;

353
        // Get the next batch from the queue
354
        // This batch might be smaller than the maximum batch size if there are not enough requests
355
        // waiting in the queue
356
        while let Some((mut entries, batch, span)) = queue
357
358
359
360
361
362
            .next_batch(
                None,
                max_batch_size,
                max_batch_prefill_tokens,
                max_batch_total_tokens,
            )
363
            .await
364
        {
365
            let mut cached_batch = prefill(&mut client, batch, &mut entries, &generation_health)
366
367
                .instrument(span)
                .await;
368
369
370
371
372
373
374
            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;
375
                let batch_max_tokens = batch.max_tokens;
376
                let mut batches = vec![batch];
377
                metrics::gauge!("tgi_batch_current_size", batch_size as f64);
378
379
380
381
382
383
384
385
386
387
388
                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)
                };

389
                let token_budget = max_batch_total_tokens.saturating_sub(batch_max_tokens);
390
                let max_size = max_batch_size.map(|max_size| max_size - batch_size as usize);
391
392

                // Try to get a new batch
393
                if let Some((mut new_entries, new_batch, span)) = queue
394
                    .next_batch(min_size, max_size, max_batch_prefill_tokens, token_budget)
395
                    .await
396
397
398
399
400
401
402
                {
                    // 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");
                    }
403

404
405
406
407
408
409
410
411
412
413
414
415
                    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
416
417
418
419
                    let new_cached_batch =
                        prefill(&mut client, new_batch, &mut new_entries, &generation_health)
                            .instrument(span)
                            .await;
420
421
422
423
424
425
                    // 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);
426
427
                    }
                }
428

429
430
431
432
433
434
                // 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
435
                    let entry_batch_span = info_span!(parent: &entry.span, "infer");
436
437
                    // Add relationships
                    next_batch_span.follows_from(&entry_batch_span);
438
439
440
441
                    entry_batch_span.follows_from(&next_batch_span);
                    // Update entry
                    entry.temp_span = Some(entry_batch_span);
                });
442

443
                cached_batch = decode(&mut client, batches, &mut entries, &generation_health)
444
445
                    .instrument(next_batch_span)
                    .await;
446
447
                waiting_tokens += 1;
            }
448
            metrics::gauge!("tgi_batch_current_size", 0.0);
449
            metrics::gauge!("tgi_batch_current_max_tokens", 0.0);
450
451
452
453
        }
    }
}

454
#[instrument(skip_all)]
455
456
457
async fn prefill(
    client: &mut ShardedClient,
    batch: Batch,
458
    entries: &mut IntMap<u64, Entry>,
459
    generation_health: &Arc<AtomicBool>,
460
) -> Option<CachedBatch> {
461
    let start_time = Instant::now();
462
    let batch_id = batch.id;
463
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "prefill");
464
465

    match client.prefill(batch).await {
466
        Ok((generations, next_batch, timings)) => {
467
468
            // Update health
            generation_health.store(true, Ordering::SeqCst);
469
470

            let start_filtering_time = Instant::now();
471
            // Send generated tokens and filter stopped entries
472
473
474
            filter_send_generations(generations, entries);

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

477
478
479
            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");
480
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "prefill");
481
482
483
484
485
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "prefill");
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
486
487
            // Update health
            generation_health.store(false, Ordering::SeqCst);
488
            let _ = client.clear_cache(Some(batch_id)).await;
489
490
491
492
493
494
495
496
497
498
            send_errors(err, entries);
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "prefill");
            None
        }
    }
}

#[instrument(skip_all)]
async fn decode(
    client: &mut ShardedClient,
499
    batches: Vec<CachedBatch>,
500
    entries: &mut IntMap<u64, Entry>,
501
    generation_health: &Arc<AtomicBool>,
502
) -> Option<CachedBatch> {
503
    let start_time = Instant::now();
504
    let batch_ids: Vec<u64> = batches.iter().map(|b| b.id).collect();
505
    metrics::increment_counter!("tgi_batch_inference_count", "method" => "decode");
506
507

    match client.decode(batches).await {
508
        Ok((generations, next_batch, timings)) => {
509
510
            // Update health
            generation_health.store(true, Ordering::SeqCst);
511
512

            let start_filtering_time = Instant::now();
513
            // Send generated tokens and filter stopped entries
514
515
516
            filter_send_generations(generations, entries);

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

519
520
521
522
523
524
            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");
525
            metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "decode");
526
            metrics::increment_counter!("tgi_batch_inference_success", "method" => "decode");
527
528
529
530
            next_batch
        }
        // If we have an error, we discard the whole batch
        Err(err) => {
531
            generation_health.store(false, Ordering::SeqCst);
532
533
534
            for id in batch_ids {
                let _ = client.clear_cache(Some(id)).await;
            }
535
            send_errors(err, entries);
536
            metrics::increment_counter!("tgi_batch_inference_failure", "method" => "decode");
537
538
539
540
541
            None
        }
    }
}

542
543
/// Filter a `batch` and remove all requests not present in `entries`
#[instrument(skip_all)]
544
545
async fn filter_batch(
    client: &mut ShardedClient,
546
    next_batch: Option<CachedBatch>,
547
    entries: &IntMap<u64, Entry>,
548
) -> Option<CachedBatch> {
549
550
551
552
553
554
555
556
557
558
    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
559
    batch.request_ids.retain(|id| entries.contains_key(id));
560

561
    if batch.request_ids.is_empty() {
562
563
564
565
566
567
568
569
570
        // 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
571
        client.filter_batch(id, batch.request_ids).await.unwrap()
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
    }
}

/// 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
593
            tracing::error!("Entry response channel error.");
594
595
596
597
598
599
600
601
602
603
604
605
606
            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
607
) -> Result<bool, Box<SendError<Result<InferStreamResponse, InferError>>>> {
608
    // Return directly if the channel is disconnected
OlivierDehaene's avatar
OlivierDehaene committed
609
610
    if entry.response_tx.is_closed() {
        metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
611
612
613
        return Ok(true);
    }

614
615
616
617
    let mut stopped = false;

    if let Some(prefill_tokens) = generation.prefill_tokens {
        // Send message
OlivierDehaene's avatar
OlivierDehaene committed
618
619
620
        entry
            .response_tx
            .send(Ok(InferStreamResponse::Prefill(prefill_tokens)))?;
621
622
623
    }

    // Create last Token
Nicolas Patry's avatar
Nicolas Patry committed
624
625
626
627
628
629
    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()
630
631
632
        .zip(tokens_.logprobs)
        .zip(tokens_.texts)
        .zip(tokens_.is_special)
Nicolas Patry's avatar
Nicolas Patry committed
633
634
635
636
637
638
639
640
641
642
        .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
643
644
            top_tokens_
                .ids
Nicolas Patry's avatar
Nicolas Patry committed
645
646
647
648
649
                .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
650
                    id,
Nicolas Patry's avatar
Nicolas Patry committed
651
                    text: text.to_string(),
Nicolas Patry's avatar
Nicolas Patry committed
652
653
                    logprob,
                    special,
Nicolas Patry's avatar
Nicolas Patry committed
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
                })
                .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
679
680
    }

681
682
683
    Ok(stopped)
}

684
/// Send errors to Infer for all `entries`
685
686
#[instrument(skip_all)]
fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
687
    entries.drain().for_each(|(_, entry)| {
688
689
690
        // 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());
691
        metrics::increment_counter!("tgi_request_failure", "err" => "generation");
692
693
        tracing::error!("{err}");

694
695
696
        // unwrap_or is valid here as we don't care if the receiver is gone.
        entry
            .response_tx
OlivierDehaene's avatar
OlivierDehaene committed
697
            .send(Err(err))
698
699
700
701
702
703
704
            .unwrap_or(());
    });
}

#[derive(Debug)]
pub(crate) enum InferStreamResponse {
    // Optional first message
Nicolas Patry's avatar
Nicolas Patry committed
705
    Prefill(Tokens),
706
    // Intermediate messages
Nicolas Patry's avatar
Nicolas Patry committed
707
708
709
710
    Intermediate {
        token: Token,
        top_tokens: Vec<Token>,
    },
711
712
713
    // Last message
    End {
        token: Token,
Nicolas Patry's avatar
Nicolas Patry committed
714
        top_tokens: Vec<Token>,
715
716
717
718
719
720
721
722
        generated_text: GeneratedText,
        start: Instant,
        queued: Instant,
    },
}

#[derive(Debug)]
pub(crate) struct InferResponse {
723
724
725
726
    /// 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,
727
    pub(crate) prefill: Vec<PrefillToken>,
728
729
730
731
    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
732
    pub(crate) top_tokens: Vec<Vec<Token>>,
733
734
735
736
737
738
739
740
741
742
743
744
}

#[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,
745
746
    #[error("Template error: {0}")]
    TemplateError(#[from] minijinja::Error),
747
}
748
749
750
751
752
753
754
755

impl InferError {
    pub(crate) fn error_type(&self) -> &str {
        match self {
            InferError::GenerationError(_) => "generation",
            InferError::Overloaded(_) => "overloaded",
            InferError::ValidationError(_) => "validation",
            InferError::IncompleteGeneration => "incomplete_generation",
756
            InferError::TemplateError(_) => "template_error",
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

// 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(),
803
                    name: None,
804
805
806
807
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
808
                    name: None,
809
810
811
812
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
813
                    name: None,
814
815
816
817
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
818
                    name: None,
819
820
821
822
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
823
            add_generation_prompt: true,
824
825
826
827
828
829
        };

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

        assert_eq!(
            result,
830
            "### User:\nHi!\n\n### Assistant:\nHello how can I help?### User:\nWhat is Deep Learning?\n\n### Assistant:\nmagic!### Assistant:\n"
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
        );
    }

    #[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(),
868
                    name: None,
869
870
871
872
                },
                Message {
                    role: "user".to_string(),
                    content: "Hi again!".to_string(),
873
                    name: None,
874
875
876
877
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
878
                    name: None,
879
880
881
882
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
883
                    name: None,
884
885
886
887
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
888
                    name: None,
889
890
891
892
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
893
            add_generation_prompt: true,
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
930
931
932
933
934
935
936
937
938
939
940
941
942
        };

        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(),
943
                    name: None,
944
945
946
947
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
948
                    name: None,
949
950
951
952
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
953
                    name: None,
954
955
956
957
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
958
                    name: None,
959
960
961
962
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
963
            add_generation_prompt: true,
964
965
966
967
968
        };

        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]");
    }
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996

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

        let source = r#"
        {% for message in messages %}
        {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
        {% endfor %}
        {% if add_generation_prompt %}
            {{ '<|im_start|>assistant\n' }}
        {% endif %}"#;

        // 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(),
997
                    name: None,
998
999
1000
1001
                },
                Message {
                    role: "assistant".to_string(),
                    content: "Hello how can I help?".to_string(),
1002
                    name: None,
1003
1004
1005
1006
                },
                Message {
                    role: "user".to_string(),
                    content: "What is Deep Learning?".to_string(),
1007
                    name: None,
1008
1009
1010
1011
                },
                Message {
                    role: "assistant".to_string(),
                    content: "magic!".to_string(),
1012
                    name: None,
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
                },
            ],
            bos_token: Some("[BOS]"),
            eos_token: Some("[EOS]"),
            add_generation_prompt: true,
        };

        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
        assert_eq!(result, "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\nHello how can I help?<|im_end|>\n<|im_start|>user\nWhat is Deep Learning?<|im_end|>\n<|im_start|>assistant\nmagic!<|im_end|>\n<|im_start|>assistant\n");
    }
1023
}