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

Some small improvements

Browse Source
This commit is contained in:
ehsanmok 2023-04-26 18:07:15 -07:00
parent b9ae7e5da1
commit 9fbf6032e5
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
router/client/src/sharded_client.rs
View File

@ -20,7 +20,7 @@ impl ShardedClient {
/// the other shards and returns all uris/unix sockets with the `service_discovery` gRPC method. /// the other shards and returns all uris/unix sockets with the `service_discovery` gRPC method.
async fn from_master_client(mut master_client: Client) -> Result<Self> { async fn from_master_client(mut master_client: Client) -> Result<Self> {
// Get all uris/unix sockets from the master client // Get all uris/unix sockets from the master client
let uris = master_client.service_discovery().await.unwrap(); let uris = master_client.service_discovery().await?;
let futures = uris.into_iter().map(Client::connect_uds); let futures = uris.into_iter().map(Client::connect_uds);
let clients: Result<Vec<Client>> = join_all(futures).await.into_iter().collect(); let clients: Result<Vec<Client>> = join_all(futures).await.into_iter().collect();
Ok(Self::new(clients?)) Ok(Self::new(clients?))

4
router/src/health.rs
View File

@ -4,6 +4,8 @@ use text_generation_client::{
Batch, NextTokenChooserParameters, Request, ShardedClient, StoppingCriteriaParameters, Batch, NextTokenChooserParameters, Request, ShardedClient, StoppingCriteriaParameters,
}; };
const LIVENESS_ID: u64 = u64::MAX;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub(crate) struct Health { pub(crate) struct Health {
client: ShardedClient, client: ShardedClient,
@ -27,7 +29,7 @@ impl Health {
// Dummy batch of 1 token and 1 generated token // Dummy batch of 1 token and 1 generated token
let liveness_request = Request { let liveness_request = Request {
id: u64::MAX, id: LIVENESS_ID,
inputs: "liveness".to_string(), inputs: "liveness".to_string(),
truncate: 10, truncate: 10,
parameters: Some(NextTokenChooserParameters { parameters: Some(NextTokenChooserParameters {