From 14fd6458a84e6048e7745badec2251404e2d3280 Mon Sep 17 00:00:00 2001 From: pyranota <92104930+pyranota@users.noreply.github.com> Date: Mon, 19 May 2025 22:31:46 +0200 Subject: [PATCH] feat(python): add annotation to skip result post-processing (#5769) * feat(python): add annotation to skip result post-processing Typically windmill will replace all NaN, Infinity and -Infinity from resulting string. We do it because JSON specification does not support these types as well as DB. However it will substitute also the cases when any of those words are used within the string. E.g. script returning "To Infinity and Beyond" will be postprocessed to "To null and Beyond". Current behaviour is done for the performance sake and now can be disabled with `#skip_result_postprocessing` annotation. * add comments * remove extra comments --- backend/windmill-common/src/worker.rs | 1 + .../windmill-worker/src/python_executor.rs | 28 +++++++++++++++---- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index ef4cfa0f46..ea50393018 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -516,6 +516,7 @@ fn parse_file(path: &str) -> Option { pub struct PythonAnnotations { pub no_cache: bool, pub no_postinstall: bool, + pub skip_result_postprocessing: bool, pub py310: bool, pub py311: bool, pub py312: bool, diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index a776bbe7c7..7721a2fc2b 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -842,6 +842,8 @@ pub async fn handle_python_job( ) -> windmill_common::error::Result> { let script_path = crate::common::use_flow_root_path(job.runnable_path()); + let annotations = PythonAnnotations::parse(inner_content); + let (py_version, mut additional_python_paths) = handle_python_deps( job_dir, requirements_o, @@ -856,10 +858,10 @@ pub async fn handle_python_job( canceled_by, &mut Some(occupancy_metrics), precomputed_agent_info, + annotations, ) .await?; - let PythonAnnotations { no_postinstall, .. } = PythonAnnotations::parse(inner_content); tracing::debug!("Finished handling python dependencies"); let python_path = get_python_path( py_version, @@ -872,7 +874,7 @@ pub async fn handle_python_job( ) .await?; - if !no_postinstall { + if !annotations.no_postinstall { if let Err(e) = postinstall(&mut additional_python_paths, job_dir, job, conn).await { tracing::error!("Postinstall stage has failed. Reason: {e}"); } @@ -938,6 +940,8 @@ pub async fn handle_python_job( "".to_string() }; + let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing); + let os_main_override = if let Some(main_override) = main_name.as_ref() { format!("os.environ[\"MAIN_OVERRIDE\"] = \"{main_override}\"\n") } else { @@ -984,7 +988,8 @@ def res_to_json(res): for k, v in res.items(): if type(v).__name__ == 'bytes': res[k] = to_b_64(v) - return re.sub(replace_invalid_fields, ' null ', json.dumps(res, separators=(',', ':'), default=str).replace('\n', '')) + unprocessed = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') + return {postprocessor} try: {preprocessor} @@ -1428,6 +1433,7 @@ async fn handle_python_deps( canceled_by: &mut Option, occupancy_metrics: &mut Option<&mut OccupancyMetrics>, precomputed_agent_info: Option, + annotations: PythonAnnotations, ) -> error::Result<(PyVersion, Vec)> { create_dependencies_dir(job_dir).await; @@ -1445,7 +1451,6 @@ async fn handle_python_deps( let mut annotated_pyv_numeric = None; let is_deployed = requirements_o.is_some(); let instance_pyv = PyVersion::from_instance_version(job_id, w_id, conn).await; - let annotations = windmill_common::worker::PythonAnnotations::parse(inner_content); let requirements = match requirements_o { Some(r) => r, None => { @@ -2326,6 +2331,15 @@ fn get_pyv_from_requirements_lines(requirements_lines: &[&str]) -> PyVersion { } } +// Returns code snippet that needs to be injected into wrapper to post-process results or leave unprocessed +fn get_result_postprocessor<'a>(skip: bool) -> &'a str { + if skip { + "unprocessed" + } else { + "re.sub(replace_invalid_fields, ' null ', unprocessed)" + } +} + #[cfg(feature = "enterprise")] use crate::JobCompletedSender; #[cfg(feature = "enterprise")] @@ -2373,6 +2387,7 @@ pub async fn start_worker( .await .to_vec(); + let annotations = PythonAnnotations::parse(inner_content); let context_envs = build_envs_map(context).await; let (_, additional_python_paths) = handle_python_deps( job_dir, @@ -2388,6 +2403,7 @@ pub async fn start_worker( &mut canceled_by, &mut None, None, + annotations, ) .await?; @@ -2405,6 +2421,7 @@ pub async fn start_worker( ) = prepare_wrapper(job_dir, false, None, None, inner_content, script_path).await?; { + let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing); let indented_transforms = transforms .lines() .map(|x| format!(" {}", x)) @@ -2456,7 +2473,8 @@ for line in sys.stdin: for k, v in res.items(): if type(v).__name__ == 'bytes': res[k] = to_b_64(v) - res_json = re.sub(replace_invalid_fields, ' null ', json.dumps(res, separators=(',', ':'), default=str).replace('\n', '')) + unprocessed = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') + res_json = {postprocessor} sys.stdout.write("wm_res[success]:" + res_json + "\n") except BaseException as e: exc_type, exc_value, exc_traceback = sys.exc_info()