stroblmeandClaude Opus 5
640654bd66
A wheel whose top-level module is `app` collides with anything else in a user's venv, so the package that is about to be published takes the name it is published under. Only the Python package moves; the repo, the Docker WORKDIR and the compose project keep theirs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
473 lines
17 KiB
Python
473 lines
17 KiB
Python
"""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 fluksio.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
|
|
|
|
#: Environment the worker is not given. ``SECRET_KEY`` decrypts every stored
|
|
#: secret, not only the ones bound to the node asking.
|
|
ENV_DENY_PREFIXES = ("POSTGRES_", "FIRST_SUPERUSER", "SENTRY_DSN")
|
|
ENV_DENY_WORDS = ("PASSWORD", "SECRET", "TOKEN")
|
|
|
|
|
|
def worker_env(extra: dict[str, str] | None = None) -> dict[str, str]:
|
|
"""The engine's environment with its credentials taken out.
|
|
|
|
A denylist, so ``PATH``, ``HOME``, the locale and whatever else a venv
|
|
needs still reach the worker. This is not a sandbox: the worker is the same
|
|
user in the same container, with the same filesystem and the same network.
|
|
It only means node code cannot read the deployment's secrets straight out
|
|
of its own environment.
|
|
|
|
``extra`` is put back afterwards, deliberately: the artifact store's
|
|
address is something the worker has to be told, and one of its names would
|
|
otherwise be caught by the denylist for containing "TOKEN".
|
|
"""
|
|
env = {
|
|
key: value
|
|
for key, value in os.environ.items()
|
|
if not key.startswith(ENV_DENY_PREFIXES)
|
|
and not any(word in key.upper() for word in ENV_DENY_WORDS)
|
|
}
|
|
env.update(extra or {})
|
|
return env
|
|
|
|
|
|
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, env: dict[str, str] | None = None
|
|
) -> None:
|
|
self.generation = generation
|
|
self.cancelled = False
|
|
# What a read took past the end of a line. A node reporting quickly
|
|
# puts several lines in one chunk, and the next one is the caller's to
|
|
# read — dropping it would lose a metric, returning it with the first
|
|
# would be unparseable.
|
|
self._buffer = bytearray()
|
|
self.proc = subprocess.Popen(
|
|
[python, str(WORKER_MAIN)],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
close_fds=True,
|
|
env=worker_env(env),
|
|
)
|
|
|
|
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 line; ``None`` past the deadline, ``""`` if the pipe closed."""
|
|
assert self.proc.stdout is not None
|
|
fd = self.proc.stdout.fileno()
|
|
while True:
|
|
end = self._buffer.find(b"\n")
|
|
if end >= 0:
|
|
line = bytes(self._buffer[: end + 1])
|
|
del self._buffer[: end + 1]
|
|
return line.decode(errors="replace")
|
|
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
|
|
# line, which is no more use than none at all.
|
|
self._buffer.clear()
|
|
return ""
|
|
self._buffer += chunk
|
|
|
|
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,
|
|
env: dict[str, str] | None = None,
|
|
) -> None:
|
|
self.python = python
|
|
self.size = size
|
|
self.events = events
|
|
# Put into every worker's environment past the denylist — where the
|
|
# artifact store is, which node code needs and cannot guess.
|
|
self.env = env or {}
|
|
self._idle: queue.Queue[_Worker | None] = queue.Queue()
|
|
# Keyed by (run, node): a sweep has the same node executing in several
|
|
# runs at once, and cancelling one of them must not kill the others.
|
|
# The live pipeline's own executions carry an empty run.
|
|
self._running: dict[tuple[str, str], _Worker] = {}
|
|
self._generation = 0
|
|
self._stopped = False
|
|
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._stopped = True
|
|
self._generation += 1
|
|
for worker in list(self._running.values()):
|
|
worker.kill()
|
|
for slot in self._drain():
|
|
if slot is not None:
|
|
slot.kill()
|
|
# A node thread can be blocked on a slot right now, and the threads are
|
|
# not daemons — without a sentinel each it waits for a worker that is
|
|
# never coming back, and the interpreter cannot exit.
|
|
for _ in range(self.size):
|
|
self._idle.put(None)
|
|
|
|
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."""
|
|
if self._stopped:
|
|
raise RemoteError("the worker pool is shutting down")
|
|
slot = self._idle.get()
|
|
if self._stopped:
|
|
# The sentinel stop() put back, so this is a wake-up rather than a slot.
|
|
raise RemoteError("the worker pool is shutting down")
|
|
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, self.env)
|
|
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 = "",
|
|
run_id: str = "",
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> dict[str, Any]:
|
|
worker = self._acquire()
|
|
key = (run_id, node_id)
|
|
if node_id:
|
|
self._running[key] = 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,
|
|
"run": run_id,
|
|
"ts": time.time(),
|
|
}
|
|
)
|
|
try:
|
|
try:
|
|
worker.send(payload)
|
|
except OSError as exc:
|
|
raise RemoteError(f"worker died: {exc}") from exc
|
|
|
|
while True:
|
|
# Reset per line, so the timeout measures how long the node has
|
|
# been silent rather than how long it has been working. A node
|
|
# that reports nothing is still held to it, which is what keeps
|
|
# the deadline meaningful for the ones that never report.
|
|
line = worker.read_line(time.monotonic() + timeout)
|
|
if line:
|
|
try:
|
|
message = dict(json.loads(line))
|
|
except (TypeError, ValueError) as exc:
|
|
# A reply we cannot read leaves this worker out of step:
|
|
# whatever is still in its pipe would be taken by the
|
|
# next request on this slot as its own answer. Retire it.
|
|
worker.cancelled = True
|
|
raise RemoteError(
|
|
f"unreadable reply from the worker: {exc}"
|
|
) from exc
|
|
if message.get("event"):
|
|
if on_event is not None:
|
|
try:
|
|
on_event(message)
|
|
except Exception:
|
|
logger.exception("Could not record a worker event")
|
|
continue
|
|
return message
|
|
if worker.cancelled:
|
|
raise NodeCancelled("cancelled while it was running")
|
|
if line is None:
|
|
worker.kill()
|
|
raise NodeTimeout(f"was silent for {timeout}s and was killed")
|
|
raise RemoteError("worker died")
|
|
finally:
|
|
# Under the lock, so a cancel that has already read this worker out
|
|
# of _running cannot kill it after another node has taken the slot.
|
|
with self._lock:
|
|
if node_id:
|
|
self._running.pop(key, 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],
|
|
node_id: str,
|
|
timeout: float,
|
|
run_id: str = "",
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> Any:
|
|
response = self._request(
|
|
{
|
|
"op": "run",
|
|
"call_id": f"{run_id}:{node_id}" if run_id else node_id,
|
|
"flow": flow,
|
|
"node": node,
|
|
"source": source,
|
|
"kwargs": kwargs,
|
|
"run": {"id": run_id} if run_id else None,
|
|
},
|
|
timeout=timeout,
|
|
node_id=node_id,
|
|
run_id=run_id,
|
|
on_event=on_event,
|
|
)
|
|
# 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,
|
|
run_id: str = "",
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> Callable[..., Any]:
|
|
"""The callable a python node runs instead of its own compiled function.
|
|
|
|
A run builds its own nodes, so the run this proxy belongs to is bound
|
|
here rather than looked up — which is also what lets two runs of one
|
|
node be told apart when one of them is cancelled.
|
|
"""
|
|
|
|
def call(**kwargs: Any) -> Any:
|
|
return self.run(
|
|
flow,
|
|
node,
|
|
source,
|
|
kwargs,
|
|
node_id,
|
|
timeout,
|
|
run_id=run_id,
|
|
on_event=on_event,
|
|
)
|
|
|
|
return call
|
|
|
|
def cancel(self, node_id: str, run_id: str = "") -> bool:
|
|
"""Stop a node that is running now. False when there was nothing to stop."""
|
|
with self._lock:
|
|
worker = self._running.get((run_id, node_id))
|
|
if worker is None:
|
|
return False
|
|
worker.cancelled = True
|
|
worker.kill()
|
|
return True
|
|
|
|
def cancel_run(self, run_id: str) -> int:
|
|
"""Stop every node this run has in a worker right now."""
|
|
with self._lock:
|
|
workers = [
|
|
worker
|
|
for (owner, _node), worker in self._running.items()
|
|
if owner == run_id
|
|
]
|
|
for worker in workers:
|
|
worker.cancelled = True
|
|
worker.kill()
|
|
return len(workers)
|
|
|
|
def _publish(self, event: dict[str, Any]) -> None:
|
|
if self.events is not None:
|
|
self.events.publish(event)
|