fix(router): fix a possible deadlock in next_batch (#1731)

This commit is contained in:
OlivierDehaene 2024-04-12 10:59:04 +02:00 committed by GitHub
parent 6c2c44b84c
commit c2c98725f8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 0 deletions

View File

@ -200,6 +200,10 @@ impl State {
}
}
// 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;
// Create span for this batch to add context to inference calls
let next_batch_span = info_span!(parent: None, "batch", batch_size = tracing::field::Empty);
next_batch_span.follows_from(&Span::current());