stroblmeandClaude Opus 5
1a9753fa9d
A worker reported its labels and nothing about the machine behind them, so the engine could route a node to a GPU box but not tell whether that box had a GPU free. Inventory — cores, GPUs, memory — now arrives with the hello frame, and the run frame carries back what the engine allocated for that call. Which is protocol 2 on both ends. GPUs are never probed: asking a vendor tool would make the one dependency two, so a GPU is what the batch job says it was given or what --gpus says. A worker that reports nothing still attaches and is scheduled by its label alone. Two things a job scheduler needs: --max-idle stops a worker started for one job rather than letting it hold its allocation to the walltime, and a refusal is now fatal instead of a reconnect loop that reads as a hang in a job's log. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
441 lines
16 KiB
Python
441 lines
16 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.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()
|
|
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,
|
|
env: dict[str, str] | None = None,
|
|
) -> Any:
|
|
worker = self.pick(label)
|
|
if worker is None:
|
|
raise NoWorker(f"no worker labelled '{label}' is attached")
|
|
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
|