Fluksio/app
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Let a worker say what machine it is

Browse Source 
A worker reported its labels and nothing about the machine behind them, so the
engine could route a node to a GPU box but not tell whether that box had a GPU
free. Inventory — cores, GPUs, memory — now arrives with the hello frame, and
the run frame carries back what the engine allocated for that call.

Which is protocol 2 on both ends. GPUs are never probed: asking a vendor tool
would make the one dependency two, so a GPU is what the batch job says it was
given or what --gpus says. A worker that reports nothing still attaches and is
scheduled by its label alone.

Two things a job scheduler needs: --max-idle stops a worker started for one job
rather than letting it hold its allocation to the walltime, and a refusal is now
fatal instead of a reconnect loop that reads as a hang in a job's log.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
This commit is contained in:
stroblme
2026-08-27 08:29:35 +02:00
co-authored by Claude Opus 5
parent 0ffcabfdb9
commit 1a9753fa9d
7 changed files with 359 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
worker/fluksio_worker/agent.py
+176 -19
View File
@@ -34,6 +34,7 @@ import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
@@ -47,13 +48,21 @@ except ImportError: # pragma: no cover - the one dependency, named plainly
log = logging.getLogger("fluksio-worker")
PROTOCOL = 1
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")
@@ -121,6 +130,8 @@ class Agent:
)
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
@@ -128,6 +139,11 @@ class Agent:
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)
@@ -146,12 +162,15 @@ class Agent:
"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":
raise RuntimeError(str(welcome.get("reason") or "refused"))
# 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,
@@ -159,28 +178,73 @@ class Agent:
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 ""))
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()
worker = Subprocess(self.args.python, self.env)
# 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:
@@ -221,6 +285,7 @@ class Agent:
finally:
heartbeat.cancel()
self.running.pop(call_id, None)
self.last_done = time.monotonic()
worker.kill()
def _cancel(self, call_id: str) -> None:
@@ -241,6 +306,68 @@ def _artifacts_from(url: str) -> str:
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:
@@ -285,6 +412,33 @@ def main(argv: list[str] | None = None) -> int:
)
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:
@@ -298,6 +452,9 @@ def main(argv: list[str] | None = None) -> int:
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