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>
370 lines
13 KiB
Python
370 lines
13 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.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 for this long is treated as gone. It sends a
|
|
#: heartbeat while it is executing, so this only ever catches a dead socket.
|
|
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)
|
|
|
|
while True:
|
|
try:
|
|
# Reset per frame: the deadline measures silence, so a node
|
|
# reporting its progress is never mistaken for a hung one.
|
|
message = inbox.get(timeout=timeout)
|
|
except queue.Empty:
|
|
self.cancel(call_id)
|
|
raise NodeTimeout(
|
|
f"'{self.name}' was silent for {timeout}s"
|
|
) 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 Exception:
|
|
logger.exception("Could not record a worker event")
|
|
continue
|
|
return message
|
|
except Exception as exc:
|
|
if isinstance(exc, (NodeTimeout, RemoteError)):
|
|
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
|