stroblmeandClaude Opus 5
c09095d369
The log tee wrote through to the real stream unguarded, and the worker pool tees a returned call's logs there after reading its result and before handing it back — so a dead stdout, which `fluksio serve` makes possible by running the engine as a child of the dashboard holding that pipe, failed the node with its outputs already in hand. The capture half runs first, so swallowing the write loses nothing. Also: `flow_events` catches the RuntimeError a peer leaving mid-send raises, which is a disconnect by another route, and the remote agent no longer raises out of the task when its subprocess died before it could be written to — the read below reports that and ends the call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TXQv6KNyyvY7Z1etYTUUAd
467 lines
17 KiB
Python
467 lines
17 KiB
Python
"""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
|
|
|
|
Install it into the environment the training code already runs in and node
|
|
code runs on that: ``--python`` defaults to the interpreter this was started
|
|
with. Point it elsewhere only when the two are meant to differ.
|
|
|
|
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
|
|
import time
|
|
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 = 2
|
|
#: 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
|
|
|
|
|
|
class Refused(RuntimeError):
|
|
"""The engine will not have this worker, and retrying will not change it."""
|
|
|
|
|
|
class Idle(Exception):
|
|
"""Nothing has run here for long enough that this worker is done."""
|
|
|
|
|
|
#: 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] = {}
|
|
self.last_done = time.monotonic()
|
|
self._idled = False
|
|
|
|
async def serve_forever(self) -> None:
|
|
backoff = 1.0
|
|
while True:
|
|
try:
|
|
await self._session()
|
|
backoff = 1.0
|
|
except Idle:
|
|
log.info("nothing to do for %.0fs — stopping", self.args.max_idle)
|
|
return
|
|
except Refused:
|
|
raise
|
|
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),
|
|
"inventory": _inventory(self.args),
|
|
}
|
|
)
|
|
)
|
|
welcome = json.loads(await socket.recv())
|
|
if welcome.get("op") != "welcome":
|
|
# Nothing about this machine will change the answer, so
|
|
# reconnecting would only be a quiet loop in the job's log.
|
|
raise Refused(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,
|
|
)
|
|
|
|
self.last_done = time.monotonic()
|
|
watchdog = (
|
|
asyncio.create_task(self._idle_watch(socket))
|
|
if self.args.max_idle > 0
|
|
else None
|
|
)
|
|
try:
|
|
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 ""))
|
|
except Exception:
|
|
# The watchdog closing the socket is a stop, not a disconnect.
|
|
if not self._idled:
|
|
raise
|
|
finally:
|
|
if watchdog is not None:
|
|
watchdog.cancel()
|
|
if self._idled:
|
|
raise Idle
|
|
|
|
async def _idle_watch(self, socket: Any) -> None:
|
|
"""Close the socket once nothing has run here for ``--max-idle``.
|
|
|
|
A worker a batch scheduler started for one node has no other way to
|
|
know it is finished, and a job that idles until its walltime is a job
|
|
somebody else was queued behind.
|
|
"""
|
|
tick = min(30.0, max(1.0, self.args.max_idle / 4))
|
|
while True:
|
|
await asyncio.sleep(tick)
|
|
if self.running:
|
|
continue
|
|
if time.monotonic() - self.last_done >= self.args.max_idle:
|
|
self._idled = True
|
|
with contextlib.suppress(Exception):
|
|
await socket.close()
|
|
return
|
|
|
|
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()
|
|
# What the engine allocated this call: thread caps and the GPUs it may
|
|
# see. Taken off the request because a library reads those once, when
|
|
# it is imported — which for this agent is a process that does not
|
|
# exist yet, so there is nothing to reconfigure and no pool to key.
|
|
extra = {str(k): str(v) for k, v in (request.pop("env", None) or {}).items()}
|
|
mask = self.env.get("CUDA_VISIBLE_DEVICES")
|
|
if mask and extra.get("CUDA_VISIBLE_DEVICES"):
|
|
extra["CUDA_VISIBLE_DEVICES"] = _remap_cuda(
|
|
mask, extra["CUDA_VISIBLE_DEVICES"]
|
|
)
|
|
worker = Subprocess(self.args.python, {**self.env, **extra})
|
|
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:
|
|
# A subprocess that died before it could be written to is reported
|
|
# by the read below, which says so and ends the call — rather than
|
|
# raising here and leaving the engine waiting out its silence.
|
|
with contextlib.suppress(OSError):
|
|
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)
|
|
self.last_done = time.monotonic()
|
|
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 _detect_cpus() -> int:
|
|
"""Cores this worker may use — what the batch job was given, or the box."""
|
|
given = os.environ.get("SLURM_CPUS_ON_NODE", "")
|
|
return int(given) if given.isdigit() else (os.cpu_count() or 1)
|
|
|
|
|
|
def _detect_ram_mb() -> int | None:
|
|
"""Memory in MB, or None where it cannot be asked — which is allowed."""
|
|
given = os.environ.get("SLURM_MEM_PER_NODE", "")
|
|
if given.isdigit():
|
|
return int(given)
|
|
try:
|
|
return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
|
|
except (AttributeError, ValueError, OSError):
|
|
return None
|
|
|
|
|
|
def _detect_gpus() -> int:
|
|
"""What this worker was *given*, never what the machine has.
|
|
|
|
Nothing is probed: asking a vendor tool would make the one dependency two,
|
|
and the engine does not probe its own GPUs either. A batch scheduler says
|
|
so in the environment; anywhere else it is ``--gpus``.
|
|
"""
|
|
for name in ("SLURM_GPUS_ON_NODE", "FLUKSIO_WORKER_GPUS"):
|
|
given = os.environ.get(name, "")
|
|
if given.isdigit():
|
|
return int(given)
|
|
listed = os.environ.get("SLURM_JOB_GPUS", "")
|
|
return len([part for part in listed.split(",") if part.strip()])
|
|
|
|
|
|
def _inventory(args: argparse.Namespace) -> dict[str, Any]:
|
|
"""What this machine holds, for the engine to schedule against."""
|
|
return {
|
|
"cpus": args.cpus or _detect_cpus(),
|
|
"gpus": _detect_gpus() if args.gpus is None else args.gpus,
|
|
"ram_mb": args.ram_mb or _detect_ram_mb(),
|
|
}
|
|
|
|
|
|
def _remap_cuda(mask: str, requested: str) -> str:
|
|
"""Turn engine-assigned device numbers into the ones this job may see.
|
|
|
|
The engine counts GPUs from the inventory reported here, so it asks for
|
|
device 0. Under a scheduler's gres this agent already runs with something
|
|
like ``CUDA_VISIBLE_DEVICES=2,3``, where 0 is a device the job was never
|
|
given — so each index is read as a position in that mask.
|
|
"""
|
|
devices = [part.strip() for part in mask.split(",") if part.strip()]
|
|
out = []
|
|
for part in requested.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
try:
|
|
out.append(devices[int(part)])
|
|
except (ValueError, IndexError):
|
|
out.append(part)
|
|
return ",".join(out)
|
|
|
|
|
|
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 (default: the one running this, "
|
|
"so installing into the venv with torch in it is enough)"
|
|
),
|
|
)
|
|
parser.add_argument("--parallel", type=int, default=1)
|
|
parser.add_argument("--artifact-url", default="")
|
|
parser.add_argument(
|
|
"--cpus",
|
|
type=int,
|
|
default=0,
|
|
help="cores to advertise (default: what the job or the machine has)",
|
|
)
|
|
parser.add_argument(
|
|
"--gpus",
|
|
type=int,
|
|
default=None,
|
|
help="GPUs to advertise — nothing is probed, so say so here or in the job",
|
|
)
|
|
parser.add_argument(
|
|
"--ram-mb",
|
|
type=int,
|
|
default=0,
|
|
help="memory in MB to advertise (default: what the job or the machine has)",
|
|
)
|
|
parser.add_argument(
|
|
"--max-idle",
|
|
type=float,
|
|
default=0.0,
|
|
help=(
|
|
"stop after this many seconds with nothing running — for a worker "
|
|
"a batch scheduler started for one job (default: never)"
|
|
),
|
|
)
|
|
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")
|
|
except Refused as exc:
|
|
log.error("the engine refused this worker: %s", exc)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|