feat: large log disk and distributed storage compaction

This commit is contained in:
Ruben Fiszel
2024-03-23 10:54:44 +01:00
parent 461243a7a5
commit 75e9e67d7a
38 changed files with 665 additions and 296 deletions
+64 -53
View File
@@ -152,7 +152,13 @@ pub async fn pip_compile(
write_file(job_dir, file, &requirements).await?;
let mut args = vec!["-q", "--no-header", file, "--resolver=backtracking", "--strip-extras"];
let mut args = vec![
"-q",
"--no-header",
file,
"--resolver=backtracking",
"--strip-extras",
];
let mut pip_args = vec![];
let pip_extra_index_url = PIP_EXTRA_INDEX_URL
.read()
@@ -776,7 +782,6 @@ pub async fn handle_python_reqs(
.await?;
};
let mut req_with_penv: Vec<(String, String)> = vec![];
for req in requirements {
@@ -800,67 +805,73 @@ pub async fn handle_python_reqs(
#[cfg(all(feature = "enterprise", feature = "parquet"))]
if req_with_penv.len() > 0 {
if let Some(os) = OBJECT_STORE_CACHE_SETTINGS.read().await.clone() {
if matches!(get_license_plan().await, LicensePlan::Pro) {
append_logs(job_id.clone(), w_id.to_string(), format!("s3 cache not available in Pro Plan"), db).await;
tracing::warn!("S3 cache not available in the pro plan");
} else {
let (done_tx, mut done_rx) = tokio::sync::mpsc::channel(1);
let job_id_2 = job_id.clone();
let db_2 = db.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {
if let Err(e) = sqlx::query_scalar!("UPDATE queue SET last_ping = now() WHERE id = $1", &job_id_2)
.execute(&db_2)
.await {
tracing::error!("failed to update last_ping: {}", e);
}
}
_ = done_rx.recv() => {
break;
let (done_tx, mut done_rx) = tokio::sync::mpsc::channel(1);
let job_id_2 = job_id.clone();
let db_2 = db.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = tokio::time::sleep(tokio::time::Duration::from_secs(5)) => {
if let Err(e) = sqlx::query_scalar!("UPDATE queue SET last_ping = now() WHERE id = $1", &job_id_2)
.execute(&db_2)
.await {
tracing::error!("failed to update last_ping: {}", e);
}
}
}
});
let start = std::time::Instant::now();
let futures = req_with_penv.clone().into_iter().map(|(req, venv_p)| {
let os = os.clone();
async move {
if pull_from_tar(os, venv_p.clone()).await.is_ok() {
PullFromTar::Pulled(venv_p.to_string())
} else {
PullFromTar::NotPulled(req.to_string(), venv_p.to_string())
}
}}).collect::<Vec<_>>();
let results = futures::future::join_all(futures).await;
req_with_penv.clear();
done_tx.send(()).await.expect("failed to send done");
let mut pulled = vec![];
for result in results {
match result {
PullFromTar::Pulled(venv_p) => {
pulled.push(venv_p.split("/").last().unwrap_or_default().to_string());
req_paths.push(venv_p);
}
PullFromTar::NotPulled(req, venv_p) => {
req_with_penv.push((req, venv_p));
_ = done_rx.recv() => {
break;
}
}
}
if pulled.len() > 0 {
append_logs(job_id.clone(), w_id.to_string(), format!("pulled {} from s3 cache in {}ms", pulled.join(", "), start.elapsed().as_millis()), db).await;
});
let start = std::time::Instant::now();
let futures = req_with_penv
.clone()
.into_iter()
.map(|(req, venv_p)| {
let os = os.clone();
async move {
if pull_from_tar(os, venv_p.clone()).await.is_ok() {
PullFromTar::Pulled(venv_p.to_string())
} else {
PullFromTar::NotPulled(req.to_string(), venv_p.to_string())
}
}
})
.collect::<Vec<_>>();
let results = futures::future::join_all(futures).await;
req_with_penv.clear();
done_tx.send(()).await.expect("failed to send done");
let mut pulled = vec![];
for result in results {
match result {
PullFromTar::Pulled(venv_p) => {
pulled.push(venv_p.split("/").last().unwrap_or_default().to_string());
req_paths.push(venv_p);
}
PullFromTar::NotPulled(req, venv_p) => {
req_with_penv.push((req, venv_p));
}
}
}
}
if pulled.len() > 0 {
append_logs(
job_id.clone(),
w_id.to_string(),
format!(
"pulled {} from s3 cache in {}ms",
pulled.join(", "),
start.elapsed().as_millis()
),
db,
)
.await;
}
}
}
for (req, venv_p) in req_with_penv {
let mut logs1 = String::new();
logs1.push_str("\n\n--- PIP INSTALL ---\n");
logs1.push_str(&format!("\n{req} is being installed for the first time.\n It will be cached for all ulterior uses."));