From 6ff56533f50bad73a6520eca60ea285ccfc30096 Mon Sep 17 00:00:00 2001 From: stroblme Date: Thu, 27 Aug 2026 08:49:36 +0200 Subject: [PATCH] Schedule a node across every machine, not just this one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The engine answered "where does this node run" twice, in two ways that could not see each other: a device sent it to a worker carrying that label, and resources were counted against the engine's own cores. Declaring both meant the second answer won and nothing was counted at all — which the data-science getting-started page and the worked example both do. One question now, in flow/placement.py: of every machine attached, which could grant what this node asked for, and which of those has it free. The books move onto each machine — one accountant per worker, built from the inventory it reported — and the waiting moves above them, where one condition variable can be woken by a release anywhere or by a worker attaching. Locks go one way: placer, then a machine's books, never back. So a node asking for a card now finds the box that has one, rather than being clamped down to none and run here. When nothing can grant the ask at all it is still cut down and run — a flow written on a cluster has to work on a laptop — but the ceiling is one real machine now, since taking the largest of each dimension separately can describe a machine nobody has. Two things fixed on the way. A device on a connector node held every batch run of its flow forever, waiting for a worker that could never run an entry point. And `prefer` falling back to the engine skipped the books, so the fallback held nothing. The bench flow's node has taken a `params` argument that with_settings has not forwarded for some time, so the benchmark could not run at all: 62 ms median submit-to-result with this, against the 61 ms on record. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW --- .../versions/b3f1a7c50d92_run_needs.py | 31 ++ backend/fluksio/api/routes/workers.py | 46 +- backend/fluksio/flow/controller.py | 117 ++++-- backend/fluksio/flow/placement.py | 390 +++++++++++++++++ backend/fluksio/flow/remote.py | 53 ++- backend/fluksio/flow/resources.py | 193 +++------ backend/fluksio/flow/runs.py | 64 ++- backend/fluksio/flow/schemas.py | 10 + backend/fluksio/main.py | 14 +- backend/fluksio/models.py | 4 + backend/scripts/bench_startup.py | 2 +- backend/tests/flow/test_placement.py | 394 ++++++++++++++++++ backend/tests/flow/test_resources.py | 193 +++------ backend/tests/flow/test_runs.py | 29 +- 14 files changed, 1214 insertions(+), 326 deletions(-) create mode 100644 backend/fluksio/alembic/versions/b3f1a7c50d92_run_needs.py create mode 100644 backend/fluksio/flow/placement.py create mode 100644 backend/tests/flow/test_placement.py diff --git a/backend/fluksio/alembic/versions/b3f1a7c50d92_run_needs.py b/backend/fluksio/alembic/versions/b3f1a7c50d92_run_needs.py new file mode 100644 index 0000000..6d34c36 --- /dev/null +++ b/backend/fluksio/alembic/versions/b3f1a7c50d92_run_needs.py @@ -0,0 +1,31 @@ +"""run.needs + +A run recorded the worker *labels* its nodes asked for, which answered whether +anything was attached to run them but nothing about whether that machine was +big enough. The size goes beside the labels, so a queued run can tell waiting +for a machine from having nowhere to run at all. + +Revision ID: b3f1a7c50d92 +Revises: e5b8c2f4a913 +Create Date: 2026-08-27 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b3f1a7c50d92" +down_revision = "e5b8c2f4a913" +branch_labels = None +depends_on = None + + +def upgrade(): + # Nullable: a run submitted before this declared nothing, and null is that + # rather than a run that needs nothing in particular. + op.add_column("run", sa.Column("needs", sa.JSON(), nullable=True)) + + +def downgrade(): + op.drop_column("run", "needs") diff --git a/backend/fluksio/api/routes/workers.py b/backend/fluksio/api/routes/workers.py index 61dd1fa..917d98b 100644 --- a/backend/fluksio/api/routes/workers.py +++ b/backend/fluksio/api/routes/workers.py @@ -90,18 +90,54 @@ def read_workers(request: Request) -> Any: ] -@router.get("/resources", dependencies=[Depends(get_current_user)]) +class ResourceLevel(BaseModel): + total: int + free: int + + +class WaitingNode(BaseModel): + node: str + reason: str + seconds: float + + +class TargetResources(BaseModel): + """One machine: this engine, or a worker attached to it.""" + + target: str + cpus: ResourceLevel + gpus: ResourceLevel + #: Absent where the machine did not say how much memory it has. + ram_mb: ResourceLevel | None = None + labels: list[str] = Field(default_factory=list) + in_flight: int = 0 + + +class ResourcesSnapshot(BaseModel): + #: This engine's own figures, kept where they have always been. + cpus: ResourceLevel + gpus: ResourceLevel + waiting: list[WaitingNode] = Field(default_factory=list) + targets: list[TargetResources] = Field(default_factory=list) + provisioners: list[dict[str, Any]] = Field(default_factory=list) + + +@router.get( + "/resources", + response_model=ResourcesSnapshot, + dependencies=[Depends(get_current_user)], +) def read_resources(request: Request) -> Any: - """What this machine has free, and which nodes are queued for it. + """Every machine, what is free of it, and which nodes are queued. A node waiting its turn looks exactly like a node that has hung — the run sits at `running` and says nothing — so what is waiting, and for what, has to be readable somewhere. """ - accountant = getattr(request.app.state, "resources", None) - if accountant is None: + placer = getattr(request.app.state, "placer", None) + if placer is None: raise HTTPException(status_code=503, detail="Resources are not accounted here") - return accountant.snapshot() + return placer.snapshot() @router.post( diff --git a/backend/fluksio/flow/controller.py b/backend/fluksio/flow/controller.py index 5701689..22f2bc1 100644 --- a/backend/fluksio/flow/controller.py +++ b/backend/fluksio/flow/controller.py @@ -62,6 +62,7 @@ from fluksio.flow.pipeline import ( ValidationIssue, ValueSource, ) +from fluksio.flow.placement import Placer from fluksio.flow.remote import RemoteWorkerHub from fluksio.flow.resources import ResourceAccountant, derive_env from fluksio.flow.schemas import ( @@ -413,6 +414,7 @@ class FlowController: workers: PythonWorkerPool | None = None, remote: RemoteWorkerHub | None = None, resources: ResourceAccountant | None = None, + placer: Placer | None = None, ) -> None: self.store = store # Without a pool, python nodes are compiled and run in this process — @@ -420,9 +422,11 @@ class FlowController: self.workers = workers # Workers on other hosts. A node without a device never touches it. self.remote = remote - # What the machine has, for the nodes that say what they need. Without - # one, a declaration is recorded and nothing is held against it. + # This machine's own books. The pool sizes its fair share off them. self.resources = resources + # Every machine there is, for the nodes that say what they need. + # Without one, a declaration is recorded and nothing is held against it. + self.placer = placer self.state = state if state is not None else MemoryState() self.events = events self.max_workers = max_workers @@ -865,9 +869,18 @@ class FlowController: # Building # ------------------------------------------------------------------------- - def _allocated( + def _runs_elsewhere(self, node_def: NodeDef) -> bool: + """Whether this node asks for more than this machine could ever give.""" + if node_def.resources is None or self.placer is None: + return False + wanted = node_def.resources + return not self.placer.local.fits(wanted.cpus, wanted.gpus, wanted.ram or 0) + + def _placed( self, wanted: Resources, + device: str | None, + policy: str, owner: str, local: str, code: str, @@ -877,32 +890,40 @@ class FlowController: run_id: str, on_event: Callable[[dict[str, Any]], None], ) -> Callable[..., Any]: - """A call that holds its share of the machine while it runs. + """A call that picks a machine and holds its share while it runs. + + Which machine is decided per call rather than when the flow was built, + so a worker that attaches later is used without anything being rebuilt. The order is load-bearing: the resources are claimed first, and only then is a worker slot taken. The other way round, a node holding a slot could sit waiting for cores that a node holding the cores cannot get a slot to release. - The worker comes from the pool whose environment this allocation - derives, so what the node is told about its share is what the library - inside it reads at import — the only moment those variables are read. + What the allocation implies is handed to the process the node runs in, + here or on the worker, because a library reads those variables when it + is imported and never again. """ - if self.workers is None or self.resources is None: - return self.workers.proxy( # type: ignore[union-attr] - owner, - local, - code, - node_id=node_id, - timeout=timeout, - run_id=run_id, - on_event=on_event, - ) - accountant, pool = self.resources, self.workers + placer, pool = self.placer, self.workers def call(**kwargs: Any) -> Any: - with accountant.claim(wanted, node=node_id, run=run_id) as allocation: - return pool.for_env(derive_env(wanted, allocation)).run( + with placer.claim( + wanted, device=device, policy=policy, node=node_id, run=run_id + ) as (target, allocation): + env = derive_env(wanted, allocation) + if target.worker is None: + return pool.for_env(env).run( + owner, + local, + code, + kwargs, + node_id, + timeout, + run_id=run_id, + on_event=on_event, + ) + return self.remote.run_on( # type: ignore[union-attr] + target.worker, owner, local, code, @@ -911,6 +932,7 @@ class FlowController: timeout, run_id=run_id, on_event=on_event, + env=env, ) return call @@ -1003,6 +1025,11 @@ class FlowController: problem = self.remote.compile( node_def.device or "", owner, local, code ) + elif self._runs_elsewhere(node_def): + # Asks for more than this machine has, so it will run on + # one that has it. Same reason as a device: checking the + # import here would fail a node that is fine there. + problem = None else: problem = self.workers.compile(owner, local, code) if problem: @@ -1016,18 +1043,15 @@ class FlowController: if node_def.timeout is not None else settings.FLOW_NODE_TIMEOUT ) - function = self.workers.proxy( - owner, - local, - code, - node_id=node_id, - timeout=timeout, - run_id=run.run_id if run else "", - on_event=emissions.handle, - ) - if node_def.resources is not None and not node_def.device: - function = self._allocated( + if node_def.resources is not None and self.placer is not None: + # Says how much of a machine it takes, so which machine + # and how much of it are one decision — including when + # it also names a device, which used to mean the two + # answers disagreed and nothing was accounted at all. + function = self._placed( node_def.resources, + node_def.device, + node_def.device_policy, owner, local, code, @@ -1036,11 +1060,9 @@ class FlowController: run_id=run.run_id if run else "", on_event=emissions.handle, ) - if node_def.device and self.remote is not None: - # A node with a device runs on a worker carrying that - # label. Which worker is decided per call, so one that - # attaches after this flow was built is used without - # anything being rebuilt. + elif node_def.device and self.remote is not None: + # A device and nothing about size: the label alone + # decides, least busy first, as it always has. function = self.remote.proxy( node_def.device, owner, @@ -1051,9 +1073,30 @@ class FlowController: run_id=run.run_id if run else "", on_event=emissions.handle, fallback=( - function if node_def.device_policy == "prefer" else None + self.workers.proxy( + owner, + local, + code, + node_id=node_id, + timeout=timeout, + run_id=run.run_id if run else "", + on_event=emissions.handle, + ) + if node_def.device_policy == "prefer" + else None ), ) + else: + # Declares nothing: the shared pool, at no extra cost. + function = self.workers.proxy( + owner, + local, + code, + node_id=node_id, + timeout=timeout, + run_id=run.run_id if run else "", + on_event=emissions.handle, + ) # Outermost, so it sees the result whichever of the three # above answered the call. function = emissions.wrap(function) diff --git a/backend/fluksio/flow/placement.py b/backend/fluksio/flow/placement.py new file mode 100644 index 0000000..9b53f46 --- /dev/null +++ b/backend/fluksio/flow/placement.py @@ -0,0 +1,390 @@ +"""Which machine a node runs on, and the waiting when none of them has room. + +The engine used to answer this twice, in two ways that could not see each +other. A node with a ``device`` went to a worker carrying that label, chosen by +which of them had least in flight — a machine's *name*, never its size. A node +with ``resources`` was accounted against the engine's own cores, and only the +engine's: declaring both meant the second answer silently won and nothing was +counted at all. + +One question, then, asked once here: of every machine attached — this one and +each worker — which could grant what this node asked for, and which of those +has it free right now. A worker reports its inventory when it attaches, so the +answer covers the whole installation rather than the host the engine happens to +be on. + +Two properties are worth stating because they are what make the waiting safe: + +*Waiting is central.* One condition variable, woken by a release on any machine +and by a worker attaching. A node waiting inside one machine's books could not +be woken by a second machine appearing, which is exactly the case a cluster is +for. + +*Locks go one way.* This condition, then a machine's books, never the reverse — +:meth:`ResourceAccountant.try_take` does not block and +:meth:`ResourceAccountant.give_back` calls back out after dropping its own +lock. The other invariant the engine already had still holds too: the claim is +taken before a worker slot, so a node holding a slot can never be waiting on +resources that a node holding those resources cannot get a slot to release. +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from fluksio.flow.events import EventBus +from fluksio.flow.resources import Allocation, ResourceAccountant +from fluksio.flow.schemas import Resources + +if TYPE_CHECKING: # pragma: no cover - imported for types only + from fluksio.flow.provision import Provisioner + from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub + +logger = logging.getLogger(__name__) + +#: How long a waiting node sleeps before looking again without being woken. +#: Everything that frees a machine wakes it, so this only catches a wakeup lost +#: to a race and gives a provisioned machine's timeout somewhere to be noticed. +WAKE_S = 5.0 + + +@dataclass(frozen=True) +class Target: + """One machine a node could run on: this engine, or an attached worker.""" + + name: str + accountant: ResourceAccountant + #: None is the engine itself, which runs nodes in its own worker pool. + worker: RemoteWorker | None = None + + def carries(self, device: str) -> bool: + """Whether a node asking for this device may run here. + + The engine carries no label: a label names a machine somebody attached + for the purpose, and the engine is where the node would have run anyway. + """ + if self.worker is None: + return False + return device == self.name or device in self.worker.labels + + +@dataclass +class _Waiting: + node: str + since: float + reason: str + + +class Placer: + """Every machine the engine can reach, and who is queued for one.""" + + def __init__( + self, + local: ResourceAccountant, + hub: RemoteWorkerHub | None = None, + events: EventBus | None = None, + ) -> None: + self.local = local + self.hub = hub + self.events = events + self.provisioners: list[Provisioner] = [] + self._condition = threading.Condition() + self._waiting: dict[int, _Waiting] = {} + self._ticket = 0 + local.on_release = self.wake + + # -- the machines ---------------------------------------------------------- + + def wake(self) -> None: + """Something changed: a machine freed up, or one attached or left.""" + for provisioner in self.provisioners: + provisioner.reconcile(self._attached_names()) + with self._condition: + self._condition.notify_all() + + def _attached_names(self) -> set[str]: + if self.hub is None: + return set() + return {worker.name for worker in self.hub.workers() if not worker.gone} + + def targets(self) -> list[Target]: + """This engine, and every worker attached right now.""" + found = [Target(name="engine", accountant=self.local)] + if self.hub is not None: + for worker in self.hub.workers(): + if worker.gone: + continue + # Assigned here rather than at attach so a worker that + # reconnects — a new object, with fresh books — is wired up + # without the hub having to know about placement at all. + worker.accountant.on_release = self.wake + found.append( + Target( + name=worker.name, + accountant=worker.accountant, + worker=worker, + ) + ) + return found + + def _candidates( + self, + cpus: int, + gpus: int, + ram_mb: int, + device: str | None, + policy: str, + ) -> list[Target]: + """The machines that could run this node, best first.""" + targets = self.targets() + remote = [target for target in targets if target.worker is not None] + engine = [target for target in targets if target.worker is None] + + if device: + labelled = [target for target in remote if target.carries(device)] + fitting = self._fitting(labelled, cpus, gpus, ram_mb) + if fitting or policy != "prefer": + # `require` waits for the labelled machine even when the engine + # could take the node: the label is a statement about where + # this node is correct, not a preference. + return fitting + # `prefer` is what makes a flow work before the GPU box exists. It + # falls back only when nothing labelled could ever take the node — + # one that is merely busy is worth queueing for. + return self._fitting(engine, cpus, gpus, ram_mb) + + # No label: the engine first when it fits, because a node that runs + # here costs no network and reuses a warm worker. + ordered = engine + sorted( + remote, key=lambda target: target.worker.in_flight if target.worker else 0 + ) + return self._fitting(ordered, cpus, gpus, ram_mb) + + @staticmethod + def _fitting( + targets: list[Target], cpus: int, gpus: int, ram_mb: int + ) -> list[Target]: + return [t for t in targets if t.accountant.fits(cpus, gpus, ram_mb)] + + # -- what can actually be granted ------------------------------------------ + + def _shapes(self, device: str | None, policy: str) -> list[tuple[int, int, int]]: + """Every machine's size, as something to measure a request against. + + A provisioner's job shapes count: a machine it can start on demand is + one this installation has, even when nothing is attached yet. + """ + shapes = [] + for target in self.targets(): + if device and not target.carries(device): + # A `prefer` node may end up here, so the engine still counts. + if policy == "require" or target.worker is not None: + continue + books = target.accountant + shapes.append((books.cpus, books.gpus, books.ram_mb or 0)) + for provisioner in self.provisioners: + shapes.extend(provisioner.shapes(device)) + return shapes or [(self.local.cpus, self.local.gpus, self.local.ram_mb or 0)] + + def _clamp( + self, wanted: Resources, node: str, device: str | None, policy: str + ) -> tuple[int, int, int]: + """What can be granted somewhere, which may be less than was asked. + + A flow written against a sixty-four core box should still run on a + laptop — waiting forever for cores that do not exist is not a smaller + machine, it is a hung run. What it must not become is an ask no single + machine can grant: taking the largest of each dimension separately can + describe a machine that is not there, and waiting on *that* is the hung + run this exists to avoid. So the GPUs are settled first, and the rest is + measured only against machines that carry that many. + """ + ram = wanted.ram or 0 + shapes = self._shapes(device, policy) + gpus = min(wanted.gpus, max(shape[1] for shape in shapes)) + capable = [shape for shape in shapes if shape[1] >= gpus] + cpus = min(wanted.cpus, max(shape[0] for shape in capable)) + # A machine that reported no memory is not a memory limit; only those + # that did have anything to say about it. + said = [shape[2] for shape in capable if shape[2]] + ram = min(ram, max(said)) if said and ram else ram + if (cpus, gpus, ram) != (wanted.cpus, wanted.gpus, wanted.ram or 0): + logger.warning( + "%s asked for %d cpu(s), %d gpu(s) and %s MB; " + "the largest machine here can give %d, %d and %s", + node or "a node", + wanted.cpus, + wanted.gpus, + wanted.ram or "no stated", + cpus, + gpus, + ram or "no stated", + ) + return cpus, gpus, ram + + # -- holding a machine for the length of a call ---------------------------- + + @contextmanager + def claim( + self, + wanted: Resources, + *, + device: str | None = None, + policy: str = "require", + node: str = "", + run: str = "", + ) -> Iterator[tuple[Target, Allocation]]: + """Hold a machine's share of itself for as long as this node runs.""" + waited_from = time.monotonic() + ticket = 0 + chosen: Target | None = None + allocation: Allocation | None = None + + with self._condition: + while chosen is None: + cpus, gpus, ram = self._clamp(wanted, node, device, policy) + for target in self._candidates(cpus, gpus, ram, device, policy): + got = target.accountant.try_take(cpus, gpus, ram) + if got is not None: + chosen, allocation = target, got + break + if chosen is not None: + break + if not ticket: + ticket = self._announce( + node, run, waited_from, self._reason(cpus, gpus, device) + ) + self._provision(cpus, gpus, ram, device) + self._condition.wait(timeout=WAKE_S) + if ticket: + self._waiting.pop(ticket, None) + + if ticket: + logger.info( + "%s waited %.1fs and runs on '%s'", + node or "a node", + time.monotonic() - waited_from, + chosen.name, + ) + assert allocation is not None + try: + yield chosen, allocation + finally: + # Every way a call can end comes through here: an answer, a + # timeout, a cancellation, or the worker going away mid-call. + chosen.accountant.give_back(allocation) + + def _reason(self, cpus: int, gpus: int, device: str | None) -> str: + parts = [] + if gpus: + parts.append(f"{gpus} gpu(s)") + parts.append(f"{cpus} cpu(s)") + wants = " and ".join(parts) + if device: + return f"waiting for a worker labelled '{device}' with {wants} free" + return f"waiting for {wants}" + + def _provision(self, cpus: int, gpus: int, ram_mb: int, device: str | None) -> None: + """Ask for a machine, if something here can start one.""" + for provisioner in self.provisioners: + if provisioner.covers(cpus, gpus, ram_mb, device): + provisioner.provision(cpus, gpus, ram_mb, device) + return + + def _announce(self, node: str, run: str, since: float, reason: str) -> int: + """Say a node is queued, not stuck. + + The failure this exists for looked identical to a hang: a run sitting + at `running` for twenty minutes with no error and no output. A node + waiting its turn has to say so somewhere a person will look. + """ + self._ticket += 1 + ticket = self._ticket + self._waiting[ticket] = _Waiting(node=node, since=since, reason=reason) + logger.info("%s is %s", node or "a node", reason) + if self.events is not None: + self.events.publish( + { + "type": "node_queued", + "flow": node.split(".", 1)[0] if node else "", + "node": node, + "run": run, + "detail": reason, + "ts": time.time(), + } + ) + return ticket + + # -- what it looks like from outside --------------------------------------- + + def satisfiable(self, needs: dict[str, Any] | None) -> str | None: + """What a run would have to wait for, or None if it would not wait. + + Only a genuine wait is worth holding a run for. A request nothing can + grant is cut down to what is here and runs anyway, so answering "no + machine has 64 cores" would hold a run that was about to work. + """ + if not needs: + return None + cpus = int(needs.get("cpus") or 1) + gpus = int(needs.get("gpus") or 0) + ram = int(needs.get("ram_mb") or 0) + device = needs.get("device") or None + for target in self.targets(): + if device and not target.carries(device): + continue + if target.accountant.fits(cpus, gpus, ram): + # Something here could take it. Busy is queueing, not waiting. + return None + if not any( + provisioner.covers(cpus, gpus, ram, device) + for provisioner in self.provisioners + ): + # Nothing attached fits and nothing can be started, so the node + # will be clamped onto what is here and run. Holding the run would + # be holding it for something that is not going to happen — and + # where the wait is a missing *label*, `required_labels` says so + # already. + return None + return f"{gpus} gpu(s) and {cpus} cpu(s)" if gpus else f"{cpus} cpu(s)" + + def snapshot(self) -> dict[str, Any]: + """Every machine, what is free of it, and who is queued. + + A node waiting its turn looks exactly like a node that has hung — the + run sits at `running` and says nothing — so what is waiting, and for + what, has to be readable somewhere. + """ + targets = [] + for target in self.targets(): + entry: dict[str, Any] = {"target": target.name} + entry.update(target.accountant.snapshot()) + if target.worker is not None: + entry["labels"] = sorted(target.worker.labels) + entry["in_flight"] = target.worker.in_flight + targets.append(entry) + with self._condition: + waiting = list(self._waiting.values()) + local = self.local.snapshot() + return { + # The engine's own figures stay where they were, so a reader that + # only ever knew about one machine still finds them. + "cpus": local["cpus"], + "gpus": local["gpus"], + "waiting": [ + { + "node": entry.node, + "reason": entry.reason, + "seconds": round(time.monotonic() - entry.since, 1), + } + for entry in waiting + ], + "targets": targets, + "provisioners": [p.status() for p in self.provisioners], + } diff --git a/backend/fluksio/flow/remote.py b/backend/fluksio/flow/remote.py index 526872f..30678ca 100644 --- a/backend/fluksio/flow/remote.py +++ b/backend/fluksio/flow/remote.py @@ -32,6 +32,7 @@ from dataclasses import dataclass from typing import Any from fluksio.flow.nodes.base import NodeOutputError +from fluksio.flow.resources import ResourceAccountant from fluksio.flow.workers import NodeTimeout, RemoteError, _remote_class logger = logging.getLogger(__name__) @@ -106,6 +107,14 @@ class RemoteWorker: self.labels = set(labels) self.info = info or {} self.inventory = inventory or WorkerInventory() + # This machine's books, held here so that a worker reconnecting — which + # is a new object — starts from a clean set rather than from counters + # the engine kept for a socket that is gone. + self.accountant = ResourceAccountant( + cpus=self.inventory.cpus, + gpus=self.inventory.gpus, + ram_mb=self.inventory.ram_mb, + ) self.attached_at = time.time() self.last_seen = time.time() self._send = send @@ -251,9 +260,13 @@ class RemoteWorker: class RemoteWorkerHub: """Every attached worker, and which of them a node may run on.""" - def __init__(self) -> None: + def __init__(self, on_change: Callable[[], None] | None = None) -> None: self._workers: dict[str, RemoteWorker] = {} self._lock = threading.Lock() + #: Told when the set of attached workers changes, so whoever is waiting + #: for a machine can look again. Called outside the lock: it takes one + #: of its own, and the two are only ever taken in that order. + self.on_change = on_change # ------------------------------------------------------------------------- # Attachment @@ -270,6 +283,7 @@ class RemoteWorkerHub: logger.info( "Worker '%s' attached with labels %s", worker.name, sorted(worker.labels) ) + self._changed() def detach(self, name: str) -> None: with self._lock: @@ -277,6 +291,11 @@ class RemoteWorkerHub: if worker is not None: worker.detach() logger.info("Worker '%s' detached", name) + self._changed() + + def _changed(self) -> None: + if self.on_change is not None: + self.on_change() def workers(self) -> list[RemoteWorker]: with self._lock: @@ -326,6 +345,38 @@ class RemoteWorkerHub: worker = self.pick(label) if worker is None: raise NoWorker(f"no worker labelled '{label}' is attached") + return self.run_on( + worker, + flow, + node, + source, + kwargs, + node_id, + timeout, + run_id=run_id, + on_event=on_event, + env=env, + ) + + def run_on( + self, + worker: RemoteWorker, + flow: str, + node: str, + source: str, + kwargs: dict[str, Any], + node_id: str, + timeout: float, + run_id: str = "", + on_event: Callable[[dict[str, Any]], None] | None = None, + env: dict[str, str] | None = None, + ) -> Any: + """Run this node on the worker the caller has already settled on. + + The placer holds that machine's cores for the length of the call, so + resolving the label a second time here could send the work somewhere + else and leave the claim on a machine doing nothing. + """ payload = { "op": "run", "call_id": f"{run_id}:{node_id}" if run_id else node_id, diff --git a/backend/fluksio/flow/resources.py b/backend/fluksio/flow/resources.py index 413906a..3c03b07 100644 --- a/backend/fluksio/flow/resources.py +++ b/backend/fluksio/flow/resources.py @@ -8,16 +8,18 @@ seconds. On a GPU the same shape is worse: three processes each preallocating most of the card deadlock at zero utilisation, with nothing failing and nothing to read. -:class:`ResourceAccountant` is the first half: a node that declares what it -needs waits until that much is free, the way it already waits for a worker -slot. :func:`derive_env` is the second: what it got is handed to the library as -the environment it reads at import, because that is the only moment those knobs -are read. +:class:`ResourceAccountant` is the first half: the books for one machine, which +hand out what is free and take it back. :func:`derive_env` is the second: what +a node got is handed to the library as the environment it reads at import, +because that is the only moment those knobs are read. + +One accountant is one machine. Which machine a node goes to, and the waiting +when none of them has room, is :mod:`fluksio.flow.placement` — the books here +never block, so a caller holding several of them can ask each in turn. Deliberately cooperative — nothing here is enforced with cgroups or rlimits, so a node that ignores its share is only accounted for, not stopped. That is the -same trust the worker pool already extends to node code, and the enforcement -half belongs with the scheduler work this is the first step of. +same trust the worker pool already extends to node code. """ from __future__ import annotations @@ -25,12 +27,9 @@ from __future__ import annotations import logging import os import threading -import time -from collections.abc import Iterator -from contextlib import contextmanager +from collections.abc import Callable from dataclasses import dataclass -from fluksio.flow.events import EventBus from fluksio.flow.schemas import Resources logger = logging.getLogger(__name__) @@ -63,146 +62,92 @@ class Allocation: cpus: int = 1 gpus: tuple[int, ...] = () - - -@dataclass -class _Waiting: - node: str - since: float - reason: str + ram_mb: int = 0 class ResourceAccountant: - """What is free on this machine, and who is waiting for it. + """The books for one machine: what it holds, and what is free of it. - A plain condition variable over two counters. Blocking is the whole - mechanism: the caller is a node thread that would otherwise be executing, - and making it wait is the same backpressure that the worker pool's slot - queue already applies — one reason the two must always be taken in the same - order, resources first, so neither can be held while waiting for the other. + Nothing here blocks. A node that has to wait waits in the placer, which + holds one of these per machine and tries each in turn — a lock per machine + could not answer "is there room anywhere", and a node waiting inside one + machine's books could not be woken by another machine attaching. + + ``on_release`` is how it tells the placer something came free. It is called + outside this object's lock, deliberately: the placer takes its own lock + there, and the two are only ever taken in that one order. """ def __init__( - self, cpus: int = 0, gpus: int = 0, events: EventBus | None = None + self, + cpus: int = 0, + gpus: int = 0, + ram_mb: int | None = None, + on_release: Callable[[], None] | None = None, ) -> None: self.cpus = cpus or machine_cpus() self.gpus = max(0, gpus) - self.events = events + #: None is memory this machine did not report, which is allowed: it is + #: then not accounted rather than accounted as zero, since zero would + #: refuse every node that asks for any. + self.ram_mb = ram_mb + self.on_release = on_release self._free_cpus = self.cpus self._free_gpus = list(range(self.gpus)) - self._waiting: dict[int, _Waiting] = {} - self._ticket = 0 - self._condition = threading.Condition() + self._free_ram = ram_mb or 0 + self._lock = threading.Lock() - # -- what a caller asks for ------------------------------------------------ + def fits(self, cpus: int, gpus: int, ram_mb: int = 0) -> bool: + """Whether this machine could ever grant that much, busy or not. - def _clamp(self, wanted: Resources, node: str) -> tuple[int, int]: - """What can be granted here, which may be less than what was asked. - - A flow written against a sixty-four core box should still run on a - laptop. Waiting forever for cores that do not exist is not a smaller - machine, it is a hung run. + The question a placer asks before it decides to wait: waiting for a + card that is merely busy is queueing, and waiting for one the machine + does not have is a hung run. """ - cpus = min(wanted.cpus, self.cpus) - gpus = min(wanted.gpus, self.gpus) - if cpus != wanted.cpus or gpus != wanted.gpus: - logger.warning( - "%s asked for %d cpu(s) and %d gpu(s); this engine has %d and %d", - node or "a node", - wanted.cpus, - wanted.gpus, - self.cpus, - self.gpus, - ) - return cpus, gpus + if cpus > self.cpus or gpus > self.gpus: + return False + return self.ram_mb is None or ram_mb <= self.ram_mb - @contextmanager - def claim( - self, wanted: Resources, node: str = "", run: str = "" - ) -> Iterator[Allocation]: - """Hold this node's share for as long as it runs.""" - cpus, gpus = self._clamp(wanted, node) - allocation = self._take(cpus, gpus, node, run) - try: - yield allocation - finally: - self._give_back(allocation) - - def _take(self, cpus: int, gpus: int, node: str, run: str) -> Allocation: - waited_from = time.monotonic() - ticket = 0 - with self._condition: - while self._free_cpus < cpus or len(self._free_gpus) < gpus: - if not ticket: - reason = ( - f"waiting for {cpus} cpu(s) ({self._free_cpus} free)" - if self._free_cpus < cpus - else f"waiting for {gpus} gpu(s) ({len(self._free_gpus)} free)" - ) - ticket = self._announce(node, run, waited_from, reason) - self._condition.wait() + def try_take(self, cpus: int, gpus: int, ram_mb: int = 0) -> Allocation | None: + """Take this much if it is free right now, or answer that it is not.""" + with self._lock: + if self._free_cpus < cpus or len(self._free_gpus) < gpus: + return None + if self.ram_mb is not None and self._free_ram < ram_mb: + return None self._free_cpus -= cpus taken = tuple(self._free_gpus.pop(0) for _ in range(gpus)) - if ticket: - self._waiting.pop(ticket, None) - if ticket: - logger.info( - "%s waited %.1fs for %d cpu(s) and %d gpu(s)", - node or "a node", - time.monotonic() - waited_from, - cpus, - gpus, - ) - return Allocation(cpus=cpus, gpus=taken) + if self.ram_mb is not None: + self._free_ram -= ram_mb + return Allocation(cpus=cpus, gpus=taken, ram_mb=ram_mb) - def _give_back(self, allocation: Allocation) -> None: - with self._condition: + def give_back(self, allocation: Allocation) -> None: + """Return what an execution held, however it ended.""" + with self._lock: self._free_cpus += allocation.cpus self._free_gpus.extend(allocation.gpus) self._free_gpus.sort() - self._condition.notify_all() - - # -- what it looks like from outside --------------------------------------- - - def _announce(self, node: str, run: str, since: float, reason: str) -> int: - """Say a node is queued, not stuck. - - The failure this exists for looked identical to a hang: a run sitting - at `running` for twenty minutes with no error and no output. A node - waiting its turn has to say so somewhere a person will look. - """ - self._ticket += 1 - ticket = self._ticket - self._waiting[ticket] = _Waiting(node=node, since=since, reason=reason) - logger.info("%s is %s", node or "a node", reason) - if self.events is not None: - self.events.publish( - { - "type": "node_queued", - "flow": node.split(".", 1)[0] if node else "", - "node": node, - "run": run, - "detail": reason, - "ts": time.time(), - } - ) - return ticket + if self.ram_mb is not None: + self._free_ram += allocation.ram_mb + if self.on_release is not None: + self.on_release() def snapshot(self) -> dict[str, object]: - """What is free and who is waiting, for the workers screen.""" - with self._condition: - free_cpus, free_gpus = self._free_cpus, len(self._free_gpus) + """What this machine holds and what is free of it.""" + with self._lock: + free_cpus, free_gpus, free_ram = ( + self._free_cpus, + len(self._free_gpus), + self._free_ram, + ) return { "cpus": {"total": self.cpus, "free": free_cpus}, "gpus": {"total": self.gpus, "free": free_gpus}, - "waiting": [ - { - "node": entry.node, - "reason": entry.reason, - "seconds": round(time.monotonic() - entry.since, 1), - } - for entry in list(self._waiting.values()) - ], + "ram_mb": ( + None + if self.ram_mb is None + else {"total": self.ram_mb, "free": free_ram} + ), } diff --git a/backend/fluksio/flow/runs.py b/backend/fluksio/flow/runs.py index 811ea3b..2f5094a 100644 --- a/backend/fluksio/flow/runs.py +++ b/backend/fluksio/flow/runs.py @@ -160,6 +160,17 @@ def _has_no_body( return not store.has_node_source(flow, node.id, draft=draft) +def _runs_on_a_worker(node: NodeDef) -> bool: + """Whether this node's body could go to a worker at all. + + Only a node whose source travels can: a connector is an entry point loaded + in this process, so a device on one is a field nothing reads. Holding a run + for a worker that could never take that node is a run that never starts. + """ + node_type = NODE_TYPES.get(node.type) + return node_type is not None and node_type.has_source + + def required_labels(flow: FlowDef) -> list[str]: """Worker labels this flow cannot run without. @@ -171,11 +182,38 @@ def required_labels(flow: FlowDef) -> list[str]: { node.device for node in flow.nodes - if node.device and node.device.strip() and node.device_policy == "require" + if node.device + and node.device.strip() + and node.device_policy == "require" + and _runs_on_a_worker(node) } ) +def required_resources(flow: FlowDef) -> dict[str, Any] | None: + """The largest single thing this flow needs a machine to have. + + Dimension by dimension rather than per node, which is deliberately blunt: + it answers "could this run start at all", not "in what order". A node that + only *prefers* its device is left out — it runs here when nothing carries + the label, so it is never what a run is waiting for. + """ + cpus, gpus, ram = 0, 0, 0 + device = "" + for node in flow.nodes: + wanted = node.resources + if wanted is None or not _runs_on_a_worker(node): + continue + cpus = max(cpus, wanted.cpus) + gpus = max(gpus, wanted.gpus) + ram = max(ram, wanted.ram or 0) + if node.device and node.device_policy == "require": + device = node.device + if not (cpus or gpus or ram): + return None + return {"cpus": cpus, "gpus": gpus, "ram_mb": ram, "device": device} + + #: What a run's output is called from outside it: ``@run:.``. RUN_REF_PREFIX = "@run:" @@ -759,6 +797,7 @@ class RunService: no_cache=no_cache, status="queued", labels=required_labels(flow), + needs=required_resources(flow), created_at=datetime.now(UTC), actor=actor, idempotency_key=idempotency_key, @@ -854,7 +893,7 @@ class RunService: if not item.run_id: self.queue.ack(item) continue - missing = self._missing_labels(item.run_id) + missing = self._missing_requirements(item.run_id) if missing: # Left in the queue rather than failed: submitting a run # before turning the GPU box on is a normal way to work, and @@ -870,21 +909,30 @@ class RunService: except RuntimeError: logger.warning("Run %s not started: shutting down", item.run_id) - def _missing_labels(self, run_id: str) -> list[str]: - """Worker labels this run needs that nothing attached carries.""" + def _missing_requirements(self, run_id: str) -> list[str]: + """What this run needs that nothing attached can give it yet.""" with Session(db_engine) as session: run = session.get(Run, run_id) needed = list(run.labels) if run else [] - if not needed: - return [] + needs = dict(run.needs) if run and run.needs else None hub = self.controller.remote available = hub.labels() | {w.name for w in hub.workers()} if hub else set() # A node that only prefers its label runs locally instead, so it is not # a reason to hold the run back; that is decided per node at call time. - return sorted(set(needed) - available) + missing = [f"labelled {label}" for label in sorted(set(needed) - available)] + placer = self.controller.placer + if placer is not None: + # Only a wait something is going to end — a machine that could take + # this is being started for it. An ask nothing can grant is cut down + # to what is here and runs, so holding the run would be holding it + # for something that is not coming. + short = placer.satisfiable(needs) + if short: + missing.append(f"with {short}") + return missing def _waiting(self, run_id: str, missing: list[str]) -> None: - reason = f"Waiting for a worker labelled {', '.join(missing)}" + reason = f"Waiting for a worker {' and '.join(missing)}" try: with Session(db_engine) as session: session.exec( diff --git a/backend/fluksio/flow/schemas.py b/backend/fluksio/flow/schemas.py index 535e3b9..0a0f702 100644 --- a/backend/fluksio/flow/schemas.py +++ b/backend/fluksio/flow/schemas.py @@ -60,6 +60,16 @@ class Resources(BaseModel): "runs, which is what keeps two preallocating processes apart." ), ) + ram: int | None = Field( + default=None, + ge=1, + description=( + "Megabytes held for the whole execution. Counted against machines " + "that said how much they have, and ignored by those that did not — " + "which is a machine with nothing to say about memory, not one with " + "none." + ), + ) env: dict[str, str] = Field( default_factory=dict, description=( diff --git a/backend/fluksio/main.py b/backend/fluksio/main.py index e1c9def..e4186bd 100644 --- a/backend/fluksio/main.py +++ b/backend/fluksio/main.py @@ -29,6 +29,7 @@ from fluksio.flow.executor import ExecutionService from fluksio.flow.metrics import MetricsCollector from fluksio.flow.nodes.http import close_shared_client from fluksio.flow.pipeline import ValueSource +from fluksio.flow.placement import Placer from fluksio.flow.plugins import load_plugins from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue from fluksio.flow.remote import RemoteWorkerHub @@ -142,10 +143,11 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # not something anyone wrote, so it has no business in the git repository. artifacts = ArtifactStore(settings.FLOWS_DIR.parent / "artifacts") app.state.artifact_store = artifacts - accountant = ResourceAccountant( - cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS, events=event_bus - ) + accountant = ResourceAccountant(cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS) app.state.resources = accountant + # Every machine a node could run on: this one, and whatever attaches. + placer = Placer(local=accountant, events=event_bus) + app.state.placer = placer pool = PythonWorkerPool( python=modules.venv_python(), size=settings.FLOW_MAX_WORKERS, @@ -163,7 +165,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: ) pool.start() app.state.worker_pool = pool - worker_hub = RemoteWorkerHub() + # Assigned rather than passed both ways: the hub tells the placer when a + # machine comes or goes, and the placer needs the hub to know what is there. + worker_hub = RemoteWorkerHub(on_change=placer.wake) + placer.hub = worker_hub app.state.worker_hub = worker_hub controller = FlowController( store=store, @@ -176,6 +181,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: workers=pool, remote=worker_hub, resources=accountant, + placer=placer, ) app.state.flow_controller = controller # A "dashboard" alert channel puts its alert into the graph. Bound here diff --git a/backend/fluksio/models.py b/backend/fluksio/models.py index 91a5d45..21eb296 100644 --- a/backend/fluksio/models.py +++ b/backend/fluksio/models.py @@ -362,6 +362,10 @@ class Run(SQLModel, table=True): status_reason: str = Field(default="", max_length=1024) #: Worker labels its nodes need, so a run with nowhere to go can say so. labels: list[str] = Field(sa_column=Column(JSON), default_factory=list) + #: The largest machine its nodes ask for, as cpus, gpus, ram_mb and device. + #: Read while the run is queued, to tell waiting for a machine from having + #: nowhere to run at all. + needs: dict[str, Any] | None = Field(sa_column=Column(JSON), default=None) created_at: datetime = Field( index=True, sa_type=UTCDateTime, diff --git a/backend/scripts/bench_startup.py b/backend/scripts/bench_startup.py index d2f95bf..5b91065 100644 --- a/backend/scripts/bench_startup.py +++ b/backend/scripts/bench_startup.py @@ -37,7 +37,7 @@ log = logging.getLogger("bench") #: A flow that computes nothing, so what is timed is the getting-started. BENCH_FLOW = "bench_startup" -NODE_SOURCE = '"""Does nothing, on purpose."""\n\n\ndef process(n, params):\n return {"out": n}\n' +NODE_SOURCE = '"""Does nothing, on purpose."""\n\n\ndef process(n):\n return {"out": n}\n' class Api: diff --git a/backend/tests/flow/test_placement.py b/backend/tests/flow/test_placement.py new file mode 100644 index 0000000..d5edeca --- /dev/null +++ b/backend/tests/flow/test_placement.py @@ -0,0 +1,394 @@ +"""Which machine a node runs on, and what happens when none of them is free. + +The two questions the engine used to answer separately — where does a node run, +and how much of that machine does it get — are one question here. What is +checked is that the answer covers every machine attached rather than only the +host the engine is on, that a claim is given back however the call ends, and +that a node with nowhere to go waits and says so instead of failing. +""" + +import asyncio +import threading +import time +from collections.abc import Iterator + +import pytest + +from fluksio.flow.placement import Placer +from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub, WorkerInventory +from fluksio.flow.resources import ResourceAccountant +from fluksio.flow.schemas import NodeDef, Resources + + +@pytest.fixture +def loop() -> Iterator[asyncio.AbstractEventLoop]: + """An event loop running in a thread, as the server's does.""" + running = asyncio.new_event_loop() + thread = threading.Thread(target=running.run_forever, daemon=True) + thread.start() + yield running + running.call_soon_threadsafe(running.stop) + thread.join(timeout=5) + running.close() + + +class FakeSocket: + def __init__(self) -> None: + self.sent: list[dict] = [] + self.arrived = threading.Event() + + async def send_json(self, payload: dict) -> None: + self.sent.append(payload) + self.arrived.set() + + +def attach( + hub: RemoteWorkerHub, + loop: asyncio.AbstractEventLoop, + name: str = "gpu1", + labels: tuple[str, ...] = ("gpu",), + cpus: int = 8, + gpus: int = 1, + ram_mb: int | None = None, +): + socket = FakeSocket() + worker = RemoteWorker( + name=name, + labels=list(labels), + send=socket.send_json, + loop=loop, + max_parallel=4, + inventory=WorkerInventory(cpus=cpus, gpus=gpus, ram_mb=ram_mb), + ) + hub.attach(worker) + return worker, socket + + +def placer_over(hub: RemoteWorkerHub | None = None, cpus: int = 4, gpus: int = 0): + return Placer(local=ResourceAccountant(cpus=cpus, gpus=gpus), hub=hub) + + +# ----------------------------------------------------------------------------- +# Where a node goes +# ----------------------------------------------------------------------------- + + +def test_a_node_asking_for_a_card_finds_the_machine_that_has_one(loop): + """No label anywhere — the size of the ask is what places it. + + This is the whole point of a worker reporting its inventory: the engine has + no GPU, so a node that needs one used to be clamped down to none and run + here anyway. + """ + hub = RemoteWorkerHub() + placer = placer_over(hub) + worker, _ = attach(hub, loop, labels=()) + + with placer.claim(Resources(cpus=2, gpus=1)) as (target, allocation): + assert target.worker is worker + assert allocation.gpus == (0,) + assert worker.accountant.snapshot()["gpus"]["free"] == 0 + + assert worker.accountant.snapshot()["gpus"]["free"] == 1 + + +def test_a_node_that_fits_here_stays_here(loop): + """No network, and a worker pool that is already warm.""" + hub = RemoteWorkerHub() + placer = placer_over(hub) + attach(hub, loop) + + with placer.claim(Resources(cpus=2)) as (target, _): + assert target.worker is None + assert target.name == "engine" + + +def test_a_required_label_waits_for_its_machine_rather_than_running_here(loop): + hub = RemoteWorkerHub() + placer = placer_over(hub) + worker, _ = attach(hub, loop, labels=("gpu",)) + + with placer.claim(Resources(cpus=2), device="gpu", policy="require") as ( + target, + _, + ): + assert target.worker is worker + + +def test_a_preferred_label_falls_back_here_and_is_still_accounted(loop): + """The fallback used to skip the books entirely, so it held nothing.""" + hub = RemoteWorkerHub() + placer = placer_over(hub) + + with placer.claim(Resources(cpus=3), device="gpu", policy="prefer") as ( + target, + allocation, + ): + assert target.worker is None + assert allocation.cpus == 3 + assert placer.local.snapshot()["cpus"]["free"] == 1 + + +def test_asking_for_more_than_anything_has_gets_what_there_is(loop): + """A flow written on a cluster still has to run on a laptop.""" + placer = placer_over(cpus=2) + + with placer.claim(Resources(cpus=64, gpus=4)) as (target, allocation): + assert target.name == "engine" + assert allocation.cpus == 2 + assert allocation.gpus == () + + +def test_what_is_clamped_to_is_a_machine_that_exists(loop): + """Each dimension taken separately can describe a machine nobody has. + + A big CPU box and a small GPU box are not one machine with both. Clamping + to the largest of each would ask for something that never comes free, which + is the hung run the clamping is there to avoid. + """ + hub = RemoteWorkerHub() + placer = placer_over(hub, cpus=2) + attach(hub, loop, name="big", labels=(), cpus=64, gpus=0) + attach(hub, loop, name="small-gpu", labels=(), cpus=4, gpus=1) + + with placer.claim(Resources(cpus=64, gpus=1)) as (target, allocation): + assert target.name == "small-gpu" + assert (allocation.cpus, allocation.gpus) == (4, (0,)) + + +# ----------------------------------------------------------------------------- +# Waiting +# ----------------------------------------------------------------------------- + + +def test_a_node_waits_for_a_busy_machine_and_says_so(loop): + hub = RemoteWorkerHub() + events: list[dict] = [] + placer = Placer( + local=ResourceAccountant(cpus=2), + hub=hub, + events=type("Bus", (), {"publish": lambda self, e: events.append(e)})(), + ) + started = threading.Event() + + def second() -> None: + with placer.claim(Resources(cpus=2), node="study.b"): + started.set() + + with placer.claim(Resources(cpus=2), node="study.a"): + thread = threading.Thread(target=second, daemon=True) + thread.start() + # Long enough to have run if nothing was holding it back. + assert not started.wait(0.3) + assert placer.snapshot()["waiting"][0]["node"] == "study.b" + assert events[0]["type"] == "node_queued" + assert "2 cpu(s)" in events[0]["detail"] + + thread.join(timeout=5) + assert started.is_set() + assert placer.snapshot()["waiting"] == [] + + +def test_a_worker_attaching_wakes_what_was_waiting_for_one(loop): + """A node waiting inside one machine's books could not hear this.""" + hub = RemoteWorkerHub() + placer = placer_over(hub) + hub.on_change = placer.wake + ran = threading.Event() + + def needs_a_card() -> None: + with placer.claim(Resources(cpus=1, gpus=1), node="study.fit") as (target, _): + if target.worker is not None: + ran.set() + + thread = threading.Thread(target=needs_a_card, daemon=True) + thread.start() + # Nothing has a GPU yet, so it is clamped to none and runs here — the wait + # only happens while a machine that *could* take it is busy. + thread.join(timeout=5) + assert not thread.is_alive() + + worker, _ = attach(hub, loop, labels=()) + held = worker.accountant.try_take(cpus=1, gpus=1) + thread = threading.Thread(target=needs_a_card, daemon=True) + thread.start() + assert not ran.wait(0.3) + + worker.accountant.give_back(held) + thread.join(timeout=5) + assert ran.is_set() + + +def test_everything_comes_back_when_a_node_fails(loop): + placer = placer_over(cpus=4, gpus=1) + with pytest.raises(ValueError): + with placer.claim(Resources(cpus=4, gpus=1)): + raise ValueError("the node raised") + + assert placer.local.snapshot()["cpus"]["free"] == 4 + assert placer.local.snapshot()["gpus"]["free"] == 1 + + +def test_a_worker_going_away_mid_call_gives_its_share_back(loop): + hub = RemoteWorkerHub() + placer = placer_over(hub) + worker, _ = attach(hub, loop, labels=()) + + with pytest.raises(RuntimeError): + with placer.claim(Resources(cpus=2, gpus=1)) as (target, _): + assert target.worker is worker + hub.detach(worker.name) + raise RuntimeError("worker went away mid-call") + + # The books are the worker's own, so they go with it — what matters is that + # nothing is left holding them and the next claim is not blocked. + assert worker.accountant.snapshot()["gpus"]["free"] == 1 + + +def test_many_nodes_at_once_all_finish(loop): + """The placer must not deadlock under contention; it is on every call.""" + placer = placer_over(cpus=4) + done = [] + + def work() -> None: + with placer.claim(Resources(cpus=2), node="study.n"): + time.sleep(0.01) + done.append(1) + + threads = [threading.Thread(target=work) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert len(done) == 8 + assert placer.local.snapshot()["cpus"]["free"] == 4 + + +# ----------------------------------------------------------------------------- +# What a run is told before it starts +# ----------------------------------------------------------------------------- + + +def test_a_run_is_only_held_for_a_wait_that_will_end(loop): + """An ask nothing can grant is cut down and runs; holding it would hang.""" + placer = placer_over(cpus=4) + + assert placer.satisfiable(None) is None + assert placer.satisfiable({"cpus": 2}) is None + # No machine here has a GPU and nothing can be started with one. + assert placer.satisfiable({"cpus": 1, "gpus": 1}) is None + + +def test_a_run_waits_for_a_machine_something_can_start(loop): + class Cluster: + def covers(self, cpus, gpus, ram_mb, label=None): + return gpus <= 2 + + def provision(self, cpus, gpus, ram_mb, label=None): + pass + + def reconcile(self, attached): + pass + + def shapes(self, label=None): + return [(16, 2, 65536)] + + def status(self): + return {"name": "hpc"} + + placer = placer_over(cpus=4) + placer.provisioners = [Cluster()] + + assert placer.satisfiable({"cpus": 1, "gpus": 1}) == "1 gpu(s) and 1 cpu(s)" + assert placer.satisfiable({"cpus": 1, "gpus": 0}) is None + + +# ----------------------------------------------------------------------------- +# The whole path, through a real worker +# ----------------------------------------------------------------------------- + + +def test_a_declared_node_is_held_to_its_share(tmp_path): + """Four nodes wanting two cores each, on a machine with four. + + What is checked is the pair: never more than the machine has in flight at + once, and every one of them told what it was given — the two halves that + together are the oversubscription this is for. Also that it does not + deadlock, since the claim is taken before a worker slot and both block. + """ + import sys + + from fluksio.flow.controller import FlowController, RunContext + from fluksio.flow.messages import DType, MessageSpec + from fluksio.flow.schemas import FlowDef + from fluksio.flow.state import MemoryState + from fluksio.flow.store import FlowStore + from fluksio.flow.workers import PythonWorkerPool + + store = FlowStore(tmp_path / "flows") + store.write_flow( + FlowDef( + name="study", + mode="batch", + nodes=[ + NodeDef( + id="fit", + provides=[MessageSpec(name="threads", dtype=DType.INT)], + resources=Resources(cpus=2, env={"XLA_FLAGS": "--x=false"}), + ) + ], + ) + ) + store.write_node_source( + "study", + "fit", + "import os, time\n\n\ndef process():\n" + " time.sleep(0.2)\n" + " return {'threads': int(os.environ['OMP_NUM_THREADS'])}\n", + ) + + pool = PythonWorkerPool(python=sys.executable, size=4) + pool.start() + accountant = ResourceAccountant(cpus=4) + placer = Placer(local=accountant) + controller = FlowController( + store, workers=pool, resources=accountant, placer=placer + ) + answers: list[object] = [] + in_flight: list[int] = [] + + def once(index: int) -> None: + pipeline = controller.build_run_pipeline( + store.read_flow("study"), + state=MemoryState(), + run=RunContext(run_id=f"r{index}"), + ) + pipeline.run({}) + answers.append(pipeline.state.get("study.threads")) + + def watch(until: threading.Event) -> None: + while not until.is_set(): + in_flight.append(4 - int(accountant.snapshot()["cpus"]["free"])) + time.sleep(0.01) + + finished = threading.Event() + watcher = threading.Thread(target=watch, args=(finished,)) + watcher.start() + threads = [threading.Thread(target=once, args=(i,)) for i in range(4)] + try: + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + assert not thread.is_alive(), "a claim and a worker slot deadlocked" + finally: + finished.set() + watcher.join(timeout=5) + pool.stop() + + assert answers == [2, 2, 2, 2] + assert max(in_flight) <= 4 + assert accountant.snapshot()["cpus"]["free"] == 4 + # One environment, so one extra pool however many nodes derived it. + assert len(pool._children) == 1 diff --git a/backend/tests/flow/test_resources.py b/backend/tests/flow/test_resources.py index 30551ad..7a043f6 100644 --- a/backend/tests/flow/test_resources.py +++ b/backend/tests/flow/test_resources.py @@ -1,15 +1,14 @@ -"""Who gets the machine: accounting, exclusivity, and what the worker is told. +"""One machine's books: what it holds, and what is free of it. The failure this exists for is not subtle — five concurrent nodes each sizing a thread pool to every core starved the engine's own event loop, and three GPU processes each preallocating most of the card deadlocked at zero utilisation. Both come down to arithmetic nobody was doing, so the arithmetic is what is -checked here. +checked here. Which machine a node goes to, and the waiting, is +``test_placement.py``. """ import os -import threading -import time import pytest @@ -25,79 +24,70 @@ from fluksio.flow.schemas import NodeDef, Resources def test_what_is_free_is_what_was_handed_out(): accountant = ResourceAccountant(cpus=4, gpus=0) - with accountant.claim(Resources(cpus=3)): - assert accountant.snapshot()["cpus"] == {"total": 4, "free": 1} + held = accountant.try_take(cpus=3, gpus=0) + + assert accountant.snapshot()["cpus"] == {"total": 4, "free": 1} + accountant.give_back(held) assert accountant.snapshot()["cpus"] == {"total": 4, "free": 4} -def test_a_node_waits_until_there_is_room(): - """Blocking is the mechanism — the same backpressure a worker slot applies.""" +def test_asking_for_more_than_is_free_is_answered_not_waited_on(): + """The books never block: the placer has other machines to try first.""" accountant = ResourceAccountant(cpus=2) - running = threading.Event() - started = threading.Event() - - def second() -> None: - with accountant.claim(Resources(cpus=2), node="study.b"): - started.set() - - with accountant.claim(Resources(cpus=2), node="study.a"): - thread = threading.Thread(target=second) - thread.start() - # Long enough to have run if nothing was holding it back. - assert not started.wait(0.2) - assert accountant.snapshot()["waiting"][0]["node"] == "study.b" - running.set() - - thread.join(timeout=5) - assert started.is_set() + assert accountant.try_take(cpus=2, gpus=0) is not None + assert accountant.try_take(cpus=2, gpus=0) is None def test_a_gpu_is_held_by_one_node_at_a_time(): """The deadlock was three processes each preallocating most of one card.""" accountant = ResourceAccountant(cpus=8, gpus=2) - with accountant.claim(Resources(cpus=1, gpus=1)) as first: - with accountant.claim(Resources(cpus=1, gpus=1)) as second: - assert set(first.gpus) & set(second.gpus) == set() - assert accountant.snapshot()["gpus"] == {"total": 2, "free": 0} + first = accountant.try_take(cpus=1, gpus=1) + second = accountant.try_take(cpus=1, gpus=1) + + assert set(first.gpus) & set(second.gpus) == set() + assert accountant.snapshot()["gpus"] == {"total": 2, "free": 0} + assert accountant.try_take(cpus=1, gpus=1) is None + + accountant.give_back(first) + accountant.give_back(second) assert accountant.snapshot()["gpus"] == {"total": 2, "free": 2} -def test_asking_for_more_than_there_is_gets_what_there_is(): - """A flow written on a big box still has to run on a laptop.""" - accountant = ResourceAccountant(cpus=2, gpus=0) - with accountant.claim(Resources(cpus=64, gpus=4)) as allocation: - assert allocation.cpus == 2 - assert allocation.gpus == () - - -def test_everything_comes_back_when_a_node_fails(): +def test_what_a_machine_could_ever_grant_is_a_different_question(): + """Busy is worth queueing for; too small is not, and reads the same.""" accountant = ResourceAccountant(cpus=4, gpus=1) - with pytest.raises(ValueError): - with accountant.claim(Resources(cpus=4, gpus=1)): - raise ValueError("the node raised") + assert accountant.fits(cpus=4, gpus=1) + assert not accountant.fits(cpus=8, gpus=0) + assert not accountant.fits(cpus=1, gpus=2) - assert accountant.snapshot()["cpus"]["free"] == 4 - assert accountant.snapshot()["gpus"]["free"] == 1 + accountant.try_take(cpus=4, gpus=1) + # Still true with nothing free: it is about the machine, not the moment. + assert accountant.fits(cpus=4, gpus=1) -def test_many_nodes_at_once_all_finish(): - """The accountant must not deadlock under contention; it is on every call.""" - accountant = ResourceAccountant(cpus=4) - done = [] +def test_a_machine_that_said_nothing_about_memory_is_not_a_machine_with_none(): + said = ResourceAccountant(cpus=4, ram_mb=2048) + assert said.fits(cpus=1, gpus=0, ram_mb=2048) + assert not said.fits(cpus=1, gpus=0, ram_mb=4096) + held = said.try_take(cpus=1, gpus=0, ram_mb=1536) + assert said.try_take(cpus=1, gpus=0, ram_mb=1024) is None + said.give_back(held) + assert said.snapshot()["ram_mb"] == {"total": 2048, "free": 2048} - def work() -> None: - with accountant.claim(Resources(cpus=2), node="study.n"): - time.sleep(0.01) - done.append(1) + quiet = ResourceAccountant(cpus=4) + assert quiet.fits(cpus=1, gpus=0, ram_mb=999_999) + assert quiet.try_take(cpus=1, gpus=0, ram_mb=999_999) is not None + assert quiet.snapshot()["ram_mb"] is None - threads = [threading.Thread(target=work) for _ in range(8)] - for thread in threads: - thread.start() - for thread in threads: - thread.join(timeout=10) - assert len(done) == 8 - assert accountant.snapshot()["cpus"]["free"] == 4 +def test_a_release_says_so_once_it_has_let_go(): + """The placer is told outside the lock, which is what keeps the order one-way.""" + seen: list[dict] = [] + accountant = ResourceAccountant(cpus=2) + accountant.on_release = lambda: seen.append(accountant.snapshot()["cpus"]) + + accountant.give_back(accountant.try_take(cpus=2, gpus=0)) + assert seen == [{"total": 2, "free": 2}] # ----------------------------------------------------------------------------- @@ -160,90 +150,3 @@ def test_declaring_nothing_stays_exactly_as_it_was(): def test_a_misspelled_resource_is_refused(): with pytest.raises(ValueError): Resources(cpu=4) - - -# ----------------------------------------------------------------------------- -# The whole path, through a real worker -# ----------------------------------------------------------------------------- - - -def test_a_declared_node_is_held_to_its_share(tmp_path): - """Four nodes wanting two cores each, on a machine with four. - - What is checked is the pair: never more than the machine has in flight at - once, and every one of them told what it was given — the two halves that - together are the oversubscription this is for. Also that it does not - deadlock, since the claim is taken before a worker slot and both block. - """ - import sys - - from fluksio.flow.controller import FlowController, RunContext - from fluksio.flow.messages import DType, MessageSpec - from fluksio.flow.schemas import FlowDef - from fluksio.flow.state import MemoryState - from fluksio.flow.store import FlowStore - from fluksio.flow.workers import PythonWorkerPool - - store = FlowStore(tmp_path / "flows") - store.write_flow( - FlowDef( - name="study", - mode="batch", - nodes=[ - NodeDef( - id="fit", - provides=[MessageSpec(name="threads", dtype=DType.INT)], - resources=Resources(cpus=2, env={"XLA_FLAGS": "--x=false"}), - ) - ], - ) - ) - store.write_node_source( - "study", - "fit", - "import os, time\n\n\ndef process():\n" - " time.sleep(0.2)\n" - " return {'threads': int(os.environ['OMP_NUM_THREADS'])}\n", - ) - - pool = PythonWorkerPool(python=sys.executable, size=4) - pool.start() - accountant = ResourceAccountant(cpus=4) - controller = FlowController(store, workers=pool, resources=accountant) - answers: list[object] = [] - in_flight: list[int] = [] - - def once(index: int) -> None: - pipeline = controller.build_run_pipeline( - store.read_flow("study"), - state=MemoryState(), - run=RunContext(run_id=f"r{index}"), - ) - pipeline.run({}) - answers.append(pipeline.state.get("study.threads")) - - def watch(until: threading.Event) -> None: - while not until.is_set(): - in_flight.append(4 - int(accountant.snapshot()["cpus"]["free"])) - time.sleep(0.01) - - finished = threading.Event() - watcher = threading.Thread(target=watch, args=(finished,)) - watcher.start() - threads = [threading.Thread(target=once, args=(i,)) for i in range(4)] - try: - for thread in threads: - thread.start() - for thread in threads: - thread.join(timeout=60) - assert not thread.is_alive(), "a claim and a worker slot deadlocked" - finally: - finished.set() - watcher.join(timeout=5) - pool.stop() - - assert answers == [2, 2, 2, 2] - assert max(in_flight) <= 4 - assert accountant.snapshot()["cpus"]["free"] == 4 - # One environment, so one extra pool however many nodes derived it. - assert len(pool._children) == 1 diff --git a/backend/tests/flow/test_runs.py b/backend/tests/flow/test_runs.py index 59aeacb..8d3ac4a 100644 --- a/backend/tests/flow/test_runs.py +++ b/backend/tests/flow/test_runs.py @@ -19,9 +19,10 @@ from fluksio.flow.runs import ( collect_result, digest_of, required_labels, + required_resources, seed_values, ) -from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef +from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef, Resources from fluksio.flow.state import MemoryState @@ -365,6 +366,32 @@ def test_a_preferred_device_does_not_hold_a_run_back(): assert required_labels(flow) == [] +def test_a_device_on_a_node_no_worker_could_run_is_not_waited_for(): + """A connector is an entry point loaded here, so a worker cannot take it. + + Holding the run for one is a run that never starts — and it looked exactly + like a GPU box somebody had forgotten to turn on. + """ + flow = double_flow() + flow.nodes[0].type = "mqtt" + flow.nodes[0].device = "gpu" + assert required_labels(flow) == [] + + +def test_what_a_run_needs_is_the_largest_single_machine_it_asks_for(): + flow = double_flow() + assert required_resources(flow) is None + + flow.nodes[0].resources = Resources(cpus=2, gpus=1) + flow.nodes.append(NodeDef(id="fit", resources=Resources(cpus=8))) + assert required_resources(flow) == { + "cpus": 8, + "gpus": 1, + "ram_mb": 0, + "device": "", + } + + def test_a_runs_seed_fills_an_input_of_that_name(): flow = double_flow() flow.inputs.append(FlowInput(spec=spec("seed", DType.INT), initial=0))