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

use join_all instead

Browse Source
This commit is contained in:
OlivierDehaene 2023-04-09 15:10:59 +02:00
parent 4267378b1f
commit af10275f46
1 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -60,9 +60,8 @@ impl ShardedClient {
.iter_mut() .iter_mut()
.map(|client| Box::pin(client.prefill(batch.clone()))) .map(|client| Box::pin(client.prefill(batch.clone())))
.collect(); .collect();
// As soon as we receive one response, we can return as all shards will return the same // all shards return the same message
let (result, _, _) = select_all(futures).await; join_all(futures).await.pop().unwrap()
result
} }
/// Generate one token for each request in the given cached batches /// Generate one token for each request in the given cached batches
@ -79,8 +78,7 @@ impl ShardedClient {
.iter_mut() .iter_mut()
.map(|client| Box::pin(client.decode(batches.clone()))) .map(|client| Box::pin(client.decode(batches.clone())))
.collect(); .collect();
// As soon as we receive one response, we can return as all shards will return the same // all shards return the same message
let (result, _, _) = select_all(futures).await; join_all(futures).await.pop().unwrap()
result
} }
} }