Move ballista standalone mode to client

Author: edrevo

ballista/rust/client/Cargo.toml

@@ -27,8 +27,14 @@ edition = "2018"
 
 [dependencies]
 ballista-core = { path = "../core" }
+ballista-executor = { path = "../executor", optional = true }
+ballista-scheduler = { path = "../scheduler", optional = true }
 futures = "0.3"
 log = "0.4"
 tokio = "1.0"
 
 datafusion = { path = "../../../datafusion" }
+
+[features]
+default = []
+standalone = ["ballista-executor", "ballista-scheduler"]

ballista/rust/client/src/columnar_batch.rs

@@ -29,9 +29,7 @@ use datafusion::scalar::ScalarValue;
 pub type MaybeColumnarBatch = Result<Option<ColumnarBatch>>;
 
 /// Batch of columnar data.
-#[allow(dead_code)]
 #[derive(Debug, Clone)]
-
 pub struct ColumnarBatch {
     schema: Arc<Schema>,
     columns: HashMap<String, ColumnarValue>,
@@ -112,9 +110,7 @@ impl ColumnarBatch {
 }
 
 /// A columnar value can either be a scalar value or an Arrow array.
-#[allow(dead_code)]
 #[derive(Debug, Clone)]
-
 pub enum ColumnarValue {
     Scalar(ScalarValue, usize),
     Columnar(ArrayRef),

ballista/rust/client/src/context.rs

@@ -44,49 +44,81 @@ use futures::future;
 use futures::StreamExt;
 use log::{error, info};
 
-#[allow(dead_code)]
 struct BallistaContextState {
     /// Scheduler host
     scheduler_host: String,
     /// Scheduler port
     scheduler_port: u16,
     /// Tables that have been registered with this context
     tables: HashMap<String, LogicalPlan>,
-    /// General purpose settings
-    settings: HashMap<String, String>,
 }
 
 impl BallistaContextState {
-    pub fn new(
-        scheduler_host: String,
-        scheduler_port: u16,
-        settings: HashMap<String, String>,
-    ) -> Self {
+    pub fn new(scheduler_host: String, scheduler_port: u16) -> Self {
         Self {
             scheduler_host,
             scheduler_port,
             tables: HashMap::new(),
-            settings,
         }
     }
-}
 
-#[allow(dead_code)]
+    #[cfg(feature = "standalone")]
+    pub async fn new_standalone(
+        concurrent_tasks: usize,
+    ) -> ballista_core::error::Result<Self> {
+        info!("Running in local mode. Scheduler will be run in-proc");
+
+        let addr = ballista_scheduler::new_standalone_scheduler().await?;
+
+        let scheduler = loop {
+            match SchedulerGrpcClient::connect(format!(
+                "http://localhost:{}",
+                addr.port()
+            ))
+            .await
+            {
+                Err(_) => {
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    info!("Attempting to connect to in-proc scheduler...");
+                }
+                Ok(scheduler) => break scheduler,
+            }
+        };
+
+        ballista_executor::new_standalone_executor(scheduler, concurrent_tasks).await?;
+        Ok(Self {
+            scheduler_host: "localhost".to_string(),
+            scheduler_port: addr.port(),
+            tables: HashMap::new(),
+        })
+    }
+}
 
 pub struct BallistaContext {
     state: Arc<Mutex<BallistaContextState>>,
 }
 
 impl BallistaContext {
     /// Create a context for executing queries against a remote Ballista scheduler instance
-    pub fn remote(host: &str, port: u16, settings: HashMap<String, String>) -> Self {
-        let state = BallistaContextState::new(host.to_owned(), port, settings);
+    pub fn remote(host: &str, port: u16) -> Self {
+        let state = BallistaContextState::new(host.to_owned(), port);
 
         Self {
             state: Arc::new(Mutex::new(state)),
         }
     }
 
+    #[cfg(feature = "standalone")]
+    pub async fn standalone(
+        concurrent_tasks: usize,
+    ) -> ballista_core::error::Result<Self> {
+        let state = BallistaContextState::new_standalone(concurrent_tasks).await?;
+
+        Ok(Self {
+            state: Arc::new(Mutex::new(state)),
+        })
+    }
+
     /// Create a DataFrame representing a Parquet table scan
 
     pub fn read_parquet(&self, path: &str) -> Result<Arc<dyn DataFrame>> {
@@ -268,3 +300,15 @@ impl BallistaContext {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[tokio::test]
+    #[cfg(feature = "standalone")]
+    async fn test_standalone_mode() {
+        use super::*;
+        let context = BallistaContext::standalone(1).await.unwrap();
+        let df = context.sql("SELECT 1;").unwrap();
+        context.collect(&df.to_logical_plan()).await.unwrap();
+    }
+}

ballista/rust/executor/Cargo.toml

@@ -32,7 +32,6 @@ snmalloc = ["snmalloc-rs"]
 anyhow = "1"
 async-trait = "0.1.36"
 ballista-core = { path = "../core" }
-ballista-scheduler = { path = "../scheduler" }
 configure_me = "0.4.0"
 env_logger = "0.8"
 futures = "0.3"

ballista/rust/executor/executor_config_spec.toml

@@ -36,10 +36,6 @@ type = "u16"
 default = "50050"
 doc = "scheduler port"
 
-[[switch]]
-name = "local"
-doc = "Running in local mode will launch a standalone scheduler inside the executor process. This will create a single-executor cluster, and is useful for development scenarios."
-
 [[param]]
 name = "bind_host"
 type = "String"
@@ -69,8 +65,3 @@ name = "concurrent_tasks"
 type = "usize"
 default = "4"
 doc = "Max concurrent tasks."
-
-[[param]]
-name = "scheduler_data_path"
-type = "String"
-doc = "Path for standalone data"

ballista/rust/executor/src/execution_loop.rs

@@ -29,9 +29,10 @@ use ballista_core::serde::protobuf::{
     self, scheduler_grpc_client::SchedulerGrpcClient, task_status, FailedTask,
     PartitionId, PollWorkParams, PollWorkResult, TaskDefinition, TaskStatus,
 };
-use ballista_executor::executor::Executor;
 use protobuf::CompletedTask;
 
+use crate::executor::Executor;
+
 pub async fn poll_loop(
     mut scheduler: SchedulerGrpcClient<Channel>,
     executor: Arc<Executor>,

ballista/rust/executor/src/lib.rs

@@ -18,5 +18,9 @@
 //! Core executor logic for executing queries and storing results in memory.
 
 pub mod collect;
+pub mod execution_loop;
 pub mod executor;
 pub mod flight_service;
+
+mod standalone;
+pub use standalone::new_standalone_executor;

ballista/rust/executor/src/main.rs

@@ -17,14 +17,11 @@
 
 //! Ballista Rust executor binary.
 
-use std::{
-    net::{IpAddr, Ipv4Addr},
-    sync::Arc,
-};
+use std::sync::Arc;
 
 use anyhow::{Context, Result};
 use arrow_flight::flight_service_server::FlightServiceServer;
-use futures::future::MaybeDone;
+use ballista_executor::execution_loop;
 use log::info;
 use tempfile::TempDir;
 use tonic::transport::Server;
@@ -34,17 +31,11 @@ use ballista_core::serde::protobuf::{
     executor_registration, scheduler_grpc_client::SchedulerGrpcClient,
     ExecutorRegistration,
 };
-use ballista_core::{
-    print_version, serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer,
-    BALLISTA_VERSION,
-};
+use ballista_core::{print_version, BALLISTA_VERSION};
 use ballista_executor::executor::Executor;
 use ballista_executor::flight_service::BallistaFlightService;
-use ballista_scheduler::{state::StandaloneClient, SchedulerServer};
 use config::prelude::*;
 
-mod execution_loop;
-
 #[macro_use]
 extern crate configure_me;
 
@@ -82,11 +73,7 @@ async fn main() -> Result<()> {
         .parse()
         .with_context(|| format!("Could not parse address: {}", addr))?;
 
-    let scheduler_host = if opt.local {
-        "localhost".to_string()
-    } else {
-        opt.scheduler_host
-    };
+    let scheduler_host = opt.scheduler_host;
     let scheduler_port = opt.scheduler_port;
     let scheduler_url = format!("http://{}:{}", scheduler_host, scheduler_port);
 
@@ -109,58 +96,6 @@ async fn main() -> Result<()> {
         port: port as u32,
     };
 
-    if opt.local {
-        info!("Running in local mode. Scheduler will be run in-proc");
-
-        let client = match opt.scheduler_data_path {
-            Some(v) => StandaloneClient::try_new(v)
-                .context("Could not create standalone config backend")?,
-            None => StandaloneClient::try_new_temporary()
-                .context("Could not create standalone config backend")?,
-        };
-
-        let server = SchedulerGrpcServer::new(SchedulerServer::new(
-            Arc::new(client),
-            "ballista".to_string(),
-            IpAddr::V4(Ipv4Addr::LOCALHOST),
-        ));
-        let addr = format!("localhost:{}", scheduler_port);
-        let addr = addr
-            .parse()
-            .with_context(|| format!("Could not parse {}", addr))?;
-        info!(
-            "Ballista v{} Rust Scheduler listening on {:?}",
-            BALLISTA_VERSION, addr
-        );
-        let scheduler_future =
-            tokio::spawn(Server::builder().add_service(server).serve(addr));
-        let mut scheduler_result = futures::future::maybe_done(scheduler_future);
-
-        // Ensure scheduler is ready to receive connections
-        while SchedulerGrpcClient::connect(scheduler_url.clone())
-            .await
-            .is_err()
-        {
-            let scheduler_future = match scheduler_result {
-                MaybeDone::Future(f) => f,
-                MaybeDone::Done(Err(e)) => return Err(e).context("Tokio error"),
-                MaybeDone::Done(Ok(Err(e))) => {
-                    return Err(e).context("Scheduler failed to initialize correctly")
-                }
-                MaybeDone::Done(Ok(Ok(()))) => {
-                    return Err(anyhow::format_err!(
-                        "Scheduler unexpectedly finished successfully"
-                    ))
-                }
-                MaybeDone::Gone => {
-                    panic!("Received Gone from recently created MaybeDone")
-                }
-            };
-            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
-            scheduler_result = futures::future::maybe_done(scheduler_future);
-        }
-    }
-
     let scheduler = SchedulerGrpcClient::connect(scheduler_url)
         .await
         .context("Could not connect to scheduler")?;

ballista/rust/executor/src/standalone.rs

@@ -0,0 +1,55 @@
+use std::sync::Arc;
+
+use arrow_flight::flight_service_server::FlightServiceServer;
+use ballista_core::{
+    error::Result,
+    serde::protobuf::{scheduler_grpc_client::SchedulerGrpcClient, ExecutorRegistration},
+    BALLISTA_VERSION,
+};
+use log::info;
+use tempfile::TempDir;
+use tokio::net::TcpListener;
+use tonic::transport::{Channel, Server};
+use uuid::Uuid;
+
+use crate::{execution_loop, executor::Executor, flight_service::BallistaFlightService};
+
+pub async fn new_standalone_executor(
+    scheduler: SchedulerGrpcClient<Channel>,
+    concurrent_tasks: usize,
+) -> Result<()> {
+    let work_dir = TempDir::new()?
+        .into_path()
+        .into_os_string()
+        .into_string()
+        .unwrap();
+    let executor = Arc::new(Executor::new(&work_dir));
+
+    let service = BallistaFlightService::new(executor.clone());
+
+    let server = FlightServiceServer::new(service);
+    // Let the OS assign a random, free port
+    let listener = TcpListener::bind("localhost:0").await?;
+    let addr = listener.local_addr()?;
+    info!(
+        "Ballista v{} Rust Executor listening on {:?}",
+        BALLISTA_VERSION, addr
+    );
+    tokio::spawn(
+        Server::builder().add_service(server).serve_with_incoming(
+            tokio_stream::wrappers::TcpListenerStream::new(listener),
+        ),
+    );
+    let executor_meta = ExecutorRegistration {
+        id: Uuid::new_v4().to_string(), // assign this executor a unique ID
+        optional_host: None,
+        port: addr.port() as u32,
+    };
+    tokio::spawn(execution_loop::poll_loop(
+        scheduler,
+        executor,
+        executor_meta,
+        concurrent_tasks,
+    ));
+    Ok(())
+}

ballista/rust/scheduler/Cargo.toml

@@ -28,7 +28,7 @@ edition = "2018"
 [features]
 default = ["etcd", "sled"]
 etcd = ["etcd-client"]
-sled = ["sled_package"]
+sled = ["sled_package", "tokio-stream"]
 
 [dependencies]
 anyhow = "1"
@@ -48,6 +48,7 @@ rand = "0.8"
 serde = {version = "1", features = ["derive"]}
 sled_package = { package = "sled", version = "0.34", optional = true }
 tokio = { version = "1.0", features = ["full"] }
+tokio-stream = { version = "0.1", features = ["net"], optional = true }
 tonic = "0.4"
 tower = { version = "0.4" }
 warp = "0.3"