kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Display memory peak while installing wheels with uv (#4889)

Browse Source 
* 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 <ruben@windmill.dev>
This commit is contained in:
pyranota
2024-12-12 11:26:02 +01:00
committed by GitHub
parent 0dab23f6bf
commit 7398c1c09c
6 changed files with 183 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
backend/.sqlx/query-497d93db931922b09f96cf73239513d7141f3d37f85ada46597079991b3bff30.json generated Normal file
View File

@@ -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"
}

12
backend/.sqlx/query-a1f8667bfd5b000dd7ba384e78f2bd7fabb6b8055f4559f77e20eef7c2b1c902.json generated Normal file
View File

@@ -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"
}

26
backend/.sqlx/query-eb125df0f64c0caa07f50dc82e8ae9c6cd2872c7afe96d41bed731c5041ab671.json generated Normal file
View File

@@ -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"
}

2
backend/ee-repo-ref.txt
View File

@@ -1 +1 @@
89a221f8b6e0dc431c668ed804066c97b428f6ef
271c210ddbbdcad0f4d6c007e650eceda5ddfa64

2
backend/windmill-worker/src/handle_child.rs
View File

@@ -437,7 +437,7 @@ pub async fn handle_child(
}
}
async fn get_mem_peak(pid: Option<u32>, nsjail: bool) -> i32 {
pub(crate) async fn get_mem_peak(pid: Option<u32>, nsjail: bool) -> i32 {
if pid.is_none() {
return -1;
}

170
backend/windmill-worker/src/python_executor.rs
View File

@@ -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<CanceledBy>,
db: &sqlx::Pool<sqlx::Postgres>,
_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