From da18a698082654767e9d038887f16c2e9fee6ced Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Wed, 8 Apr 2026 17:59:10 +0000 Subject: [PATCH] perf: add agent-batch mode with batch pull from server Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/src/main.rs | 8 ++--- backend/windmill-common/src/utils.rs | 17 +++++++++ backend/windmill-worker/src/agent_workers.rs | 9 +++++ backend/windmill-worker/src/worker.rs | 36 +++++++++++++++++--- 4 files changed, 62 insertions(+), 8 deletions(-) diff --git a/backend/src/main.rs b/backend/src/main.rs index af22a203c6..976dc39fe7 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -754,7 +754,7 @@ async fn windmill_main() -> anyhow::Result<()> { .and_then(|x| x.parse().ok()) .unwrap_or(IpAddr::from(default_bind_addr)); - let (conn, first_suffix, agent_config) = if mode == Mode::Agent { + let (conn, first_suffix, agent_config) = if matches!(mode, Mode::Agent | Mode::AgentBatch) { let agent_config = match AgentConfig::from_env() { Ok(config) => config, Err(e) => { @@ -829,7 +829,7 @@ async fn windmill_main() -> anyhow::Result<()> { let _guard = windmill_common::tracing_init::initialize_tracing(&hostname, &mode, &environment); - let is_agent = mode == Mode::Agent; + let is_agent = matches!(mode, Mode::Agent | Mode::AgentBatch); let mut migration_handle: Option> = None; #[cfg(feature = "parquet")] @@ -953,7 +953,7 @@ async fn windmill_main() -> anyhow::Result<()> { } } - let conn = if mode == Mode::Agent { + let conn = if matches!(mode, Mode::Agent | Mode::AgentBatch) { conn } else { // Drop the initial connection pool before creating the main one. @@ -1289,7 +1289,7 @@ Windmill Community Edition {GIT_VERSION} ) }, worker_name: worker_name_with_suffix( - mode == Mode::Agent, + matches!(mode, Mode::Agent | Mode::AgentBatch), WORKER_GROUP.as_str(), &suffix, ), diff --git a/backend/windmill-common/src/utils.rs b/backend/windmill-common/src/utils.rs index 10b1a3b408..6404f85eed 100644 --- a/backend/windmill-common/src/utils.rs +++ b/backend/windmill-common/src/utils.rs @@ -122,6 +122,21 @@ lazy_static::lazy_static! { } #[cfg(feature = "enterprise")] Mode::Agent + } else if &x == "agent-batch" { + println!("Binary is in 'agent-batch' mode with BASE_INTERNAL_URL={}", std::env::var("BASE_INTERNAL_URL").unwrap_or_default()); + if std::env::var("BASE_INTERNAL_URL").is_err() { + panic!("BASE_INTERNAL_URL is required in agent-batch mode") + } + if std::env::var("AGENT_TOKEN").is_err() { + println!("AGENT_TOKEN is not passed. This is required for the agent to work and contains the JWT to authenticate with the server.") + } + + #[cfg(not(feature = "enterprise"))] + { + panic!("Agent-batch mode is only available in the EE, ignoring..."); + } + #[cfg(feature = "enterprise")] + Mode::AgentBatch } else if &x == "indexer" { tracing::info!("Binary is in 'indexer' mode"); #[cfg(not(feature = "tantivy"))] @@ -474,6 +489,7 @@ pub fn map_string_to_number(s: &str, max_number: u64) -> u64 { pub enum Mode { Worker, Agent, + AgentBatch, Server, Standalone, Indexer, @@ -485,6 +501,7 @@ impl std::fmt::Display for Mode { match self { Mode::Worker => write!(f, "worker"), Mode::Agent => write!(f, "agent"), + Mode::AgentBatch => write!(f, "agent-batch"), Mode::Server => write!(f, "server"), Mode::Standalone => write!(f, "standalone"), Mode::Indexer => write!(f, "indexer"), diff --git a/backend/windmill-worker/src/agent_workers.rs b/backend/windmill-worker/src/agent_workers.rs index 9bf0ea1834..99bd08e0b8 100644 --- a/backend/windmill-worker/src/agent_workers.rs +++ b/backend/windmill-worker/src/agent_workers.rs @@ -37,6 +37,15 @@ pub async fn pull_job( .await } +pub async fn batch_pull_jobs( + client: &HttpClient, + batch_size: i32, +) -> anyhow::Result> { + client + .post("/api/agent_workers/batch_pull", None, &batch_size) + .await +} + pub async fn send_result(client: &HttpClient, jc: JobCompleted) -> anyhow::Result { client .post( diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index d7b8566e4b..f90a131e3f 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2078,6 +2078,8 @@ pub async fn run_worker( let batch_pull_size = *windmill_common::worker::BATCH_PULL_SIZE; let mut batch_pull_buffer: std::collections::VecDeque = std::collections::VecDeque::new(); + let mut agent_batch_buffer: std::collections::VecDeque = + std::collections::VecDeque::new(); let mut killpill_rx2 = killpill_rx.resubscribe(); @@ -2504,10 +2506,36 @@ pub async fn run_worker( } } - Connection::Http(client) => crate::agent_workers::pull_job(&client, None, None) - .await - .map_err(|e| error::Error::InternalErr(e.to_string())) - .map(|x| x.map(|y| NextJob::Http(y))), + Connection::Http(client) => { + if batch_pull_size > 0 { + // Agent-batch mode: serve from buffer, refill via batch pull + if let Some(job) = agent_batch_buffer.pop_front() { + Ok(Some(NextJob::Http(job))) + } else { + match crate::agent_workers::batch_pull_jobs( + &client, + batch_pull_size, + ) + .await + { + Ok(mut jobs) if !jobs.is_empty() => { + let first = jobs.remove(0); + for job in jobs { + agent_batch_buffer.push_back(job); + } + Ok(Some(NextJob::Http(first))) + } + Ok(_) => Ok(None), + Err(e) => Err(error::Error::InternalErr(e.to_string())), + } + } + } else { + crate::agent_workers::pull_job(&client, None, None) + .await + .map_err(|e| error::Error::InternalErr(e.to_string())) + .map(|x| x.map(|y| NextJob::Http(y))) + } + } } } };