stroblmeandClaude Opus 5
93374a310e
Four things the python SDK turned up, each fixed where every client sees it. A key no port declares is now an error rather than a silent drop, on the return, the yield and the emit alike — the contract the docs already stated. The SDK reads literal yields at sync time, so a typo fails before anything runs, and an emission of one fails the call rather than being logged where nobody looks. NaN and infinity are refused at the port. JSON cannot spell either, so one that travelled came back as a 500, a socket frame that stopped the canvas, or a metric batch the database dropped whole. An artifact input takes `@run:<id>.<output>` or a bare digest, resolved on the engine — so the CLI, the run dialog and a python caller mean the same thing, and a sweep can pass one at all. Node timeouts are off by default. The clock measured silence, which a training node is full of, and remote workers had already stopped enforcing it — their heartbeat reset it. Now a heartbeat proves the agent rather than the node, ninety seconds of nothing fails the call either way, and the engine touches work it is still running so a long node is not redelivered at sixty seconds. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019V5bsYGNxcgPs4xXmTPx69
392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""Workers on other hosts, reached over a socket they opened themselves.
|
|
|
|
The engine runs where the automations are; a GPU sits somewhere else. Those
|
|
two are usually not on the same network, and the one that can be dialled is
|
|
the engine — so a worker connects *out* to it and the engine answers over the
|
|
connection it was given. That also means nothing has to expose Redis, which is
|
|
the thing a remote worker must never be handed.
|
|
|
|
What travels is the protocol the local worker pool already speaks: one JSON
|
|
object per line becomes one JSON frame, the node's source rides along with
|
|
every call so no code has to be distributed, and the reports a node makes
|
|
while it runs arrive the same way they do over a pipe. A node cannot tell
|
|
which kind of worker it is running on, which is the point — the same flow runs
|
|
in both places.
|
|
|
|
The awkward part is that the socket lives on the event loop and a node
|
|
executes on a worker thread. A call therefore hands its frame to the loop and
|
|
blocks on a queue of its own until the loop puts the answer there.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import logging
|
|
import queue
|
|
import sys
|
|
import threading
|
|
import time
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from fluksio.flow.nodes.base import NodeOutputError
|
|
from fluksio.flow.workers import NodeTimeout, RemoteError, _remote_class
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: How long the loop is given to accept a frame we are handing it.
|
|
SEND_TIMEOUT_S = 30.0
|
|
#: A worker that has said nothing at all for this long — not even a heartbeat —
|
|
#: is treated as gone. It beats every ten seconds while it is executing, so this
|
|
#: catches a dead socket rather than a slow node, and it is what bounds a call
|
|
#: whose node has no timeout of its own.
|
|
SILENCE_S = 90.0
|
|
#: Protocol version this engine speaks. A worker announcing anything else is
|
|
#: refused rather than half-understood.
|
|
PROTOCOL = 1
|
|
|
|
|
|
class NoWorker(RemoteError):
|
|
"""Nothing is attached that carries the label this node asked for."""
|
|
|
|
|
|
class RemoteWorker:
|
|
"""One attached worker, and the calls it has in flight."""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
labels: list[str],
|
|
send: Callable[[dict[str, Any]], Any],
|
|
loop: asyncio.AbstractEventLoop,
|
|
max_parallel: int = 1,
|
|
info: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self.name = name
|
|
self.labels = set(labels)
|
|
self.info = info or {}
|
|
self.attached_at = time.time()
|
|
self.last_seen = time.time()
|
|
self._send = send
|
|
self._loop = loop
|
|
self._slots = threading.Semaphore(max_parallel)
|
|
self.max_parallel = max_parallel
|
|
self._pending: dict[str, queue.Queue[dict[str, Any] | None]] = {}
|
|
self._lock = threading.Lock()
|
|
self._gone = False
|
|
# Source this worker has already loaded, by digest. A rebuild asks
|
|
# every device-bound node whether it compiles, and without this that
|
|
# is a network round trip each time — which is how attaching a GPU box
|
|
# makes rebuilding the flows slow enough to fail a health check. A
|
|
# worker that reconnects is a new object, so this empties with it.
|
|
self._compiled: set[str] = set()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# From the socket's side, on the event loop
|
|
# -------------------------------------------------------------------------
|
|
|
|
def deliver(self, message: dict[str, Any]) -> None:
|
|
"""Hand a frame to whichever call is waiting for it."""
|
|
self.last_seen = time.time()
|
|
call_id = str(message.get("call_id") or "")
|
|
with self._lock:
|
|
inbox = self._pending.get(call_id)
|
|
if inbox is not None:
|
|
inbox.put(message)
|
|
|
|
def detach(self) -> None:
|
|
"""The socket closed: wake everything still waiting on it."""
|
|
self._gone = True
|
|
with self._lock:
|
|
inboxes = list(self._pending.values())
|
|
for inbox in inboxes:
|
|
# None is "no more answers are coming", which the caller turns into
|
|
# a failed node rather than a wait that never ends.
|
|
inbox.put(None)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# From a node's side, on a worker thread
|
|
# -------------------------------------------------------------------------
|
|
|
|
def request(
|
|
self,
|
|
payload: dict[str, Any],
|
|
timeout: float,
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> dict[str, Any]:
|
|
if self._gone:
|
|
raise RemoteError(f"worker '{self.name}' is no longer attached")
|
|
call_id = str(payload["call_id"])
|
|
inbox: queue.Queue[dict[str, Any] | None] = queue.Queue()
|
|
# Blocking here is the backpressure, exactly as taking a slot is in the
|
|
# local pool.
|
|
self._slots.acquire()
|
|
with self._lock:
|
|
self._pending[call_id] = inbox
|
|
try:
|
|
future = asyncio.run_coroutine_threadsafe(self._send(payload), self._loop)
|
|
future.result(timeout=SEND_TIMEOUT_S)
|
|
|
|
# The node's own deadline measures silence, so a node reporting its
|
|
# progress is never mistaken for a hung one. A heartbeat is not
|
|
# progress: it says the agent is alive, which is what SILENCE_S
|
|
# asks, and says nothing about the node — so it feeds the liveness
|
|
# bound below and never the node's own.
|
|
deadline = time.monotonic() + timeout if timeout > 0 else None
|
|
while True:
|
|
wait = SILENCE_S
|
|
if deadline is not None:
|
|
wait = min(SILENCE_S, deadline - time.monotonic())
|
|
try:
|
|
message = inbox.get(timeout=max(wait, 0.0))
|
|
except queue.Empty:
|
|
self.cancel(call_id)
|
|
if deadline is not None and time.monotonic() >= deadline:
|
|
raise NodeTimeout(
|
|
f"'{self.name}' was silent for {timeout}s"
|
|
) from None
|
|
raise RemoteError(
|
|
f"worker '{self.name}' sent nothing for "
|
|
f"{SILENCE_S:.0f}s and is presumed gone"
|
|
) from None
|
|
if message is None:
|
|
raise RemoteError(f"worker '{self.name}' went away mid-call")
|
|
kind = message.get("event")
|
|
if kind == "heartbeat":
|
|
continue
|
|
if kind:
|
|
if on_event is not None:
|
|
try:
|
|
on_event(message)
|
|
except NodeOutputError:
|
|
# A port the node never declared. Stop the call
|
|
# rather than let the rest of its emissions arrive.
|
|
self.cancel(call_id)
|
|
raise
|
|
except Exception:
|
|
logger.exception("Could not record a worker event")
|
|
if deadline is not None:
|
|
deadline = time.monotonic() + timeout
|
|
continue
|
|
return message
|
|
except Exception as exc:
|
|
if isinstance(exc, (NodeTimeout, RemoteError, NodeOutputError)):
|
|
raise
|
|
raise RemoteError(f"worker '{self.name}': {exc}") from exc
|
|
finally:
|
|
with self._lock:
|
|
self._pending.pop(call_id, None)
|
|
self._slots.release()
|
|
|
|
@property
|
|
def compiled(self) -> set[str]:
|
|
"""Digests of the source this worker has already loaded."""
|
|
return self._compiled
|
|
|
|
@property
|
|
def gone(self) -> bool:
|
|
return self._gone
|
|
|
|
@property
|
|
def in_flight(self) -> int:
|
|
"""Calls this worker has not answered yet."""
|
|
with self._lock:
|
|
return len(self._pending)
|
|
|
|
def calls_of(self, run_id: str) -> list[str]:
|
|
with self._lock:
|
|
return [call for call in self._pending if call.startswith(f"{run_id}:")]
|
|
|
|
def cancel(self, call_id: str) -> None:
|
|
"""Ask the worker to kill what it is running for this call."""
|
|
try:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self._send({"op": "cancel", "call_id": call_id}), self._loop
|
|
).result(timeout=SEND_TIMEOUT_S)
|
|
except Exception:
|
|
logger.warning("Could not cancel '%s' on '%s'", call_id, self.name)
|
|
|
|
|
|
class RemoteWorkerHub:
|
|
"""Every attached worker, and which of them a node may run on."""
|
|
|
|
def __init__(self) -> None:
|
|
self._workers: dict[str, RemoteWorker] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Attachment
|
|
# -------------------------------------------------------------------------
|
|
|
|
def attach(self, worker: RemoteWorker) -> None:
|
|
with self._lock:
|
|
existing = self._workers.get(worker.name)
|
|
if existing is not None:
|
|
# A worker that reconnects after a network drop: the old socket
|
|
# is dead whether or not it has noticed yet.
|
|
existing.detach()
|
|
self._workers[worker.name] = worker
|
|
logger.info(
|
|
"Worker '%s' attached with labels %s", worker.name, sorted(worker.labels)
|
|
)
|
|
|
|
def detach(self, name: str) -> None:
|
|
with self._lock:
|
|
worker = self._workers.pop(name, None)
|
|
if worker is not None:
|
|
worker.detach()
|
|
logger.info("Worker '%s' detached", name)
|
|
|
|
def workers(self) -> list[RemoteWorker]:
|
|
with self._lock:
|
|
return list(self._workers.values())
|
|
|
|
def labels(self) -> set[str]:
|
|
"""Every label something attached right now carries."""
|
|
with self._lock:
|
|
return {
|
|
label for worker in self._workers.values() for label in worker.labels
|
|
}
|
|
|
|
def pick(self, label: str) -> RemoteWorker | None:
|
|
"""A worker carrying this label, least busy first.
|
|
|
|
Resolved per call rather than when the flow was built, so a worker that
|
|
attaches after a run was submitted picks the work up without anything
|
|
being rebuilt.
|
|
"""
|
|
with self._lock:
|
|
candidates = [
|
|
worker
|
|
for worker in self._workers.values()
|
|
if not worker.gone and (label == worker.name or label in worker.labels)
|
|
]
|
|
if not candidates:
|
|
return None
|
|
return min(candidates, key=lambda worker: worker.in_flight)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Running a node on one
|
|
# -------------------------------------------------------------------------
|
|
|
|
def run(
|
|
self,
|
|
label: str,
|
|
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:
|
|
worker = self.pick(label)
|
|
if worker is None:
|
|
raise NoWorker(f"no worker labelled '{label}' is attached")
|
|
response = worker.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,
|
|
},
|
|
timeout=timeout,
|
|
on_event=on_event,
|
|
)
|
|
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 compile(
|
|
self, label: str, flow: str, node: str, source: str, timeout: float = 60.0
|
|
) -> str | None:
|
|
"""Load this source on the worker that will run it.
|
|
|
|
Which machine compiles matters: a node importing torch is fine on the
|
|
GPU box and a ``ModuleNotFoundError`` on the engine, so checking it
|
|
here would fail a node that is perfectly correct. When nothing is
|
|
attached there is nothing to check against, and ``None`` says so — the
|
|
node is not broken, it is waiting.
|
|
"""
|
|
worker = self.pick(label)
|
|
if worker is None:
|
|
return None
|
|
digest = hashlib.md5(f"{flow}.{node}:{source}".encode()).hexdigest()
|
|
if digest in worker.compiled:
|
|
return None
|
|
try:
|
|
response = worker.request(
|
|
{
|
|
"op": "compile",
|
|
"call_id": f"compile:{flow}.{node}",
|
|
"flow": flow,
|
|
"node": node,
|
|
"source": source,
|
|
},
|
|
timeout=timeout,
|
|
)
|
|
except RemoteError as exc:
|
|
return f"{type(exc).__name__}: {exc}"
|
|
if response.get("ok"):
|
|
worker.compiled.add(digest)
|
|
return None
|
|
error = response.get("error") or {}
|
|
return str(error.get("short") or "The node could not be loaded.")
|
|
|
|
def proxy(
|
|
self,
|
|
label: str,
|
|
flow: str,
|
|
node: str,
|
|
source: str,
|
|
node_id: str,
|
|
timeout: float,
|
|
run_id: str = "",
|
|
on_event: Callable[[dict[str, Any]], None] | None = None,
|
|
fallback: Callable[..., Any] | None = None,
|
|
) -> Callable[..., Any]:
|
|
"""What a node with a device runs instead of its own function.
|
|
|
|
``fallback`` is the local pool's proxy, used when the node only prefers
|
|
the label rather than requiring it.
|
|
"""
|
|
|
|
def call(**kwargs: Any) -> Any:
|
|
if fallback is not None and self.pick(label) is None:
|
|
return fallback(**kwargs)
|
|
return self.run(
|
|
label,
|
|
flow,
|
|
node,
|
|
source,
|
|
kwargs,
|
|
node_id,
|
|
timeout,
|
|
run_id=run_id,
|
|
on_event=on_event,
|
|
)
|
|
|
|
return call
|
|
|
|
def cancel_run(self, run_id: str) -> int:
|
|
"""Kill whatever this run has executing on any attached worker."""
|
|
stopped = 0
|
|
for worker in self.workers():
|
|
for call in worker.calls_of(run_id):
|
|
worker.cancel(call)
|
|
stopped += 1
|
|
return stopped
|