Split the worker into a distribution of its own
A cluster or GPU host installs `pip install fluksio-worker` and gets the agent and the runner, not psycopg, numpy and the MCP SDK. The engine depends on it as a workspace member, so the file it launches node code with is the same file a remote worker runs — which is what keeps a node unable to tell the difference. Copying the two files by hand still works. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -0,0 +1,298 @@
|
||||
"""The agent that runs Fluksio nodes on a machine the engine cannot reach.
|
||||
|
||||
Point it at the engine on the box with the GPU and it dials in::
|
||||
|
||||
pip install fluksio-worker
|
||||
fluksio-worker --url wss://api.example.com/api/v1/workers/attach \\
|
||||
--token "$FLUKSIO_WORKER_TOKEN" --labels gpu --python /opt/venv/bin/python
|
||||
|
||||
It connects *out*, so the engine needs no route back and nothing has to expose
|
||||
Redis. What it then does is what the engine's own worker pool does: hold a few
|
||||
subprocesses running ``worker_main.py``, hand each call to one, and pass back
|
||||
everything that comes out — including the metrics a training loop reports
|
||||
while it is still running.
|
||||
|
||||
Deliberately two files with one dependency. Nothing of the engine is imported
|
||||
here; a worker host installs Python, ``websockets``, and whatever the nodes
|
||||
themselves need. Where pip is not an option, copying this file and
|
||||
``worker_main.py`` into one directory and running ``python agent.py`` is the
|
||||
same thing — the engine serves the runner at ``GET /api/v1/workers/runtime``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import websockets
|
||||
except ImportError: # pragma: no cover - the one dependency, named plainly
|
||||
print(
|
||||
"This needs the 'websockets' package: pip install websockets", file=sys.stderr
|
||||
)
|
||||
raise SystemExit(1) from None
|
||||
|
||||
log = logging.getLogger("fluksio-worker")
|
||||
|
||||
PROTOCOL = 1
|
||||
#: Sent while a call is running, so the engine can tell working from wedged.
|
||||
HEARTBEAT_S = 10.0
|
||||
#: Reconnection backs off to this and no further.
|
||||
MAX_BACKOFF_S = 30.0
|
||||
|
||||
|
||||
#: The user-code runner, beside this file — installed together, or copied
|
||||
#: together onto a host where pip is not an option.
|
||||
WORKER_MAIN = Path(__file__).resolve().with_name("worker_main.py")
|
||||
|
||||
|
||||
class Subprocess:
|
||||
"""One user-code process and the framing of one call over its pipes."""
|
||||
|
||||
def __init__(self, python: str, env: dict[str, str]) -> None:
|
||||
self.proc = subprocess.Popen(
|
||||
[python, str(WORKER_MAIN)],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
close_fds=True,
|
||||
env=env,
|
||||
)
|
||||
self._buffer = bytearray()
|
||||
|
||||
def alive(self) -> bool:
|
||||
return self.proc.poll() is None
|
||||
|
||||
def send(self, request: dict[str, Any]) -> None:
|
||||
assert self.proc.stdin is not None
|
||||
self.proc.stdin.write((json.dumps(request) + "\n").encode())
|
||||
self.proc.stdin.flush()
|
||||
|
||||
def read_line(self) -> str:
|
||||
"""One line, blocking. Empty when the process is gone.
|
||||
|
||||
Read at the file-descriptor level so a report arriving mid-call is
|
||||
passed on the moment it is written rather than when a buffer fills.
|
||||
"""
|
||||
assert self.proc.stdout is not None
|
||||
fd = self.proc.stdout.fileno()
|
||||
while True:
|
||||
end = self._buffer.find(b"\n")
|
||||
if end >= 0:
|
||||
line = bytes(self._buffer[: end + 1])
|
||||
del self._buffer[: end + 1]
|
||||
return line.decode(errors="replace")
|
||||
try:
|
||||
chunk = os.read(fd, 65536)
|
||||
except OSError:
|
||||
chunk = b""
|
||||
if not chunk:
|
||||
self._buffer.clear()
|
||||
return ""
|
||||
self._buffer += chunk
|
||||
|
||||
def kill(self) -> None:
|
||||
with contextlib.suppress(OSError):
|
||||
self.proc.send_signal(signal.SIGKILL)
|
||||
with contextlib.suppress(subprocess.TimeoutExpired):
|
||||
self.proc.wait(timeout=5)
|
||||
|
||||
|
||||
class Agent:
|
||||
"""Holds the connection, and one subprocess per call in flight."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace) -> None:
|
||||
self.args = args
|
||||
self.env = dict(os.environ)
|
||||
self.env["FLUKSIO_ARTIFACT_URL"] = args.artifact_url or _artifacts_from(
|
||||
args.url
|
||||
)
|
||||
self.env["FLUKSIO_ARTIFACT_TOKEN"] = args.token
|
||||
self.running: dict[str, Subprocess] = {}
|
||||
|
||||
async def serve_forever(self) -> None:
|
||||
backoff = 1.0
|
||||
while True:
|
||||
try:
|
||||
await self._session()
|
||||
backoff = 1.0
|
||||
except Exception as exc:
|
||||
log.warning("disconnected: %s — retrying in %.0fs", exc, backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(MAX_BACKOFF_S, backoff * 2)
|
||||
|
||||
async def _session(self) -> None:
|
||||
url = f"{self.args.url}?token={self.args.token}"
|
||||
async with websockets.connect(url, max_size=None, ping_interval=20) as socket:
|
||||
await socket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"op": "hello",
|
||||
"protocol": PROTOCOL,
|
||||
"name": self.args.name,
|
||||
"labels": self.args.labels,
|
||||
"python": self.args.python,
|
||||
"max_parallel": self.args.parallel,
|
||||
"venv_digest": _venv_digest(self.args.python),
|
||||
}
|
||||
)
|
||||
)
|
||||
welcome = json.loads(await socket.recv())
|
||||
if welcome.get("op") != "welcome":
|
||||
raise RuntimeError(str(welcome.get("reason") or "refused"))
|
||||
log.info(
|
||||
"attached to %s as '%s' with labels %s",
|
||||
self.args.url,
|
||||
welcome.get("name"),
|
||||
self.args.labels,
|
||||
)
|
||||
|
||||
async for raw in socket:
|
||||
message = json.loads(raw)
|
||||
op = message.get("op")
|
||||
if op in ("run", "compile"):
|
||||
# Compiling is loading the source, which is the same trip
|
||||
# through a subprocess a call is — and has to happen here
|
||||
# rather than on the engine, because "does this import"
|
||||
# is a question about *this* machine's packages.
|
||||
log.info("%s %s", op, message.get("call_id"))
|
||||
task = asyncio.create_task(self._run(socket, message))
|
||||
# Without this a failure in the task is only noticed when
|
||||
# it is garbage collected, which reads as a call that
|
||||
# vanished.
|
||||
task.add_done_callback(_report_failure)
|
||||
elif op == "cancel":
|
||||
self._cancel(str(message.get("call_id") or ""))
|
||||
|
||||
async def _run(self, socket: Any, request: dict[str, Any]) -> None:
|
||||
"""Execute one call in a subprocess, streaming what it says back."""
|
||||
call_id = str(request.get("call_id") or "")
|
||||
loop = asyncio.get_running_loop()
|
||||
worker = Subprocess(self.args.python, self.env)
|
||||
self.running[call_id] = worker
|
||||
|
||||
async def beat() -> None:
|
||||
while True:
|
||||
await asyncio.sleep(HEARTBEAT_S)
|
||||
with contextlib.suppress(Exception):
|
||||
await socket.send(
|
||||
json.dumps({"call_id": call_id, "event": "heartbeat"})
|
||||
)
|
||||
|
||||
heartbeat = asyncio.create_task(beat())
|
||||
try:
|
||||
await loop.run_in_executor(None, worker.send, request)
|
||||
while True:
|
||||
line = await loop.run_in_executor(None, worker.read_line)
|
||||
if not line:
|
||||
await socket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"call_id": call_id,
|
||||
"ok": False,
|
||||
"error": {
|
||||
"type": "NodeCancelled"
|
||||
if worker.proc.returncode
|
||||
else "RemoteError",
|
||||
"message": "the node process stopped",
|
||||
"short": "the node process stopped",
|
||||
"traceback": "",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return
|
||||
await socket.send(line.strip())
|
||||
# Anything without an `event` is the answer; the call is over.
|
||||
if not json.loads(line).get("event"):
|
||||
return
|
||||
finally:
|
||||
heartbeat.cancel()
|
||||
self.running.pop(call_id, None)
|
||||
worker.kill()
|
||||
|
||||
def _cancel(self, call_id: str) -> None:
|
||||
worker = self.running.get(call_id)
|
||||
if worker is not None:
|
||||
log.info("cancelling %s", call_id)
|
||||
worker.kill()
|
||||
|
||||
|
||||
def _report_failure(task: asyncio.Task[Any]) -> None:
|
||||
if not task.cancelled() and task.exception() is not None:
|
||||
log.exception("call failed", exc_info=task.exception())
|
||||
|
||||
|
||||
def _artifacts_from(url: str) -> str:
|
||||
"""The artifact endpoint beside the socket, so one URL configures both."""
|
||||
base = url.replace("wss://", "https://").replace("ws://", "http://")
|
||||
return base.rsplit("/workers/attach", 1)[0] + "/artifacts"
|
||||
|
||||
|
||||
def _venv_digest(python: str) -> str:
|
||||
"""What is installed here, so the engine can say when it has drifted."""
|
||||
try:
|
||||
listing = subprocess.run(
|
||||
[python, "-m", "pip", "freeze"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
check=False,
|
||||
).stdout
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return ""
|
||||
import hashlib
|
||||
|
||||
return hashlib.sha256(listing.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
|
||||
)
|
||||
parser = argparse.ArgumentParser(description="Run Fluksio nodes on this machine.")
|
||||
parser.add_argument("--url", required=True, help="wss://…/api/v1/workers/attach")
|
||||
parser.add_argument(
|
||||
"--token",
|
||||
default=os.environ.get("FLUKSIO_WORKER_TOKEN", ""),
|
||||
help="issued by POST /api/v1/workers/tokens",
|
||||
)
|
||||
parser.add_argument("--name", default=os.uname().nodename)
|
||||
parser.add_argument(
|
||||
"--labels",
|
||||
default="",
|
||||
help="comma-separated, e.g. gpu,cuda12 — what a node's device matches",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--python",
|
||||
default=sys.executable,
|
||||
help="the interpreter node code runs on; point it at the venv with torch",
|
||||
)
|
||||
parser.add_argument("--parallel", type=int, default=1)
|
||||
parser.add_argument("--artifact-url", default="")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if not args.token:
|
||||
parser.error("a token is required (--token or FLUKSIO_WORKER_TOKEN)")
|
||||
args.labels = [part.strip() for part in args.labels.split(",") if part.strip()]
|
||||
if not WORKER_MAIN.exists():
|
||||
parser.error(f"{WORKER_MAIN} is missing — copy it beside this file")
|
||||
|
||||
agent = Agent(args)
|
||||
try:
|
||||
asyncio.run(agent.serve_forever())
|
||||
except KeyboardInterrupt:
|
||||
log.info("stopping")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user