andygrove commented on code in PR #2669:
URL:
https://github.com/apache/datafusion-ballista/pull/2669#discussion_r4418632577
##########
ballista/executor/src/standalone.rs:
##########
@@ -98,12 +102,31 @@ pub async fn new_standalone_executor_from_builder(
let address = listener.local_addr()?;
info!("Ballista v{BALLISTA_VERSION} Rust Executor listening on
{address:?}");
+ let config = config_producer();
+ let max_message_size = config.ballista_grpc_client_max_message_size();
+ let push_staged = matches!(
+ config
+ .ballista_config()
+ .standalone_task_scheduling_policy()?,
+ ballista_core::config::TaskSchedulingPolicy::PushStaged
+ );
+
+ // A push-staged executor serves the scheduler's task launches on a gRPC
+ // port of its own, which `executor_server::startup` binds. Take a free
one.
+ let grpc_port = if push_staged {
+ std::net::TcpListener::bind("127.0.0.1:0")?
+ .local_addr()?
+ .port()
+ } else {
+ // TODO Make it configurable
+ 50020
+ };
Review Comment:
Good catch, thanks! The standalone executor now keeps its listener bound and
passes it to a new `pub(crate) startup_with_incoming`. `startup` binds as
before and delegates to it, so nothing new is public.
##########
ballista/executor/src/standalone.rs:
##########
@@ -143,14 +163,46 @@ pub async fn new_standalone_executor_from_builder(
)),
);
- tokio::spawn(execution_loop::poll_loop(
- scheduler,
- executor,
- codec,
- None,
- None,
- crate::health::ExecutorHealth::new(),
- ));
+ if push_staged {
+ let process_config = crate::executor_process::ExecutorProcessConfig {
+ bind_host: "127.0.0.1".to_owned(),
+ port: address.port(),
+ grpc_port,
+ vcores,
+ task_scheduling_policy:
+ ballista_core::config::TaskSchedulingPolicy::PushStaged,
+ grpc_max_decoding_message_size: max_message_size as u32,
+ grpc_max_encoding_message_size: max_message_size as u32,
+ ..Default::default()
+ };
+ let shutdown = crate::shutdown::ShutdownNotifier::new();
+ let (stop_send, stop_receive) = tokio::sync::mpsc::channel::<bool>(10);
+ crate::executor_server::startup(
+ scheduler,
+ Arc::new(process_config),
+ executor,
+ codec,
+ stop_send,
+ &shutdown,
+ crate::health::ExecutorHealth::new(),
+ )
+ .await?;
+ // The executor's loops stop when the notifier or the stop channel is
+ // dropped, so keep both for the life of the runtime.
Review Comment:
Thanks, applied your wording.
##########
ballista/client/tests/standalone_push_staged.rs:
##########
@@ -0,0 +1,145 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+mod common;
+
+/// End-to-end tests of a standalone cluster running push-staged, selected with
+/// `ballista.standalone.task_scheduling_policy`. Each query's results are
+/// compared with the same query run on a pull-staged standalone cluster.
+#[cfg(test)]
+#[cfg(feature = "standalone")]
+mod standalone_push_staged {
+ use crate::common::setup_test_cluster_with_executors;
+ use ballista::prelude::SessionContextExt;
+ use ballista_core::config::{
+ BALLISTA_STANDALONE_PARALLELISM,
BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY,
+ };
+ use ballista_core::extension::SessionConfigExt;
+ use datafusion::arrow::util::pretty::pretty_format_batches;
+ use datafusion::execution::{SessionState, SessionStateBuilder};
+ use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext};
+ use std::io::Write;
+ use std::path::Path;
+ use tempfile::TempDir;
+
+ /// Writes an `orders` table split over four files, so its scan has four
+ /// partitions, and a small `customers` table.
+ fn write_tables(dir: &Path) {
+ let orders = dir.join("orders");
+ std::fs::create_dir(&orders).unwrap();
+ for file in 0..4 {
+ let mut out =
+
std::fs::File::create(orders.join(format!("part-{file}.csv"))).unwrap();
+ writeln!(out, "o_id,o_cust,o_amount").unwrap();
+ for i in 0..250 {
+ let id = file * 250 + i;
+ writeln!(out, "{id},{},{}", id % 97, (id * 7) % 1000).unwrap();
+ }
+ }
+ let mut out =
std::fs::File::create(dir.join("customers.csv")).unwrap();
+ writeln!(out, "c_id,c_name,c_region").unwrap();
+ for id in 0..100 {
+ writeln!(out, "{id},customer-{id},region-{}", id % 5).unwrap();
+ }
+ }
+
+ fn state(policy: &str) -> SessionState {
+ let config = SessionConfig::new_with_ballista()
+ .with_target_partitions(4)
+ .set_str(BALLISTA_STANDALONE_PARALLELISM, "4")
+ .set_str(BALLISTA_STANDALONE_TASK_SCHEDULING_POLICY, policy);
+ SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build()
+ }
+
+ async fn register_tables(ctx: &SessionContext, dir: &Path) {
+ ctx.register_csv(
+ "orders",
+ dir.join("orders").to_str().unwrap(),
+ CsvReadOptions::new(),
+ )
+ .await
+ .unwrap();
+ ctx.register_csv(
+ "customers",
+ dir.join("customers.csv").to_str().unwrap(),
+ CsvReadOptions::new(),
+ )
+ .await
+ .unwrap();
+ }
+
+ async fn rows(ctx: &SessionContext, sql: &str) -> String {
+ let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+ pretty_format_batches(&batches).unwrap().to_string()
+ }
+
+ const QUERIES: &[&str] = &[
+ "SELECT count(*), sum(o_amount) FROM orders",
+ "SELECT o_cust, count(*) AS n, sum(o_amount) AS total \
+ FROM orders GROUP BY o_cust ORDER BY o_cust",
+ "SELECT c_region, count(*) AS n, sum(o_amount) AS total \
+ FROM orders JOIN customers ON o_cust = c_id \
+ GROUP BY c_region ORDER BY c_region",
+ ];
+
+ async fn expected(dir: &Path) -> Vec<String> {
+ let ctx = SessionContext::standalone_with_state(state("pull-staged"))
+ .await
+ .unwrap();
+ register_tables(&ctx, dir).await;
+ let mut out = Vec::new();
+ for sql in QUERIES {
+ out.push(rows(&ctx, sql).await);
+ }
+ out
+ }
+
+ #[tokio::test]
+ async fn push_staged_standalone_matches_pull_staged() {
+ let dir = TempDir::new().unwrap();
+ write_tables(dir.path());
+ let expected = expected(dir.path()).await;
+
+ let ctx = SessionContext::standalone_with_state(state("push-staged"))
+ .await
+ .unwrap();
+ register_tables(&ctx, dir.path()).await;
+ for (sql, expected) in QUERIES.iter().zip(&expected) {
+ assert_eq!(&rows(&ctx, sql).await, expected, "{sql}");
+ }
+ }
+
+ #[tokio::test]
+ async fn push_staged_cluster_of_two_executors_matches_pull_staged() {
+ let dir = TempDir::new().unwrap();
+ write_tables(dir.path());
+ let expected = expected(dir.path()).await;
+
+ let state = state("push-staged");
+ let (host, port) = setup_test_cluster_with_executors(&state, 2,
2).await;
Review Comment:
Nice find. The two-executor test now asserts `poll_work` gets
`FailedPrecondition`. I checked it against a pull-staged cluster and it fails
there with `InvalidArgument`, so it catches the setting being ignored.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
