From 7398c1c09c55a0e220cbffa8f59cd924ad8cef79 Mon Sep 17 00:00:00 2001 From: pyranota <92104930+pyranota@users.noreply.github.com> Date: Thu, 12 Dec 2024 11:26:02 +0100 Subject: [PATCH] Display memory peak while installing wheels with uv (#4889) * fix no_uv not affecting deploy Before this fix no_uv, no_uv_compile and no_uv_install were not affecting Dependency jobs These jobs are only affected if used USE_PIP_COMPILE or USE_PIP_INSTALL env variables To make it more consistant, no_uv should also affect dep jobs. Also make ansible use uv by default * Make it build * Add no_uv_usage stats * Provide build-env for Samael through shell.nix * Run update_sqlx.sh * Update ee-repo-ref * Display memory peak while installing wheels with uv * Invert DISABLE_NSJAIL * Make it safer --------- Co-authored-by: Ruben Fiszel --- ...513d7141f3d37f85ada46597079991b3bff30.json | 26 +++ ...2bd7fabb6b8055f4559f77e20eef7c2b1c902.json | 12 ++ ...ae9c6cd2872c7afe96d41bed731c5041ab671.json | 26 +++ backend/ee-repo-ref.txt | 2 +- backend/windmill-worker/src/handle_child.rs | 2 +- .../windmill-worker/src/python_executor.rs | 170 ++++++++++++------ 6 files changed, 183 insertions(+), 55 deletions(-) create mode 100644 backend/.sqlx/query-497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30.json create mode 100644 backend/.sqlx/query-a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902.json create mode 100644 backend/.sqlx/query-eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671.json diff --git a/backend/.sqlx/query-497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30.json b/backend/.sqlx/query-497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30.json new file mode 100644 index 0000000000..98e7d5d6ae --- /dev/null +++ b/backend/.sqlx/query-497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30.json @@ -0,0 +1,26 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT \n MAX (created_at) as last_deploy, \n COUNT (*) as deploys_count \n FROM metrics \n WHERE id = 'no_uv_usage_py'", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "last_deploy", + "type_info": "Timestamptz" + }, + { + "ordinal": 1, + "name": "deploys_count", + "type_info": "Int8" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + null, + null + ] + }, + "hash": "497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30" +} diff --git a/backend/.sqlx/query-a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902.json b/backend/.sqlx/query-a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902.json new file mode 100644 index 0000000000..2f20d99512 --- /dev/null +++ b/backend/.sqlx/query-a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902.json @@ -0,0 +1,12 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO metrics (id, value) \n VALUES ('no_uv_usage_py', ''::text::jsonb)\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [] + }, + "nullable": [] + }, + "hash": "a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902" +} diff --git a/backend/.sqlx/query-eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671.json b/backend/.sqlx/query-eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671.json new file mode 100644 index 0000000000..6b82368d66 --- /dev/null +++ b/backend/.sqlx/query-eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671.json @@ -0,0 +1,26 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT \n MAX (created_at) as last_deploy, \n COUNT (*) as deploys_count \n FROM metrics \n WHERE id = 'no_uv_usage_ansible'", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "last_deploy", + "type_info": "Timestamptz" + }, + { + "ordinal": 1, + "name": "deploys_count", + "type_info": "Int8" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + null, + null + ] + }, + "hash": "eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671" +} diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 9d6ca729e0..52d406af20 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -89a221f8b6e0dc431c668ed804066c97b428f6ef +271c210ddbbdcad0f4d6c007e650eceda5ddfa64 diff --git a/backend/windmill-worker/src/handle_child.rs b/backend/windmill-worker/src/handle_child.rs index f2fead39c9..5215b05b83 100644 --- a/backend/windmill-worker/src/handle_child.rs +++ b/backend/windmill-worker/src/handle_child.rs @@ -437,7 +437,7 @@ pub async fn handle_child( } } -async fn get_mem_peak(pid: Option, nsjail: bool) -> i32 { +pub(crate) async fn get_mem_peak(pid: Option, nsjail: bool) -> i32 { if pid.is_none() { return -1; } diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index 3c31fb1913..09caf863e2 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -82,7 +82,7 @@ use crate::{ create_args_and_out_file, get_main_override, get_reserved_variables, read_file, read_result, start_child_process, OccupancyMetrics, }, - handle_child::handle_child, + handle_child::{get_mem_peak, handle_child}, AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, LOCK_CACHE_DIR, NSJAIL_PATH, PATH_ENV, PIP_CACHE_DIR, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS, PY311_CACHE_DIR, TZ_ENV, UV_CACHE_DIR, @@ -1330,7 +1330,7 @@ pub async fn handle_python_reqs( requirements: Vec<&str>, job_id: &Uuid, w_id: &str, - _mem_peak: &mut i32, + mem_peak: &mut i32, _canceled_by: &mut Option, db: &sqlx::Pool, _worker_name: &str, @@ -1500,58 +1500,105 @@ pub async fn handle_python_reqs( let db_2 = db.clone(); let w_id_2 = w_id.to_string(); - tokio::spawn(async move { - loop { - tokio::select! { - _ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => { - // Notify server that we are still alive - // Detect if job has been canceled - let canceled = - sqlx::query_scalar::<_, bool> - (r#" - - UPDATE queue - SET last_ping = now() - WHERE id = $1 - RETURNING canceled - - "#) - .bind(job_id_2) - .fetch_optional(&db_2) - .await - .unwrap_or_else(|e| { - tracing::error!(%e, "error updating job {job_id_2}: {e:#}"); - Some(false) - }) - .unwrap_or_else(|| { - // if the job is not in queue, it can only be in the completed_job so it is already complete - false - }); - - if canceled { - - tracing::info!( - // If there is listener on other side, - workspace_id = %w_id_2, - "cancelling installations", - ); - - if let Err(ref e) = kill_tx.send(()){ - tracing::error!( - // If there is listener on other side, - workspace_id = %w_id_2, - "failed to send done: Probably receiving end closed too early or have not opened yet\n{}", - // If there is no listener, it will be dropped safely - e - ); + // Wheels to install + let total_to_install = req_with_penv.len(); + let pids = Arc::new(tokio::sync::Mutex::new(vec![None; total_to_install])); + let mem_peak_thread_safe = Arc::new(tokio::sync::Mutex::new(0)); + { + let pids = pids.clone(); + let mem_peak_thread_safe = mem_peak_thread_safe.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => { + let mut local_mem_peak = 0; + for pid_o in pids.lock().await.iter() { + if pid_o.is_some(){ + let mem = get_mem_peak(*pid_o, !*DISABLE_NSJAIL).await; + if mem < 0 { + tracing::warn!( + workspace_id = %w_id_2, + "Cannot get memory peak for pid: {:?}, job_id: {:?}, exit code: {mem}", + pid_o, + job_id_2 + ); + } else { + local_mem_peak += mem; + } + } } + + let mem_peak_actual = { + let mut mem_peak_lock = mem_peak_thread_safe.lock().await; + + if local_mem_peak > *mem_peak_lock{ + *mem_peak_lock = local_mem_peak; + } else { + tracing::debug!( + workspace_id = %w_id_2, + "Local mem_peak {:?}mb is smaller then global one {:?}mb, ignoring. job_id: {:?}", + local_mem_peak / 1000, + *mem_peak_lock / 1000, + job_id_2 + ); + + } + // Get the copy of value and drop lock itself, to release it as fast as possible + *mem_peak_lock + }; + + // Notify server that we are still alive + // Detect if job has been canceled + let canceled = + sqlx::query_scalar::<_, bool> + (r#" + + UPDATE queue + SET last_ping = now() + , mem_peak = $1 + WHERE id = $2 + RETURNING canceled + + "#) + .bind(mem_peak_actual) + .bind(job_id_2) + .fetch_optional(&db_2) + .await + .unwrap_or_else(|e| { + tracing::error!(%e, "error updating job {job_id_2}: {e:#}"); + Some(false) + }) + .unwrap_or_else(|| { + // if the job is not in queue, it can only be in the completed_job so it is already complete + false + }); + + if canceled { + + tracing::info!( + // If there is listener on other side, + workspace_id = %w_id_2, + "cancelling installations", + ); + + if let Err(ref e) = kill_tx.send(()){ + tracing::error!( + // If there is listener on other side, + workspace_id = %w_id_2, + "failed to send done: Probably receiving end closed too early or have not opened yet\n{}", + // If there is no listener, it will be dropped safely + e + ); + } + } } + // Once done_tx is dropped, this will be fired + _ = done_rx.recv() => break } - // Once done_tx is dropped, this will be fired - _ = done_rx.recv() => break } - } - }); + }); + + } // tl = total_length // "small".len == 5 @@ -1559,8 +1606,6 @@ pub async fn handle_python_reqs( // "largest".len == 7 // ==> req_tl = 7 let mut req_tl = 0; - // Wheels to install - let total_to_install = req_with_penv.len(); if total_to_install > 0 { let mut logs = String::new(); // Do we use UV? @@ -1595,13 +1640,14 @@ pub async fn handle_python_reqs( let semaphore = Arc::new(Semaphore::new(parallel_limit)); let mut handles = Vec::with_capacity(total_to_install); + // let mem_peak_thread_safe = Arc::new(tokio::sync::Mutex::new(0)); #[cfg(all(feature = "enterprise", feature = "parquet", unix))] let is_not_pro = !matches!(get_license_plan().await, LicensePlan::Pro); let total_time = std::time::Instant::now(); let has_work = req_with_penv.len() > 0; - for ((req, venv_p), mut kill_rx) in req_with_penv.iter().zip(kill_rxs.into_iter()) { + for ((i, (req, venv_p)), mut kill_rx) in req_with_penv.iter().enumerate().zip(kill_rxs.into_iter()) { let permit = semaphore.clone().acquire_owned().await; // Acquire a permit if let Err(_) = permit { @@ -1627,6 +1673,7 @@ pub async fn handle_python_reqs( let venv_p = venv_p.clone(); let counter_arc = counter_arc.clone(); let pip_indexes = pip_indexes.clone(); + let pids = pids.clone(); handles.push(task::spawn(async move { // permit will be dropped anyway if this thread exits at any point @@ -1667,6 +1714,7 @@ pub async fn handle_python_reqs( start, db ).await; + pids.lock().await.get_mut(i).and_then(|e| e.take()); return Ok(()); } } @@ -1693,6 +1741,7 @@ pub async fn handle_python_reqs( db, ) .await; + pids.lock().await.get_mut(i).and_then(|e| e.take()); return Err(e.into()); } }; @@ -1702,10 +1751,20 @@ pub async fn handle_python_reqs( .take() .ok_or(anyhow!("Cannot take stderr from uv_install_proccess"))?; + if let Some(pid) = pids.lock().await.get_mut(i) { + *pid = uv_install_proccess.id(); + } else { + tracing::error!( + workspace_id = %w_id, + "Index out of range for uv pids", + ); + } + tokio::select! { // Canceled _ = kill_rx.recv() => { uv_install_proccess.kill().await?; + pids.lock().await.get_mut(i).and_then(|e| e.take()); return Err(anyhow::anyhow!("uv pip install was canceled")); } // Finished @@ -1734,6 +1793,7 @@ pub async fn handle_python_reqs( db, ) .await; + pids.lock().await.get_mut(i).and_then(|e| e.take()); return Err(anyhow!(buf)); }, Err(e) => { @@ -1741,6 +1801,7 @@ pub async fn handle_python_reqs( workspace_id = %w_id, "Cannot wait for uv_install_proccess, ExitStatus is Err: {e:?}", ); + pids.lock().await.get_mut(i).and_then(|e| e.take()); return Err(e.into()); } } @@ -1781,6 +1842,7 @@ pub async fn handle_python_reqs( job_id ); + pids.lock().await.get_mut(i).and_then(|e| e.take()); Ok(()) })); } @@ -1814,6 +1876,8 @@ pub async fn handle_python_reqs( append_logs(&job_id, w_id, format!("\nenv set in {}ms", total_time), db).await; } + *mem_peak = *mem_peak_thread_safe.lock().await; + // Usually done_tx will drop after this return // If there is listener on other side, // it will be triggered