perf(flotilla): Use Worker Affinity with Pre-Shuffle Merge (#5112)

Author: srilman

src/daft-distributed/src/pipeline_node/into_batches.rs

@@ -114,6 +114,7 @@ impl IntoBatchesNode {
                                 StatsState::NotMaterialized,
                             )
                         },
+                        None,
                     )?;
                     if result_tx.send(task).await.is_err() {
                         break;
@@ -136,6 +137,7 @@ impl IntoBatchesNode {
                         StatsState::NotMaterialized,
                     )
                 },
+                None,
             )?;
             let _ = result_tx.send(task).await;
         }

src/daft-distributed/src/pipeline_node/into_partitions.rs

@@ -99,7 +99,7 @@ impl IntoPartitionsNode {
         // Remainder: 10 % 3 = 1 (one task gets an extra input)
         let num_partitions_with_extra_input = tasks.len() % self.num_partitions;
 
-        let mut output_futures = OrderedJoinSet::new();
+        let mut tasks_per_partition = Vec::new();
 
         let mut task_iter = tasks.into_iter();
         for partition_idx in 0..self.num_partitions {
@@ -115,8 +115,13 @@ impl IntoPartitionsNode {
                 .take(chunk_size)
                 .map(|task| task.submit(scheduler_handle))
                 .collect::<DaftResult<Vec<_>>>()?;
+            tasks_per_partition.push(submitted_tasks);
+        }
+
+        let mut output_futures = OrderedJoinSet::new();
+        for tasks in tasks_per_partition {
             output_futures.spawn(async move {
-                let materialized_output = futures::future::try_join_all(submitted_tasks)
+                let materialized_output = futures::future::try_join_all(tasks)
                     .await?
                     .into_iter()
                     .flatten()
@@ -136,6 +141,7 @@ impl IntoPartitionsNode {
                 move |input| {
                     LocalPhysicalPlan::into_partitions(input, 1, StatsState::NotMaterialized)
                 },
+                None,
             )?;
             if result_tx.send(task).await.is_err() {
                 break;
@@ -167,7 +173,7 @@ impl IntoPartitionsNode {
         // Remainder: 10 % 3 = 1 (one partition will split into 4 outputs)
         let num_partitions_with_extra_output = self.num_partitions % tasks.len();
 
-        let mut output_futures = OrderedJoinSet::new();
+        let mut submitted_tasks = Vec::new();
 
         for (input_partition_idx, task) in tasks.into_iter().enumerate() {
             let mut num_outputs = base_splits_per_partition;
@@ -187,7 +193,12 @@ impl IntoPartitionsNode {
                 },
             );
             let submitted_task = into_partitions_task.submit(scheduler_handle)?;
-            output_futures.spawn(submitted_task);
+            submitted_tasks.push(submitted_task);
+        }
+
+        let mut output_futures = OrderedJoinSet::new();
+        for task in submitted_tasks {
+            output_futures.spawn(task);
         }
 
         // Collect all the outputs and emit a new task for each output.
@@ -200,6 +211,7 @@ impl IntoPartitionsNode {
                         TaskContext::from((&self.context, task_id_counter.next())),
                         vec![output],
                         &(self_arc as Arc<dyn DistributedPipelineNode>),
+                        None,
                     )?;
                     if result_tx.send(task).await.is_err() {
                         break;

src/daft-distributed/src/pipeline_node/limit.rs

@@ -106,6 +106,7 @@ impl LimitNode {
                                 input
                             }
                         },
+                        None,
                     )?
                 }
                 Ordering::Greater => {
@@ -122,6 +123,7 @@ impl LimitNode {
                                 StatsState::NotMaterialized,
                             )
                         },
+                        None,
                     )?;
                     *remaining_take = 0;
                     task

src/daft-distributed/src/pipeline_node/mod.rs

@@ -317,6 +317,7 @@ fn make_new_task_from_materialized_outputs<F>(
     materialized_outputs: Vec<MaterializedOutput>,
     node: &Arc<dyn DistributedPipelineNode>,
     plan_builder: F,
+    scheduling_strategy: Option<SchedulingStrategy>,
 ) -> DaftResult<SubmittableTask<SwordfishTask>>
 where
     F: FnOnce(LocalPhysicalPlanRef) -> LocalPhysicalPlanRef + Send + Sync + 'static,
@@ -338,7 +339,7 @@ where
         plan,
         node.config().execution_config.clone(),
         psets,
-        SchedulingStrategy::Spread,
+        scheduling_strategy.unwrap_or(SchedulingStrategy::Spread),
         node.context().to_hashmap(),
     );
     Ok(SubmittableTask::new(task))
@@ -348,8 +349,15 @@ fn make_in_memory_task_from_materialized_outputs(
     task_context: TaskContext,
     materialized_outputs: Vec<MaterializedOutput>,
     node: &Arc<dyn DistributedPipelineNode>,
+    scheduling_strategy: Option<SchedulingStrategy>,
 ) -> DaftResult<SubmittableTask<SwordfishTask>> {
-    make_new_task_from_materialized_outputs(task_context, materialized_outputs, node, |input| input)
+    make_new_task_from_materialized_outputs(
+        task_context,
+        materialized_outputs,
+        node,
+        |input| input,
+        scheduling_strategy,
+    )
 }
 
 fn append_plan_to_existing_task<F>(

src/daft-distributed/src/pipeline_node/shuffles/gather.rs

@@ -83,6 +83,7 @@ impl GatherNode {
             TaskContext::from((&self_clone.context, task_id_counter.next())),
             materialized,
             &(self_clone as Arc<dyn DistributedPipelineNode>),
+            None,
         )?;
 
         let _ = result_tx.send(task).await;

src/daft-distributed/src/pipeline_node/shuffles/pre_shuffle_merge.rs

@@ -12,7 +12,7 @@ use crate::{
     },
     scheduling::{
         scheduler::{SchedulerHandle, SubmittableTask},
-        task::{SwordfishTask, TaskContext},
+        task::{SchedulingStrategy, SwordfishTask, TaskContext},
         worker::WorkerId,
     },
     stage::{StageConfig, StageExecutionContext, TaskIDCounter},
@@ -172,6 +172,10 @@ impl PreShuffleMergeNode {
                         TaskContext::from((self.context(), task_id_counter.next())),
                         materialized_outputs,
                         &(self_clone as Arc<dyn DistributedPipelineNode>),
+                        Some(SchedulingStrategy::WorkerAffinity {
+                            worker_id,
+                            soft: false,
+                        }),
                     )?;
 
                     // Send the task directly to result_tx
@@ -183,13 +187,17 @@ impl PreShuffleMergeNode {
         }
 
         // Handle any remaining buckets that haven't reached the threshold
-        for (_, materialized_outputs) in worker_buckets {
+        for (worker_id, materialized_outputs) in worker_buckets {
             if !materialized_outputs.is_empty() {
                 let self_clone = self.clone();
                 let task = make_in_memory_task_from_materialized_outputs(
                     TaskContext::from((self.context(), task_id_counter.next())),
                     materialized_outputs,
                     &(self_clone as Arc<dyn DistributedPipelineNode>),
+                    Some(SchedulingStrategy::WorkerAffinity {
+                        worker_id,
+                        soft: false,
+                    }),
                 )?;
 
                 if result_tx.send(task).await.is_err() {

src/daft-distributed/src/pipeline_node/shuffles/repartition.rs

@@ -99,6 +99,7 @@ impl RepartitionNode {
                 TaskContext::from((&self_clone.context, task_id_counter.next())),
                 partition_group,
                 &(self_clone as Arc<dyn DistributedPipelineNode>),
+                None,
             )?;
 
             let _ = result_tx.send(task).await;

src/daft-distributed/src/pipeline_node/sink.rs

@@ -133,6 +133,7 @@ impl SinkNode {
                     StatsState::NotMaterialized,
                 )
             },
+            None,
         )?;
         let _ = sender.send(task).await;
         Ok(())

src/daft-distributed/src/pipeline_node/sort.rs

@@ -211,6 +211,7 @@ impl SortNode {
                         StatsState::NotMaterialized,
                     )
                 },
+                None,
             )?;
             let _ = result_tx.send(task).await;
             return Ok(());
@@ -242,6 +243,7 @@ impl SortNode {
                             StatsState::NotMaterialized,
                         )
                     },
+                    None,
                 )?;
                 let submitted_task = task.submit(&scheduler_handle)?;
                 Ok(submitted_task)
@@ -281,6 +283,7 @@ impl SortNode {
                             StatsState::NotMaterialized,
                         )
                     },
+                    None,
                 )?;
                 let submitted_task = task.submit(&scheduler_handle)?;
                 Ok(submitted_task)
@@ -312,6 +315,7 @@ impl SortNode {
                         StatsState::NotMaterialized,
                     )
                 },
+                None,
             )?;
             let _ = result_tx.send(task).await;
         }

src/daft-shuffles/src/client/flight_client.rs

@@ -40,8 +40,12 @@ impl ShuffleFlightClient {
             let channel = endpoint.connect().await.map_err(|e| {
                 DaftError::External(format!("Failed to connect to endpoint: {}", e).into())
             })?;
-            self.inner =
-                ClientState::Initialized(std::mem::take(address), FlightClient::new(channel));
+            let client = FlightClient::new(channel);
+            let inner = client.into_inner().max_decoding_message_size(usize::MAX);
+            self.inner = ClientState::Initialized(
+                std::mem::take(address),
+                FlightClient::new_from_inner(inner),
+            );
         }
         match &mut self.inner {
             ClientState::Uninitialized(_) => unreachable!("Client should be initialized"),