huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 19:34:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

improved comments

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-02 13:01:59 +01:00
parent 3f963d8a00
commit 0c93da571b
1 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
router/src/db.rs
View File

@ -4,6 +4,7 @@ use crate::infer::InferStreamResponse;
use crate::validation::ValidGenerateRequest; use crate::validation::ValidGenerateRequest;
use nohash_hasher::{BuildNoHashHasher, IntMap}; use nohash_hasher::{BuildNoHashHasher, IntMap};
use std::cmp::min; use std::cmp::min;
use async_stream::reexport::next;
use text_generation_client::{Batch, Request}; use text_generation_client::{Batch, Request};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit}; use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit};
@ -121,25 +122,24 @@ impl State {
// Get the next batch // Get the next batch
fn next_batch(&mut self, min_size: Option<usize>, max_size: usize) -> Option<NextBatch> { fn next_batch(&mut self, min_size: Option<usize>, max_size: usize) -> Option<NextBatch> {
// Check if we have enough entries in DB by comparing next batch id and current id if self.entries.is_empty() {
return None;
}
// Check if we have enough entries in DB
if let Some(min_size) = min_size { if let Some(min_size) = min_size {
if self.entries.len() < min_size { if self.entries.len() < min_size {
return None; return None;
} }
} }
// If both ids are equal, the DB is empty
if self.entries.is_empty() {
return None;
}
let next_batch_size = min(self.entries.len(), max_size); let next_batch_size = min(self.entries.len(), max_size);
// Iterates for max_size over the BTreemap starting from next_batch_start_id let mut batch_requests = Vec::with_capacity(next_batch_size);
let mut batch_requests = Vec::new();
let mut batch_entries = let mut batch_entries =
IntMap::with_capacity_and_hasher(next_batch_size, BuildNoHashHasher::default()); IntMap::with_capacity_and_hasher(next_batch_size, BuildNoHashHasher::default());
// Drain next_batch_size entries
self.entries self.entries
.drain(..next_batch_size) .drain(..next_batch_size)
.for_each(|(id, mut entry)| { .for_each(|(id, mut entry)| {
@ -152,7 +152,7 @@ impl State {
}); });
// Set batch_time // Set batch_time
entry.batch_time = Some(Instant::now()); entry.batch_time = Some(Instant::now());
// Insert in entries IntMap // Insert in batch_entries IntMap
batch_entries.insert(id, entry); batch_entries.insert(id, entry);
}); });