diff --git a/backend/.vscode/settings.json b/backend/.vscode/settings.json index dc3f727cad..de6dac71e9 100644 --- a/backend/.vscode/settings.json +++ b/backend/.vscode/settings.json @@ -1,3 +1,4 @@ { - "python.analysis.typeCheckingMode": "basic" + "python.analysis.typeCheckingMode": "basic", + "rust-analyzer.linkedProjects": ["./windmill-common/Cargo.toml"] } diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 950f86a9a3..48c758ee7d 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -5105,6 +5105,9 @@ dependencies = [ "dotenv", "futures", "git-version", + "lazy_static", + "once_cell", + "prometheus", "rand 0.8.5", "reqwest", "rsa", @@ -5117,6 +5120,7 @@ dependencies = [ "tokio-metrics", "tracing", "url", + "uuid 1.3.1", "windmill-api", "windmill-api-client", "windmill-common", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index f09939b746..93103b4fa2 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -50,6 +50,10 @@ base64.workspace = true sha2.workspace = true rsmq_async.workspace = true url.workspace = true +lazy_static.workspace = true +once_cell.workspace = true +prometheus.workspace = true +uuid.workspace = true [dev-dependencies] diff --git a/backend/src/main.rs b/backend/src/main.rs index 27a3fe084b..2b08c7e980 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -9,6 +9,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use git_version::git_version; +use monitor::handle_zombie_jobs_periodically; use sqlx::{Pool, Postgres}; use windmill_common::{utils::rd_string, METRICS_ADDR}; @@ -18,6 +19,7 @@ const DEFAULT_PORT: u16 = 8000; const DEFAULT_SERVER_BIND_ADDR: Ipv4Addr = Ipv4Addr::new(0, 0, 0, 0); mod ee; +mod monitor; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -230,7 +232,7 @@ pub fn monitor_db let rx2 = rx.resubscribe(); let base_internal_url = base_internal_url.to_string(); tokio::spawn(async move { - windmill_worker::handle_zombie_jobs_periodically(&db1, rx, &base_internal_url, rsmq).await + handle_zombie_jobs_periodically(&db1, rx, &base_internal_url, rsmq).await }); tokio::spawn(async move { windmill_api::delete_expired_items_perdiodically(&db2, rx2).await }); } diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs new file mode 100644 index 0000000000..f604ed9844 --- /dev/null +++ b/backend/src/monitor.rs @@ -0,0 +1,148 @@ +use std::time::Duration; + +use once_cell::sync::OnceCell; +use sqlx::{Pool, Postgres}; +use tokio::sync::mpsc; +use uuid::Uuid; +use windmill_common::{error, jobs::{JobKind, QueuedJob}, METRICS_ENABLED}; +use windmill_worker::{ + create_token_for_owner, handle_job_error, AuthedClient, SESSION_TOKEN_EXPIRY, +}; + +lazy_static::lazy_static! { + static ref ZOMBIE_JOB_TIMEOUT: String = std::env::var("ZOMBIE_JOB_TIMEOUT") + .ok() + .and_then(|x| x.parse::().ok()) + .unwrap_or_else(|| "30".to_string()); + + + pub static ref RESTART_ZOMBIE_JOBS: bool = std::env::var("RESTART_ZOMBIE_JOBS") + .ok() + .and_then(|x| x.parse::().ok()) + .unwrap_or(true); + + static ref QUEUE_ZOMBIE_RESTART_COUNT: prometheus::IntCounter = prometheus::register_int_counter!( + "queue_zombie_restart_count", + "Total number of jobs restarted due to ping timeout." + ) + .unwrap(); + static ref QUEUE_ZOMBIE_DELETE_COUNT: prometheus::IntCounter = prometheus::register_int_counter!( + "queue_zombie_delete_count", + "Total number of jobs deleted due to their ping timing out in an unrecoverable state." + ) + .unwrap(); +} + +pub async fn handle_zombie_jobs_periodically< + R: rsmq_async::RsmqConnection + Send + Sync + Clone, +>( + db: &Pool, + mut rx: tokio::sync::broadcast::Receiver<()>, + base_internal_url: &str, + rsmq: Option, +) { + loop { + handle_zombie_jobs(db, base_internal_url, rsmq.clone()).await; + + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(30)) => (), + _ = rx.recv() => { + println!("received killpill for monitor job"); + break; + } + } + } +} + +async fn handle_zombie_jobs( + db: &Pool, + base_internal_url: &str, + rsmq: Option, +) { + if *RESTART_ZOMBIE_JOBS { + let restarted = sqlx::query!( + "UPDATE queue SET running = false, started_at = null, logs = logs || '\nRestarted job after not receiving job''s ping for too long the ' || now() || '\n\n' WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 AND job_kind != $3 AND same_worker = false RETURNING id, workspace_id, last_ping", + *ZOMBIE_JOB_TIMEOUT, + JobKind::Flow: JobKind, + JobKind::FlowPreview: JobKind, + ) + .fetch_all(db) + .await + .ok() + .unwrap_or_else(|| vec![]); + + if *METRICS_ENABLED { + QUEUE_ZOMBIE_RESTART_COUNT.inc_by(restarted.len() as _); + } + for r in restarted { + tracing::info!( + "restarted zombie job {} {} {}", + r.id, + r.workspace_id, + r.last_ping + ); + } + } + + let mut timeout_query = "SELECT * FROM queue WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2".to_string(); + if *RESTART_ZOMBIE_JOBS { + timeout_query.push_str(" AND same_worker = true"); + }; + let timeouts = sqlx::query_as::<_, QueuedJob>(&timeout_query) + .bind(ZOMBIE_JOB_TIMEOUT.as_str()) + .bind(JobKind::Flow) + .fetch_all(db) + .await + .ok() + .unwrap_or_else(|| vec![]); + + if *METRICS_ENABLED { + QUEUE_ZOMBIE_DELETE_COUNT.inc_by(timeouts.len() as _); + } + + for job in timeouts { + tracing::info!("timedout zombie job {} {}", job.id, job.workspace_id,); + + // since the job is unrecoverable, the same worker queue should never be sent anything + let (same_worker_tx_never_used, _same_worker_rx_never_used) = mpsc::channel::(1); + + let token = create_token_for_owner( + &db, + &job.workspace_id, + &job.permissioned_as, + "ephemeral-zombie-jobs", + *SESSION_TOKEN_EXPIRY, + &job.email, + ) + .await + .expect("could not create job token"); + + let client = AuthedClient { + base_internal_url: base_internal_url.to_string(), + token, + workspace: job.workspace_id.to_string(), + client: OnceCell::new(), + }; + + let last_ping = job.last_ping.clone(); + let _ = handle_job_error( + db, + &client, + job, + error::Error::ExecutionErr(format!( + "Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {})", + last_ping + .map(|x| x.to_string()) + .unwrap_or_else(|| "no ping".to_string()), + *ZOMBIE_JOB_TIMEOUT + )), + None, + true, + same_worker_tx_never_used, + "", + base_internal_url, + rsmq.clone(), + ) + .await; + } +} diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index 4dfb704330..f8d8edd6b5 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -6,9 +6,10 @@ use windmill_api::jobs::{CompletedJob, Job}; use windmill_common::{ flow_status::{FlowStatus, FlowStatusModule}, flows::{FlowModule, FlowModuleValue, FlowValue, InputTransform}, + jobs::{JobPayload, RawCode}, scripts::ScriptLang, }; -use windmill_queue::{get_queued_job, JobPayload, RawCode}; +use windmill_queue::get_queued_job; async fn initialize_tracing() { use std::sync::Once; @@ -128,8 +129,7 @@ mod suspend_resume { use futures::{Stream, StreamExt}; use serde_json::json; use sqlx::{query_scalar, types::Uuid}; - use windmill_common::flows::FlowValue; - use windmill_queue::JobPayload; + use windmill_common::{flows::FlowValue, jobs::JobPayload}; use super::*; @@ -393,7 +393,6 @@ mod retry { use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use windmill_common::flows::FlowValue; - use windmill_queue::JobPayload; use super::*; diff --git a/backend/windmill-api/src/apps.rs b/backend/windmill-api/src/apps.rs index 8f3c36134c..3a8f590e13 100644 --- a/backend/windmill-api/src/apps.rs +++ b/backend/windmill-api/src/apps.rs @@ -32,12 +32,13 @@ use windmill_audit::{audit_log, ActionKind}; use windmill_common::{ apps::ListAppQuery, error::{to_anyhow, Error, JsonResult, Result}, + jobs::{JobPayload, RawCode}, users::username_to_permissioned_as, utils::{ http_get_from_hub, list_elems_from_hub, not_found_if_none, paginate, Pagination, StripPath, }, }; -use windmill_queue::{push, JobPayload, QueueTransaction, RawCode}; +use windmill_queue::{push, QueueTransaction}; pub fn workspaced_service() -> Router { Router::new() diff --git a/backend/windmill-api/src/flows.rs b/backend/windmill-api/src/flows.rs index fc8f20b709..65bbd4441b 100644 --- a/backend/windmill-api/src/flows.rs +++ b/backend/windmill-api/src/flows.rs @@ -29,9 +29,9 @@ use windmill_common::{ schedule::Schedule, utils::{ http_get_from_hub, list_elems_from_hub, not_found_if_none, paginate, Pagination, StripPath, - }, + }, jobs::JobPayload, }; -use windmill_queue::{push, schedule::push_scheduled_job, JobPayload, QueueTransaction}; +use windmill_queue::{push, schedule::push_scheduled_job, QueueTransaction}; pub fn workspaced_service() -> Router { Router::new() diff --git a/backend/windmill-api/src/inputs.rs b/backend/windmill-api/src/inputs.rs index a5be6f70b8..80e8741c7c 100644 --- a/backend/windmill-api/src/inputs.rs +++ b/backend/windmill-api/src/inputs.rs @@ -22,11 +22,10 @@ use std::{ }; use windmill_common::{ error::JsonResult, + jobs::JobKind, scripts::to_i64, utils::{paginate, Pagination}, }; -use windmill_queue::JobKind; - pub fn workspaced_service() -> Router { Router::new() .route("/history", get(get_input_history)) diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index cc7aaf7a9b..93f14bf933 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -31,14 +31,13 @@ use windmill_common::{ error::{self, to_anyhow, Error}, flow_status::{Approval, FlowStatus, FlowStatusModule}, flows::FlowValue, + jobs::{JobKind, JobPayload, QueuedJob, RawCode}, oauth2::HmacSha256, scripts::{ScriptHash, ScriptLang}, users::username_to_permissioned_as, utils::{not_found_if_none, now_from_db, paginate, require_admin, Pagination, StripPath}, }; -use windmill_queue::{ - get_queued_job, push, JobKind, JobPayload, QueueTransaction, QueuedJob, RawCode, -}; +use windmill_queue::{get_queued_job, push, QueueTransaction}; pub fn workspaced_service() -> Router { Router::new() diff --git a/backend/windmill-api/src/oauth2.rs b/backend/windmill-api/src/oauth2.rs index 99d29a93f7..bb30d74868 100644 --- a/backend/windmill-api/src/oauth2.rs +++ b/backend/windmill-api/src/oauth2.rs @@ -29,6 +29,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use sqlx::{Postgres, Transaction}; use tower_cookies::{Cookie, Cookies}; use windmill_audit::{audit_log, ActionKind}; +use windmill_common::jobs::JobPayload; use windmill_common::users::username_to_permissioned_as; use windmill_common::utils::{not_found_if_none, now_from_db}; @@ -44,7 +45,7 @@ use crate::{BASE_URL, HTTP_CLIENT, IS_SECURE, OAUTH_CLIENTS, SLACK_SIGNING_SECRE use windmill_common::error::{self, to_anyhow, Error}; use windmill_common::oauth2::*; -use windmill_queue::{JobPayload, QueueTransaction}; +use windmill_queue::QueueTransaction; use std::{fs, str}; diff --git a/backend/windmill-api/src/schedule.rs b/backend/windmill-api/src/schedule.rs index 9ab73e4272..6bea75bbf5 100644 --- a/backend/windmill-api/src/schedule.rs +++ b/backend/windmill-api/src/schedule.rs @@ -23,9 +23,9 @@ use windmill_audit::{audit_log, ActionKind}; use windmill_common::{ error::{Error, JsonResult, Result}, schedule::Schedule, - utils::{not_found_if_none, paginate, Pagination, StripPath}, + utils::{not_found_if_none, paginate, Pagination, StripPath}, jobs::JobKind, }; -use windmill_queue::{self, schedule::push_scheduled_job, JobKind, QueueTransaction}; +use windmill_queue::{self, schedule::push_scheduled_job, QueueTransaction}; pub fn workspaced_service() -> Router { Router::new() diff --git a/backend/windmill-api/src/scripts.rs b/backend/windmill-api/src/scripts.rs index 3830a7578c..9fb0ec193d 100644 --- a/backend/windmill-api/src/scripts.rs +++ b/backend/windmill-api/src/scripts.rs @@ -32,6 +32,7 @@ use std::{ use windmill_audit::{audit_log, ActionKind}; use windmill_common::{ error::{Error, JsonResult, Result}, + jobs::JobPayload, schedule::Schedule, scripts::{ to_i64, HubScript, ListScriptQuery, ListableScript, NewScript, Script, ScriptHash, @@ -428,7 +429,7 @@ async fn create_script( let (_, new_tx) = windmill_queue::push( tx, &w_id, - windmill_queue::JobPayload::Dependencies { hash, dependencies, language: ns.language }, + JobPayload::Dependencies { hash, dependencies, language: ns.language }, serde_json::Map::new(), &authed.username, &authed.email, diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs index 8b13789179..7833aa7c8c 100644 --- a/backend/windmill-common/src/jobs.rs +++ b/backend/windmill-common/src/jobs.rs @@ -1 +1,122 @@ +use serde::{Deserialize, Serialize}; +use uuid::Uuid; +use crate::{ + flow_status::FlowStatus, + flows::FlowValue, + scripts::{ScriptHash, ScriptLang}, +}; + +#[derive(sqlx::Type, Serialize, Deserialize, Debug, PartialEq, Clone)] +#[sqlx(type_name = "JOB_KIND", rename_all = "lowercase")] +#[serde(rename_all(serialize = "lowercase"))] +pub enum JobKind { + Script, + #[allow(non_camel_case_types)] + Script_Hub, + Preview, + Dependencies, + Flow, + FlowPreview, + Identity, + FlowDependencies, +} + +#[derive(Debug, sqlx::FromRow, Serialize, Clone)] +pub struct QueuedJob { + pub workspace_id: String, + pub id: Uuid, + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_job: Option, + pub created_by: String, + pub created_at: chrono::DateTime, + #[serde(skip_serializing_if = "Option::is_none")] + pub started_at: Option>, + pub scheduled_for: chrono::DateTime, + pub running: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub script_hash: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub script_path: Option, + pub args: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub logs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub raw_code: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub raw_lock: Option, + pub canceled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub canceled_by: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub canceled_reason: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub last_ping: Option>, + pub job_kind: JobKind, + #[serde(skip_serializing_if = "Option::is_none")] + pub schedule_path: Option, + pub permissioned_as: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub flow_status: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub raw_flow: Option, + pub is_flow_step: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub language: Option, + pub same_worker: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub pre_run_error: Option, + pub email: String, + pub visible_to_owner: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub suspend: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub mem_peak: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub root_job: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub leaf_jobs: Option, +} + +impl QueuedJob { + pub fn script_path(&self) -> &str { + self.script_path + .as_ref() + .map(String::as_str) + .unwrap_or("tmp/main") + } +} + +impl QueuedJob { + pub fn parse_raw_flow(&self) -> Option { + self.raw_flow + .as_ref() + .and_then(|v| serde_json::from_value::(v.clone()).ok()) + } + + pub fn parse_flow_status(&self) -> Option { + self.flow_status + .as_ref() + .and_then(|v| serde_json::from_value::(v.clone()).ok()) + } +} + +#[derive(Debug, Clone)] +pub enum JobPayload { + ScriptHub { path: String }, + ScriptHash { hash: ScriptHash, path: String }, + Code(RawCode), + Dependencies { hash: ScriptHash, dependencies: String, language: ScriptLang }, + FlowDependencies { path: String }, + Flow(String), + RawFlow { value: FlowValue, path: Option }, + Identity, +} + +#[derive(Clone, Serialize, Deserialize, Debug)] +pub struct RawCode { + pub content: String, + pub path: Option, + pub language: ScriptLang, + pub lock: Option, +} diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index a0411efd1b..14a46cba95 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -15,6 +15,7 @@ pub mod error; pub mod external_ip; pub mod flow_status; pub mod flows; +pub mod jobs; pub mod more_serde; pub mod oauth2; pub mod schedule; diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 17bdb9a5b2..391c654ce1 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -10,7 +10,6 @@ use std::collections::HashMap; use anyhow::Context; use reqwest::Client; -use serde::{Deserialize, Serialize}; use sqlx::{Pool, Postgres, Transaction}; use tracing::{instrument, Instrument}; use ulid::Ulid; @@ -20,6 +19,7 @@ use windmill_common::{ error::{self, Error}, flow_status::{FlowStatus, JobResult, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL}, flows::{FlowModule, FlowModuleValue, FlowValue}, + jobs::{JobKind, JobPayload, QueuedJob, RawCode}, scripts::{get_full_hub_script_by_path, HubScript, ScriptHash, ScriptLang}, utils::StripPath, METRICS_ENABLED, @@ -655,117 +655,3 @@ pub async fn get_hub_script( .await .map(|e| e) } - -#[derive(Debug, sqlx::FromRow, Serialize, Clone)] -pub struct QueuedJob { - pub workspace_id: String, - pub id: Uuid, - #[serde(skip_serializing_if = "Option::is_none")] - pub parent_job: Option, - pub created_by: String, - pub created_at: chrono::DateTime, - #[serde(skip_serializing_if = "Option::is_none")] - pub started_at: Option>, - pub scheduled_for: chrono::DateTime, - pub running: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub script_hash: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub script_path: Option, - pub args: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub logs: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub raw_code: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub raw_lock: Option, - pub canceled: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub canceled_by: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub canceled_reason: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub last_ping: Option>, - pub job_kind: JobKind, - #[serde(skip_serializing_if = "Option::is_none")] - pub schedule_path: Option, - pub permissioned_as: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub flow_status: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub raw_flow: Option, - pub is_flow_step: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub language: Option, - pub same_worker: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub pre_run_error: Option, - pub email: String, - pub visible_to_owner: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub suspend: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub mem_peak: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub root_job: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub leaf_jobs: Option, -} - -impl QueuedJob { - pub fn script_path(&self) -> &str { - self.script_path - .as_ref() - .map(String::as_str) - .unwrap_or("tmp/main") - } -} - -impl QueuedJob { - pub fn parse_raw_flow(&self) -> Option { - self.raw_flow - .as_ref() - .and_then(|v| serde_json::from_value::(v.clone()).ok()) - } - - pub fn parse_flow_status(&self) -> Option { - self.flow_status - .as_ref() - .and_then(|v| serde_json::from_value::(v.clone()).ok()) - } -} - -#[derive(sqlx::Type, Serialize, Deserialize, Debug, PartialEq, Clone)] -#[sqlx(type_name = "JOB_KIND", rename_all = "lowercase")] -#[serde(rename_all(serialize = "lowercase"))] -pub enum JobKind { - Script, - #[allow(non_camel_case_types)] - Script_Hub, - Preview, - Dependencies, - Flow, - FlowPreview, - Identity, - FlowDependencies, -} - -#[derive(Debug, Clone)] -pub enum JobPayload { - ScriptHub { path: String }, - ScriptHash { hash: ScriptHash, path: String }, - Code(RawCode), - Dependencies { hash: ScriptHash, dependencies: String, language: ScriptLang }, - FlowDependencies { path: String }, - Flow(String), - RawFlow { value: FlowValue, path: Option }, - Identity, -} - -#[derive(Clone, Serialize, Deserialize, Debug)] -pub struct RawCode { - pub content: String, - pub path: Option, - pub language: ScriptLang, - pub lock: Option, -} diff --git a/backend/windmill-queue/src/schedule.rs b/backend/windmill-queue/src/schedule.rs index a59b1b2e86..e58c217ac4 100644 --- a/backend/windmill-queue/src/schedule.rs +++ b/backend/windmill-queue/src/schedule.rs @@ -6,8 +6,10 @@ * LICENSE-AGPL for a copy of the license. */ -use crate::{push, JobPayload}; +use crate::push; +use crate::QueueTransaction; use sqlx::{query_scalar, Postgres, Transaction}; +use windmill_common::jobs::JobPayload; use std::str::FromStr; use windmill_common::{ error::{self, Result}, @@ -15,7 +17,6 @@ use windmill_common::{ users::username_to_permissioned_as, utils::{now_from_db, StripPath}, }; -use crate::{QueueTransaction}; pub async fn push_scheduled_job<'c, R: rsmq_async::RsmqConnection + Send + 'c>( mut tx: QueueTransaction<'c, R>, diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs new file mode 100644 index 0000000000..a86e07539a --- /dev/null +++ b/backend/windmill-worker/src/common.rs @@ -0,0 +1,41 @@ +use sqlx::{Pool, Postgres}; +use tokio::{fs::File, io::AsyncReadExt}; +use windmill_common::error::{self, Error}; +use windmill_queue::CLOUD_HOSTED; + +use crate::MAX_RESULT_SIZE; + +pub async fn read_result(job_dir: &str) -> error::Result { + let mut file = File::open(format!("{job_dir}/result.json")).await?; + let mut content = "".to_string(); + file.read_to_string(&mut content).await?; + if *CLOUD_HOSTED && content.len() > MAX_RESULT_SIZE { + return Err(Error::ExecutionErr("Result is too large for the cloud app (limit 2MB). + If using this script as part of the flow, use the shared folder to pass heavy data between steps.".to_owned())); + } + serde_json::from_str(&content) + .map_err(|e| Error::ExecutionErr(format!("Error parsing result: {e}"))) +} + +#[tracing::instrument(level = "trace", skip_all)] +pub async fn set_logs(logs: &str, id: &uuid::Uuid, db: &Pool) { + if sqlx::query!( + "UPDATE queue SET logs = $1 WHERE id = $2", + logs.to_owned(), + id + ) + .execute(db) + .await + .is_err() + { + tracing::error!(%id, "error updating logs for id {id}") + }; +} + +pub fn capitalize(s: &str) -> String { + let mut c = s.chars(); + match c.next() { + None => String::new(), + Some(f) => f.to_uppercase().collect::() + c.as_str(), + } +} diff --git a/backend/windmill-worker/src/global_cache.rs b/backend/windmill-worker/src/global_cache.rs new file mode 100644 index 0000000000..088173bde4 --- /dev/null +++ b/backend/windmill-worker/src/global_cache.rs @@ -0,0 +1,249 @@ +#[cfg(feature = "enterprise")] +use crate::{ + DENO_CACHE_DIR, DENO_TMP_CACHE_DIR, GO_CACHE_DIR, GO_TMP_CACHE_DIR, PIP_CACHE_DIR, + PIP_TMP_CACHE_DIR, +}; + +use crate::{ROOT_CACHE_DIR, ROOT_TMP_CACHE_DIR}; +#[cfg(feature = "enterprise")] +use std::process::Stdio; +#[cfg(feature = "enterprise")] +use tokio::{fs::DirBuilder, process::Command, sync::mpsc::Sender, time::Instant}; +use windmill_common::error; + +#[cfg(feature = "enterprise")] +const TAR_CACHE_FILENAME: &str = "entirecache.tar"; + +#[cfg(feature = "enterprise")] +pub async fn copy_cache_from_bucket( + bucket: &str, + tx: Option>, +) -> Option> { + tracing::info!("Copying cache from bucket in the background {bucket}"); + let bucket = bucket.to_string(); + let tx_is_some = tx.is_some(); + let f = async move { + let elapsed = Instant::now(); + + match Command::new("rclone") + .arg("copy") + .arg(format!(":s3,env_auth=true:{bucket}")) + .arg(if tx_is_some { + ROOT_TMP_CACHE_DIR + } else { + ROOT_CACHE_DIR + }) + .arg("--size-only") + .arg("--fast-list") + .arg("--exclude") + .arg(format!("\"{TAR_CACHE_FILENAME},/deno/gen/file/**\"")) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + h.wait().await.unwrap(); + } + Err(e) => tracing::warn!("Failed to run periodic job pull. Error: {:?}", e), + } + tracing::info!( + "Finished copying cache from bucket {bucket}, took {:?}s", + elapsed.elapsed().as_secs() + ); + + for x in if !tx_is_some { + [PIP_CACHE_DIR, DENO_CACHE_DIR, GO_CACHE_DIR] + } else { + [PIP_TMP_CACHE_DIR, DENO_TMP_CACHE_DIR, GO_TMP_CACHE_DIR] + } { + DirBuilder::new() + .recursive(true) + .create(x) + .await + .expect("could not create initial worker dir"); + } + + if let Some(tx) = tx { + tx.send(()).await.expect("can send copy cache signal"); + } + }; + if tx_is_some { + return Some(tokio::spawn(f)); + } else { + f.await; + return None; + } +} + +#[cfg(feature = "enterprise")] +pub async fn copy_cache_to_bucket(bucket: &str) { + tracing::info!("Copying cache to bucket {bucket}"); + let elapsed = Instant::now(); + match Command::new("rclone") + .arg("copy") + .arg(ROOT_CACHE_DIR) + .arg(format!(":s3,env_auth=true:{bucket}")) + .arg("--size-only") + .arg("--fast-list") + .arg("--exclude") + .arg(format!("\"{TAR_CACHE_FILENAME},/deno/gen/file/**\"")) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + h.wait().await.unwrap(); + } + Err(e) => tracing::info!("Failed to run periodic job push. Error: {:?}", e), + } + tracing::info!( + "Finished copying cache to bucket {bucket}, took: {:?}s", + elapsed.elapsed().as_secs() + ); +} + +#[cfg(feature = "enterprise")] +pub async fn copy_cache_to_bucket_as_tar(bucket: &str) { + tracing::info!("Copying cache to bucket {bucket} as tar"); + let elapsed = Instant::now(); + + match Command::new("tar") + .current_dir(ROOT_CACHE_DIR) + .arg("-c") + .arg("-f") + .arg(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")) + .args(&["pip", "go", "deno"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + if !h.wait().await.unwrap().success() { + tracing::info!("Failed to tar cache"); + return; + } + } + Err(e) => { + tracing::info!("Failed tar cache. Error: {e:?}"); + return; + } + } + + let tar_metadata = tokio::fs::metadata(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")).await; + if tar_metadata.is_err() || tar_metadata.as_ref().unwrap().len() == 0 { + tracing::info!("Failed to tar cache"); + return; + } + + match Command::new("rclone") + .current_dir(ROOT_CACHE_DIR) + .arg("copyto") + .arg(TAR_CACHE_FILENAME) + .arg(format!(":s3,env_auth=true:{bucket}/{TAR_CACHE_FILENAME}")) + .arg("-vv") + .arg("--size-only") + .arg("--fast-list") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + h.wait().await.unwrap(); + } + Err(e) => tracing::info!("Failed to copy tar cache to bucket. Error: {:?}", e), + } + + if let Err(e) = tokio::fs::remove_file(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")).await { + tracing::info!("Failed to remove tar cache. Error: {:?}", e); + }; + + tracing::info!( + "Finished copying cache to bucket {bucket} as tar, took: {:?}s. Size of new tar: {}", + elapsed.elapsed().as_secs(), + tar_metadata.unwrap().len() + ); +} + +#[cfg(feature = "enterprise")] +pub async fn copy_cache_from_bucket_as_tar(bucket: &str) -> bool { + tracing::info!("Copying cache from bucket {bucket} as tar"); + let elapsed = Instant::now(); + + match Command::new("rclone") + .arg("copyto") + .arg(format!(":s3,env_auth=true:{bucket}/{TAR_CACHE_FILENAME}")) + .arg(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")) + .arg("-vv") + .arg("--size-only") + .arg("--fast-list") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + if !h.wait().await.unwrap().success() { + tracing::info!("Failed to download tar cache, continuing nonetheless"); + return false; + } + } + Err(e) => { + tracing::info!("Failed to download tar cache, continuing nonetheless. Error: {e:?}"); + return false; + } + } + + match Command::new("tar") + .current_dir(ROOT_TMP_CACHE_DIR) + .arg("-xpvf") + .arg(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .spawn() + { + Ok(mut h) => { + if !h.wait().await.unwrap().success() { + tracing::info!("Failed to untar cache, continuing nonetheless"); + return false; + } + } + Err(e) => { + tracing::warn!("Failed to untar cache, continuing nonetheless. Error: {e:?}"); + return false; + } + } + + if let Err(e) = + tokio::fs::remove_file(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")).await + { + tracing::info!("Failed to remove tar cache. Error: {:?}", e); + }; + + let r = move_tmp_cache_to_cache().await.is_ok(); + + tracing::info!( + "Finished copying cache from bucket {bucket} as tar, took: {:?}s. copy success: {r}", + elapsed.elapsed().as_secs() + ); + return r; +} + +// async fn check_if_bucket_syncable(bucket: &str) -> bool { +// match Command::new("rclone") +// .arg("lsf") +// .arg(format!(":s3,env_auth=true:{bucket}/NOSYNC")) + +// .arg("-vv") +// .arg("--fast-list") +// .stdin(Stdio::null()) +// .stdout(Stdio::null()) +// .output() +// .await; +// return true; +// } + +pub async fn move_tmp_cache_to_cache() -> error::Result<()> { + tokio::fs::remove_dir_all(ROOT_CACHE_DIR).await?; + tokio::fs::rename(ROOT_TMP_CACHE_DIR, ROOT_CACHE_DIR).await?; + tracing::info!("Finished moving tmp cache to cache"); + Ok(()) +} diff --git a/backend/windmill-worker/src/go_executor.rs b/backend/windmill-worker/src/go_executor.rs new file mode 100644 index 0000000000..a08e2f6796 --- /dev/null +++ b/backend/windmill-worker/src/go_executor.rs @@ -0,0 +1,362 @@ +use std::process::Stdio; + +use itertools::Itertools; +use tokio::{ + fs::{DirBuilder, File}, + io::AsyncReadExt, + process::Command, +}; +use uuid::Uuid; +use windmill_common::{ + error::{self, Error}, + jobs::QueuedJob, + utils::calculate_hash, +}; +use windmill_parser_go::parse_go_imports; + +use crate::{ + common::{capitalize, read_result, set_logs}, + create_args_and_out_file, get_reserved_variables, handle_child, write_file, + AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, GOPRIVATE, GO_CACHE_DIR, HOME_ENV, + NETRC, NSJAIL_PATH, PATH_ENV, +}; + +const GO_REQ_SPLITTER: &str = "//go.sum\n"; +const NSJAIL_CONFIG_RUN_GO_CONTENT: &str = include_str!("../nsjail/run.go.config.proto"); + +lazy_static::lazy_static! { + static ref GO_PATH: String = std::env::var("GO_PATH").unwrap_or_else(|_| "/usr/bin/go".to_string()); +} + +#[tracing::instrument(level = "trace", skip_all)] +pub async fn handle_go_job( + logs: &mut String, + job: &QueuedJob, + db: &sqlx::Pool, + client: &AuthedClientBackgroundTask, + inner_content: &str, + job_dir: &str, + requirements_o: Option, + shared_mount: &str, + base_internal_url: &str, + worker_name: &str, +) -> Result { + //go does not like executing modules at temp root + let job_dir = &format!("{job_dir}/go"); + let (skip_go_mod, skip_tidy) = if let Some(requirements) = requirements_o { + gen_go_mod(inner_content, job_dir, &requirements).await? + } else { + (false, false) + }; + + logs.push_str("\n\n--- GO DEPENDENCIES SETUP ---\n"); + set_logs(logs, &job.id, db).await; + + install_go_dependencies( + &job.id, + inner_content, + logs, + job_dir, + db, + true, + skip_go_mod, + skip_tidy, + worker_name, + &job.workspace_id, + ) + .await?; + + logs.push_str("\n\n--- GO CODE EXECUTION ---\n"); + set_logs(logs, &job.id, db).await; + let client = &client.get_authed().await; + create_args_and_out_file(client, job, job_dir).await?; + { + let sig = windmill_parser_go::parse_go_sig(&inner_content)?; + drop(inner_content); + + const WRAPPER_CONTENT: &str = r#"package main + +import ( + "encoding/json" + "os" + "fmt" + "mymod/inner" +) + +func main() {{ + + dat, err := os.ReadFile("args.json") + if err != nil {{ + fmt.Println(err) + os.Exit(1) + }} + + var req inner.Req + + if err := json.Unmarshal(dat, &req); err != nil {{ + fmt.Println(err) + os.Exit(1) + }} + + res, err := inner.Run(req) + if err != nil {{ + fmt.Println(err) + os.Exit(1) + }} + res_json, err := json.Marshal(res) + if err != nil {{ + fmt.Println(err) + os.Exit(1) + }} + f, err := os.OpenFile("result.json", os.O_APPEND|os.O_WRONLY, os.ModeAppend) + if err != nil {{ + fmt.Println(err) + os.Exit(1) + }} + _, err = f.WriteString(string(res_json)) + if err != nil {{ + fmt.Println(err) + os.Exit(1) + }} +}}"#; + + write_file(job_dir, "main.go", WRAPPER_CONTENT).await?; + + { + let spread = &sig + .args + .clone() + .into_iter() + .map(|x| format!("req.{}", capitalize(&x.name))) + .join(", "); + let req_body = &sig + .args + .into_iter() + .map(|x| { + format!( + "{} {} `json:\"{}\"`", + capitalize(&x.name), + windmill_parser_go::otyp_to_string(x.otyp), + x.name + ) + }) + .join("\n"); + let runner_content: String = format!( + r#"package inner +type Req struct {{ + {req_body} +}} + +func Run(req Req) (interface{{}}, error){{ + return main({spread}) +}} + +"#, + ); + write_file(&format!("{job_dir}/inner"), "runner.go", &runner_content).await?; + } + } + let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?; + reserved_variables.insert("RUST_LOG".to_string(), "info".to_string()); + + let child = if !*DISABLE_NSJAIL { + let _ = write_file( + job_dir, + "run.config.proto", + &NSJAIL_CONFIG_RUN_GO_CONTENT + .replace("{JOB_DIR}", job_dir) + .replace("{CACHE_DIR}", GO_CACHE_DIR) + .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()) + .replace("{SHARED_MOUNT}", shared_mount), + ) + .await?; + let build_go = Command::new(GO_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + .env("PATH", PATH_ENV.as_str()) + .env("BASE_INTERNAL_URL", base_internal_url) + .env("GOPATH", GO_CACHE_DIR) + .env("HOME", HOME_ENV.as_str()) + .args(vec!["build", "main.go"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + handle_child( + &job.id, + db, + logs, + build_go, + false, + worker_name, + &job.workspace_id, + ) + .await?; + + Command::new(NSJAIL_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + .envs(reserved_variables) + .env("PATH", PATH_ENV.as_str()) + .env("BASE_INTERNAL_URL", base_internal_url) + .args(vec!["--config", "run.config.proto", "--", "/tmp/go/main"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + } else { + if let Some(ref netrc) = *NETRC { + write_file(&HOME_ENV, ".netrc", netrc).await?; + } + Command::new(GO_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + .envs(reserved_variables) + .env("PATH", PATH_ENV.as_str()) + .env("BASE_INTERNAL_URL", base_internal_url) + .env("GOPATH", GO_CACHE_DIR) + .env("GOPRIVATE", GOPRIVATE.as_ref().unwrap_or(&String::new())) + .env("HOME", HOME_ENV.as_str()) + .args(vec!["run", "main.go"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + }; + handle_child( + &job.id, + db, + logs, + child, + !*DISABLE_NSJAIL, + worker_name, + &job.workspace_id, + ) + .await?; + read_result(job_dir).await +} + +async fn gen_go_mod( + inner_content: &str, + job_dir: &str, + requirements: &str, +) -> error::Result<(bool, bool)> { + gen_go_mymod(inner_content, job_dir).await?; + + let md = requirements.split_once(GO_REQ_SPLITTER); + if let Some((req, sum)) = md { + write_file(job_dir, "go.mod", &req).await?; + write_file(job_dir, "go.sum", &sum).await?; + Ok((true, true)) + } else { + write_file(job_dir, "go.mod", &requirements).await?; + Ok((true, false)) + } +} + +pub async fn install_go_dependencies( + job_id: &Uuid, + code: &str, + logs: &mut String, + job_dir: &str, + db: &sqlx::Pool, + non_dep_job: bool, + skip_go_mod: bool, + has_sum: bool, + worker_name: &str, + w_id: &str, +) -> error::Result { + if !skip_go_mod { + gen_go_mymod(code, job_dir).await?; + let child = Command::new("go") + .current_dir(job_dir) + .args(vec!["mod", "init", "mymod"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + + handle_child(job_id, db, logs, child, false, worker_name, w_id).await?; + } + + let mut new_lockfile = false; + + let hash = if !has_sum { + calculate_hash(parse_go_imports(&code)?.iter().join("\n").as_str()) + } else { + "".to_string() + }; + + let mut skip_tidy = has_sum; + + if !has_sum { + if let Some(cached) = sqlx::query_scalar!( + "SELECT lockfile FROM pip_resolution_cache WHERE hash = $1", + hash + ) + .fetch_optional(db) + .await? + { + logs.push_str(&format!("\nfound cached resolution")); + gen_go_mod(code, job_dir, &cached).await?; + skip_tidy = true; + new_lockfile = false; + } else { + new_lockfile = true; + } + } + + let mod_command = if skip_tidy { "download" } else { "tidy" }; + let child = Command::new(GO_PATH.as_str()) + .current_dir(job_dir) + .env("GOPATH", GO_CACHE_DIR) + .args(vec!["mod", mod_command]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + handle_child(job_id, db, logs, child, false, worker_name, &w_id) + .await + .map_err(|e| Error::ExecutionErr(format!("Lock file generation failed: {e:?}")))?; + + if (!new_lockfile || has_sum) && non_dep_job { + return Ok("".to_string()); + } + + let mut req_content = "".to_string(); + + let mut file = File::open(format!("{job_dir}/go.mod")).await?; + file.read_to_string(&mut req_content).await?; + req_content.push_str(GO_REQ_SPLITTER); + let sum_path = format!("{job_dir}/go.sum"); + if tokio::fs::metadata(&sum_path).await.is_ok() { + let mut file = File::open(sum_path).await?; + file.read_to_string(&mut req_content).await?; + } + + if non_dep_job { + sqlx::query!( + "INSERT INTO pip_resolution_cache (hash, lockfile, expiration) VALUES ($1, $2, now() + ('3 days')::interval) ON CONFLICT (hash) DO UPDATE SET lockfile = $2", + hash, + req_content + ).fetch_optional(db).await?; + + return Ok(String::new()); + } else { + Ok(req_content) + } +} + +async fn gen_go_mymod(code: &str, job_dir: &str) -> error::Result<()> { + let code = if code.trim_start().starts_with("package") { + code.to_string() + } else { + format!("package inner; {code}") + }; + + let mymod_dir = format!("{job_dir}/inner"); + DirBuilder::new() + .recursive(true) + .create(&mymod_dir) + .await + .expect("could not create go's mymod dir"); + + write_file(&mymod_dir, "inner_main.go", &code).await?; + + Ok(()) +} diff --git a/backend/windmill-worker/src/jobs.rs b/backend/windmill-worker/src/jobs.rs index 826d96859b..1f7d342cf5 100644 --- a/backend/windmill-worker/src/jobs.rs +++ b/backend/windmill-worker/src/jobs.rs @@ -10,11 +10,13 @@ use sqlx::{Pool, Postgres}; use tracing::instrument; use uuid::Uuid; use windmill_common::{ - error::Error, flow_status::FlowStatusModule, schedule::Schedule, METRICS_ENABLED, -}; -use windmill_queue::{ - delete_job, schedule::get_schedule_opt, JobKind, QueueTransaction, QueuedJob, CLOUD_HOSTED, + error::Error, + flow_status::FlowStatusModule, + jobs::{JobKind, QueuedJob}, + schedule::Schedule, + METRICS_ENABLED, }; +use windmill_queue::{delete_job, schedule::get_schedule_opt, QueueTransaction, CLOUD_HOSTED}; #[instrument(level = "trace", skip_all)] pub async fn add_completed_job_error( diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs index a9abf4728f..86cb82d4ef 100644 --- a/backend/windmill-worker/src/lib.rs +++ b/backend/windmill-worker/src/lib.rs @@ -1,5 +1,9 @@ +mod common; +mod global_cache; +mod go_executor; mod jobs; mod js_eval; +mod python_executor; mod worker; mod worker_flow; diff --git a/backend/windmill-worker/src/main2.rs b/backend/windmill-worker/src/main2.rs deleted file mode 100644 index fb79d6ba3e..0000000000 --- a/backend/windmill-worker/src/main2.rs +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Author: Ruben Fiszel - * Copyright: Windmill Labs, Inc 2022 - * This file and its contents are licensed under the AGPLv3 License. - * Please see the included NOTICE for copyright information and - * LICENSE-AGPL for a copy of the license. - */ - -// use std::{net::SocketAddr, time::Duration}; - -// use anyhow::Context; -// use sqlx::{postgres::PgPoolOptions, Pool, Postgres}; -// use windmill_common::{ -// error::{self, Error}, -// utils::rd_string, -// }; - -// #[tokio::main] -// async fn main() -> anyhow::Result<()> { -// // dotenv().ok(); - -// windmill_common::tracing_init::initialize_tracing(); - -// let db = async { -// let database_url = std::env::var("DATABASE_URL") -// .map_err(|_| Error::BadConfig("DATABASE_URL env var is missing".to_string()))?; - -// let max_connections = match std::env::var("DATABASE_CONNECTIONS") { -// Ok(n) => n.parse::().context("invalid DATABASE_CONNECTIONS")?, -// Err(_) => 10, -// }; - -// Ok::, error::Error>( -// PgPoolOptions::new() -// .max_connections(max_connections) -// .max_lifetime(Duration::from_secs(30 * 60)) // 30 mins -// .connect(&database_url) -// .await -// .map_err(|err| Error::ConnectingToDatabase(err.to_string()))?, -// ) -// } -// .await?; - -// let metrics_addr: Option = std::env::var("METRICS_ADDR") -// .ok() -// .map(|s| { -// s.parse::() -// .map(|b| b.then(|| SocketAddr::from(([0, 0, 0, 0], 8001)))) -// .or_else(|_| s.parse::().map(Some)) -// }) -// .transpose()? -// .flatten(); - -// let (tx, rx) = tokio::sync::broadcast::channel::<()>(3); -// let shutdown_signal = windmill_common::shutdown_signal(tx); - -// let workers_f = async { -// let instance_name = rd_string(5); - -// let ip = windmill_common::external_ip::get_ip() -// .await -// .unwrap_or_else(|e| { -// tracing::warn!(error = e.to_string(), "failed to get external IP"); -// "unretrievable IP".to_string() -// }); -// let worker_name = format!("dt-worker-{}-{}", &instance_name, rd_string(5)); -// windmill_worker::run_worker( -// &db.clone(), -// &instance_name, -// worker_name, -// 1, -// 1, -// &ip, -// rx.resubscribe(), -// ) -// .await; -// Ok(()) as anyhow::Result<()> -// }; - -// let metrics_f = async { -// match metrics_addr { -// Some(addr) => windmill_common::serve_metrics(addr, rx.resubscribe()) -// .await -// .map_err(anyhow::Error::from), -// None => Ok(()), -// } -// }; - -// futures::try_join!(shutdown_signal, workers_f, metrics_f)?; - -// Ok(()) -// } diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs new file mode 100644 index 0000000000..101183c3c5 --- /dev/null +++ b/backend/windmill-worker/src/python_executor.rs @@ -0,0 +1,541 @@ +use std::process::Stdio; + +use itertools::Itertools; +use regex::Regex; +use sqlx::{Pool, Postgres}; +use tokio::{ + fs::{metadata, DirBuilder, File}, + io::AsyncReadExt, + process::Command, +}; +use uuid::Uuid; +use windmill_common::{ + error::{self, Error}, + utils::calculate_hash, jobs::QueuedJob, +}; + +lazy_static::lazy_static! { + static ref PYTHON_PATH: String = + std::env::var("PYTHON_PATH").unwrap_or_else(|_| "/usr/local/bin/python3".to_string()); + + static ref PIP_INDEX_URL: Option = std::env::var("PIP_INDEX_URL").ok(); + static ref PIP_EXTRA_INDEX_URL: Option = std::env::var("PIP_EXTRA_INDEX_URL").ok(); + static ref PIP_TRUSTED_HOST: Option = std::env::var("PIP_TRUSTED_HOST").ok(); + static ref PIP_LOCAL_DEPENDENCIES: Option> = { + let pip_local_dependencies = std::env::var("PIP_LOCAL_DEPENDENCIES") + .ok() + .map(|x| x.split(',').map(|x| x.to_string()).collect()); + if pip_local_dependencies == Some(vec!["".to_string()]) { + None + } else { + pip_local_dependencies + } + }; + + static ref ADDITIONAL_PYTHON_PATHS: Option> = std::env::var("ADDITIONAL_PYTHON_PATHS") + .ok() + .map(|x| x.split(':').map(|x| x.to_string()).collect()); + + static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(((u|f)\.)|\.)"#).unwrap(); + +} + +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"); + +use crate::{ + common::{read_result, set_logs}, + create_args_and_out_file, get_reserved_variables, handle_child, write_file, + AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, NSJAIL_PATH, PATH_ENV, + PIP_CACHE_DIR, +}; + +pub async fn create_dependencies_dir(job_dir: &str) { + DirBuilder::new() + .recursive(true) + .create(&format!("{job_dir}/dependencies")) + .await + .expect("could not create dependencies dir"); +} + +pub async fn pip_compile( + job_id: &Uuid, + requirements: &str, + logs: &mut String, + job_dir: &str, + db: &Pool, + worker_name: &str, + w_id: &str, +) -> error::Result { + logs.push_str(&format!("\nresolving dependencies...")); + set_logs(logs, job_id, db).await; + logs.push_str(&format!("\ncontent of requirements:\n{}", requirements)); + let req_hash = calculate_hash(&requirements); + if let Some(cached) = sqlx::query_scalar!( + "SELECT lockfile FROM pip_resolution_cache WHERE hash = $1", + req_hash + ) + .fetch_optional(db) + .await? + { + logs.push_str(&format!("\nfound cached resolution")); + return Ok(cached); + } + let file = "requirements.in"; + let requirements = if let Some(pip_local_dependencies) = PIP_LOCAL_DEPENDENCIES.as_ref() { + let deps = pip_local_dependencies.clone(); + requirements + .lines() + .filter(|s| !deps.contains(&s.to_string())) + .join("\n") + } else { + requirements.to_string() + }; + write_file(job_dir, file, &requirements).await?; + + let mut args = vec!["-q", "--no-header", file, "--resolver=backtracking"]; + if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { + args.extend(["--extra-index-url", url]); + } + if let Some(url) = PIP_INDEX_URL.as_ref() { + args.extend(["--index-url", url]); + } + if let Some(host) = PIP_TRUSTED_HOST.as_ref() { + args.extend(["--trusted-host", host]); + } + let child = Command::new("pip-compile") + .current_dir(job_dir) + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + handle_child(job_id, db, logs, child, false, worker_name, &w_id) + .await + .map_err(|e| Error::ExecutionErr(format!("Lock file generation failed: {e:?}")))?; + let path_lock = format!("{job_dir}/requirements.txt"); + let mut file = File::open(path_lock).await?; + let mut req_content = "".to_string(); + file.read_to_string(&mut req_content).await?; + let lockfile = req_content + .lines() + .filter(|x| !x.trim_start().starts_with('#')) + .map(|x| x.to_string()) + .collect::>() + .join("\n"); + sqlx::query!( + "INSERT INTO pip_resolution_cache (hash, lockfile, expiration) VALUES ($1, $2, now() + ('3 days')::interval) ON CONFLICT (hash) DO UPDATE SET lockfile = $2", + req_hash, + lockfile + ).fetch_optional(db).await?; + Ok(lockfile) +} + +#[tracing::instrument(level = "trace", skip_all)] +pub async fn handle_python_job( + requirements_o: Option, + job_dir: &str, + worker_dir: &str, + worker_name: &str, + job: &QueuedJob, + logs: &mut String, + db: &sqlx::Pool, + client: &AuthedClientBackgroundTask, + inner_content: &String, + shared_mount: &str, + base_internal_url: &str, +) -> windmill_common::error::Result { + create_dependencies_dir(job_dir).await; + + let mut additional_python_paths: Vec = + ADDITIONAL_PYTHON_PATHS.to_owned().unwrap_or_else(|| vec![]); + + let requirements = match requirements_o { + Some(r) => r, + None => { + let requirements = windmill_parser_py::parse_python_imports(&inner_content)?.join("\n"); + if requirements.is_empty() { + "".to_string() + } else { + pip_compile( + &job.id, + &requirements, + logs, + job_dir, + db, + worker_name, + &job.workspace_id, + ) + .await + .map_err(|e| { + Error::ExecutionErr(format!("pip compile failed: {}", e.to_string())) + })? + } + } + }; + + if requirements.len() > 0 { + if !*DISABLE_NSJAIL { + let _ = write_file( + job_dir, + "download.config.proto", + &NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT + .replace("{WORKER_DIR}", &worker_dir) + .replace("{CACHE_DIR}", PIP_CACHE_DIR) + .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()), + ) + .await?; + } + + additional_python_paths = handle_python_reqs( + requirements + .split("\n") + .filter(|x| !x.starts_with("--")) + .collect(), + job, + logs, + db, + worker_name, + job_dir, + ) + .await?; + } + logs.push_str("\n\n--- PYTHON CODE EXECUTION ---\n"); + + set_logs(logs, &job.id, db).await; + + let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content); + + let script_path_splitted = &job.script_path().split("/"); + let dirs_full = script_path_splitted + .clone() + .take(script_path_splitted.clone().count() - 1) + .join("/") + .replace("-", "_"); + let dirs = if dirs_full.len() > 0 { + dirs_full + } else { + "tmp".to_string() + }; + let last = script_path_splitted + .clone() + .last() + .unwrap() + .replace("-", "_") + .replace(" ", "_") + .to_lowercase(); + let module_dir = format!("{}/{}", job_dir, dirs); + tokio::fs::create_dir_all(format!("{module_dir}/")).await?; + let _ = write_file(&module_dir, &format!("{last}.py"), inner_content).await?; + if relative_imports { + let _ = write_file(&job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?; + } + + let sig = windmill_parser_py::parse_python_signature(inner_content)?; + let transforms = sig + .args + .iter() + .map(|x| match x.typ { + windmill_parser::Typ::Bytes => { + format!( + "if \"{}\" in kwargs and kwargs[\"{}\"] is not None:\n \ + kwargs[\"{}\"] = base64.b64decode(kwargs[\"{}\"])\n", + x.name, x.name, x.name, x.name + ) + } + windmill_parser::Typ::Datetime => { + format!( + "if \"{}\" in kwargs and kwargs[\"{}\"] is not None:\n \ + kwargs[\"{}\"] = datetime.strptime(kwargs[\"{}\"], \ + '%Y-%m-%dT%H:%M')\n", + x.name, x.name, x.name, x.name + ) + } + _ => "".to_string(), + }) + .collect::>() + .join(""); + let client = client.get_authed().await; + create_args_and_out_file(&client, job, job_dir).await?; + + let import_loader = if relative_imports { + "import loader" + } else { + "" + }; + let import_base64 = if sig + .args + .iter() + .any(|x| x.typ == windmill_parser::Typ::Bytes) + { + "import base64" + } else { + "" + }; + let import_datetime = if sig + .args + .iter() + .any(|x| x.typ == windmill_parser::Typ::Datetime) + { + "from datetime import datetime" + } else { + "" + }; + let spread = if sig.star_kwargs { + "args = kwargs".to_string() + } else { + sig.args + .into_iter() + .map(|x| format!("args[\"{}\"] = kwargs.get(\"{}\")", x.name, x.name)) + .join("\n") + }; + + let module_dir_dot = dirs.replace("/", ".").replace("-", "_"); + let wrapper_content: String = format!( + r#" +import json +{import_loader} +{import_base64} +{import_datetime} +import traceback +import sys +from {module_dir_dot} import {last} as inner_script + + +with open("args.json") as f: + kwargs = json.load(f, strict=False) +args = {{}} +{spread} +{transforms} +for k, v in list(args.items()): + if v == '': + del args[k] + +try: + res = inner_script.main(**args) + typ = type(res) + if typ.__name__ == 'DataFrame': + if typ.__module__ == 'pandas.core.frame': + res = res.values.tolist() + elif typ.__module__ == 'polars.dataframe.frame': + res = res.rows() + res_json = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') + with open("result.json", 'w') as f: + f.write(res_json) +except Exception as e: + exc_type, exc_value, exc_traceback = sys.exc_info() + tb = traceback.format_tb(exc_traceback) + with open("result.json", 'w') as f: + err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '') + f.write(err_json) + sys.exit(1) +"#, + ); + write_file(job_dir, "wrapper.py", &wrapper_content).await?; + + let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?; + let additional_python_paths_folders = additional_python_paths.iter().join(":"); + if !*DISABLE_NSJAIL { + let shared_deps = additional_python_paths + .into_iter() + .map(|pp| { + format!( + r#" +mount {{ + src: "{pp}" + dst: "{pp}" + is_bind: true + rw: false +}} + "# + ) + }) + .join("\n"); + let _ = write_file( + job_dir, + "run.config.proto", + &NSJAIL_CONFIG_RUN_PYTHON3_CONTENT + .replace("{JOB_DIR}", job_dir) + .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()) + .replace("{SHARED_MOUNT}", shared_mount) + .replace("{SHARED_DEPENDENCIES}", shared_deps.as_str()) + .replace("{MAIN}", format!("{dirs}/{last}").as_str()) + .replace( + "{ADDITIONAL_PYTHON_PATHS}", + additional_python_paths_folders.as_str(), + ), + ) + .await?; + } else { + reserved_variables.insert("PYTHONPATH".to_string(), additional_python_paths_folders); + } + + tracing::info!( + worker_name = %worker_name, + job_id = %job.id, + workspace_id = %job.workspace_id, + "started python code execution {}", + job.id + ); + let child = if !*DISABLE_NSJAIL { + Command::new(NSJAIL_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + // inject PYTHONPATH here - for some reason I had to do it in nsjail conf + .envs(reserved_variables) + .env("PATH", PATH_ENV.as_str()) + .env("BASE_INTERNAL_URL", base_internal_url) + .args(vec![ + "--config", + "run.config.proto", + "--", + PYTHON_PATH.as_str(), + "-u", + "-m", + "wrapper", + ]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + } else { + Command::new(PYTHON_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + .envs(reserved_variables) + .env("PATH", PATH_ENV.as_str()) + .env("BASE_INTERNAL_URL", base_internal_url) + .args(vec!["-u", "-m", "wrapper"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + }; + + handle_child( + &job.id, + db, + logs, + child, + !*DISABLE_NSJAIL, + worker_name, + &job.workspace_id, + ) + .await?; + read_result(job_dir).await +} + +async fn handle_python_reqs( + requirements: Vec<&str>, + job: &QueuedJob, + logs: &mut String, + db: &sqlx::Pool, + worker_name: &str, + job_dir: &str, +) -> error::Result> { + let mut req_paths: Vec = vec![]; + let mut vars = vec![("PATH", PATH_ENV.as_str())]; + if !*DISABLE_NSJAIL { + if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { + vars.push(("EXTRA_INDEX_URL", url)); + } + if let Some(url) = PIP_INDEX_URL.as_ref() { + vars.push(("INDEX_URL", url)); + } + if let Some(host) = PIP_TRUSTED_HOST.as_ref() { + vars.push(("TRUSTED_HOST", host)); + } + }; + + for req in requirements { + // todo: handle many reqs + let venv_p = format!("{PIP_CACHE_DIR}/{req}"); + if metadata(&venv_p).await.is_ok() { + req_paths.push(venv_p); + continue; + } + + logs.push_str("\n--- PIP INSTALL ---\n"); + logs.push_str(&format!("\n{req} is being installed for the first time.\n It will be cached for all ulterior uses.")); + + tracing::info!( + worker_name = %worker_name, + job_id = %job.id, + workspace_id = %job.workspace_id, + "started setup python dependencies" + ); + + let child = if !*DISABLE_NSJAIL { + tracing::info!( + worker_name = %worker_name, + job_id = %job.id, + workspace_id = %job.workspace_id, + "starting nsjail" + ); + let mut vars = vars.clone(); + let req = req.to_string(); + vars.push(("REQ", &req)); + vars.push(("TARGET", &venv_p)); + Command::new(NSJAIL_PATH.as_str()) + .current_dir(job_dir) + .env_clear() + .envs(vars) + .args(vec!["--config", "download.config.proto"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + } else { + let mut args = vec![ + "-m", + "pip", + "install", + &req, + "-I", + "--no-deps", + "--no-color", + "--isolated", + "--no-warn-conflicts", + "--disable-pip-version-check", + "-t", + venv_p.as_str(), + ]; + if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { + args.extend(["--extra-index-url", url]); + } + if let Some(url) = PIP_INDEX_URL.as_ref() { + args.extend(["--index-url", url]); + } + if let Some(host) = PIP_TRUSTED_HOST.as_ref() { + args.extend(["--trusted-host", &host]); + } + Command::new(PYTHON_PATH.as_str()) + .env_clear() + .env("PATH", PATH_ENV.as_str()) + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + }; + + let child = handle_child( + &job.id, + db, + logs, + child, + false, + worker_name, + &job.workspace_id, + ) + .await; + tracing::info!( + worker_name = %worker_name, + job_id = %job.id, + workspace_id = %job.workspace_id, + is_ok = child.is_ok(), + "finished setting up python dependencies {}", + job.id + ); + child?; + + req_paths.push(venv_p); + } + Ok(req_paths) +} diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 3f4fe6cd93..2e7c5106f3 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -9,12 +9,9 @@ use anyhow::Result; use const_format::concatcp; use itertools::Itertools; -use lazy_static::lazy_static; use once_cell::sync::OnceCell; -use regex::Regex; use sqlx::{Pool, Postgres}; use windmill_api_client::Client; -use windmill_parser_go::parse_go_imports; use std::{ borrow::Borrow, collections::HashMap, io, os::unix::process::ExitStatusExt, panic, process::Stdio, time::{Duration}, sync::atomic::Ordering, @@ -27,16 +24,16 @@ use windmill_common::{ error::{self, to_anyhow, Error}, flows::{FlowModuleValue, FlowValue}, scripts::{ScriptHash, ScriptLang}, - utils::{rd_string, calculate_hash}, - variables, BASE_URL, users::SUPERADMIN_SECRET_EMAIL, IS_READY, METRICS_ENABLED, + utils::{rd_string}, + variables, BASE_URL, users::SUPERADMIN_SECRET_EMAIL, IS_READY, METRICS_ENABLED, jobs::{JobKind, QueuedJob}, }; -use windmill_queue::{canceled_job_to_result, get_queued_job, pull, JobKind, QueuedJob, CLOUD_HOSTED}; +use windmill_queue::{canceled_job_to_result, get_queued_job, pull, CLOUD_HOSTED}; use serde_json::{json, Value}; use tokio::{ fs::{metadata, symlink, DirBuilder, File}, - io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}, + io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, process::{Child, Command}, sync::{ mpsc::{self, Sender}, watch, broadcast, RwLock @@ -50,244 +47,20 @@ use futures::{ }; use async_recursion::async_recursion; +#[cfg(feature = "enterprise")] +use rand::Rng; + +#[cfg(feature = "enterprise")] +use crate::global_cache::{copy_cache_from_bucket_as_tar, copy_cache_from_bucket, copy_cache_to_bucket, copy_cache_to_bucket_as_tar}; use crate::{ jobs::{add_completed_job, add_completed_job_error}, worker_flow::{ handle_flow, update_flow_status_after_job_completion, update_flow_status_in_progress, - }, + }, python_executor::{create_dependencies_dir, pip_compile, handle_python_job}, common::{read_result, set_logs}, global_cache::{move_tmp_cache_to_cache}, go_executor::{handle_go_job, install_go_dependencies}, }; -#[cfg(feature = "enterprise")] -use rand::Rng; -#[cfg(feature = "enterprise")] -const TAR_CACHE_FILENAME: &str = "entirecache.tar"; - -#[cfg(feature = "enterprise")] -async fn copy_cache_from_bucket(bucket: &str, tx: Option>) -> Option::> { - tracing::info!("Copying cache from bucket in the background {bucket}"); - let bucket = bucket.to_string(); - let tx_is_some = tx.is_some(); - let f = async move { - let elapsed = Instant::now(); - - match Command::new("rclone") - .arg("copy") - .arg(format!(":s3,env_auth=true:{bucket}")) - .arg(if tx_is_some { ROOT_TMP_CACHE_DIR } else { ROOT_CACHE_DIR }) - .arg("--size-only") - .arg("--fast-list") - .arg("--exclude") - .arg(format!("\"{TAR_CACHE_FILENAME},/deno/gen/file/**\"")) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - h.wait().await.unwrap(); - } - Err(e) => tracing::warn!("Failed to run periodic job pull. Error: {:?}", e), - } - tracing::info!( - "Finished copying cache from bucket {bucket}, took {:?}s", - elapsed.elapsed().as_secs() - ); - - for x in - if !tx_is_some - { [PIP_CACHE_DIR, DENO_CACHE_DIR, GO_CACHE_DIR] } - else { [PIP_TMP_CACHE_DIR, DENO_TMP_CACHE_DIR, GO_TMP_CACHE_DIR] } { - DirBuilder::new() - .recursive(true) - .create(x) - .await - .expect("could not create initial worker dir"); - } - - if let Some(tx) = tx { - tx.send(()).await.expect("can send copy cache signal"); - } - }; - if tx_is_some { - return Some(tokio::spawn(f)); - } else { - f.await; - return None; - } -} - -#[cfg(feature = "enterprise")] -async fn copy_cache_to_bucket(bucket: &str) { - tracing::info!("Copying cache to bucket {bucket}"); - let elapsed = Instant::now(); - match Command::new("rclone") - .arg("copy") - .arg(ROOT_CACHE_DIR) - .arg(format!(":s3,env_auth=true:{bucket}")) - .arg("--size-only") - .arg("--fast-list") - .arg("--exclude") - .arg(format!("\"{TAR_CACHE_FILENAME},/deno/gen/file/**\"")) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - h.wait().await.unwrap(); - } - Err(e) => tracing::info!("Failed to run periodic job push. Error: {:?}", e), - } - tracing::info!( - "Finished copying cache to bucket {bucket}, took: {:?}s", - elapsed.elapsed().as_secs() - ); -} - -#[cfg(feature = "enterprise")] -async fn copy_cache_to_bucket_as_tar(bucket: &str) { - tracing::info!("Copying cache to bucket {bucket} as tar"); - let elapsed = Instant::now(); - - match Command::new("tar") - .current_dir(ROOT_CACHE_DIR) - .arg("-c") - .arg("-f") - .arg(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")) - .args(&["pip", "go", "deno"]) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - if !h.wait().await.unwrap().success() { - tracing::info!("Failed to tar cache"); - return; - } - } - Err(e) => { - tracing::info!("Failed tar cache. Error: {e:?}"); - return; - } - } - - let tar_metadata = tokio::fs::metadata(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")) - .await; - if tar_metadata.is_err() || tar_metadata.as_ref().unwrap().len() == 0 { - tracing::info!("Failed to tar cache"); - return; - } - - match Command::new("rclone") - .current_dir(ROOT_CACHE_DIR) - .arg("copyto") - .arg(TAR_CACHE_FILENAME) - .arg(format!(":s3,env_auth=true:{bucket}/{TAR_CACHE_FILENAME}")) - .arg("-vv") - .arg("--size-only") - .arg("--fast-list") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - h.wait().await.unwrap(); - } - Err(e) => tracing::info!("Failed to copy tar cache to bucket. Error: {:?}", e), - } - - if let Err(e) = tokio::fs::remove_file(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")).await { - tracing::info!("Failed to remove tar cache. Error: {:?}", e); - }; - - tracing::info!( - "Finished copying cache to bucket {bucket} as tar, took: {:?}s. Size of new tar: {}", - elapsed.elapsed().as_secs(), - tar_metadata.unwrap().len() - ); -} - -#[cfg(feature = "enterprise")] -async fn copy_cache_from_bucket_as_tar(bucket: &str) -> bool { - tracing::info!("Copying cache from bucket {bucket} as tar"); - let elapsed = Instant::now(); - - match Command::new("rclone") - .arg("copyto") - .arg(format!(":s3,env_auth=true:{bucket}/{TAR_CACHE_FILENAME}")) - .arg(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")) - .arg("-vv") - .arg("--size-only") - .arg("--fast-list") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - if !h.wait().await.unwrap().success() { - tracing::info!("Failed to download tar cache, continuing nonetheless"); - return false; - } - } - Err(e) => { - tracing::info!("Failed to download tar cache, continuing nonetheless. Error: {e:?}"); - return false; - } - } - - match Command::new("tar") - .current_dir(ROOT_TMP_CACHE_DIR) - .arg("-xpvf") - .arg(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .spawn() - { - Ok(mut h) => { - if !h.wait().await.unwrap().success() { - tracing::info!("Failed to untar cache, continuing nonetheless"); - return false; - } - } - Err(e) => { - tracing::warn!("Failed to untar cache, continuing nonetheless. Error: {e:?}"); - return false; - } - } - - if let Err(e) = tokio::fs::remove_file(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")).await { - tracing::info!("Failed to remove tar cache. Error: {:?}", e); - }; - - let r = move_tmp_cache_to_cache().await.is_ok(); - - tracing::info!( - "Finished copying cache from bucket {bucket} as tar, took: {:?}s. copy success: {r}", - elapsed.elapsed().as_secs() - ); - return r; -} - -// async fn check_if_bucket_syncable(bucket: &str) -> bool { -// match Command::new("rclone") -// .arg("lsf") -// .arg(format!(":s3,env_auth=true:{bucket}/NOSYNC")) - -// .arg("-vv") -// .arg("--fast-list") -// .stdin(Stdio::null()) -// .stdout(Stdio::null()) -// .output() -// .await; -// return true; -// } - -async fn move_tmp_cache_to_cache() -> Result<()> { - tokio::fs::remove_dir_all(ROOT_CACHE_DIR).await?; - tokio::fs::rename(ROOT_TMP_CACHE_DIR, ROOT_CACHE_DIR).await?; - tracing::info!("Finished moving tmp cache to cache"); - Ok(()) -} pub async fn create_token_for_owner_in_bg(db: &Pool, job: &QueuedJob) -> Arc> { let rw_lock = Arc::new(RwLock::new(String::new())); @@ -344,26 +117,23 @@ pub async fn create_token_for_owner( } const TMP_DIR: &str = "/tmp/windmill"; -const ROOT_CACHE_DIR: &str = "/tmp/windmill/cache/"; -const ROOT_TMP_CACHE_DIR: &str = "/tmp/windmill/tmpcache/"; -const PIP_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "pip"); -const DENO_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "deno"); -const GO_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "go"); -const PIP_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "pip"); -const DENO_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "deno"); -const GO_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "go"); +pub const ROOT_CACHE_DIR: &str = "/tmp/windmill/cache/"; +pub const ROOT_TMP_CACHE_DIR: &str = "/tmp/windmill/tmpcache/"; +pub const PIP_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "pip"); +pub const DENO_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "deno"); +pub const GO_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "go"); +pub const PIP_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "pip"); +pub const DENO_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "deno"); +pub const GO_TMP_CACHE_DIR: &str = concatcp!(ROOT_TMP_CACHE_DIR, "go"); + const NUM_SECS_PING: u64 = 5; + +#[cfg(feature = "enterprise")] const NUM_SECS_SYNC: u64 = 60 * 10; const INCLUDE_DEPS_PY_SH_CONTENT: &str = include_str!("../nsjail/download_deps.py.sh"); -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 NSJAIL_CONFIG_RUN_GO_CONTENT: &str = include_str!("../nsjail/run.go.config.proto"); const NSJAIL_CONFIG_RUN_BASH_CONTENT: &str = include_str!("../nsjail/run.bash.config.proto"); -const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py"); - -const GO_REQ_SPLITTER: &str = "//go.sum\n"; #[derive(Clone)] pub struct Metrics { @@ -375,23 +145,21 @@ pub const DEFAULT_TIMEOUT: u16 = 300; pub const DEFAULT_SLEEP_QUEUE: u64 = 50; lazy_static::lazy_static! { - - static ref SLEEP_QUEUE: u64 = std::env::var("SLEEP_QUEUE") .ok() .and_then(|x| x.parse::().ok()) .unwrap_or(DEFAULT_SLEEP_QUEUE); - static ref DISABLE_NUSER: bool = std::env::var("DISABLE_NUSER") + pub static ref DISABLE_NUSER: bool = std::env::var("DISABLE_NUSER") .ok() .and_then(|x| x.parse::().ok()) .unwrap_or(false); - static ref DISABLE_NSJAIL: bool = std::env::var("DISABLE_NSJAIL") - .ok() - .and_then(|x| x.parse::().ok()) - .unwrap_or(true); + pub static ref DISABLE_NSJAIL: bool = std::env::var("DISABLE_NSJAIL") + .ok() + .and_then(|x| x.parse::().ok()) + .unwrap_or(true); pub static ref KEEP_JOB_DIR: bool = std::env::var("KEEP_JOB_DIR") .ok() @@ -403,18 +171,13 @@ lazy_static::lazy_static! { .map(|e| Some(e)) .unwrap_or(None); - static ref DENO_PATH: String = std::env::var("DENO_PATH").unwrap_or_else(|_| "/usr/bin/deno".to_string()); - static ref GO_PATH: String = std::env::var("GO_PATH").unwrap_or_else(|_| "/usr/bin/go".to_string()); - static ref PYTHON_PATH: String = - std::env::var("PYTHON_PATH").unwrap_or_else(|_| "/usr/local/bin/python3".to_string()); - static ref NSJAIL_PATH: String = std::env::var("NSJAIL_PATH").unwrap_or_else(|_| "nsjail".to_string()); - static ref PATH_ENV: String = std::env::var("PATH").unwrap_or_else(|_| String::new()); - static ref HOME_ENV: String = std::env::var("HOME").unwrap_or_else(|_| String::new()); - static ref PIP_INDEX_URL: Option = std::env::var("PIP_INDEX_URL").ok(); - static ref GOPRIVATE: Option = std::env::var("GOPRIVATE").ok(); - static ref NETRC: Option = std::env::var("NETRC").ok(); - static ref PIP_EXTRA_INDEX_URL: Option = std::env::var("PIP_EXTRA_INDEX_URL").ok(); - static ref PIP_TRUSTED_HOST: Option = std::env::var("PIP_TRUSTED_HOST").ok(); + pub static ref DENO_PATH: String = std::env::var("DENO_PATH").unwrap_or_else(|_| "/usr/bin/deno".to_string()); + pub static ref NSJAIL_PATH: String = std::env::var("NSJAIL_PATH").unwrap_or_else(|_| "nsjail".to_string()); + pub static ref PATH_ENV: String = std::env::var("PATH").unwrap_or_else(|_| String::new()); + pub static ref HOME_ENV: String = std::env::var("HOME").unwrap_or_else(|_| String::new()); + pub static ref GOPRIVATE: Option = std::env::var("GOPRIVATE").ok(); + pub static ref NETRC: Option = std::env::var("NETRC").ok(); + static ref DENO_AUTH_TOKENS: String = std::env::var("DENO_AUTH_TOKENS") .ok() .map(|x| format!(";{x}")) @@ -426,17 +189,8 @@ lazy_static::lazy_static! { .ok() .map(|x| x.split(' ').map(|x| x.to_string()).collect()); - static ref PIP_LOCAL_DEPENDENCIES: Option> = { - let pip_local_dependencies = std::env::var("PIP_LOCAL_DEPENDENCIES") - .ok() - .map(|x| x.split(',').map(|x| x.to_string()).collect()); - if pip_local_dependencies == Some(vec!["".to_string()]) { - None - } else { - pip_local_dependencies - } - - }; + + static ref WHITELIST_WORKSPACES: Option> = std::env::var("WHITELIST_WORKSPACES") .ok() .map(|x| x.split(',').map(|x| x.to_string()).collect()); @@ -449,26 +203,12 @@ lazy_static::lazy_static! { .and_then(|x| x.parse::().ok()) .unwrap_or(100); - static ref ADDITIONAL_PYTHON_PATHS: Option> = std::env::var("ADDITIONAL_PYTHON_PATHS") - .ok() - .map(|x| x.split(':').map(|x| x.to_string()).collect()); - - static ref WORKER_STARTED: prometheus::IntGauge = prometheus::register_int_gauge!( "worker_started", "Total number of workers started." ) .unwrap(); - static ref QUEUE_ZOMBIE_RESTART_COUNT: prometheus::IntCounter = prometheus::register_int_counter!( - "queue_zombie_restart_count", - "Total number of jobs restarted due to ping timeout." - ) - .unwrap(); - static ref QUEUE_ZOMBIE_DELETE_COUNT: prometheus::IntCounter = prometheus::register_int_counter!( - "queue_zombie_delete_count", - "Total number of jobs deleted due to their ping timing out in an unrecoverable state." - ) - .unwrap(); + static ref WORKER_UPTIME_OPTS: prometheus::Opts = prometheus::opts!( "worker_uptime", "Total number of seconds since the worker has started" @@ -481,22 +221,11 @@ lazy_static::lazy_static! { static ref TIMEOUT_DURATION: Duration = Duration::from_secs(*TIMEOUT as u64); - static ref ZOMBIE_JOB_TIMEOUT: String = std::env::var("ZOMBIE_JOB_TIMEOUT") - .ok() - .and_then(|x| x.parse::().ok()) - .unwrap_or_else(|| "30".to_string()); - - - pub static ref RESTART_ZOMBIE_JOBS: bool = std::env::var("RESTART_ZOMBIE_JOBS") - .ok() - .and_then(|x| x.parse::().ok()) - .unwrap_or(true); - - static ref SESSION_TOKEN_EXPIRY: i32 = (*TIMEOUT as i32) * 2; + pub static ref SESSION_TOKEN_EXPIRY: i32 = (*TIMEOUT as i32) * 2; } //only matter if CLOUD_HOSTED -const MAX_RESULT_SIZE: usize = 1024 * 1024 * 2; // 2MB +pub const MAX_RESULT_SIZE: usize = 1024 * 1024 * 2; // 2MB pub struct AuthedClientBackgroundTask { pub base_internal_url: String, @@ -905,7 +634,7 @@ pub async fn run_worker( } } -async fn handle_job_error( +pub async fn handle_job_error( db: &Pool, client: &AuthedClient, job: QueuedJob, @@ -1192,7 +921,7 @@ async fn handle_queued_job( } #[tracing::instrument(level = "trace", skip_all)] -async fn write_file(dir: &str, path: &str, content: &str) -> error::Result { +pub async fn write_file(dir: &str, path: &str, content: &str) -> error::Result { let path = format!("{}/{}", dir, path); let mut file = File::create(&path).await?; file.write_all(content.as_bytes()).await?; @@ -1389,206 +1118,6 @@ mount {{ result } -async fn gen_go_mod(inner_content: &str, job_dir: &str, requirements: &str) -> error::Result<(bool, bool)> { - gen_go_mymod(inner_content, job_dir).await?; - - let md = requirements - .split_once(GO_REQ_SPLITTER); - if let Some((req, sum)) = md { - write_file(job_dir, "go.mod", &req).await?; - write_file(job_dir, "go.sum", &sum).await?; - Ok((true, true)) - } else { - write_file(job_dir, "go.mod", &requirements).await?; - Ok((true, false)) - } -} -#[tracing::instrument(level = "trace", skip_all)] -async fn handle_go_job( - logs: &mut String, - job: &QueuedJob, - db: &sqlx::Pool, - client: &AuthedClientBackgroundTask, - inner_content: &str, - job_dir: &str, - requirements_o: Option, - shared_mount: &str, - base_internal_url: &str, - worker_name: &str, -) -> Result { - //go does not like executing modules at temp root - let job_dir = &format!("{job_dir}/go"); - let (skip_go_mod, skip_tidy) = if let Some(requirements) = requirements_o { - gen_go_mod(inner_content, job_dir, &requirements).await? - } else { - (false, false) - }; - - logs.push_str("\n\n--- GO DEPENDENCIES SETUP ---\n"); - set_logs(logs, &job.id, db).await; - - install_go_dependencies( - &job.id, - inner_content, - logs, - job_dir, - db, - true, - skip_go_mod, - skip_tidy, - worker_name, - &job.workspace_id, - ) - .await?; - - logs.push_str("\n\n--- GO CODE EXECUTION ---\n"); - set_logs(logs, &job.id, db).await; - let client = &client.get_authed().await; - create_args_and_out_file(client, job, job_dir).await?; - { - let sig = windmill_parser_go::parse_go_sig(&inner_content)?; - drop(inner_content); - - const WRAPPER_CONTENT: &str = r#"package main - -import ( - "encoding/json" - "os" - "fmt" - "mymod/inner" -) - -func main() {{ - - dat, err := os.ReadFile("args.json") - if err != nil {{ - fmt.Println(err) - os.Exit(1) - }} - - var req inner.Req - - if err := json.Unmarshal(dat, &req); err != nil {{ - fmt.Println(err) - os.Exit(1) - }} - - res, err := inner.Run(req) - if err != nil {{ - fmt.Println(err) - os.Exit(1) - }} - res_json, err := json.Marshal(res) - if err != nil {{ - fmt.Println(err) - os.Exit(1) - }} - f, err := os.OpenFile("result.json", os.O_APPEND|os.O_WRONLY, os.ModeAppend) - if err != nil {{ - fmt.Println(err) - os.Exit(1) - }} - _, err = f.WriteString(string(res_json)) - if err != nil {{ - fmt.Println(err) - os.Exit(1) - }} -}}"#; - - write_file(job_dir, "main.go", WRAPPER_CONTENT).await?; - - { - let spread = &sig - .args - .clone() - .into_iter() - .map(|x| format!("req.{}", capitalize(&x.name))) - .join(", "); - let req_body = &sig - .args - .into_iter() - .map(|x| { - format!( - "{} {} `json:\"{}\"`", - capitalize(&x.name), - windmill_parser_go::otyp_to_string(x.otyp), - x.name - ) - }) - .join("\n"); - let runner_content: String = format!( - r#"package inner -type Req struct {{ - {req_body} -}} - -func Run(req Req) (interface{{}}, error){{ - return main({spread}) -}} - -"#, - ); - write_file(&format!("{job_dir}/inner"), "runner.go", &runner_content).await?; - } - } - let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?; - reserved_variables.insert("RUST_LOG".to_string(), "info".to_string()); - - let child = if !*DISABLE_NSJAIL { - let _ = write_file( - job_dir, - "run.config.proto", - &NSJAIL_CONFIG_RUN_GO_CONTENT - .replace("{JOB_DIR}", job_dir) - .replace("{CACHE_DIR}", GO_CACHE_DIR) - .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()) - .replace("{SHARED_MOUNT}", shared_mount), - ) - .await?; - let build_go = Command::new(GO_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - .env("PATH", PATH_ENV.as_str()) - .env("BASE_INTERNAL_URL", base_internal_url) - .env("GOPATH", GO_CACHE_DIR) - .env("HOME", HOME_ENV.as_str()) - .args(vec!["build", "main.go"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; - handle_child(&job.id, db, logs, build_go, false, worker_name, &job.workspace_id).await?; - - Command::new(NSJAIL_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - .envs(reserved_variables) - .env("PATH", PATH_ENV.as_str()) - .env("BASE_INTERNAL_URL", base_internal_url) - .args(vec!["--config", "run.config.proto", "--", "/tmp/go/main"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - } else { - if let Some(ref netrc) = *NETRC { - write_file(&HOME_ENV, ".netrc", netrc).await?; - } - Command::new(GO_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - .envs(reserved_variables) - .env("PATH", PATH_ENV.as_str()) - .env("BASE_INTERNAL_URL", base_internal_url) - .env("GOPATH", GO_CACHE_DIR) - .env("GOPRIVATE", GOPRIVATE.as_ref().unwrap_or(&String::new())) - .env("HOME", HOME_ENV.as_str()) - .args(vec!["run", "main.go"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - }; - handle_child(&job.id, db, logs, child, !*DISABLE_NSJAIL, worker_name, &job.workspace_id).await?; - read_result(job_dir).await -} #[tracing::instrument(level = "trace", skip_all)] async fn handle_bash_job( @@ -1673,14 +1202,6 @@ async fn handle_bash_job( .unwrap_or_else(String::new))) } -fn capitalize(s: &str) -> String { - let mut c = s.chars(); - match c.next() { - None => String::new(), - Some(f) => f.to_uppercase().collect::() + c.as_str(), - } -} - fn get_common_deno_proc_envs(token: &str, base_internal_url: &str) -> HashMap { let hostname_base = BASE_URL.split("://").last().unwrap_or("localhost"); let hostname_internal = base_internal_url.split("://").last().unwrap_or("localhost"); @@ -1812,7 +1333,7 @@ async fn handle_deno_job( let child = async { let script_path = format!("{job_dir}/wrapper.ts"); let import_map_path = format!("{job_dir}/import_map.json"); - let mut args = Vec::new(); + let mut args = Vec::with_capacity(12); args.push("run"); args.push("--no-check"); args.push("--import-map"); @@ -1852,7 +1373,7 @@ async fn handle_deno_job( } #[tracing::instrument(level = "trace", skip_all)] -async fn create_args_and_out_file( +pub async fn create_args_and_out_file( client: &AuthedClient, job: &QueuedJob, job_dir: &str, @@ -1868,290 +1389,8 @@ async fn create_args_and_out_file( Ok(()) } -lazy_static! { - static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(((u|f)\.)|\.)"#).unwrap(); -} - -#[tracing::instrument(level = "trace", skip_all)] -async fn handle_python_job( - requirements_o: Option, - job_dir: &str, - worker_dir: &str, - worker_name: &str, - job: &QueuedJob, - logs: &mut String, - db: &sqlx::Pool, - client: &AuthedClientBackgroundTask, - inner_content: &String, - shared_mount: &str, - base_internal_url: &str -) -> error::Result { - create_dependencies_dir(job_dir).await; - - let mut additional_python_paths: Vec = - ADDITIONAL_PYTHON_PATHS.to_owned().unwrap_or_else(|| vec![]); - - let requirements = match requirements_o { - Some(r) => r, - None => { - let requirements = windmill_parser_py::parse_python_imports(&inner_content)?.join("\n"); - if requirements.is_empty() { - "".to_string() - } else { - pip_compile(&job.id, &requirements, logs, job_dir, db, worker_name, &job.workspace_id) - .await - .map_err(|e| { - Error::ExecutionErr(format!("pip compile failed: {}", e.to_string())) - })? - } - } - }; - - if requirements.len() > 0 { - if !*DISABLE_NSJAIL { - let _ = write_file( - job_dir, - "download.config.proto", - &NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT - .replace("{WORKER_DIR}", &worker_dir) - .replace("{CACHE_DIR}", PIP_CACHE_DIR) - .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()), - ) - .await?; - } - - additional_python_paths = handle_python_reqs( - requirements - .split("\n") - .filter(|x| !x.starts_with("--")) - .collect(), - job, - logs, - db, - worker_name, - job_dir, - ) - .await?; - } - logs.push_str("\n\n--- PYTHON CODE EXECUTION ---\n"); - - set_logs(logs, &job.id, db).await; - - let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content); - - let script_path_splitted = &job.script_path().split("/"); - let dirs_full = script_path_splitted.clone().take(script_path_splitted.clone().count() - 1).join("/").replace("-", "_"); - let dirs = if dirs_full.len() > 0 { dirs_full } else { "tmp".to_string() }; - let last = script_path_splitted.clone().last().unwrap().replace("-", "_").replace(" ", "_").to_lowercase(); - let module_dir = format!("{}/{}", job_dir, dirs ); - tokio::fs::create_dir_all(format!("{module_dir}/")).await?; - let _ = write_file(&module_dir, &format!("{last}.py"), inner_content).await?; - if relative_imports { - let _ = write_file(&job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?; - } - - let sig = windmill_parser_py::parse_python_signature(inner_content)?; - let transforms = sig - .args - .iter() - .map(|x| match x.typ { - windmill_parser::Typ::Bytes => { - format!( - "if \"{}\" in kwargs and kwargs[\"{}\"] is not None:\n \ - kwargs[\"{}\"] = base64.b64decode(kwargs[\"{}\"])\n", - x.name, x.name, x.name, x.name - ) - } - windmill_parser::Typ::Datetime => { - format!( - "if \"{}\" in kwargs and kwargs[\"{}\"] is not None:\n \ - kwargs[\"{}\"] = datetime.strptime(kwargs[\"{}\"], \ - '%Y-%m-%dT%H:%M')\n", - x.name, x.name, x.name, x.name - ) - } - _ => "".to_string(), - }) - .collect::>() - .join(""); - let client = client.get_authed().await; - create_args_and_out_file(&client, job, job_dir).await?; - - let import_loader = if relative_imports { - "import loader" - } else { - "" - }; - let import_base64 = if sig - .args - .iter() - .any(|x| x.typ == windmill_parser::Typ::Bytes) - { - "import base64" - } else { - "" - }; - let import_datetime = if sig - .args - .iter() - .any(|x| x.typ == windmill_parser::Typ::Datetime) - { - "from datetime import datetime" - } else { - "" - }; - let spread = if sig.star_kwargs { - "args = kwargs".to_string() - } else { - sig.args - .into_iter() - .map(|x| format!("args[\"{}\"] = kwargs.get(\"{}\")", x.name, x.name)) - .join("\n") - }; - - let module_dir_dot = dirs.replace("/", ".").replace("-", "_"); - let wrapper_content: String = format!( - r#" -import json -{import_loader} -{import_base64} -{import_datetime} -import traceback -import sys -from {module_dir_dot} import {last} as inner_script -with open("args.json") as f: - kwargs = json.load(f, strict=False) -args = {{}} -{spread} -{transforms} -for k, v in list(args.items()): - if v == '': - del args[k] - -try: - res = inner_script.main(**args) - typ = type(res) - if typ.__name__ == 'DataFrame': - if typ.__module__ == 'pandas.core.frame': - res = res.values.tolist() - elif typ.__module__ == 'polars.dataframe.frame': - res = res.rows() - res_json = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') - with open("result.json", 'w') as f: - f.write(res_json) -except Exception as e: - exc_type, exc_value, exc_traceback = sys.exc_info() - tb = traceback.format_tb(exc_traceback) - with open("result.json", 'w') as f: - err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '') - f.write(err_json) - sys.exit(1) -"#, - ); - write_file(job_dir, "wrapper.py", &wrapper_content).await?; - - let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?; - let additional_python_paths_folders = additional_python_paths.iter().join(":"); - if !*DISABLE_NSJAIL { - let shared_deps = additional_python_paths - .into_iter() - .map(|pp| { - format!( - r#" -mount {{ - src: "{pp}" - dst: "{pp}" - is_bind: true - rw: false -}} - "# - ) - }) - .join("\n"); - let _ = write_file( - job_dir, - "run.config.proto", - &NSJAIL_CONFIG_RUN_PYTHON3_CONTENT - .replace("{JOB_DIR}", job_dir) - .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()) - .replace("{SHARED_MOUNT}", shared_mount) - .replace("{SHARED_DEPENDENCIES}", shared_deps.as_str()) - .replace("{MAIN}", format!("{dirs}/{last}").as_str()) - .replace( - "{ADDITIONAL_PYTHON_PATHS}", - additional_python_paths_folders.as_str(), - ), - ) - .await?; - } else { - reserved_variables.insert("PYTHONPATH".to_string(), additional_python_paths_folders); - } - - tracing::info!( - worker_name = %worker_name, - job_id = %job.id, - workspace_id = %job.workspace_id, - "started python code execution {}", - job.id - ); - let child = if !*DISABLE_NSJAIL { - Command::new(NSJAIL_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - // inject PYTHONPATH here - for some reason I had to do it in nsjail conf - .envs(reserved_variables) - .env("PATH", PATH_ENV.as_str()) - .env("BASE_INTERNAL_URL", base_internal_url) - .args(vec![ - "--config", - "run.config.proto", - "--", - PYTHON_PATH.as_str(), - "-u", - "-m", - "wrapper", - ]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - } else { - Command::new(PYTHON_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - .envs(reserved_variables) - .env("PATH", PATH_ENV.as_str()) - .env("BASE_INTERNAL_URL", base_internal_url) - .args(vec!["-u", "-m", "wrapper"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - }; - - handle_child(&job.id, db, logs, child, !*DISABLE_NSJAIL, worker_name, &job.workspace_id).await?; - read_result(job_dir).await -} - -async fn create_dependencies_dir(job_dir: &str) { - DirBuilder::new() - .recursive(true) - .create(&format!("{job_dir}/dependencies")) - .await - .expect("could not create dependencies dir"); -} - -async fn read_result(job_dir: &str) -> error::Result { - let mut file = File::open(format!("{job_dir}/result.json")).await?; - let mut content = "".to_string(); - file.read_to_string(&mut content).await?; - if *CLOUD_HOSTED && content.len() > MAX_RESULT_SIZE { - return Err(Error::ExecutionErr("Result is too large for the cloud app (limit 2MB). - If using this script as part of the flow, use the shared folder to pass heavy data between steps.".to_owned())); - } - serde_json::from_str(&content) - .map_err(|e| Error::ExecutionErr(format!("Error parsing result: {e}"))) -} #[tracing::instrument(level = "trace", skip_all)] async fn handle_dependency_job( @@ -2306,52 +1545,6 @@ async fn handle_flow_dependency_job( Ok(()) } -// #[cfg(not(feature = "deno-lock"))] -// async fn generate_deno_lock( -// _job_id: &Uuid, -// _code: &str, -// _logs: &mut String, -// _job_dir: &str, -// _db: &sqlx::Pool, -// _timeout: i32, -// ) -> error::Result { -// Ok(String::new()) -// } - -// #[cfg(feature = "deno-lock")] -// async fn generate_deno_lock( -// job_id: &Uuid, -// code: &str, -// logs: &mut String, -// job_dir: &str, -// db: &sqlx::Pool, -// timeout: i32, -// ) -> error::Result { -// let _ = write_file(job_dir, "main.ts", code).await?; - -// let child = Command::new(deno_path) -// .current_dir(job_dir) -// .args(vec![ -// "cache", -// "--unstable", -// "--lock=lock.json", -// "--lock-write", -// "main.ts", -// ]) -// .env("NO_COLOR", "1") -// .stdout(Stdio::piped()) -// .stderr(Stdio::piped()) -// .spawn()?; - -// handle_child(job_id, db, logs, timeout, child).await?; - -// let path_lock = format!("{job_dir}/lock.json"); -// let mut file = File::open(path_lock).await?; -// let mut req_content = "".to_string(); -// file.read_to_string(&mut req_content).await?; -// Ok(req_content) -// } - async fn capture_dependency_job( job_id: &Uuid, job_language: &ScriptLang, @@ -2390,189 +1583,9 @@ async fn capture_dependency_job( } } -async fn pip_compile( - job_id: &Uuid, - requirements: &str, - logs: &mut String, - job_dir: &str, - db: &Pool, - worker_name: &str, - w_id: &str, -) -> error::Result { - logs.push_str(&format!("\nresolving dependencies...")); - set_logs(logs, job_id, db).await; - logs.push_str(&format!("\ncontent of requirements:\n{}", requirements)); - let req_hash = calculate_hash(&requirements) ; - if let Some(cached) = sqlx::query_scalar!( - "SELECT lockfile FROM pip_resolution_cache WHERE hash = $1", - req_hash - ).fetch_optional(db).await? { - logs.push_str(&format!("\nfound cached resolution")); - return Ok(cached); - } - let file = "requirements.in"; - let requirements = if let Some(pip_local_dependencies) = PIP_LOCAL_DEPENDENCIES.as_ref() { - let deps = pip_local_dependencies.clone(); - requirements - .lines() - .filter(|s| !deps.contains(&s.to_string())) - .join("\n") - } else { - requirements.to_string() - }; - write_file(job_dir, file, &requirements).await?; - - let mut args = vec!["-q", "--no-header", file, "--resolver=backtracking"]; - if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { - args.extend(["--extra-index-url", url]); - } - if let Some(url) = PIP_INDEX_URL.as_ref() { - args.extend(["--index-url", url]); - } - if let Some(host) = PIP_TRUSTED_HOST.as_ref() { - args.extend(["--trusted-host", host]); - } - let child = Command::new("pip-compile") - .current_dir(job_dir) - .args(args) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; - handle_child(job_id, db, logs, child, false, worker_name, &w_id) - .await - .map_err(|e| Error::ExecutionErr(format!("Lock file generation failed: {e:?}")))?; - let path_lock = format!("{job_dir}/requirements.txt"); - let mut file = File::open(path_lock).await?; - let mut req_content = "".to_string(); - file.read_to_string(&mut req_content).await?; - let lockfile = req_content - .lines() - .filter(|x| !x.trim_start().starts_with('#')) - .map(|x| x.to_string()) - .collect::>() - .join("\n"); - sqlx::query!( - "INSERT INTO pip_resolution_cache (hash, lockfile, expiration) VALUES ($1, $2, now() + ('3 days')::interval) ON CONFLICT (hash) DO UPDATE SET lockfile = $2", - req_hash, - lockfile - ).fetch_optional(db).await?; - Ok(lockfile) -} - -async fn install_go_dependencies( - job_id: &Uuid, - code: &str, - logs: &mut String, - job_dir: &str, - db: &sqlx::Pool, - non_dep_job: bool, - skip_go_mod: bool, - has_sum: bool, - worker_name: &str, - w_id: &str -) -> error::Result { - if !skip_go_mod { - gen_go_mymod(code, job_dir).await?; - let child = Command::new("go") - .current_dir(job_dir) - .args(vec!["mod", "init", "mymod"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; - - handle_child(job_id, db, logs, child, false, worker_name, w_id).await?; - } - - let mut new_lockfile = false; - - let hash = if !has_sum { - calculate_hash(parse_go_imports(&code)?.iter().join("\n").as_str()) - } else { - "".to_string() - }; - - let mut skip_tidy = has_sum; - - if !has_sum { - if let Some(cached) = sqlx::query_scalar!( - "SELECT lockfile FROM pip_resolution_cache WHERE hash = $1", - hash - ).fetch_optional(db).await? { - logs.push_str(&format!("\nfound cached resolution")); - gen_go_mod(code, job_dir, &cached).await?; - skip_tidy = true; - new_lockfile = false; - } else { - new_lockfile = true; - } - } - - let mod_command = if skip_tidy { - "download" - } else { - "tidy" - }; - let child = Command::new(GO_PATH.as_str()) - .current_dir(job_dir) - .env("GOPATH", GO_CACHE_DIR) - .args(vec!["mod", mod_command]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; - handle_child(job_id, db, logs, child, false, worker_name, &w_id) - .await - .map_err(|e| Error::ExecutionErr(format!("Lock file generation failed: {e:?}")))?; - - if (!new_lockfile || has_sum) && non_dep_job{ - return Ok("".to_string()); - } - - - let mut req_content = "".to_string(); - - let mut file = File::open(format!("{job_dir}/go.mod")).await?; - file.read_to_string(&mut req_content).await?; - req_content.push_str(GO_REQ_SPLITTER); - let sum_path = format!("{job_dir}/go.sum"); - if tokio::fs::metadata(&sum_path).await.is_ok() { - let mut file = File::open(sum_path).await?; - file.read_to_string(&mut req_content).await?; - } - - if non_dep_job { - sqlx::query!( - "INSERT INTO pip_resolution_cache (hash, lockfile, expiration) VALUES ($1, $2, now() + ('3 days')::interval) ON CONFLICT (hash) DO UPDATE SET lockfile = $2", - hash, - req_content - ).fetch_optional(db).await?; - - return Ok(String::new()) - } else { - Ok(req_content) - } -} - -async fn gen_go_mymod(code: &str, job_dir: &str) -> error::Result<()> { - let code = if code.trim_start().starts_with("package") { - code.to_string() - } else { - format!("package inner; {code}") - }; - - let mymod_dir = format!("{job_dir}/inner"); - DirBuilder::new() - .recursive(true) - .create(&mymod_dir) - .await - .expect("could not create go's mymod dir"); - - write_file(&mymod_dir, "inner_main.go", &code).await?; - - Ok(()) -} #[tracing::instrument(level = "trace", skip_all)] -async fn get_reserved_variables( +pub async fn get_reserved_variables( job: &QueuedJob, token: &str, db: &sqlx::Pool, @@ -2636,7 +1649,7 @@ async fn get_mem_peak(pid: Option, nsjail: bool) -> i32 { /// - update "queue"."last_ping" every five seconds /// - kill process if we exceed timeout or "queue"."canceled" is set #[tracing::instrument(level = "trace", skip_all)] -async fn handle_child( +pub async fn handle_child( job_id: &Uuid, db: &Pool, logs: &mut String, @@ -2958,20 +1971,7 @@ fn append_with_limit(dst: &mut String, src: &str, limit: &mut usize) { } } -#[tracing::instrument(level = "trace", skip_all)] -async fn set_logs(logs: &str, id: &uuid::Uuid, db: &Pool) { - if sqlx::query!( - "UPDATE queue SET logs = $1 WHERE id = $2", - logs.to_owned(), - id - ) - .execute(db) - .await - .is_err() - { - tracing::error!(%id, "error updating logs for id {id}") - }; -} + /* TODO retry this? */ #[tracing::instrument(level = "trace", skip_all)] @@ -2991,214 +1991,3 @@ async fn append_logs(job_id: uuid::Uuid, logs: impl AsRef, db: impl Borrow< tracing::error!(%job_id, %err, "error updating logs for job {job_id}: {err}"); } } - -pub async fn handle_zombie_jobs_periodically( - db: &Pool, - mut rx: tokio::sync::broadcast::Receiver<()>, - base_internal_url: &str, - rsmq: Option, -) { - loop { - handle_zombie_jobs(db, base_internal_url, rsmq.clone()).await; - - tokio::select! { - _ = tokio::time::sleep(Duration::from_secs(30)) => (), - _ = rx.recv() => { - println!("received killpill for monitor job"); - break; - } - } - } -} - -async fn handle_zombie_jobs(db: &Pool, base_internal_url: &str, rsmq: Option) { - if *RESTART_ZOMBIE_JOBS { - let restarted = sqlx::query!( - "UPDATE queue SET running = false, started_at = null, logs = logs || '\nRestarted job after not receiving job''s ping for too long the ' || now() || '\n\n' WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2 AND job_kind != $3 AND same_worker = false RETURNING id, workspace_id, last_ping", - *ZOMBIE_JOB_TIMEOUT, - JobKind::Flow: JobKind, - JobKind::FlowPreview: JobKind, - ) - .fetch_all(db) - .await - .ok() - .unwrap_or_else(|| vec![]); - - if *METRICS_ENABLED { - QUEUE_ZOMBIE_RESTART_COUNT.inc_by(restarted.len() as _); - } - for r in restarted { - tracing::info!( - "restarted zombie job {} {} {}", - r.id, - r.workspace_id, - r.last_ping - ); - } - } - - let mut timeout_query = "SELECT * FROM queue WHERE last_ping < now() - ($1 || ' seconds')::interval AND running = true AND job_kind != $2".to_string(); - if *RESTART_ZOMBIE_JOBS { - timeout_query.push_str(" AND same_worker = true"); - }; - let timeouts = sqlx::query_as::<_, QueuedJob>( - &timeout_query - ) - .bind(ZOMBIE_JOB_TIMEOUT.as_str()) - .bind(JobKind::Flow) - .fetch_all(db) - .await - .ok() - .unwrap_or_else(|| vec![]); - - if *METRICS_ENABLED { - QUEUE_ZOMBIE_DELETE_COUNT.inc_by(timeouts.len() as _); - } - for job in timeouts { - tracing::info!( - "timedout zombie job {} {}", - job.id, - job.workspace_id, - ); - - // since the job is unrecoverable, the same worker queue should never be sent anything - let (same_worker_tx_never_used, _same_worker_rx_never_used) = mpsc::channel::(1); - - let token = create_token_for_owner( - &db, - &job.workspace_id, - &job.permissioned_as, - "ephemeral-zombie-jobs", - *SESSION_TOKEN_EXPIRY, - &job.email, - ) - .await - .expect("could not create job token"); - - let client = AuthedClient { base_internal_url: base_internal_url.to_string(), token, workspace: job.workspace_id.to_string(), client: OnceCell::new() }; - - let last_ping = job.last_ping.clone(); - let _ = handle_job_error( - db, - &client, - job, - error::Error::ExecutionErr(format!("Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {})", - last_ping.map(|x| x.to_string()).unwrap_or_else(|| "no ping".to_string()), *ZOMBIE_JOB_TIMEOUT)), - None, - true, - same_worker_tx_never_used, - "", - base_internal_url, - rsmq.clone() - ) - .await; - } -} - -async fn handle_python_reqs( - requirements: Vec<&str>, - job: &QueuedJob, - logs: &mut String, - db: &sqlx::Pool, - worker_name: &str, - job_dir: &str, -) -> error::Result> { - let mut req_paths: Vec = vec![]; - let mut vars = vec![("PATH", PATH_ENV.as_str())]; - if !*DISABLE_NSJAIL { - if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { - vars.push(("EXTRA_INDEX_URL", url)); - } - if let Some(url) = PIP_INDEX_URL.as_ref() { - vars.push(("INDEX_URL", url)); - } - if let Some(host) = PIP_TRUSTED_HOST.as_ref() { - vars.push(("TRUSTED_HOST", host)); - } - }; - - for req in requirements { - // todo: handle many reqs - let venv_p = format!("{PIP_CACHE_DIR}/{req}"); - if metadata(&venv_p).await.is_ok() { - req_paths.push(venv_p); - continue; - } - - logs.push_str("\n--- PIP INSTALL ---\n"); - logs.push_str(&format!("\n{req} is being installed for the first time.\n It will be cached for all ulterior uses.")); - - tracing::info!( - worker_name = %worker_name, - job_id = %job.id, - workspace_id = %job.workspace_id, - "started setup python dependencies" - ); - - let child = if !*DISABLE_NSJAIL { - tracing::info!( - worker_name = %worker_name, - job_id = %job.id, - workspace_id = %job.workspace_id, - "starting nsjail" - ); - let mut vars = vars.clone(); - let req = req.to_string(); - vars.push(("REQ", &req)); - vars.push(("TARGET", &venv_p)); - Command::new(NSJAIL_PATH.as_str()) - .current_dir(job_dir) - .env_clear() - .envs(vars) - .args(vec!["--config", "download.config.proto"]) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - } else { - let mut args = vec![ - "-m", - "pip", - "install", - &req, - "-I", - "--no-deps", - "--no-color", - "--isolated", - "--no-warn-conflicts", - "--disable-pip-version-check", - "-t", - venv_p.as_str(), - ]; - if let Some(url) = PIP_EXTRA_INDEX_URL.as_ref() { - args.extend(["--extra-index-url", url]); - } - if let Some(url) = PIP_INDEX_URL.as_ref() { - args.extend(["--index-url", url]); - } - if let Some(host) = PIP_TRUSTED_HOST.as_ref() { - args.extend(["--trusted-host", &host]); - } - Command::new(PYTHON_PATH.as_str()) - .env_clear() - .env("PATH", PATH_ENV.as_str()) - .args(args) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()? - }; - - let child = handle_child(&job.id, db, logs, child, false, worker_name, &job.workspace_id).await; - tracing::info!( - worker_name = %worker_name, - job_id = %job.id, - workspace_id = %job.workspace_id, - is_ok = child.is_ok(), - "finished setting up python dependencies {}", - job.id - ); - child?; - - req_paths.push(venv_p); - } - Ok(req_paths) -} diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 837945362a..d1ddc2a4df 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -19,6 +19,7 @@ use tokio::sync::mpsc::Sender; use tracing::instrument; use uuid::Uuid; use windmill_common::flow_status::{FlowStatusModuleWParent, Iterator, JobResult}; +use windmill_common::jobs::{QueuedJob, JobPayload, RawCode}; use windmill_common::{ error::{self, to_anyhow, Error}, flow_status::{ @@ -30,9 +31,7 @@ use windmill_common::{ type DB = sqlx::Pool; -use windmill_queue::{ - canceled_job_to_result, get_queued_job, push, JobPayload, QueueTransaction, QueuedJob, RawCode, -}; +use windmill_queue::{canceled_job_to_result, get_queued_job, push, QueueTransaction}; // #[instrument(level = "trace", skip_all)] pub async fn update_flow_status_after_job_completion<