stroblme
3a9fc43707
The hello carries the fluksio-worker version. An agent refused for an older protocol stays listed as Needs upgrade until it attaches again. Docs gain an Upgrading workers section.
540 lines
20 KiB
Python
540 lines
20 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 --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 importlib.metadata
|
|
import ipaddress
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import urlsplit
|
|
|
|
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
|
|
#: Environment node code is not given; the engine's lists, see `_node_env`.
|
|
ENV_DENY_PREFIXES = ("DATABASE_URL", "FIRST_SUPERUSER", "SENTRY_DSN")
|
|
ENV_DENY_WORDS = ("PASSWORD", "SECRET", "TOKEN")
|
|
|
|
|
|
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 = _node_env()
|
|
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:
|
|
# A header rather than the query string, which access logs keep.
|
|
async with websockets.connect(
|
|
self.args.url,
|
|
additional_headers={"Authorization": f"Bearer {self.args.token}"},
|
|
max_size=None,
|
|
ping_interval=20,
|
|
) as socket:
|
|
await socket.send(
|
|
json.dumps(
|
|
{
|
|
"op": "hello",
|
|
"protocol": PROTOCOL,
|
|
"version": _version(),
|
|
"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 _node_env() -> dict[str, str]:
|
|
"""This host's environment with its credentials taken out, for node code.
|
|
|
|
The engine's own denylist (``fluksio.flow.workers.worker_env``), copied
|
|
because nothing of the engine is imported here: ``PATH``, ``HOME`` and
|
|
whatever a venv needs still arrive, a ``*_TOKEN`` or ``*_PASSWORD`` does
|
|
not — this worker's own token included. Not a sandbox: node code runs as
|
|
the same user on the same machine.
|
|
"""
|
|
return {
|
|
key: value
|
|
for key, value in os.environ.items()
|
|
if not key.startswith(ENV_DENY_PREFIXES)
|
|
and not any(word in key.upper() for word in ENV_DENY_WORDS)
|
|
}
|
|
|
|
|
|
def _encrypted_or_local(url: str) -> bool:
|
|
"""Whether the token and node secrets sent to *url* stay off the wire.
|
|
|
|
TLS, or this machine. An engine across a network on plain ``ws://`` hands
|
|
both to anything on the path.
|
|
"""
|
|
parts = urlsplit(url)
|
|
if parts.scheme in ("wss", "https"):
|
|
return True
|
|
host = parts.hostname or ""
|
|
if host == "localhost":
|
|
return True
|
|
try:
|
|
return ipaddress.ip_address(host).is_loopback
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
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 _version() -> str:
|
|
"""The installed fluksio-worker, or empty when this file was copied over."""
|
|
try:
|
|
return importlib.metadata.version("fluksio-worker")
|
|
except importlib.metadata.PackageNotFoundError:
|
|
return ""
|
|
|
|
|
|
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(
|
|
"--insecure",
|
|
action="store_true",
|
|
help=(
|
|
"allow ws:// and http:// to an engine off this machine — the token "
|
|
"and every secret a node is handed then cross the network in clear"
|
|
),
|
|
)
|
|
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)")
|
|
for url in (args.url, args.artifact_url):
|
|
if url and not args.insecure and not _encrypted_or_local(url):
|
|
parser.error(
|
|
f"{url} is not encrypted: use wss:// (https:// for --artifact-url), "
|
|
"or --insecure on a network you trust"
|
|
)
|
|
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())
|