From c7a30f7c9db26252f0ee69e1276ddccc0d52acb3 Mon Sep 17 00:00:00 2001 From: Stephan Fitzpatrick Date: Sun, 19 Nov 2023 17:07:20 -0800 Subject: [PATCH] feat(python): Add functionality and resiliency to wmill python client (#2650) * Refine wmill client.py `run_script_sync` and `run_script_by_path_sync` with more features This commit enriches the functionality of the `run_script_sync` and `run_script_by_path_sync` functions. New features introduced include script cancellation upon exit, logging capabilities, and script execution timeout. These enhancements improve script execution control and provide better debug information. The `get_result` function was also adjusted to enable the control of 'is not None' assertion on the job result. Consequently, user flexibility is enhanced, and the method can cater to cases where a `None` result is within the expected behavior. * remove unnecessary local import and rename cancel_atexit to cleanup --- python-client/wmill/wmill/client.py | 120 ++++++++++++++++++++++++++-- 1 file changed, 112 insertions(+), 8 deletions(-) diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index aed99520d3..127cc163c1 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -1,8 +1,12 @@ from typing import Any, Union, Dict -from typing import Generic, TypeVar +from typing import Generic, TypeVar, Optional import os import json +from datetime import timedelta +import logging +import atexit +import time from time import sleep from windmill_api.models.whoami_response_200 import WhoamiResponse200 @@ -28,6 +32,8 @@ class JobStatus(Enum): _client: "AuthenticatedClient | None" = None +logger = logging.getLogger("wmill_client") + def create_client(base_url: "str | None" = None, token: "str | None" = None) -> AuthenticatedClient: env_base_url = os.environ.get("BASE_INTERNAL_URL") @@ -81,13 +87,48 @@ def run_script_async( ).content.decode("us-ascii") -def run_script_sync(hash: str, args: Dict[str, Any] = {}, verbose: bool = False) -> Dict[str, Any]: +def run_script_sync( + hash: str, + args: Optional[Dict[str, Any]] = None, + verbose: bool = False, + assert_result_is_not_none: bool = True, + cleanup: bool = True, + timeout: Optional[timedelta] = None, +) -> Dict[str, Any]: """ Run a script, wait for it to complete and return the result of the launched script """ + args = args or {} job_id = run_script_async(hash, args, None) + + def cancel_job(): + from windmill_api.api.job.cancel_queued_job import ( + sync_detailed, + CancelQueuedJobJsonBody, + ) + logger.warning(f"cancelling job {job_id}") + return sync_detailed( + workspace=get_workspace(), + id=job_id, + client=create_client(), + json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"), + ) + + if cleanup: + atexit.register(cancel_job) + nb_iter = 0 + + start_time = time.time() + timeout_seconds = timeout.total_seconds() if timeout else None + while get_job_status(job_id) != JobStatus.COMPLETED: + if timeout_seconds is not None: + elapsed_time = time.time() - start_time + if elapsed_time > timeout_seconds: + msg = f"Script execution timed out after {timeout_seconds} seconds" + logger.warning(msg) + raise TimeoutError(msg) if verbose: print(f"Waiting for {job_id} to complete...") if nb_iter < 10: @@ -95,7 +136,21 @@ def run_script_sync(hash: str, args: Dict[str, Any] = {}, verbose: bool = False) else: sleep(5.0) nb_iter += 1 - return get_result(job_id) + + result = get_result( + job_id, + assert_result_is_not_none=assert_result_is_not_none, + ) + + # the job finished--we don't need to cancel it anymore + if cleanup: + atexit.unregister(cancel_job) + + error = isinstance(result, dict) and result.get("error") + + assert not error, error + + return result def run_script_by_path_async( @@ -120,13 +175,48 @@ def run_script_by_path_async( ).content.decode("us-ascii") -def run_script_by_path_sync(path: str, args: Dict[str, Any] = {}, verbose: bool = False) -> Dict[str, Any]: +def run_script_by_path_sync( + path: str, + args: Dict[str, Any] = {}, + verbose: bool = False, + assert_result_is_not_none: bool = True, + cleanup: bool = True, + timeout: Optional[timedelta] = None, +) -> Dict[str, Any]: """ Run a script, wait for it to complete and return the result of the launched script """ + args = args or {} job_id = run_script_by_path_async(path, args, None) + + def cancel_job(): + from windmill_api.api.job.cancel_queued_job import ( + sync_detailed, + CancelQueuedJobJsonBody, + ) + logger.warning(f"cancelling job {job_id}") + return sync_detailed( + workspace=get_workspace(), + id=job_id, + client=create_client(), + json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"), + ) + + if cleanup: + atexit.register(cancel_job) + nb_iter = 0 + + start_time = time.time() + timeout_seconds = timeout.total_seconds() if timeout else None + while get_job_status(job_id) != JobStatus.COMPLETED: + if timeout_seconds is not None: + elapsed_time = time.time() - start_time + if elapsed_time > timeout_seconds: + msg = f"Script execution timed out after {timeout_seconds} seconds" + logger.warning(msg) + raise TimeoutError(msg) if verbose: print(f"Waiting for {job_id} to complete...") if nb_iter < 10: @@ -134,7 +224,21 @@ def run_script_by_path_sync(path: str, args: Dict[str, Any] = {}, verbose: bool else: sleep(5.0) nb_iter += 1 - return get_result(job_id) + + result = get_result( + job_id, + assert_result_is_not_none=assert_result_is_not_none, + ) + + # the job finished--we don't need to cancel it anymore + if cleanup: + atexit.unregister(cancel_job) + + error = isinstance(result, dict) and result.get("error") + + assert not error, error + + return result def get_job_status(job_id: str) -> JobStatus: @@ -160,7 +264,7 @@ def get_job_status(job_id: str) -> JobStatus: return JobStatus.WAITING -def get_result(job_id: str) -> Dict[str, Any]: +def get_result(job_id: str, assert_result_is_not_none: bool = True) -> Dict[str, Any]: """ Returns the result of a completed job """ @@ -169,8 +273,8 @@ def get_result(job_id: str) -> Dict[str, Any]: res = get_completed_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed if not res: raise Exception(f"Job {job_id} not found") - if not res.result: - raise Exception(f"Unexpected result not found for completed job {job_id}") + if assert_result_is_not_none and res.result is None: + raise Exception(f"result was null for completed job {job_id}") else: return res.result