stroblmeandClaude Opus 5
40f8ad378d
Slurm is not a machine that attaches and stays; it is a queue somebody else owns. So nothing here submits a node to it. It submits a job whose payload is an ordinary worker dialling back in, and everything downstream — the protocol, the artifacts, cancellation, the books — already worked and did not have to learn what Slurm is. The alternative, which Covalent takes, is to stage a serialized call and a runner onto the login node, poll squeue and copy the result back: a second way of running a node beside the one that exists. The cost of not doing that is one assumption, that a compute node can open a connection outward. Where that is false, _payload is the single method a staged variant would replace. Clusters are configured in provisioners.json beside the alerts, since this is infrastructure an operator writes rather than anything a flow says. The script is generated with the system ssh and no new dependency, and prerun owns the environment — deliberately no pip install, because what is on a cluster is somebody's decision. One outstanding request per profile, cancelled if it never attaches and on the way out. Nothing autoscales. The run gate needed the same hook: a run held before it starts never reaches the placer's own wait, so it would have queued forever on a machine nothing had asked for. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
407 lines
16 KiB
Python
407 lines
16 KiB
Python
"""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 provision_for(self, needs: dict[str, Any] | None) -> None:
|
|
"""Ask for a machine a queued run is waiting on.
|
|
|
|
A run held before it starts never reaches the wait above, so this is
|
|
where the asking happens for it — otherwise it would sit waiting for a
|
|
machine that nothing ever requested.
|
|
"""
|
|
if not needs:
|
|
return
|
|
self._provision(
|
|
int(needs.get("cpus") or 1),
|
|
int(needs.get("gpus") or 0),
|
|
int(needs.get("ram_mb") or 0),
|
|
needs.get("device") or None,
|
|
)
|
|
|
|
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],
|
|
}
|