134 lines
4.0 KiB
Python
Executable File
134 lines
4.0 KiB
Python
Executable File
from queue import Empty
|
|
import multiprocessing
|
|
import traceback
|
|
import textwrap
|
|
|
|
from .errors import (
|
|
TaskResultMissingError,
|
|
TaskRuntimeError,
|
|
TaskTimeoutError,
|
|
TaskProcessExitError,
|
|
)
|
|
|
|
from .message_types.broker import NodeMode, Items
|
|
from .constants import EXECUTOR_USER_OUTPUT_KEY
|
|
|
|
from multiprocessing.context import SpawnProcess
|
|
|
|
MULTIPROCESSING_CONTEXT = multiprocessing.get_context("spawn")
|
|
|
|
|
|
class TaskExecutor:
|
|
"""Responsible for executing Python code tasks in isolated subprocesses."""
|
|
|
|
@staticmethod
|
|
def create_process(code: str, node_mode: NodeMode, items: Items):
|
|
"""Create a subprocess for executing a Python code task and a queue for communication."""
|
|
|
|
fn = (
|
|
TaskExecutor._all_items
|
|
if node_mode == "all_items"
|
|
else TaskExecutor._per_item
|
|
)
|
|
|
|
queue = MULTIPROCESSING_CONTEXT.Queue()
|
|
process = MULTIPROCESSING_CONTEXT.Process(target=fn, args=(code, items, queue))
|
|
|
|
return process, queue
|
|
|
|
@staticmethod
|
|
def execute_process(
|
|
process: SpawnProcess,
|
|
queue: multiprocessing.Queue,
|
|
task_timeout: int,
|
|
continue_on_fail: bool,
|
|
):
|
|
"""Execute a subprocess for a Python code task."""
|
|
|
|
try:
|
|
process.start()
|
|
process.join(timeout=task_timeout)
|
|
|
|
if process.is_alive():
|
|
TaskExecutor.stop_process(process)
|
|
raise TaskTimeoutError(task_timeout)
|
|
|
|
if process.exitcode != 0:
|
|
assert process.exitcode is not None
|
|
raise TaskProcessExitError(process.exitcode)
|
|
|
|
try:
|
|
returned = queue.get_nowait()
|
|
except Empty:
|
|
raise TaskResultMissingError()
|
|
|
|
if "error" in returned:
|
|
raise TaskRuntimeError(returned["error"])
|
|
|
|
return returned["result"] or []
|
|
|
|
except Exception as e:
|
|
if continue_on_fail:
|
|
return [{"json": {"error": str(e)}}]
|
|
raise
|
|
|
|
@staticmethod
|
|
def stop_process(process: SpawnProcess | None):
|
|
"""Stop a running subprocess, gracefully else force-killing."""
|
|
|
|
if process is None or not process.is_alive():
|
|
return
|
|
|
|
process.terminate()
|
|
process.join(timeout=1) # 1s grace period
|
|
|
|
if process.is_alive():
|
|
process.kill()
|
|
|
|
@staticmethod
|
|
def _all_items(raw_code: str, items: Items, queue: multiprocessing.Queue):
|
|
"""Execute a Python code task in all-items mode."""
|
|
|
|
try:
|
|
code = TaskExecutor._wrap_code(raw_code)
|
|
globals = {"__builtins__": __builtins__, "_items": items}
|
|
exec(code, globals)
|
|
queue.put({"result": globals[EXECUTOR_USER_OUTPUT_KEY]})
|
|
|
|
except Exception as e:
|
|
TaskExecutor._put_error(queue, e)
|
|
|
|
@staticmethod
|
|
def _per_item(raw_code: str, items: Items, queue: multiprocessing.Queue):
|
|
"""Execute a Python code task in per-item mode."""
|
|
|
|
try:
|
|
wrapped_code = TaskExecutor._wrap_code(raw_code)
|
|
compiled_code = compile(wrapped_code, "<per_item_task_execution>", "exec")
|
|
|
|
result = []
|
|
for index, item in enumerate(items):
|
|
globals = {"__builtins__": __builtins__, "_item": item}
|
|
exec(compiled_code, globals)
|
|
user_output = globals[EXECUTOR_USER_OUTPUT_KEY]
|
|
|
|
if user_output is None:
|
|
continue
|
|
|
|
user_output["pairedItem"] = {"item": index}
|
|
result.append(user_output)
|
|
|
|
queue.put({"result": result})
|
|
|
|
except Exception as e:
|
|
TaskExecutor._put_error(queue, e)
|
|
|
|
@staticmethod
|
|
def _wrap_code(raw_code: str) -> str:
|
|
indented_code = textwrap.indent(raw_code, " ")
|
|
return f"def _user_function():\n{indented_code}\n\n{EXECUTOR_USER_OUTPUT_KEY} = _user_function()"
|
|
|
|
@staticmethod
|
|
def _put_error(queue: multiprocessing.Queue, e: Exception):
|
|
queue.put({"error": {"message": str(e), "stack": traceback.format_exc()}})
|