Run python nodes out of process, with modules of their own
User code no longer execs in the engine. A pool of persistent worker subprocesses speaks one JSON object per line; the controller installs a proxy as the node's function, so every execution path funnels through it and the pipeline is untouched. A crash costs one subprocess, a per-node timeout is a kill, and cancelling from the canvas is that same kill. The workers run a venv of the user's own on the data volume, filled from a pip manifest versioned beside the flows. Applying it retires the workers and rebuilds, so a package lands without restarting the engine. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017MeiWk3Yq12n2pTvnQWYvt
This commit is contained in:
Claude Fable 5
@@ -13,6 +13,7 @@ from app.core.config import settings
|
||||
from app.core.db import engine
|
||||
from app.flow.controller import FlowController
|
||||
from app.flow.dashboards import DashboardStore
|
||||
from app.flow.workers import PythonWorkerPool
|
||||
from app.models import TokenPayload, User
|
||||
|
||||
reusable_oauth2 = OAuth2PasswordBearer(
|
||||
@@ -113,6 +114,16 @@ def get_dashboard_store(request: Request) -> DashboardStore:
|
||||
DashboardStoreDep = Annotated[DashboardStore, Depends(get_dashboard_store)]
|
||||
|
||||
|
||||
def get_worker_pool(request: Request) -> PythonWorkerPool:
|
||||
pool: PythonWorkerPool | None = getattr(request.app.state, "worker_pool", None)
|
||||
if pool is None:
|
||||
raise HTTPException(status_code=503, detail="The flow engine is not running")
|
||||
return pool
|
||||
|
||||
|
||||
WorkerPoolDep = Annotated[PythonWorkerPool, Depends(get_worker_pool)]
|
||||
|
||||
|
||||
def get_current_active_superuser(current_user: CurrentUser) -> User:
|
||||
if not current_user.is_superuser:
|
||||
raise HTTPException(
|
||||
|
||||
@@ -6,6 +6,7 @@ from app.api.routes import (
|
||||
flows,
|
||||
login,
|
||||
messages,
|
||||
modules,
|
||||
oauth,
|
||||
private,
|
||||
secrets,
|
||||
@@ -23,6 +24,7 @@ api_router.include_router(secrets.router)
|
||||
api_router.include_router(alerts.router)
|
||||
api_router.include_router(dashboards.router)
|
||||
api_router.include_router(messages.router)
|
||||
api_router.include_router(modules.router)
|
||||
# Always mounted so the generated SDK stays the same shape; the endpoints
|
||||
# themselves refuse to work unless MCP is switched on.
|
||||
api_router.include_router(oauth.router)
|
||||
|
||||
@@ -606,6 +606,19 @@ async def trigger_node(
|
||||
return _flow_state(controller, name)
|
||||
|
||||
|
||||
@router.post("/{name}/nodes/{node_id}/cancel", response_model=Message)
|
||||
def cancel_node(name: str, node_id: str, controller: FlowControllerDep) -> Any:
|
||||
"""Stop a node that is running right now, by killing the worker running it.
|
||||
|
||||
Idempotent on purpose: by the time a click reaches here the node may well
|
||||
have finished, and that is the outcome that was asked for.
|
||||
"""
|
||||
pool = controller.workers
|
||||
if pool is not None and pool.cancel(f"{name}.{node_id}"):
|
||||
return Message(message=f"Stopped '{node_id}'")
|
||||
return Message(message=f"'{node_id}' was not running")
|
||||
|
||||
|
||||
@router.get("/{name}/state", response_model=FlowStatePublic)
|
||||
def read_flow_state(name: str, controller: FlowControllerDep) -> Any:
|
||||
"""The last value seen on every message of this flow."""
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""The python packages node code may import.
|
||||
|
||||
A manifest in the flow store, a venv on the data volume, and one button that
|
||||
brings the second in line with the first. Nothing here restarts the engine:
|
||||
the worker pool retires its processes, and the next node call picks up the new
|
||||
packages.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
|
||||
from app.api.deps import FlowControllerDep, WorkerPoolDep, get_current_user
|
||||
from app.flow import modules
|
||||
from app.flow.schemas import ApplyRequest, ApplyResult, ModulesInfo
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/modules", tags=["modules"], dependencies=[Depends(get_current_user)]
|
||||
)
|
||||
|
||||
|
||||
@router.get("/", response_model=ModulesInfo)
|
||||
async def read_modules(controller: FlowControllerDep) -> Any:
|
||||
"""What node code can import, and whether it matches the manifest."""
|
||||
return await run_in_threadpool(modules.info, controller.store)
|
||||
|
||||
|
||||
@router.post("/apply", response_model=ApplyResult)
|
||||
async def apply_modules(
|
||||
body: ApplyRequest, controller: FlowControllerDep, pool: WorkerPoolDep
|
||||
) -> Any:
|
||||
"""Install exactly these requirements, then hand them to the workers.
|
||||
|
||||
A manifest that does not resolve changes nothing: the venv is left as it
|
||||
was and the stored manifest is only written once the install succeeded.
|
||||
"""
|
||||
ok, output = await run_in_threadpool(modules.sync, body.requirements)
|
||||
if not ok:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=output or "These requirements could not be installed",
|
||||
)
|
||||
await run_in_threadpool(controller.store.write_requirements, body.requirements)
|
||||
# Retire the workers first, so the rebuild compiles every node against the
|
||||
# packages that were just installed — a node that could not import one is
|
||||
# the reason this was called, and it stays red until it is built again.
|
||||
pool.respawn_all()
|
||||
await controller.reload()
|
||||
return ApplyResult(ok=True, output=output)
|
||||
@@ -59,6 +59,10 @@ class Settings(BaseSettings):
|
||||
MCP_TOKEN_EXPIRE_MINUTES: int = 60
|
||||
MCP_REFRESH_EXPIRE_DAYS: int = 30
|
||||
FLOW_MAX_WORKERS: int = 4
|
||||
# How long a python node may run before its worker is killed, unless the
|
||||
# node sets its own. Long enough for a slow HTTP call, short enough that a
|
||||
# runaway loop is not a wedged flow.
|
||||
FLOW_NODE_TIMEOUT: float = 30.0
|
||||
# Without a Redis host the engine keeps its state in memory.
|
||||
REDIS_HOST: str | None = None
|
||||
REDIS_PORT: int = 6379
|
||||
|
||||
@@ -9,19 +9,16 @@ never stops the rest.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import sys
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from types import ModuleType
|
||||
from typing import Any, cast
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
|
||||
from app.core.config import settings
|
||||
from app.flow.alerts import AlertManager
|
||||
from app.flow.events import EventBus
|
||||
from app.flow.executor import ExecutionService
|
||||
@@ -55,6 +52,8 @@ from app.flow.secrets import SecretNotFound, resolve_params
|
||||
from app.flow.state import MemoryState, StateBackend
|
||||
from app.flow.store import LIB_DIR, FlowNotFound, FlowStore, LibNotFound
|
||||
from app.flow.supervision import Supervisor
|
||||
from app.flow.worker_main import load_function
|
||||
from app.flow.workers import PythonWorkerPool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -223,38 +222,6 @@ def node_type_info() -> list[NodeTypeInfo]:
|
||||
]
|
||||
|
||||
|
||||
def _load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
|
||||
"""Compile a node's source and return the function to run.
|
||||
|
||||
A node file defines ``process(...)``; if it defines exactly one public
|
||||
function under another name, that one is used.
|
||||
"""
|
||||
digest = hashlib.md5(code.encode()).hexdigest()[:8]
|
||||
module_name = f"_fluksio_node_{flow}_{node_id}_{digest}"
|
||||
|
||||
module = ModuleType(module_name)
|
||||
module.__dict__["__name__"] = module_name
|
||||
sys.modules[module_name] = module
|
||||
exec(compile(code, f"<node {flow}.{node_id}>", "exec"), module.__dict__)
|
||||
|
||||
if callable(getattr(module, "process", None)):
|
||||
return cast(Callable[..., Any], module.process)
|
||||
|
||||
functions = [
|
||||
value
|
||||
for name, value in vars(module).items()
|
||||
if callable(value)
|
||||
and not name.startswith("_")
|
||||
and getattr(value, "__module__", None) == module_name
|
||||
]
|
||||
if len(functions) == 1:
|
||||
return cast(Callable[..., Any], functions[0])
|
||||
raise ValueError(
|
||||
"Define a function named 'process' — this file has "
|
||||
f"{len(functions)} functions to choose from."
|
||||
)
|
||||
|
||||
|
||||
class FlowController:
|
||||
"""Owns the running pipeline and keeps it in step with the store."""
|
||||
|
||||
@@ -267,8 +234,12 @@ class FlowController:
|
||||
fastapi_app: FastAPI | None = None,
|
||||
execution: ExecutionService | None = None,
|
||||
alerts: AlertManager | None = None,
|
||||
workers: PythonWorkerPool | None = None,
|
||||
) -> None:
|
||||
self.store = store
|
||||
# Without a pool, python nodes are compiled and run in this process —
|
||||
# which is what the tests do, and what a bare `Pipeline` has always done.
|
||||
self.workers = workers
|
||||
self.state = state if state is not None else MemoryState()
|
||||
self.events = events
|
||||
self.max_workers = max_workers
|
||||
@@ -448,11 +419,29 @@ class FlowController:
|
||||
# A shared node runs the library's copy, compiled once under
|
||||
# the library's own name so every flow using it agrees.
|
||||
if node_def.source_ref:
|
||||
owner, local = LIB_DIR, node_def.source_ref
|
||||
code = self.store.read_lib_source(node_def.source_ref)
|
||||
function = _load_function(LIB_DIR, node_def.source_ref, code)
|
||||
else:
|
||||
owner, local = flow, node_def.id
|
||||
code = self.store.read_node_source(flow, node_def.id, draft=draft)
|
||||
function = _load_function(flow, node_def.id, code)
|
||||
|
||||
if self.workers is None:
|
||||
function = load_function(owner, local, code)
|
||||
else:
|
||||
# The code never runs here: it is loaded in a worker, and
|
||||
# the node calls that worker instead of a local function.
|
||||
problem = self.workers.compile(owner, local, code)
|
||||
if problem:
|
||||
entry.status = NodeStatus.ERROR
|
||||
entry.error = problem
|
||||
return entry
|
||||
function = self.workers.proxy(
|
||||
owner,
|
||||
local,
|
||||
code,
|
||||
node_id=node_id,
|
||||
timeout=node_def.timeout or settings.FLOW_NODE_TIMEOUT,
|
||||
)
|
||||
node = Node(
|
||||
f=function,
|
||||
requires=_bound(node_def.requires),
|
||||
@@ -559,8 +548,10 @@ class FlowController:
|
||||
|
||||
def compile_check(self, flow: str, node_id: str, code: str) -> str | None:
|
||||
"""Does this source load? Returns what to show the author, or None."""
|
||||
if self.workers is not None:
|
||||
return self.workers.compile(flow, node_id, code)
|
||||
try:
|
||||
_load_function(flow, node_id, code)
|
||||
load_function(flow, node_id, code)
|
||||
except Exception as exc:
|
||||
return _short_error(exc)
|
||||
return None
|
||||
|
||||
@@ -71,6 +71,11 @@ def node_traceback() -> str:
|
||||
exc_type, exc, tb = sys.exc_info()
|
||||
if exc is None:
|
||||
return ""
|
||||
# A node running out of process already trimmed its own; the frames on this
|
||||
# side are the RPC that carried it.
|
||||
remote = getattr(exc, "remote_traceback", "")
|
||||
if remote:
|
||||
return str(remote)
|
||||
frames = traceback.extract_tb(tb)
|
||||
start = next(
|
||||
(i for i, frame in enumerate(frames) if frame.filename.startswith("<node ")),
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
"""The python packages node code may import, in a venv of the user's own.
|
||||
|
||||
A pip manifest lives beside the flows in the same git repository, so what a
|
||||
deployment installed is versioned with what uses it. The packages themselves go
|
||||
into a venv on the data volume rather than into the engine's environment: a
|
||||
pin here can never shadow — or be shadowed by — what the app itself runs on,
|
||||
and the worker processes that import them need nothing from the app.
|
||||
|
||||
``uv pip sync`` rather than install, so a line taken out of the manifest is
|
||||
uninstalled. The manifest is written only after a sync succeeds, which is all
|
||||
the rollback a failed resolve needs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import importlib.metadata
|
||||
import logging
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from app.core.config import settings
|
||||
from app.flow.schemas import ModulePackage, ModulesInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.flow.store import FlowStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Beside the flow store rather than in it: this is installed state, not source.
|
||||
VENV_DIR = settings.FLOWS_DIR.parent / "user-venv"
|
||||
|
||||
#: A resolve that takes longer than this is not going to finish.
|
||||
SYNC_TIMEOUT = 300
|
||||
|
||||
|
||||
def venv_python() -> str:
|
||||
"""The interpreter node code runs on.
|
||||
|
||||
Falls back to the engine's own when there is no venv — a deployment without
|
||||
``uv`` still runs python nodes, it just cannot add packages to them.
|
||||
"""
|
||||
path = VENV_DIR / "bin" / "python"
|
||||
return str(path) if path.exists() else sys.executable
|
||||
|
||||
|
||||
def _marker() -> Path:
|
||||
return VENV_DIR / ".applied"
|
||||
|
||||
|
||||
def _digest(requirements: str) -> str:
|
||||
version = f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
return hashlib.sha256(f"{version}\n{requirements}".encode()).hexdigest()
|
||||
|
||||
|
||||
def ensure_venv() -> None:
|
||||
"""Create the venv if it is missing or its interpreter has gone."""
|
||||
if (VENV_DIR / "bin" / "python").exists():
|
||||
return
|
||||
VENV_DIR.parent.mkdir(parents=True, exist_ok=True)
|
||||
# From the base interpreter, not from the engine's venv: nesting one venv
|
||||
# inside another is how a user pin ends up resolving against app packages.
|
||||
subprocess.run(
|
||||
[
|
||||
"uv",
|
||||
"venv",
|
||||
"--python",
|
||||
str(Path(sys.base_prefix, "bin", "python3")),
|
||||
str(VENV_DIR),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
timeout=SYNC_TIMEOUT,
|
||||
)
|
||||
|
||||
|
||||
def sync(requirements: str) -> tuple[bool, str]:
|
||||
"""Make the venv hold exactly these packages. Returns success and uv's output."""
|
||||
ensure_venv()
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as handle:
|
||||
handle.write(requirements)
|
||||
manifest = handle.name
|
||||
try:
|
||||
result = subprocess.run(
|
||||
# Empty means empty: without the flag uv refuses to clear a venv,
|
||||
# so deleting the last line would leave the package installed.
|
||||
[
|
||||
"uv",
|
||||
"pip",
|
||||
"sync",
|
||||
"--allow-empty-requirements",
|
||||
"--python",
|
||||
venv_python(),
|
||||
manifest,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=SYNC_TIMEOUT,
|
||||
)
|
||||
finally:
|
||||
Path(manifest).unlink(missing_ok=True)
|
||||
|
||||
output = (result.stdout + result.stderr).strip()
|
||||
if result.returncode == 0:
|
||||
_marker().write_text(_digest(requirements))
|
||||
return result.returncode == 0, output
|
||||
|
||||
|
||||
def reconcile(store: FlowStore) -> None:
|
||||
"""Bring the venv in line with the stored manifest. Blocking.
|
||||
|
||||
Called at startup, where nothing may be fatal: a deployment whose packages
|
||||
cannot be installed still runs, with the nodes needing them reporting an
|
||||
import error each time they execute.
|
||||
"""
|
||||
try:
|
||||
# First, and unconditionally: the workers are started against this
|
||||
# interpreter, and it has to be the venv's one before anything is
|
||||
# installed into it, not after.
|
||||
ensure_venv()
|
||||
requirements = store.read_requirements()
|
||||
if not requirements.strip() and not _marker().exists():
|
||||
return
|
||||
if _marker().exists() and _marker().read_text() == _digest(requirements):
|
||||
return
|
||||
ok, output = sync(requirements)
|
||||
if not ok:
|
||||
logger.warning("Could not install the stored modules: %s", output)
|
||||
except Exception:
|
||||
logger.exception("Could not reconcile the module venv")
|
||||
|
||||
|
||||
def info(store: FlowStore) -> ModulesInfo:
|
||||
"""What is installed, what was asked for, and whether the two agree."""
|
||||
requirements = store.read_requirements()
|
||||
config = VENV_DIR / "pyvenv.cfg"
|
||||
version = ""
|
||||
if config.exists():
|
||||
for line in config.read_text().splitlines():
|
||||
if line.startswith("version"):
|
||||
version = line.split("=", 1)[1].strip()
|
||||
|
||||
packages: list[ModulePackage] = []
|
||||
site = sorted(VENV_DIR.glob("lib/python*/site-packages"))
|
||||
if site:
|
||||
packages = sorted(
|
||||
(
|
||||
ModulePackage(name=dist.metadata["Name"] or "", version=dist.version)
|
||||
for dist in importlib.metadata.distributions(path=[str(site[0])])
|
||||
),
|
||||
key=lambda package: package.name.lower(),
|
||||
)
|
||||
|
||||
return ModulesInfo(
|
||||
python_version=version,
|
||||
venv_path=str(VENV_DIR),
|
||||
requirements=requirements,
|
||||
packages=packages,
|
||||
applied=(
|
||||
_marker().read_text() == _digest(requirements)
|
||||
if _marker().exists()
|
||||
# Nothing asked for and nothing installed is already in step.
|
||||
else not requirements.strip()
|
||||
),
|
||||
)
|
||||
@@ -44,6 +44,14 @@ class NodeDef(BaseModel):
|
||||
#: Name of a shared source in the library, instead of this node's own file.
|
||||
#: Editing it edits the copy every flow using it runs.
|
||||
source_ref: str | None = None
|
||||
timeout: float | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
description=(
|
||||
"Seconds this node's code may run before it is stopped. Above 60 "
|
||||
"the engine may deliver its work again while it is still running."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("id")
|
||||
@classmethod
|
||||
@@ -147,6 +155,35 @@ class FlowStatePublic(BaseModel):
|
||||
nodes: list[NodeStatusPublic] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ModulePackage(BaseModel):
|
||||
"""One package installed in the venv node code runs on."""
|
||||
|
||||
name: str
|
||||
version: str
|
||||
|
||||
|
||||
class ModulesInfo(BaseModel):
|
||||
"""The venv node code imports from, and the manifest that describes it."""
|
||||
|
||||
python_version: str = ""
|
||||
venv_path: str = ""
|
||||
requirements: str = ""
|
||||
packages: list[ModulePackage] = Field(default_factory=list)
|
||||
#: Whether what is installed matches the manifest.
|
||||
applied: bool = False
|
||||
|
||||
|
||||
class ApplyRequest(BaseModel):
|
||||
"""A pip manifest, one requirement per line."""
|
||||
|
||||
requirements: str = ""
|
||||
|
||||
|
||||
class ApplyResult(BaseModel):
|
||||
ok: bool
|
||||
output: str = ""
|
||||
|
||||
|
||||
class NodeTypeInfo(BaseModel):
|
||||
"""A node type the editor can offer, with its parameter schema."""
|
||||
|
||||
|
||||
@@ -160,6 +160,21 @@ class FlowStore:
|
||||
def _lib_file(self, name: str) -> Path:
|
||||
return self.root / LIB_DIR / f"{name}.py"
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Module requirements
|
||||
#
|
||||
# What node code may import, versioned with the flows importing it. A file
|
||||
# in the root cannot collide with a flow, which is always a directory.
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def read_requirements(self) -> str:
|
||||
path = self.root / "requirements.txt"
|
||||
return path.read_text() if path.exists() else ""
|
||||
|
||||
def write_requirements(self, text: str) -> None:
|
||||
(self.root / "requirements.txt").write_text(text)
|
||||
self._commit("Update module requirements")
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Shared node sources
|
||||
#
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
"""One user-code process: compiles a node's source once, then runs it on demand.
|
||||
|
||||
This runs under the *user* venv's interpreter, so a node's imports resolve
|
||||
against what the Modules page installed rather than against the engine's own
|
||||
packages. Nothing of the app is importable here — the file is handed to the
|
||||
interpreter by path and is deliberately standard library only.
|
||||
|
||||
The protocol is one JSON object per line: a request arrives on stdin, the reply
|
||||
goes out on a private duplicate of fd 1 taken before anything else can write to
|
||||
it. fd 1 itself is pointed at stderr, so a stray ``write(1, ...)`` — from a
|
||||
native library, or a node printing during an import — lands in the server log
|
||||
instead of corrupting the reply stream.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from types import ModuleType
|
||||
from typing import Any, cast
|
||||
|
||||
#: Enough of a print to debug with. A node printing in a loop must not fill the
|
||||
#: pipe or the reply.
|
||||
MAX_LOG = 16 * 1024
|
||||
|
||||
|
||||
def load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
|
||||
"""Compile a node's source and return the function to run.
|
||||
|
||||
A node file defines ``process(...)``; if it defines exactly one public
|
||||
function under another name, that one is used.
|
||||
"""
|
||||
digest = hashlib.md5(code.encode()).hexdigest()[:8]
|
||||
module_name = f"_fluksio_node_{flow}_{node_id}_{digest}"
|
||||
|
||||
module = ModuleType(module_name)
|
||||
module.__dict__["__name__"] = module_name
|
||||
sys.modules[module_name] = module
|
||||
exec(compile(code, f"<node {flow}.{node_id}>", "exec"), module.__dict__)
|
||||
|
||||
if callable(getattr(module, "process", None)):
|
||||
return cast(Callable[..., Any], module.process)
|
||||
|
||||
functions = [
|
||||
value
|
||||
for name, value in vars(module).items()
|
||||
if callable(value)
|
||||
and not name.startswith("_")
|
||||
and getattr(value, "__module__", None) == module_name
|
||||
]
|
||||
if len(functions) == 1:
|
||||
return cast(Callable[..., Any], functions[0])
|
||||
raise ValueError(
|
||||
"Define a function named 'process' — this file has "
|
||||
f"{len(functions)} functions to choose from."
|
||||
)
|
||||
|
||||
|
||||
class _Capped(io.StringIO):
|
||||
"""Keeps the first ``MAX_LOG`` characters and quietly drops the rest."""
|
||||
|
||||
def write(self, text: str) -> int:
|
||||
room = MAX_LOG - self.tell()
|
||||
if room > 0:
|
||||
super().write(text[:room])
|
||||
return len(text)
|
||||
|
||||
|
||||
def _short_error(exc: BaseException) -> str:
|
||||
"""One line a node author can act on — the engine's rule, applied here."""
|
||||
if isinstance(exc, SyntaxError):
|
||||
# Its own message already names the compiled file, which is noise here.
|
||||
return f"{type(exc).__name__}: {exc.msg} (line {exc.lineno})"
|
||||
|
||||
frames = [
|
||||
frame
|
||||
for frame in traceback.extract_tb(exc.__traceback__)
|
||||
if frame.filename.startswith("<node ")
|
||||
]
|
||||
where = f" (line {frames[-1].lineno})" if frames else ""
|
||||
return f"{type(exc).__name__}: {exc}{where}"
|
||||
|
||||
|
||||
def _node_traceback(exc: BaseException) -> str:
|
||||
"""The traceback from the node's own code onward; the rest is this loop."""
|
||||
frames = traceback.extract_tb(exc.__traceback__)
|
||||
start = next(
|
||||
(i for i, frame in enumerate(frames) if frame.filename.startswith("<node ")),
|
||||
0,
|
||||
)
|
||||
return "".join(
|
||||
["Traceback (most recent call last):\n"]
|
||||
+ traceback.format_list(frames[start:])
|
||||
+ traceback.format_exception_only(type(exc), exc)
|
||||
)
|
||||
|
||||
|
||||
def _handle(request: dict[str, Any], cache: dict[tuple[str, str], Any]) -> Any:
|
||||
"""Compile if needed, then run — the part that may raise the user's error."""
|
||||
flow, node = request["flow"], request["node"]
|
||||
source = request.get("source") or ""
|
||||
digest = hashlib.md5(source.encode()).hexdigest()
|
||||
|
||||
cached = cache.get((flow, node))
|
||||
if cached is None or cached[0] != digest:
|
||||
cache[(flow, node)] = (digest, load_function(flow, node, source))
|
||||
function = cache[(flow, node)][1]
|
||||
|
||||
if request["op"] == "compile":
|
||||
return None
|
||||
|
||||
result = function(
|
||||
**(request.get("kwargs") or {}), params=request.get("params") or {}
|
||||
)
|
||||
try:
|
||||
json.dumps(result)
|
||||
except (TypeError, ValueError):
|
||||
# The typed-message contract only carries JSON, and a node running out
|
||||
# of process is where that stops being a formality.
|
||||
raise ValueError(
|
||||
f"returned {type(result).__name__}, which cannot be sent back as "
|
||||
"JSON — return numbers, strings, booleans, lists or dicts."
|
||||
) from None
|
||||
return result
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rpc = os.fdopen(os.dup(1), "w")
|
||||
# Everything the node writes to the real stdout now goes to the server log.
|
||||
os.dup2(2, 1)
|
||||
|
||||
cache: dict[tuple[str, str], Any] = {}
|
||||
for line in sys.stdin:
|
||||
if not line.strip():
|
||||
continue
|
||||
request = json.loads(line)
|
||||
captured = _Capped()
|
||||
response: dict[str, Any] = {"id": request.get("id")}
|
||||
try:
|
||||
with (
|
||||
contextlib.redirect_stdout(captured),
|
||||
contextlib.redirect_stderr(captured),
|
||||
):
|
||||
response["result"] = _handle(request, cache)
|
||||
response["ok"] = True
|
||||
except Exception as exc:
|
||||
response["ok"] = False
|
||||
response["error"] = {
|
||||
"type": type(exc).__name__,
|
||||
"message": str(exc),
|
||||
"short": _short_error(exc),
|
||||
"traceback": _node_traceback(exc),
|
||||
}
|
||||
response["logs"] = captured.getvalue()
|
||||
rpc.write(json.dumps(response) + "\n")
|
||||
rpc.flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,342 @@
|
||||
"""A pool of subprocesses that run the code people write in python nodes.
|
||||
|
||||
User code used to be ``exec``'d in the engine process, where an ``os._exit``,
|
||||
a segfaulting C extension or a ``while True`` took the whole engine with it.
|
||||
Here each node call is an RPC to a long-lived worker: a crash costs one
|
||||
subprocess, a timeout is a kill, and a cancel is the same kill on request.
|
||||
|
||||
The workers run the user venv's interpreter, so what the Modules page installs
|
||||
is what a node can import. Only JSON crosses the boundary, which the typed
|
||||
message contract already guarantees for everything a node consumes or provides.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import select
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from app.flow.events import EventBus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
WORKER_MAIN = Path(__file__).with_name("worker_main.py")
|
||||
|
||||
#: Importing what a node needs can be slow the first time; compiling is not
|
||||
#: something a person is watching a spinner for.
|
||||
COMPILE_TIMEOUT = 60.0
|
||||
|
||||
|
||||
class RemoteError(Exception):
|
||||
"""Something that went wrong inside a worker, re-raised on this side.
|
||||
|
||||
``remote_traceback`` is the worker's own traceback, trimmed to the node's
|
||||
code — the frames here are the RPC, which the node's author did not write.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, remote_traceback: str = "") -> None:
|
||||
super().__init__(message)
|
||||
self.remote_traceback = remote_traceback
|
||||
|
||||
|
||||
class NodeTimeout(RemoteError):
|
||||
"""The node ran past its timeout, so its worker was killed."""
|
||||
|
||||
|
||||
class NodeCancelled(RemoteError):
|
||||
"""Someone asked for this node to stop while it was running."""
|
||||
|
||||
|
||||
_remote_types: dict[str, type[RemoteError]] = {}
|
||||
|
||||
|
||||
def _remote_class(name: str) -> type[RemoteError]:
|
||||
"""A ``RemoteError`` wearing the remote exception's name.
|
||||
|
||||
The engine renders a node failure as ``f"{type(exc).__name__}: {exc}"``, so
|
||||
the author still reads ``ValueError: bad input`` rather than the name of
|
||||
the transport that carried it.
|
||||
"""
|
||||
cls = _remote_types.get(name)
|
||||
if cls is None:
|
||||
cls = type(name, (RemoteError,), {})
|
||||
_remote_types[name] = cls
|
||||
return cls
|
||||
|
||||
|
||||
class _Worker:
|
||||
"""One subprocess, and the framing of one request/response over its pipes."""
|
||||
|
||||
def __init__(self, python: str, generation: int) -> None:
|
||||
self.generation = generation
|
||||
self.cancelled = False
|
||||
self.proc = subprocess.Popen(
|
||||
[python, str(WORKER_MAIN)],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
close_fds=True,
|
||||
)
|
||||
|
||||
def alive(self) -> bool:
|
||||
return self.proc.poll() is None
|
||||
|
||||
def send(self, request: dict[str, Any]) -> None:
|
||||
assert self.proc.stdin is not None
|
||||
self.proc.stdin.write((json.dumps(request) + "\n").encode())
|
||||
self.proc.stdin.flush()
|
||||
|
||||
def read_line(self, deadline: float) -> str | None:
|
||||
"""One reply line; ``None`` past the deadline, ``""`` if the pipe closed."""
|
||||
assert self.proc.stdout is not None
|
||||
fd = self.proc.stdout.fileno()
|
||||
buffer = bytearray()
|
||||
while True:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
return None
|
||||
ready, _, _ = select.select([fd], [], [], remaining)
|
||||
if not ready:
|
||||
return None
|
||||
chunk = os.read(fd, 65536)
|
||||
if not chunk:
|
||||
# Partial output before the pipe closed is a half-written
|
||||
# reply, which is no more use than none at all.
|
||||
return ""
|
||||
buffer += chunk
|
||||
if buffer.endswith(b"\n"):
|
||||
return buffer.decode(errors="replace")
|
||||
|
||||
def kill(self) -> None:
|
||||
"""SIGKILL: user code has no cleanup we can trust to run.
|
||||
|
||||
The pipes are deliberately left alone — a cancel runs on another
|
||||
thread than the one blocked reading this worker, and closing the fd
|
||||
out from under it is how you get a reader on somebody else's socket.
|
||||
Killing the process closes the far end, which is what wakes the reader.
|
||||
"""
|
||||
try:
|
||||
self.proc.send_signal(signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
# Reap it, so a killed worker does not linger as a zombie.
|
||||
try:
|
||||
self.proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("Worker %s did not die", self.proc.pid)
|
||||
|
||||
|
||||
class PythonWorkerPool:
|
||||
"""Fixed set of worker slots, handed out one call at a time.
|
||||
|
||||
A slot holds a live worker or nothing; a slot that is empty, dead or from
|
||||
before the last ``respawn_all`` spawns a fresh process the next time it is
|
||||
taken. There is no background thread — the pool only does work while a node
|
||||
is calling it.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, python: str, size: int = 4, events: EventBus | None = None
|
||||
) -> None:
|
||||
self.python = python
|
||||
self.size = size
|
||||
self.events = events
|
||||
self._idle: queue.Queue[_Worker | None] = queue.Queue()
|
||||
# ponytail: _running is last-wins; two concurrent runs of one node mean
|
||||
# cancel kills the newest. Key by run id if that ever matters.
|
||||
self._running: dict[str, _Worker] = {}
|
||||
self._generation = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
"""Open the slots. Processes are spawned by the first call that needs one."""
|
||||
for _ in range(self.size):
|
||||
self._idle.put(None)
|
||||
|
||||
def stop(self) -> None:
|
||||
self._generation += 1
|
||||
for worker in list(self._running.values()):
|
||||
worker.kill()
|
||||
for slot in self._drain():
|
||||
if slot is not None:
|
||||
slot.kill()
|
||||
|
||||
def respawn_all(self) -> None:
|
||||
"""Retire every worker, so the next call picks up a changed venv.
|
||||
|
||||
Idle ones go now; a busy one is replaced when it comes back, because
|
||||
its generation no longer matches.
|
||||
"""
|
||||
with self._lock:
|
||||
self._generation += 1
|
||||
for slot in self._drain():
|
||||
if slot is not None:
|
||||
slot.kill()
|
||||
self._idle.put(None)
|
||||
|
||||
def _drain(self) -> list[_Worker | None]:
|
||||
slots = []
|
||||
while True:
|
||||
try:
|
||||
slots.append(self._idle.get_nowait())
|
||||
except queue.Empty:
|
||||
return slots
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Slots
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _acquire(self) -> _Worker:
|
||||
"""Take a slot, blocking while every worker is busy — that is the backpressure."""
|
||||
slot = self._idle.get()
|
||||
if (
|
||||
slot is not None
|
||||
and slot.alive()
|
||||
and slot.generation == self._generation
|
||||
and not slot.cancelled
|
||||
):
|
||||
return slot
|
||||
if slot is not None:
|
||||
slot.kill()
|
||||
try:
|
||||
return _Worker(self.python, self._generation)
|
||||
except Exception as exc:
|
||||
self._idle.put(None)
|
||||
raise RemoteError(f"worker unavailable: {exc}") from exc
|
||||
|
||||
def _release(self, worker: _Worker) -> None:
|
||||
reusable = (
|
||||
worker.alive()
|
||||
and worker.generation == self._generation
|
||||
and not worker.cancelled
|
||||
)
|
||||
if reusable:
|
||||
self._idle.put(worker)
|
||||
return
|
||||
worker.kill()
|
||||
self._idle.put(None)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Calls
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _request(
|
||||
self, payload: dict[str, Any], timeout: float, node_id: str = ""
|
||||
) -> dict[str, Any]:
|
||||
worker = self._acquire()
|
||||
if node_id:
|
||||
self._running[node_id] = worker
|
||||
self._publish(
|
||||
{
|
||||
"type": "node_started",
|
||||
# A flow name cannot contain a dot, so this is exact.
|
||||
"flow": node_id.split(".", 1)[0],
|
||||
"node": node_id,
|
||||
"ts": time.time(),
|
||||
}
|
||||
)
|
||||
try:
|
||||
try:
|
||||
worker.send(payload)
|
||||
except OSError as exc:
|
||||
raise RemoteError(f"worker died: {exc}") from exc
|
||||
|
||||
line = worker.read_line(time.monotonic() + timeout)
|
||||
if line:
|
||||
return dict(json.loads(line))
|
||||
if worker.cancelled:
|
||||
raise NodeCancelled("cancelled while it was running")
|
||||
if line is None:
|
||||
worker.kill()
|
||||
raise NodeTimeout(f"exceeded {timeout}s and was killed")
|
||||
raise RemoteError("worker died")
|
||||
finally:
|
||||
if node_id:
|
||||
self._running.pop(node_id, None)
|
||||
self._release(worker)
|
||||
|
||||
def compile(self, flow: str, node: str, source: str) -> str | None:
|
||||
"""Load this source in a worker. Returns what to show the author, or None."""
|
||||
try:
|
||||
response = self._request(
|
||||
{"op": "compile", "flow": flow, "node": node, "source": source},
|
||||
timeout=COMPILE_TIMEOUT,
|
||||
)
|
||||
except RemoteError as exc:
|
||||
return f"{type(exc).__name__}: {exc}"
|
||||
if response.get("ok"):
|
||||
return None
|
||||
error = response.get("error") or {}
|
||||
return str(error.get("short") or "The node could not be loaded.")
|
||||
|
||||
def run(
|
||||
self,
|
||||
flow: str,
|
||||
node: str,
|
||||
source: str,
|
||||
kwargs: dict[str, Any],
|
||||
params: dict[str, Any] | None,
|
||||
node_id: str,
|
||||
timeout: float,
|
||||
) -> Any:
|
||||
response = self._request(
|
||||
{
|
||||
"op": "run",
|
||||
"flow": flow,
|
||||
"node": node,
|
||||
"source": source,
|
||||
"kwargs": kwargs,
|
||||
"params": params or {},
|
||||
},
|
||||
timeout=timeout,
|
||||
node_id=node_id,
|
||||
)
|
||||
# Into the tee, from the thread the engine is capturing on: this is
|
||||
# what puts a node's prints in the log panel, so it has to happen
|
||||
# before the value comes back or the error goes up.
|
||||
logs = response.get("logs")
|
||||
if logs:
|
||||
sys.stdout.write(logs)
|
||||
if response.get("ok"):
|
||||
return response.get("result")
|
||||
error = response.get("error") or {}
|
||||
raise _remote_class(str(error.get("type") or "RemoteError"))(
|
||||
str(error.get("message") or "the node failed"),
|
||||
str(error.get("traceback") or ""),
|
||||
)
|
||||
|
||||
def proxy(
|
||||
self, flow: str, node: str, source: str, node_id: str, timeout: float
|
||||
) -> Callable[..., Any]:
|
||||
"""The callable a python node runs instead of its own compiled function."""
|
||||
|
||||
def call(params: dict[str, Any] | None = None, **kwargs: Any) -> Any:
|
||||
return self.run(flow, node, source, kwargs, params, node_id, timeout)
|
||||
|
||||
return call
|
||||
|
||||
def cancel(self, node_id: str) -> bool:
|
||||
"""Stop a node that is running now. False when there was nothing to stop."""
|
||||
worker = self._running.get(node_id)
|
||||
if worker is None:
|
||||
return False
|
||||
worker.cancelled = True
|
||||
worker.kill()
|
||||
return True
|
||||
|
||||
def _publish(self, event: dict[str, Any]) -> None:
|
||||
if self.events is not None:
|
||||
self.events.publish(event)
|
||||
+16
-2
@@ -5,6 +5,7 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
|
||||
import sentry_sdk
|
||||
from fastapi import FastAPI
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.routing import APIRoute
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
@@ -13,7 +14,7 @@ from app.api.main import api_router
|
||||
from app.api.routes.alerts import read_config as read_alerts_config
|
||||
from app.core import security
|
||||
from app.core.config import settings
|
||||
from app.flow import logs
|
||||
from app.flow import logs, modules
|
||||
from app.flow.alerts import AlertManager
|
||||
from app.flow.controller import FlowController
|
||||
from app.flow.dashboards import DashboardStore
|
||||
@@ -26,6 +27,7 @@ from app.flow.secrets import init_secrets
|
||||
from app.flow.state import MemoryState, RedisState, StateBackend
|
||||
from app.flow.store import FlowStore
|
||||
from app.flow.watchdog import LoopWatchdog
|
||||
from app.flow.workers import PythonWorkerPool
|
||||
|
||||
|
||||
def custom_generate_unique_id(route: APIRoute) -> str:
|
||||
@@ -74,14 +76,25 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
max_workers=settings.FLOW_MAX_WORKERS,
|
||||
events=event_bus,
|
||||
)
|
||||
store = FlowStore(settings.FLOWS_DIR)
|
||||
# The packages node code imports, before anything tries to import them.
|
||||
await run_in_threadpool(modules.reconcile, store)
|
||||
pool = PythonWorkerPool(
|
||||
python=modules.venv_python(),
|
||||
size=settings.FLOW_MAX_WORKERS,
|
||||
events=event_bus,
|
||||
)
|
||||
pool.start()
|
||||
app.state.worker_pool = pool
|
||||
controller = FlowController(
|
||||
store=FlowStore(settings.FLOWS_DIR),
|
||||
store=store,
|
||||
state=_state_backend(),
|
||||
events=event_bus,
|
||||
max_workers=settings.FLOW_MAX_WORKERS,
|
||||
fastapi_app=app,
|
||||
execution=execution,
|
||||
alerts=alerts,
|
||||
workers=pool,
|
||||
)
|
||||
app.state.flow_controller = controller
|
||||
dashboards = DashboardStore(controller.store)
|
||||
@@ -104,6 +117,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
watchdog_task.cancel()
|
||||
alerts_task.cancel()
|
||||
await controller.stop()
|
||||
pool.stop()
|
||||
close_shared_client()
|
||||
if settings.MCP_ENABLED:
|
||||
from app.mcp.http import aclose
|
||||
|
||||
@@ -238,3 +238,30 @@ async def pause_flow(name: str) -> Any:
|
||||
async def resume_flow(name: str) -> Any:
|
||||
"""Let a paused flow carry on, running whatever was held back."""
|
||||
return await _call("POST", f"/flows/{name}/resume")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def cancel_node(name: str, node_id: str) -> Any:
|
||||
"""Stop a node's code while it is running. Nothing to stop is not an error."""
|
||||
return await _call("POST", f"/flows/{name}/nodes/{node_id}/cancel")
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Modules
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_modules() -> Any:
|
||||
"""The python packages node code can import, and the manifest asking for them."""
|
||||
return await _call("GET", "/modules/")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def apply_modules(requirements: str) -> Any:
|
||||
"""Install exactly these requirements, one pip line each.
|
||||
|
||||
This replaces the whole manifest: a package left out is uninstalled. A
|
||||
manifest that does not resolve changes nothing.
|
||||
"""
|
||||
return await _call("POST", "/modules/apply", json={"requirements": requirements})
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
"""The venv node code imports from is built and kept in step with a manifest."""
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.flow import modules
|
||||
from app.flow.store import FlowStore
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
shutil.which("uv") is None, reason="module management needs uv"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
directory = tmp_path / "user-venv"
|
||||
monkeypatch.setattr(modules, "VENV_DIR", directory)
|
||||
return directory
|
||||
|
||||
|
||||
def test_an_empty_manifest_gives_a_venv_of_its_own(venv: Path, tmp_path: Path):
|
||||
store = FlowStore(tmp_path / "flows")
|
||||
store.write_requirements("")
|
||||
modules.sync("")
|
||||
|
||||
assert Path(modules.venv_python()).exists()
|
||||
assert modules.venv_python().startswith(str(venv))
|
||||
|
||||
information = modules.info(store)
|
||||
assert information.applied is True
|
||||
assert information.python_version
|
||||
|
||||
|
||||
def test_a_manifest_already_applied_is_not_installed_again(
|
||||
venv: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
store = FlowStore(tmp_path / "flows")
|
||||
store.write_requirements("")
|
||||
modules.sync("")
|
||||
marker = (venv / ".applied").read_text()
|
||||
|
||||
# Recorded rather than refused: reconcile swallows what it fails on, so an
|
||||
# exception raised in here would never reach the test.
|
||||
calls: list[object] = []
|
||||
monkeypatch.setattr(subprocess, "run", lambda *args, **kwargs: calls.append(args))
|
||||
modules.reconcile(store)
|
||||
|
||||
assert calls == []
|
||||
assert (venv / ".applied").read_text() == marker
|
||||
|
||||
|
||||
def test_a_manifest_that_does_not_resolve_leaves_the_venv_alone(venv: Path):
|
||||
modules.sync("")
|
||||
ok, output = modules.sync("fluksio-no-such-package-anywhere==9.9.9")
|
||||
|
||||
assert ok is False
|
||||
assert output
|
||||
# The marker still describes the manifest that actually installed.
|
||||
assert (venv / ".applied").read_text() == modules._digest("")
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Python nodes run in a worker process, and stay there when things go wrong."""
|
||||
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
|
||||
from app.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pool() -> Iterator[PythonWorkerPool]:
|
||||
# One worker: a respawn is then provably the same slot coming back.
|
||||
worker_pool = PythonWorkerPool(python=sys.executable, size=1)
|
||||
worker_pool.start()
|
||||
yield worker_pool
|
||||
worker_pool.stop()
|
||||
|
||||
|
||||
def run(pool: PythonWorkerPool, code: str, node: str = "demo", **kwargs):
|
||||
return pool.run(
|
||||
"demo", node, code, kwargs, {"factor": 2}, f"demo.{node}", timeout=5
|
||||
)
|
||||
|
||||
|
||||
def test_a_node_returns_its_value_and_what_it_printed(pool, capsys):
|
||||
result = run(
|
||||
pool,
|
||||
"def process(value, params):\n"
|
||||
" print('seen', value)\n"
|
||||
" return {'out': value * params['factor']}\n",
|
||||
value=21,
|
||||
)
|
||||
assert result == {"out": 42}
|
||||
# The proxy writes them to stdout, which is where the engine's tee is.
|
||||
assert "seen 21" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_a_failure_keeps_its_class_and_points_at_the_node(pool):
|
||||
with pytest.raises(Exception) as caught:
|
||||
run(pool, "def process(params):\n raise ValueError('bad input')\n")
|
||||
|
||||
# The engine renders a node error as "<class>: <message>", so both have to
|
||||
# survive the trip.
|
||||
assert type(caught.value).__name__ == "ValueError"
|
||||
assert str(caught.value) == "bad input"
|
||||
assert "<node demo." in caught.value.remote_traceback
|
||||
assert "ValueError: bad input" in caught.value.remote_traceback
|
||||
|
||||
|
||||
def test_a_node_that_kills_its_worker_is_an_ordinary_error(pool):
|
||||
with pytest.raises(Exception, match="worker died"):
|
||||
run(pool, "import os\n\n\ndef process(params):\n os._exit(1)\n")
|
||||
|
||||
assert run(pool, "def process(params):\n return {'out': 1}\n") == {"out": 1}
|
||||
|
||||
|
||||
def test_a_node_that_runs_too_long_is_killed_and_the_pool_recovers(pool):
|
||||
started = time.monotonic()
|
||||
with pytest.raises(NodeTimeout):
|
||||
pool.run(
|
||||
"demo",
|
||||
"slow",
|
||||
"import time\n\n\ndef process(params):\n time.sleep(30)\n",
|
||||
{},
|
||||
{},
|
||||
"demo.slow",
|
||||
timeout=1,
|
||||
)
|
||||
assert time.monotonic() - started < 10
|
||||
|
||||
# The killed worker's slot is refilled on the next call.
|
||||
assert run(pool, "def process(params):\n return {'out': 2}\n") == {"out": 2}
|
||||
|
||||
|
||||
def test_a_running_node_can_be_cancelled(pool):
|
||||
def stop_it() -> None:
|
||||
for _ in range(100):
|
||||
if pool.cancel("demo.slow"):
|
||||
return
|
||||
time.sleep(0.05)
|
||||
|
||||
stopper = threading.Thread(target=stop_it)
|
||||
stopper.start()
|
||||
try:
|
||||
with pytest.raises(NodeCancelled):
|
||||
pool.run(
|
||||
"demo",
|
||||
"slow",
|
||||
"import time\n\n\ndef process(params):\n time.sleep(30)\n",
|
||||
{},
|
||||
{},
|
||||
"demo.slow",
|
||||
timeout=30,
|
||||
)
|
||||
finally:
|
||||
stopper.join()
|
||||
|
||||
|
||||
def test_a_result_that_is_not_json_is_refused(pool):
|
||||
with pytest.raises(Exception, match="cannot be sent back as JSON"):
|
||||
run(pool, "def process(params):\n return {'out': {1, 2}}\n")
|
||||
|
||||
|
||||
def test_compiling_reports_where_the_source_is_wrong(pool):
|
||||
assert pool.compile("demo", "broken", "def process(params)\n return {}\n")
|
||||
assert pool.compile("demo", "fine", "def process(params):\n return {}\n") is None
|
||||
Reference in New Issue
Block a user