stroblmeandClaude Opus 5
608d30d884
Five concurrent training nodes, each sizing its thread pool to every core,
left the engine's own event loop unscheduled: the API stopped answering
within 10 s and every client died. The same shape on a GPU deadlocked a run
for 21 minutes at 0% utilisation with nothing failing and nothing to read --
it just sat in `running`.
@node(resources={"cpus": 2}) is the declaration. The engine holds that much
for the length of the execution, so more of them than the machine has room
for wait their turn rather than oversubscribing it, and a `gpus` node holds
its card exclusively. FLOW_CPUS defaults to every core but two, and those two
are what keeps the engine answering.
Because a thread cap is read when the process imports the library, a warm
worker cannot be told a different one -- so an environment gets a pool of its
own and nodes deriving the same one share it, rather than paying a cold start
per call on exactly the nodes whose imports are slowest. XLA_FLAGS is never
derived: it is a composed, version-dependent string, so it travels in
resources.env where it is visible.
A node that declares nothing is not accounted for and behaves as it always
did -- it just gets FLOW_CPUS/FLOW_MAX_WORKERS as a thread cap, which is the
half of this that fixes the reported incident without anybody declaring
anything. An operator who set OMP_NUM_THREADS themselves still wins.
Resources are claimed strictly before a worker slot, so the two blocking
waits cannot deadlock. A node queued for them publishes node_queued and shows
on GET /workers/resources, because waiting and hanging looked identical.
Accounted, not enforced: no cgroups, no rlimits. Scheduling across machines,
flavours and enforcement are the next steps.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
643 lines
25 KiB
Python
643 lines
25 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 hashlib
|
|
import itertools
|
|
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
|
|
|
|
import orjson
|
|
from fluksio_worker import worker_main as _worker_main
|
|
|
|
from fluksio.flow.events import EventBus
|
|
from fluksio.flow.nodes.base import NodeOutputError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: The runner, in the worker distribution — the same file a remote worker runs,
|
|
#: so a node cannot tell which kind of worker it is on.
|
|
WORKER_MAIN = Path(_worker_main.__file__)
|
|
|
|
#: 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
|
|
|
|
#: How often to wake up while waiting on a node with no timeout. Nothing is
|
|
#: checked on a schedule — a worker that dies closes its pipe and wakes the
|
|
#: read immediately — so this is only here to notice a pipe held open by
|
|
#: something that outlived the worker it belonged to.
|
|
POLL_S = 30.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 = ("DATABASE_URL", "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."""
|
|
|
|
|
|
#: How many distinct exception names to keep classes for. The key is a name
|
|
#: node code chose, so an unbounded cache is a node raising a class named after
|
|
#: its loop counter.
|
|
MAX_REMOTE_TYPES = 256
|
|
|
|
_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 not None:
|
|
return cls
|
|
cls = type(name, (RemoteError,), {})
|
|
# Past the cap the class is still built, just not kept: the names already
|
|
# here are the ones a deployment actually raises, and evicting them to make
|
|
# room for generated ones would be the wrong way round.
|
|
if len(_remote_types) < MAX_REMOTE_TYPES:
|
|
_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
|
|
# Mirrors the worker's own module cache: (flow, node) -> source digest.
|
|
# What it is for is knowing whether a call has to pay for imports
|
|
# before the node's clock starts; see ``PythonWorkerPool._warm``.
|
|
self.loaded: dict[tuple[str, str], str] = {}
|
|
# 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
|
|
# Node inputs may be keyed by something other than a string, which
|
|
# the stdlib encoder this replaced turned into strings rather than
|
|
# refusing.
|
|
self.proc.stdin.write(
|
|
orjson.dumps(request, option=orjson.OPT_NON_STR_KEYS) + b"\n"
|
|
)
|
|
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 {}
|
|
# Pools for nodes that declared resources, keyed by the environment
|
|
# their allocation derived. Empty on all but the one pool the engine
|
|
# builds; see ``for_env``.
|
|
self._children: dict[frozenset[tuple[str, str]], PythonWorkerPool] = {}
|
|
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._requests = itertools.count(1)
|
|
# Not a lock over the pool — the slot queue is what serialises calls.
|
|
# It covers the two places where bookkeeping must not interleave: a
|
|
# cancel taking a worker out of ``_running`` while ``_request`` is
|
|
# handing that same worker back to a slot, and two respawns draining
|
|
# the idle queue at once. Registering into ``_running``, the generation
|
|
# counter and the ``_stopped`` flag are deliberately outside it — each
|
|
# is one atomic operation, and a cancel that misses a node by a
|
|
# microsecond is a cancel that arrived a microsecond early.
|
|
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 for_env(self, extra: dict[str, str]) -> PythonWorkerPool:
|
|
"""The pool whose workers were started with this environment.
|
|
|
|
A thread cap or a GPU assignment is read once, when the process
|
|
imports the library, so a worker that is already warm cannot be told a
|
|
different one — and retiring a worker per call would give back the cold
|
|
starts ``_warm`` exists to avoid, on exactly the nodes whose imports
|
|
are slowest. So an environment gets a pool of its own instead, and
|
|
nodes deriving the same one share it.
|
|
|
|
How many of them may run at once is the accountant's business, not this
|
|
pool's: the slots here only bound one environment's own concurrency.
|
|
"""
|
|
if not extra:
|
|
return self
|
|
key = frozenset(extra.items())
|
|
with self._lock:
|
|
child = self._children.get(key)
|
|
if child is None:
|
|
child = PythonWorkerPool(
|
|
self.python,
|
|
size=self.size,
|
|
events=self.events,
|
|
env={**self.env, **extra},
|
|
)
|
|
child._generation = self._generation
|
|
child.start()
|
|
self._children[key] = child
|
|
return child
|
|
|
|
def stop(self) -> None:
|
|
self._stopped = True
|
|
self._generation += 1
|
|
for child in list(self._children.values()):
|
|
child.stop()
|
|
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)
|
|
children = list(self._children.values())
|
|
for child in children:
|
|
child.respawn_all()
|
|
|
|
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:
|
|
if payload["op"] == "run":
|
|
refused = self._warm(worker, payload)
|
|
if refused is not None:
|
|
return refused
|
|
return self._exchange(worker, payload, timeout, on_event)
|
|
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 _warm(self, worker: _Worker, payload: dict[str, Any]) -> dict[str, Any] | None:
|
|
"""Load the node's source in this worker before its clock starts.
|
|
|
|
A node's timeout is what its *body* may take. Importing torch on a cold
|
|
worker is not the body, and charging it to the same budget cannot be
|
|
won: the timeout kills the worker, so the retry is cold again and pays
|
|
the imports from the start. ``compile`` warms only the one worker it
|
|
happened to land on, and a pool has several.
|
|
|
|
Returns the failed reply when the source does not load, so the caller
|
|
raises the author's own error rather than a timeout.
|
|
"""
|
|
key = (str(payload["flow"]), str(payload["node"]))
|
|
digest = hashlib.md5(str(payload.get("source") or "").encode()).hexdigest()
|
|
if worker.loaded.get(key) == digest:
|
|
return None
|
|
response = self._exchange(worker, {**payload, "op": "compile"}, COMPILE_TIMEOUT)
|
|
if not response.get("ok"):
|
|
return response
|
|
worker.loaded[key] = digest
|
|
return None
|
|
|
|
def _exchange(
|
|
self,
|
|
worker: _Worker,
|
|
payload: dict[str, Any],
|
|
timeout: float,
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""One request down the pipe, and the reply that belongs to it."""
|
|
request_id = next(self._requests)
|
|
try:
|
|
worker.send({**payload, "id": request_id})
|
|
except OSError as exc:
|
|
# A broken pipe is the truthful signal here: poll() still reports a
|
|
# child that has exited but not yet been reaped as alive, so
|
|
# _release would hand this corpse to the next call. Retire it.
|
|
worker.cancelled = True
|
|
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 timeout > 0 else POLL_S)
|
|
)
|
|
if line:
|
|
try:
|
|
message = dict(orjson.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 NodeOutputError:
|
|
# An emission on a port the node never declared.
|
|
# The call is already wrong, and its generator has
|
|
# more frames coming down this pipe — so retire the
|
|
# worker and let the author see what they emitted.
|
|
worker.cancelled = True
|
|
worker.kill()
|
|
raise
|
|
except Exception:
|
|
logger.exception("Could not record a worker event")
|
|
continue
|
|
if message.get("id") != request_id:
|
|
# The pipe is a call behind. Returning this would hand one
|
|
# node another node's result, which is worse than any
|
|
# failure — so say so and retire the worker rather than
|
|
# log a warning nobody reads.
|
|
worker.cancelled = True
|
|
raise RemoteError(
|
|
f"the worker answered request {message.get('id')!r} "
|
|
f"while {request_id} was outstanding"
|
|
)
|
|
return message
|
|
if worker.cancelled:
|
|
raise NodeCancelled("cancelled while it was running")
|
|
if line is None:
|
|
if timeout > 0:
|
|
worker.kill()
|
|
raise NodeTimeout(f"was silent for {timeout}s and was killed")
|
|
# No limit, so silence is the node working. What ends the call
|
|
# is the worker dying, and that arrives as the pipe closing
|
|
# rather than as a deadline — unless something else inherited
|
|
# the write end and is holding it open, which is what this
|
|
# asks about.
|
|
if worker.alive():
|
|
continue
|
|
worker.cancelled = True
|
|
raise RemoteError("worker died")
|
|
worker.cancelled = True
|
|
raise RemoteError("worker died")
|
|
|
|
def compile(
|
|
self, flow: str, node: str, source: str, keep: bool = True
|
|
) -> str | None:
|
|
"""Load this source in a worker. Returns what to show the author, or None.
|
|
|
|
``keep`` off is a question rather than a load: the source is compiled
|
|
to see whether it would run and then dropped. That is what a draft
|
|
wants — caching it would evict the published source the node's next
|
|
call is about to need.
|
|
"""
|
|
try:
|
|
response = self._request(
|
|
{
|
|
"op": "compile",
|
|
"flow": flow,
|
|
"node": node,
|
|
"source": source,
|
|
"keep": keep,
|
|
},
|
|
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))
|
|
children = list(self._children.values())
|
|
if worker is not None:
|
|
worker.cancelled = True
|
|
worker.kill()
|
|
return True
|
|
# A node with declared resources runs in a pool of its own, and it is
|
|
# the same node to whoever pressed cancel.
|
|
return any(child.cancel(node_id, run_id) for child in children)
|
|
|
|
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()
|
|
children = list(self._children.values())
|
|
return len(workers) + sum(child.cancel_run(run_id) for child in children)
|
|
|
|
def _publish(self, event: dict[str, Any]) -> None:
|
|
if self.events is not None:
|
|
self.events.publish(event)
|