queue.rs 20.6 KB
Newer Older
1
2
3
4
use crate::infer::InferError;
use crate::infer::InferStreamResponse;
use crate::validation::ValidGenerateRequest;
use nohash_hasher::{BuildNoHashHasher, IntMap};
5
use std::cmp::min;
6
use std::collections::VecDeque;
7
use text_generation_client::{Batch, Request};
OlivierDehaene's avatar
OlivierDehaene committed
8
use tokio::sync::{mpsc, oneshot};
9
use tokio::time::Instant;
10
use tracing::{info_span, instrument, Span};
11
12
13
14
15
16
17

/// Queue entry
#[derive(Debug)]
pub(crate) struct Entry {
    /// Request
    pub request: ValidGenerateRequest,
    /// Response sender to communicate between the Infer struct and the batching_task
OlivierDehaene's avatar
OlivierDehaene committed
18
    pub response_tx: mpsc::UnboundedSender<Result<InferStreamResponse, InferError>>,
19
20
21
22
23
24
    /// Span that will live as long as entry
    pub span: Span,
    /// Temporary span used as a guard when logging inference, wait times...
    pub temp_span: Option<Span>,
    /// Instant when this entry was queued
    pub queue_time: Instant,
25
26
27
28
29
30
31
32
    /// Instant when this entry was added to a batch
    pub batch_time: Option<Instant>,
}

/// Request Queue
#[derive(Debug, Clone)]
pub(crate) struct Queue {
    /// Channel to communicate with the background queue task
OlivierDehaene's avatar
OlivierDehaene committed
33
    queue_sender: mpsc::UnboundedSender<QueueCommand>,
34
35
36
}

impl Queue {
Nicolas Patry's avatar
Nicolas Patry committed
37
38
39
40
41
42
    pub(crate) fn new(
        requires_padding: bool,
        block_size: u32,
        window_size: Option<u32>,
        speculate: u32,
    ) -> Self {
43
        // Create channel
OlivierDehaene's avatar
OlivierDehaene committed
44
        let (queue_sender, queue_receiver) = mpsc::unbounded_channel();
45
46

        // Launch background queue task
47
48
49
50
        tokio::spawn(queue_task(
            requires_padding,
            block_size,
            window_size,
Nicolas Patry's avatar
Nicolas Patry committed
51
            speculate,
52
53
            queue_receiver,
        ));
54
55
56
57
58

        Self { queue_sender }
    }

    /// Append an entry to the queue
59
    #[instrument(skip_all)]
60
61
62
    pub(crate) fn append(&self, entry: Entry) {
        // Send append command to the background task managing the state
        // Unwrap is safe here
63
        self.queue_sender
64
            .send(QueueCommand::Append(Box::new(entry), Span::current()))
65
            .unwrap();
66
67
68
    }

    // Get the next batch
69
    #[instrument(skip(self))]
70
71
72
    pub(crate) async fn next_batch(
        &self,
        min_size: Option<usize>,
73
        max_size: Option<usize>,
74
        prefill_token_budget: u32,
75
        token_budget: u32,
76
77
78
79
80
81
82
83
    ) -> Option<NextBatch> {
        // Create response channel
        let (response_sender, response_receiver) = oneshot::channel();
        // Send next batch command to the background task managing the state
        // Unwrap is safe here
        self.queue_sender
            .send(QueueCommand::NextBatch {
                min_size,
84
                max_size,
85
                prefill_token_budget,
86
                token_budget,
87
                response_sender,
88
                span: Span::current(),
89
90
91
92
93
94
95
96
97
            })
            .unwrap();
        // Await on response channel
        // Unwrap is safe here
        response_receiver.await.unwrap()
    }
}

// Background task responsible of the queue state
98
99
100
async fn queue_task(
    requires_padding: bool,
    block_size: u32,
101
    window_size: Option<u32>,
Nicolas Patry's avatar
Nicolas Patry committed
102
    speculate: u32,
OlivierDehaene's avatar
OlivierDehaene committed
103
    mut receiver: mpsc::UnboundedReceiver<QueueCommand>,
104
) {
Nicolas Patry's avatar
Nicolas Patry committed
105
    let mut state = State::new(requires_padding, block_size, window_size, speculate);
106

OlivierDehaene's avatar
OlivierDehaene committed
107
    while let Some(cmd) = receiver.recv().await {
108
        match cmd {
109
            QueueCommand::Append(entry, span) => {
110
                span.in_scope(|| state.append(*entry));
111
112
                metrics::increment_gauge!("tgi_queue_size", 1.0);
            }
113
114
            QueueCommand::NextBatch {
                min_size,
115
                max_size,
116
                prefill_token_budget,
117
                token_budget,
118
                response_sender,
119
120
                span,
            } => span.in_scope(|| {
121
122
                let next_batch =
                    state.next_batch(min_size, max_size, prefill_token_budget, token_budget);
123
                response_sender.send(next_batch).unwrap();
124
                metrics::gauge!("tgi_queue_size", state.entries.len() as f64);
125
            }),
126
127
128
129
130
131
132
133
        }
    }
}

/// Queue State
#[derive(Debug)]
struct State {
    /// Queue entries organized in a Vec
134
    entries: VecDeque<(u64, Entry)>,
135
136
137
138
139
140

    /// Id of the next entry
    next_id: u64,

    /// Id of the next batch
    next_batch_id: u64,
141
142
143

    /// Whether the model is using padding
    requires_padding: bool,
144
145
146

    /// Paged Attention block size
    block_size: u32,
147
148
149

    /// Sliding window
    window_size: Option<u32>,
Nicolas Patry's avatar
Nicolas Patry committed
150
151
152

    /// Speculation amount
    speculate: u32,
153
154
155
}

impl State {
Nicolas Patry's avatar
Nicolas Patry committed
156
157
158
159
160
161
    fn new(
        requires_padding: bool,
        block_size: u32,
        window_size: Option<u32>,
        speculate: u32,
    ) -> Self {
162
        Self {
163
            entries: VecDeque::with_capacity(128),
164
165
            next_id: 0,
            next_batch_id: 0,
166
            requires_padding,
167
            block_size,
168
            window_size,
Nicolas Patry's avatar
Nicolas Patry committed
169
            speculate,
170
171
172
173
        }
    }

    /// Append an entry to the queue
174
175
176
177
178
179
    fn append(&mut self, mut entry: Entry) {
        // Create a span that will live as long as the entry is in the queue waiting to be batched
        let queue_span = info_span!(parent: &entry.span, "queued");
        entry.temp_span = Some(queue_span);

        // Push entry in the queue
180
        self.entries.push_back((self.next_id, entry));
181
182
183
184
        self.next_id += 1;
    }

    // Get the next batch
185
186
187
    fn next_batch(
        &mut self,
        min_size: Option<usize>,
188
        max_size: Option<usize>,
189
190
191
        prefill_token_budget: u32,
        token_budget: u32,
    ) -> Option<NextBatch> {
192
        if self.entries.is_empty() {
193
            tracing::debug!("No queue");
194
195
196
197
198
199
            return None;
        }

        // Check if we have enough entries
        if let Some(min_size) = min_size {
            if self.entries.len() < min_size {
200
                tracing::debug!("Not enough entries");
201
202
203
204
                return None;
            }
        }

205
206
207
208
        // Pad prefill_token_budget to be a multiple of block size
        let prefill_token_budget =
            ((prefill_token_budget + self.block_size - 1) / self.block_size) * self.block_size;

209
        // Create span for this batch to add context to inference calls
210
        let next_batch_span = info_span!(parent: None, "batch", batch_size = tracing::field::Empty);
211
212
        next_batch_span.follows_from(&Span::current());

213
        let mut batch_requests = Vec::with_capacity(self.entries.len());
214
        let mut batch_entries =
215
            IntMap::with_capacity_and_hasher(self.entries.len(), BuildNoHashHasher::default());
216

217
218
219
220
221
        let mut max_input_length = 0;
        let mut prefill_tokens: u32 = 0;
        let mut decode_tokens: u32 = 0;

        // Pop entries starting from the front of the queue
222
223
224
        while let Some((id, mut entry)) = self.entries.pop_front() {
            // Filter entries where the response receiver was dropped (== entries where the request
            // was dropped by the client)
OlivierDehaene's avatar
OlivierDehaene committed
225
            if entry.response_tx.is_closed() {
226
                metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
227
                tracing::debug!("Dropping entry");
228
229
230
                continue;
            }

231
232
233
234
235
236
            if self.requires_padding {
                // We pad to max input length in the Python shards
                // We need to take these padding tokens into the equation
                max_input_length = max_input_length.max(entry.request.input_length);
                prefill_tokens = (batch_requests.len() + 1) as u32 * max_input_length
            } else {
237
238
239
240
                // pad to block size
                prefill_tokens += ((entry.request.input_length + self.block_size - 1)
                    / self.block_size)
                    * self.block_size;
241
242
            }

243
244
245
            if self.requires_padding {
                decode_tokens += entry.request.stopping_parameters.max_new_tokens;
            } else {
246
247
248
249
250
251
252
253
                let max_new_tokens = match self.window_size {
                    None => entry.request.stopping_parameters.max_new_tokens,
                    Some(window_size) => min(
                        window_size.saturating_sub(entry.request.input_length),
                        entry.request.stopping_parameters.max_new_tokens,
                    ),
                };

254
255
                // pad to block size
                decode_tokens +=
256
                    ((max_new_tokens + self.block_size - 1) / self.block_size) * self.block_size;
257
            }
258

259
            if prefill_tokens > prefill_token_budget
Nicolas Patry's avatar
Nicolas Patry committed
260
                || (prefill_tokens + decode_tokens + self.speculate) > token_budget
261
            {
262
263
                // Entry is over budget
                // Add it back to the front
264
                tracing::debug!("Over budget: prefill_tokens={prefill_tokens} > {prefill_token_budget} || {prefill_tokens} + {decode_tokens} + {} > {token_budget}", self.speculate);
265
266
267
268
                self.entries.push_front((id, entry));
                break;
            }

269
            tracing::debug!("Accepting entry");
270
271
272
273
274
275
276
277
278
279
            // Create a new span to link the batch back to this entry
            let entry_batch_span = info_span!(parent: &entry.span, "infer");
            // Add relationships
            next_batch_span.follows_from(&entry_batch_span);
            entry_batch_span.follows_from(&next_batch_span);
            // Update entry
            entry.temp_span = Some(entry_batch_span);

            batch_requests.push(Request {
                id,
280
                prefill_logprobs: entry.request.decoder_input_details,
281
282
283
284
                inputs: entry.request.inputs.clone(),
                truncate: entry.request.truncate,
                parameters: Some(entry.request.parameters.clone()),
                stopping_parameters: Some(entry.request.stopping_parameters.clone()),
Nicolas Patry's avatar
Nicolas Patry committed
285
                top_n_tokens: entry.request.top_n_tokens,
286
            });
287
288
289
290
            // Set batch_time
            entry.batch_time = Some(Instant::now());
            // Insert in batch_entries IntMap
            batch_entries.insert(id, entry);
291
292
293
294
295

            // Check if max_size
            if Some(batch_requests.len()) == max_size {
                break;
            }
296
297
        }

298
        // Empty batch
299
        if batch_requests.is_empty() {
300
            tracing::debug!("Filterered out all entries");
301
302
303
            return None;
        }

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        // Check if our batch is big enough
        if let Some(min_size) = min_size {
            // Batch is too small
            if batch_requests.len() < min_size {
                // Add back entries to the queue in the correct order
                for r in batch_requests.into_iter().rev() {
                    let id = r.id;
                    let entry = batch_entries.remove(&id).unwrap();
                    self.entries.push_front((id, entry));
                }

                return None;
            }
        }

        // Final batch size
320
321
        let size = batch_requests.len() as u32;
        next_batch_span.record("batch_size", size);
322
323
324
325

        let batch = Batch {
            id: self.next_batch_id,
            requests: batch_requests,
326
            size,
327
            max_tokens: (prefill_tokens + decode_tokens),
328
329
330
331
        };
        // Increment batch id
        self.next_batch_id += 1;

332
        metrics::histogram!("tgi_batch_next_size", batch.size as f64);
333

334
        Some((batch_entries, batch, next_batch_span))
335
336
337
    }
}

338
type NextBatch = (IntMap<u64, Entry>, Batch, Span);
339
340
341

#[derive(Debug)]
enum QueueCommand {
342
    Append(Box<Entry>, Span),
343
344
    NextBatch {
        min_size: Option<usize>,
345
        max_size: Option<usize>,
346
        prefill_token_budget: u32,
347
        token_budget: u32,
348
        response_sender: oneshot::Sender<Option<NextBatch>>,
349
        span: Span,
350
351
352
353
354
355
    },
}

#[cfg(test)]
mod tests {
    use super::*;
drbh's avatar
drbh committed
356
357
358
    use text_generation_client::{
        GrammarType as ProtoGrammarType, NextTokenChooserParameters, StoppingCriteriaParameters,
    };
359
    use tracing::info_span;
360

361
362
    fn default_entry() -> (
        Entry,
OlivierDehaene's avatar
OlivierDehaene committed
363
        mpsc::UnboundedReceiver<Result<InferStreamResponse, InferError>>,
364
    ) {
OlivierDehaene's avatar
OlivierDehaene committed
365
        let (response_tx, receiver_tx) = mpsc::unbounded_channel();
366

367
        let entry = Entry {
368
            request: ValidGenerateRequest {
drbh's avatar
drbh committed
369
                inputs: String::new(),
370
                input_length: 0,
371
                truncate: 0,
372
                decoder_input_details: false,
373
374
375
376
                parameters: NextTokenChooserParameters {
                    temperature: 0.0,
                    top_k: 0,
                    top_p: 0.0,
377
                    typical_p: 0.0,
378
379
380
                    do_sample: false,
                    seed: 0,
                    repetition_penalty: 0.0,
381
                    frequency_penalty: 0.0,
382
                    watermark: false,
drbh's avatar
drbh committed
383
384
                    grammar: String::new(),
                    grammar_type: ProtoGrammarType::None as i32,
385
386
                },
                stopping_parameters: StoppingCriteriaParameters {
387
                    ignore_eos_token: false,
388
                    max_new_tokens: 1,
389
390
                    stop_sequences: vec![],
                },
Nicolas Patry's avatar
Nicolas Patry committed
391
                top_n_tokens: 0,
392
393
            },
            response_tx,
394
395
396
            span: info_span!("entry"),
            temp_span: None,
            queue_time: Instant::now(),
397
            batch_time: None,
398
399
        };
        (entry, receiver_tx)
400
401
402
403
    }

    #[test]
    fn test_append() {
Nicolas Patry's avatar
Nicolas Patry committed
404
        let mut state = State::new(false, 1, None, 0);
405
        let (entry, _guard) = default_entry();
406
407
408
409
410
411
412
413

        assert_eq!(state.next_id, 0);
        assert_eq!(state.entries.len(), 0);

        state.append(entry);

        assert_eq!(state.next_id, 1);
        assert_eq!(state.entries.len(), 1);
414
        let (id, _) = state.entries.remove(0).unwrap();
415
416
417
418
419
        assert_eq!(id, 0);
    }

    #[test]
    fn test_next_batch_empty() {
Nicolas Patry's avatar
Nicolas Patry committed
420
        let mut state = State::new(false, 1, None, 0);
421

422
423
        assert!(state.next_batch(None, None, 1, 1).is_none());
        assert!(state.next_batch(Some(1), None, 1, 1).is_none());
424
425
426
427
    }

    #[test]
    fn test_next_batch_min_size() {
Nicolas Patry's avatar
Nicolas Patry committed
428
        let mut state = State::new(false, 1, None, 0);
429
430
431
432
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        state.append(entry1);
        state.append(entry2);
433

434
        let (entries, batch, _) = state.next_batch(None, None, 2, 2).unwrap();
435
436
437
438
439
440
441
442
443
444
445
446
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key(&0));
        assert!(entries.contains_key(&1));
        assert!(entries.get(&0).unwrap().batch_time.is_some());
        assert!(entries.get(&1).unwrap().batch_time.is_some());
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 2);

        assert_eq!(state.next_id, 2);
        assert_eq!(state.entries.len(), 0);
        assert_eq!(state.next_batch_id, 1);

447
448
        let (entry3, _guard3) = default_entry();
        state.append(entry3);
449

450
        assert!(state.next_batch(Some(2), None, 2, 2).is_none());
451
452
453

        assert_eq!(state.next_id, 3);
        assert_eq!(state.entries.len(), 1);
454
        let (id, _) = state.entries.remove(0).unwrap();
455
456
457
        assert_eq!(id, 2);
    }

458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
    #[test]
    fn test_next_batch_max_size() {
        let mut state = State::new(false, 1, None, 0);
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        state.append(entry1);
        state.append(entry2);

        let (entries, batch, _) = state.next_batch(None, Some(1), 2, 2).unwrap();
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key(&0));
        assert!(entries.get(&0).unwrap().batch_time.is_some());
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 1);

        assert_eq!(state.next_id, 2);
        assert_eq!(state.entries.len(), 1);
        assert_eq!(state.next_batch_id, 1);
    }

478
    #[test]
479
    fn test_next_batch_token_budget() {
Nicolas Patry's avatar
Nicolas Patry committed
480
        let mut state = State::new(false, 1, None, 0);
481
482
483
484
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        state.append(entry1);
        state.append(entry2);
485

486
        let (entries, batch, _) = state.next_batch(None, None, 1, 1).unwrap();
487
488
489
490
491
492
493
494
495
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key(&0));
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 1);

        assert_eq!(state.next_id, 2);
        assert_eq!(state.entries.len(), 1);
        assert_eq!(state.next_batch_id, 1);

496
497
        let (entry3, _guard3) = default_entry();
        state.append(entry3);
498

499
        let (entries, batch, _) = state.next_batch(None, None, 3, 3).unwrap();
500
501
502
503
504
505
506
507
508
509
510
511
512
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key(&1));
        assert!(entries.contains_key(&2));
        assert_eq!(batch.id, 1);
        assert_eq!(batch.size, 2);

        assert_eq!(state.next_id, 3);
        assert_eq!(state.entries.len(), 0);
        assert_eq!(state.next_batch_id, 2);
    }

    #[tokio::test]
    async fn test_queue_append() {
Nicolas Patry's avatar
Nicolas Patry committed
513
        let queue = Queue::new(false, 1, None, 0);
514
515
        let (entry, _guard) = default_entry();
        queue.append(entry);
516
517
518
519
    }

    #[tokio::test]
    async fn test_queue_next_batch_empty() {
Nicolas Patry's avatar
Nicolas Patry committed
520
        let queue = Queue::new(false, 1, None, 0);
521

522
523
        assert!(queue.next_batch(None, None, 1, 1).await.is_none());
        assert!(queue.next_batch(Some(1), None, 1, 1).await.is_none());
524
525
526
527
    }

    #[tokio::test]
    async fn test_queue_next_batch_min_size() {
Nicolas Patry's avatar
Nicolas Patry committed
528
        let queue = Queue::new(false, 1, None, 0);
529
530
531
532
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);
533

534
        let (entries, batch, _) = queue.next_batch(None, None, 2, 2).await.unwrap();
535
536
537
538
539
540
541
542
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key(&0));
        assert!(entries.contains_key(&1));
        assert!(entries.get(&0).unwrap().batch_time.is_some());
        assert!(entries.get(&1).unwrap().batch_time.is_some());
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 2);

543
544
        let (entry3, _guard3) = default_entry();
        queue.append(entry3);
545

546
        // Not enough requests pending
547
        assert!(queue.next_batch(Some(2), None, 2, 2).await.is_none());
548
        // Not enough token budget
549
        assert!(queue.next_batch(Some(1), None, 0, 0).await.is_none());
550
        // Ok
551
        let (entries2, batch2, _) = queue.next_batch(Some(1), None, 2, 2).await.unwrap();
552
553
554
555
556
        assert_eq!(entries2.len(), 1);
        assert!(entries2.contains_key(&2));
        assert!(entries2.get(&2).unwrap().batch_time.is_some());
        assert_eq!(batch2.id, 1);
        assert_eq!(batch2.size, 1);
557
558
    }

559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
    #[tokio::test]
    async fn test_queue_next_batch_max_size() {
        let queue = Queue::new(false, 1, None, 0);
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);

        let (entries, batch, _) = queue.next_batch(None, Some(1), 2, 2).await.unwrap();
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key(&0));
        assert!(entries.get(&0).unwrap().batch_time.is_some());
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 1);
    }

575
    #[tokio::test]
576
    async fn test_queue_next_batch_token_budget() {
Nicolas Patry's avatar
Nicolas Patry committed
577
        let queue = Queue::new(false, 1, None, 0);
578
579
580
581
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);
582

583
        let (entries, batch, _) = queue.next_batch(None, None, 1, 1).await.unwrap();
584
585
586
587
588
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key(&0));
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 1);

589
590
        let (entry3, _guard3) = default_entry();
        queue.append(entry3);
591

592
        let (entries, batch, _) = queue.next_batch(None, None, 3, 3).await.unwrap();
593
594
595
596
597
598
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key(&1));
        assert!(entries.contains_key(&2));
        assert_eq!(batch.id, 1);
        assert_eq!(batch.size, 2);
    }
599

Nicolas Patry's avatar
Nicolas Patry committed
600
601
602
603
604
605
606
607
608
    #[tokio::test]
    async fn test_queue_next_batch_token_speculate() {
        let queue = Queue::new(false, 1, None, 2);
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);

        // Budget of 1 is not enough
609
        assert!(queue.next_batch(None, None, 1, 1).await.is_none());
Nicolas Patry's avatar
Nicolas Patry committed
610

611
        let (entries, batch, _) = queue.next_batch(None, None, 6, 6).await.unwrap();
Nicolas Patry's avatar
Nicolas Patry committed
612
613
614
615
616
617
618
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key(&0));
        assert!(entries.contains_key(&1));
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 2);
    }

619
620
    #[tokio::test]
    async fn test_queue_next_batch_dropped_receiver() {
Nicolas Patry's avatar
Nicolas Patry committed
621
        let queue = Queue::new(false, 1, None, 0);
622
623
624
        let (entry, _) = default_entry();
        queue.append(entry);

625
        assert!(queue.next_batch(None, None, 1, 1).await.is_none());
626
    }
627
}