fix(python): serialize concurrent installs into shared wheel cache dir (#9787)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

backend/windmill-worker/src/python_executor.rs

@@ -62,6 +62,13 @@ lazy_static::lazy_static! {
     static ref PY_CONCURRENT_DOWNLOADS: usize =
     var("PY_CONCURRENT_DOWNLOADS").ok().map(|flag| flag.parse().unwrap_or(20)).unwrap_or(20);
 
+    // uv's HTTP request timeout (seconds). spawn_uv_install uses env_clear(), so a
+    // UV_HTTP_TIMEOUT set on the worker is dropped unless forwarded explicitly.
+    // Only forwarded when set; otherwise uv keeps its own default. Lets operators
+    // raise it for slow/contended private registries ("operation timed out").
+    static ref UV_HTTP_TIMEOUT: Option<String> =
+    var("UV_HTTP_TIMEOUT").ok().filter(|v| !v.is_empty());
+
 
     static ref NON_ALPHANUM_CHAR: Regex = regex::Regex::new(r"[^0-9A-Za-z=.-]").unwrap();
 
@@ -2059,6 +2066,24 @@ lazy_static::lazy_static! {
     /// in-memory marker and pays only the original single stat.
     static ref VERIFIED_VENVS: tokio::sync::Mutex<HashSet<String>> =
         tokio::sync::Mutex::new(HashSet::new());
+
+    // In-process locks serializing concurrent installs into the same shared
+    // `venv_p` cache dir; `uv --reinstall` removes a package's .dist-info/RECORD
+    // before rewriting it, so a sibling install/verify racing it corrupts the
+    // dir. Keyed by venv_p so distinct deps still install in parallel.
+    static ref PY_INSTALL_LOCKS: tokio::sync::Mutex<std::collections::HashMap<String, Arc<tokio::sync::Mutex<()>>>> =
+        tokio::sync::Mutex::new(std::collections::HashMap::new());
+}
+
+/// Returns the in-process install lock for a given target cache dir, creating it
+/// on first use. Idle entries (only the map holds a reference) are pruned each
+/// call so the map stays bounded by the number of in-flight installs.
+async fn get_venv_install_lock(venv_p: &str) -> Arc<tokio::sync::Mutex<()>> {
+    let mut map = PY_INSTALL_LOCKS.lock().await;
+    map.retain(|_, v| Arc::strong_count(v) > 1);
+    map.entry(venv_p.to_string())
+        .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
+        .clone()
 }
 
 /// Spawn process of uv install
@@ -2104,6 +2129,9 @@ async fn spawn_uv_install(
         if *NATIVE_CERT {
             vars.push(("UV_NATIVE_TLS", "true"));
         }
+        if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() {
+            vars.push(("UV_HTTP_TIMEOUT", timeout.as_str()));
+        }
 
         let _owner;
         if let Some(py_path) = py_path.as_ref() {
@@ -2207,6 +2235,9 @@ async fn spawn_uv_install(
         let mut envs = vec![("PATH", PATH_ENV.as_str())];
         envs.push(("HOME", HOME_ENV.as_str()));
         envs.push(("UV_INDEX_STRATEGY", uv_index_strategy));
+        if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() {
+            envs.push(("UV_HTTP_TIMEOUT", timeout.as_str()));
+        }
         if let Some(mirror) = uv_python_install_mirror.as_ref() {
             envs.push(("UV_PYTHON_INSTALL_MIRROR", mirror));
         }
@@ -2780,6 +2811,96 @@ pub async fn handle_python_reqs(
             );
 
             let start = std::time::Instant::now();
+
+            // Lock the shared target dir (see PY_INSTALL_LOCKS). In-process lock
+            // first; only one task per dir then contends the cross-process file
+            // lock below. Both guards drop on every return path.
+            let venv_lock = get_venv_install_lock(&venv_p).await;
+            let _venv_guard = tokio::select! {
+                _ = kill_rx.recv() => {
+                    pids.lock().await.get_mut(i).and_then(|e| e.take());
+                    return Err(Error::from(anyhow::anyhow!(
+                        "install of {venv_p} canceled while waiting for venv lock"
+                    )));
+                }
+                guard = venv_lock.lock_owned() => guard,
+            };
+
+            // Cross-process advisory lock. Best-effort: if the filesystem doesn't
+            // support flock we log and proceed — verify_wheel_record + job retry
+            // still guard correctness, just without the dedup.
+            #[cfg(unix)]
+            let _venv_file_lock: Option<std::fs::File> = {
+                use std::os::unix::io::AsRawFd;
+                let lock_path = format!("{venv_p}.lock");
+                if let Some(parent) = std::path::Path::new(&lock_path).parent() {
+                    let _ = std::fs::create_dir_all(parent);
+                }
+                match std::fs::OpenOptions::new().create(true).write(true).open(&lock_path) {
+                    Ok(f) => {
+                        // Bounded wait: a holder that crashes releases the lock (the
+                        // kernel drops it on fd close), but a live-but-stuck holder
+                        // (e.g. uv wedged on a hung mount) would otherwise block us
+                        // forever. After the cap, proceed degraded rather than hang —
+                        // verify_wheel_record + retry still guard correctness.
+                        const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(300);
+                        let waited_since = std::time::Instant::now();
+                        loop {
+                            match nix::fcntl::flock(f.as_raw_fd(), nix::fcntl::FlockArg::LockExclusiveNonblock) {
+                                Ok(()) => break Some(f),
+                                // EWOULDBLOCK == EAGAIN on Linux: another holder has the lock.
+                                Err(nix::errno::Errno::EWOULDBLOCK) => {
+                                    if waited_since.elapsed() >= MAX_WAIT {
+                                        tracing::warn!(
+                                            workspace_id = %w_id,
+                                            "venv install lock {lock_path} still held after {}s, proceeding without cross-process install lock",
+                                            MAX_WAIT.as_secs()
+                                        );
+                                        break Some(f);
+                                    }
+                                    tokio::select! {
+                                        _ = kill_rx.recv() => {
+                                            pids.lock().await.get_mut(i).and_then(|e| e.take());
+                                            return Err(Error::from(anyhow::anyhow!(
+                                                "install of {venv_p} canceled while waiting for venv file lock"
+                                            )));
+                                        }
+                                        _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {}
+                                    }
+                                }
+                                Err(e) => {
+                                    tracing::warn!(
+                                        workspace_id = %w_id,
+                                        "could not flock {lock_path}, proceeding without cross-process install lock: {e}"
+                                    );
+                                    break Some(f);
+                                }
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        tracing::warn!(
+                            workspace_id = %w_id,
+                            "could not open install lock file {lock_path}, proceeding without cross-process install lock: {e}"
+                        );
+                        None
+                    }
+                }
+            };
+
+            // Double-checked: another job (this process or another sharing the
+            // mount) may have installed this exact dep while we waited on the
+            // locks. Reuse it instead of reinstalling.
+            if metadata(format!("{venv_p}/.valid.windmill")).await.is_ok() {
+                print_success(
+                    false, false, &job_id, &w_id, &req, req_tl, counter_arc,
+                    total_to_install, start, &conn,
+                )
+                .await;
+                pids.lock().await.get_mut(i).and_then(|e| e.take());
+                return Ok(());
+            }
+
             #[cfg(all(feature = "enterprise", feature = "parquet"))]
             if is_not_pro {
                 if let Some(os) = windmill_object_store::get_object_store().await {
@@ -3533,11 +3654,7 @@ mod tests {
         }
         let dist_info = root.join("pkg-1.0.0.dist-info");
         std::fs::create_dir_all(&dist_info).unwrap();
-        std::fs::write(
-            dist_info.join("RECORD"),
-            record_entries.join("\n") + "\n",
-        )
-        .unwrap();
+        std::fs::write(dist_info.join("RECORD"), record_entries.join("\n") + "\n").unwrap();
     }
 
     #[tokio::test]
@@ -3605,4 +3722,99 @@ mod tests {
             .await
             .is_ok());
     }
+
+    // Regression tests for the concurrent-install guard. Two jobs installing the
+    // same uncached dep into the shared `venv_p` used to race uv's `--reinstall`,
+    // corrupting the on-disk wheel and failing with "Env installation did not
+    // succeed". The guard serializes those installs.
+
+    #[tokio::test]
+    async fn test_venv_install_lock_serializes_same_path() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        // Same target path => one shared lock => no two tasks install at once.
+        let active = Arc::new(AtomicUsize::new(0));
+        let max_seen = Arc::new(AtomicUsize::new(0));
+        let mut handles = vec![];
+        for _ in 0..8 {
+            let active = active.clone();
+            let max_seen = max_seen.clone();
+            handles.push(tokio::spawn(async move {
+                let lock = get_venv_install_lock("/cache/py/3.11/samedep==1.0").await;
+                let _g = lock.lock_owned().await;
+                let cur = active.fetch_add(1, Ordering::SeqCst) + 1;
+                max_seen.fetch_max(cur, Ordering::SeqCst);
+                // Yield so any concurrency would be observed by another task.
+                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
+                active.fetch_sub(1, Ordering::SeqCst);
+            }));
+        }
+        for h in handles {
+            h.await.unwrap();
+        }
+        assert_eq!(
+            max_seen.load(Ordering::SeqCst),
+            1,
+            "installs into the same target dir must be serialized"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_venv_install_lock_distinct_paths_are_independent() {
+        // Different target paths get different locks and never block each other.
+        let a = get_venv_install_lock("/cache/py/3.11/depA==1.0").await;
+        let b = get_venv_install_lock("/cache/py/3.11/depB==1.0").await;
+        let _ga = a.lock_owned().await;
+        // Holding depA's lock must not prevent acquiring depB's.
+        assert!(
+            b.try_lock().is_ok(),
+            "distinct deps must install in parallel"
+        );
+        // Same path returns the same underlying lock.
+        let a2 = get_venv_install_lock("/cache/py/3.11/depA==1.0").await;
+        assert!(
+            a2.try_lock().is_err(),
+            "same target dir must map to the same lock"
+        );
+    }
+
+    #[cfg(unix)]
+    #[tokio::test]
+    async fn test_venv_file_lock_excludes_across_descriptions() {
+        // The cross-process layer: flock on a sibling `.lock` excludes a second
+        // independent open file description (i.e. another worker process) while
+        // held, and frees it on close. Mirrors the loop in handle_python_reqs.
+        use nix::fcntl::{flock, FlockArg};
+        use std::os::unix::io::AsRawFd;
+
+        let dir = std::env::temp_dir().join("wm_venv_lock_test");
+        std::fs::create_dir_all(&dir).unwrap();
+        let lock_path = dir.join("dep==1.0.lock");
+
+        let f1 = std::fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .open(&lock_path)
+            .unwrap();
+        flock(f1.as_raw_fd(), FlockArg::LockExclusiveNonblock).unwrap();
+
+        // A second descriptor (stand-in for another process) cannot take it.
+        let f2 = std::fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .open(&lock_path)
+            .unwrap();
+        assert_eq!(
+            flock(f2.as_raw_fd(), FlockArg::LockExclusiveNonblock),
+            Err(nix::errno::Errno::EWOULDBLOCK),
+            "a second holder must be blocked while the lock is held"
+        );
+
+        // Releasing the first lets the second acquire it.
+        drop(f1);
+        flock(f2.as_raw_fd(), FlockArg::LockExclusiveNonblock)
+            .expect("lock must be acquirable once the holder releases it");
+
+        drop(f2);
+        let _ = std::fs::remove_file(&lock_path);
+    }
 }