diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index 101a6c96..7333ccd0 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -5,11 +5,13 @@ use crate::distributed_planner::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost, inject_network_boundaries, }; +use crate::events::DynamicStageBuiltHandlers; use crate::events::TaskCountAnnotation::{Desired, Maximum}; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; use crate::{ - BytesCounterMetric, LoadInfo, MaxGaugeMetric, NetworkBoundaryExt, NetworkCoalesceExec, Stage, + BytesCounterMetric, DynamicStageBuiltEvent, LoadInfo, MaxGaugeMetric, NetworkBoundaryExt, + NetworkCoalesceExec, Stage, }; use dashmap::DashMap; use datafusion::common::stats::Precision; @@ -69,6 +71,15 @@ pub(super) async fn prepare_dynamic_plan( .task_count(&input_stage.plan)? .merge(Desired(compute_based_task_count)); + let ev = DynamicStageBuiltEvent { + session_config: nb_ctx.cfg, + cost, + plan: &input_stage.plan, + }; + if let Some(response) = DynamicStageBuiltHandlers::handle(ev).transpose()? { + input_stage.plan = response.plan + }; + // Propagate the final task_count inferred based on runtime statistics and compute cost. // Here is where leaf nodes are scaled up by ScaleUpLeafNodeHandler, and the // plan is finally left ready for distribution. diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 800d973c..6e1887a9 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -3,17 +3,17 @@ use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; use crate::events::{ - DesiredTaskCountHandler, DesiredTaskCountHandlers, RouteTasksHandler, RouteTasksHandlers, - ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, WorkerPlanRewriteHandler, - WorkerPlanRewriteHandlers, + DesiredTaskCountHandler, DesiredTaskCountHandlers, DynamicStageBuiltHandlers, + RouteTasksHandler, RouteTasksHandlers, ScaleUpLeafNodeHandler, ScaleUpLeafNodeHandlers, + WorkerPlanRewriteHandler, WorkerPlanRewriteHandlers, }; use crate::passthrough_headers::set_passthrough_headers; 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, LocalWorkerContext, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, get_distributed_worker_resolver, + ChannelResolver, DistributedConfig, DynamicStageBuiltHandler, LocalWorkerContext, WorkUnitFeed, + WorkUnitFeedProvider, WorkerResolver, get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -737,6 +737,39 @@ pub trait DistributedExt: Sized { &mut self, handler: T, ); + + /// Register an event handler that fires every time a new stage is built during dynamic + /// planning. + /// + /// In this handler, users can inspect the plan that is about to be sent to a remote worker, + /// optimize it, or short circuit execution. + /// + /// ```rust + /// # use datafusion::common::{exec_err, Result}; + /// # use datafusion::execution::SessionStateBuilder; + /// # use datafusion_distributed::{DistributedExt, DynamicStageBuiltEvent, DynamicStageBuiltEventResponse}; + /// + /// fn handle_dynamic_stage_built(event: DynamicStageBuiltEvent) -> Option> { + /// if *event.cost.cpu.get_value().unwrap_or(&0) > 1024 * 1024 * 1024 { + /// return Some(exec_err!("Plan is too expensive to execute")) + /// } + /// None + /// } + /// + /// SessionStateBuilder::new() + /// .with_distributed_dynamic_stage_built_handler(handle_dynamic_stage_built); + /// ``` + fn with_distributed_dynamic_stage_built_handler( + self, + handler: T, + ) -> Self; + + /// Same as [DistributedExt::with_distributed_dynamic_stage_built_handler] but with an + /// in-place mutation. + fn set_distributed_dynamic_stage_built_handler( + &mut self, + handler: T, + ); } /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context. @@ -917,6 +950,10 @@ impl DistributedExt for SessionConfig { WorkerPlanRewriteHandlers::push_custom(self, Arc::new(h)); } + fn set_distributed_dynamic_stage_built_handler(&mut self, h: T) { + DynamicStageBuiltHandlers::push_custom(self, Arc::new(h)); + } + delegate! { to self { #[call(set_distributed_option_extension)] @@ -1024,6 +1061,10 @@ impl DistributedExt for SessionConfig { #[call(set_distributed_worker_plan_rewrite_handler)] #[expr($;self)] fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; + + #[call(set_distributed_dynamic_stage_built_handler)] + #[expr($;self)] + fn with_distributed_dynamic_stage_built_handler(mut self, h: H) -> Self; } } } @@ -1172,6 +1213,11 @@ impl DistributedExt for SessionStateBuilder { #[call(set_distributed_worker_plan_rewrite_handler)] #[expr($;self)] fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; + + fn set_distributed_dynamic_stage_built_handler(&mut self, h: H); + #[call(set_distributed_dynamic_stage_built_handler)] + #[expr($;self)] + fn with_distributed_dynamic_stage_built_handler(mut self, h: H) -> Self; } } } @@ -1322,6 +1368,11 @@ impl DistributedExt for SessionState { #[call(set_distributed_worker_plan_rewrite_handler)] #[expr($;self)] fn with_distributed_worker_plan_rewrite_handler(mut self, h: H) -> Self; + + fn set_distributed_dynamic_stage_built_handler(&mut self, h: H); + #[call(set_distributed_dynamic_stage_built_handler)] + #[expr($;self)] + fn with_distributed_dynamic_stage_built_handler(mut self, h: H) -> Self; } } } @@ -1465,6 +1516,11 @@ impl DistributedExt for SessionContext { #[call(set_distributed_worker_plan_rewrite_handler)] #[expr($;self)] fn with_distributed_worker_plan_rewrite_handler(self, h: H) -> Self; + + fn set_distributed_dynamic_stage_built_handler(&mut self, h: H); + #[call(set_distributed_dynamic_stage_built_handler)] + #[expr($;self)] + fn with_distributed_dynamic_stage_built_handler(self, h: H) -> Self; } } } diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 75e6813d..60ea2264 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -158,8 +158,8 @@ pub(crate) async fn inject_network_boundaries( #[derive(Clone)] pub(crate) struct InjectNetworkBoundaryContext<'a> { pub(crate) d_cfg: &'a DistributedConfig, + pub(crate) cfg: &'a SessionConfig, - cfg: &'a SessionConfig, worker_resolver: Arc, nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync), task_counts: &'a Mutex>, diff --git a/src/distributed_planner/statistics/cost.rs b/src/distributed_planner/statistics/cost.rs index dbb76e56..4b5330b2 100644 --- a/src/distributed_planner/statistics/cost.rs +++ b/src/distributed_planner/statistics/cost.rs @@ -1,3 +1,4 @@ +use crate::Cost; use crate::distributed_planner::statistics::complexity_cpu::complexity_cpu; use crate::distributed_planner::statistics::complexity_memory::complexity_memory; use crate::distributed_planner::statistics::complexity_network::complexity_network; @@ -8,13 +9,6 @@ use datafusion::physical_plan::{ExecutionPlan, Statistics}; use std::ops::AddAssign; use std::sync::Arc; -#[derive(Default, Debug)] -pub(crate) struct Cost { - pub(crate) cpu: Precision, - pub(crate) memory: Precision, - pub(crate) network: Precision, -} - impl AddAssign for Cost { fn add_assign(&mut self, rhs: Self) { self.cpu = sum_precision(self.cpu, rhs.cpu); diff --git a/src/distributed_planner/statistics/mod.rs b/src/distributed_planner/statistics/mod.rs index 59f1ec04..3f64e091 100644 --- a/src/distributed_planner/statistics/mod.rs +++ b/src/distributed_planner/statistics/mod.rs @@ -6,5 +6,4 @@ mod cost; mod default_bytes_for_datatype; mod plan_statistics; -#[allow(unused)] // will be used in a follow-up PR. pub(crate) use cost::calculate_cost; diff --git a/src/events/dynamic_stage_built.rs b/src/events/dynamic_stage_built.rs new file mode 100644 index 00000000..b35426dd --- /dev/null +++ b/src/events/dynamic_stage_built.rs @@ -0,0 +1,67 @@ +use crate::events::common::EventHandlerChain; +use datafusion::common::{Result, stats::Precision}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// Estimated cost associated to the fragment of the plan. Calculated based on the amount of data +/// that will flow through each node and the static cpu, memory and network complexity of the +/// different nodes based on the algorithms that will run inside them. +#[derive(Default, Debug, Clone, Copy)] +pub struct Cost { + /// Estimated cost of CPU measured in bytes. + pub cpu: Precision, + /// Estimated amount of memory consumed measured in bytes. + pub memory: Precision, + /// Estimated amount of data that will need to be transferred over the wire in bytes. + pub network: Precision, +} + +/// Event definition fired for each individual stage during dynamic planning. +#[derive(Copy, Clone)] +pub struct DynamicStageBuiltEvent<'a> { + /// Cost associated to the plan in the field below and all its children recursively until + /// network boundaries. + pub cost: Cost, + /// Plan that is about to be sent to a remote worker for runtime sampling. + pub plan: &'a Arc, + /// SessionConfig in scope at the moment of building the stage. + pub session_config: &'a SessionConfig, +} + +pub struct DynamicStageBuiltEventResponse { + /// A potentially modified plan (e.g., optimizations applied). + pub plan: Arc, +} + +impl DynamicStageBuiltEventResponse { + pub fn new(plan: Arc) -> Self { + Self { plan } + } +} + +pub trait DynamicStageBuiltHandler: Send + Sync + 'static { + fn handle(&self, ev: DynamicStageBuiltEvent) -> Option>; +} + +impl DynamicStageBuiltHandler for F +where + F: Send + Sync + 'static, + F: for<'a> Fn(DynamicStageBuiltEvent<'a>) -> Option>, +{ + fn handle(&self, ev: DynamicStageBuiltEvent) -> Option> { + self(ev) + } +} + +pub(crate) type DynamicStageBuiltHandlers = EventHandlerChain; + +impl DynamicStageBuiltHandlers { + pub(crate) fn handle( + ev: DynamicStageBuiltEvent, + ) -> Option> { + ev.session_config + .get_extension::()? + .find_map(|handler| handler.handle(ev)) + } +} diff --git a/src/events/mod.rs b/src/events/mod.rs index 9eb4b019..d8279e65 100644 --- a/src/events/mod.rs +++ b/src/events/mod.rs @@ -1,6 +1,7 @@ mod common; mod defaults; mod desired_task_count; +mod dynamic_stage_built; mod route_tasks; mod scale_up_leaf_node; mod worker_plan_rewrite; @@ -14,6 +15,10 @@ pub use desired_task_count::{ DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, TaskCountAnnotation, }; +pub(crate) use dynamic_stage_built::DynamicStageBuiltHandlers; +pub use dynamic_stage_built::{ + Cost, DynamicStageBuiltEvent, DynamicStageBuiltEventResponse, DynamicStageBuiltHandler, +}; pub(crate) use route_tasks::RouteTasksHandlers; pub use route_tasks::{RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler}; pub(crate) use scale_up_leaf_node::ScaleUpLeafNodeHandlers; diff --git a/src/lib.rs b/src/lib.rs index 8273be18..4af3fd47 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,10 +23,11 @@ pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, }; pub use events::{ - DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, RouteTasksEvent, - RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, - ScaleUpLeafNodeHandler, TaskCountAnnotation, WorkerPlanRewriteEvent, - WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler, + Cost, DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, + DynamicStageBuiltEvent, DynamicStageBuiltEventResponse, DynamicStageBuiltHandler, + RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, + ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler, TaskCountAnnotation, + WorkerPlanRewriteEvent, WorkerPlanRewriteEventResponse, WorkerPlanRewriteHandler, }; pub use execution_plans::{ BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec,