stroblmeandClaude Opus 5
058f16ec1d
Each was a loose end recorded under `### SDK` in the notepad. `serve` takes its own pidfile down on SIGTERM. uvicorn restores the handler it found and re-raises the signal it stopped on, so the default handler ended the process without unwinding and the `finally` never ran — which is what a stop sends, and what left `serve.pid` behind. `serve.log` is cut back past 5 MB by the engine rather than by the screen that started it, so an adopted engine is bounded too. Gated on its own stdout being an appended regular file, which is what makes the cut safe: the kernel then puts the next write at the new end. Cards are counted from `/dev/nvidia[0-9]*`, so `FLOW_GPUS`/`--gpus` of 0 means "work it out" the way `FLOW_CPUS` always has. The engine counts, not the accountant — a remote worker builds one of those from its own inventory, and detecting there would hand it the engine host's cards. The worker counts last: what a batch job says it was granted still wins. `GET /runs/metrics/names` is the distinct over a selection that `--list` and the terminal's metric picker were approximating by reading the newest run that had measured anything, which missed a name only an older run ever wrote. `MetricSink` announces each batch it has written (`run_metric`, carrying the names). Not a per-point event: one covers up to 500 points or two seconds of them, and the rows stay the record. The terminal comparison fills in as the first readings land instead of staying blank until reopened, and the browser refetches the run and any comparison rather than the list behind them. `retry --group` pages the list route by `before` instead of stopping at 500. The terminal dashboard takes the terminal's colours (`ansi-dark`), and the web UI can re-pair from Settings without disconnecting first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PRQ9bmTvCbqCwXo9mxZzzV
472 lines
17 KiB
Python
472 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 glob
|
|
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*, else what the box appears to have.
|
|
|
|
No vendor tool is asked — that would make the one dependency two. A batch
|
|
scheduler says what the job was given in the environment and that always
|
|
wins, since a node with eight cards may have granted this job one. With
|
|
nothing said, NVIDIA's device nodes are counted, which is what the engine
|
|
does for its own machine.
|
|
"""
|
|
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", "")
|
|
if listed.strip():
|
|
return len([part for part in listed.split(",") if part.strip()])
|
|
return len(glob.glob("/dev/nvidia[0-9]*"))
|
|
|
|
|
|
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())
|