mbutrovich commented on code in PR #2669: URL: https://github.com/apache/datafusion-ballista/pull/2669#discussion_r4417195342
########## 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: Should these tests check that the cluster really runs push-staged? Their results match pull-staged by design, so they would also pass if the setting were ignored. I tried that locally on the head commit. I hard-coded `PullStaged` in `new_standalone_scheduler_with_builder`, forced `push_staged` to `false` in the executor, and both tests still passed. A push-staged scheduler rejects `poll_work` with `FailedPrecondition` ([grpc.rs](https://github.com/apache/datafusion-ballista/blob/986c07d790c69ce148704bb33cbe2683e14e19a2/ballista/scheduler/src/scheduler_server/grpc.rs#L129-L134)), so the two-executor test could assert that after the cluster starts: ```rust use ballista_core::serde::protobuf::PollWorkParams; use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient; let mut scheduler = SchedulerGrpcClient::connect(format!("http://{host}:{port}";)) .await .unwrap(); let status = scheduler .poll_work(PollWorkParams::default()) .await .unwrap_err(); assert_eq!(status.code(), tonic::Code::FailedPrecondition); ``` ########## 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: Could this comment describe what happens if the receiver is dropped? Dropping the notifier does stop the loops, because `Shutdown::recv` returns once the broadcast sender is gone. Dropping `stop_receive` doesn't stop them. Instead, the next `StopExecutor` request panics at `stop_sender.send(force).await.unwrap()` ([executor_server.rs](https://github.com/apache/datafusion-ballista/blob/986c07d790c69ce148704bb33cbe2683e14e19a2/ballista/executor/src/executor_server.rs#L992-L993)). ```suggestion // The executor's loops stop when the notifier is dropped, and its // `stop_executor` handler panics if the stop channel's receiver is // dropped, so keep both for the life of the runtime. ``` ########## 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: Could the executor keep the gRPC listener bound instead of releasing the port and binding it again in `executor_server::startup`? Between the drop on line 117 and the bind in `startup`, any other bind to port 0 in the process can get the same number. The two new tests run in parallel and each binds several port-0 listeners, so this can happen in this test binary. When it does, `startup` fails with `AddrInUse` and the cluster doesn't come up. The repo already ran into this with `free_ports` in chaos-testing: "Releasing each before taking the next lets the kernel hand the same number back" ([cluster.rs](https://github.com/apache/datafusion-ballista/blob/c1de8e5b15604bce4800f5e396cd42a97ccee06b/chaos-testing/src/cluster.rs#L32-L38)). One option is a variant of `startup` that takes an already-bound `TcpIncoming`, with `startup` calling it. That variant would be a new public item in `executor_server`, so it might be better to keep it `pub(crate)` if the standalone executor is its only caller. -- 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]
