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
backend/fluksio/api/routes/workers.py
+13 -1
View File
@@ -21,7 +21,12 @@ from pydantic import BaseModel, Field
from fluksio.api.deps import get_current_active_superuser, get_current_user
from fluksio.core import security
from fluksio.flow.remote import PROTOCOL, RemoteWorker, RemoteWorkerHub
from fluksio.flow.remote import (
PROTOCOL,
RemoteWorker,
RemoteWorkerHub,
WorkerInventory,
)
logger = logging.getLogger(__name__)
@@ -40,6 +45,9 @@ class WorkerInfo(BaseModel):
last_seen: float = 0.0
python: str = ""
venv_digest: str = ""
cpus: int = 1
gpus: int = 0
ram_mb: int | None = None
class TokenRequest(BaseModel):
@@ -74,6 +82,9 @@ def read_workers(request: Request) -> Any:
last_seen=worker.last_seen,
python=str(worker.info.get("python") or ""),
venv_digest=str(worker.info.get("venv_digest") or ""),
cpus=worker.inventory.cpus,
gpus=worker.inventory.gpus,
ram_mb=worker.inventory.ram_mb,
)
for worker in _hub(request.app).workers()
]
@@ -163,6 +174,7 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
"python": hello.get("python"),
"venv_digest": hello.get("venv_digest"),
},
inventory=WorkerInventory.from_hello(hello.get("inventory")),
)
hub.attach(worker)
await websocket.send_json({"op": "welcome", "protocol": PROTOCOL, "name": name})
backend/fluksio/flow/remote.py
+53 -4
View File
@@ -28,6 +28,7 @@ import sys
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
from fluksio.flow.nodes.base import NodeOutputError
@@ -44,13 +45,50 @@ SEND_TIMEOUT_S = 30.0
SILENCE_S = 90.0
#: Protocol version this engine speaks. A worker announcing anything else is
#: refused rather than half-understood.
PROTOCOL = 1
PROTOCOL = 2
class NoWorker(RemoteError):
"""Nothing is attached that carries the label this node asked for."""
def _whole(value: Any, default: int) -> int:
"""An int from whatever a worker sent, or the default when it was junk."""
try:
return int(value)
except (TypeError, ValueError):
return default
@dataclass(frozen=True)
class WorkerInventory:
"""What an attached worker says it holds.
The defaults are a machine that runs one thing at a time and has no GPU,
so a worker reporting nothing is scheduled by its label alone — which is
how every worker was scheduled before any of them reported anything.
"""
cpus: int = 1
gpus: int = 0
ram_mb: int | None = None
@classmethod
def from_hello(cls, payload: Any) -> WorkerInventory:
"""Read what is understood and ignore the rest.
Unknown members are deliberately not an error: a worker that learns to
report something new must not need this engine taught about it first.
"""
if not isinstance(payload, dict):
return cls()
return cls(
cpus=max(1, _whole(payload.get("cpus"), 1)),
gpus=max(0, _whole(payload.get("gpus"), 0)),
ram_mb=_whole(payload.get("ram_mb"), 0) or None,
)
class RemoteWorker:
"""One attached worker, and the calls it has in flight."""
@@ -62,10 +100,12 @@ class RemoteWorker:
loop: asyncio.AbstractEventLoop,
max_parallel: int = 1,
info: dict[str, Any] | None = None,
inventory: WorkerInventory | None = None,
) -> None:
self.name = name
self.labels = set(labels)
self.info = info or {}
self.inventory = inventory or WorkerInventory()
self.attached_at = time.time()
self.last_seen = time.time()
self._send = send
@@ -281,12 +321,12 @@ class RemoteWorkerHub:
timeout: float,
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
env: dict[str, str] | None = None,
) -> Any:
worker = self.pick(label)
if worker is None:
raise NoWorker(f"no worker labelled '{label}' is attached")
response = worker.request(
{
payload = {
"op": "run",
"call_id": f"{run_id}:{node_id}" if run_id else node_id,
"flow": flow,
@@ -295,7 +335,14 @@ class RemoteWorkerHub:
"kwargs": kwargs,
"run": {"id": run_id} if run_id else None,
"timeout": timeout,
},
}
if env:
# What this call was allocated: the thread caps and the devices the
# node may see. The worker starts a process per call, so it applies
# them where a library still reads them — before the import.
payload["env"] = env
response = worker.request(
payload,
timeout=timeout,
on_event=on_event,
)
@@ -357,6 +404,7 @@ class RemoteWorkerHub:
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
fallback: Callable[..., Any] | None = None,
env: dict[str, str] | None = None,
) -> Callable[..., Any]:
"""What a node with a device runs instead of its own function.
@@ -377,6 +425,7 @@ class RemoteWorkerHub:
timeout,
run_id=run_id,
on_event=on_event,
env=env,
)
return call
backend/tests/flow/test_remote.py
+70 -1
View File
@@ -15,7 +15,12 @@ from collections.abc import Iterator
import pytest
from fluksio.flow import remote
from fluksio.flow.remote import NoWorker, RemoteWorker, RemoteWorkerHub
from fluksio.flow.remote import (
NoWorker,
RemoteWorker,
RemoteWorkerHub,
WorkerInventory,
)
from fluksio.flow.workers import NodeTimeout, RemoteError
@@ -295,3 +300,67 @@ def test_with_no_timeout_a_beating_worker_is_left_to_finish(loop, monkeypatch):
thread.join(timeout=5)
assert result["value"] == {"done": True}
def test_an_allocation_rides_along_with_the_call(loop):
"""Thread caps and devices reach the worker on the frame that needs them.
The worker starts a process per call, so this is the only moment it can
apply them: a library reads them when it is imported and never again.
"""
hub = RemoteWorkerHub()
worker, socket = attach(hub, loop)
thread = call_in_thread(
lambda: hub.run(
"gpu",
"flow",
"node",
"src",
{},
"flow.node",
timeout=5,
env={"OMP_NUM_THREADS": "4", "CUDA_VISIBLE_DEVICES": "0"},
)
)
assert socket.arrived.wait(5)
assert socket.sent[0]["env"] == {
"OMP_NUM_THREADS": "4",
"CUDA_VISIBLE_DEVICES": "0",
}
worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
thread.join(timeout=5)
def test_a_call_with_nothing_allocated_carries_no_env(loop):
hub = RemoteWorkerHub()
worker, socket = attach(hub, loop)
thread = call_in_thread(
lambda: hub.run("gpu", "flow", "node", "src", {}, "flow.node", timeout=5)
)
assert socket.arrived.wait(5)
assert "env" not in socket.sent[0]
worker.deliver({"call_id": socket.sent[0]["call_id"], "ok": True, "result": None})
thread.join(timeout=5)
def test_a_worker_that_reports_nothing_is_the_machine_it_used_to_be():
"""Inventory is read leniently: absent is a default, junk is a default.
A worker saying nothing has to keep being scheduled by its label, and one
that learns to report something new must not need this engine taught about
it first.
"""
assert WorkerInventory.from_hello({}) == WorkerInventory(cpus=1, gpus=0)
assert WorkerInventory.from_hello(None) == WorkerInventory()
read = WorkerInventory.from_hello(
{"cpus": "8", "gpus": 2, "ram_mb": 64000, "tpus": 4}
)
assert (read.cpus, read.gpus, read.ram_mb) == (8, 2, 64000)
junk = WorkerInventory.from_hello({"cpus": "many", "gpus": None, "ram_mb": 0})
assert (junk.cpus, junk.gpus, junk.ram_mb) == (1, 0, None)
docs/code/workers.md
+38 -2
View File
@@ -35,11 +35,37 @@ driver in order to run a training step. An engine host already has it, and
| `--python` | this interpreter | the interpreter node code runs on |
| `--parallel` | `1` | how many node calls it will take at once |
| `--artifact-url` | derived from `--url` | where the artifact store is, if not beside the socket |
| `--cpus` | what the job or the machine has | cores to advertise |
| `--gpus` | what the job says, else none | GPUs to advertise; never probed |
| `--ram-mb` | what the job or the machine has | memory to advertise, in MB |
| `--max-idle` | never | stop after this many seconds with nothing running |
`--python` is the important one. It is how this machine keeps its own wheels —
the CUDA build, the vendor SDK, the thing that will not install anywhere else —
without the engine ever installing them or knowing about them.
### What it says it has
A worker reports its inventory when it attaches — cores, GPUs and memory — and
the engine schedules against it: a node asking for two cores and a GPU goes to
a machine that has them free, not merely to one carrying the right label.
Cores and memory are read off the machine, or off the batch job that started
this worker (`SLURM_CPUS_ON_NODE`, `SLURM_MEM_PER_NODE`). **GPUs are never
probed.** Asking a vendor tool would make the one dependency two, so a GPU is
something the job says it was given (`SLURM_GPUS_ON_NODE`, `SLURM_JOB_GPUS`,
or `FLUKSIO_WORKER_GPUS`) or something you say with `--gpus`. A worker that
reports nothing still attaches and is scheduled by its label alone, as every
worker was before any of them reported anything.
The engine tells each call what it may use — thread caps, and the devices it
may see. The worker starts a process per call, so it applies them at the only
moment a numerical library still reads them: before the import.
`--max-idle` is for a worker something else started for one job — a batch
scheduler, say. It exits when nothing has run for that long, so the allocation
goes back rather than idling until its walltime.
## Mint the token
On the engine, as a superuser:
@@ -122,8 +148,18 @@ curl -s $FLUKSIO/workers -H "Authorization: Bearer $TOKEN" | jq
```
Name, labels, how many calls it will take at once, how many are in flight, when
it attached, when it was last seen, its Python version, and a digest of its
environment.
it attached, when it was last seen, its Python version, a digest of its
environment, and what it says it has: cores, GPUs and memory.
## Upgrading
The engine and the worker speak a version-matched protocol, and a worker
announcing anything else is refused rather than half-understood. Protocol 2 —
the one that carries inventory — is `fluksio-worker` 0.2.0. An older agent is
told so on the socket and stops, rather than retrying against an engine that
will never accept it; `pip install -U fluksio-worker` on that host is the whole
upgrade. Nothing changed in the runner served at `GET /api/v1/workers/runtime`,
so a host that copies its two files copies the same one as before.
## What a worker is not
uv.lock
Generated
+1 -1
View File
@@ -948,7 +948,7 @@ dev = [
[[package]]
name = "fluksio-worker"
version = "0.1.4"
version = "0.2.0"
source = { editable = "worker" }
dependencies = [
{ name = "websockets" },
worker/fluksio_worker/agent.py
+167 -10
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,
)
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.
# 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.
# 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
worker/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "fluksio-worker"
version = "0.1.4"
version = "0.2.0"
description = "Runs Fluksio nodes on a machine the engine cannot reach"
readme = "README.md"
license = "AGPL-3.0-or-later"