Update benchmarks in README.md and make ballista work (#557)

Makes ballista work for tpch sf100 and updates the charts in the
benchmarks

Author: gabotechs

README.md

@@ -49,12 +49,12 @@ relative to DataFusion Distributed
 <details>
 <summary>Per-dataset totals</summary>
 
-| Dataset     | df-dist |                                                                                Ballista | Spark | Trino | Queries compared |
-|-------------|--------:|----------------------------------------------------------------------------------------:|------:|------:|-----------------:|
-| TPC-H SF1   |  **7s** |                                                                                     11s |   30s |   18s |               22 |
-| TPC-H SF10  | **10s** |                                                                                     57s |   51s |   33s |               22 |
-| TPC-H SF100 | **42s** | N/A ([#1836](https://github.com/datafusion-contrib/datafusion-distributed/issues/1836)) |  261s |   93s |               19 |
-| TPC-DS SF1  | **29s** |                                                                                     72s |  101s |   85s |               67 |
+| Dataset     | df-dist | Ballista | Spark | Trino | Queries compared |
+|-------------|--------:|---------:|------:|------:|-----------------:|
+| TPC-H SF1   |  **7s** |      11s |   30s |   18s |               22 |
+| TPC-H SF10  | **10s** |      42s |   51s |   33s |               22 |
+| TPC-H SF100 | **42s** |     237s |  261s |   93s |               19 |
+| TPC-DS SF1  | **29s** |      72s |  101s |   85s |               67 |
 
 ![TPC-H SF1](./docs/source/_static/images/tpch_sf1.png)
 ![TPC-H SF10](./docs/source/_static/images/tpch_sf10.png)

benchmarks/cdk/ballista/src/ballista_http.rs

@@ -31,6 +31,10 @@ struct Cmd {
     /// The bucket name.
     #[structopt(long, default_value = "datafusion-distributed-benchmarks")]
     bucket: String,
+
+    /// Number of partitions used for scans and distributed shuffle stages.
+    #[structopt(long, default_value = "96")]
+    target_partitions: usize,
 }
 
 #[tokio::main]
@@ -59,7 +63,13 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let runtime_env = Arc::new(RuntimeEnv::default());
     runtime_env.register_object_store(&s3_url, s3);
 
-    let config = SessionConfig::new_with_ballista().with_ballista_job_name("Benchmarks");
+    info!(
+        "Configuring Ballista with {} target partitions",
+        cmd.target_partitions
+    );
+    let config = SessionConfig::new_with_ballista()
+        .with_ballista_job_name("Benchmarks")
+        .with_target_partitions(cmd.target_partitions);
 
     let state = SessionStateBuilder::new()
         .with_config(config)

benchmarks/cdk/lib/ballista.ts

@@ -16,6 +16,10 @@ let ballistaExecutorBinary: s3assets.Asset
 
 // Ballista is built as a standalone project, separate from the main workspace
 const BALLISTA_DIR = path.join(ROOT, 'benchmarks/cdk/ballista')
+const BALLISTA_TARGET_PARTITIONS = 96 // 12 c5n.2xlarge instances * 8 vCPUs
+const BALLISTA_EXECUTOR_MEMORY_POOL = '16GiB' // leave ~4 GiB for the OS and untracked allocations
+const BALLISTA_JOB_DATA_CLEAN_UP_INTERVAL_SECONDS = 300
+const BALLISTA_JOB_DATA_TTL_SECONDS = 3600
 
 export const BALLISTA_ENGINE: QueryEngine = {
     beforeEc2Machines(ctx: BeforeEc2MachinesContext): void {
@@ -95,7 +99,10 @@ ExecStart=/usr/local/bin/ballista-executor \\
   --bind-port 50051 \\
   --work-dir /var/ballista/executor \\
   --scheduler-host localhost \\
-  --scheduler-port 50050
+  --scheduler-port 50050 \\
+  --memory-pool-size ${BALLISTA_EXECUTOR_MEMORY_POOL} \\
+  --job-data-clean-up-interval-seconds ${BALLISTA_JOB_DATA_CLEAN_UP_INTERVAL_SECONDS} \\
+  --job-data-ttl-seconds ${BALLISTA_JOB_DATA_TTL_SECONDS}
 Restart=on-failure
 RestartSec=5
 User=root
@@ -118,7 +125,9 @@ After=network.target ballista-scheduler.service
 Requires=ballista-scheduler.service
 [Service]
 Type=simple
-ExecStart=/usr/local/bin/ballista-http --bucket ${ctx.bucketName}
+ExecStart=/usr/local/bin/ballista-http \\
+  --bucket ${ctx.bucketName} \\
+  --target-partitions ${BALLISTA_TARGET_PARTITIONS}
 Restart=on-failure
 RestartSec=5
 User=root
@@ -159,11 +168,13 @@ BALLISTA_EOF`
 
         // Reconfigure executors on worker nodes to point to scheduler. The executor in the machine holding the scheduler
         // communicates to it using localhost, so no need to update it with scheduler.instancePrivateIp.
-        sendCommandsUnconditionally(
+        const updateExecutors = sendCommandsUnconditionally(
             ctx.scope,
             "ConfigureBallistaExecutors",
             [scheduler, ...executors],
             [
+                `aws s3 cp s3://${ballistaExecutorBinary.s3BucketName}/${ballistaExecutorBinary.s3ObjectKey} /usr/local/bin/ballista-executor`,
+                'chmod +x /usr/local/bin/ballista-executor',
                 `cat > /etc/systemd/system/ballista-executor.service << 'BALLISTA_EOF'
 [Unit]
 Description=Ballista Executor
@@ -175,7 +186,10 @@ ExecStart=/usr/local/bin/ballista-executor \\
   --bind-port 50051 \\
   --work-dir /var/ballista/executor \\
   --scheduler-host ${scheduler.instancePrivateIp} \\
-  --scheduler-port 50050
+  --scheduler-port 50050 \\
+  --memory-pool-size ${BALLISTA_EXECUTOR_MEMORY_POOL} \\
+  --job-data-clean-up-interval-seconds ${BALLISTA_JOB_DATA_CLEAN_UP_INTERVAL_SECONDS} \\
+  --job-data-ttl-seconds ${BALLISTA_JOB_DATA_TTL_SECONDS}
 Restart=on-failure
 RestartSec=5
 User=root
@@ -187,8 +201,21 @@ StandardError=append:/var/ballista/logs/executor.log
 WantedBy=multi-user.target
 BALLISTA_EOF`,
                 'systemctl daemon-reload',
+                'systemctl enable --now ballista-executor',
                 'systemctl restart ballista-executor',
             ]
         )
+
+        const updateHttp = sendCommandsUnconditionally(
+            ctx.scope,
+            "UpdateBallistaHttp",
+            [scheduler],
+            [
+                `aws s3 cp s3://${ballistaServerBinary.s3BucketName}/${ballistaServerBinary.s3ObjectKey} /usr/local/bin/ballista-http`,
+                'chmod +x /usr/local/bin/ballista-http',
+                'systemctl restart ballista-http',
+            ]
+        )
+        updateHttp.node.addDependency(updateExecutors)
     }
 }

benchmarks/cdk/lib/cdk-stack.ts

@@ -219,7 +219,7 @@ export function sendCommandsUnconditionally(
     name: string,
     instances: ec2.Instance[],
     commands: string[]
-) {
+): cr.AwsCustomResource {
     const cmd = new cr.AwsCustomResource(construct, name, {
         onUpdate: {
             service: 'SSM',
@@ -247,4 +247,5 @@ export function sendCommandsUnconditionally(
 
     // Ensure instances are created before restarting
     cmd.node.addDependency(...instances)
+    return cmd
 }