wip

Author: peasee

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -170,7 +170,8 @@ impl ExecutionPlan for ShuffleReaderExec {
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         let task_id = context.task_id().unwrap_or_else(|| partition.to_string());
-        info!("ShuffleReaderExec::execute({task_id}) partition={partition}, num_locations={}, stage_id={}",
+        info!(
+            "ShuffleReaderExec::execute({task_id}) partition={partition}, num_locations={}, stage_id={}",
             self.partition.get(partition).map(|p| p.len()).unwrap_or(0),
             self.stage_id,
         );
@@ -496,7 +497,10 @@ fn send_fetch_partitions(
     );
     info!(
         "send_fetch_partitions: {} total locations (memory={}, local={}, object_store={}, remote={})",
-        locations.memory.len() + locations.local.len() + locations.object_store.len() + locations.remote.len(),
+        locations.memory.len()
+            + locations.local.len()
+            + locations.object_store.len()
+            + locations.remote.len(),
         locations.memory.len(),
         locations.local.len(),
         locations.object_store.len(),
@@ -522,8 +526,18 @@ fn send_fetch_partitions(
     let customize_endpoint_c = customize_endpoint.clone();
     let metrics_callback_c = metrics_callback.clone();
     let local_locations = locations.local;
+    let local_count = local_locations.len();
     spawned_tasks.push(SpawnedTask::spawn(async move {
-        for p in local_locations {
+        for (i, p) in local_locations.into_iter().enumerate() {
+            info!(
+                "fetch_local[{}/{}]: reading {}/{}/{} from {}",
+                i + 1,
+                local_count,
+                p.partition_id.job_id,
+                p.partition_id.stage_id,
+                p.partition_id.partition_id,
+                p.path
+            );
             let start_time = std::time::Instant::now();
             let r = PartitionReaderEnum::Local
                 .fetch_partition(
@@ -534,6 +548,17 @@ fn send_fetch_partitions(
                     use_tls,
                 )
                 .await;
+            let ok = r.is_ok();
+            info!(
+                "fetch_local[{}/{}]: {}/{}/{} completed in {:.3}s, ok={}",
+                i + 1,
+                local_count,
+                p.partition_id.job_id,
+                p.partition_id.stage_id,
+                p.partition_id.partition_id,
+                start_time.elapsed().as_secs_f64(),
+                ok
+            );
 
             // Record local read metrics if callback is set and read succeeded
             if r.is_ok()
@@ -712,8 +737,7 @@ async fn fetch_partition_remote(
     let port = metadata.port;
     info!(
         "fetch_partition_remote: fetching {}/{}/{} from {}:{}",
-        partition_id.job_id, partition_id.stage_id, partition_id.partition_id,
-        host, port
+        partition_id.job_id, partition_id.stage_id, partition_id.partition_id, host, port
     );
     let mut ballista_client = BallistaClient::try_new(
         host,

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -68,7 +68,7 @@ use datafusion::arrow::error::ArrowError;
 use datafusion::execution::context::TaskContext;
 use datafusion::physical_plan::repartition::BatchPartitioner;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use log::{debug, info};
+use log::{debug, info, warn};
 
 use super::shuffle_writer_trait::ShuffleWriter;
 
@@ -348,7 +348,30 @@ impl ShuffleWriterExec {
                 now.elapsed().as_secs_f64()
             );
 
-            if use_memory {
+            // Watchdog: log periodically if no progress is made
+            let watchdog_job_id = job_id.clone();
+            let watchdog_flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
+            let watchdog_flag_clone = watchdog_flag.clone();
+            let _watchdog = tokio::task::spawn(async move {
+                let mut interval =
+                    tokio::time::interval(std::time::Duration::from_secs(30));
+                interval.tick().await; // skip first immediate tick
+                loop {
+                    interval.tick().await;
+                    if watchdog_flag_clone.load(std::sync::atomic::Ordering::Relaxed) {
+                        break;
+                    }
+                    warn!(
+                        "ShuffleWriter {}/{} partition {}: STALLED - no first batch received after {:.0}s",
+                        watchdog_job_id,
+                        stage_id,
+                        input_partition,
+                        now.elapsed().as_secs_f64()
+                    );
+                }
+            });
+
+            let result = if use_memory {
                 // Use in-memory shuffle storage with configurable format
                 Self::execute_shuffle_write_memory(
                     &job_id,
@@ -393,7 +416,9 @@ impl ShuffleWriterExec {
                     file_ext,
                 )
                 .await
-            }
+            };
+            watchdog_flag.store(true, std::sync::atomic::Ordering::Relaxed);
+            result
         }
     }
 
@@ -470,9 +495,25 @@ impl ShuffleWriterExec {
                 )?;
 
                 let schema = stream.schema();
+                let mut batch_count: u64 = 0;
+                let mut total_rows: u64 = 0;
 
                 while let Some(result) = stream.next().await {
                     let input_batch = result?;
+                    batch_count += 1;
+                    total_rows += input_batch.num_rows() as u64;
+                    if batch_count == 1 {
+                        info!(
+                            "ShuffleWriter partition {input_partition}: received first batch ({} rows) after {:.2}s",
+                            input_batch.num_rows(),
+                            now.elapsed().as_secs_f64()
+                        );
+                    } else if batch_count % 100 == 0 {
+                        info!(
+                            "ShuffleWriter partition {input_partition}: processed {batch_count} batches ({total_rows} rows) in {:.2}s",
+                            now.elapsed().as_secs_f64()
+                        );
+                    }
 
                     write_metrics.input_rows.add(input_batch.num_rows());
 

ballista/executor/src/executor_server.rs

@@ -606,7 +606,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskRunnerPool<T,
         let mut task_runner_shutdown = shutdown_noti.subscribe_for_shutdown();
         let task_runner_complete = shutdown_noti.shutdown_complete_tx.clone();
         tokio::spawn(async move {
-            info!("Starting the task runner pool");
+            info!(
+                "Starting the task runner pool with {} threads",
+                executor_server.executor.concurrent_tasks
+            );
 
             // Use a dedicated executor for CPU bound tasks so that the main tokio
             // executor can still answer requests even when under load