From 619e2784f06ad38a115fd292dc39fdd8d581dbd3 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 4 Mar 2024 20:19:28 +0100 Subject: [PATCH] feat: workflow as code v0 --- backend/ee-repo-ref.txt | 2 +- backend/parsers/windmill-parser-py/src/lib.rs | 45 +++++--- .../parsers/windmill-parser-wasm/src/lib.rs | 2 +- backend/windmill-api/openapi.yaml | 47 ++++++++ backend/windmill-api/src/jobs.rs | 104 +++++++++++++++-- backend/windmill-common/src/jobs.rs | 2 + backend/windmill-queue/src/jobs.rs | 51 +++++---- .../windmill-worker/src/python_executor.rs | 32 +++++- backend/windmill-worker/src/worker.rs | 2 +- backend/windmill-worker/src/worker_flow.rs | 2 +- python-client/wmill/wmill/client.py | 108 +++++++++++++----- 11 files changed, 310 insertions(+), 87 deletions(-) diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 0369bef3a5..140fdd7904 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -a99faf22f440ea9c5732a5dd559686aa881a13c6 \ No newline at end of file +7733647d5d3026654d3dcc279ec0cae934532f8a \ No newline at end of file diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index 72124e5b21..07a128a93f 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -20,17 +20,17 @@ use rustpython_parser::{ Parse, }; -const DEF_MAIN: &str = "def main("; const FUNCTION_CALL: &str = ""; -fn filter_non_main(code: &str) -> String { +fn filter_non_main(code: &str, main_name: &str) -> String { + let def_main = format!("def {}(", main_name); let mut filtered_code = String::new(); let mut code_iter = code.split("\n"); let mut remaining: String = String::new(); while let Some(line) = code_iter.next() { - if line.starts_with(DEF_MAIN) { - filtered_code += DEF_MAIN; - remaining += line.strip_prefix(DEF_MAIN).unwrap(); + if line.starts_with(&def_main) { + filtered_code += &def_main; + remaining += line.strip_prefix(&def_main).unwrap(); remaining += &code_iter.join("\n"); break; } @@ -57,15 +57,21 @@ fn filter_non_main(code: &str) -> String { return filtered_code; } -pub fn parse_python_signature(code: &str) -> anyhow::Result { - let filtered_code = filter_non_main(code); +pub fn parse_python_signature( + code: &str, + override_main: Option, +) -> anyhow::Result { + let main_name = override_main.unwrap_or("main".to_string()); + + let filtered_code = filter_non_main(code, &main_name); if filtered_code.is_empty() { return Err(anyhow::anyhow!("No main function found".to_string(),)); } let ast = Suite::parse(&filtered_code, "main.py") .map_err(|e| anyhow::anyhow!("Error parsing code: {}", e.to_string()))?; + let param = ast.into_iter().find_map(|x| match x { - Stmt::FunctionDef(StmtFunctionDef { name, args, .. }) if &name == "main" => Some(*args), + Stmt::FunctionDef(StmtFunctionDef { name, args, .. }) if &name == &main_name => Some(*args), _ => None, }); if let Some(params) = param { @@ -122,11 +128,16 @@ pub fn parse_python_signature(code: &str) -> anyhow::Result { fn parse_expr(e: &Box) -> Typ { match e.as_ref() { Expr::Name(ExprName { id, .. }) => parse_typ(id.as_ref()), - Expr::Attribute(x) => if x.value.as_name_expr().is_some_and(|x| x.id.as_str() == "wmill") { - parse_typ(x.attr.as_str()) - } else { - Typ::Unknown - }, + Expr::Attribute(x) => { + if x.value + .as_name_expr() + .is_some_and(|x| x.id.as_str() == "wmill") + { + parse_typ(x.attr.as_str()) + } else { + Typ::Unknown + } + } Expr::Subscript(x) => match x.value.as_ref() { Expr::Name(ExprName { id, .. }) => match id.as_str() { "Literal" => { @@ -243,7 +254,7 @@ def main(test1: str, name: datetime.datetime = datetime.now(), byte: bytes = byt "; //println!("{}", serde_json::to_string()?); assert_eq!( - parse_python_signature(code)?, + parse_python_signature(code, None)?, MainArgSignature { star_args: false, star_kwargs: false, @@ -323,7 +334,7 @@ def main(test1: str, "; //println!("{}", serde_json::to_string()?); assert_eq!( - parse_python_signature(code)?, + parse_python_signature(code, None)?, MainArgSignature { star_args: false, star_kwargs: false, @@ -377,7 +388,7 @@ def main(test1: str, "; //println!("{}", serde_json::to_string()?); assert_eq!( - parse_python_signature(code)?, + parse_python_signature(code, None)?, MainArgSignature { star_args: false, star_kwargs: false, @@ -428,7 +439,7 @@ def main(test1: Literal["foo", "bar"], test2: List[Literal["foo", "bar"]]): retu "#; //println!("{}", serde_json::to_string()?); assert_eq!( - parse_python_signature(code)?, + parse_python_signature(code, None)?, MainArgSignature { star_args: false, star_kwargs: false, diff --git a/backend/parsers/windmill-parser-wasm/src/lib.rs b/backend/parsers/windmill-parser-wasm/src/lib.rs index 34c33eddd0..0ecfc32727 100644 --- a/backend/parsers/windmill-parser-wasm/src/lib.rs +++ b/backend/parsers/windmill-parser-wasm/src/lib.rs @@ -55,7 +55,7 @@ pub fn parse_go(code: &str) -> String { #[wasm_bindgen] pub fn parse_python(code: &str) -> String { - wrap_sig(windmill_parser_py::parse_python_signature(code)) + wrap_sig(windmill_parser_py::parse_python_signature(code, None)) } #[wasm_bindgen] diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index f3673d7fb9..8c50f0e74c 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -4840,6 +4840,44 @@ paths: type: string format: uuid + + /w/{workspace}/jobs/workflow_as_code/{job_id}/{entrypoint}: + post: + summary: run code-workflow task + operationId: runCodeWorkflowTask + tags: + - job + parameters: + - $ref: "#/components/parameters/WorkspaceId" + + - name: job_id + in: path + required: true + schema: + type: string + - name: entrypoint + in: path + required: true + schema: + type: string + + requestBody: + description: preview + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/WorkflowTask" + + responses: + "201": + description: job created + content: + text/plain: + schema: + type: string + format: uuid + /w/{workspace}/jobs/run/dependencies: post: summary: run a one-off dependencies job @@ -8857,6 +8895,15 @@ components: required: - args + WorkflowTask: + type: object + properties: + args: + $ref: "#/components/schemas/ScriptArgs" + required: + - args + + CreateResource: type: object properties: diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 0a363fe2c3..93a1d6109f 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -14,6 +14,7 @@ use std::sync::atomic::Ordering; #[cfg(feature = "prometheus")] use tokio::time::Instant; use windmill_common::flow_status::{JobResult, RestartedFrom}; +use windmill_common::jobs::ENTRYPOINT_OVERRIDE; use windmill_common::variables::get_workspace_key; use crate::db::ApiAuthed; @@ -58,7 +59,6 @@ use windmill_common::{ #[cfg(feature = "prometheus")] use windmill_common::{METRICS_DEBUG_ENABLED, METRICS_ENABLED}; - use windmill_common::{get_latest_deployed_hash_for_path, BASE_URL}; use windmill_queue::{ add_completed_job_error, get_queued_job, get_result_by_id_from_running_flow, job_is_complete, @@ -110,6 +110,12 @@ pub fn workspaced_service() -> Router { .head(|| async { "" }) .layer(cors.clone()), ) + .route( + "/run/workflow_as_code/:job_id/:entrypoint", + post(run_workflow_as_code) + .head(|| async { "" }) + .layer(cors.clone()), + ) .route( "/restart/f/:job_id/from/:step_id", post(restart_flow).head(|| async { "" }).layer(cors.clone()), @@ -120,7 +126,7 @@ pub fn workspaced_service() -> Router { ) .route( "/run/p/*script_path", - post(run_job_by_path) + post(run_script_by_path) .head(|| async { "" }) .layer(cors.clone()), ) @@ -150,7 +156,7 @@ pub fn workspaced_service() -> Router { .head(|| async { "" }) .layer(cors.clone()), ) - .route("/run/preview", post(run_preview_job)) + .route("/run/preview", post(run_preview_script)) .route("/add_batch_jobs/:n", post(add_batch_jobs)) .route("/run/preview_flow", post(run_preview_flow_job)) .route( @@ -267,10 +273,10 @@ async fn get_root_job( } async fn compute_root_job_for_flow(db: &DB, w_id: &str, job_id: Uuid) -> error::Result { - let mut job = get_queued_job(job_id, w_id, db).await?; + let mut job = get_queued_job(&job_id, w_id, db).await?; while let Some(j) = job { if let Some(uuid) = j.parent_job { - job = get_queued_job(uuid, w_id, db).await?; + job = get_queued_job(&uuid, w_id, db).await?; } else { return Ok(j.id.to_string()); } @@ -480,7 +486,7 @@ async fn get_flow_job_debug_info( Extension(db): Extension, Path((w_id, id)): Path<(String, Uuid)>, ) -> error::Result { - let job = get_queued_job(id, w_id.as_str(), &db).await?; + let job = get_queued_job(&id, w_id.as_str(), &db).await?; if let Some(job) = job { let is_flow = job.is_flow(); if job.is_flow_step || !is_flow { @@ -641,6 +647,7 @@ pub struct RunJobQuery { payload: Option, job_id: Option, tag: Option, + timeout: Option, } impl RunJobQuery { @@ -858,7 +865,7 @@ async fn cancel_all( for j in jobs.iter() { if !j.running && !j.is_flow_step.unwrap_or(false) { let e = serde_json::json!({"message": format!("Job canceled: cancel_all by {username}"), "name": "Canceled", "reason": "cancel_all", "canceler": username}); - let job_running = get_queued_job(j.id, &w_id, &db).await?; + let job_running = get_queued_job(&j.id, &w_id, &db).await?; if let Some(job_running) = job_running { let add_job = add_completed_job_error( @@ -1784,6 +1791,11 @@ struct Preview { lock: Option, } +#[derive(Deserialize)] +pub struct WorkflowTask { + pub args: Option>>, +} + #[derive(Deserialize)] struct PreviewFlow { value: FlowValue, @@ -2086,7 +2098,7 @@ pub async fn restart_flow( Ok((StatusCode::CREATED, uuid.to_string())) } -pub async fn run_job_by_path( +pub async fn run_script_by_path( authed: ApiAuthed, Extension(db): Extension, Extension(user_db): Extension, @@ -2138,6 +2150,78 @@ pub async fn run_job_by_path( Ok((StatusCode::CREATED, uuid.to_string())) } +pub async fn run_workflow_as_code( + authed: ApiAuthed, + Extension(db): Extension, + Extension(user_db): Extension, + Extension(rsmq): Extension>, + Path((w_id, job_id, entrypoint)): Path<(String, Uuid, String)>, + Query(run_query): Query, + Json(task): Json, +) -> error::Result<(StatusCode, String)> { + #[cfg(feature = "enterprise")] + check_license_key_valid().await?; + check_tag_available_for_workspace(&w_id, &run_query.tag).await?; + + let job = get_queued_job(&job_id, &w_id, &db).await?; + let job = not_found_if_none(job, "Queued Job", &job_id.to_string())?; + let (job_payload, tag, _delete_after_use, timeout) = match job.job_kind { + JobKind::Preview => ( + JobPayload::Code(RawCode { + content: job.raw_code.unwrap_or_default(), + path: job.script_path, + language: job.language.unwrap_or_else(|| ScriptLang::Deno), + lock: job.raw_lock, + concurrent_limit: job.concurrent_limit, + concurrency_time_window_s: job.concurrency_time_window_s, + cache_ttl: job.cache_ttl, + dedicated_worker: None, + }), + Some(job.tag.clone()), + None, + run_query.timeout, + ), + JobKind::Script => script_path_to_payload(job.script_path(), &db, &w_id).await?, + _ => return Err(anyhow::anyhow!("Not supported").into()), + }; + + let mut extra = HashMap::new(); + extra.insert(ENTRYPOINT_OVERRIDE.to_string(), to_raw_value(&entrypoint)); + + let args = PushArgs { args: sqlx::types::Json(task.args), extra }; + let scheduled_for = run_query.get_scheduled_for(&db).await?; + + let tag = run_query.tag.clone().or(tag).or(Some(job.tag)); + let tx = PushIsolationLevel::Isolated(user_db.clone(), authed.clone().into(), rsmq); + + let (uuid, tx) = push( + &db, + tx, + &w_id, + job_payload, + args, + &authed.username, + &authed.email, + username_to_permissioned_as(&authed.username), + scheduled_for, + None, + run_query.parent_job, + run_query.parent_job, + run_query.job_id, + false, + false, + None, + !run_query.invisible_to_owner.unwrap_or(false), + tag, + timeout, + None, + None, + ) + .await?; + tx.commit().await?; + Ok((StatusCode::CREATED, uuid.to_string())) +} + struct Guard { done: bool, id: Uuid, @@ -2689,7 +2773,7 @@ async fn run_wait_result_flow_by_path_internal( run_wait_result(&db, uuid, w_id, early_return).await } -async fn run_preview_job( +async fn run_preview_script( authed: ApiAuthed, Extension(db): Extension, Extension(user_db): Extension, @@ -2744,7 +2828,7 @@ async fn run_preview_job( None, true, tag, - None, + run_query.timeout, None, None, ) diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs index 5e3f770098..3f5bb6dacc 100644 --- a/backend/windmill-common/src/jobs.rs +++ b/backend/windmill-common/src/jobs.rs @@ -5,6 +5,8 @@ use serde_json::value::RawValue; use sqlx::{types::Json, Pool, Postgres, Transaction}; use uuid::Uuid; +pub const ENTRYPOINT_OVERRIDE: &str = "_ENTRYPOINT_OVERRIDE"; + use crate::{ error::{self, Error}, flow_status::{FlowStatus, RestartedFrom}, diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 903ded95e6..43ca83bbdd 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -7,7 +7,9 @@ */ use std::{ - collections::{HashMap, HashSet}, sync::Arc, vec + collections::{HashMap, HashSet}, + sync::Arc, + vec, }; use anyhow::Context; @@ -97,7 +99,6 @@ lazy_static::lazy_static! { } - lazy_static::lazy_static! { pub static ref HTTP_CLIENT: Client = reqwest::ClientBuilder::new() .user_agent("windmill/beta") @@ -106,7 +107,7 @@ lazy_static::lazy_static! { pub static ref HTTP_CLIENT_WORKER: Client = reqwest::ClientBuilder::new() .user_agent("windmill/beta") .build().unwrap(); - + } #[cfg(feature = "enterprise")] @@ -171,7 +172,7 @@ pub async fn cancel_job<'c: 'async_recursion>( e, rsmq.clone(), "server", - false + false, ) .await; if let Err(e) = add_job { @@ -362,7 +363,7 @@ pub async fn add_completed_job_error, _worker_name: &str, - flow_is_done: bool + flow_is_done: bool, ) -> Result { #[cfg(feature = "prometheus")] register_metric( @@ -568,11 +569,13 @@ pub async fn add_completed_job< tx = delete_job(tx, &queued_job.workspace_id, job_id).await?; // tracing::error!("3 {:?}", start.elapsed()); - if queued_job.is_flow_step - { - if let Some(parent_job) = queued_job.parent_job { + if queued_job.is_flow_step { + if let Some(parent_job) = queued_job.parent_job { // persist the flow last progress timestamp to avoid zombie flow jobs - tracing::debug!("Persisting flow last progress timestamp to flow job: {:?}", parent_job); + tracing::debug!( + "Persisting flow last progress timestamp to flow job: {:?}", + parent_job + ); sqlx::query!( "UPDATE queue SET last_ping = now() WHERE id = $1 AND workspace_id = $2", parent_job, @@ -587,11 +590,13 @@ pub async fn add_completed_job< &queued_job.id ).fetch_optional(&mut tx).await?; if r.is_some() { - tracing::info!("parallel flow is done, setting parallel monitor last ping lock for job {}", &queued_job.id); + tracing::info!( + "parallel flow is done, setting parallel monitor last ping lock for job {}", + &queued_job.id + ); } } } - } if !queued_job.is_flow_step @@ -1891,7 +1896,6 @@ pub struct ResultWithId { id: Uuid, } - pub async fn get_result_by_id( db: Pool, w_id: String, @@ -1908,11 +1912,8 @@ pub async fn get_result_by_id( ) .await { - Ok(res) => { - Ok(res) - }, + Ok(res) => Ok(res), Err(_) => { - let running_flow_job = sqlx::query_as::<_, QueuedJob>( "SELECT * FROM queue WHERE COALESCE((SELECT root_job FROM queue WHERE id = $1), $1) = id AND workspace_id = $2" ).bind(flow_id) @@ -2021,7 +2022,6 @@ async fn get_result_by_id_from_original_flow( ) .await?; - if !leaf_jobs_for_flow.contains_key(&node_id.to_string()) { // if the flow is itself a restart flow, the step job might be from the upstream flow let restarted_from = windmill_common::utils::not_found_if_none( @@ -2151,13 +2151,16 @@ async fn extract_result_from_job_result( .fetch_all(db) .await? .into_iter() - .filter_map(|x| ResultWithId::from_row(&x).ok().and_then(|x| x.result.map(|y| (x.id, y)))) + .filter_map(|x| { + ResultWithId::from_row(&x) + .ok() + .and_then(|x| x.result.map(|y| (x.id, y))) + }) .collect::>>>(); let result = job_ids .into_iter() .map(|id| { - rows - .get(&id) + rows.get(&id) .map(|x| x.0.clone()) .unwrap_or_else(|| to_raw_value(&serde_json::Value::Null)) }) @@ -2250,7 +2253,7 @@ pub async fn get_queued_job_tx<'c>( } } -pub async fn get_queued_job(id: Uuid, w_id: &str, db: &DB) -> error::Result> { +pub async fn get_queued_job(id: &Uuid, w_id: &str, db: &DB) -> error::Result> { let r = sqlx::query( "SELECT * FROM queue WHERE id = $1 AND workspace_id = $2", @@ -2306,6 +2309,12 @@ pub struct PushArgs { pub args: Json, } +impl PushArgs { + pub fn insert, V: Into>>(&mut self, k: K, v: V) { + self.extra.insert(k.into(), v.into()); + } +} + #[derive(Deserialize)] pub struct RequestQuery { pub raw: Option, diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index f7b73c3b1f..6c13c090ac 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, process::Stdio}; use itertools::Itertools; use regex::Regex; use serde_json::value::RawValue; -use sqlx::{Pool, Postgres}; +use sqlx::{types::Json, Pool, Postgres}; use tokio::{ fs::{metadata, DirBuilder, File}, io::AsyncReadExt, @@ -14,7 +14,7 @@ use uuid::Uuid; use windmill_common::ee::{get_license_plan, LicensePlan}; use windmill_common::{ error::{self, Error}, - jobs::QueuedJob, + jobs::{QueuedJob, ENTRYPOINT_OVERRIDE}, utils::calculate_hash, worker::WORKER_CONFIG, DB, @@ -272,10 +272,17 @@ pub async fn handle_python_job( last, transforms, spread, - ) = prepare_wrapper(job_dir, inner_content, script_path).await?; + main_name, + ) = prepare_wrapper(job_dir, inner_content, script_path, job.args.as_ref()).await?; create_args_and_out_file(&client, job, job_dir, db).await?; + let os_main_override = if let Some(main_override) = main_name.as_ref() { + format!("import os\nos.environ[\"MAIN_OVERRIDE\"] = \"{main_override}\"\n") + } else { + String::new() + }; + let main_override = main_name.unwrap_or_else(|| "main".to_string()); let wrapper_content: String = format!( r#" import json @@ -284,6 +291,7 @@ import json {import_datetime} import traceback import sys +{os_main_override} from {module_dir_dot} import {last} as inner_script import re @@ -303,7 +311,7 @@ def to_b_64(v: bytes): replace_nan = re.compile(r'(?:\bNaN\b|\\u0000)') try: - res = inner_script.main(**args) + res = inner_script.{main_override}(**args) typ = type(res) if typ.__name__ == 'DataFrame': if typ.__module__ == 'pandas.core.frame': @@ -443,6 +451,7 @@ async fn prepare_wrapper( job_dir: &str, inner_content: &str, script_path: &str, + args: Option<&Json>>>, ) -> error::Result<( &'static str, &'static str, @@ -452,7 +461,15 @@ async fn prepare_wrapper( String, String, String, + Option, )> { + let main_override = args + .map(|x| { + x.0.get(ENTRYPOINT_OVERRIDE) + .map(|x| x.get().to_string().replace("\"", "")) + }) + .flatten(); + let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content); let script_path_splitted = script_path.split("/"); @@ -491,7 +508,7 @@ async fn prepare_wrapper( let _ = write_file(job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?; } - let sig = windmill_parser_py::parse_python_signature(inner_content)?; + let sig = windmill_parser_py::parse_python_signature(inner_content, main_override.clone())?; let transforms = sig .args .iter() @@ -569,6 +586,7 @@ if args["{name}"] is None: last, transforms, spread, + main_override, )) } @@ -991,6 +1009,7 @@ pub async fn start_worker( logs.push_str("\n\n--- PYTHON CODE EXECUTION ---\n"); set_logs(&mut logs, &Uuid::nil(), db).await; + let _args = None; let ( import_loader, import_base64, @@ -1000,7 +1019,8 @@ pub async fn start_worker( last, transforms, spread, - ) = prepare_wrapper(job_dir, inner_content, script_path).await?; + _, + ) = prepare_wrapper(job_dir, inner_content, script_path, _args.as_ref()).await?; { let indented_transforms = transforms diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 5579beaef8..0b04a817f6 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2432,7 +2432,7 @@ pub async fn handle_job_error Any: """Run script synchronously and return its result.""" args = args or {} @@ -131,12 +131,21 @@ class Windmill: logger.info(f"running `{path}` synchronously with {args = }") if isinstance(timeout, dt.timedelta): - timeout = timeout.total_seconds - - start_time = time.time() + timeout = timeout.total_seconds() job_id = self.run_script_async(path=path, hash_=hash_, args=args) + return self.wait_job( + job_id, timeout, verbose, cleanup, assert_result_is_not_none + ) + def wait_job( + self, + job_id, + timeout: dt.timedelta | int | float | None = None, + verbose: bool = False, + cleanup: bool = True, + assert_result_is_not_none: bool = False, + ): def cancel_job(): logger.warning(f"cancelling job: {job_id}") self.post( @@ -147,8 +156,33 @@ class Windmill: if cleanup: atexit.register(cancel_job) + start_time = time.time() + + if isinstance(timeout, dt.timedelta): + timeout = timeout.total_seconds() + while True: - job = self.get_job(job_id) + result_res = self.get(f"/w/{self.workspace}/jobs_u/completed/get_result_maybe/{job_id}", False).json() + + started = result_res["started"] + completed = result_res["completed"] + success = result_res["success"] + + if not started and verbose: + logger.info(f"job {job_id} has not started yet") + + if cleanup and completed: + atexit.unregister(cancel_job) + + if completed: + result = result_res["result"] + if success: + if result is None and assert_result_is_not_none: + raise Exception("Result was none") + return result + else: + error = result["error"] + raise Exception(f"Job {job_id} was not successful: {str(error)}") if timeout and ((time.time() - start_time) > timeout): msg = "reached timeout" @@ -158,33 +192,13 @@ class Windmill: json={"reason": msg}, ) raise TimeoutError(msg) - - result = job.get("result") - canceled, canceled_reason = job.get("canceled"), job.get("canceled_reason") - success = job.get("success") - job_type = job.get("type", "") - completed = job_type.lower() == "completedjob" - - if cleanup and completed: - atexit.unregister(cancel_job) - - if completed: - if success: - if assert_result_is_not_none and result is None: - raise Exception(f"result is None for {job_id = }") - return result - else: - if canceled: - raise Exception(f"job canceled: {canceled_reason}") - else: - error = result.get("error") - raise Exception(f"job failed: {error}") - if verbose: logger.info(f"sleeping 0.5 seconds for {job_id = }") + time.sleep(0.5) + def cancel_running(self) -> dict: """Cancel currently running executions of the same script.""" logger.info("canceling running executions of this script") @@ -226,7 +240,7 @@ class Windmill: return self.get(f"/w/{self.workspace}/jobs_u/get_root_job_id/{job_id}").json() - def get_id_token(self, audience: str) -> dict: + def get_id_token(self, audience: str) -> str: return self.post(f"/w/{self.workspace}/oidc/token/{audience}").text def get_job_status(self, job_id: str) -> JobStatus: @@ -865,3 +879,39 @@ def run_script( cleanup=cleanup, timeout=timeout, ) + + +def task(*args, **kwargs): + + def f(func, tag: str | None = None): + if os.environ.get("WM_JOB_ID") is None or os.environ.get("MAIN_OVERRIDE") == func.__name__: + def inner(*args, **kwargs): + return func(*args, **kwargs) + + return inner + else: + + def inner(*args, **kwargs): + _client=Windmill() + w_id = os.environ.get("WM_WORKSPACE") + job_id = os.environ.get("WM_JOB_ID") + f_name = func.__name__ + json = kwargs + for i, v in enumerate(signature(func).parameters): + json[v[0]] = args[i] + tag_str = f"?tag={tag}" if tag is not None else "" + r = _client.post( + f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{f_name}{tag_str}", + json={"args": json}, + ) + job_id = r.text + logger.info(f"Executing task {func.__name__} on job {job_id}") + return _client.wait_job(job_id) + + return inner + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + return f(args[0], None) + else: + return lambda x: f(x, kwargs.get("tag")) + +