Rename the import package app to fluksio
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>
This commit is contained in:
Claude Opus 5
@@ -0,0 +1,369 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user