Use `futures::future::join_all` instead of `futures::stream` in assistant_eval (#26974)

Michael Sloan created 9 months ago

Release Notes:

- N/A

Change summary

crates/assistant_eval/src/main.rs | 22 +++++++++++++---------
1 file changed, 13 insertions(+), 9 deletions(-)

Detailed changes

crates/assistant_eval/src/main.rs

@@ -4,7 +4,7 @@ mod judge;
 
 use clap::Parser;
 use eval::{Eval, EvalOutput};
-use futures::{stream, StreamExt};
+use futures::future;
 use gpui::{Application, AsyncApp};
 use headless_assistant::{authenticate_model_provider, find_model, HeadlessAppState};
 use itertools::Itertools;
@@ -126,12 +126,13 @@ fn main() {
                 .await
                 .unwrap();
 
-            let loaded_evals = stream::iter(evals_to_run)
+            let eval_load_futures = evals_to_run
+                .into_iter()
                 .map(|eval_name| {
                     let eval_path = evaluation_data_dir.join(&eval_name);
-                    let repos_dir = repos_dir.clone();
+                    let load_future = Eval::load(eval_name.clone(), eval_path, &repos_dir);
                     async move {
-                        match Eval::load(eval_name.clone(), eval_path, &repos_dir).await {
+                        match load_future.await {
                             Ok(eval) => Some(eval),
                             Err(err) => {
                                 // TODO: Persist errors / surface errors at the end.
@@ -141,8 +142,9 @@ fn main() {
                         }
                     }
                 })
-                .buffer_unordered(args.concurrency)
-                .collect::<Vec<_>>()
+                .collect::<Vec<_>>();
+
+            let loaded_evals = future::join_all(eval_load_futures)
                 .await
                 .into_iter()
                 .flatten()
@@ -160,7 +162,8 @@ fn main() {
             // Sort groups in descending order, so that bigger groups start first.
             evals_grouped_by_url.sort_by_key(|evals| cmp::Reverse(evals.len()));
 
-            let results = stream::iter(evals_grouped_by_url)
+            let result_futures = evals_grouped_by_url
+                .into_iter()
                 .map(|evals| {
                     let model = model.clone();
                     let judge_model = judge_model.clone();
@@ -185,8 +188,9 @@ fn main() {
                         results
                     }
                 })
-                .buffer_unordered(args.concurrency)
-                .collect::<Vec<_>>()
+                .collect::<Vec<_>>();
+
+            let results = future::join_all(result_futures)
                 .await
                 .into_iter()
                 .flatten()