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
379 lines
14 KiB
Python
379 lines
14 KiB
Python
"""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
|