Mirrors
/
preemo_text-generation-inference
mirror of https://github.com/Preemo-Inc/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(rust-client): use join_all instead of select_all to hopefully fix nccl issues (#162)

This commit is contained in:
OlivierDehaene 2023-04-09 20:07:02 +02:00 committed by GitHub
parent e63a21eb4d
commit 5cddc055e6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2,7 +2,6 @@
use crate::Result;
use crate::{Batch, Client, Generation};
use futures::future::join_all;
use futures::future::select_all;
use tonic::transport::Uri;
use tracing::instrument;
@ -60,9 +59,8 @@ impl ShardedClient {
.iter_mut()
.map(|client| Box::pin(client.prefill(batch.clone())))
.collect();
// As soon as we receive one response, we can return as all shards will return the same
let (result, _, _) = select_all(futures).await;
result
// all shards return the same message
join_all(futures).await.pop().unwrap()
}
/// Generate one token for each request in the given cached batches
@ -79,8 +77,7 @@ impl ShardedClient {
.iter_mut()
.map(|client| Box::pin(client.decode(batches.clone())))
.collect();
// As soon as we receive one response, we can return as all shards will return the same
let (result, _, _) = select_all(futures).await;
result
// all shards return the same message
join_all(futures).await.pop().unwrap()
}
}