Fluksio/app
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Ask a cluster for a machine when nothing here will do

Browse Source 
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:
stroblme
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
Download Patch File Download Diff File Expand all files Collapse all files
backend/fluksio/flow/provision.py
+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