oxidelake_runtime/
cluster.rs

1//! Cluster mode on Apache DataFusion Ballista (ADR-0013).
2//!
3//! OxideLake does not write a scheduler. This module builds the pieces Ballista
4//! lets us plug in — a session builder that installs the placement rule for the
5//! declared cluster capability, a physical codec that ships `Gpu*Exec` nodes,
6//! and the executor function registry — and wires them into Ballista's
7//! scheduler and executor processes, plus an in-process standalone cluster for
8//! tests and the CLI.
9
10use std::net::SocketAddr;
11use std::sync::Arc;
12
13use ballista_core::extension::{
14    SessionConfigExt, ballista_aggregate_functions, ballista_scalar_functions,
15    ballista_window_functions,
16};
17use ballista_core::registry::BallistaFunctionRegistry;
18use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
19use ballista_core::serde::{
20    BallistaCodec, BallistaLogicalExtensionCodec, BallistaPhysicalExtensionCodec,
21};
22use ballista_core::{ConfigProducer, RuntimeProducer};
23use ballista_executor::executor_process::{ExecutorProcessConfig, start_executor_process};
24use ballista_executor::new_standalone_executor_from_builder;
25use ballista_scheduler::cluster::BallistaCluster;
26use ballista_scheduler::config::SchedulerConfig;
27use ballista_scheduler::scheduler_process::start_server;
28use ballista_scheduler::scheduler_server::SessionBuilder;
29use ballista_scheduler::standalone::new_standalone_scheduler_with_builder;
30use datafusion::execution::SessionStateBuilder;
31use datafusion::execution::runtime_env::RuntimeEnvBuilder;
32use datafusion::execution::session_state::SessionState;
33use datafusion::prelude::SessionConfig;
34use datafusion_proto::physical_plan::PhysicalExtensionCodec;
35use oxidelake_compute::oxide_udfs;
36use oxidelake_core::{BackendKind, EngineError};
37use oxidelake_planner::{HardwarePlacementRule, OxidePhysicalCodec, physical_optimizer_rules};
38use oxidelake_storage::{with_gpu_batch_size, with_pruning};
39
40fn ballista_error(err: impl std::fmt::Display) -> EngineError {
41    EngineError::execution(format!("ballista: {err}"))
42}
43
44/// The physical codec every OxideLake process installs: our `Gpu*Exec`
45/// encoding on top of Ballista's own codec (shuffle nodes).
46pub fn oxide_codec() -> Arc<dyn PhysicalExtensionCodec> {
47    Arc::new(OxidePhysicalCodec::new(Arc::new(
48        BallistaPhysicalExtensionCodec::default(),
49    )))
50}
51
52/// Ballista's combined logical + physical codec with [`oxide_codec`] as the physical half.
53pub fn ballista_codec() -> BallistaCodec {
54    BallistaCodec::new(
55        Arc::new(BallistaLogicalExtensionCodec::default()),
56        oxide_codec(),
57    )
58}
59
60/// Session configuration for every process: Ballista defaults, Parquet pruning
61/// on, and the OxideLake codec.
62pub fn session_config() -> SessionConfig {
63    with_pruning(SessionConfig::new_with_ballista())
64        .with_ballista_physical_extension_codec(oxide_codec())
65}
66
67/// Ballista `ConfigProducer` yielding [`session_config`].
68pub fn config_producer() -> ConfigProducer {
69    Arc::new(session_config)
70}
71
72/// Builds the scheduler-side session state: Ballista's defaults plus the
73/// placement rule targeting the declared cluster capability. This is where
74/// `Gpu*Exec` nodes enter a distributed plan.
75pub fn build_session_state(
76    config: SessionConfig,
77    target: BackendKind,
78) -> datafusion::error::Result<SessionState> {
79    let mut scalar_functions = ballista_scalar_functions();
80    scalar_functions.extend(oxide_udfs());
81    let config = if target.is_gpu() {
82        with_gpu_batch_size(config)
83    } else {
84        config
85    };
86    Ok(SessionStateBuilder::new()
87        .with_default_features()
88        .with_config(config)
89        .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build()?))
90        .with_scalar_functions(scalar_functions)
91        .with_aggregate_functions(ballista_aggregate_functions())
92        .with_window_functions(ballista_window_functions())
93        .with_physical_optimizer_rules(physical_optimizer_rules(HardwarePlacementRule::new(target)))
94        .build())
95}
96
97/// Ballista `SessionBuilder` wrapping [`build_session_state`].
98pub fn session_builder(target: BackendKind) -> SessionBuilder {
99    Arc::new(move |config: SessionConfig| build_session_state(config, target))
100}
101
102/// Ballista `RuntimeProducer` for executors.
103pub fn runtime_producer() -> RuntimeProducer {
104    Arc::new(|_config: &SessionConfig| Ok(Arc::new(RuntimeEnvBuilder::new().build()?)))
105}
106
107/// Functions registered on every executor: Ballista's defaults plus
108/// OxideLake's SQL UDFs (`l2_distance`, `cosine_distance`), so a distance
109/// projection that was not lowered to a `Gpu*Exec` still runs on workers.
110pub fn function_registry() -> BallistaFunctionRegistry {
111    let mut registry = BallistaFunctionRegistry::default();
112    for udf in oxide_udfs() {
113        registry.scalar_functions.insert(udf.name().to_owned(), udf);
114    }
115    registry
116}
117
118/// Scheduler configuration with OxideLake's hooks installed.
119pub fn scheduler_config(bind_host: &str, port: u16, target: BackendKind) -> SchedulerConfig {
120    SchedulerConfig {
121        bind_host: bind_host.to_owned(),
122        bind_port: port,
123        override_session_builder: Some(session_builder(target)),
124        override_config_producer: Some(config_producer()),
125        override_physical_codec: Some(oxide_codec()),
126        ..SchedulerConfig::default()
127    }
128}
129
130/// Runs a Ballista scheduler until it exits.
131pub async fn run_scheduler(config: SchedulerConfig) -> Result<(), EngineError> {
132    let address: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port)
133        .parse()
134        .map_err(|e| EngineError::plan(format!("invalid scheduler bind address: {e}")))?;
135    let config = Arc::new(config);
136    let cluster = BallistaCluster::new_from_config(&config)
137        .await
138        .map_err(ballista_error)?;
139    start_server(cluster, address, config)
140        .await
141        .map_err(ballista_error)
142}
143
144/// Executor configuration with OxideLake's hooks installed.
145pub fn executor_config(
146    scheduler_host: &str,
147    scheduler_port: u16,
148    port: u16,
149    grpc_port: u16,
150    concurrent_tasks: Option<usize>,
151    work_dir: Option<String>,
152) -> ExecutorProcessConfig {
153    let defaults = ExecutorProcessConfig::default();
154    ExecutorProcessConfig {
155        scheduler_host: scheduler_host.to_owned(),
156        scheduler_port,
157        port,
158        grpc_port,
159        concurrent_tasks: concurrent_tasks.unwrap_or(defaults.concurrent_tasks),
160        work_dir,
161        override_physical_codec: Some(oxide_codec()),
162        override_function_registry: Some(Arc::new(function_registry())),
163        override_config_producer: Some(config_producer()),
164        ..defaults
165    }
166}
167
168/// Runs a Ballista executor until it exits.
169pub async fn run_executor(config: ExecutorProcessConfig) -> Result<(), EngineError> {
170    start_executor_process(Arc::new(config))
171        .await
172        .map_err(ballista_error)
173}
174
175/// Starts an in-process scheduler plus `executors` executors on ephemeral
176/// localhost ports and returns the scheduler address. Everything runs on the
177/// current tokio runtime and stops with it.
178pub async fn start_standalone(
179    target: BackendKind,
180    executors: usize,
181    concurrent_tasks: usize,
182) -> Result<SocketAddr, EngineError> {
183    let addr = new_standalone_scheduler_with_builder(
184        session_builder(target),
185        config_producer(),
186        ballista_codec(),
187    )
188    .await
189    .map_err(ballista_error)?;
190    for _ in 0..executors.max(1) {
191        let client = SchedulerGrpcClient::connect(format!("http://{addr}"))
192            .await
193            .map_err(ballista_error)?;
194        new_standalone_executor_from_builder(
195            client,
196            concurrent_tasks.max(1),
197            config_producer(),
198            runtime_producer(),
199            ballista_codec(),
200            function_registry(),
201        )
202        .await
203        .map_err(ballista_error)?;
204    }
205    tracing::info!(%addr, executors, "standalone Ballista cluster started");
206    Ok(addr)
207}
oxidelake_runtime/cluster.rs

oxidelake_runtime/
cluster.rs
