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

improved naming

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-02 14:12:05 +01:00
parent 0c93da571b
commit d2d5394991
1 changed files with 10 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
router/src/db.rs
View File

@ -4,7 +4,6 @@ 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};
@ -29,7 +28,7 @@ pub(crate) struct Entry {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct Db { pub(crate) struct Db {
/// Channel to communicate with the background database task /// Channel to communicate with the background database task
sender: UnboundedSender<DatabaseCommand>, db_sender: UnboundedSender<DatabaseCommand>,
} }
impl Db { impl Db {
@ -40,14 +39,14 @@ impl Db {
// Launch background database task // Launch background database task
tokio::spawn(database_task(db_receiver)); tokio::spawn(database_task(db_receiver));
Self { sender: db_sender } Self { db_sender }
} }
/// Append an entry to the database /// Append an entry to the database
pub(crate) fn append(&self, entry: Entry) { pub(crate) fn append(&self, entry: Entry) {
// Send append command to the background task managing the state // Send append command to the background task managing the state
// Unwrap is safe here // Unwrap is safe here
self.sender.send(DatabaseCommand::Append(entry)).unwrap(); self.db_sender.send(DatabaseCommand::Append(entry)).unwrap();
} }
// Get the next batch // Get the next batch
@ -57,19 +56,19 @@ impl Db {
max_size: usize, max_size: usize,
) -> Option<NextBatch> { ) -> Option<NextBatch> {
// Create response channel // Create response channel
let (sender, receiver) = oneshot::channel(); let (response_sender, response_receiver) = oneshot::channel();
// Send next batch command to the background task managing the state // Send next batch command to the background task managing the state
// Unwrap is safe here // Unwrap is safe here
self.sender self.db_sender
.send(DatabaseCommand::NextBatch { .send(DatabaseCommand::NextBatch {
min_size, min_size,
max_size, max_size,
response_rx: sender, response_sender,
}) })
.unwrap(); .unwrap();
// Await on response channel // Await on response channel
// Unwrap is safe here // Unwrap is safe here
receiver.await.unwrap() response_receiver.await.unwrap()
} }
} }
@ -83,10 +82,10 @@ async fn database_task(mut receiver: UnboundedReceiver<DatabaseCommand>) {
DatabaseCommand::NextBatch { DatabaseCommand::NextBatch {
min_size, min_size,
max_size, max_size,
response_rx, response_sender,
} => { } => {
let next_batch = state.next_batch(min_size, max_size); let next_batch = state.next_batch(min_size, max_size);
response_rx.send(next_batch).unwrap_or(()); response_sender.send(next_batch).unwrap_or(());
} }
} }
} }
@ -176,7 +175,7 @@ enum DatabaseCommand {
NextBatch { NextBatch {
min_size: Option<usize>, min_size: Option<usize>,
max_size: usize, max_size: usize,
response_rx: oneshot::Sender<Option<NextBatch>>, response_sender: oneshot::Sender<Option<NextBatch>>,
}, },
} }