mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-20 00:02:19 +00:00
Stephan Fitzpatrick
d1b1459ebb
* 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
528 lines
15 KiB
Python
528 lines
15 KiB
Python
from typing import Any, Union, Dict
|
|
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
|
|
|
|
from windmill_api.client import AuthenticatedClient
|
|
|
|
from enum import Enum
|
|
|
|
from windmill_api.types import Unset
|
|
|
|
S = TypeVar("S")
|
|
|
|
|
|
class Resource(Generic[S]):
|
|
pass
|
|
|
|
|
|
class JobStatus(Enum):
|
|
WAITING = 1
|
|
RUNNING = 2
|
|
COMPLETED = 3
|
|
|
|
|
|
_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")
|
|
|
|
if env_base_url is not None:
|
|
env_base_url = env_base_url + "/api"
|
|
|
|
base_url_: str = base_url or env_base_url or "http://localhost:8000/api"
|
|
token_: str = token or os.environ.get("WM_TOKEN") or ""
|
|
global _client
|
|
if _client is None:
|
|
_client = AuthenticatedClient(base_url=base_url_, token=token_, timeout=30, verify_ssl=False)
|
|
return _client
|
|
|
|
|
|
def get_workspace() -> str:
|
|
from_env = os.environ.get("WM_WORKSPACE")
|
|
if from_env is None:
|
|
raise Exception("Workspace not passed as WM_WORKSPACE")
|
|
return from_env
|
|
|
|
|
|
def get_version() -> str:
|
|
"""
|
|
Returns the current version of the backend
|
|
"""
|
|
from windmill_api.api.settings import backend_version
|
|
|
|
return backend_version.sync_detailed(client=create_client()).content.decode("us-ascii")
|
|
|
|
|
|
def run_script_async(
|
|
hash: str,
|
|
args: Dict[str, Any] = {},
|
|
scheduled_in_secs: Union[None, int] = None,
|
|
) -> str:
|
|
"""
|
|
Launch the run of a script and return immediately its job id
|
|
"""
|
|
from windmill_api.api.job import run_script_by_hash
|
|
|
|
from windmill_api.models.run_script_by_hash_json_body import RunScriptByHashJsonBody
|
|
|
|
return run_script_by_hash.sync_detailed(
|
|
client=create_client(),
|
|
workspace=get_workspace(),
|
|
hash_=hash,
|
|
json_body=RunScriptByHashJsonBody.from_dict(args),
|
|
scheduled_in_secs=scheduled_in_secs,
|
|
parent_job=os.environ.get("DT_JOB_ID"),
|
|
).content.decode("us-ascii")
|
|
|
|
|
|
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:
|
|
sleep(2.0)
|
|
else:
|
|
sleep(5.0)
|
|
nb_iter += 1
|
|
|
|
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(
|
|
path: str,
|
|
args: Dict[str, Any] = {},
|
|
scheduled_in_secs: Union[None, int] = None,
|
|
) -> str:
|
|
"""
|
|
Launch the run of a script and return immediately its job id
|
|
"""
|
|
from windmill_api.api.job import run_script_by_path
|
|
|
|
from windmill_api.models.run_script_by_path_json_body import RunScriptByPathJsonBody
|
|
|
|
return run_script_by_path.sync_detailed(
|
|
client=create_client(),
|
|
workspace=get_workspace(),
|
|
path=path,
|
|
json_body=RunScriptByPathJsonBody.from_dict(args),
|
|
scheduled_in_secs=scheduled_in_secs,
|
|
parent_job=os.environ.get("DT_JOB_ID"),
|
|
).content.decode("us-ascii")
|
|
|
|
|
|
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:
|
|
sleep(2.0)
|
|
else:
|
|
sleep(5.0)
|
|
nb_iter += 1
|
|
|
|
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:
|
|
"""
|
|
Returns the status of a queued or completed job
|
|
"""
|
|
from windmill_api.models.get_job_response_200_type import GetJobResponse200Type
|
|
from windmill_api.api.job import get_job
|
|
|
|
res = get_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed
|
|
if not res:
|
|
raise Exception(f"Job {job_id} not found")
|
|
elif not res.type:
|
|
raise Exception(f"Unexpected type not found for job {job_id}")
|
|
elif res.type == GetJobResponse200Type.COMPLETEDJOB:
|
|
return JobStatus.COMPLETED
|
|
else:
|
|
if not "running" in res.additional_properties:
|
|
raise Exception(f"Unexpected running not found for completed job {job_id}")
|
|
elif bool(res.additional_properties["running"]):
|
|
return JobStatus.RUNNING
|
|
else:
|
|
return JobStatus.WAITING
|
|
|
|
|
|
def get_result(job_id: str, assert_result_is_not_none: bool = True) -> Dict[str, Any]:
|
|
"""
|
|
Returns the result of a completed job
|
|
"""
|
|
from windmill_api.api.job import get_completed_job
|
|
|
|
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 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
|
|
|
|
|
|
def get_resource(path: Union[str, None] = None, none_if_undefined: bool = False) -> Any:
|
|
"""
|
|
Returns the resource at a given path
|
|
"""
|
|
from windmill_api.api.resource import (
|
|
get_resource_value_interpolated as get_resource_api,
|
|
)
|
|
|
|
path = path or get_state_path()
|
|
parsed = get_resource_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client())
|
|
try:
|
|
content = parsed.content.decode("utf-8")
|
|
parsed = json.loads(content)
|
|
except:
|
|
parsed = None
|
|
|
|
if parsed is None:
|
|
if none_if_undefined:
|
|
return None
|
|
else:
|
|
raise Exception(
|
|
f"Resource at path {path} does not exist or you do not have read permissions on it: {content}"
|
|
)
|
|
|
|
return parsed
|
|
|
|
|
|
def duckdb_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Union[str, None]:
|
|
"""
|
|
Convenient helpers that takes an S3 resource as input and returns the settings necessary to
|
|
initiate an S3 connection from DuckDB
|
|
"""
|
|
from windmill_api.api.helpers import duckdb_connection_settings
|
|
from windmill_api.models.s3_resource import S3Resource
|
|
|
|
parsed = duckdb_connection_settings.sync_detailed(
|
|
workspace=get_workspace(),
|
|
client=create_client(),
|
|
json_body={
|
|
"s3_resource": s3_resource,
|
|
},
|
|
)
|
|
try:
|
|
content = parsed.content.decode("utf-8")
|
|
parsed = json.loads(content)
|
|
except:
|
|
parsed = None
|
|
|
|
if parsed is None:
|
|
if none_if_undefined:
|
|
return None
|
|
else:
|
|
raise Exception(f"Could not generate DuckDB S3 connection settings from the provided resource")
|
|
return parsed
|
|
|
|
|
|
def polars_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Any:
|
|
"""
|
|
Convenient helpers that takes an S3 resource as input and returns the settings necessary to
|
|
initiate an S3 connection from Polars
|
|
"""
|
|
from windmill_api.api.helpers import polars_connection_settings
|
|
|
|
parsed = polars_connection_settings.sync_detailed(
|
|
workspace=get_workspace(),
|
|
client=create_client(),
|
|
json_body={
|
|
"s3_resource": s3_resource,
|
|
},
|
|
)
|
|
try:
|
|
content = parsed.content.decode("utf-8")
|
|
parsed = json.loads(content)
|
|
except:
|
|
parsed = None
|
|
|
|
if parsed is None:
|
|
if none_if_undefined:
|
|
return None
|
|
else:
|
|
raise Exception(f"Could not generate Polars S3 connection settings from the provided resource")
|
|
return parsed
|
|
|
|
|
|
def whoami() -> Union[WhoamiResponse200, None]:
|
|
"""
|
|
Returns the current user
|
|
"""
|
|
from windmill_api.api.user import whoami
|
|
|
|
return whoami.sync(client=create_client(), workspace=get_workspace())
|
|
|
|
|
|
def get_state() -> Any:
|
|
"""
|
|
Get the state
|
|
"""
|
|
return get_resource(None, True)
|
|
|
|
|
|
def set_resource(value: Any, path: Union[str, None] = None, resource_type: str = "state") -> None:
|
|
"""
|
|
Set the resource at a given path as a string, creating it if it does not exist
|
|
"""
|
|
from windmill_api.models.create_resource_json_body import CreateResourceJsonBody
|
|
from windmill_api.models.update_resource_value_json_body import (
|
|
UpdateResourceValueJsonBody,
|
|
)
|
|
from windmill_api.api.resource import (
|
|
exists_resource,
|
|
update_resource_value,
|
|
create_resource,
|
|
)
|
|
|
|
path = path or get_state_path()
|
|
workspace = get_workspace()
|
|
client = create_client()
|
|
if not exists_resource.sync_detailed(workspace=workspace, path=path, client=client).parsed:
|
|
create_resource.sync_detailed(
|
|
workspace=workspace,
|
|
client=client,
|
|
json_body=CreateResourceJsonBody(path=path, value=value, resource_type=resource_type),
|
|
)
|
|
else:
|
|
update_resource_value.sync_detailed(
|
|
workspace=get_workspace(),
|
|
client=client,
|
|
path=path,
|
|
json_body=UpdateResourceValueJsonBody(value=value),
|
|
)
|
|
|
|
|
|
def set_state(value: Any) -> None:
|
|
"""
|
|
Set the state
|
|
"""
|
|
set_resource(value, None)
|
|
|
|
|
|
def set_shared_state_pickle(value: Any, path: str = "state.pickle") -> None:
|
|
"""
|
|
Set the state in the shared folder using pickle
|
|
"""
|
|
import pickle
|
|
|
|
with open(f"/shared/{path}", "wb") as handle:
|
|
pickle.dump(value, handle, protocol=pickle.HIGHEST_PROTOCOL)
|
|
|
|
|
|
def get_shared_state_pickle(path: str = "state.pickle") -> Any:
|
|
"""
|
|
Get the state in the shared folder using pickle
|
|
"""
|
|
import pickle
|
|
|
|
with open(f"/shared/{path}", "rb") as handle:
|
|
return pickle.load(handle)
|
|
|
|
|
|
def set_shared_state(value: Any, path: str = "state.json") -> None:
|
|
"""
|
|
Set the state in the shared folder using pickle
|
|
"""
|
|
import json
|
|
|
|
with open(f"/shared/{path}", "w", encoding="utf-8") as f:
|
|
json.dump(value, f, ensure_ascii=False, indent=4)
|
|
|
|
|
|
def get_shared_state(path: str = "state.json") -> None:
|
|
"""
|
|
Set the state in the shared folder using pickle
|
|
"""
|
|
import json
|
|
|
|
with open(f"/shared/{path}", "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def get_variable(path: str) -> str:
|
|
"""
|
|
Returns the variable at a given path as a string
|
|
"""
|
|
from windmill_api.api.variable import get_variable_value as get_variable_api
|
|
|
|
res = get_variable_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client())
|
|
parsed = res.parsed
|
|
if parsed is None:
|
|
raise Exception(
|
|
f"Variable at path {path} does not exist or you do not have read permissions on it: {res.content.decode('utf-8')}"
|
|
)
|
|
|
|
return parsed
|
|
|
|
|
|
def set_variable(path: str, value: str) -> None:
|
|
"""
|
|
Set the variable at a given path as a string, creating it if it does not exist
|
|
"""
|
|
from windmill_api.api.variable import (
|
|
exists_variable,
|
|
update_variable,
|
|
create_variable,
|
|
)
|
|
|
|
from windmill_api.models.update_variable_json_body import UpdateVariableJsonBody
|
|
from windmill_api.models.create_variable_json_body import CreateVariableJsonBody
|
|
|
|
workspace = get_workspace()
|
|
client = create_client()
|
|
if not exists_variable.sync_detailed(workspace=workspace, path=path, client=client).parsed:
|
|
create_variable.sync_detailed(
|
|
workspace=workspace,
|
|
client=client,
|
|
json_body=CreateVariableJsonBody(path=path, value=value, is_secret=False, description=""),
|
|
)
|
|
else:
|
|
update_variable.sync_detailed(
|
|
workspace=get_workspace(),
|
|
path=path,
|
|
client=client,
|
|
json_body=UpdateVariableJsonBody(value=value),
|
|
)
|
|
|
|
|
|
def get_state_path() -> str:
|
|
state_path = os.environ.get("WM_STATE_PATH_NEW") or os.environ.get("WM_STATE_PATH")
|
|
if state_path is None:
|
|
raise Exception("State path not found")
|
|
return state_path
|
|
|
|
|
|
def get_resume_urls(approver: Union[str, None] = None) -> Dict:
|
|
from windmill_api.api.job import get_resume_urls as get_resume_urls_api
|
|
|
|
workspace = get_workspace()
|
|
client = create_client()
|
|
job_id = os.environ.get("WM_JOB_ID") or "NO_ID"
|
|
import random
|
|
|
|
nonce = random.randint(0, 1000000000)
|
|
res = get_resume_urls_api.sync_detailed(workspace, job_id, nonce, client=client, approver=approver)
|
|
if res.parsed is not None:
|
|
return res.parsed.to_dict()
|
|
else:
|
|
raise Exception("Failed to get resume urls")
|