From 93c34c6d769dd56da02064c0abdaa060efe0259d Mon Sep 17 00:00:00 2001 From: hugocasa Date: Tue, 25 Nov 2025 17:25:34 +0100 Subject: [PATCH] feat: flow for loop squashing (#7107) * feat: flow for loop optimization * fix: make dedicated flows work with flow nodes * fix: transform dedicated inputs + better error handling + differentiate squash from same worker * fix builds * fix build * fix build * make dedicated worker private * update ee ref * frontend nits * nit * add trace to dedicated worker * update ee ref * Update ee-repo-ref.txt --------- Co-authored-by: windmill-internal-app[bot] --- backend/ee-repo-ref.txt | 2 +- backend/src/monitor.rs | 14 +- backend/tests/worker.rs | 2 + backend/windmill-api/src/flows.rs | 1 + backend/windmill-common/src/flows.rs | 7 + backend/windmill-queue/src/jobs.rs | 80 +- backend/windmill-queue/src/lib.rs | 2 +- backend/windmill-worker/src/ai/tools.rs | 4 +- backend/windmill-worker/src/bun_executor.rs | 40 +- .../windmill-worker/src/dedicated_worker.rs | 759 ------------------ .../src/dedicated_worker_oss.rs | 26 + backend/windmill-worker/src/deno_executor.rs | 15 +- backend/windmill-worker/src/lib.rs | 6 +- .../windmill-worker/src/python_executor.rs | 22 +- .../windmill-worker/src/result_processor.rs | 37 +- backend/windmill-worker/src/worker.rs | 239 ++++-- backend/windmill-worker/src/worker_flow.rs | 103 ++- .../windmill-worker/src/worker_lockfiles.rs | 5 +- .../lib/components/FlowGraphViewerStep.svelte | 4 +- .../components/flows/content/FlowLoop.svelte | 29 + .../flows/content/FlowWhileLoop.svelte | 31 +- .../lib/components/flows/map/MapItem.svelte | 2 +- openflow.openapi.yaml | 4 + 23 files changed, 548 insertions(+), 886 deletions(-) delete mode 100644 backend/windmill-worker/src/dedicated_worker.rs create mode 100644 backend/windmill-worker/src/dedicated_worker_oss.rs diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index ee9c8abbc4..bc6e0fecec 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -1de62d3f0901aac6ef9ddd7595e1032be56b1938 \ No newline at end of file +f9434efa127dadb0a04c87d6fc368cc289d93501 diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs index eaf83c2e4c..d85f5f3d58 100644 --- a/backend/src/monitor.rs +++ b/backend/src/monitor.rs @@ -2316,13 +2316,14 @@ async fn handle_zombie_jobs(db: &Pool, base_internal_url: &str, node_n ); let error_message = format!( - "Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {}, reason: {:?}).\nThis likely means that the job died on worker {}, OOM are a common reason for worker crashes.\nCheck the workers around the time of the last ping and the exit code if any.", - job.last_ping.unwrap_or_default(), - *ZOMBIE_JOB_TIMEOUT, - error_kind.to_string(), - job.worker.clone().unwrap_or_default(), - ); + "Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {}, reason: {:?}).\nThis likely means that the job died on worker {}, OOM are a common reason for worker crashes.\nCheck the workers around the time of the last ping and the exit code if any.", + job.last_ping.unwrap_or_default(), + *ZOMBIE_JOB_TIMEOUT, + error_kind.to_string(), + job.worker.clone().unwrap_or_default(), + ); let memory_peak = job.memory_peak.unwrap_or(0); + let (_, killpill_rx_never_used) = KillpillSender::new(1); let _ = handle_job_error( db, &client, @@ -2335,6 +2336,7 @@ async fn handle_zombie_jobs(db: &Pool, base_internal_url: &str, node_n "", node_name, send_result_never_used, + &killpill_rx_never_used, #[cfg(feature = "benchmark")] &mut windmill_common::bench::BenchmarkIter::new(), ) diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index 0a92dfc74d..bb18d6e397 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -217,6 +217,7 @@ async fn test_deno_flow(db: Pool) -> anyhow::Result<()> { iterator: InputTransform::Javascript { expr: "result".to_string() }, skip_failures: false, parallel: false, + squash: None, parallelism: None, modules: vec![FlowModule { id: "c".to_string(), @@ -396,6 +397,7 @@ async fn test_deno_flow_same_worker(db: Pool) -> anyhow::Result<()> { iterator: InputTransform::Static { value: windmill_common::worker::to_raw_value(&[1, 2, 3]) }, skip_failures: false, parallel: false, + squash: None, parallelism: None, modules: vec![ FlowModule { diff --git a/backend/windmill-api/src/flows.rs b/backend/windmill-api/src/flows.rs index 2a75045f9f..98130e2d2d 100644 --- a/backend/windmill-api/src/flows.rs +++ b/backend/windmill-api/src/flows.rs @@ -1583,6 +1583,7 @@ mod tests { skip_failures: true, parallel: false, parallelism: None, + squash: None, }), stop_after_if: Some(StopAfterIf { expr: "previous.isEmpty()".to_string(), diff --git a/backend/windmill-common/src/flows.rs b/backend/windmill-common/src/flows.rs index 41e02b7ee8..104f99b301 100644 --- a/backend/windmill-common/src/flows.rs +++ b/backend/windmill-common/src/flows.rs @@ -869,6 +869,8 @@ pub enum FlowModuleValue { parallel: bool, #[serde(skip_serializing_if = "Option::is_none")] parallelism: Option, + #[serde(skip_serializing_if = "Option::is_none")] + squash: Option, }, /// While loop node @@ -878,6 +880,8 @@ pub enum FlowModuleValue { modules_node: Option, #[serde(default = "default_false")] skip_failures: bool, + #[serde(skip_serializing_if = "Option::is_none")] + squash: Option, }, /// Branch-one node @@ -989,6 +993,7 @@ struct UntaggedFlowModuleValue { assets: Option>, tools: Option>, pass_flow_input_directly: Option, + squash: Option, } impl<'de> Deserialize<'de> for FlowModuleValue { @@ -1027,6 +1032,7 @@ impl<'de> Deserialize<'de> for FlowModuleValue { skip_failures: untagged.skip_failures.unwrap_or(true), parallel: untagged.parallel.unwrap_or(false), parallelism: untagged.parallelism, + squash: untagged.squash, }), "whileloopflow" => Ok(FlowModuleValue::WhileloopFlow { modules: untagged @@ -1034,6 +1040,7 @@ impl<'de> Deserialize<'de> for FlowModuleValue { .ok_or_else(|| serde::de::Error::missing_field("modules"))?, modules_node: untagged.modules_node, skip_failures: untagged.skip_failures.unwrap_or(false), + squash: untagged.squash, }), "branchone" => Ok(FlowModuleValue::BranchOne { branches: untagged diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index c050d56659..cc29b94f1b 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -24,6 +24,9 @@ use serde::{ser::SerializeMap, Serialize}; use serde_json::{json, value::RawValue}; use sqlx::{types::Json, Pool, Postgres, Transaction}; use sqlx::{Encode, PgExecutor}; +use tokio::sync::mpsc::Sender; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; use tokio::{sync::RwLock, time::sleep}; use ulid::Ulid; use uuid::Uuid; @@ -137,7 +140,7 @@ pub struct CanceledBy { pub reason: Option, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize)] pub struct JobCompleted { pub job: MiniCompletedJob, pub preprocessed_args: Option>>, @@ -151,6 +154,10 @@ pub struct JobCompleted { pub duration: Option, pub has_stream: Option, pub from_cache: Option, + #[serde(skip)] + pub flow_runners: Option>, + #[serde(skip)] + pub done_tx: Option>, } pub async fn cancel_single_job<'c>( @@ -2271,6 +2278,8 @@ pub struct JobAndPerms { pub parent_runnable_path: Option, pub token: String, pub precomputed_agent_info: Option, + #[serde(skip)] + pub flow_runners: Option>, } impl PulledJob { pub async fn get_job_and_perms(self, db: &DB) -> JobAndPerms { @@ -2302,6 +2311,7 @@ impl PulledJob { parent_runnable_path: self.parent_runnable_path, token, precomputed_agent_info: None, + flow_runners: None, } } } @@ -2492,6 +2502,8 @@ impl PulledJobResult { duration: None, has_stream: Some(false), from_cache: None, + flow_runners: None, + done_tx: None, }), ), PulledJobResult { job: Some(job), error_while_preprocessing: Some(e), .. } => Err( @@ -2511,6 +2523,8 @@ impl PulledJobResult { duration: None, has_stream: Some(false), from_cache: None, + flow_runners: None, + done_tx: None, }), ), PulledJobResult { job, .. } => Ok(job), @@ -5715,10 +5729,74 @@ async fn restarted_flows_resolution( )) } + + +// Wrapper struct to send both job and optional flow_runners to dedicated workers +pub struct DedicatedWorkerJob { + pub job: Arc, + pub flow_runners: Option>, + pub done_tx: Option>, +} + +#[derive(Debug)] +pub struct FlowRunners { + pub runners: HashMap>, + pub handles: Vec>, + pub job_id: Uuid, +} + +impl Drop for FlowRunners { + fn drop(&mut self) { + let total_runners = self.handles.len(); + tracing::info!("dropping {} flow runners for job {}", total_runners, self.job_id); + + // First, drop all senders to signal workers to stop gracefully + self.runners.clear(); + + // Spawn a background task to wait with timeout and abort if needed + let handles = std::mem::take(&mut self.handles); + let job_id = self.job_id; + + tokio::spawn(async move { + // Extract abort handles before consuming the join handles + let abort_handles: Vec<_> = handles.iter().map(|h| h.abort_handle()).collect(); + + // Wait up to 5 seconds for natural termination + let timeout_result = tokio::time::timeout( + tokio::time::Duration::from_secs(5), + futures::future::join_all(handles) + ).await; + + match timeout_result { + Ok(_) => { + tracing::info!("all {} flow runners for job {} terminated gracefully", total_runners, job_id); + } + Err(_) => { + // Timeout reached, abort only the handles that haven't finished + let mut aborted = 0; + for abort_handle in abort_handles { + if !abort_handle.is_finished() { + abort_handle.abort(); + aborted += 1; + } + } + let graceful = total_runners - aborted; + tracing::warn!( + "flow runners for job {}: {} terminated gracefully, {} aborted after 5s timeout", + job_id, graceful, aborted + ); + } + } + }); + } +} + #[derive(Debug, Serialize, Deserialize)] pub struct SameWorkerPayload { pub job_id: Uuid, pub recoverable: bool, + #[serde(skip)] + pub flow_runners: Option>, } pub async fn get_same_worker_job( diff --git a/backend/windmill-queue/src/lib.rs b/backend/windmill-queue/src/lib.rs index 3ea2e2c1aa..ffd102c081 100644 --- a/backend/windmill-queue/src/lib.rs +++ b/backend/windmill-queue/src/lib.rs @@ -6,7 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ -mod jobs; +pub mod jobs; #[cfg(feature = "private")] pub mod jobs_ee; pub mod jobs_oss; diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs index 478eca5e54..84304f232c 100644 --- a/backend/windmill-worker/src/ai/tools.rs +++ b/backend/windmill-worker/src/ai/tools.rs @@ -529,6 +529,7 @@ async fn execute_windmill_tool( &mut occupancy_metrics_spawn, &mut killpill_rx_spawn, None, + None, #[cfg(feature = "benchmark")] &mut bench_spawn, ) @@ -656,8 +657,9 @@ async fn handle_tool_execution_success( .. }) = send_result.as_ref() { + let result = result.clone(); ctx.job_completed_tx - .send(send_result.as_ref().unwrap().result.clone(), true) + .send(send_result.unwrap().result, true) .await .map_err(to_anyhow)?; result diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index b8723ad815..26c59b77ba 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -9,10 +9,8 @@ use serde_json::value::RawValue; use uuid::Uuid; use windmill_parser_ts::remove_pinned_imports; -use windmill_queue::{append_logs, CanceledBy, MiniPulledJob, PrecomputedAgentInfo}; -#[cfg(feature = "enterprise")] -use crate::common::build_envs_map; +use windmill_queue::{append_logs, CanceledBy, MiniPulledJob, PrecomputedAgentInfo}; use crate::{ common::{ @@ -38,12 +36,6 @@ use tokio::{fs::File, process::Command}; use tokio::io::AsyncReadExt; -#[cfg(feature = "enterprise")] -use tokio::sync::mpsc::Receiver; - -#[cfg(feature = "enterprise")] -use windmill_common::variables; - use windmill_common::{ error::{self, Result}, get_latest_hash_for_path, @@ -1567,10 +1559,18 @@ pub async fn get_common_bun_proc_envs(base_internal_url: Option<&str>) -> HashMa return bun_envs; } -#[cfg(feature = "enterprise")] -use crate::{dedicated_worker::handle_dedicated_process, JobCompletedSender}; +#[cfg(feature = "private")] +use crate::{ + common::build_envs_map, dedicated_worker_oss::handle_dedicated_process, JobCompletedSender, +}; +#[cfg(feature = "private")] +use tokio::sync::mpsc::Receiver; +#[cfg(feature = "private")] +use windmill_common::variables; +#[cfg(feature = "private")] +use windmill_queue::DedicatedWorkerJob; -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] pub async fn start_worker( requirements_o: Option, codebase: Option, @@ -1584,8 +1584,9 @@ pub async fn start_worker( script_path: &str, token: &str, job_completed_tx: JobCompletedSender, - jobs_rx: Receiver>, + jobs_rx: Receiver, killpill_rx: tokio::sync::broadcast::Receiver<()>, + client: windmill_common::client::AuthedClient, ) -> Result<()> { let mut logs = "".to_string(); let mut mem_peak: i32 = 0; @@ -1737,6 +1738,12 @@ BigInt.prototype.toJSON = function () {{ console.log('start'); +function getArgs(line) {{ + let {{ {spread} }} = JSON.parse(line) + {dates} + return [ {spread} ]; +}} + for await (const line of Readline.createInterface({{ input: process.stdin }})) {{ {print_lines} @@ -1744,9 +1751,8 @@ for await (const line of Readline.createInterface({{ input: process.stdin }})) { process.exit(0); }} try {{ - let {{ {spread} }} = JSON.parse(line) - {dates} - let res = await Main.main(...[ {spread} ]); + const args = getArgs(line); + const res = await Main.main(...args); console.log("wm_res[success]:" + JSON.stringify(res ?? null, (key, value) => typeof value === 'undefined' ? null : value)); }} catch (e) {{ console.log("wm_res[error]:" + JSON.stringify({{ message: e.message, name: e.name, stack: e.stack, line: line }})); @@ -1812,6 +1818,7 @@ for await (const line of Readline.createInterface({{ input: process.stdin }})) { db, &script_path, "nodejs", + client, ) .await } else { @@ -1838,6 +1845,7 @@ for await (const line of Readline.createInterface({{ input: process.stdin }})) { db, script_path, "bun", + client, ) .await } diff --git a/backend/windmill-worker/src/dedicated_worker.rs b/backend/windmill-worker/src/dedicated_worker.rs deleted file mode 100644 index e65655c94d..0000000000 --- a/backend/windmill-worker/src/dedicated_worker.rs +++ /dev/null @@ -1,759 +0,0 @@ -// use tokio::sync::mpsc; - -// pub fn create_dedicated_worker() { -// let (job_completed_tx, mut new_job) = mpsc::channel::(100); -// } - -use async_recursion::async_recursion; -use std::{collections::VecDeque, process::Stdio, sync::Arc}; -use tokio::sync::mpsc::Sender; -use tokio::{ - io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, - process::Command, - task::JoinHandle, -}; -use windmill_common::error::Error; -use windmill_common::flows::FlowValue; -use windmill_common::worker::WORKER_CONFIG; -use windmill_common::KillpillSender; -use windmill_common::{ - cache, error, - flows::{FlowModule, FlowModuleValue}, - scripts::{ScriptHash, ScriptLang}, - variables, - worker::to_raw_value, - DB, -}; -use windmill_queue::append_logs; -use windmill_queue::MiniPulledJob; - -use anyhow::Context; - -use crate::{common::start_child_process, JobCompletedSender, MAX_BUFFERED_DEDICATED_JOBS}; - -use futures::{future, Future}; -use std::{collections::HashMap, task::Poll}; - -use tokio::sync::mpsc::Receiver; - -fn conditional_polling( - fut: impl Future, - predicate: bool, -) -> impl Future { - let mut fut = Box::pin(fut); - future::poll_fn(move |cx| { - if predicate { - fut.as_mut().poll(cx) - } else { - Poll::Pending - } - }) -} - -async fn write_stdin(stdin: &mut tokio::process::ChildStdin, s: &str) -> error::Result<()> { - let _ = &stdin.write_all(format!("{s}\n").as_bytes()).await?; - stdin.flush().await.context("stdin flush")?; - Ok(()) -} - -#[cfg(feature = "enterprise")] -pub async fn handle_dedicated_process( - command_path: &String, - job_dir: &str, - context_envs: HashMap, - envs: HashMap, - reserved_variables: Vec, - common_bun_proc_envs: HashMap, - args: Vec<&str>, - mut killpill_rx: tokio::sync::broadcast::Receiver<()>, - job_completed_tx: JobCompletedSender, - token: &str, - mut jobs_rx: Receiver>, - worker_name: &str, - db: &DB, - script_path: &str, - mode: &str, -) -> std::result::Result<(), error::Error> { - //do not cache local dependencies - - use windmill_queue::{JobCompleted, MiniCompletedJob}; - - use crate::{handle_child::process_status, PROXY_ENVS}; - let cmd_name = format!("dedicated {command_path}"); - let mut child = { - let mut cmd = Command::new(command_path); - cmd.current_dir(job_dir) - .env_clear() - .envs(context_envs) - .envs(envs) - .envs(PROXY_ENVS.clone()) - .envs( - reserved_variables - .iter() - .map(|x| (x.name.clone(), x.value.clone())) - .collect::>(), - ) - .envs(common_bun_proc_envs) - .args(args) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()); - start_child_process(cmd, command_path, false).await? - }; - - let stdout = child - .stdout() - .take() - .expect("child did not have a handle to stdout"); - - let stderr = child - .stderr() - .take() - .expect("child did not have a handle to stderr"); - - let mut reader = BufReader::new(stdout).lines(); - - let mut err_reader = BufReader::new(stderr).lines(); - - let mut stdin = child - .stdin() - .take() - .expect("child did not have a handle to stdin"); - - // Ensure the child process is spawned in the runtime so it can - // make progress on its own while we await for any output. - let child = tokio::spawn(async move { - let status = Box::into_pin(child.wait()) - .await - .expect("child process encountered an error"); - if let Err(e) = process_status(&cmd_name, status, vec![]) { - tracing::error!("child exit status was not success: {e:#}"); - } else { - tracing::info!("child exit status was success"); - } - }); - - let mut jobs: VecDeque = VecDeque::with_capacity(MAX_BUFFERED_DEDICATED_JOBS); - // let mut i = 0; - // let mut j = 0; - let mut alive = true; - - let init_log = format!("dedicated worker {mode}: {worker_name}\n\n"); - let mut logs = init_log.clone(); - loop { - tokio::select! { - biased; - _ = killpill_rx.recv(), if alive => { - println!("received killpill for dedicated worker"); - alive = false; - if let Err(e) = write_stdin(&mut stdin, "end").await { - tracing::info!("Could not write end message to stdin: {e:?}") - } - stdin.flush().await.context("stdin flush")?; - }, - line = err_reader.next_line() => { - if let Some(line) = line.expect("line is ok") { - tracing::error!("stderr dedicated worker: {line}"); - logs.push_str("[stderr] "); - logs.push_str(&line); - logs.push_str("\n"); - } else { - tracing::info!("dedicated worker process exited {script_path}"); - let mut last_stdout = "".to_string(); - while let Some(line) = reader.next_line().await.ok().flatten() { - last_stdout = line; - last_stdout.push_str("\n"); - } - tracing::info!("Last stdout for {script_path}: {last_stdout}"); - break; - } - }, - line = reader.next_line() => { - // j += 1; - - if let Some(line) = line.expect("line is ok") { - if line == "start" { - tracing::info!("dedicated worker process started {script_path}"); - continue; - } - tracing::debug!("processed job: |{line}|"); - if line.starts_with("wm_res[") { - let job = jobs.pop_front().expect("pop"); - tracing::info!("job completed on dedicated worker {script_path}: {}", job.id); - match serde_json::from_str::>(&line.replace("wm_res[success]:", "").replace("wm_res[error]:", "")) { - Ok(result) => { - let result = Arc::new(result); - append_logs(&job.id, &job.workspace_id, logs.clone(), &db.into()).await; - if line.starts_with("wm_res[success]:") { - job_completed_tx.send_job(JobCompleted { job , result, result_columns: None, mem_peak: 0, canceled_by: None, success: true, cached_res_path: None, token: token.to_string(), duration: None, preprocessed_args: None, has_stream: Some(false), from_cache: None }, true).await.unwrap() - } else { - job_completed_tx.send_job(JobCompleted { job , result, result_columns: None, mem_peak: 0, canceled_by: None, success: false, cached_res_path: None, token: token.to_string(), duration: None, preprocessed_args: None, has_stream: Some(false), from_cache: None }, true).await.unwrap() - } - }, - Err(e) => { - tracing::error!("Could not deserialize job result `{line}`: {e:?}"); - job_completed_tx.send_job(JobCompleted { job , result: Arc::new(to_raw_value(&serde_json::json!({"error": format!("Could not deserialize job result `{line}`: {e:?}")}))), result_columns: None, mem_peak: 0, canceled_by: None, success: false, cached_res_path: None, token: token.to_string(), duration: None, preprocessed_args: None, has_stream: Some(false), from_cache: None }, true).await.unwrap(); - }, - }; - logs = init_log.clone(); - } else { - logs.push_str(&line); - logs.push_str("\n"); - } - } else { - tracing::info!("dedicated worker {script_path} process exited"); - let mut last_stderr = "".to_string(); - while let Some(line) = err_reader.next_line().await.ok().flatten() { - last_stderr = line; - last_stderr.push_str("\n"); - } - tracing::info!("Last stderr for {script_path}: {last_stderr}"); - break; - } - }, - job = conditional_polling(jobs_rx.recv(), alive && jobs.len() < MAX_BUFFERED_DEDICATED_JOBS) => { - // i += 1; - if let Some(job) = job { - let id = job.id; - let args = serde_json::to_string(&job.args).expect("serialize"); - jobs.push_back(MiniCompletedJob::from(job)); - tracing::info!("received job and adding to queue on dedicated worker for {script_path}: {} (queue_size: {})", id, jobs.len()); - - // write_stdin(&mut stdin, &serde_json::to_string(&job.args.unwrap_or_else(|| serde_json::json!({"x": job.id}))).expect("serialize")).await?; - write_stdin(&mut stdin, &args).await?; - stdin.flush().await.context("stdin flush")?; - // tracing::info!("wrote job to stdin for {script_path}: {} (queue_size: {})", id, jobs.len()); - } else { - tracing::debug!("job channel closed"); - alive = false; - if let Err(e) = write_stdin(&mut stdin, "end").await { - tracing::error!("Could not write end message to stdin: {e:?}") - } - } - } - } - } - - child - .await - .map_err(|e| anyhow::anyhow!("child process {script_path} encountered an error: {e:#}"))?; - - tracing::info!("dedicated worker {script_path} child process exited successfully"); - Ok(()) -} - -type DedicatedWorker = ( - String, - Sender>, - Option>, -); - -// spawn one dedicated worker per compatible steps of the flow, associating the node id to the dedicated worker channel send -#[async_recursion] -#[cfg(feature = "enterprise")] -async fn spawn_dedicated_workers_for_flow( - modules: &Vec, - w_id: &str, - path: &str, - killpill_tx: KillpillSender, - killpill_rx: &tokio::sync::broadcast::Receiver<()>, - db: &DB, - worker_dir: &str, - base_internal_url: &str, - worker_name: &str, - job_completed_tx: &JobCompletedSender, -) -> Vec { - let mut workers = vec![]; - let mut script_path_to_worker: HashMap>> = - HashMap::new(); - for module in modules.iter() { - let value = module.get_value(); - if let Ok(value) = value { - match &value { - FlowModuleValue::Script { path, hash, .. } => { - let key = format!( - "{}:{}", - path, - hash.clone() - .map(|x| x.to_string()) - .unwrap_or_else(|| "".to_string()) - ); - if let Some(sender) = script_path_to_worker.get(&key) { - workers.push((module.id.clone(), sender.clone(), None)); - } else { - if let Some(dedi_w) = spawn_dedicated_worker( - SpawnWorker::Script { path: path.to_string(), hash: hash.clone() }, - w_id, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - Some(module.id.clone()), - ) - .await - { - script_path_to_worker.insert(key, dedi_w.1.clone()); - workers.push(dedi_w); - } - } - } - FlowModuleValue::ForloopFlow { modules, .. } => { - let w = spawn_dedicated_workers_for_flow( - &modules, - w_id, - path, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - ) - .await; - workers.extend(w); - } - FlowModuleValue::WhileloopFlow { modules, .. } => { - let w = spawn_dedicated_workers_for_flow( - &modules, - w_id, - path, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - ) - .await; - workers.extend(w); - } - FlowModuleValue::BranchOne { branches, default, .. } => { - for modules in branches - .iter() - .map(|x| &x.modules) - .chain(std::iter::once(default)) - { - let w = spawn_dedicated_workers_for_flow( - &modules, - w_id, - path, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - ) - .await; - workers.extend(w); - } - } - FlowModuleValue::BranchAll { branches, .. } => { - for branch in branches { - let w = spawn_dedicated_workers_for_flow( - &branch.modules, - w_id, - path, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - ) - .await; - workers.extend(w); - } - } - FlowModuleValue::RawScript { content, lock, path: spath, language, .. } => { - if let Some(dedi_w) = spawn_dedicated_worker( - SpawnWorker::RawScript { - path: spath.clone().unwrap_or(path.to_string()), - content: content.to_string(), - lock: lock.clone(), - lang: language.clone(), - }, - w_id, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - Some(module.id.clone()), - ) - .await - { - workers.push(dedi_w); - } - } - FlowModuleValue::FlowScript { id, language, .. } => { - let spawn = cache::flow::fetch_script( - &windmill_common::worker::Connection::Sql(db.clone()), - *id, - ) - .await - .map(|data| SpawnWorker::RawScript { - path: "".to_string(), - content: data.code.clone(), - lock: data.lock.clone(), - lang: *language, - }); - match spawn { - Ok(spawn) => { - if let Some(dedi_w) = spawn_dedicated_worker( - spawn, - w_id, - killpill_tx.clone(), - killpill_rx, - db, - worker_dir, - base_internal_url, - worker_name, - job_completed_tx, - Some(module.id.clone()), - ) - .await - { - workers.push(dedi_w); - } - } - Err(err) => tracing::error!( - "failed to get script for module: {:?}, err: {:?}", - module, - err - ), - } - } - FlowModuleValue::Flow { .. } => (), - FlowModuleValue::Identity => (), - FlowModuleValue::AIAgent { .. } => (), - } - } else { - tracing::error!("failed to get value for module: {:?}", module); - } - } - workers -} - -pub async fn create_dedicated_worker_map( - killpill_tx: &KillpillSender, - killpill_rx: &tokio::sync::broadcast::Receiver<()>, - db: &DB, - worker_dir: &str, - base_internal_url: &str, - worker_name: &str, - job_completed_tx: &JobCompletedSender, -) -> ( - HashMap>>, - bool, - Vec>, -) { - let mut dedicated_handles = vec![]; - if let Some(_wp) = WORKER_CONFIG.read().await.dedicated_worker.clone() { - let mut hm = HashMap::new(); - let is_flow_worker; - if let Some(flow_path) = _wp.path.strip_prefix("flow/") { - is_flow_worker = true; - let value = sqlx::query_scalar!( - "SELECT flow_version.value AS \"value!: sqlx::types::Json>\" - FROM flow - LEFT JOIN flow_version - ON flow_version.id = flow.versions[array_upper(flow.versions, 1)] - WHERE flow.path = $1 AND flow.workspace_id = $2", - flow_path, - _wp.workspace_id - ) - .fetch_optional(db) - .await; - if let Ok(v) = value { - if let Some(v) = v { - let value = serde_json::from_str::(v.get()).map_err(|err| { - Error::internal_err(format!( - "could not convert json to flow for {flow_path}: {err:?}" - )) - }); - if let Ok(flow) = value { - let workers = spawn_dedicated_workers_for_flow( - &flow.modules, - &_wp.workspace_id, - &_wp.path, - killpill_tx.clone(), - &killpill_rx, - db, - &worker_dir, - base_internal_url, - &worker_name, - &job_completed_tx, - ) - .await; - workers.into_iter().for_each(|(path, sender, handle)| { - tracing::info!("spawned dedicated worker for flow: {}", path.as_str()); - if let Some(h) = handle { - dedicated_handles.push(h); - } - hm.insert(path, sender); - }); - } - } else { - tracing::error!( - "flow present but value not found for dedicated worker. {}", - flow_path - ); - } - } else { - tracing::error!("flow not found for dedicated worker: {}. Waiting for dependency job and expected to restart.", flow_path); - } - } else { - is_flow_worker = false; - if let Some((path, sender, handle)) = spawn_dedicated_worker( - SpawnWorker::Script { path: _wp.path.clone(), hash: None }, - &_wp.workspace_id, - killpill_tx.clone(), - &killpill_rx, - db, - &worker_dir, - base_internal_url, - &worker_name, - &job_completed_tx, - None, - ) - .await - { - if let Some(h) = handle { - dedicated_handles.push(h); - } - hm.insert(path, sender); - } else { - tracing::error!( - "failed to spawn dedicated worker for {}, script not found", - _wp.path - ); - } - } - (hm, is_flow_worker, dedicated_handles) - } else { - (HashMap::new(), false, dedicated_handles) - } -} - -#[derive(Debug, Clone)] -pub enum SpawnWorker { - Script { path: String, hash: Option }, - RawScript { path: String, content: String, lock: Option, lang: ScriptLang }, -} - -// spawn one dedicated worker and return the key, the channel sender and the join handle -// note that for it will return none for language that do not support dedicated workers -// note that go using cache binary does not need dedicated workers so all languages are supported -#[cfg(feature = "enterprise")] -async fn spawn_dedicated_worker( - sw: SpawnWorker, - w_id: &str, - killpill_tx: KillpillSender, - killpill_rx: &tokio::sync::broadcast::Receiver<()>, - db: &DB, - worker_dir: &str, - base_internal_url: &str, - worker_name: &str, - job_completed_tx: &JobCompletedSender, - node_id: Option, -) -> Option { - use windmill_common::{ - error::Error, - scripts::{ScriptHash, ScriptLang}, - utils::rd_string, - }; - use windmill_queue::MiniPulledJob; - - use crate::{build_envs, get_script_content_by_hash, ContentReqLangEnvs}; - - #[cfg(not(feature = "enterprise"))] - { - tracing::error!("Dedicated worker is an enterprise feature"); - killpill_tx.send(()).expect("send"); - return None; - } - - #[cfg(feature = "enterprise")] - { - let (dedicated_worker_tx, dedicated_worker_rx) = tokio::sync::mpsc::channel::< - std::sync::Arc, - >(MAX_BUFFERED_DEDICATED_JOBS); - let killpill_rx = killpill_rx.resubscribe(); - let db2 = db.clone(); - let base_internal_url = base_internal_url.to_string(); - let worker_name = worker_name.to_string(); - let job_completed_tx = job_completed_tx.clone(); - let job_dir = format!( - "{}/dedicated{}", - worker_dir, - node_id - .as_ref() - .map(|x| format!("-{x}")) - .unwrap_or_else(|| "".to_string()) - ); - tokio::fs::create_dir_all(&job_dir) - .await - .expect("create dir"); - - let path = match &sw { - SpawnWorker::RawScript { path, .. } => path.to_string(), - SpawnWorker::Script { path, .. } => path.to_string(), - }; - - let path2 = path.clone(); - let w_id = w_id.to_string(); - - let (content, lock, language, envs, codebase) = match sw.clone() { - SpawnWorker::Script { path, hash } => { - let q = if let Some(hash) = hash { - get_script_content_by_hash(&hash, &w_id, &db2.into()) - .await - .map(|r: ContentReqLangEnvs| { - Some((r.content, r.lockfile, r.language, r.envs, r.codebase)) - }) - } else { - sqlx::query_as::<_, (String, Option, Option, Option>, bool, Option)>( - "SELECT content, lock, language, envs, codebase IS NOT NULL, hash FROM script WHERE path = $1 AND workspace_id = $2 AND - archived = false AND lock IS not NULL AND lock_error_logs IS NULL ORDER BY created_at DESC LIMIT 1", - ) - .bind(&path) - .bind(&w_id) - .fetch_optional(&db2) - .await - .map_err(|e| Error::internal_err(format!("expected content and lock: {e:#}"))) - .map(|x| x.map(|y| (y.0, y.1, y.2, y.3, if y.4 { y.5.map(|z| z.to_string()) } else { None }))) - }; - if let Ok(q) = q { - if let Some(wp) = q { - wp - } else { - tracing::error!( - "Failed to fetch script `{}` in workspace {} for dedicated worker.", - path, - w_id - ); - return None; - } - } else { - tracing::error!("Failed to fetch script for dedicated worker"); - killpill_tx.send(); - return None; - } - } - SpawnWorker::RawScript { content, lock, lang, .. } => { - (content, lock, Some(lang), None, None) - } - }; - - match language { - Some(ScriptLang::Python3) | Some(ScriptLang::Bun) | Some(ScriptLang::Deno) => {} - _ => return None, - } - - let db = db.clone(); - let handle = tokio::spawn(async move { - let token = { - let token = rd_string(32); - if let Err(e) = sqlx::query_scalar!( - "INSERT INTO token - (token, label, super_admin, email) - VALUES ($1, $2, $3, $4)", - token, - "dedicated_worker", - true, - "dedicated_worker@windmill.dev" - ) - .execute(&db) - .await - { - tracing::error!("failed to create token for dedicated worker: {:?}", e); - killpill_tx.clone().send(); - }; - token - }; - - let worker_envs = build_envs(envs.as_ref()).expect("failed to build envs"); - - if let Err(e) = match language { - Some(ScriptLang::Python3) => { - #[cfg(not(feature = "python"))] - { - tracing::error!("Python requires the python feature to be enabled"); - killpill_tx.send(); - return; - } - - #[cfg(feature = "python")] - crate::python_executor::start_worker( - lock.as_ref(), - &db, - &content, - &base_internal_url, - &job_dir, - &worker_name, - worker_envs, - &w_id, - &path, - &token, - job_completed_tx, - dedicated_worker_rx, - killpill_rx, - ) - .await - } - Some(ScriptLang::Bun) => { - crate::bun_executor::start_worker( - lock, - codebase, - &db, - &content, - &base_internal_url, - &job_dir, - &worker_name, - worker_envs, - &w_id, - &path, - &token, - job_completed_tx, - dedicated_worker_rx, - killpill_rx, - ) - .await - } - Some(ScriptLang::Deno) => { - crate::deno_executor::start_worker( - &content, - &base_internal_url, - &job_dir, - &worker_name, - worker_envs, - &w_id, - &path, - &token, - job_completed_tx, - dedicated_worker_rx, - killpill_rx, - &db, - ) - .await - } - _ => unreachable!("Non supported language for dedicated worker"), - } { - tracing::error!("error in dedicated worker for {sw:#?}: {:?}", e); - }; - killpill_tx.clone().send(); - }); - return Some((node_id.unwrap_or(path2), dedicated_worker_tx, Some(handle))); - // (Some(dedi_path), Some(dedicated_worker_tx), Some(handle)) - } -} diff --git a/backend/windmill-worker/src/dedicated_worker_oss.rs b/backend/windmill-worker/src/dedicated_worker_oss.rs new file mode 100644 index 0000000000..09e86193e0 --- /dev/null +++ b/backend/windmill-worker/src/dedicated_worker_oss.rs @@ -0,0 +1,26 @@ +#[cfg(feature = "private")] +// #[allow(unused)] +pub use crate::dedicated_worker_ee::*; + +#[cfg(not(feature = "private"))] +pub async fn spawn_flow_module_runners( + _job: &windmill_queue::MiniPulledJob, + _module: &windmill_common::flows::FlowModule, + _failure_module: Option<&Box>, + _killpill_rx: &tokio::sync::broadcast::Receiver<()>, + _db: &windmill_common::db::DB, + _worker_dir: &str, + _base_internal_url: &str, + _worker_name: &str, + _job_completed_tx: &crate::JobCompletedSender, +) -> windmill_common::error::Result<( + std::collections::HashMap< + String, + tokio::sync::mpsc::Sender, + >, + Vec>, +)> { + Err(windmill_common::error::Error::internal_err( + "Squashing flow loops is not available in OSS".to_string(), + )) +} diff --git a/backend/windmill-worker/src/deno_executor.rs b/backend/windmill-worker/src/deno_executor.rs index 332cd142db..e80eec265d 100644 --- a/backend/windmill-worker/src/deno_executor.rs +++ b/backend/windmill-worker/src/deno_executor.rs @@ -514,13 +514,14 @@ async fn build_import_map( Ok(()) as error::Result<()> } -#[cfg(feature = "enterprise")] -use crate::{dedicated_worker::handle_dedicated_process, JobCompletedSender}; - -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] +use crate::{dedicated_worker_oss::handle_dedicated_process, JobCompletedSender}; +#[cfg(feature = "private")] use tokio::sync::mpsc::Receiver; +#[cfg(feature = "private")] +use windmill_queue::DedicatedWorkerJob; -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] pub async fn start_worker( inner_content: &str, base_internal_url: &str, @@ -531,9 +532,10 @@ pub async fn start_worker( script_path: &str, token: &str, job_completed_tx: JobCompletedSender, - jobs_rx: Receiver>, + jobs_rx: Receiver, killpill_rx: tokio::sync::broadcast::Receiver<()>, db: &sqlx::Pool, + client: windmill_common::client::AuthedClient, ) -> Result<()> { use windmill_common::variables; @@ -654,6 +656,7 @@ for await (const chunk of Deno.stdin.readable) {{ db, script_path, "deno", + client, ) .await } diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs index 4bf60f6d4a..c5db591c42 100644 --- a/backend/windmill-worker/src/lib.rs +++ b/backend/windmill-worker/src/lib.rs @@ -23,8 +23,10 @@ mod bun_executor; pub mod common; mod config; mod csharp_executor; -#[cfg(feature = "enterprise")] -mod dedicated_worker; + +#[cfg(feature = "private")] +mod dedicated_worker_ee; +mod dedicated_worker_oss; mod deno_executor; #[cfg(feature = "duckdb")] mod duckdb_executor; diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index bc84a2e19f..b417014d85 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -18,6 +18,8 @@ use tokio::{ sync::Semaphore, task, }; +use windmill_queue::MiniPulledJob; + use uuid::Uuid; #[cfg(all(feature = "enterprise", feature = "parquet", unix))] use windmill_common::ee_oss::{get_license_plan, LicensePlan}; @@ -546,7 +548,6 @@ pub async fn handle_python_job( precomputed_agent_info: Option, has_stream: &mut bool, ) -> windmill_common::error::Result> { - let script_path = crate::common::use_flow_root_path(job.runnable_path()); let annotations = PythonAnnotations::parse(inner_content); @@ -2130,16 +2131,16 @@ fn get_result_postprocessor<'a>(skip: bool) -> &'a str { } } -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] use crate::JobCompletedSender; -#[cfg(feature = "enterprise")] -use crate::{common::build_envs_map, dedicated_worker::handle_dedicated_process}; -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] +use crate::{common::build_envs_map, dedicated_worker_oss::handle_dedicated_process}; +#[cfg(feature = "private")] use windmill_common::variables; +#[cfg(feature = "private")] +use windmill_queue::DedicatedWorkerJob; -use windmill_queue::MiniPulledJob; - -#[cfg(feature = "enterprise")] +#[cfg(feature = "private")] pub async fn start_worker( requirements_o: Option<&String>, db: &sqlx::Pool, @@ -2152,10 +2153,12 @@ pub async fn start_worker( script_path: &str, token: &str, job_completed_tx: JobCompletedSender, - jobs_rx: tokio::sync::mpsc::Receiver>, + jobs_rx: tokio::sync::mpsc::Receiver, killpill_rx: tokio::sync::broadcast::Receiver<()>, + client: windmill_common::client::AuthedClient, ) -> error::Result<()> { use crate::{PyV, PyVAlias}; + tracing::info!("script path: {}", script_path); let mut mem_peak: i32 = 0; let mut canceled_by: Option = None; @@ -2346,6 +2349,7 @@ for line in sys.stdin: db, script_path, "python", + client, ) .await } diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs index 1e168d3a30..e5a7a0052c 100644 --- a/backend/windmill-worker/src/result_processor.rs +++ b/backend/windmill-worker/src/result_processor.rs @@ -31,8 +31,8 @@ use windmill_common::{ use windmill_common::bench::{BenchmarkInfo, BenchmarkIter}; use windmill_queue::{ - append_logs, get_mini_completed_job, CanceledBy, JobCompleted, MiniCompletedJob, MiniPulledJob, - ValidableJson, WrappedError, INIT_SCRIPT_TAG, + append_logs, get_mini_completed_job, CanceledBy, FlowRunners, JobCompleted, MiniCompletedJob, + MiniPulledJob, ValidableJson, WrappedError, INIT_SCRIPT_TAG, }; use serde_json::{json, value::RawValue, Value}; @@ -66,6 +66,7 @@ async fn process_jc( same_worker_tx: Option<&SameWorkerSender>, job_completed_sender: &JobCompletedSender, stats_map: &JobStatsMap, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) { let success: bool = jc.success; @@ -155,6 +156,7 @@ async fn process_jc( same_worker_tx, &worker_name, job_completed_sender.clone(), + killpill_rx, #[cfg(feature = "benchmark")] bench, ) @@ -285,6 +287,7 @@ pub fn start_background_processor( Some(&same_worker_tx), &job_completed_sender, &stats_map, + &killpill_rx, #[cfg(feature = "benchmark")] &mut bench, ) @@ -350,6 +353,8 @@ pub fn start_background_processor( stop_early_override, &worker_name, job_completed_sender.clone(), + None, + &killpill_rx, #[cfg(feature = "benchmark")] &mut bench, ) @@ -413,6 +418,7 @@ pub async fn process_result( conn: &Connection, duration: Option, has_stream: bool, + flow_runners: Option>, ) -> error::Result { match result { Ok(result) => { @@ -431,6 +437,8 @@ pub async fn process_result( duration, has_stream: Some(has_stream), from_cache: None, + flow_runners, + done_tx: None, }, ) .with_context(windmill_common::otel_oss::otel_ctx()) @@ -495,6 +503,8 @@ pub async fn process_result( duration, has_stream: Some(has_stream), from_cache: None, + flow_runners, + done_tx: None, }, ) .with_context(windmill_common::otel_oss::otel_ctx()) @@ -512,6 +522,7 @@ pub async fn handle_receive_completed_job( same_worker_tx: Option<&SameWorkerSender>, worker_name: &str, job_completed_tx: JobCompletedSender, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) -> Option> { let token = jc.token.clone(); @@ -529,6 +540,7 @@ pub async fn handle_receive_completed_job( same_worker_tx.clone(), worker_name, job_completed_tx.clone(), + killpill_rx, #[cfg(feature = "benchmark")] bench, ) @@ -548,6 +560,7 @@ pub async fn handle_receive_completed_job( &worker_dir, worker_name, job_completed_tx, + killpill_rx, #[cfg(feature = "benchmark")] bench, ) @@ -571,6 +584,8 @@ pub async fn process_completed_job( preprocessed_args, has_stream, from_cache, + flow_runners, + done_tx, .. }: JobCompleted, client: &AuthedClient, @@ -579,6 +594,7 @@ pub async fn process_completed_job( same_worker_tx: Option<&SameWorkerSender>, worker_name: &str, job_completed_tx: JobCompletedSender, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) -> error::Result>> { if success { @@ -661,12 +677,19 @@ pub async fn process_completed_job( None, worker_name, job_completed_tx, + flow_runners, + killpill_rx, #[cfg(feature = "benchmark")] bench, ) .warn_after_seconds(10) .await?; add_time!(bench, "updated flow status END"); + if let Some(done_tx) = done_tx { + done_tx + .send(()) + .expect("done receiver should still be alive"); + } return Ok(r); } } @@ -707,11 +730,18 @@ pub async fn process_completed_job( None, worker_name, job_completed_tx, + flow_runners, + killpill_rx, #[cfg(feature = "benchmark")] bench, ) .warn_after_seconds(10) .await?; + if let Some(done_tx) = done_tx { + done_tx + .send(()) + .expect("done receiver should still be alive"); + } return Ok(r); } } @@ -761,6 +791,7 @@ pub async fn handle_job_error( worker_dir: &str, worker_name: &str, job_completed_tx: JobCompletedSender, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) { let err_string = format!("{}: {}", err.name(), err.to_string()); @@ -809,6 +840,8 @@ pub async fn handle_job_error( None, worker_name, job_completed_tx.clone(), + None, + killpill_rx, #[cfg(feature = "benchmark")] bench, ) diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 9b47aa5113..2919243dbc 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -11,6 +11,7 @@ use anyhow::anyhow; use futures::TryFutureExt; +use tokio::sync::Mutex; use tokio::time::sleep; use tokio::time::timeout; use windmill_common::client::AuthedClient; @@ -59,6 +60,8 @@ use std::{ time::Duration, }; use windmill_parser::MainArgSignature; +use windmill_queue::DedicatedWorkerJob; +use windmill_queue::FlowRunners; use windmill_queue::MiniCompletedJob; use windmill_queue::PulledJobResultToJobErr; @@ -167,8 +170,8 @@ use crate::duckdb_executor::do_duckdb; #[cfg(all(feature = "enterprise", feature = "oracledb"))] use crate::oracledb_executor::do_oracledb; -#[cfg(feature = "enterprise")] -use crate::dedicated_worker::create_dedicated_worker_map; +#[cfg(all(feature = "private", feature = "enterprise"))] +use crate::dedicated_worker_oss::create_dedicated_worker_map; #[cfg(feature = "enterprise")] use crate::snowflake_executor::do_snowflake; @@ -532,6 +535,8 @@ lazy_static::lazy_static! { .ok() .and_then(|x| x.parse::().ok()) .unwrap_or(1000); + + pub static ref FLOW_RUNNER_RUNNING: Mutex = Mutex::new(false); } type Envs = Vec<(String, String)>; @@ -558,14 +563,14 @@ lazy_static::lazy_static! { #[derive(Debug)] pub enum NextJob { - Sql(PulledJob), + Sql { job: PulledJob, flow_runners: Option> }, Http(JobAndPerms), } impl NextJob { pub fn job(self) -> MiniPulledJob { match self { - NextJob::Sql(job) => job.job, + NextJob::Sql { job, .. } => job.job, NextJob::Http(job) => job.job, } } @@ -575,7 +580,7 @@ impl std::ops::Deref for NextJob { type Target = MiniPulledJob; fn deref(&self) -> &Self::Target { match self { - NextJob::Sql(job) => &job.job, + NextJob::Sql { job, .. } => &job.job, NextJob::Http(job) => &job.job, } } @@ -733,10 +738,10 @@ impl JobCompletedSender { impl SameWorkerSender { pub async fn send( &self, - payload: SameWorkerPayload, + message: SameWorkerPayload, ) -> Result<(), tokio::sync::mpsc::error::SendError> { self.1.fetch_add(1, Ordering::Relaxed); - self.0.send(payload).await + self.0.send(message).await } } @@ -829,18 +834,36 @@ fn add_outstanding_wait_time( async fn extract_job_and_perms(job: NextJob, conn: &Connection) -> JobAndPerms { match (job, conn) { - (NextJob::Sql(job), Connection::Sql(db)) => job.get_job_and_perms(db).await, - (NextJob::Sql(_), Connection::Http(_)) => panic!("sql job on http connection"), + (NextJob::Sql { job, flow_runners, .. }, Connection::Sql(db)) => { + JobAndPerms { flow_runners, ..job.get_job_and_perms(db).await } + } + (NextJob::Sql { .. }, Connection::Http(_)) => panic!("sql job on http connection"), (NextJob::Http(job), _) => job, } } -fn create_span(arc_job: &Arc, worker_name: &str, hostname: &str) -> Span { - let span = tracing::span!(tracing::Level::INFO, "job", - job_id = %arc_job.id, root_job = field::Empty, workspace_id = %arc_job.workspace_id, worker = %worker_name, hostname = %hostname, tag = %arc_job.tag, +pub fn create_span_with_name( + arc_job: &MiniPulledJob, + worker_name: &str, + hostname: Option<&str>, + span_name: &str, +) -> Span { + // The span macro requires a literal, so we use a fixed name and set otel.name dynamically + let span = tracing::span!( + tracing::Level::INFO, + "job", + job_id = %arc_job.id, + root_job = field::Empty, + workspace_id = %arc_job.workspace_id, + worker = %worker_name, + hostname = field::Empty, + tag = %arc_job.tag, language = field::Empty, - script_path = field::Empty, flow_step_id = field::Empty, parent_job = field::Empty, - otel.name = field::Empty); + script_path = field::Empty, + flow_step_id = field::Empty, + parent_job = field::Empty, + otel.name = field::Empty + ); let rj = arc_job.flow_innermost_root_job.unwrap_or(arc_job.id); @@ -848,10 +871,10 @@ fn create_span(arc_job: &Arc, worker_name: &str, hostname: &str) span.record("language", lg.as_str()); } if let Some(step_id) = arc_job.flow_step_id.as_ref() { - span.record("otel.name", format!("job {}", step_id).as_str()); + span.record("otel.name", format!("{} {}", span_name, step_id).as_str()); span.record("flow_step_id", step_id.as_str()); } else { - span.record("otel.name", "job"); + span.record("otel.name", span_name); } if let Some(parent_job) = arc_job.parent_job.as_ref() { span.record("parent_job", parent_job.to_string().as_str()); @@ -862,6 +885,9 @@ fn create_span(arc_job: &Arc, worker_name: &str, hostname: &str) if let Some(root_job) = arc_job.flow_innermost_root_job.as_ref() { span.record("root_job", root_job.to_string().as_str()); } + if let Some(hostname) = hostname { + span.record("hostname", hostname); + } windmill_common::otel_oss::set_span_parent(&span, &rj); span @@ -876,6 +902,7 @@ pub async fn handle_all_job_kind_error( worker_dir: &str, worker_name: &str, job_completed_tx: JobCompletedSender, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) { match conn { @@ -892,6 +919,7 @@ pub async fn handle_all_job_kind_error( &worker_dir, &worker_name, job_completed_tx.clone(), + &killpill_rx, #[cfg(feature = "benchmark")] bench, ) @@ -915,6 +943,8 @@ pub async fn handle_all_job_kind_error( duration: None, has_stream: Some(false), from_cache: None, + flow_runners: None, + done_tx: None, }, false, ) @@ -962,27 +992,21 @@ pub fn start_interactive_worker_shell( .await; match job { - Ok(j) => match j.to_pulled_job() { - Ok(j) => Ok(j.clone().map(NextJob::Sql)), - ref e @ (Err(PulledJobResultToJobErr::MissingConcurrencyKey( - ref jc, - )) - | Err(PulledJobResultToJobErr::ErrorWhilePreprocessing( - ref jc, - ))) => { - if let Err(err) = - job_completed_tx.send_job(jc.clone(), true).await - { - let e_fmt = match e { - Ok(_) => "unknown error".to_owned(), - Err(e) => e.to_string(), - }; - - tracing::error!("An error occurred while sending job completed ({e_fmt}): {:#?}", err) + Ok(j) => { + match j.to_pulled_job() { + Ok(j) => Ok(j + .clone() + .map(|job| NextJob::Sql { flow_runners: None, job })), + Err(PulledJobResultToJobErr::MissingConcurrencyKey(jc)) + | Err(PulledJobResultToJobErr::ErrorWhilePreprocessing(jc)) => { + if let Err(err) = job_completed_tx.send_job(jc, true).await + { + tracing::error!("An error occurred while sending job completed: {:#?}", err) + } + Ok(None) } - Ok(None) } - }, + } Err(err) => Err(err), } } @@ -1010,6 +1034,7 @@ pub fn start_interactive_worker_shell( parent_runnable_path, token, precomputed_agent_info: precomputed_bundle, + flow_runners, } = extract_job_and_perms(job, &conn).await; let authed_client = AuthedClient::new( @@ -1039,6 +1064,7 @@ pub fn start_interactive_worker_shell( &mut occupancy_metrics, &mut killpill_rx, precomputed_bundle, + flow_runners, #[cfg(feature = "benchmark")] &mut bench, ) @@ -1476,9 +1502,9 @@ pub async fn run_worker( // Option>>, // Option>, - #[cfg(feature = "enterprise")] + #[cfg(all(feature = "private", feature = "enterprise"))] let (dedicated_workers, is_flow_worker, dedicated_handles): ( - HashMap>>, + HashMap>, bool, Vec>, ) = match conn { @@ -1497,9 +1523,9 @@ pub async fn run_worker( Connection::Http(_) => (HashMap::new(), false, vec![]), }; - #[cfg(not(feature = "enterprise"))] + #[cfg(any(not(feature = "private"), not(feature = "enterprise")))] let (dedicated_workers, is_flow_worker, dedicated_handles): ( - HashMap>>, + HashMap>, bool, Vec>, ) = (HashMap::new(), false, vec![]); @@ -1656,7 +1682,7 @@ pub async fn run_worker( if let Ok(same_worker_job) = same_worker_rx.try_recv() { same_worker_queue_size.fetch_sub(1, Ordering::SeqCst); - tracing::debug!( + tracing::info!( worker = %worker_name, hostname = %hostname, "received {} from same worker channel", same_worker_job.job_id @@ -1676,7 +1702,12 @@ pub async fn run_worker( .expect("send kill to job completed tx"); break; } else { - job.map(|x| x.map(NextJob::Sql)) + job.map(|x| { + x.map(|job| NextJob::Sql { + flow_runners: same_worker_job.flow_runners, + job, + }) + }) } } Connection::Http(client) => client @@ -1827,22 +1858,14 @@ pub async fn run_worker( } match job { Ok(pulled_job_result) => match pulled_job_result.to_pulled_job() { - Ok(j) => Ok(j.map(NextJob::Sql)), - ref e @ (Err(PulledJobResultToJobErr::MissingConcurrencyKey( - ref jc, - )) - | Err(PulledJobResultToJobErr::ErrorWhilePreprocessing( - ref jc, - ))) => { - if let Err(err) = - job_completed_tx.send_job(jc.clone(), true).await - { - let e_fmt = match e { - Ok(_) => "unknown error".to_owned(), - Err(e) => e.to_string(), - }; - - tracing::error!("An error occurred while sending job completed ({e_fmt}): {:#?}", err) + Ok(j) => Ok(j.map(|job| NextJob::Sql { flow_runners: None, job })), + Err(PulledJobResultToJobErr::MissingConcurrencyKey(jc)) + | Err(PulledJobResultToJobErr::ErrorWhilePreprocessing(jc)) => { + if let Err(err) = job_completed_tx.send_job(jc, true).await { + tracing::error!( + "An error occurred while sending job completed: {:#?}", + err + ) } Ok(None) } @@ -1874,7 +1897,10 @@ pub async fn run_worker( tracing::debug!(worker = %worker_name, hostname = %hostname, "started handling of job {}", job.id); - if matches!(job.kind, JobKind::Script | JobKind::Preview) { + if matches!( + job.kind, + JobKind::Script | JobKind::Preview | JobKind::FlowScript + ) { if !dedicated_workers.is_empty() { let key_o = if is_flow_worker { job.flow_step_id.as_ref().map(|x| x.to_string()) @@ -1883,8 +1909,12 @@ pub async fn run_worker( }; if let Some(key) = key_o { if let Some(dedicated_worker_tx) = dedicated_workers.get(&key) { - if let Err(e) = dedicated_worker_tx.send(Arc::new(job.job())).await - { + let dedicated_job = DedicatedWorkerJob { + job: Arc::new(job.job()), + flow_runners: None, + done_tx: None, + }; + if let Err(e) = dedicated_worker_tx.send(dedicated_job).await { tracing::info!("failed to send jobs to dedicated workers. Likely dedicated worker has been shut down. This is normal: {e:?}"); } @@ -1898,6 +1928,77 @@ pub async fn run_worker( } } } + + // Extract flow_runners early to use in both dedicated workers and flow runners + let flow_runners = match &job { + NextJob::Sql { flow_runners, .. } => flow_runners.clone(), + NextJob::Http(_) => None, + }; + + if let Some(flow_runners) = flow_runners { + let key_o = job.flow_step_id.as_ref().map(|x| x.to_string()); + if let Some(key) = key_o { + if let Some(flow_runner_tx) = flow_runners.runners.get(&key) { + tracing::info!( + "sending job {} to flow runner step {}", + job.id, + key + ); + let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>(); + let flow_runners = flow_runners.clone(); + + let job = job.job(); + + let dedicated_job = DedicatedWorkerJob { + job: Arc::new(job.clone()), + flow_runners: Some(flow_runners), + done_tx: Some(done_tx), + }; + + if let Err(e) = flow_runner_tx.send(dedicated_job).await { + let token = match &conn { + Connection::Sql(db) => { + windmill_queue::jobs::create_token(db, &job, None).await + } + _ => "".to_string(), + }; + handle_all_job_kind_error( + &conn, + &AuthedClient::new( + base_internal_url.to_owned(), + job.workspace_id.clone(), + token, + None, + ), + MiniCompletedJob::from(job), + error::Error::InternalErr(format!( + "failed to send jobs to flow runners: {e:?}" + )), + Some(&same_worker_tx), + &worker_dir, + &worker_name, + job_completed_tx.clone(), + &killpill_rx, + #[cfg(feature = "benchmark")] + &mut bench, + ) + .await; + } else { + if let Err(err) = done_rx.await { + tracing::error!("Flow runner done channel has been dropped without being received: {err:?}"); + } + } + + #[cfg(feature = "benchmark")] + { + add_time!(bench, "sent to flow runner"); + infos.add_iter(bench, true); + } + + continue; + } + } + } } if matches!(job.kind, JobKind::Noop) { @@ -1917,6 +2018,8 @@ pub async fn run_worker( duration: None, has_stream: Some(false), from_cache: None, + flow_runners: None, + done_tx: None, }, true, ) @@ -2048,6 +2151,7 @@ pub async fn run_worker( parent_runnable_path, token, precomputed_agent_info: precomputed_bundle, + flow_runners, } = extract_job_and_perms(job, &conn).await; let authed_client = AuthedClient::new( @@ -2059,7 +2163,7 @@ pub async fn run_worker( let arc_job = Arc::new(job); - let span = create_span(&arc_job, &worker_name, hostname); + let span = create_span_with_name(&arc_job, &worker_name, Some(hostname), "job"); let job_result = handle_queued_job( arc_job.clone(), @@ -2079,6 +2183,7 @@ pub async fn run_worker( &mut occupancy_metrics, &mut killpill_rx2, precomputed_bundle, + flow_runners, #[cfg(feature = "benchmark")] &mut bench, ) @@ -2138,6 +2243,7 @@ pub async fn run_worker( &worker_dir, &worker_name, job_completed_tx.clone(), + &killpill_rx, #[cfg(feature = "benchmark")] &mut bench, ) @@ -2302,7 +2408,7 @@ async fn queue_init_bash_maybe<'c>( }; if let Some((uuid, content)) = uuid_content { same_worker_tx - .send(SameWorkerPayload { job_id: uuid, recoverable: false }) + .send(SameWorkerPayload { job_id: uuid, recoverable: false, flow_runners: None }) .await .map_err(to_anyhow)?; tracing::info!("Creating initial job {uuid} from initial script script: {content}"); @@ -2420,7 +2526,7 @@ async fn queue_periodic_script_bash_maybe<'c>( }; same_worker_tx - .send(SameWorkerPayload { job_id: uuid, recoverable: false }) + .send(SameWorkerPayload { job_id: uuid, recoverable: false, flow_runners: None }) .await .map_err(to_anyhow)?; tracing::info!("Creating periodic script job {uuid} from periodic script: {content}"); @@ -2432,7 +2538,6 @@ pub struct SendResult { pub time: Instant, } -#[derive(Clone)] pub enum SendResultPayload { JobCompleted(JobCompleted), UpdateFlow(UpdateFlow), @@ -2519,6 +2624,7 @@ pub async fn handle_queued_job( occupancy_metrics: &mut OccupancyMetrics, killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, precomputed_agent_info: Option, + flow_runners: Option>, #[cfg(feature = "benchmark")] _bench: &mut BenchmarkIter, ) -> windmill_common::error::Result { // Extract the active span from the context @@ -2646,6 +2752,8 @@ pub async fn handle_queued_job( duration: None, has_stream: Some(false), from_cache: Some(true), + flow_runners: None, + done_tx: None, }, true, ) @@ -2681,6 +2789,8 @@ pub async fn handle_queued_job( worker_dir, job_completed_tx.clone(), worker_name, + flow_runners, + &killpill_rx, )) .warn_after_seconds(10) .await?; @@ -2911,6 +3021,7 @@ pub async fn handle_queued_job( conn, Some(started.elapsed().as_millis() as i64), has_stream, + flow_runners, ) .await } diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index f0bc4cfc1e..dd45d3b786 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -64,7 +64,7 @@ use windmill_common::{ use windmill_queue::schedule::get_schedule_opt; use windmill_queue::{ add_completed_job, add_completed_job_error, append_logs, get_mini_pulled_job, - handle_maybe_scheduled_job, insert_concurrency_key, interpolate_args, CanceledBy, + handle_maybe_scheduled_job, insert_concurrency_key, interpolate_args, CanceledBy, FlowRunners, MiniCompletedJob, MiniPulledJob, PushArgs, PushIsolationLevel, SameWorkerPayload, WrappedError, }; @@ -88,6 +88,8 @@ pub async fn update_flow_status_after_job_completion( stop_early_override: Option, worker_name: &str, job_completed_tx: JobCompletedSender, + flow_runners: Option>, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) -> error::Result>> { // this is manual tailrecursion because async_recursion blows up the stack @@ -121,6 +123,8 @@ pub async fn update_flow_status_after_job_completion( rec.has_triggered_error_handler, worker_name, job_completed_tx.clone(), + flow_runners.clone(), + killpill_rx, #[cfg(feature = "benchmark")] bench, )) @@ -147,6 +151,8 @@ pub async fn update_flow_status_after_job_completion( rec.has_triggered_error_handler, worker_name, job_completed_tx.clone(), + flow_runners.clone(), + killpill_rx, #[cfg(feature = "benchmark")] bench, )) @@ -285,6 +291,8 @@ pub async fn update_flow_status_after_job_completion_internal( has_triggered_error_handler: bool, worker_name: &str, job_completed_tx: JobCompletedSender, + flow_runners: Option>, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, #[cfg(feature = "benchmark")] bench: &mut BenchmarkIter, ) -> error::Result { let mut has_triggered_error_handler = has_triggered_error_handler; @@ -434,7 +442,6 @@ pub async fn update_flow_status_after_job_completion_internal( args })); - let from_result_to_args = |args: &Result>, sqlx::Error>| { let args = args .as_ref() @@ -473,7 +480,7 @@ pub async fn update_flow_status_after_job_completion_internal( _ => None, }; let args = from_result_to_args(args.as_ref().await.get_ref())?; - + compute_bool_from_expr( &expr, Marc::new(args), @@ -742,7 +749,6 @@ pub async fn update_flow_status_after_job_completion_internal( &mut stop_early_err_msg, &mut nresult, args, - ) .await?; } @@ -1565,6 +1571,8 @@ pub async fn update_flow_status_after_job_completion_internal( worker_dir, job_completed_tx, worker_name, + flow_runners, + &killpill_rx, )) .warn_after_seconds(10) .await @@ -2044,6 +2052,8 @@ pub async fn handle_flow( worker_dir: &str, job_completed_tx: JobCompletedSender, worker_name: &str, + flow_runners: Option>, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, ) -> anyhow::Result<()> { let flow = flow_data.value(); let status = flow_job @@ -2101,6 +2111,9 @@ pub async fn handle_flow( same_worker_tx, worker_dir, worker_name, + flow_runners.clone(), + job_completed_tx.clone(), + &killpill_rx, )) .warn_after_seconds(10) .await?; @@ -2185,6 +2198,7 @@ struct PushNextFlowJobRec { flow_job: Arc, status: FlowStatus, } + // #[async_recursion] // #[instrument(level = "trace", skip_all)] async fn push_next_flow_job( @@ -2197,6 +2211,9 @@ async fn push_next_flow_job( same_worker_tx: &SameWorkerSender, worker_dir: &str, worker_name: &str, + flow_runners: Option>, + job_completed_tx: JobCompletedSender, + killpill_rx: &tokio::sync::broadcast::Receiver<()>, ) -> error::Result { let job_root = flow_job .flow_innermost_root_job @@ -3028,6 +3045,22 @@ async fn push_next_flow_job( } }; + // only start runners if we're not already in a squash for loop + let start_runners = flow_runners.is_none() + && matches!( + next_status, + NextStatus::NextLoopIteration { start_runners: true, .. } + ); + + let do_not_pass_runners = matches!(next_status, NextStatus::NextStep { .. }) + && flow_runners + .as_ref() + .is_some_and(|fr| fr.job_id == flow_job.id); + + let continue_with_runners = (start_runners || (flow_runners.is_some() && !do_not_pass_runners)) + && module.suspend.is_none() + && module.sleep.is_none(); + // Also check `flow_job.same_worker` for [`JobKind::Flow`] jobs as it's no // more reflected to the flow value on push. let job_same_worker = flow_job.same_worker @@ -3083,6 +3116,7 @@ async fn push_next_flow_job( NextStatus::NextLoopIteration { next: ForloopNextIteration { new_args, .. }, simple_input_transforms, + .. } => { let mut args = if let Ok(args) = nargs { args.as_ref().clone() @@ -3300,7 +3334,7 @@ async fn push_next_flow_job( Some(module.id.clone()), new_job_priority_override, job_perms.as_ref(), - false, + continue_with_runners, None, None, None, @@ -3308,7 +3342,7 @@ async fn push_next_flow_job( .warn_after_seconds(2) .await?; - if continue_on_same_worker { + if continue_on_same_worker || continue_with_runners { let _ = sqlx::query!( "UPDATE v2_job_queue SET worker = $2 WHERE id = $1", uuid, @@ -3583,7 +3617,7 @@ async fn push_next_flow_job( .warn_after_seconds(3) .await?; - if continue_on_same_worker { + if continue_on_same_worker || continue_with_runners { if !is_one_uuid { return Err(Error::BadRequest( "Cannot continue on same worker with multiple jobs, parallel cannot be used in conjunction with same_worker".to_string(), @@ -3593,9 +3627,43 @@ async fn push_next_flow_job( tx.commit().warn_after_seconds(3).await?; tracing::info!(id = %flow_job.id, root_id = %job_root, "all next flow jobs pushed: {uuids:?}"); - if continue_on_same_worker { + if continue_on_same_worker || continue_with_runners { + let flow_runners = if start_runners { + tracing::info!(id = %flow_job.id, "starting flow runners for module {}", module.id); + let (new_flow_runners, new_flow_runner_handles) = + crate::dedicated_worker_oss::spawn_flow_module_runners( + &flow_job, + module, + flow.failure_module.as_ref(), + &killpill_rx, + db, + worker_dir, + &client.base_internal_url, + worker_name, + &job_completed_tx, + ) + .await + .with_context(|| { + format!( + "failed to spawn flow module runners for job: {}", + flow_job.id + ) + })?; + + let flow_runners = FlowRunners { + runners: new_flow_runners, + handles: new_flow_runner_handles, + job_id: flow_job.id, + }; + Some(Arc::new(flow_runners)) + } else if !do_not_pass_runners { + flow_runners + } else { + None + }; + same_worker_tx - .send(SameWorkerPayload { job_id: first_uuid, recoverable: true }) + .send(SameWorkerPayload { job_id: first_uuid, recoverable: true, flow_runners }) .warn_after_seconds(3) .await .map_err(to_anyhow)?; @@ -3702,6 +3770,7 @@ enum NextStatus { NextLoopIteration { next: ForloopNextIteration, simple_input_transforms: Option>, + start_runners: bool, }, AllFlowJobs { branchall: Option, @@ -3943,7 +4012,7 @@ async fn compute_next_flow_transform( NextStatus::NextStep, )) } - FlowModuleValue::WhileloopFlow { modules, modules_node, .. } => { + FlowModuleValue::WhileloopFlow { modules, modules_node, squash, .. } => { // if it's a simple single step flow, we will collapse it as an optimization and need to pass flow_input as an arg let is_simple = is_simple_modules(&modules, flow.failure_module.as_ref()); let (flow_jobs, flow_jobs_success, flow_jobs_duration) = match status_module { @@ -3960,6 +4029,8 @@ async fn compute_next_flow_transform( _ => (vec![], Some(vec![]), Some(FlowJobsDuration::new(0))), }; let next_loop_idx = flow_jobs.len(); + + let start_runners = next_loop_idx == 0 && squash.unwrap_or(false); next_loop_iteration( flow, status, @@ -3982,11 +4053,14 @@ async fn compute_next_flow_transform( db, module, delete_after_use, + start_runners, ) .await } /* forloop modules are expected set `iter: { value: Value, index: usize }` as job arguments */ - FlowModuleValue::ForloopFlow { modules, modules_node, iterator, parallel, .. } => { + FlowModuleValue::ForloopFlow { + modules, modules_node, iterator, parallel, squash, .. + } => { // if it's a simple single step flow, we will collapse it as an optimization and need to pass flow_input as an arg let is_simple = !matches!(flow_job.kind, JobKind::FlowPreview) && !parallel @@ -4025,6 +4099,7 @@ async fn compute_next_flow_transform( Ok(NextFlowTransform::EmptyInnerFlows { branch_chosen: None }) } ForLoopStatus::NextIteration(ns) => { + let start_runners = ns.index == 0 && squash.unwrap_or(false); next_loop_iteration( flow, status, @@ -4036,6 +4111,7 @@ async fn compute_next_flow_transform( db, module, delete_after_use, + start_runners, ) .await } @@ -4300,6 +4376,7 @@ async fn next_loop_iteration( db: &sqlx::Pool, module: &FlowModule, delete_after_use: bool, + start_runners: bool, ) -> Result { let inner_path = || format!("{}/loop-{}", flow_job.runnable_path(), ns.index); if is_simple { @@ -4317,7 +4394,7 @@ async fn next_loop_iteration( ContinuePayload::SingleJob( payload_from_simple_module(value, db, flow_job, module, inner_path()).await?, ), - NextStatus::NextLoopIteration { next: ns, simple_input_transforms }, + NextStatus::NextLoopIteration { next: ns, simple_input_transforms, start_runners }, )); } @@ -4341,7 +4418,7 @@ async fn next_loop_iteration( timeout: None, on_behalf_of: None, }), - NextStatus::NextLoopIteration { next: ns, simple_input_transforms: None }, + NextStatus::NextLoopIteration { next: ns, simple_input_transforms: None, start_runners }, )) } diff --git a/backend/windmill-worker/src/worker_lockfiles.rs b/backend/windmill-worker/src/worker_lockfiles.rs index 824a45c661..4ec0fa648e 100644 --- a/backend/windmill-worker/src/worker_lockfiles.rs +++ b/backend/windmill-worker/src/worker_lockfiles.rs @@ -1193,6 +1193,7 @@ async fn lock_modules<'c>( skip_failures, parallel, parallelism, + squash, } => { let nmodules; (nmodules, tx, nmodified_ids, nerrors) = Box::pin(lock_modules( @@ -1222,6 +1223,7 @@ async fn lock_modules<'c>( skip_failures, parallel, parallelism, + squash, } .into() } @@ -1258,7 +1260,7 @@ async fn lock_modules<'c>( } e.value = FlowModuleValue::BranchAll { branches: nbranches, parallel }.into() } - FlowModuleValue::WhileloopFlow { modules, modules_node, skip_failures } => { + FlowModuleValue::WhileloopFlow { modules, modules_node, skip_failures, squash } => { let nmodules; (nmodules, tx, nmodified_ids, nerrors) = Box::pin(lock_modules( modules, @@ -1284,6 +1286,7 @@ async fn lock_modules<'c>( modules: nmodules, modules_node, skip_failures, + squash, } .into() } diff --git a/frontend/src/lib/components/FlowGraphViewerStep.svelte b/frontend/src/lib/components/FlowGraphViewerStep.svelte index 6fcb5c902c..07ad7756c0 100644 --- a/frontend/src/lib/components/FlowGraphViewerStep.svelte +++ b/frontend/src/lib/components/FlowGraphViewerStep.svelte @@ -121,6 +121,7 @@ {:else if stepDetail.value.type == 'forloopflow'} For loop {#if stepDetail.value.parallel}(parallel){/if} {#if stepDetail.value.skip_failures}(skip failures){/if} + {#if stepDetail.value.squash}(squash){/if} {:else if stepDetail.value.type == 'branchall'} Run all branches {#if stepDetail.value.parallel}(parallel){/if} {:else if stepDetail.value.type == 'branchone'} @@ -128,7 +129,8 @@ {:else if stepDetail.value.type == 'flow'} Inner flow {:else if stepDetail.value.type == 'whileloopflow'} - While loop + While loop {#if stepDetail.value.skip_failures}(skip failures){/if} + {#if stepDetail.value.squash}(squash){/if} {:else if stepDetail.id === 'failure'} Error handler {:else if stepDetail.id === 'preprocessor'} diff --git a/frontend/src/lib/components/flows/content/FlowLoop.svelte b/frontend/src/lib/components/flows/content/FlowLoop.svelte index 50e7bd2a82..4169f49f34 100644 --- a/frontend/src/lib/components/flows/content/FlowLoop.svelte +++ b/frontend/src/lib/components/flows/content/FlowLoop.svelte @@ -32,6 +32,7 @@ import { slide } from 'svelte/transition' import ToggleButtonGroup from '$lib/components/common/toggleButton-v2/ToggleButtonGroup.svelte' import ToggleButton from '$lib/components/common/toggleButton-v2/ToggleButton.svelte' + import Badge from '$lib/components/common/badge/Badge.svelte' const { previewArgs, flowStateStore, flowStore, currentEditor } = getContext('FlowEditorContext') @@ -186,6 +187,33 @@ class="whitespace-nowrap" /> +
+
Squash + + Beta + This can result in unexpected behavior, use at your own risk for now.
+ Squashing a for loop runs all iterations on the same worker, using a single runner + per step for the entire loop. This eliminates cold starts between iterations for supported + languages (Bun, Deno, and Python). + + +
+ { + ;(mod.value as ForloopFlow).squash = detail + }} + options={{ + right: 'Squash' + }} + class="whitespace-nowrap" + disabled={mod.value.parallel} + /> +
Run in parallel
diff --git a/frontend/src/lib/components/flows/content/FlowWhileLoop.svelte b/frontend/src/lib/components/flows/content/FlowWhileLoop.svelte index 81bf6c4c11..6983f7c1c6 100644 --- a/frontend/src/lib/components/flows/content/FlowWhileLoop.svelte +++ b/frontend/src/lib/components/flows/content/FlowWhileLoop.svelte @@ -14,12 +14,13 @@ import FlowModuleSleep from './FlowModuleSleep.svelte' import FlowModuleMock from './FlowModuleMock.svelte' import { Play } from 'lucide-svelte' - import type { FlowModule, Job } from '$lib/gen' + import type { FlowModule, Job, WhileloopFlow } from '$lib/gen' import FlowLoopIterationPreview from '$lib/components/FlowLoopIterationPreview.svelte' import FlowModuleDeleteAfterUse from './FlowModuleDeleteAfterUse.svelte' import FlowModuleSkip from './FlowModuleSkip.svelte' import TabsV2 from '$lib/components/common/tabs/TabsV2.svelte' import { useUiIntent } from '$lib/components/copilot/chat/flow/useUiIntent' + import Badge from '$lib/components/common/badge/Badge.svelte' const { flowStateStore } = getContext('FlowEditorContext') @@ -86,7 +87,7 @@ {#if mod.value.type === 'whileloopflow'}
-
+
Skip failures
+
+
Squash + + Beta + This can result in unexpected behavior, use at your own risk for now.
+ Squashing a for loop runs all iterations on the same worker, using a single runner + per step for the entire loop. This eliminates cold starts between iterations for supported + languages (Bun, Deno, and Python). + + +
+ { + ;(mod.value as WhileloopFlow).squash = detail + }} + options={{ + right: 'Squash' + }} + class="whitespace-nowrap" + /> +
diff --git a/frontend/src/lib/components/flows/map/MapItem.svelte b/frontend/src/lib/components/flows/map/MapItem.svelte index 08d88f26aa..4e8d6302a2 100644 --- a/frontend/src/lib/components/flows/map/MapItem.svelte +++ b/frontend/src/lib/components/flows/map/MapItem.svelte @@ -164,7 +164,7 @@ mod.summary || (mod.value.type == 'forloopflow' ? 'For loop' : 'While loop') } ${mod.value.parallel ? '(parallel)' : ''} ${ mod.value.skip_failures ? '(skip failures)' : '' - }`} + } ${mod.value.squash ? '(squash)' : ''}`} id={mod.id} on:changeId on:move diff --git a/openflow.openapi.yaml b/openflow.openapi.yaml index b724c0acc9..3fcf48b801 100644 --- a/openflow.openapi.yaml +++ b/openflow.openapi.yaml @@ -446,6 +446,8 @@ components: type: boolean parallelism: $ref: "#/components/schemas/InputTransform" + squash: + type: boolean required: - modules - iterator @@ -469,6 +471,8 @@ components: type: boolean parallelism: $ref: "#/components/schemas/InputTransform" + squash: + type: boolean required: - modules - skip_failures