mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-10-21 08:03:33 +00:00
feat: runner groups for shared-process multi-script dedicated workers (#8434)
* feat: add runner groups for shared-process multi-script dedicated workers Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: unify dedicated worker and runner group wrappers into single multi-script wrapper Replace per-language single-script wrappers with the unified load/exec/exec_preprocess/end protocol. Each start_worker() now writes scripts to scripts/<safe_name>/ and uses generate_multi_script_wrapper(). handle_dedicated_process() sends load: on start and exec: per job instead of raw JSON args. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: merge runner groups into dedicated workers with inline arg metadata Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref to match EE branch Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: gate EE-only functions behind cfg(feature = "private") to fix OSS dead_code errors Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: auto-detect runner groups from workspace dependency annotations - New endpoint GET /scripts/list_dedicated_with_deps: returns dedicated scripts with parsed workspace dependency names from content annotations - Frontend: show dep badges in DedicatedWorkersSelector with links to workspace settings, warn when referenced dep doesn't exist, group scripts sharing deps into "Shared runner" sections - Remove manual "Runner groups" tab and RunnerGroupSelector component - Remove runner_groups from WorkerConfigOpt/WorkerConfig (auto-detected) - Fix Node.js single dedicated workers: transpile main.ts -> main.js via Bun.build so the multi-script wrapper's dynamic import() works under Node - Add package.json with type:module in scripts dir to silence Node warning Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: unify dedicated worker wrappers with baked-in codegen and routing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * test: add e2e tests for multi-script dedicated worker routing (bun, deno, python) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: remove dead generate_dedicated_worker_wrapper function Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: add dependency installation to runner groups + make dep functions pub(crate) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: prevent bun loader from intercepting absolute paths within cwd When a plugin's onResolve returns an absolute path, Bun re-invokes the resolver with that path. The loader was then routing it through the remote URL resolver, breaking runner group script imports. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use _wm_ prefix for runner group scripts to avoid bun loader interception Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: extract DENO_UNSTABLE_ARGS constant to avoid repeating flags Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: regenerate system prompts Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: gate private-only exports behind cfg(feature = "private") for OSS build Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: move format strings before handle_dedicated_process to fix lifetime Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: regenerate sqlx offline cache Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix sqlx * fix: skip empty lines in deno e2e tests (double newline from console.log + '\n') Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use dict() instead of {{}} in python wrapper to avoid set literal {{{{}}}} in format!() produces {{}} which Python interprets as an empty set, not a dict. Use dict() which is unambiguous. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: remove deno from runner groups and associated tests Deno resolves dependencies at runtime via URLs/import maps, so there's no shared node_modules/pip install to benefit from runner groups. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: revert deno wrapper to inline old-style with exec: protocol Since deno doesn't support runner groups, the unified multi-script wrapper is unnecessary. Reverted to the old inline wrapper from main but adapted to use the exec:<path>:<args> protocol. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: extract deno wrapper into reusable function and add e2e tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use codebase presence (not nodejs annotation) to determine wrapper import extension On main, codebase scripts import ./main.js (pre-bundled JS). The wrapper_ext was incorrectly based on annotation.nodejs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: improve dedicated workers UI - combine lists, better badges, tooltips - Merge shared runners section with selected tags into one unified list - Move language tag to right side of selector for alignment - Change dep badge color from dark-gray to indigo - Add tooltip on yellow warning badge explaining missing workspace dep Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: group shared runners visually in dedicated workers list - Runner groups shown with a header (Shared runner · language · dep badge) - Scripts in the same group nested under the header - Standalone scripts/flows shown after groups - Used Svelte snippet for reusable tag row rendering Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: improve visual separation between shared runner groups and standalone items Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: give standalone runners same header style as shared runners - Each standalone script/flow gets its own header row with bg-surface-secondary - Header shows "Dedicated runner" / "Flow runner" label, dep link, language badge - Shared runner header: swapped language and dep badge positions - Dep shown as inline link instead of badge in headers for cleaner look Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: inline standalone runner path in header, language badge on right edge, no max height - Standalone items: path shown directly in header row (no sub-row) - Language badge placed after flex-1 spacer (right-aligned) - Removed max-h-64 overflow constraint from the list Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat: consistent badges across runner list - dep+language on right, depBadge snippet - Shared runner scripts: show (workspace) and language badge on right - Standalone items: dep badges and language badge on right (after flex-1) - Shared runner header: dep badge and language badge on right - Extract depBadge snippet to deduplicate dep badge rendering - Picker selector also uses depBadge snippet Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: show language badge on standalone items, hide from shared runner sub-items - Fetch script language from API when not available from workspace deps - Hide dep+language badges from tagRow when script is inside a runner group (already shown in the group header) - Standalone items now always show language badge Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: differentiate badge colors - gray for language, indigo for workspace deps Matches codebase convention: gray for metadata (like script hashes), indigo for linkable features/entities. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use transparent (bordered) badge for language - visible on all backgrounds Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use gray badge for language everywhere Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: revert skills.ts and AI files, add _wm_ exclusion to Windows loader - Revert cli/src/guidance/skills.ts to main (not our change) - Revert AI provider formatting changes (not our change) - Add _wm_ prefix exclusion to loader.bun.windows.js filterResolve Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: update ee-repo-ref and regenerate system prompts after merge Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * perf: use DISTINCT ON in list_dedicated_with_deps to dedup at DB level Avoids fetching all script versions and deduplicating in Rust. Addresses PR review feedback. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use sqlx query! macro for list_dedicated_with_deps and regenerate cache Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: dedicated worker review fixes and test coverage - Fix Python relative imports in dedicated workers (write loader.py, add import loader to wrapper when needed) - Move Python colon parsing inside try/except to prevent crashes on malformed stdin - Add indexOf guard in Bun/Deno wrappers for malformed protocol messages - Add stderr logging for unrecognized stdin commands in all wrappers - Remove asyncio handling from Python wrapper (consistent with normal path) - Add exec_preprocess protocol tests for Bun, Deno, and Python - Add argument transformation tests (dates, bytes, kwargs, sentinel) - Add relative import detection test for Python wrapper - Add PreprocessedArgs variant to DedicatedWorkerResult test helper Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: remove symlink from git and gate has_relative_imports behind private feature Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: update ee-repo-ref for dedicated_worker_ee.rs changes Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: add mixed exec+preprocess test to use ProtocolCmd::Exec variant Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: remove hanging deno missing-preprocessor test The Deno wrapper only generates the exec_preprocess handler when the script has a preprocessor function. Without one, the message is unrecognized and the test hangs reading stdout. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * chore: update ee-repo-ref to 182943e5ad9bf2a905ccdf07d4e346437fb329a9 This commit updates the EE repository reference after PR #466 was merged in windmill-ee-private. Previous ee-repo-ref: 995f701fe3754be6260fc6b679e5de8fc636e68a New ee-repo-ref: 182943e5ad9bf2a905ccdf07d4e346437fb329a9 Automated by sync-ee-ref workflow. --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>
This commit is contained in:
hugocasa
co-authored by
Claude Opus 4.6
windmill-internal-app[bot]
Claude Opus 4.6
windmill-internal-app[bot]
parent
4c8edd5e94
commit
c28314f424
@@ -117,7 +117,12 @@ async fn handle_piptar_uploads(mut rx: tokio::sync::mpsc::UnboundedReceiver<Pipt
|
||||
|
||||
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");
|
||||
pub const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py");
|
||||
|
||||
#[cfg(any(feature = "private", test))]
|
||||
pub fn has_relative_imports(content: &str) -> bool {
|
||||
RELATIVE_IMPORT_REGEX.is_match(content)
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "enterprise", feature = "parquet", unix))]
|
||||
use crate::global_cache::pull_from_tar;
|
||||
@@ -1084,6 +1089,320 @@ fn python_preprocessor_spread(sig: windmill_parser::MainArgSignature, indent: &s
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-computed codegen data for a Python script.
|
||||
/// Computed in Rust from the parsed signature, then baked into the wrapper template.
|
||||
#[cfg(any(feature = "private", test))]
|
||||
pub struct PyScriptCodegen {
|
||||
/// Python module directory dot notation (e.g., "f.my")
|
||||
pub module_dir_dot: String,
|
||||
/// Python module name / last path component (e.g., "script")
|
||||
pub module_name: String,
|
||||
/// Directory path for the module (e.g., "f/my")
|
||||
pub dirs: String,
|
||||
/// Inline Python code for type transforms (dates, bytes, etc.)
|
||||
pub transforms: String,
|
||||
/// Inline Python code for arg spread / filtering
|
||||
pub spread: String,
|
||||
/// Inline Python code for preprocessor arg spread (if applicable)
|
||||
pub pre_spread: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse a Python script and compute the codegen data.
|
||||
/// This reuses the same logic that was used on main in `prepare_wrapper`.
|
||||
#[cfg(any(feature = "private", test))]
|
||||
pub fn compute_py_codegen(content: &str, script_path: &str) -> PyScriptCodegen {
|
||||
let dirs = compute_python_module_dir(script_path);
|
||||
let last = script_path
|
||||
.split("/")
|
||||
.map(|x| {
|
||||
if x.starts_with(|x: char| x.is_ascii_digit()) {
|
||||
format!("_{}", x)
|
||||
} else {
|
||||
x.to_string()
|
||||
}
|
||||
})
|
||||
.last()
|
||||
.unwrap()
|
||||
.replace("-", "_")
|
||||
.replace(" ", "_")
|
||||
.to_lowercase();
|
||||
|
||||
let sig = windmill_parser_py::parse_python_signature(content, None, false).unwrap_or_default();
|
||||
let pre_sig = windmill_parser_py::parse_python_signature(
|
||||
content,
|
||||
Some("preprocessor".to_string()),
|
||||
false,
|
||||
)
|
||||
.ok()
|
||||
.filter(|s| !s.args.is_empty());
|
||||
|
||||
let init_sig = pre_sig.as_ref().unwrap_or(&sig);
|
||||
|
||||
let transforms = init_sig
|
||||
.args
|
||||
.iter()
|
||||
.map(|x| match x.typ {
|
||||
windmill_parser::Typ::Bytes => {
|
||||
let name = &x.name;
|
||||
format!(
|
||||
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
||||
kwargs[\"{name}\"] = base64.b64decode(kwargs[\"{name}\"])\n",
|
||||
)
|
||||
}
|
||||
windmill_parser::Typ::Datetime => {
|
||||
let name = &x.name;
|
||||
format!(
|
||||
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
||||
kwargs[\"{name}\"] = datetime.fromisoformat(kwargs[\"{name}\"])\n",
|
||||
)
|
||||
}
|
||||
windmill_parser::Typ::Date => {
|
||||
let name = &x.name;
|
||||
format!(
|
||||
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
||||
try:\n \
|
||||
kwargs[\"{name}\"] = date.fromisoformat(kwargs[\"{name}\"])\n \
|
||||
except ValueError:\n \
|
||||
for _fmt in (\"%d-%m-%Y\", \"%m/%d/%Y\", \"%d/%m/%Y\", \"%Y/%m/%d\"):\n \
|
||||
try:\n \
|
||||
kwargs[\"{name}\"] = datetime.strptime(kwargs[\"{name}\"], _fmt).date()\n \
|
||||
break\n \
|
||||
except ValueError:\n \
|
||||
continue\n",
|
||||
)
|
||||
}
|
||||
_ => "".to_string(),
|
||||
})
|
||||
.collect::<Vec<String>>()
|
||||
.join("");
|
||||
|
||||
let spread = if sig.star_kwargs {
|
||||
"args = kwargs".to_string()
|
||||
} else {
|
||||
sig.args
|
||||
.into_iter()
|
||||
.map(|x| {
|
||||
let name = &x.name;
|
||||
if x.default.is_none() {
|
||||
format!("args[\"{name}\"] = kwargs.get(\"{name}\")")
|
||||
} else {
|
||||
format!(
|
||||
r#"args["{name}"] = kwargs.get("{name}")
|
||||
if args["{name}"] is None:
|
||||
del args["{name}"]"#
|
||||
)
|
||||
}
|
||||
})
|
||||
.join("\n ")
|
||||
};
|
||||
|
||||
let pre_spread = pre_sig.map(|sig| python_preprocessor_spread(sig, " "));
|
||||
|
||||
let module_dir_dot = dirs.replace("/", ".").replace("-", "_");
|
||||
|
||||
PyScriptCodegen { module_dir_dot, module_name: last, dirs, transforms, spread, pre_spread }
|
||||
}
|
||||
|
||||
/// Script entry for the Python unified wrapper generator.
|
||||
#[cfg(any(feature = "private", test))]
|
||||
pub struct PyScriptEntry<'a> {
|
||||
pub original_path: &'a str,
|
||||
pub codegen: &'a PyScriptCodegen,
|
||||
}
|
||||
|
||||
/// Generate a wrapper for Python dedicated workers and runner groups.
|
||||
/// All scripts are baked in at codegen time with proper Python imports and inline arg handling.
|
||||
/// Protocol:
|
||||
/// exec:<path>:<json_args> -> wm_res[success]:<result> | wm_res[error]:<err>
|
||||
/// exec_preprocess:<path>:<json> -> wm_res[preprocessed_args]:<result> then wm_res[success]:<result> | wm_res[error]:<err>
|
||||
/// end -> exit
|
||||
#[cfg(any(feature = "private", test))]
|
||||
pub fn generate_multi_script_wrapper(
|
||||
scripts: &[PyScriptEntry<'_>],
|
||||
skip_result_postprocessing: bool,
|
||||
any_relative_imports: bool,
|
||||
) -> String {
|
||||
let postprocessor = get_result_postprocessor(skip_result_postprocessing);
|
||||
let res_to_json_body = python_res_to_json_body(postprocessor);
|
||||
|
||||
let imports: String = scripts
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, e)| {
|
||||
format!(
|
||||
"from {module_dir_dot} import {module_name} as _s{i}",
|
||||
module_dir_dot = e.codegen.module_dir_dot,
|
||||
module_name = e.codegen.module_name,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let mut functions = String::new();
|
||||
let mut registrations = String::new();
|
||||
|
||||
for (i, entry) in scripts.iter().enumerate() {
|
||||
let cg = entry.codegen;
|
||||
let indented_transforms = cg
|
||||
.transforms
|
||||
.split('\n')
|
||||
.map(|line| {
|
||||
if line.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" {}", line)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
functions.push_str(&format!(
|
||||
r#"
|
||||
def transform_{i}(kwargs):
|
||||
{indented_transforms}
|
||||
args = dict()
|
||||
{spread}
|
||||
for k, v in list(args.items()):
|
||||
if v == '<function call>':
|
||||
del args[k]
|
||||
return args
|
||||
"#,
|
||||
spread = cg.spread
|
||||
));
|
||||
|
||||
let pre_fn = if let Some(ref pre_spread) = cg.pre_spread {
|
||||
functions.push_str(&format!(
|
||||
r#"
|
||||
def pre_transform_{i}(kwargs):
|
||||
pre_args = dict()
|
||||
{pre_spread}
|
||||
for k, v in list(pre_args.items()):
|
||||
if v == '<function call>':
|
||||
del pre_args[k]
|
||||
return pre_args
|
||||
"#,
|
||||
));
|
||||
format!("pre_transform_{i}")
|
||||
} else {
|
||||
"None".to_string()
|
||||
};
|
||||
|
||||
registrations.push_str(&format!(
|
||||
"scripts[\"{path}\"] = {{ 'mod': _s{i}, 'transform': transform_{i}, 'pre_transform': {pre_fn} }}\n",
|
||||
path = entry.original_path,
|
||||
));
|
||||
}
|
||||
|
||||
let import_loader = if any_relative_imports {
|
||||
"import loader"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
format!(
|
||||
r#"
|
||||
import json
|
||||
import sys
|
||||
import traceback
|
||||
import re
|
||||
import base64
|
||||
from datetime import datetime, date
|
||||
{import_loader}
|
||||
|
||||
{imports}
|
||||
|
||||
scripts = {{}}
|
||||
|
||||
def to_b_64(v: bytes):
|
||||
b64 = base64.b64encode(v)
|
||||
return b64.decode('ascii')
|
||||
|
||||
replace_invalid_fields = re.compile(r'(?:\bNaN\b|\\u0000|Infinity|\-Infinity)')
|
||||
|
||||
def res_to_json(res, typ):
|
||||
{res_to_json_body}
|
||||
{functions}
|
||||
{registrations}
|
||||
|
||||
sys.stdout.write('start\n')
|
||||
sys.stdout.flush()
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if line == 'end':
|
||||
break
|
||||
|
||||
if line.startswith('exec_preprocess:'):
|
||||
try:
|
||||
rest = line[len('exec_preprocess:'):]
|
||||
colon_idx = rest.index(':')
|
||||
script_path = rest[:colon_idx]
|
||||
args_json = rest[colon_idx + 1:]
|
||||
|
||||
entry = scripts.get(script_path)
|
||||
if not entry:
|
||||
err_json = json.dumps({{ "message": "Script not found: " + script_path, "name": "Error" }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
mod = entry['mod']
|
||||
if not hasattr(mod, 'preprocessor') or not callable(mod.preprocessor):
|
||||
err_json = json.dumps({{"message": "preprocessor function is missing", "name": "Error"}}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
kwargs = json.loads(args_json, strict=False)
|
||||
pre_args = entry['pre_transform'](kwargs)
|
||||
preprocessed = mod.preprocessor(**pre_args)
|
||||
preprocessed_json = json.dumps(preprocessed, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[preprocessed_args]:" + preprocessed_json + "\n")
|
||||
main_args = entry['transform'](preprocessed if preprocessed else {{}})
|
||||
res = mod.main(**main_args)
|
||||
typ = type(res)
|
||||
res_json = res_to_json(res, typ)
|
||||
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
||||
except BaseException as e:
|
||||
exc_type, exc_value, exc_traceback = sys.exc_info()
|
||||
tb = traceback.format_tb(exc_traceback)
|
||||
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
if line.startswith('exec:'):
|
||||
try:
|
||||
rest = line[len('exec:'):]
|
||||
colon_idx = rest.index(':')
|
||||
script_path = rest[:colon_idx]
|
||||
args_json = rest[colon_idx + 1:]
|
||||
|
||||
entry = scripts.get(script_path)
|
||||
if not entry:
|
||||
err_json = json.dumps({{ "message": "Script not found: " + script_path, "name": "Error" }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
kwargs = json.loads(args_json, strict=False)
|
||||
args = entry['transform'](kwargs)
|
||||
res = entry['mod'].main(**args)
|
||||
typ = type(res)
|
||||
res_json = res_to_json(res, typ)
|
||||
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
||||
except BaseException as e:
|
||||
exc_type, exc_value, exc_traceback = sys.exc_info()
|
||||
tb = traceback.format_tb(exc_traceback)
|
||||
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
sys.stderr.write("Unknown command: " + line + "\n")
|
||||
"#
|
||||
)
|
||||
}
|
||||
|
||||
async fn prepare_wrapper(
|
||||
job_dir: &str,
|
||||
job_flow_step_id: Option<&str>,
|
||||
@@ -1303,7 +1622,7 @@ async fn replace_pip_secret(
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_python_deps(
|
||||
pub(crate) async fn handle_python_deps(
|
||||
job_dir: &str,
|
||||
requirements_o: Option<&String>,
|
||||
inner_content: &str,
|
||||
@@ -2419,6 +2738,22 @@ pub async fn start_worker(
|
||||
use crate::PyV;
|
||||
tracing::info!("script path: {}", script_path);
|
||||
|
||||
let codegen = compute_py_codegen(inner_content, script_path);
|
||||
|
||||
// Write script to proper module path (e.g., f/my/script.py)
|
||||
let module_dir = format!("{}/{}", job_dir, codegen.dirs);
|
||||
tokio::fs::create_dir_all(&module_dir).await?;
|
||||
write_file(
|
||||
&module_dir,
|
||||
&format!("{}.py", codegen.module_name),
|
||||
inner_content,
|
||||
)?;
|
||||
|
||||
let any_relative_imports = RELATIVE_IMPORT_REGEX.is_match(inner_content);
|
||||
if any_relative_imports {
|
||||
let _ = write_file(job_dir, "loader.py", RELATIVE_PYTHON_LOADER)?;
|
||||
}
|
||||
|
||||
let mut mem_peak: i32 = 0;
|
||||
let mut canceled_by: Option<CanceledBy> = None;
|
||||
let context = variables::get_reserved_variables(
|
||||
@@ -2463,122 +2798,12 @@ pub async fn start_worker(
|
||||
)
|
||||
.await?;
|
||||
|
||||
let (
|
||||
import_loader,
|
||||
import_base64,
|
||||
import_datetime,
|
||||
module_dir_dot,
|
||||
_dirs,
|
||||
last,
|
||||
transforms,
|
||||
spread,
|
||||
_,
|
||||
_,
|
||||
) = prepare_wrapper(job_dir, None, None, None, inner_content, script_path).await?;
|
||||
|
||||
// Parse preprocessor signature if the script has one
|
||||
let pre_spread = windmill_parser_py::parse_python_signature(
|
||||
inner_content,
|
||||
Some("preprocessor".to_string()),
|
||||
false,
|
||||
)
|
||||
.ok()
|
||||
.filter(|sig| !sig.args.is_empty())
|
||||
.map(|sig| python_preprocessor_spread(sig, " "));
|
||||
|
||||
{
|
||||
let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing);
|
||||
let indented_transforms = transforms
|
||||
.lines()
|
||||
.map(|x| format!(" {}", x))
|
||||
.collect::<Vec<String>>()
|
||||
.join("\n");
|
||||
|
||||
let preprocessor_logic = if let Some(ref pre_spread) = pre_spread {
|
||||
format!(
|
||||
r#"
|
||||
if line.startswith('preprocess:'):
|
||||
pre_input = line[len('preprocess:'):]
|
||||
kwargs = json.loads(pre_input, strict=False)
|
||||
if not hasattr(inner_script, 'preprocessor') or not callable(inner_script.preprocessor):
|
||||
err_json = json.dumps({{"message": "preprocessor function is missing", "name": "Error"}}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
try:
|
||||
pre_args = {{}}
|
||||
{pre_spread}
|
||||
for k, v in list(pre_args.items()):
|
||||
if v == '<function call>':
|
||||
del pre_args[k]
|
||||
preprocessed_kwargs = inner_script.preprocessor(**pre_args)
|
||||
preprocessed_json = json.dumps(preprocessed_kwargs, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[preprocessed_args]:" + preprocessed_json + "\n")
|
||||
transform_and_run(preprocessed_kwargs)
|
||||
except BaseException as e:
|
||||
exc_type, exc_value, exc_traceback = sys.exc_info()
|
||||
tb = traceback.format_tb(exc_traceback)
|
||||
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
"#
|
||||
)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
let res_to_json_body = python_res_to_json_body(postprocessor);
|
||||
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
|
||||
import re
|
||||
|
||||
|
||||
def to_b_64(v: bytes):
|
||||
import base64
|
||||
b64 = base64.b64encode(v)
|
||||
return b64.decode('ascii')
|
||||
|
||||
def res_to_json(res, typ):
|
||||
{res_to_json_body}
|
||||
|
||||
def transform_and_run(kwargs):
|
||||
args = {{}}
|
||||
{indented_transforms}
|
||||
{spread}
|
||||
for k, v in list(args.items()):
|
||||
if v == '<function call>':
|
||||
del args[k]
|
||||
res = inner_script.main(**args)
|
||||
typ = type(res)
|
||||
res_json = res_to_json(res, typ)
|
||||
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
||||
|
||||
replace_invalid_fields = re.compile(r'(?:\bNaN\b|\\u0000|Infinity|\-Infinity)')
|
||||
sys.stdout.write('start\n')
|
||||
|
||||
for line in sys.stdin:
|
||||
if line == 'end\n':
|
||||
break
|
||||
line = line.strip()
|
||||
{preprocessor_logic}
|
||||
kwargs = json.loads(line, strict=False)
|
||||
try:
|
||||
transform_and_run(kwargs)
|
||||
except BaseException as e:
|
||||
exc_type, exc_value, exc_traceback = sys.exc_info()
|
||||
tb = traceback.format_tb(exc_traceback)
|
||||
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
||||
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
||||
sys.stdout.flush()
|
||||
"#,
|
||||
let scripts = [PyScriptEntry { original_path: script_path, codegen: &codegen }];
|
||||
let wrapper_content = generate_multi_script_wrapper(
|
||||
&scripts,
|
||||
annotations.skip_result_postprocessing,
|
||||
any_relative_imports,
|
||||
);
|
||||
write_file(job_dir, "wrapper.py", &wrapper_content)?;
|
||||
}
|
||||
@@ -2632,6 +2857,7 @@ for line in sys.stdin:
|
||||
&mut None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
handle_dedicated_process(
|
||||
&python_path,
|
||||
job_dir,
|
||||
@@ -2704,4 +2930,43 @@ mod tests {
|
||||
// @ is replaced with .
|
||||
assert_eq!(compute_python_module_dir("u/@admin/script"), "u/.admin");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_py_codegen_basic_args() {
|
||||
let code = "def main(x: str, y: int):\n return x\n";
|
||||
let cg = compute_py_codegen(code, "f/test/script");
|
||||
assert!(cg.spread.contains("args[\"x\"]"));
|
||||
assert!(cg.spread.contains("args[\"y\"]"));
|
||||
assert!(cg.transforms.is_empty());
|
||||
assert!(cg.pre_spread.is_none());
|
||||
assert_eq!(cg.module_name, "script");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_py_codegen_with_datetime_and_bytes() {
|
||||
let code = "import datetime\n\ndef main(name: str, created_at: datetime.datetime, file: bytes):\n return name\n";
|
||||
let cg = compute_py_codegen(code, "f/my/handler");
|
||||
assert!(cg.transforms.contains("datetime.fromisoformat"));
|
||||
assert!(cg.transforms.contains("base64.b64decode"));
|
||||
assert!(cg.spread.contains("args[\"name\"]"));
|
||||
assert_eq!(cg.module_dir_dot, "f.my");
|
||||
assert_eq!(cg.module_name, "handler");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_py_codegen_star_kwargs() {
|
||||
let code = "def main(**kwargs):\n return kwargs\n";
|
||||
let cg = compute_py_codegen(code, "f/test/star");
|
||||
assert_eq!(cg.spread, "args = kwargs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_py_codegen_with_preprocessor() {
|
||||
let code = "import datetime\n\ndef main(x: str, ts: datetime.datetime):\n return x\n\ndef preprocessor(input: str, when: datetime.datetime):\n return {\"x\": input, \"ts\": when}\n";
|
||||
let cg = compute_py_codegen(code, "f/test/pre");
|
||||
assert!(cg.spread.contains("args[\"x\"]"));
|
||||
assert!(cg.pre_spread.is_some());
|
||||
let pre = cg.pre_spread.as_ref().unwrap();
|
||||
assert!(pre.contains("pre_args[\"input\"]"));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user