queue.rs 15.4 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::collections::VecDeque;
6
use text_generation_client::{Batch, Request};
7
use tokio::sync::oneshot;
8
use tokio::time::Instant;
9
use tracing::{info_span, instrument, Span};
10
11
12
13
14
15
16

/// Queue entry
#[derive(Debug)]
pub(crate) struct Entry {
    /// Request
    pub request: ValidGenerateRequest,
    /// Response sender to communicate between the Infer struct and the batching_task
17
    pub response_tx: flume::Sender<Result<InferStreamResponse, InferError>>,
18
19
20
21
22
23
    /// 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,
24
25
26
27
28
29
30
31
    /// 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
32
    queue_sender: flume::Sender<QueueCommand>,
33
34
35
}

impl Queue {
36
    pub(crate) fn new(requires_padding: bool) -> Self {
37
        // Create channel
38
        let (queue_sender, queue_receiver) = flume::unbounded();
39
40

        // Launch background queue task
41
        tokio::spawn(queue_task(requires_padding, queue_receiver));
42
43
44
45
46

        Self { queue_sender }
    }

    /// Append an entry to the queue
47
    #[instrument(skip_all)]
48
49
50
    pub(crate) fn append(&self, entry: Entry) {
        // Send append command to the background task managing the state
        // Unwrap is safe here
51
52
53
        self.queue_sender
            .send(QueueCommand::Append(entry, Span::current()))
            .unwrap();
54
55
56
    }

    // Get the next batch
57
    #[instrument(skip(self))]
58
59
60
    pub(crate) async fn next_batch(
        &self,
        min_size: Option<usize>,
61
        token_budget: u32,
62
63
64
65
66
67
68
69
    ) -> 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,
70
                token_budget,
71
                response_sender,
72
                span: Span::current(),
73
74
75
76
77
78
79
80
81
            })
            .unwrap();
        // Await on response channel
        // Unwrap is safe here
        response_receiver.await.unwrap()
    }
}

// Background task responsible of the queue state
82
83
async fn queue_task(requires_padding: bool, receiver: flume::Receiver<QueueCommand>) {
    let mut state = State::new(requires_padding);
84

85
    while let Ok(cmd) = receiver.recv_async().await {
86
        match cmd {
87
88
89
90
            QueueCommand::Append(entry, span) => {
                span.in_scope(|| state.append(entry));
                metrics::increment_gauge!("tgi_queue_size", 1.0);
            }
91
92
            QueueCommand::NextBatch {
                min_size,
93
                token_budget,
94
                response_sender,
95
96
                span,
            } => span.in_scope(|| {
97
                let next_batch = state.next_batch(min_size, token_budget);
98
                response_sender.send(next_batch).unwrap_or(());
99
                metrics::gauge!("tgi_queue_size", state.entries.len() as f64);
100
            }),
101
102
103
104
105
106
107
108
        }
    }
}

/// Queue State
#[derive(Debug)]
struct State {
    /// Queue entries organized in a Vec
109
    entries: VecDeque<(u64, Entry)>,
110
111
112
113
114
115

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

    /// Id of the next batch
    next_batch_id: u64,
116
117
118

    /// Whether the model is using padding
    requires_padding: bool,
119
120
121
}

impl State {
122
    fn new(requires_padding: bool) -> Self {
123
        Self {
124
            entries: VecDeque::with_capacity(128),
125
126
            next_id: 0,
            next_batch_id: 0,
127
            requires_padding,
128
129
130
131
        }
    }

    /// Append an entry to the queue
132
133
134
135
136
137
    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
138
        self.entries.push_back((self.next_id, entry));
139
140
141
142
        self.next_id += 1;
    }

    // Get the next batch
143
    fn next_batch(&mut self, min_size: Option<usize>, token_budget: u32) -> Option<NextBatch> {
144

145
146
147
148
149
150
151
152
153
154
155
        if self.entries.is_empty() {
            return None;
        }

        // Check if we have enough entries
        if let Some(min_size) = min_size {
            if self.entries.len() < min_size {
                return None;
            }
        }

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

160
        let mut batch_requests = Vec::with_capacity(self.entries.len());
161
        let mut batch_entries =
162
            IntMap::with_capacity_and_hasher(self.entries.len(), BuildNoHashHasher::default());
163

164
165
166
167
168
        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
169
170
171
172
173
174
175
176
        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)
            if entry.response_tx.is_disconnected() {
                metrics::increment_counter!("tgi_request_failure", "err" => "dropped");
                continue;
            }

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
            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 {
                prefill_tokens += entry.request.input_length;
            }

            decode_tokens += entry.request.stopping_parameters.max_new_tokens;

            if (prefill_tokens + decode_tokens) > token_budget {
                // Entry is over budget
                // Add it back to the front
                self.entries.push_front((id, entry));
                break;
            }

195
196
197
198
199
200
201
202
203
204
205
206
207
208
            // 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,
                inputs: entry.request.inputs.clone(),
                truncate: entry.request.truncate,
                parameters: Some(entry.request.parameters.clone()),
                stopping_parameters: Some(entry.request.stopping_parameters.clone()),
209
            });
210
211
212
213
214
215
            // Set batch_time
            entry.batch_time = Some(Instant::now());
            // Insert in batch_entries IntMap
            batch_entries.insert(id, entry);
        }

216
        // Empty batch
217
218
219
220
        if batch_requests.is_empty() {
            return None;
        }

221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        // 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
237
238
        let size = batch_requests.len() as u32;
        next_batch_span.record("batch_size", size);
239
240
241
242

        let batch = Batch {
            id: self.next_batch_id,
            requests: batch_requests,
243
            size,
244
            max_tokens: (prefill_tokens + decode_tokens),
245
246
247
248
        };
        // Increment batch id
        self.next_batch_id += 1;

249
        metrics::histogram!("tgi_batch_next_size", batch.size as f64);
250

251
        Some((batch_entries, batch, next_batch_span))
252
253
254
    }
}

255
type NextBatch = (IntMap<u64, Entry>, Batch, Span);
256
257
258

#[derive(Debug)]
enum QueueCommand {
259
    Append(Entry, Span),
260
261
    NextBatch {
        min_size: Option<usize>,
262
        token_budget: u32,
263
        response_sender: oneshot::Sender<Option<NextBatch>>,
264
        span: Span,
265
266
267
268
269
270
271
    },
}

#[cfg(test)]
mod tests {
    use super::*;
    use text_generation_client::{NextTokenChooserParameters, StoppingCriteriaParameters};
272
    use tracing::info_span;
273

274
275
276
277
278
    fn default_entry() -> (
        Entry,
        flume::Receiver<Result<InferStreamResponse, InferError>>,
    ) {
        let (response_tx, receiver_tx) = flume::unbounded();
279

280
        let entry = Entry {
281
282
            request: ValidGenerateRequest {
                inputs: "".to_string(),
283
                input_length: 0,
284
                truncate: 0,
285
286
287
288
                parameters: NextTokenChooserParameters {
                    temperature: 0.0,
                    top_k: 0,
                    top_p: 0.0,
289
                    typical_p: 0.0,
290
291
292
                    do_sample: false,
                    seed: 0,
                    repetition_penalty: 0.0,
293
                    watermark: false,
294
295
                },
                stopping_parameters: StoppingCriteriaParameters {
296
                    ignore_eos_token: false,
297
                    max_new_tokens: 1,
298
299
300
301
                    stop_sequences: vec![],
                },
            },
            response_tx,
302
303
304
            span: info_span!("entry"),
            temp_span: None,
            queue_time: Instant::now(),
305
            batch_time: None,
306
307
        };
        (entry, receiver_tx)
308
309
310
311
    }

    #[test]
    fn test_append() {
312
        let mut state = State::new(false);
313
        let (entry, _guard) = default_entry();
314
315
316
317
318
319
320
321

        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);
322
        let (id, _) = state.entries.remove(0).unwrap();
323
324
325
326
327
        assert_eq!(id, 0);
    }

    #[test]
    fn test_next_batch_empty() {
328
        let mut state = State::new(false);
329
330
331
332
333
334
335

        assert!(state.next_batch(None, 1).is_none());
        assert!(state.next_batch(Some(1), 1).is_none());
    }

    #[test]
    fn test_next_batch_min_size() {
336
        let mut state = State::new(false);
337
338
339
340
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        state.append(entry1);
        state.append(entry2);
341

342
        let (entries, batch, _) = state.next_batch(None, 2).unwrap();
343
344
345
346
347
348
349
350
351
352
353
354
        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);

355
356
        let (entry3, _guard3) = default_entry();
        state.append(entry3);
357
358
359
360
361

        assert!(state.next_batch(Some(2), 2).is_none());

        assert_eq!(state.next_id, 3);
        assert_eq!(state.entries.len(), 1);
362
        let (id, _) = state.entries.remove(0).unwrap();
363
364
365
366
        assert_eq!(id, 2);
    }

    #[test]
367
368
    fn test_next_batch_token_budget() {
        let mut state = State::new(false);
369
370
371
372
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        state.append(entry1);
        state.append(entry2);
373

374
        let (entries, batch, _) = state.next_batch(None, 1).unwrap();
375
376
377
378
379
380
381
382
383
        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);

384
385
        let (entry3, _guard3) = default_entry();
        state.append(entry3);
386

387
        let (entries, batch, _) = state.next_batch(None, 3).unwrap();
388
389
390
391
392
393
394
395
396
397
398
399
400
        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() {
401
        let queue = Queue::new(false);
402
403
        let (entry, _guard) = default_entry();
        queue.append(entry);
404
405
406
407
    }

    #[tokio::test]
    async fn test_queue_next_batch_empty() {
408
        let queue = Queue::new(false);
409
410
411
412
413
414
415

        assert!(queue.next_batch(None, 1).await.is_none());
        assert!(queue.next_batch(Some(1), 1).await.is_none());
    }

    #[tokio::test]
    async fn test_queue_next_batch_min_size() {
416
        let queue = Queue::new(false);
417
418
419
420
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);
421

422
        let (entries, batch, _) = queue.next_batch(None, 2).await.unwrap();
423
424
425
426
427
428
429
430
        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);

431
432
        let (entry3, _guard3) = default_entry();
        queue.append(entry3);
433

434
        // Not enough requests pending
435
        assert!(queue.next_batch(Some(2), 2).await.is_none());
436
437
438
439
440
441
442
443
444
        // Not enough token budget
        assert!(queue.next_batch(Some(1), 0).await.is_none());
        // Ok
        let (entries2, batch2, _) = queue.next_batch(Some(1), 2).await.unwrap();
        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);
445
446
447
    }

    #[tokio::test]
448
449
    async fn test_queue_next_batch_token_budget() {
        let queue = Queue::new(false);
450
451
452
453
        let (entry1, _guard1) = default_entry();
        let (entry2, _guard2) = default_entry();
        queue.append(entry1);
        queue.append(entry2);
454

455
        let (entries, batch, _) = queue.next_batch(None, 1).await.unwrap();
456
457
458
459
460
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key(&0));
        assert_eq!(batch.id, 0);
        assert_eq!(batch.size, 1);

461
462
        let (entry3, _guard3) = default_entry();
        queue.append(entry3);
463

464
        let (entries, batch, _) = queue.next_batch(None, 3).await.unwrap();
465
466
467
468
469
470
        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);
    }
471
472
473

    #[tokio::test]
    async fn test_queue_next_batch_dropped_receiver() {
474
        let queue = Queue::new(false);
475
476
477
478
479
        let (entry, _) = default_entry();
        queue.append(entry);

        assert!(queue.next_batch(None, 1).await.is_none());
    }
480
}