diff --git a/backend/fluksio/api/routes/workers.py b/backend/fluksio/api/routes/workers.py index 88c4b02..61dd1fa 100644 --- a/backend/fluksio/api/routes/workers.py +++ b/backend/fluksio/api/routes/workers.py @@ -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}) diff --git a/backend/fluksio/flow/remote.py b/backend/fluksio/flow/remote.py index 8554f6c..526872f 100644 --- a/backend/fluksio/flow/remote.py +++ b/backend/fluksio/flow/remote.py @@ -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,21 +321,28 @@ 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") + payload = { + "op": "run", + "call_id": f"{run_id}:{node_id}" if run_id else node_id, + "flow": flow, + "node": node, + "source": source, + "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( - { - "op": "run", - "call_id": f"{run_id}:{node_id}" if run_id else node_id, - "flow": flow, - "node": node, - "source": source, - "kwargs": kwargs, - "run": {"id": run_id} if run_id else None, - "timeout": timeout, - }, + 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 diff --git a/backend/tests/flow/test_remote.py b/backend/tests/flow/test_remote.py index 2b0c52b..8edeb4b 100644 --- a/backend/tests/flow/test_remote.py +++ b/backend/tests/flow/test_remote.py @@ -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) diff --git a/docs/code/workers.md b/docs/code/workers.md index 885285e..715b138 100644 --- a/docs/code/workers.md +++ b/docs/code/workers.md @@ -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 diff --git a/uv.lock b/uv.lock index e3ae1a9..80ba2a0 100644 --- a/uv.lock +++ b/uv.lock @@ -948,7 +948,7 @@ dev = [ [[package]] name = "fluksio-worker" -version = "0.1.4" +version = "0.2.0" source = { editable = "worker" } dependencies = [ { name = "websockets" }, diff --git a/worker/fluksio_worker/agent.py b/worker/fluksio_worker/agent.py index 37fb4bb..f2e7058 100644 --- a/worker/fluksio_worker/agent.py +++ b/worker/fluksio_worker/agent.py @@ -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 diff --git a/worker/pyproject.toml b/worker/pyproject.toml index e5f2979..620a5ba 100644 --- a/worker/pyproject.toml +++ b/worker/pyproject.toml @@ -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"