starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-11 00:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: workflow as code v0

Browse Source
This commit is contained in:
Ruben Fiszel
2024-03-04 20:19:34 +01:00
parent fd579bc51a
commit ff3aa8ab06
11 changed files with 310 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/ee-repo-ref.txt
+1 -1
View File
@@ -1 +1 @@
a99faf22f440ea9c5732a5dd559686aa881a13c6
7733647d5d3026654d3dcc279ec0cae934532f8a
backend/parsers/windmill-parser-py/src/lib.rs
+28 -17
View File
@@ -20,17 +20,17 @@ use rustpython_parser::{
Parse,
};
const DEF_MAIN: &str = "def main(";
const FUNCTION_CALL: &str = "<function call>";
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<MainArgSignature> {
let filtered_code = filter_non_main(code);
pub fn parse_python_signature(
code: &str,
override_main: Option<String>,
) -> anyhow::Result<MainArgSignature> {
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<MainArgSignature> {
fn parse_expr(e: &Box<Expr>) -> 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,
backend/parsers/windmill-parser-wasm/src/lib.rs
+1 -1
View File
@@ -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]
backend/windmill-api/openapi.yaml
+47
View File
@@ -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:
backend/windmill-api/src/jobs.rs
+94 -10
View File
@@ -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<String> {
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<DB>,
Path((w_id, id)): Path<(String, Uuid)>,
) -> error::Result<Response> {
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<String>,
job_id: Option<Uuid>,
tag: Option<String>,
timeout: Option<i32>,
}
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<String>,
}
#[derive(Deserialize)]
pub struct WorkflowTask {
pub args: Option<HashMap<String, Box<JsonRawValue>>>,
}
#[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<DB>,
Extension(user_db): Extension<UserDB>,
@@ -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<DB>,
Extension(user_db): Extension<UserDB>,
Extension(rsmq): Extension<Option<rsmq_async::MultiplexedRsmq>>,
Path((w_id, job_id, entrypoint)): Path<(String, Uuid, String)>,
Query(run_query): Query<RunJobQuery>,
Json(task): Json<WorkflowTask>,
) -> 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<DB>,
Extension(user_db): Extension<UserDB>,
@@ -2744,7 +2828,7 @@ async fn run_preview_job(
None,
true,
tag,
None,
run_query.timeout,
None,
None,
)
backend/windmill-common/src/jobs.rs
+2
View File
@@ -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},
backend/windmill-queue/src/jobs.rs
+30 -21
View File
@@ -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<R: rsmq_async::RsmqConnection + Clone + Sen
e: serde_json::Value,
rsmq: Option<R>,
_worker_name: &str,
flow_is_done: bool
flow_is_done: bool,
) -> Result<WrappedError, Error> {
#[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<Postgres>,
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::<HashMap<Uuid, Json<Box<RawValue>>>>();
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<Option<QueuedJob>> {
pub async fn get_queued_job(id: &Uuid, w_id: &str, db: &DB) -> error::Result<Option<QueuedJob>> {
let r = sqlx::query(
"SELECT *
FROM queue WHERE id = $1 AND workspace_id = $2",
@@ -2306,6 +2309,12 @@ pub struct PushArgs<T> {
pub args: Json<T>,
}
impl<T> PushArgs<T> {
pub fn insert<K: Into<String>, V: Into<Box<RawValue>>>(&mut self, k: K, v: V) {
self.extra.insert(k.into(), v.into());
}
}
#[derive(Deserialize)]
pub struct RequestQuery {
pub raw: Option<bool>,
backend/windmill-worker/src/python_executor.rs
+26 -6
View File
@@ -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<HashMap<String, Box<RawValue>>>>,
) -> error::Result<(
&'static str,
&'static str,
@@ -452,7 +461,15 @@ async fn prepare_wrapper(
String,
String,
String,
Option<String>,
)> {
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
backend/windmill-worker/src/worker.rs
+1 -1
View File
@@ -2432,7 +2432,7 @@ pub async fn handle_job_error<R: rsmq_async::RsmqConnection + Send + Sync + Clon
if let Err(err) = updated_flow {
if let Some(parent_job_id) = job.parent_job {
if let Ok(Some(parent_job)) =
get_queued_job(parent_job_id, &job.workspace_id, &db).await
get_queued_job(&parent_job_id, &job.workspace_id, &db).await
{
let e = json!({"message": err.to_string(), "name": "InternalErr"});
let _ = add_completed_job_error(
backend/windmill-worker/src/worker_flow.rs
+1 -1
View File
@@ -234,7 +234,7 @@ pub async fn update_flow_status_after_job_completion_internal<
let (mut stop_early, skip_if_stop_early) = if let Some(se) = stop_early_override {
//do not stop early if module is a flow step
let flow_job = get_queued_job(flow, w_id, db)
let flow_job = get_queued_job(&flow, w_id, db)
.await?
.ok_or_else(|| Error::InternalErr(format!("requiring flow to be in the queue")))?;
let module = get_module(&flow_job, module_index);
python-client/wmill/wmill/client.py
+79 -29
View File
@@ -119,10 +119,10 @@ class Windmill:
path: str = None,
hash_: str = None,
args: dict = None,
timeout: dt.timedelta | int | float = None,
timeout: dt.timedelta | int | float | None = None,
verbose: bool = False,
cleanup: bool = True,
assert_result_is_not_none: bool = True,
assert_result_is_not_none: bool = False,
) -> 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"))