Local worker allocations (#559)

Small optimization that allows the dynamic planner to collocate
single-tasked stages into the same machine that is acting as
coordinator.

During dynamic planning, if all stages are small, they will just use 1
task, and before this PR, there was already some logic for co-locating
all stages in a single worker. The challenge:

The single worker was a random worker in the cluster.

This means that if the coordinating context also happens to be in the
scope of a worker, there's a missed chance for running everything
locally, almost as if it was just single-node DataFusion.

This PRs unlocks this scenario by allowing users to inject their own
`LocalWorkerContext` into the coordinating `SessionContext`. That way,
the coordinator itself can recognize it's also a worker, and can send to
itself all these single-tasked stages, avoiding network hops and data
serialization.

## Benchmarks

TPCH-SF1: 1.12 faster ✔
TPCH-SF10: 1.06 faster ✔
TPCH-SF100: 1.01 faster ✔
TPC-DS: 1.08 faster ✔
ClickBench: 1.01 faster ✔

Author: gabotechs

benchmarks/cdk/bin/worker.rs

@@ -88,6 +88,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
     info!("Starting HTTP listener on {LISTENER_ADDR}...");
     let listener = tokio::net::TcpListener::bind(LISTENER_ADDR).await?;
 
+    let self_url = get_self_url().await?;
+    info!("Resolved self URL as {self_url}");
+
     // Register S3 object store
     let s3_url = Url::parse(&format!("s3://{}", cmd.bucket))?;
 
@@ -100,9 +103,18 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let runtime_env = Arc::new(RuntimeEnv::default());
     runtime_env.register_object_store(&s3_url, s3);
 
+    let worker = Worker::from_session_builder(|ctx: WorkerQueryContext| async move {
+        Ok(ctx
+            .builder
+            .with_distributed_user_codec(WorkUnitFileScanCodec)
+            .build())
+    })
+    .with_runtime_env(Arc::clone(&runtime_env));
+
     let state_builder = SessionStateBuilder::new()
         .with_default_features()
-        .with_runtime_env(Arc::clone(&runtime_env))
+        .with_runtime_env(runtime_env)
+        .with_distributed_local_worker_context(worker.to_local_worker_context(self_url))
         .with_distributed_worker_resolver(Ec2WorkerResolver::new())
         .with_distributed_planner()
         .with_distributed_broadcast_joins(cmd.broadcast_joins)?
@@ -119,14 +131,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let ctx = SessionContext::from(state);
     let ctx_clone = ctx.clone();
 
-    let worker = Worker::from_session_builder(|ctx: WorkerQueryContext| async move {
-        Ok(ctx
-            .builder
-            .with_distributed_user_codec(WorkUnitFileScanCodec)
-            .build())
-    })
-    .with_runtime_env(runtime_env);
-
     let http_server = axum::serve(
         listener,
         Router::new()
@@ -295,6 +299,38 @@ fn err(s: impl Display) -> (StatusCode, String) {
     (StatusCode::INTERNAL_SERVER_ERROR, s.to_string())
 }
 
+async fn get_self_url() -> Result<Url, Box<dyn Error>> {
+    let client = reqwest::Client::new();
+    // AWS EC2 Instance Metadata Service (IMDS) IPv4 endpoint.
+    // AWS docs: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+    // We read the instance's private IPv4 from here so the worker can build its own
+    // gRPC URL as http://<private-ip>:9001.
+    let metadata_base = "http://169.254.169.254/latest/meta-data";
+
+    // Try IMDSv2 first, then fall back to unauthenticated metadata reads if token
+    // acquisition is disabled in the instance configuration.
+    let token = match client
+        .put("http://169.254.169.254/latest/api/token")
+        .header("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+        .send()
+        .await
+    {
+        Ok(response) => match response.error_for_status() {
+            Ok(response) => response.text().await.ok(),
+            Err(_) => None,
+        },
+        Err(_) => None,
+    };
+
+    let mut request = client.get(format!("{metadata_base}/local-ipv4"));
+    if let Some(token) = token {
+        request = request.header("X-aws-ec2-metadata-token", token);
+    }
+
+    let local_ip = request.send().await?.error_for_status()?.text().await?;
+    Ok(Url::parse(&format!("http://{}:9001", local_ip.trim()))?)
+}
+
 #[derive(Clone)]
 struct Ec2WorkerResolver {
     urls: Arc<RwLock<Vec<Url>>>,

src/coordinator/prepare_dynamic_plan.rs

@@ -85,10 +85,16 @@ pub(super) async fn prepare_dynamic_plan(
             let mut load_info_rxs = Vec::with_capacity(input_stage.tasks);
 
             let routed_urls = if input_stage.tasks == 1 {
-                // If there's an input stage with a single worker, and the current stage is also
-                // going to run in a single worker, we want to co-locate them so that unnecessary
-                // network transfers are avoided.
-                match stage_coordinator.find_input_stage_with_single_url() {
+                match stage_coordinator
+                    // If the current coordinating context is running within the scope of a local
+                    // worker (same coordinating machine happens to also be a worker), we prefer to
+                    // co-locate single-tasked stages on it.
+                    .find_self_url()
+                    // If there's an input stage with a single worker, and the current stage is also
+                    // going to run in a single worker, we want to co-locate them so that unnecessary
+                    // network transfers are avoided.
+                    .or_else(|| stage_coordinator.find_input_stage_with_single_url())
+                {
                     Some(single_url) => vec![single_url],
                     None => stage_coordinator.routed_urls()?,
                 }

src/coordinator/query_coordinator.rs

@@ -8,6 +8,7 @@ use crate::passthrough_headers::get_passthrough_headers;
 use crate::stage::LocalStage;
 use crate::work_unit_feed::WorkUnitFeedRegistry;
 use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time};
+use crate::worker::LocalWorkerContext;
 use crate::{
     BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg,
     DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext,
@@ -408,6 +409,13 @@ impl<'a> StageCoordinator<'a> {
         Ok(routed_urls)
     }
 
+    pub(super) fn find_self_url(&self) -> Option<Url> {
+        self.task_ctx
+            .session_config()
+            .get_extension::<LocalWorkerContext>()
+            .map(|v| v.self_url.clone())
+    }
+
     pub(super) fn find_input_stage_with_single_url(&self) -> Option<Url> {
         let mut single_stage_url = None;
         self.plan

src/distributed_ext.rs

@@ -8,8 +8,8 @@ use crate::protocol::set_distributed_channel_resolver;
 use crate::work_unit_feed::set_distributed_work_unit_feed;
 use crate::worker_resolver::set_distributed_worker_resolver;
 use crate::{
-    ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider,
-    WorkerResolver, get_distributed_worker_resolver,
+    ChannelResolver, DistributedConfig, LocalWorkerContext, TaskEstimator, WorkUnitFeed,
+    WorkUnitFeedProvider, WorkerResolver, get_distributed_worker_resolver,
 };
 use datafusion::common::DataFusionError;
 use datafusion::config::ConfigExtension;
@@ -600,6 +600,17 @@ pub trait DistributedExt: Sized {
         &mut self,
         dynamic_bytes_per_partition: usize,
     ) -> Result<(), DataFusionError>;
+
+    /// Target throughput in bytes per partition per second used by the dynamic task count
+    /// allocator to decide how many tasks to assign to each stage based on runtime statistics.
+    fn with_distributed_local_worker_context(
+        self,
+        local_worker_context: LocalWorkerContext,
+    ) -> Self;
+
+    /// Same as [DistributedExt::with_distributed_local_worker_context] but with an
+    /// in-place mutation.
+    fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
 }
 
 /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context.
@@ -768,6 +779,10 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext) {
+        self.set_extension(Arc::new(local_worker_context));
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -859,6 +874,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_dynamic_bytes_per_partition)]
             #[expr($?;Ok(self))]
             fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_local_worker_context)]
+            #[expr($;self)]
+            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
         }
     }
 }
@@ -987,6 +1006,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_dynamic_bytes_per_partition)]
             #[expr($?;Ok(self))]
             fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
+            #[call(set_distributed_local_worker_context)]
+            #[expr($;self)]
+            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
         }
     }
 }
@@ -1117,6 +1141,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_dynamic_bytes_per_partition)]
             #[expr($?;Ok(self))]
             fn with_distributed_dynamic_bytes_per_partition(mut self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
+            #[call(set_distributed_local_worker_context)]
+            #[expr($;self)]
+            fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self;
         }
     }
 }
@@ -1240,6 +1269,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_dynamic_bytes_per_partition)]
             #[expr($?;Ok(self))]
             fn with_distributed_dynamic_bytes_per_partition(self, dynamic_bytes_per_partition: usize) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext);
+            #[call(set_distributed_local_worker_context)]
+            #[expr($;self)]
+            fn with_distributed_local_worker_context(self, local_worker_context: LocalWorkerContext) -> Self;
         }
     }
 }

src/lib.rs

@@ -33,6 +33,7 @@ pub use metrics::{
     MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric,
     P99LatencyMetric, rewrite_distributed_plan_with_metrics,
 };
+pub use worker::LocalWorkerContext;
 
 #[cfg(any(feature = "integration", test))]
 pub mod test_utils;

src/test_utils/localhost.rs

@@ -61,12 +61,14 @@ where
         });
     }
     tokio::time::sleep(Duration::from_millis(100)).await;
+    let first_worker_url = Url::parse(&format!("http://127.0.0.1:{}", ports[0])).unwrap();
 
     let worker_resolver = LocalHostWorkerResolver::new(ports);
     let state = SessionStateBuilder::new()
         .with_default_features()
         .with_config(SessionConfig::new().with_target_partitions(3))
         .with_distributed_planner()
+        .with_distributed_local_worker_context(workers[0].to_local_worker_context(first_worker_url))
         .with_distributed_worker_resolver(worker_resolver)
         // Test datasets are tiny, so budget one byte per partition: the estimator then asks for far
         // more partitions than exist, which gets capped at the worker count, fanning every scan out

src/worker/mod.rs

@@ -9,11 +9,12 @@ mod worker_connection_pool;
 mod worker_service;
 
 pub(crate) use single_write_multi_read::SingleWriteMultiRead;
-pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool};
+pub(crate) use worker_connection_pool::WorkerConnectionPool;
 
 pub use session_builder::{
     DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt,
     WorkerQueryContext, WorkerSessionBuilder,
 };
 pub use task_data::TaskData;
+pub use worker_connection_pool::LocalWorkerContext;
 pub use worker_service::Worker;

src/worker/worker_connection_pool.rs

@@ -25,7 +25,7 @@ use url::Url;
 ///
 /// This information can be used for executing tasks locally bypassing gRPC comms if the tasks that
 /// needs to be remotely executed happens to be owned by this same worker.
-pub(crate) struct LocalWorkerContext {
+pub struct LocalWorkerContext {
     /// The registry of in-flight tasks the [crate::Worker] in the current scope owns.
     pub(crate) task_data_entries: Arc<TaskDataEntries>,
     /// The URL of the [crate::Worker] in scope. When trying to reach to a target URL that happens

src/worker/worker_service.rs

@@ -1,4 +1,4 @@
-use crate::worker::{SingleWriteMultiRead, WorkerSessionBuilder};
+use crate::worker::{LocalWorkerContext, SingleWriteMultiRead, WorkerSessionBuilder};
 use crate::{DefaultSessionBuilder, TaskData, TaskKey};
 use datafusion::common::DataFusionError;
 use datafusion::execution::runtime_env::RuntimeEnv;
@@ -8,6 +8,7 @@ use moka::future::Cache;
 use std::borrow::Cow;
 use std::sync::Arc;
 use std::time::Duration;
+use url::Url;
 
 const TASK_CACHE_TTI: Duration = Duration::from_mins(10);
 
@@ -121,6 +122,17 @@ impl Worker {
         &self.version
     }
 
+    /// Builds a [LocalWorkerContext] suitable to be injecting into a coordinating [SessionContext].
+    /// Having a [LocalWorkerContext] present in the coordinating [SessionContext] is not strictly
+    /// necessary, but it allows the planner to better colocate small stages near it, avoiding
+    /// unnecessary network hops.
+    pub fn to_local_worker_context(&self, self_url: Url) -> LocalWorkerContext {
+        LocalWorkerContext {
+            task_data_entries: Arc::clone(&self.task_data_entries),
+            self_url,
+        }
+    }
+
     /// Returns the number of cached task entries currently held by this worker.
     #[cfg(any(test, feature = "integration"))]
     pub async fn tasks_running(&self) -> usize {

tests/local_connections.rs

@@ -0,0 +1,51 @@
+#[cfg(test)]
+mod tests {
+    use datafusion::common::internal_err;
+    use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+    use datafusion::physical_plan::collect;
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
+    use datafusion_distributed::test_utils::parquet::register_parquet_tables;
+    use datafusion_distributed::{
+        DefaultSessionBuilder, DistributedExt, DistributedMetricsFormat, NetworkBoundaryExt,
+        display_plan_ascii, rewrite_distributed_plan_with_metrics,
+    };
+    use std::sync::Arc;
+
+    /// During dynamic planning, if the planner decides that all the stages in the query are small
+    /// enough to fit in a single machine, it should co-locate them on in the machine that contains
+    /// the coordinating context, avoiding all network jumps.
+    #[tokio::test]
+    async fn all_local_connections_dynamic_planner() -> Result<(), Box<dyn std::error::Error>> {
+        let (mut ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await;
+        register_parquet_tables(&ctx).await?;
+        ctx.set_distributed_dynamic_task_count(true)?;
+        ctx.set_distributed_file_scan_config_bytes_per_partition(1024 * 1024)?;
+
+        let query =
+            r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#;
+
+        let df = ctx.sql(query).await?;
+        let plan = df.create_physical_plan().await?;
+        collect(Arc::clone(&plan), ctx.task_ctx()).await?;
+        let format = DistributedMetricsFormat::Aggregated;
+        let plan = rewrite_distributed_plan_with_metrics(plan, format).await?;
+        println!("{}", display_plan_ascii(plan.as_ref(), true));
+        plan.apply(|plan| {
+            if !plan.is_network_boundary() {
+                return Ok(TreeNodeRecursion::Continue);
+            };
+
+            let metrics = plan.metrics().unwrap();
+            let local_connections_used = metrics
+                .sum(|v| v.value().name() == "local_connections_used")
+                .map_or(0, |v| v.as_usize());
+            if local_connections_used == 0 {
+                return internal_err!("local_connections_used==0");
+            };
+
+            Ok(TreeNodeRecursion::Continue)
+        })?;
+
+        Ok(())
+    }
+}