Better logs for the benchmarks

gabotechs · gabotechs · commit 75398747a284 · 2025-11-18T12:00:12.000+01:00
diff --git a/benchmarks/cdk/bin/datafusion-bench.ts b/benchmarks/cdk/bin/datafusion-bench.ts
@@ -15,6 +15,7 @@ async function main () {
     .option('-i, --iterations <number>', 'Number of iterations', '3')
     .option('--files-per-task <number>', 'Files per task', '4')
     .option('--cardinality-task-sf <number>', 'Cardinality task scale factor', '2')
+    .option('--query <number>', 'A specific query to run', undefined)
     .parse(process.argv);
 
   const options = program.opts();
@@ -36,6 +37,10 @@ async function main () {
   `)
 
   for (let id of IDS) {
+    if (options.query && parseInt(options.query) !== id) {
+      continue
+    }
+
     const queryId = `q${id}`;
     const filePath = path.join(queriesPath, `${queryId}.sql`)
     const content = await fs.readFile(filePath, 'utf-8')
diff --git a/benchmarks/cdk/bin/worker.rs b/benchmarks/cdk/bin/worker.rs
@@ -6,6 +6,7 @@ use axum::{Json, Router, extract::Query, http::StatusCode, routing::get};
 use dashmap::{DashMap, Entry};
 use datafusion::common::DataFusionError;
 use datafusion::common::instant::Instant;
+use datafusion::common::runtime::SpawnedTask;
 use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
@@ -14,15 +15,17 @@ use datafusion_distributed::{
     DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
     create_flight_client, display_plan_ascii,
 };
-use futures::{TryFutureExt, TryStreamExt};
-use log::{error, info};
+use futures::{StreamExt, TryFutureExt};
+use log::{error, info, warn};
 use object_store::ObjectStore;
 use object_store::aws::AmazonS3Builder;
 use serde::Serialize;
 use std::collections::HashMap;
 use std::error::Error;
 use std::fmt::Display;
+use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, RwLock};
+use std::time::Duration;
 use structopt::StructOpt;
 use tonic::transport::{Channel, Server};
 use url::Url;
@@ -119,7 +122,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
                         if sql.trim().is_empty() {
                             continue;
                         }
-                        info!("Executing query: {sql}");
                         let df = ctx.sql(sql).await.map_err(err)?;
                         df_opt = Some(df);
                     }
@@ -129,15 +131,32 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
                     let start = Instant::now();
 
+                    info!("Executing query...");
+                    let abort_notifier = AbortNotifier::new("Query aborted");
+                    let abort_notifier_clone = abort_notifier.clone();
+                    let task = SpawnedTask::spawn(async move {
+                        let _ = abort_notifier_clone;
+                        loop {
+                            tokio::time::sleep(Duration::from_secs(5)).await;
+                            info!("Query still running...");
+                        }
+                    });
                     let physical = df.create_physical_plan().await.map_err(err)?;
-                    let stream = execute_stream(physical.clone(), ctx.task_ctx()).map_err(err)?;
-                    let batches = stream.try_collect::<Vec<_>>().await.map_err(err)?;
-                    let count = batches.iter().map(|b| b.num_rows()).sum::<usize>();
+                    let mut stream =
+                        execute_stream(physical.clone(), ctx.task_ctx()).map_err(err)?;
+                    let mut count = 0;
+                    while let Some(batch) = stream.next().await {
+                        count += batch.map_err(err)?.num_rows();
+                        info!("Gathered {count} rows, query still in progress..")
+                    }
                     let plan = display_plan_ascii(physical.as_ref(), true);
+                    drop(task);
 
                     let elapsed = start.elapsed();
                     let ms = elapsed.as_secs_f64() * 1000.0;
+                    info!("Finished executing query:\n{sql}\n\n{plan}");
                     info!("Returned {count} rows in {ms} ms");
+                    abort_notifier.finished();
 
                     Ok::<_, (StatusCode, String)>(Json(QueryResult { count, plan }))
                 }
@@ -162,6 +181,33 @@ async fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+struct AbortNotifier {
+    aborted: AtomicBool,
+    msg: String,
+}
+
+impl AbortNotifier {
+    fn new(msg: impl Display) -> Arc<Self> {
+        Arc::new(AbortNotifier {
+            aborted: AtomicBool::new(true),
+            msg: msg.to_string(),
+        })
+    }
+
+    fn finished(&self) {
+        self.aborted
+            .store(false, std::sync::atomic::Ordering::Relaxed)
+    }
+}
+
+impl Drop for AbortNotifier {
+    fn drop(&mut self) {
+        if self.aborted.load(std::sync::atomic::Ordering::Relaxed) {
+            warn!("{}", self.msg);
+        }
+    }
+}
+
 fn err(s: impl Display) -> (StatusCode, String) {
     (StatusCode::INTERNAL_SERVER_ERROR, s.to_string())
 }
@@ -222,7 +268,7 @@ async fn background_ec2_worker_resolver(urls: Arc<RwLock<Vec<Url>>>) {
                 );
                 *urls.write().unwrap() = workers;
             }
-            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+            tokio::time::sleep(Duration::from_secs(1)).await;
         }
     });
 }
diff --git a/benchmarks/cdk/lib/cdk-stack.ts b/benchmarks/cdk/lib/cdk-stack.ts
@@ -141,26 +141,23 @@ EOF`,
       instances.push(instance);
     }
 
+    const firstInstance = instances[0]?.instanceId
     // Output Session Manager commands for all instances
     new CfnOutput(this, 'ConnectCommands', {
-      value: instances.map((inst, i) =>
-        `
-# instance-${i}
-aws ssm start-session --target ${inst.instanceId}
-`
-      ).join(''),
-      description: 'Session Manager commands to connect to instances',
-    });
+      value: `\n
+export INSTANCE_ID=${firstInstance}
+
+# === port forward the HTTP ===
+aws ssm start-session --target $INSTANCE_ID --document-name AWS-StartPortForwardingSession --parameters "portNumber=9000,localPortNumber=9000"
 
-    // Output port forwarding commands
-    new CfnOutput(this, 'PortForwardCommands', {
-      value: instances.map((inst, i) =>
-        `
-# instance-${i} (forward port 9000 to localhost:${9000 + i})
-aws ssm start-session --target ${inst.instanceId} --document-name AWS-StartPortForwardingSession --parameters "portNumber=9000,localPortNumber=${9000 + i}"
-`
-      ).join(''),
-      description: 'Port forwarding commands (HTTP server on port 9000)',
+# === open a sh session in the remote machine ===
+aws ssm start-session --target $INSTANCE_ID
+
+# === See worker logs inside a sh session ===
+sudo journalctl -u worker.service -f -o cat
+
+`,
+      description: 'Session Manager commands to connect to instances',
     });
 
     // Custom resource to restart worker service on every deploy
@@ -180,6 +177,7 @@ aws ssm start-session --target ${inst.instanceId} --document-name AWS-StartPortF
           },
         },
         physicalResourceId: cr.PhysicalResourceId.of(`restart-${Date.now()}`),
+        ignoreErrorCodesMatching: '.*',
       },
       policy: cr.AwsCustomResourcePolicy.fromStatements([
         new iam.PolicyStatement({
