stroblmeandClaude Opus 5
6ff56533f5
The engine answered "where does this node run" twice, in two ways that could not see each other: a device sent it to a worker carrying that label, and resources were counted against the engine's own cores. Declaring both meant the second answer won and nothing was counted at all — which the data-science getting-started page and the worked example both do. One question now, in flow/placement.py: of every machine attached, which could grant what this node asked for, and which of those has it free. The books move onto each machine — one accountant per worker, built from the inventory it reported — and the waiting moves above them, where one condition variable can be woken by a release anywhere or by a worker attaching. Locks go one way: placer, then a machine's books, never back. So a node asking for a card now finds the box that has one, rather than being clamped down to none and run here. When nothing can grant the ask at all it is still cut down and run — a flow written on a cluster has to work on a laptop — but the ceiling is one real machine now, since taking the largest of each dimension separately can describe a machine nobody has. Two things fixed on the way. A device on a connector node held every batch run of its flow forever, waiting for a worker that could never run an entry point. And `prefer` falling back to the engine skipped the books, so the fallback held nothing. The bench flow's node has taken a `params` argument that with_settings has not forwarded for some time, so the benchmark could not run at all: 62 ms median submit-to-result with this, against the 61 ms on record. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
492 lines
18 KiB
Python
492 lines
18 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 dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from fluksio.flow.nodes.base import NodeOutputError
|
|
from fluksio.flow.resources import ResourceAccountant
|
|
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 = 2
|
|
|
|
|
|
class NoWorker(RemoteError):
|
|
"""Nothing is attached that carries the label this node asked for."""
|
|
|
|
|
|
def _whole(value: Any, default: int) -> int:
|
|
"""An int from whatever a worker sent, or the default when it was junk."""
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkerInventory:
|
|
"""What an attached worker says it holds.
|
|
|
|
The defaults are a machine that runs one thing at a time and has no GPU,
|
|
so a worker reporting nothing is scheduled by its label alone — which is
|
|
how every worker was scheduled before any of them reported anything.
|
|
"""
|
|
|
|
cpus: int = 1
|
|
gpus: int = 0
|
|
ram_mb: int | None = None
|
|
|
|
@classmethod
|
|
def from_hello(cls, payload: Any) -> WorkerInventory:
|
|
"""Read what is understood and ignore the rest.
|
|
|
|
Unknown members are deliberately not an error: a worker that learns to
|
|
report something new must not need this engine taught about it first.
|
|
"""
|
|
if not isinstance(payload, dict):
|
|
return cls()
|
|
return cls(
|
|
cpus=max(1, _whole(payload.get("cpus"), 1)),
|
|
gpus=max(0, _whole(payload.get("gpus"), 0)),
|
|
ram_mb=_whole(payload.get("ram_mb"), 0) or None,
|
|
)
|
|
|
|
|
|
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,
|
|
inventory: WorkerInventory | None = None,
|
|
) -> None:
|
|
self.name = name
|
|
self.labels = set(labels)
|
|
self.info = info or {}
|
|
self.inventory = inventory or WorkerInventory()
|
|
# This machine's books, held here so that a worker reconnecting — which
|
|
# is a new object — starts from a clean set rather than from counters
|
|
# the engine kept for a socket that is gone.
|
|
self.accountant = ResourceAccountant(
|
|
cpus=self.inventory.cpus,
|
|
gpus=self.inventory.gpus,
|
|
ram_mb=self.inventory.ram_mb,
|
|
)
|
|
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, on_change: Callable[[], None] | None = None) -> None:
|
|
self._workers: dict[str, RemoteWorker] = {}
|
|
self._lock = threading.Lock()
|
|
#: Told when the set of attached workers changes, so whoever is waiting
|
|
#: for a machine can look again. Called outside the lock: it takes one
|
|
#: of its own, and the two are only ever taken in that order.
|
|
self.on_change = on_change
|
|
|
|
# -------------------------------------------------------------------------
|
|
# 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)
|
|
)
|
|
self._changed()
|
|
|
|
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)
|
|
self._changed()
|
|
|
|
def _changed(self) -> None:
|
|
if self.on_change is not None:
|
|
self.on_change()
|
|
|
|
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,
|
|
env: dict[str, str] | None = None,
|
|
) -> Any:
|
|
worker = self.pick(label)
|
|
if worker is None:
|
|
raise NoWorker(f"no worker labelled '{label}' is attached")
|
|
return self.run_on(
|
|
worker,
|
|
flow,
|
|
node,
|
|
source,
|
|
kwargs,
|
|
node_id,
|
|
timeout,
|
|
run_id=run_id,
|
|
on_event=on_event,
|
|
env=env,
|
|
)
|
|
|
|
def run_on(
|
|
self,
|
|
worker: RemoteWorker,
|
|
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,
|
|
env: dict[str, str] | None = None,
|
|
) -> Any:
|
|
"""Run this node on the worker the caller has already settled on.
|
|
|
|
The placer holds that machine's cores for the length of the call, so
|
|
resolving the label a second time here could send the work somewhere
|
|
else and leave the claim on a machine doing nothing.
|
|
"""
|
|
payload = {
|
|
"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,
|
|
}
|
|
if env:
|
|
# What this call was allocated: the thread caps and the devices the
|
|
# node may see. The worker starts a process per call, so it applies
|
|
# them where a library still reads them — before the import.
|
|
payload["env"] = env
|
|
response = worker.request(
|
|
payload,
|
|
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,
|
|
env: dict[str, str] | 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,
|
|
env=env,
|
|
)
|
|
|
|
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
|