From 5cddc055e6c0521e4ef11c487fd5f1d9e7ef1e62 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Sun, 9 Apr 2023 20:07:02 +0200
Subject: [PATCH] fix(rust-client): use join_all instead of select_all to
 hopefully fix nccl issues (#162)

---
 router/client/src/sharded_client.rs | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/router/client/src/sharded_client.rs b/router/client/src/sharded_client.rs
index 7f0ec6f2..31f76312 100644
--- a/router/client/src/sharded_client.rs
+++ b/router/client/src/sharded_client.rs
@@ -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()
     }
 }