From c4adaeeabd287ca1c4f3522bcd8bcea30b00fe6d Mon Sep 17 00:00:00 2001 From: "claude[bot]" <209825114+claude[bot]@users.noreply.github.com> Date: Thu, 19 Jun 2025 15:32:21 +0200 Subject: [PATCH] fix: improve piptar upload - sequential uploads via background task queue (#5994) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: improve piptar upload with sequential background task queue Replace simultaneous piptar uploads with sequential processing via background job queue while keeping dependency installation parallelized. Key changes: - Add PiptarUpload job kind and database migration - Queue piptar uploads instead of using tokio::spawn - Implement job handler for sequential S3 uploads - Maintain parallel dependency installation as requested 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Ruben Fiszel * Remove PiptarUpload job kind migration files Refactoring piptar uploads to use tokio channel instead of Windmill jobs. Migration files are no longer needed. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Ruben Fiszel * Refactor piptar uploads from Windmill jobs to tokio channel Replace the complex job queue infrastructure for piptar uploads with a simple tokio channel approach as requested. This maintains sequential upload behavior while keeping dependency installation parallel. Key changes: - Add PIPTAR_UPLOAD_CHANNEL global channel for sequential processing - Replace JobPayload::PiptarUpload with simple channel send - Remove PiptarUpload from JobKind enum and all job handling code - Remove job dispatcher case from worker.rs - Simplify upload logic while maintaining same functionality Benefits: - Reduced complexity by removing unnecessary job infrastructure - Sequential uploads without blocking dependency installation - Better separation of concerns 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Ruben Fiszel * Update worker.rs * Update python_executor.rs --------- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Ruben Fiszel Co-authored-by: Ruben Fiszel --- backend/windmill-queue/src/jobs.rs | 1 + .../windmill-worker/src/python_executor.rs | 55 ++++++++++++++++++- 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 1676148b00..3e89cd98b2 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -4028,6 +4028,7 @@ pub async fn push<'c, 'd>( ), }; + let final_priority: Option; #[cfg(not(feature = "enterprise"))] { diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index de5bedc5b1..90ef0afccd 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -62,12 +62,52 @@ lazy_static::lazy_static! { static ref EPHEMERAL_TOKEN_CMD: Option = var("EPHEMERAL_TOKEN_CMD").ok(); } +#[cfg(all(feature = "enterprise", feature = "parquet", unix))] +lazy_static::lazy_static! { + static ref PIPTAR_UPLOAD_CHANNEL: tokio::sync::mpsc::UnboundedSender = { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + + // Spawn background task to handle uploads sequentially + tokio::spawn(handle_piptar_uploads(rx)); + + tx + }; +} + +#[cfg(all(feature = "enterprise", feature = "parquet", unix))] +#[derive(Debug)] +struct PiptarUploadTask { + venv_path: String, + cache_dir: String, +} + +#[cfg(all(feature = "enterprise", feature = "parquet", unix))] +async fn handle_piptar_uploads(mut rx: tokio::sync::mpsc::UnboundedReceiver) { + use crate::global_cache::build_tar_and_push; + use windmill_common::s3_helpers::get_object_store; + + while let Some(task) = rx.recv().await { + if let Some(os) = get_object_store().await { + match build_tar_and_push(os, task.venv_path.clone(), task.cache_dir, None, false).await { + Ok(()) => { + tracing::info!("Successfully uploaded piptar for {}", task.venv_path); + } + Err(e) => { + tracing::error!("Failed to upload piptar for {}: {}", task.venv_path, e); + } + } + } else { + tracing::warn!("S3 object store not available for piptar upload: {}", task.venv_path); + } + } +} + const NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT: &str = include_str!("../nsjail/download.py.config.proto"); const NSJAIL_CONFIG_RUN_PYTHON3_CONTENT: &str = include_str!("../nsjail/run.python3.config.proto"); const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py"); #[cfg(all(feature = "enterprise", feature = "parquet", unix))] -use crate::global_cache::{build_tar_and_push, pull_from_tar}; +use crate::global_cache::pull_from_tar; #[cfg(all(feature = "enterprise", feature = "parquet", unix))] use windmill_common::s3_helpers::OBJECT_STORE_SETTINGS; @@ -1897,8 +1937,16 @@ pub async fn handle_python_reqs( #[cfg(all(feature = "enterprise", feature = "parquet", unix))] if s3_push { - if let Some(os) = windmill_common::s3_helpers::get_object_store().await { - tokio::spawn(build_tar_and_push(os, venv_p.clone(), py_version.to_cache_dir_top_level(false), None, false)); + // Send to upload channel for sequential processing + let upload_task = PiptarUploadTask { + venv_path: venv_p.clone(), + cache_dir: py_version.to_cache_dir_top_level(false), + }; + + if let Err(e) = PIPTAR_UPLOAD_CHANNEL.send(upload_task) { + tracing::warn!("Failed to queue piptar upload for {venv_p}: {e}"); + } else { + tracing::info!("Queued piptar upload for {venv_p}"); } } @@ -2200,3 +2248,4 @@ for line in sys.stdin: ) .await } +