Ask a cluster for a machine when nothing here will do

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
This commit is contained in:
2026-08-27 09:09:44 +02:00
co-authored by Claude Opus 5
parent a82f88cf0a
commit 40f8ad378d
7 changed files with 680 additions and 0 deletions
+5
View File
@@ -24,6 +24,7 @@ DERIVED_PATHS = {
"PANELS_FILE": "panels.json",
"OAUTH_PRIVATE_KEY_FILE": "oauth-key.pem",
"CLOUD_CONFIG_FILE": "cloud.json",
"PROVISIONERS_FILE": "provisioners.json",
}
@@ -65,6 +66,10 @@ class Settings(BaseSettings):
# Which failures reach which channel. Beside the flows, not in them:
# alerting is the deployment's concern, not any one flow's.
ALERTS_FILE: Path = Path("flow-data/alerts.json")
# Where machines can be started from when a node needs one and nothing that
# could take it is attached. Operator-authored, like the alerts beside it,
# and absent on an installation that has nowhere to start one.
PROVISIONERS_FILE: Path = Path("flow-data/provisioners.json")
# Which dashboards each device shows. Beside the flows for the same reason
# alerting is: where a screen hangs is the deployment's concern rather than
# any one dashboard's.
+16
View File
@@ -297,6 +297,22 @@ class Placer:
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.
+378
View File
@@ -0,0 +1,378 @@
"""Machines started on demand, for a node that has nowhere to run.
A cluster is not a machine that attaches and stays. It is a queue somebody else
owns, and what you get from it is an allocation for as long as your job holds
one. So a provisioner does not run nodes: it *asks* for a machine, and what
comes back is an ordinary worker dialling in on the ordinary socket. Everything
downstream of that — the protocol, the artifacts, cancellation, the books —
already works, and none of it had to learn what Slurm is.
That is the whole design. The alternative, which Covalent takes, is to stage a
serialized call and a runner script onto the login node, submit *that*, poll
``squeue``, and copy the result back: a second way of running a node, beside
the one that already exists, with its own transport and its own failure modes.
The cost of doing it this way instead is one assumption — that a compute node
can open a connection outward — which is true of most clusters and false of
air-gapped ones. Where it is false, :meth:`SlurmProvisioner._payload` is the
one method a staged variant would need to replace.
Nothing here autoscales. One outstanding request per profile, dropped when the
machine attaches or when it has taken too long, and the job is cancelled on the
way out.
"""
from __future__ import annotations
import json
import logging
import subprocess
import threading
import time
import uuid
from dataclasses import dataclass, field
from datetime import timedelta
from pathlib import Path
from typing import Any, Protocol
from fluksio.core import security
from fluksio.flow.events import EventBus
logger = logging.getLogger(__name__)
#: How long a submission is given before it is treated as failed.
SUBMIT_TIMEOUT_S = 30.0
#: How long a worker's credential is good for. Short: it sits in a scheduler's
#: spool file, and the machine it is for is expected within the hour.
TOKEN_HOURS = 24
class Provisioner(Protocol):
"""Somewhere machines can be asked for.
Implemented once, for Slurm. The shape is deliberately small — asking is
not scheduling, and the placer already does the scheduling.
"""
name: str
def covers(
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
) -> bool:
"""Whether a machine this could start would fit this node."""
def provision(
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
) -> None:
"""Ask for one. Returns at once; the machine arrives by attaching."""
def shapes(self, label: str | None = None) -> list[tuple[int, int, int]]:
"""The sizes it can start, as cpus, gpus and MB."""
def reconcile(self, attached: set[str]) -> None:
"""Forget requests that arrived, and give up on ones that did not."""
def shutdown(self) -> None:
"""Cancel whatever is outstanding."""
def status(self) -> dict[str, Any]:
"""What it is called, what is outstanding, and what last went wrong."""
@dataclass(frozen=True)
class SlurmProfile:
"""One kind of machine this cluster can be asked for.
Named like a flavor and deliberately not the same thing: a flavor is what a
node asks for, a profile is what a scheduler is asked for. They agree when
somebody sets them up to.
"""
name: str
cpus: int = 1
gpus: int = 0
ram_mb: int = 0
#: What a worker started this way advertises, so a node bound to a device
#: can be the reason one is started.
labels: list[str] = field(default_factory=list)
#: ``#SBATCH`` lines, verbatim: partition, gres, walltime, account.
sbatch: list[str] = field(default_factory=list)
#: Shell run before the worker starts. Where the environment comes from —
#: `module load`, a venv with fluksio-worker already in it. Deliberately
#: not pip: what is installed on a cluster is somebody's decision, not this.
prerun: list[str] = field(default_factory=list)
@dataclass
class _Job:
job_id: str
worker: str
since: float
class SlurmProvisioner:
"""A Slurm cluster, asked for machines over ssh.
Submission is an ``sbatch`` whose payload is a worker dialling back here.
The engine needs no route to the compute node and no share of its
filesystem; what it needs is for the compute node to reach the engine.
"""
def __init__(
self,
name: str,
login: str,
engine_url: str,
profiles: list[SlurmProfile],
ssh_key: str = "",
artifact_url: str = "",
max_idle_s: float = 300.0,
provision_timeout_s: float = 900.0,
events: EventBus | None = None,
) -> None:
self.name = name
self.login = login
self.engine_url = engine_url
self.profiles = sorted(profiles, key=lambda p: (p.gpus, p.cpus, p.ram_mb))
self.ssh_key = ssh_key
self.artifact_url = artifact_url
self.max_idle_s = max_idle_s
self.provision_timeout_s = provision_timeout_s
self.events = events
self.last_error = ""
self._outstanding: dict[str, _Job] = {}
self._lock = threading.Lock()
# -- what it can be asked for ----------------------------------------------
def _matching(
self, cpus: int, gpus: int, ram_mb: int, label: str | None
) -> SlurmProfile | None:
"""The smallest profile that would fit, or None."""
for profile in self.profiles:
if label and label not in profile.labels and label != profile.name:
continue
if profile.cpus >= cpus and profile.gpus >= gpus:
if not ram_mb or not profile.ram_mb or profile.ram_mb >= ram_mb:
return profile
return None
def covers(
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
) -> bool:
return self._matching(cpus, gpus, ram_mb, label) is not None
def shapes(self, label: str | None = None) -> list[tuple[int, int, int]]:
return [
(profile.cpus, profile.gpus, profile.ram_mb)
for profile in self.profiles
if not label or label in profile.labels or label == profile.name
]
# -- asking ----------------------------------------------------------------
def provision(
self, cpus: int, gpus: int, ram_mb: int, label: str | None = None
) -> None:
profile = self._matching(cpus, gpus, ram_mb, label)
if profile is None:
return
with self._lock:
if profile.name in self._outstanding:
# One at a time per profile. The placer asks on every pass of
# its wait, and a queue full of jobs nobody needed is worse
# than a node waiting a little longer.
return
self._outstanding[profile.name] = _Job("", "", time.monotonic())
# Off the caller's thread: it is holding the placer's condition, and
# ssh to a login node is not something to hold a lock across.
threading.Thread(
target=self._submit,
args=(profile,),
name=f"provision-{self.name}",
daemon=True,
).start()
def _submit(self, profile: SlurmProfile) -> None:
worker = f"{self.name}-{profile.name}-{uuid.uuid4().hex[:8]}"
token = security.create_worker_token(worker, timedelta(hours=TOKEN_HOURS))
try:
result = self._ssh(["sbatch", "--parsable"], self._payload(profile, token))
except (OSError, subprocess.SubprocessError) as exc:
self._failed(profile, str(exc))
return
if result.returncode != 0:
self._failed(profile, (result.stderr or result.stdout).strip()[:500])
return
job_id = (result.stdout or "").strip().splitlines()[0] if result.stdout else ""
with self._lock:
self._outstanding[profile.name] = _Job(job_id, worker, time.monotonic())
logger.info(
"Asked %s for a '%s' machine: job %s, worker '%s'",
self.name,
profile.name,
job_id or "?",
worker,
)
self._publish("worker_provisioned", profile, job=job_id, worker=worker)
def _payload(self, profile: SlurmProfile, token: str) -> str:
"""The batch script: a worker that dials in, and stops when it is done.
A staged executor — for a cluster whose compute nodes cannot reach the
engine — is this method and nothing else.
"""
worker = [
"exec fluksio-worker",
f"--url {self.engine_url}",
f"--token {token}",
f"--cpus {profile.cpus}",
f"--gpus {profile.gpus}",
f"--max-idle {self.max_idle_s:.0f}",
]
if profile.ram_mb:
worker.append(f"--ram-mb {profile.ram_mb}")
if profile.labels:
worker.append(f"--labels {','.join(profile.labels)}")
if self.artifact_url:
worker.append(f"--artifact-url {self.artifact_url}")
lines = [
"#!/bin/bash",
f"#SBATCH --job-name=fluksio-{profile.name}",
*[f"#SBATCH {option}" for option in profile.sbatch],
"",
*profile.prerun,
" ".join(worker),
"",
]
return "\n".join(lines)
# -- keeping the outstanding list honest -----------------------------------
def reconcile(self, attached: set[str]) -> None:
expired = []
with self._lock:
for name, job in list(self._outstanding.items()):
if job.worker and job.worker in attached:
# It arrived. Asking again is somebody else's decision.
del self._outstanding[name]
elif time.monotonic() - job.since > self.provision_timeout_s:
del self._outstanding[name]
if job.job_id:
expired.append(job.job_id)
for job_id in expired:
logger.warning(
"%s: job %s never attached, cancelling it", self.name, job_id
)
self._cancel(job_id)
def shutdown(self) -> None:
with self._lock:
jobs = [job.job_id for job in self._outstanding.values() if job.job_id]
self._outstanding.clear()
for job_id in jobs:
self._cancel(job_id)
def _cancel(self, job_id: str) -> None:
try:
self._ssh(["scancel", job_id])
except (OSError, subprocess.SubprocessError) as exc:
logger.warning("%s: could not cancel job %s: %s", self.name, job_id, exc)
def status(self) -> dict[str, Any]:
with self._lock:
outstanding = [
{
"profile": name,
"job": job.job_id,
"worker": job.worker,
"seconds": round(time.monotonic() - job.since, 1),
}
for name, job in self._outstanding.items()
]
return {
"name": self.name,
"kind": "slurm",
"profiles": [profile.name for profile in self.profiles],
"outstanding": outstanding,
"last_error": self.last_error,
}
# -- the one thing that talks to the cluster -------------------------------
def _ssh(
self, command: list[str], stdin: str = ""
) -> subprocess.CompletedProcess[str]:
"""Run a command on the login node. The system ssh, and no new library."""
argv = ["ssh", "-o", "BatchMode=yes"]
if self.ssh_key:
argv += ["-i", self.ssh_key]
argv += [self.login, *command]
return subprocess.run(
argv,
input=stdin,
capture_output=True,
text=True,
timeout=SUBMIT_TIMEOUT_S,
check=False,
)
def _failed(self, profile: SlurmProfile, detail: str) -> None:
self.last_error = detail
with self._lock:
self._outstanding.pop(profile.name, None)
logger.error(
"%s: could not start a '%s' machine: %s", self.name, profile.name, detail
)
self._publish("worker_provision_failed", profile, detail=detail)
def _publish(self, kind: str, profile: SlurmProfile, **extra: Any) -> None:
if self.events is None:
return
self.events.publish(
{
"type": kind,
"provisioner": self.name,
"profile": profile.name,
"ts": time.time(),
**extra,
}
)
def load_provisioners(path: Path, events: EventBus | None = None) -> list[Any]:
"""Read the configured clusters. An installation with none has no file."""
if not path.exists():
return []
try:
entries = json.loads(path.read_text())
except (OSError, ValueError):
logger.exception("Could not read %s; no machines can be started", path)
return []
found = []
for entry in entries if isinstance(entries, list) else []:
kind = str(entry.get("type") or "slurm")
if kind != "slurm":
logger.error("%s: no provisioner of kind '%s'", path, kind)
continue
try:
found.append(
SlurmProvisioner(
name=str(entry["name"]),
login=str(entry["login"]),
engine_url=str(entry["engine_url"]),
profiles=[
SlurmProfile(**profile) for profile in entry.get("profiles", [])
],
ssh_key=str(entry.get("ssh_key") or ""),
artifact_url=str(entry.get("artifact_url") or ""),
max_idle_s=float(entry.get("max_idle_s", 300)),
provision_timeout_s=float(entry.get("provision_timeout_s", 900)),
events=events,
)
)
except (KeyError, TypeError, ValueError):
logger.exception("%s: could not read a provisioner", path)
if found:
logger.info("%d provisioner(s) can start machines", len(found))
return found
+4
View File
@@ -929,6 +929,10 @@ class RunService:
short = placer.satisfiable(needs)
if short:
missing.append(f"with {short}")
# Held here rather than at the placer's own wait, so this is
# where the machine has to be asked for. Asked again on every
# pass; one outstanding request is the provisioner's business.
placer.provision_for(needs)
return missing
def _waiting(self, run_id: str, missing: list[str]) -> None:
+7
View File
@@ -31,6 +31,7 @@ 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.provision import load_provisioners
from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue
from fluksio.flow.remote import RemoteWorkerHub
from fluksio.flow.resources import ResourceAccountant, fair_share_env
@@ -147,6 +148,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
app.state.resources = accountant
# Every machine a node could run on: this one, and whatever attaches.
placer = Placer(local=accountant, events=event_bus)
# Where more machines can be asked for when nothing attached will do.
placer.provisioners = load_provisioners(settings.PROVISIONERS_FILE, event_bus)
app.state.placer = placer
pool = PythonWorkerPool(
python=modules.venv_python(),
@@ -251,6 +254,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
await run_in_threadpool(run_service.stop)
await controller.stop()
pool.stop()
# A machine asked for and not yet arrived would hold an allocation
# nobody is going to use.
for provisioner in placer.provisioners:
await run_in_threadpool(provisioner.shutdown)
close_shared_client()
if settings.MCP_ENABLED:
from fluksio.mcp.http import aclose
+219
View File
@@ -0,0 +1,219 @@
"""Asking a batch scheduler for a machine, and what comes back.
The point of doing it this way is that almost nothing here is new: what a
cluster gives back is an ordinary worker on the ordinary socket. So what is
worth checking is the asking — that it happens once, that the script says what
it should, and that a job nobody ever attached does not sit in the queue.
"""
import subprocess
import threading
from fluksio.flow.placement import Placer
from fluksio.flow.provision import SlurmProfile, SlurmProvisioner, load_provisioners
from fluksio.flow.resources import ResourceAccountant
from fluksio.flow.schemas import Resources
class FakeSsh:
"""Stands in for the login node: records argv, answers like sbatch."""
def __init__(self, returncode: int = 0, stdout: str = "4711\n") -> None:
self.calls: list[tuple[list[str], str]] = []
self.returncode = returncode
self.stdout = stdout
self.ran = threading.Event()
def __call__(self, argv, **kwargs):
self.calls.append((argv, kwargs.get("input") or ""))
self.ran.set()
return subprocess.CompletedProcess(
argv, self.returncode, stdout=self.stdout, stderr="sbatch: bad partition"
)
def cluster(events=None, **kwargs) -> SlurmProvisioner:
return SlurmProvisioner(
name="hpc",
login="me@login.cluster",
engine_url="wss://engine.example.com/api/v1/workers/attach",
profiles=[
SlurmProfile(name="cpu", cpus=16, ram_mb=32768),
SlurmProfile(
name="gpu-small",
cpus=8,
gpus=1,
ram_mb=65536,
labels=["gpu"],
sbatch=["--partition=gpu", "--gres=gpu:1", "--time=04:00:00"],
prerun=["module load cuda/12", "source ~/venvs/flux/bin/activate"],
),
],
events=events,
**kwargs,
)
def test_the_smallest_machine_that_would_do_is_the_one_asked_for(monkeypatch):
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster()
hpc.provision(cpus=2, gpus=0, ram_mb=0)
assert ssh.ran.wait(5)
argv, script = ssh.calls[0]
assert argv[:3] == ["ssh", "-o", "BatchMode=yes"]
assert argv[-2:] == ["sbatch", "--parsable"]
assert "--job-name=fluksio-cpu" in script
def test_the_script_starts_a_worker_that_dials_back_and_stops_when_done(monkeypatch):
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster(max_idle_s=120)
hpc.provision(cpus=4, gpus=1, ram_mb=0, label="gpu")
assert ssh.ran.wait(5)
_, script = ssh.calls[0]
assert "#SBATCH --gres=gpu:1" in script
# The environment is the cluster's business, not this module's.
assert "module load cuda/12" in script
assert "pip install" not in script
assert "exec fluksio-worker" in script
assert "--url wss://engine.example.com/api/v1/workers/attach" in script
assert "--gpus 1" in script
assert "--labels gpu" in script
# Without this the job holds its allocation until the walltime.
assert "--max-idle 120" in script
def test_one_job_per_profile_however_often_it_is_asked_for(monkeypatch):
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster()
for _ in range(5):
hpc.provision(cpus=2, gpus=0, ram_mb=0)
assert ssh.ran.wait(5)
assert len(ssh.calls) == 1
assert hpc.status()["outstanding"][0]["job"] == "4711"
def test_a_machine_that_arrives_clears_the_way_for_the_next_ask(monkeypatch):
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster()
hpc.provision(cpus=2, gpus=0, ram_mb=0)
assert ssh.ran.wait(5)
arrived = hpc.status()["outstanding"][0]["worker"]
hpc.reconcile({arrived})
assert hpc.status()["outstanding"] == []
# No scancel: it did what it was asked to do.
assert [argv for argv, _ in ssh.calls if "scancel" in argv] == []
def test_a_job_that_never_attaches_is_cancelled(monkeypatch):
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster(provision_timeout_s=0)
hpc.provision(cpus=2, gpus=0, ram_mb=0)
assert ssh.ran.wait(5)
hpc.reconcile(set())
assert hpc.status()["outstanding"] == []
assert ["scancel", "4711"] == ssh.calls[-1][0][-2:]
def test_a_refused_submission_says_what_the_cluster_said(monkeypatch):
ssh = FakeSsh(returncode=1)
monkeypatch.setattr(subprocess, "run", ssh)
events: list[dict] = []
hpc = cluster(
events=type("Bus", (), {"publish": lambda self, e: events.append(e)})()
)
hpc.provision(cpus=2, gpus=0, ram_mb=0)
assert ssh.ran.wait(5)
for _ in range(100):
if events:
break
threading.Event().wait(0.02)
assert "bad partition" in hpc.last_error
assert events[-1]["type"] == "worker_provision_failed"
# And it is not left outstanding, so asking again actually asks.
assert hpc.status()["outstanding"] == []
def test_a_node_waiting_for_a_machine_asks_for_one_once(monkeypatch):
"""The placer asks on every pass of its wait; the cluster is asked once."""
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster()
placer = Placer(local=ResourceAccountant(cpus=4, gpus=0))
placer.provisioners = [hpc]
ran = threading.Event()
def wants_a_card() -> None:
with placer.claim(Resources(cpus=2, gpus=1), node="study.fit"):
ran.set()
thread = threading.Thread(target=wants_a_card, daemon=True)
thread.start()
# It waits rather than being cut down to what is here, because something
# can start a machine that would take it.
assert not ran.wait(0.5)
assert len(ssh.calls) == 1
placer.provisioners = []
placer.wake()
thread.join(timeout=5)
assert ran.is_set()
def test_an_installation_with_nowhere_to_start_one_has_no_file(tmp_path):
assert load_provisioners(tmp_path / "provisioners.json") == []
def test_a_configured_cluster_is_read(tmp_path):
path = tmp_path / "provisioners.json"
path.write_text(
"""[{"type": "slurm", "name": "hpc", "login": "me@login",
"engine_url": "wss://e/api/v1/workers/attach",
"profiles": [{"name": "gpu-small", "cpus": 8, "gpus": 1,
"labels": ["gpu"], "sbatch": ["--gres=gpu:1"]}]}]"""
)
found = load_provisioners(path)
assert [p.name for p in found] == ["hpc"]
assert found[0].covers(cpus=4, gpus=1, ram_mb=0, label="gpu")
assert not found[0].covers(cpus=4, gpus=2, ram_mb=0)
def test_a_file_that_cannot_be_read_is_not_a_failure_to_start(tmp_path):
path = tmp_path / "provisioners.json"
path.write_text("{ not json")
assert load_provisioners(path) == []
def test_a_run_held_before_it_starts_still_gets_a_machine_asked_for(monkeypatch):
"""The gate holds the run above the placer, so the asking has to happen there.
Without this the run waits on a machine nothing ever requested, which is a
queue entry that never moves and no error anywhere.
"""
ssh = FakeSsh()
monkeypatch.setattr(subprocess, "run", ssh)
hpc = cluster()
placer = Placer(local=ResourceAccountant(cpus=4, gpus=0))
placer.provisioners = [hpc]
needs = {"cpus": 2, "gpus": 1, "ram_mb": 0, "device": ""}
assert placer.satisfiable(needs) == "1 gpu(s) and 2 cpu(s)"
placer.provision_for(needs)
assert ssh.ran.wait(5)
assert "--gpus 1" in ssh.calls[0][1]